refactor : grad_norm 指标简化,clip_grad_norm 移至 executor

- metrics 默认加入 grad_norm,移除 grad_std/max/min/mean/nan_num
- grad_norm 默认返回总 L2 范数,per_param=True 返回各参数范数
- clip_grad_norm 从 callback 移至 BaseExecutor/FSDPExecutor
- FSDPExecutor 覆盖为 model.clip_grad_norm_() 保证分布式正确
- ctx_get_grad_norm 改为读取 context.grad_norm
This commit is contained in:
2026-06-30 14:59:43 +08:00
parent 84d4769163
commit 0f1fcb079f
6 changed files with 38 additions and 69 deletions
+14
View File
@@ -132,6 +132,12 @@ class BaseExecutor:
def grad_accum_steps(self) -> int:
return self.gradient_state.num_steps
def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
if isinstance(total_norm, torch.Tensor):
return total_norm.item()
return total_norm
class ExecutorFactory(BaseFactory[BaseExecutor]):
pass
@@ -260,6 +266,14 @@ class FSDPExecutor(BaseExecutor):
return model.no_sync()
return contextlib.nullcontext()
def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
if isinstance(model, FSDP) and self.use_distributed:
total_norm = model.clip_grad_norm_(max_norm)
if isinstance(total_norm, torch.Tensor):
return total_norm.item()
return total_norm
return super().clip_grad_norm(model, max_norm)
def unwrap_model(self, model: nn.Module):
if isinstance(model, FSDP) and self.use_distributed:
with FSDP.state_dict_type(