已开启
clip_grad_norm_ 混合精度场景下 dtype 转换缺失和返回值为本地 norm 的 bug #64
qiang7680219创建于  3月31日
qiang7680219
3月31日 创建

clip_grad_norm_ 混合精度场景下 dtype 转换缺失和返回值为本地 norm 的 bug

该问题是怎么引起的?

hyper_parallel/platform/torch/clip_grad.py 存在两处 bug:

Bug 1:_clip_grads_with_norm_ grouped 路径缺少 dtype 转换

_group_tensors_by_device_and_dtype 按 (device, dtype) 分组后,主路径对
clip_coef_clamped 只做了 .to(device) 设备迁移,没有转换 dtype。当梯度为
fp16/bf16 时,fp16_tensor.mul_(fp32_coef) 触发 PyTorch 类型错误。而 fallback
路径正确使用了 .to(grad.device, grad.dtype),两条路径行为不一致。

Bug 2:返回值为本地 norm 而非全局 norm

有限 p-norm 场景下,_get_total_norm 返回 (total_norm, local_combined),原代码将
各 rank 的本地 local_combined(all-reduce 前)作为函数返回值,导致多卡下各 rank
返回值不同,与文档描述 "Returns: The total (unclipped) gradient norm" 不符,也与
inf/-inf/0 norm 路径返回全局 norm 的行为不一致。

重现步骤

Bug 1:

import torch                                              
from hyper_parallel.platform.torch.clip_grad import clip_grad_norm_
                                                                                                                                                               
# 混合精度:部分参数有 fp32 main_grad,部分参数 grad 为 fp16                                                                                                   
param_fp32 = torch.nn.Parameter(torch.randn(4, 4))                                                                                                             
param_fp32.main_grad = torch.randn(4, 4)  # fp32                                                                                                               
                                                                                                                                                               
param_fp16 = torch.nn.Parameter(torch.randn(4, 4))                                                                                                             
param_fp16.grad = torch.randn(4, 4).half()  # fp16                                                                                                             
                                                                                                                                                               
clip_grad_norm_([param_fp32, param_fp16], max_norm=1.0)                                                                                                        
                                                                                                                                                               
Bug 2(分布式环境):                                                                                                                                          
                                                          
# 8 卡 HSDP 训练                                                                                                                                               
norm = clip_grad_norm_(model.parameters(), max_norm=1.0)                                                                                                       
# norm 在不同 rank 上值不同,无法用于全局梯度监控                                                                                                              
print(f"rank {dist.get_rank()}: norm={norm.item()}")  # 各 rank 输出不一致                                                                                     
                                                                                                                                                               
报错信息                                                                                                                                                       
                                                                                                                                                               
Bug 1:                                                                                                                                                        
                                                          
RuntimeError: result type Float can't be cast to the desired output type Half                                                                                  
                                                                                                                                                               
发生位置:_clip_grads_with_norm_ 函数 g.mul_(clip_coef_clamped_device) 处。                                                                                    
                                                                                                                                                               
Bug 2:                                                                                                                                                        
                                                          
无崩溃报错,但行为错误:多卡训练时 clip_grad_norm_ 返回值在不同 rank 上不一致,                                                                                
无法作为全局梯度 norm 用于日志记录或自适应学习率调整。    
likedislike
Qqiang7680219
3月31日 issue类型由 Documentation 改变为 Bug