共计 3226 个字符，预计需要花费 9 分钟才能阅读完成。

背景痛点

很多 AI 工程师在使用 NVIDIA 5070 显卡进行 fp16 半精度计算时，往往会遇到一些常见误区。首先是对精度损失的过度担忧，导致不敢大胆使用 fp16；其次是未能充分利用 tensor core 的计算能力，使得性能提升有限；还有就是对内存带宽优化不够重视，导致整体推理性能无法达到最佳状态。

技术对比

让我们先来看看 fp32、fp16 和 int8 三种精度在计算速度、内存占用和模型精度上的差异：

• fp32：32 位浮点，计算精度最高，但计算速度最慢，内存占用最大
• fp16：16 位浮点，计算精度适中，计算速度快，内存占用小，能充分发挥 tensor core 优势
• int8：8 位整数，计算速度最快，内存占用最小，但精度损失较大

在实际应用中，fp16 通常能在保持较好模型精度的同时，提供接近 int8 的计算速度，是一个很好的平衡点。

核心实现

1. PyTorch AMP 自动混合精度训练

PyTorch 的 AMP(Automatic Mixed Precision)功能可以自动管理 fp16 和 fp32 的转换，让训练过程既保持精度又提升速度。关键配置如下：

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()  # 用于防止梯度下溢

for epoch in range(epochs):
    for data, target in train_loader:
        optimizer.zero_grad()

        # 开启自动混合精度
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # 缩放损失并反向传播
        scaler.scale(loss).backward()
        # 更新权重
        scaler.step(optimizer)
        # 更新缩放因子
        scaler.update()

2. TensorRT 构建 fp16 引擎

使用 TensorRT 构建 fp16 引擎时，需要注意几个关键参数：

import tensorrt as trt

# 创建 builder 和 network
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 设置 fp16 模式
builder.fp16_mode = True
builder.strict_type_constraints = True  # 强制使用 fp16

# 构建引擎
engine = builder.build_cuda_engine(network)

3. CUDA 核函数优化

在编写 CUDA 核函数时，使用 half2 数据类型可以进一步提升性能：

__global__ void my_kernel(half* input, half* output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size/2) {half2 in = *reinterpret_cast<half2*>(&input[idx*2]);
        half2 out;
        out.x = __hadd(in.x, in.y);
        out.y = __hsub(in.x, in.y);
        *reinterpret_cast<half2*>(&output[idx*2]) = out;
    }
}

代码示例

PyTorch AMP 完整示例

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast

# 定义简单模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# 初始化
model = SimpleModel().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scaler = GradScaler()

# 训练循环
for epoch in range(10):
    for data, target in train_loader:
        data, target = data.cuda(), target.cuda()

        optimizer.zero_grad()

        # 混合精度上下文
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # 缩放梯度并反向传播
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

TensorRT fp16 引擎构建

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# ONNX 模型转换
def build_engine(onnx_path):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)

    # 显式 batch size
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)

    # 解析 ONNX 模型
    with open(onnx_path, 'rb') as model:
        if not parser.parse(model.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    # 配置 fp16
    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.FP16)

    # 构建引擎
    engine = builder.build_engine(network, config)
    return engine

Nsight Compute 性能分析

# 使用 Nsight Compute 分析 kernel 性能
nsys profile -o output_report ./your_inference_program

避坑指南

1. 动态 shape 处理：
2. 问题：TensorRT 对动态 shape 支持有限

3. 解决方案：在构建引擎时明确指定可能的 shape 范围

4. fp16 溢出检测：

5. 问题：fp16 范围较小容易溢出

6. 解决方案：使用梯度缩放 (GradScaler) 和数值范围检查

7. 层类型不支持：

8. 问题：某些层可能不支持 fp16
9. 解决方案：使用混合精度或强制指定某些层使用 fp32

性能验证

动手实验

现在，你可以尝试使用 NVIDIA Triton 部署一个 fp16 模型：

1. 将你的模型转换为 ONNX 格式
2. 使用上面提供的代码构建 TensorRT fp16 引擎
3. 配置 Triton 模型仓库
4. 启动 Triton 服务器
5. 使用客户端测试推理性能

期待你在实践中获得更好的性能表现！