第三章 精度对齐工具链总览：PyTorch精度对比工具、ONNX Runtime验证、自定义C++测试框架搭建

精度对齐这件事，说白了就是让芯片算出来的结果和标准答案一模一样。我见过太多团队，模型在GPU上跑得飞起，一上芯片就崩了。嗯，这其实就是精度没对齐。

这一章，我带你把工具链捋一遍。我个人习惯把工具链分成三层：PyTorch层做初步验证，ONNX Runtime做中间检查，最后用C++框架做最终确认。三层都过了，基本就稳了。

3.1 PyTorch精度对比工具

PyTorch自带的对比工具其实够用。我刚开始做算子开发时，总喜欢自己写对比脚本，后来发现torch.allclose就挺好使。

import torch

def compare_tensors(golden, result, rtol=1e-5, atol=1e-8):
    """
    golden: 标准结果（GPU上跑的）
    result: 芯片上跑的结果
    """
    if not torch.allclose(golden, result, rtol=rtol, atol=atol):
        diff = torch.abs(golden - result)
        max_diff = torch.max(diff).item()
        mean_diff = torch.mean(diff).item()
        print(f"最大误差: {max_diff:.2e}")
        print(f"平均误差: {mean_diff:.2e}")
        
        # 定位异常位置
        outlier_mask = diff > atol
        outlier_ratio = torch.sum(outlier_mask).item() / diff.numel()
        print(f"异常点比例: {outlier_ratio:.2%}")
        return False
    return True

这里有个坑。我曾经遇到过一个问题：两个tensor看起来一模一样，但allclose就是返回False。后来发现是数据类型不一致——一个是float32，一个是float16。你想想看，精度都不一样，怎么比？

对于更细粒度的分析，我一般会用逐元素对比：

def detailed_comparison(golden, result):
    diff = golden - result
    stats = {
        'max_abs': torch.max(torch.abs(diff)).item(),
        'mean_abs': torch.mean(torch.abs(diff)).item(),
        'std': torch.std(diff).item(),
        'cosine_sim': torch.cosine_similarity(
            golden.flatten().unsqueeze(0), 
            result.flatten().unsqueeze(0)
        ).item()
    }
    return stats

为什么要看余弦相似度？因为有时候绝对值误差不大，但方向错了，这对后续计算影响很大。我在做卷积算子时遇到过这种情况，绝对值误差1e-6，但余弦相似度只有0.99，结果模型精度直接掉了2个点。

3.2 ONNX Runtime验证

PyTorch验证完了，下一步就是ONNX Runtime。为什么需要这一步？因为芯片的推理框架通常走ONNX路线，PyTorch的算子实现和ONNX Runtime的实现可能有差异。

我建议的做法是：先把模型导出成ONNX，然后用ONNX Runtime跑一遍，和PyTorch结果对比。

import onnxruntime as ort
import torch
import numpy as np

def validate_with_onnx(onnx_path, input_tensor, pytorch_output):
    # 创建ONNX Runtime session
    session = ort.InferenceSession(onnx_path)
    
    # 获取输入输出名称
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    
    # 运行推理
    onnx_output = session.run(
        [output_name], 
        {input_name: input_tensor.numpy()}
    )[0]
    
    # 转换为torch tensor
    onnx_tensor = torch.from_numpy(onnx_output)
    
    # 对比
    return compare_tensors(pytorch_output, onnx_tensor)

我记得有一次，一个LayerNorm算子，PyTorch和ONNX Runtime的误差到了1e-4。查了半天，发现是PyTorch的LayerNorm实现里用了快速近似算法，而ONNX Runtime用的是标准实现。最后我们决定以ONNX Runtime为准，因为芯片的算子库也是按标准实现的。

3.3 自定义C++测试框架搭建

到了这一步，才是真正的硬核环节。PyTorch和ONNX Runtime都只是参考，最终还是要用C++框架来验证芯片的算子。

我搭建C++测试框架时，一般包含这几个模块：

1. 数据生成器：生成各种形状、各种数据类型的输入
2. 参考实现：用CPU实现一个标准的算子版本
3. 芯片调用器：封装芯片的API调用
4. 对比器：支持多种精度标准对比

// 测试框架核心结构
struct TestCase {
    std::string name;
    std::vector<int> shape;
    DataType dtype;
    float rtol;
    float atol;
};

class PrecisionTestFramework {
public:
    // 注册测试用例
    void registerTestCase(const TestCase& tc);
    
    // 运行所有测试
    bool runAllTests();
    
    // 生成报告
    void generateReport();
    
private:
    // CPU参考实现
    Tensor cpuReference(const Tensor& input);
    
    // 芯片实现
    Tensor chipImplementation(const Tensor& input);
    
    // 精度对比
    bool compare(const Tensor& golden, const Tensor& result, 
                 float rtol, float atol);
};

这里有个设计要点：数据生成器要能生成边界值。比如对于float16，要能生成最大值、最小值、次正规数等。我曾经因为没测试次正规数，结果芯片上跑出来全是0，查了两天才发现是芯片的浮点单元对次正规数处理有问题。

对于对比器，我一般支持三种模式：

模式	适用场景	精度要求
严格模式	数学运算（加减乘除）	rtol=1e-7, atol=1e-10
宽松模式	激活函数、归一化	rtol=1e-5, atol=1e-7
统计模式	池化、采样	只看统计分布

为什么要分模式？因为不同算子的精度要求不一样。比如加法，误差应该极小；但Softmax这种，因为涉及指数运算，误差天然就大一些。

3.4 工具链的整体流程

我把整个流程画了张图，你看一眼就明白了：

你看，整个流程是逐层递进的。PyTorch验证过了，不代表ONNX Runtime能过；ONNX Runtime过了，不代表C++框架能过。每一层都会发现新的问题。

我个人习惯是：先在PyTorch层把大问题解决掉，再到ONNX Runtime层处理中间问题，最后在C++框架层做精细调优。这样效率最高，不会在底层被一些低级问题卡住。

想要更多资料，联系可可