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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user