【Bug已解决】How to select half precision (BFLOAT16 vs FLOAT16) for your trained model? 解决方案

问题描述

在深度学习训练和推理中，半精度浮点数（Half Precision）已经成为加速计算、减少内存占用的关键技术。PyTorch 支持两种半精度格式：FLOAT16（FP16）和BFLOAT16（BF16）。许多开发者在选择使用哪种格式时感到困惑，经常遇到以下问题：

• 使用 FP16 训练时出现 NaN（Not a Number）或 Inf（Infinity）损失值
• 梯度下溢（gradient underflow）导致训练无法收敛
• 使用 FP16 时需要复杂的梯度缩放（GradScaler），而 BF16 不需要
• 推理时模型输出精度下降
• 不同 GPU 架构对 FP16 和 BF16 的支持程度不同
• 混合精度训练（AMP）中应该选择哪种格式

这些问题的根源在于 FP16 和 BF16 在浮点数表示上有本质区别，各有优劣。选择错误的格式可能导致训练不稳定、精度下降或性能没有提升。本文将深入分析两种格式的区别，并提供完整的选择指南和代码实现。

错误复现

错误场景一：FP16 训练出现 NaN

import torch import torch.nn as nn import torch.optim as optim # 创建一个较大的模型 model = nn.Sequential( nn.Linear(1000, 4096), nn.ReLU(), nn.Linear(4096, 4096), nn.ReLU(), nn.Linear(4096, 1000) ).cuda() # 错误：直接使用 FP16 而不使用梯度缩放 model = model.half() # 将模型转为 FP16 optimizer = optim.Adam(model.parameters(), lr=0.001) criterion = nn.MSELoss() # 训练数据 x = torch.randn(32, 1000).cuda().half() y = torch.randn(32, 1000).cuda().half() for epoch in range(10): optimizer.zero_grad() output = model(x) loss = criterion(output, y) loss.backward() optimizer.step() if torch.isnan(loss) or torch.isinf(loss): print(f"Epoch {epoch}: Loss is NaN/Inf! (loss={loss.item()})") break print(f"Epoch {epoch}: Loss = {loss.item():.6f}")

运行结果：

Epoch 0: Loss = 1.001953 Epoch 1: Loss = nan

FP16 的动态范围有限（最大值约 65504，最小正常值约 6e-8），梯度在反向传播中很容易下溢到零，导致训练崩溃。

错误场景二：FP16 推理精度下降

import torch import torch.nn as nn model = nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ).cuda() # 先在 FP32 下训练 x_train = torch.randn(1000, 100).cuda() y_train = torch.randint(0, 10, (1000,)).cuda() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) for epoch in range(50): optimizer.zero_grad() output = model(x_train) loss = criterion(output, y_train) loss.backward() optimizer.step() # FP32 推理 model.eval() with torch.no_grad(): fp32_output = model(x_train[:5]) print(f"FP32 output: {fp32_output[0]}") # 转为 FP16 推理 model_half = model.half() x_half = x_train[:5].half() with torch.no_grad(): fp16_output = model_half(x_half) print(f"FP16 output: {fp16_output[0]}") # 对比差异 diff = (fp32_output.float() - fp16_output.float()).abs().mean() print(f"\nMean absolute difference: {diff.item():.6f}")

错误场景三：在不支持 BF16 的 GPU 上使用 BF16

import torch # 在较老的 GPU（如 GTX 1080, V100）上尝试使用 BF16 try: x = torch.randn(10, 10, dtype=torch.bfloat16, device='cuda') y = torch.randn(10, 10, dtype=torch.bfloat16, device='cuda') z = torch.mm(x, y) print(f"BF16 supported, result shape: {z.shape}") except RuntimeError as e: print(f"Error: {e}") # 可能的输出: "not implemented for 'BFloat16'" 或设备不支持

根因分析

1. FP16 与 BF16 的位结构对比

浮点数由三部分组成：符号位（Sign）、指数位（Exponent）、尾数位（Mantissa/Fraction）。

关键区别：

• FP16：与 FP32 相比，指数位从 8 减到 5（动态范围大幅缩小），尾数位从 23 减到 10（精度降低但尚可接受）。
• BF16：与 FP32 相比，指数位保持 8 不变（动态范围相同），尾数位从 23 减到 7（精度降低较多）。

2. 动态范围对比

FP16 的核心问题：动态范围只有约 9 个数量级（从 6e-8 到 65504），而深度学习中的梯度值经常跨越 10+ 个数量级。小梯度会下溢为零（丢失信息），大激活值会上溢为 Inf（导致 NaN）。

BF16 的核心优势：与 FP32 相同的动态范围（约 80 个数量级），不会出现 FP16 的上溢/下溢问题。但精度较低（只有 7 位尾数），在某些对精度敏感的操作中可能不够。

3. GPU 硬件支持

4. 混合精度训练（AMP）的工作原理

PyTorch 的自动混合精度（AMP）通过torch.cuda.amp实现，有两种模式：

FP16 AMP（传统模式）：

• 需要GradScaler来缩放梯度，防止小梯度下溢
• 在反向传播前将 loss 乘以一个大的缩放因子
• 在优化器更新前将梯度除以缩放因子
• 如果检测到 Inf/NaN，跳过该步并减小缩放因子

BF16 AMP（新模式）：

• 不需要GradScaler，因为 BF16 的动态范围与 FP32 相同
• 代码更简洁，训练更稳定
• 但需要 Ampere 或更新的 GPU

5. 选择决策矩阵

解决方案

方案一：使用 BF16 混合精度训练（推荐，Ampere+ GPU）

import torch import torch.nn as nn import torch.optim as optim from torch.cuda.amp import autocast class SimpleModel(nn.Module): def __init__(self, input_dim=1000, hidden_dim=4096, output_dim=1000): super().__init__() self.network = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_dim, output_dim) ) def forward(self, x): return self.network(x) def train_with_bf16(): """使用 BF16 混合精度训练""" device = torch.device('cuda') model = SimpleModel().to(device) optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) criterion = nn.MSELoss() # 生成数据 x = torch.randn(64, 1000, device=device) y = torch.randn(64, 1000, device=device) # BF16 AMP：不需要 GradScaler！ for epoch in range(20): optimizer.zero_grad() # 使用 autocast 自动选择哪些操作用 BF16 with autocast(dtype=torch.bfloat16): output = model(x) loss = criterion(output, y) # 直接反向传播，无需缩放 loss.backward() # 梯度裁剪（BF16 下仍然推荐） torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() if torch.isnan(loss) or torch.isinf(loss): print(f"Epoch {epoch}: NaN/Inf detected!") break print(f"Epoch {epoch}: Loss = {loss.item():.6f}") # 检查 GPU 是否支持 BF16 if torch.cuda.is_available(): gpu_name = torch.cuda.get_device_name() major, minor = torch.cuda.get_device_capability() print(f"GPU: {gpu_name}, Compute Capability: {major}.{minor}") if major >= 8: # Ampere 及以上 print("BF16 is supported!") train_with_bf16() else: print("BF16 not supported, use FP16 instead.")

方案二：使用 FP16 混合精度训练（Volta/Turing GPU）

import torch import torch.nn as nn import torch.optim as optim from torch.cuda.amp import autocast, GradScaler def train_with_fp16(): """使用 FP16 混合精度训练，需要 GradScaler""" device = torch.device('cuda') model = SimpleModel().to(device) optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) criterion = nn.MSELoss() # 关键：创建 GradScaler scaler = GradScaler() x = torch.randn(64, 1000, device=device) y = torch.randn(64, 1000, device=device) for epoch in range(20): optimizer.zero_grad() # FP16 AMP with autocast(dtype=torch.float16): output = model(x) loss = criterion(output, y) # 使用 scaler 缩放 loss 并反向传播 scaler.scale(loss).backward() # 梯度裁剪（需要在 unscale 之前或使用特殊方式） scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 更新参数 scaler.step(optimizer) scaler.update() if torch.isnan(loss) or torch.isinf(loss): print(f"Epoch {epoch}: NaN/Inf detected! Scaler may adjust scale.") continue print(f"Epoch {epoch}: Loss = {loss.item():.6f}, Scale = {scaler.get_scale()}") train_with_fp16()

方案三：BF16 与 FP16 性能和精度对比

import torch import torch.nn as nn import time import numpy as np def compare_precision_formats(): """对比 FP32, FP16, BF16 的性能和精度""" device = torch.device('cuda') # 创建大模型 model = nn.Sequential( nn.Linear(2048, 8192), nn.ReLU(), nn.Linear(8192, 8192), nn.ReLU(), nn.Linear(8192, 2048) ).to(device) x = torch.randn(256, 2048, device=device) # FP32 基准 model_fp32 = model.float() model_fp32.eval() with torch.no_grad(): start = time.time() for _ in range(100): out_fp32 = model_fp32(x) torch.cuda.synchronize() fp32_time = (time.time() - start) / 100 # FP16 model_fp16 = model.half() x_fp16 = x.half() model_fp16.eval() with torch.no_grad(): start = time.time() for _ in range(100): out_fp16 = model_fp16(x_fp16) torch.cuda.synchronize() fp16_time = (time.time() - start) / 100 # BF16 model_bf16 = model.bfloat16() x_bf16 = x.bfloat16() model_bf16.eval() with torch.no_grad(): start = time.time() for _ in range(100): out_bf16 = model_bf16(x_bf16) torch.cuda.synchronize() bf16_time = (time.time() - start) / 100 # 精度对比 fp16_diff = (out_fp32.float() - out_fp16.float()).abs().mean().item() bf16_diff = (out_fp32.float() - out_bf16.float()).abs().mean().item() print(f"{'Format':<10} {'Time (ms)':<15} {'Speedup':<10} {'MAE vs FP32':<15}") print(f"{'FP32':<10} {fp32_time*1000:<15.2f} {'1.00x':<10} {'0.000000':<15}") print(f"{'FP16':<10} {fp16_time*1000:<15.2f} {fp32_time/fp16_time:.2f}x{'':<6} {fp16_diff:<15.8f}") print(f"{'BF16':<10} {bf16_time*1000:<15.2f} {fp32_time/bf16_time:.2f}x{'':<6} {bf16_diff:<15.8f}") compare_precision_formats()

完整修复代码

以下是一个完整的、可运行的半精度训练框架，自动检测硬件支持并选择最佳格式：

import torch import torch.nn as nn import torch.optim as optim from torch.cuda.amp import autocast, GradScaler from torch.utils.data import Dataset, DataLoader import numpy as np import time import os # ============================================================ # 1. 数据集 # ============================================================ class RegressionDataset(Dataset): """回归任务数据集""" def __init__(self, num_samples=10000, input_dim=512, output_dim=256): self.data = torch.randn(num_samples, input_dim) # 创建有意义的标签（非线性变换） w = torch.randn(input_dim, output_dim) self.labels = torch.mm(self.data, w) + torch.randn(num_samples, output_dim) * 0.1 def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.labels[idx] # ============================================================ # 2. 模型 # ============================================================ class MLPModel(nn.Module): def __init__(self, input_dim=512, hidden_dim=2048, output_dim=256): super().__init__() self.layers = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim, output_dim) ) def forward(self, x): return self.layers(x) # ============================================================ # 3. 半精度训练管理器 # ============================================================ class HalfPrecisionTrainer: """ 自动选择最佳半精度格式的训练管理器。 - 自动检测 GPU 是否支持 BF16 - 自动选择 FP16+GradScaler 或 BF16 - 提供统一的训练接口 """ def __init__(self, model, optimizer, criterion, device='cuda', preferred_dtype=None): self.model = model.to(device) self.optimizer = optimizer self.criterion = criterion self.device = device # 自动检测最佳格式 self.dtype, self.use_scaler, self.scaler = self._select_dtype(preferred_dtype) print(f"Half Precision Configuration:") print(f" Selected dtype: {self.dtype}") print(f" Using GradScaler: {self.use_scaler}") if self.use_scaler: print(f" Initial scale: {self.scaler.get_scale()}") def _select_dtype(self, preferred=None): """选择最佳半精度格式""" if preferred: if preferred == torch.bfloat16: return torch.bfloat16, False, None elif preferred == torch.float16: return torch.float16, True, GradScaler() # 自动检测 if torch.cuda.is_available(): major, minor = torch.cuda.get_device_capability() if major >= 8: # Ampere+ return torch.bfloat16, False, None else: return torch.float16, True, GradScaler() else: print("Warning: CUDA not available, using FP32") return torch.float32, False, None def train_step(self, batch_x, batch_y): """执行一个训练步骤""" batch_x = batch_x.to(self.device) batch_y = batch_y.to(self.device) self.optimizer.zero_grad() # 混合精度前向传播 with autocast(dtype=self.dtype): output = self.model(batch_x) loss = self.criterion(output, batch_y) if self.use_scaler: # FP16: 使用 GradScaler self.scaler.scale(loss).backward() self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.scaler.step(self.optimizer) self.scaler.update() else: # BF16: 直接反向传播 loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.optimizer.step() return loss.item() def eval_step(self, batch_x, batch_y): """执行一个评估步骤""" batch_x = batch_x.to(self.device) batch_y = batch_y.to(self.device) with torch.no_grad(): with autocast(dtype=self.dtype): output = self.model(batch_x) loss = self.criterion(output, batch_y) return loss.item() def train(self, train_loader, val_loader, num_epochs=20): """完整训练流程""" best_val_loss = float('inf') for epoch in range(num_epochs): # 训练 self.model.train() train_losses = [] epoch_start = time.time() for batch_x, batch_y in train_loader: loss = self.train_step(batch_x, batch_y) train_losses.append(loss) train_loss = np.mean(train_losses) # 验证 self.model.eval() val_losses = [] with torch.no_grad(): for batch_x, batch_y in val_loader: loss = self.eval_step(batch_x, batch_y) val_losses.append(loss) val_loss = np.mean(val_losses) epoch_time = time.time() - epoch_start # 保存最佳模型 if val_loss < best_val_loss: best_val_loss = val_loss # 保存为 FP32 以保证兼容性 torch.save({ 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'val_loss': val_loss, 'dtype': str(self.dtype), }, 'best_model_half_precision.pth') print(f"Epoch [{epoch+1}/{num_epochs}] ({epoch_time:.1f}s) " f"Train: {train_loss:.6f} | Val: {val_loss:.6f} | " f"Best: {best_val_loss:.6f}") return best_val_loss # ============================================================ # 4. 推理时的半精度管理 # ============================================================ class HalfPrecisionInference: """ 半精度推理管理器。 支持在 FP16/BF16/FP32 之间切换并对比。 """ def __init__(self, model_path, model_class, model_kwargs, device='cuda'): self.device = device self.model_path = model_path self.model_class = model_class self.model_kwargs = model_kwargs def load_model(self, dtype=torch.float32): """加载模型并转换为指定精度""" model = self.model_class(**self.model_kwargs).to(self.device) checkpoint = torch.load(self.model_path, map_location=self.device) model.load_state_dict(checkpoint['model_state_dict']) model.eval() if dtype != torch.float32: model = model.to(dtype) return model def benchmark(self, test_data, num_runs=100): """对比不同精度的推理性能""" results = {} for dtype_name, dtype in [("FP32", torch.float32), ("FP16", torch.float16), ("BF16", torch.bfloat16)]: try: model = self.load_model(dtype) x = test_data.to(self.device).to(dtype) # 预热 with torch.no_grad(): for _ in range(10): _ = model(x) torch.cuda.synchronize() # 正式测试 start = time.time() with torch.no_grad(): for _ in range(num_runs): output = model(x) torch.cuda.synchronize() elapsed = (time.time() - start) / num_runs # 精度对比（与 FP32 对比） if dtype_name == "FP32": fp32_output = output.float() mae = 0.0 else: mae = (fp32_output - output.float()).abs().mean().item() results[dtype_name] = { 'time_ms': elapsed * 1000, 'mae': mae, 'output_shape': output.shape } print(f"{dtype_name}: {elapsed*1000:.2f}ms/run, MAE vs FP32: {mae:.8f}") except RuntimeError as e: print(f"{dtype_name}: Not supported - {e}") results[dtype_name] = None return results # ============================================================ # 5. 主程序 # ============================================================ def main(): torch.manual_seed(42) np.random.seed(42) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.cuda.is_available(): gpu_name = torch.cuda.get_device_name() major, minor = torch.cuda.get_device_capability() print(f"GPU: {gpu_name}") print(f"Compute Capability: {major}.{minor}") print(f"BF16 Support: {'Yes' if major >= 8 else 'No'}") print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB") print() # 超参数 INPUT_DIM = 512 HIDDEN_DIM = 2048 OUTPUT_DIM = 256 BATCH_SIZE = 128 NUM_EPOCHS = 15 LR = 1e-3 # 数据 full_dataset = RegressionDataset(10000, INPUT_DIM, OUTPUT_DIM) train_size = int(0.8 * len(full_dataset)) val_size = len(full_dataset) - train_size train_dataset, val_dataset = torch.utils.data.random_split( full_dataset, [train_size, val_size] ) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4) # 模型 model = MLPModel(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM) optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4) criterion = nn.MSELoss() print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}\n") # 训练 print("=" * 60) print("Starting Half Precision Training") print("=" * 60) trainer = HalfPrecisionTrainer(model, optimizer, criterion, device) best_loss = trainer.train(train_loader, val_loader, NUM_EPOCHS) print(f"\nBest validation loss: {best_loss:.6f}") # 推理基准测试 print("\n" + "=" * 60) print("Inference Benchmark") print("=" * 60) test_data = torch.randn(64, INPUT_DIM) inferencer = HalfPrecisionInference( 'best_model_half_precision.pth', MLPModel, {'input_dim': INPUT_DIM, 'hidden_dim': HIDDEN_DIM, 'output_dim': OUTPUT_DIM}, device ) inferencer.benchmark(test_data) # 内存占用对比 print("\n" + "=" * 60) print("Memory Usage Comparison") print("=" * 60) for dtype_name, dtype in [("FP32", torch.float32), ("FP16", torch.float16), ("BF16", torch.bfloat16)]: try: model = MLPModel(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM).to(device).to(dtype) param_memory = sum(p.nelement() * p.element_size() for p in model.parameters()) print(f"{dtype_name}: {param_memory / 1024 / 1024:.2f} MB") del model torch.cuda.empty_cache() except: print(f"{dtype_name}: Not supported") if __name__ == '__main__': main()

常见陷阱与注意事项

1. BF16 的精度较低

BF16 只有 7 位尾数（约 2.4 位有效十进制数字），而 FP16 有 10 位尾数（约 3.3 位有效十进制数字）。在对精度敏感的操作中（如累加、小数值计算），BF16 的误差可能比 FP16 更大。

# 累加精度对比 x = torch.ones(10000, dtype=torch.float32) print(f"FP32 sum: {x.sum().item()}") # 10000.0 x_fp16 = x.half() print(f"FP16 sum: {x_fp16.sum().item()}") # 接近 10000.0 x_bf16 = x.bfloat16() print(f"BF16 sum: {x_bf16.sum().item()}") # 误差可能更大

2. 某些操作不支持半精度

并非所有 PyTorch 操作都支持 FP16/BF16。以下操作通常需要 FP32：

• 损失函数中的数值稳定操作
• BatchNorm/LayerNorm 的统计量计算
• 累加操作
• 某些自定义操作

autocast会自动处理这些情况，但如果手动转换精度，需要注意。

3. 模型保存和加载

保存模型时建议使用 FP32，以保证兼容性和精度：

# 保存时转回 FP32 torch.save(model.float().state_dict(), 'model.pth') # 加载后再转换 model.load_state_dict(torch.load('model.pth')) model = model.bfloat16() # 或 .half()

4. GradScaler 的动态调整

FP16 训练中，GradScaler会动态调整缩放因子。如果频繁出现 Inf/NaN，缩放因子会自动减小；如果训练稳定，缩放因子会逐渐增大。不要手动设置缩放因子，除非你确切知道自己在做什么。

5. 分布式训练中的半精度

在分布式训练（DDP）中，所有进程必须使用相同的精度格式。梯度同步在 FP32 下进行（即使训练使用半精度），所以不需要额外处理。

6. ONNX 导出

导出 ONNX 模型时，BF16 的支持不如 FP16 广泛。如果需要部署到 ONNX Runtime 或 TensorRT，建议使用 FP16。

7. CPU 上的半精度

PyTorch 在 CPU 上对 BF16 的支持较好（从 PyTorch 1.12 开始），但 FP16 的支持有限。在 CPU 上推理时，BF16 通常是更好的选择。

8. 不要将优化器状态转为半精度

优化器状态（如 Adam 的动量和方差）应始终保持 FP32。autocast会自动处理这一点，但如果手动管理精度，需要确保优化器状态不被转换。

9. 学习率可能需要调整

半精度训练有时需要不同的学习率。BF16 训练通常可以使用与 FP32 相同的学习率，但 FP16 训练可能需要更小的学习率以避免数值不稳定。

10. 检查 Tensor Core 利用率

半精度的性能提升主要来自 Tensor Core。确保你的张量维度是 8 的倍数（Ampere 及以上）或 16 的倍数（Volta/Turing），以充分利用 Tensor Core。

# 好的维度 linear = nn.Linear(1024, 2048) # 都是 8/16 的倍数 # 不好的维度（Tensor Core 无法利用） linear = nn.Linear(1000, 2047) # 不是 8/16 的倍数

总结

选择 BFLOAT16 还是 FLOAT16 取决于你的具体场景和硬件：

选择 BF16 的场景：

• 使用 Ampere 或更新的 GPU（A100, RTX 30xx+, H100）
• 训练大型模型（如 LLM），梯度范围大
• 希望简化代码（不需要 GradScaler）
• 训练稳定性是首要考虑

选择 FP16 的场景：

• 使用 Volta 或 Turing GPU（V100, RTX 20xx）
• 推理任务，对精度要求较高
• 部署到边缘设备或 ONNX/TensorRT
• 需要更广泛的硬件兼容性

最佳实践：

1. 优先使用autocast自动混合精度，而非手动转换
2. 训练时保存 FP32 权重，推理时再转换
3. 确保张量维度是 8 或 16 的倍数以利用 Tensor Core
4. 使用 PyTorch Profiler 验证半精度是否真正加速了你的模型
5. 在切换精度格式后，重新调整学习率和超参数
6. 监控训练过程中的 loss 曲线，及时发现数值不稳定

通过理解 FP16 和 BF16 的本质区别，并根据硬件能力和任务需求做出正确选择，你可以充分利用半精度计算的优势，同时避免常见的数值稳定性问题。