ViperEkura
0f1fcb079f
- 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
39 lines
896 B
Python
39 lines
896 B
Python
from typing import Dict
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, float]:
|
|
grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
|
|
if not grads:
|
|
return 0.0
|
|
|
|
total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
|
|
if per_param:
|
|
norms = {}
|
|
for name, param in model.named_parameters():
|
|
if param.grad is not None:
|
|
norms[name] = param.grad.norm(2).item()
|
|
else:
|
|
norms[name] = 0.0
|
|
norms["total"] = total_sq.sqrt().item()
|
|
return norms
|
|
return total_sq.sqrt().item()
|
|
|
|
|
|
def ctx_get_loss(ctx):
|
|
return ctx.loss
|
|
|
|
|
|
def ctx_get_lr(ctx):
|
|
return ctx.optimizer.param_groups[-1]["lr"]
|
|
|
|
|
|
def ctx_get_val_loss(ctx):
|
|
return ctx.val_loss
|
|
|
|
|
|
def ctx_get_grad_norm(ctx):
|
|
return ctx.grad_norm
|