diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index d9086de..846c67b 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -72,7 +72,7 @@ class TrainConfig(BaseConfig): metadata={"help": "Number of batch iterations between metric logs."}, ) metrics: List[str] = field( - default_factory=lambda: ["loss", "lr"], + default_factory=lambda: ["loss", "lr", "grad_norm"], metadata={"help": "Metrics to record during training."}, ) diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index 4987823..571365d 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -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( diff --git a/astrai/trainer/metric_util.py b/astrai/trainer/metric_util.py index c66fc44..e81d9e9 100644 --- a/astrai/trainer/metric_util.py +++ b/astrai/trainer/metric_util.py @@ -1,42 +1,25 @@ -from typing import Any, Callable, Dict +from typing import Dict import torch import torch.nn as nn -def _grad_stat( - model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any -) -> dict: - results = {} - for name, param in model.named_parameters(): - results[name] = default - if param.grad is not None: - results[name] = fn(param.grad.data) - return results +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 - -def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]: - return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0) - - -def grad_std(model: nn.Module) -> Dict[str, float]: - return _grad_stat(model, lambda g: g.std().item(), 0.0) - - -def grad_max(model: nn.Module) -> Dict[str, float]: - return _grad_stat(model, lambda g: g.max().item(), -float("inf")) - - -def grad_min(model: nn.Module) -> Dict[str, float]: - return _grad_stat(model, lambda g: g.min().item(), float("inf")) - - -def grad_mean(model: nn.Module) -> Dict[str, float]: - return _grad_stat(model, lambda g: g.mean().item(), 0.0) - - -def grad_nan_num(model: nn.Module) -> Dict[str, int]: - return _grad_stat(model, lambda g: g.isnan().sum().item(), 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): @@ -52,24 +35,4 @@ def ctx_get_val_loss(ctx): def ctx_get_grad_norm(ctx): - return grad_norm(ctx.model) - - -def ctx_get_grad_std(ctx): - return grad_std(ctx.model) - - -def ctx_get_grad_max(ctx): - return grad_max(ctx.model) - - -def ctx_get_grad_min(ctx): - return grad_min(ctx.model) - - -def ctx_get_grad_mean(ctx): - return grad_mean(ctx.model) - - -def ctx_get_grad_nan_num(ctx): - return grad_nan_num(ctx.model) + return ctx.grad_norm diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index eeb2b21..a052122 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -9,7 +9,6 @@ from typing import IO, Callable, List, Optional, Protocol, runtime_checkable import torch import torch.distributed as dist import torch.nn as nn -from torch.nn.utils import clip_grad_norm_ from torch.utils.checkpoint import checkpoint as torch_checkpoint from tqdm import tqdm @@ -18,12 +17,7 @@ from astrai.parallel import only_on_rank from astrai.parallel.setup import get_current_device, get_rank from astrai.serialization import Checkpoint from astrai.trainer.metric_util import ( - ctx_get_grad_max, - ctx_get_grad_mean, - ctx_get_grad_min, - ctx_get_grad_nan_num, ctx_get_grad_norm, - ctx_get_grad_std, ctx_get_loss, ctx_get_lr, ctx_get_val_loss, @@ -86,7 +80,9 @@ class GradientClippingCallback(TrainCallback): self.max_grad_norm = max_grad_norm def on_optimizer_step(self, context: TrainContext): - clip_grad_norm_(context.model.parameters(), self.max_grad_norm) + context.grad_norm = context.executor.clip_grad_norm( + context.model, self.max_grad_norm + ) @CallbackFactory.register("gradient_checkpointing") @@ -252,11 +248,6 @@ class MetricLoggerCallback(TrainCallback): "lr": ctx_get_lr, "val_loss": ctx_get_val_loss, "grad_norm": ctx_get_grad_norm, - "grad_std": ctx_get_grad_std, - "grad_max": ctx_get_grad_max, - "grad_min": ctx_get_grad_min, - "grad_mean": ctx_get_grad_mean, - "grad_nan_num": ctx_get_grad_nan_num, } def _metrics(self, context: TrainContext, names): diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index c7c7877..783cd68 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -31,6 +31,7 @@ class TrainContext: epoch: int = field(default=0) iteration: int = field(default=0) loss: float = field(default=0.0) + grad_norm: Optional[float] = field(default=None) val_dataloader: Optional[DataLoader] = field(default=None) val_loss: Optional[float] = field(default=None) diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 5c8611a..e1a5291 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -150,8 +150,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--metrics", nargs="*", - default=["loss", "lr"], - help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr.", + default=["loss", "lr", "grad_norm"], + help="Metrics to log (e.g. --metrics loss lr val_loss). Default: loss lr grad_norm.", ) parser.add_argument( "--log_dir",