fix: make FSDP2 executor work with ABC+Generic model hierarchy

- Wrap each child module individually, skip root (CPython layout
  conflict between ABC+Generic and FSDP2 __class__ assignment)
- Remove manual unshard in clip_grad_norm (DTensor compatible)
- Fix _no_sync to iterate modules() instead of checking root
- Add reshard after unwrap_model
- Guard __init_subclass__ type resolution against dynamic subclasses
- Add fsdp2 to --parallel_mode CLI choices

astrai/factory.py

@@ -67,7 +67,10 @@ class BaseFactory(ABC, Generic[T]):
             if _get_origin(orig_base) is BaseFactory:
                 (arg,) = _get_args(orig_base)
                 cls._entries = {}
+                try:
                     cls._component_base = _resolve_type(arg, cls)
+                except Exception:
+                    cls._component_base = None
                 return
 
     @classmethod

astrai/parallel/executor.py

@@ -317,9 +317,11 @@ class FSDPExecutor(BaseExecutor):
 class FSDP2Executor(BaseExecutor):
     """FSDP2 executor using `torch.distributed.fsdp.fully_shard` (per-module API).
 
-    Unlike FSDP1's `FSDP(model, ...)` wrapper, FSDP2 wraps each submodule
+    Wraps each child module individually via ``fully_shard``.
-    individually via `fully_shard(module)`. Original `Parameter` objects are
+    Skips the root model because ``ABC + Generic[T]`` in the MRO makes
-    preserved (as DTensors) — no `FlatParameter`, no `use_orig_params=True` hack.
+    FSDP2's dynamic ``__class__`` assignment fail at the CPython level.
+    Original ``Parameter`` objects are preserved (as DTensors) — no
+    ``FlatParameter``, no ``use_orig_params=True`` hack.
     """
 
     def __init__(
@@ -346,26 +348,37 @@ class FSDP2Executor(BaseExecutor):
         )
         kwargs = {k: v for k, v in kwargs.items() if v is not None}
 
-        model = fully_shard(model, **kwargs)
+        for child in model.children():
-        logger.info("Model wrapped with FSDP2 (world_size=%d)", get_world_size())
+            if isinstance(child, nn.ModuleList):
+                for sub in child:
+                    fully_shard(sub, **kwargs)
+            else:
+                fully_shard(child, **kwargs)
+
+        logger.info(
+            "FSDP2 wrapping applied to %d direct children (root skipped for ABC compat)",
+            len(list(model.children())),
+        )
         return model
 
     @contextmanager
     def _no_sync(self, model: nn.Module):
-        if isinstance(model, FSDPModule):
+        fsdp_modules = [
-            model.set_requires_gradient_sync(False, recurse=True)
+            m for m in model.modules() if isinstance(m, FSDPModule)
+        ]
+        if fsdp_modules:
+            for m in fsdp_modules:
+                m.set_requires_gradient_sync(False, recurse=True)
             try:
                 yield
             finally:
-                model.set_requires_gradient_sync(True, recurse=True)
+                for m in fsdp_modules:
+                    m.set_requires_gradient_sync(True, recurse=True)
         else:
             yield
 
     def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
-        if isinstance(model, FSDPModule) and self.use_distributed:
+        if self.use_distributed:
-            for module in model.modules():
-                if isinstance(module, FSDPModule):
-                    module.unshard()
             total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
             if isinstance(total_norm, torch.Tensor):
                 return total_norm.item()
@@ -373,7 +386,7 @@ class FSDP2Executor(BaseExecutor):
         return super().clip_grad_norm(model, max_norm)
 
     def unwrap_model(self, model: nn.Module):
-        if not self.use_distributed or not isinstance(model, FSDPModule):
+        if not self.use_distributed:
             return model.state_dict()
 
         if get_rank() != 0:
@@ -384,7 +397,13 @@ class FSDP2Executor(BaseExecutor):
                 module.unshard()
 
         state_dict = model.state_dict()
-        return {
+        result = {
             k: (v.full_tensor() if isinstance(v, DTensor) else v)
             for k, v in state_dict.items()
         }
+
+        for module in model.modules():
+            if isinstance(module, FSDPModule):
+                module.reshard()
+
+        return result

scripts/tools/train.py

@@ -327,8 +327,8 @@ def parse_args() -> argparse.Namespace:
         "--parallel_mode",
         type=str,
         default="none",
-        choices=["none", "ddp", "fsdp"],
+        choices=["none", "ddp", "fsdp", "fsdp2"],
-        help="Parallel training strategy (none, ddp, fsdp).",
+        help="Parallel training strategy (none, ddp, fsdp, fsdp2).",
     )
     parser.add_argument(
         "--device_type", type=str, default="cuda", help="Device type to use."
@@ -472,7 +472,7 @@ def train(
     ]
     assert os.path.exists(param_path)
     if nprocs > 1 and parallel_mode == "none":
-        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
+        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'")
 
     # Load config
     config_path = os.path.join(param_path, "config.json")