refactor: remove redundant strategy/executor code

- Drop BaseStrategy.model_fn (stored but never read)
- Drop model_fn= passed to StrategyFactory.create in train_context
- Simplify FSDPExecutor.clip_grad_norm None branch to delegate to super()
- Remove DDPExecutor._gather_state_dict override (identical to base)

astrai/parallel/executor.py

@@ -231,13 +231,6 @@ class DDPExecutor(BaseExecutor):
             return model.module.state_dict()
         return model.state_dict()
 
-    def _gather_state_dict(self, model: nn.Module):
-        if not self.use_distributed:
-            return self.unwrap_model(model)
-        if get_rank() != 0:
-            return None
-        return self.unwrap_model(model)
-
 
 @ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
@@ -298,12 +291,7 @@ class FSDPExecutor(BaseExecutor):
 
     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()
+            return super().clip_grad_norm(model, max_norm)
         if isinstance(model, FSDP) and self.use_distributed:
             total_norm = model.clip_grad_norm_(max_norm)
             if isinstance(total_norm, torch.Tensor):

astrai/trainer/strategy.py

@@ -98,7 +98,6 @@ class BaseStrategy(ABC):
         self.model = model
         self.device = device
         self.executor = kwargs.pop("executor", None)
-        self.model_fn = kwargs.pop("model_fn", None)
         self.extra_kwargs = kwargs
 
     @abstractmethod

astrai/trainer/train_context.py

@@ -211,7 +211,6 @@ class TrainContextBuilder:
             model=context.model,
             device=device,
             executor=executor,
-            model_fn=cfg.model_fn,
             **strategy_kwargs,
         )