执行摘要
- 一句话:修正 grad_norm 日志为裁剪前范数
- 推荐动作:值得快速浏览,无需深入精读。可关注两点:Megatron clip_grad_norm() 的返回值语义(先计算后缩放),以及日志指标文档与 UI 提示保持一致性的维护方式。
功能与动机
PR body 指出,日志中的 train/grad_norm 来自 Megatron optimizer.step() 中的 clip_grad_norm(),该函数先计算全局 L2 范数并返回,再原地缩放梯度,因此记录值永远不会被裁剪修改。原 docstring 声称是 after clipping,会误导用户以为范数不会超过 --clip-grad,从而忽略裁剪每步都在触发的事实。dashboard tooltip 已正确描述为 before clipping,docstring 是仓库唯一错误表述处。
实现拆解
- 定位错误表述:在 miles/backends/training_utils/log_utils.py 的 log_train_step docstring Args 块中,grad_norm 描述为 "Gradient norm after clipping"。
- 修正描述:改为 "Global gradient L2 norm before clipping",直接对应 Megatron clip_grad_norm() 先计算、后缩放的语义。
- 收尾调整:第二个提交删除提交信息中关于 multi-LoRA 的额外说明,保持 docstring 聚焦“裁剪前”这一单一事实,避免描述仍在重构的代码路径。
- 验证与配套:纯 docstring/注释修改,无运行时行为变化,无需测试变更;作者核对 Megatron 优化器源码确认语义,并确认这是仓库中唯一错误表述。
关键文件:
miles/backends/training_utils/log_utils.py(模块 训练日志;类别 source;类型 documentation;符号 log_train_step): 修改 log_train_step 的 docstring,修正 grad_norm 描述为裁剪前的全局 L2 范数,是本次唯一变更文件。
关键符号:log_train_step
关键源码片段
miles/backends/training_utils/log_utils.py
修改 log_train_step 的 docstring,修正 grad_norm 描述为裁剪前的全局 L2 范数,是本次唯一变更文件。
# 文件:miles/backends/training_utils/log_utils.py
# 本次修改仅涉及 log_train_step 的 docstring 中对 grad_norm 的描述。
# 原描述为 “Gradient norm after clipping”,现修正为 “Global gradient L2 norm before clipping”,
# 因为 Megatron optimizer.step() 中的 clip_grad_norm() 先计算全局 L2 范数并返回,
# 再原地缩放梯度,所以此处收到的值从未被裁剪修改。
def log_train_step(
args,
loss_dict,
grad_norm,
rollout_id,
step_id,
num_steps_per_rollout,
role='actor',
extra_metrics=None,
should_log=None,
):
'''Log training metrics for one step.
Args:
args: Configuration.
loss_dict: Dictionary of loss metrics from aggregate_train_losses.
grad_norm: Global gradient L2 norm before clipping. # 修正后的表述
rollout_id: Rollout ID.
step_id: Step ID within the rollout.
num_steps_per_rollout: Total number of steps per rollout.
role: Role name (e.g., 'actor', 'critic').
extra_metrics: Optional extra metrics to log (e.g., learning rates, MTP loss).
should_log: Optional override for logging condition. If None, uses rank == 0.
Returns:
The formatted log_dict (for CI tests or other uses).
'''
# 组装日志字典,grad_norm 直接透传调用方传入的值
accumulated_step_id = rollout_id * num_steps_per_rollout + step_id
role_tag = '' if role == 'actor' else f'{role}-'
log_dict_out = {
f'train/{role_tag}{key}': val.mean().item() if isinstance(val, torch.Tensor) else val
for key, val in loss_dict.items()
}
# 这里记录的 grad_norm 就是 Megatron clip_grad_norm() 返回的“裁剪前”值
log_dict_out[f'train/{role_tag}grad_norm'] = float(grad_norm)
# ...(其余日志组装与输出逻辑不变)
评论区精华
无实质技术讨论。唯一评论来自 gemini-code-assist[bot],提示其消费版审查活动已停止;Zhichenzzz 直接批准,未提出修改意见。
- docstring 修正获批 (other): PR 以纯文档修正合并,无未决问题。
风险与影响
- 风险:风险极低:改动仅限 docstring,不影响任何运行时逻辑。对用户的影响是消除了 train/grad_norm 指标解读的歧义;对系统无影响。注意提交信息中提到的 multi-LoRA 场景(调用方传入各 adapter slot 的最大范数)未写入 docstring,待 multi-LoRA 重构完成后可能需要补充。
- 影响:影响面很小:仅影响阅读代码与日志的用户,帮助他们正确理解 train/grad_norm 是裁剪前的全局 L2 范数。对训练行为、系统性能、配置均无影响,团队内部观测指标语义得到统一。
- 风险标记:纯文档变更, 无运行时影响
关联脉络
- PR #1965 dashboard: fix phase visibility for manager events and idle processes: 同一训练指标观测链路:dashboard 侧 tooltip 已正确描述 grad_norm 为裁剪前,本次修正确认并统一了该指标语义。
参与讨论