fix: default max_grad_norm to 1.0 and drop None branch

assets/docs/params.md

@@ -26,7 +26,7 @@
 |-----------|-------------|---------|
 | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
-| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | None |
+| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | 1.0 |
 
 ### Optimizer (MuonMix)
 

astrai/config/train_config.py

@@ -38,7 +38,7 @@ class TrainConfig(BaseConfig):
         default=1, metadata={"help": "Number of iterations between steps."}
     )
     max_grad_norm: Optional[float] = field(
-        default=None,
+        default=1.0,
         metadata={"help": "Maximum gradient norm. None disables clipping."},
     )
     gradient_checkpointing_modules: List[str] = field(

astrai/parallel/executor.py

@@ -148,14 +148,7 @@ class BaseExecutor:
     def grad_accum_steps(self) -> int:
         return self.gradient_state.num_steps
 
-    def clip_grad_norm(self, model: nn.Module, max_norm: Optional[float]) -> float:
-        if max_norm is None:
-            total_norm = torch.norm(
-                torch.stack(
-                    [p.grad.norm(2) for p in model.parameters() if p.grad is not None]
-                )
-            )
-            return total_norm.item()
+    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()
@@ -289,9 +282,7 @@ class FSDPExecutor(BaseExecutor):
             return model.no_sync()
         return contextlib.nullcontext()
 
-    def clip_grad_norm(self, model: nn.Module, max_norm: Optional[float]) -> float:
-        if max_norm is None:
-            return super().clip_grad_norm(model, max_norm)
+    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):

scripts/tools/train.py

@@ -148,7 +148,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--max_grad_norm",
         type=float,
-        default=None,
+        default=1.0,
         help="Max gradient norm for clipping. None disables clipping.",
     )
     parser.add_argument(