feat: make max_grad_norm optional (None disables clipping)

- TrainConfig.max_grad_norm defaults to None
- executor.clip_grad_norm returns grad norm without clipping when None
- train.py --max_grad_norm defaults to None

astrai/config/train_config.py

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

astrai/parallel/executor.py

@@ -148,7 +148,14 @@ 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:
+    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()
         total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
         if isinstance(total_norm, torch.Tensor):
             return total_norm.item()
@@ -289,7 +296,14 @@ class FSDPExecutor(BaseExecutor):
             return model.no_sync()
         return contextlib.nullcontext()
 
-    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
+    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()
         if isinstance(model, FSDP) and self.use_distributed:
             total_norm = model.clip_grad_norm_(max_norm)
             if isinstance(total_norm, torch.Tensor):