refactor: replace FSDP with FSDP2 as default parallel backend

- Remove FSDPExecutor (FullyShardedDataParallel wrapper)
- Rename FSDP2Executor to FSDPExecutor, register as 'fsdp'
- Remove 'fsdp2' from CLI choices, make 'fsdp' the default parallel_mode
- Pass after_wrap to executor.prepare for compile-after-wrap ordering
- Update architecture.md, params.md, AGENTS.md references
- FSDP2 uses per-module fully_shard: no FlatParameter, better compile compat

assets/docs/architecture.md

@@ -1195,11 +1195,6 @@ classDiagram
         }
 
         class FSDPExecutor {
-            -_prepare_model(model) nn.Module
-            +unwrap_model(model) dict
-        }
-
-        class FSDP2Executor {
             -_prepare_model(model) nn.Module
             -_no_sync(model) context manager
             +unwrap_model(model) dict
@@ -1302,7 +1297,6 @@ classDiagram
     BaseExecutor <|-- NoneExecutor
     BaseExecutor <|-- DDPExecutor
     BaseExecutor <|-- FSDPExecutor
-    BaseExecutor <|-- FSDP2Executor
     ResponseBuilder <|-- OpenAIResponseBuilder
     ResponseBuilder <|-- AnthropicResponseBuilder
     BaseToolParser <|-- SimpleJsonToolParser
@@ -1396,7 +1390,6 @@ classDiagram
     ExecutorFactory ..> NoneExecutor : creates
     ExecutorFactory ..> DDPExecutor : creates
     ExecutorFactory ..> FSDPExecutor : creates
-    ExecutorFactory ..> FSDP2Executor : creates
     ToolParserFactory ..> BaseToolParser : creates
     TrainContextBuilder ..> ExecutorFactory : creates
     Trainer ..> TrainContextBuilder : uses
@@ -1444,7 +1437,7 @@ classDiagram
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–ContiguousCache/PageCache, CacheView–ContiguousCacheView/PageCacheView, Allocator–Storage, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
-| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, FSDP2Executor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
+| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
 
@@ -1461,7 +1454,7 @@ classDiagram
 | **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
-| **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor`, `FSDP2Executor` | Gradient accumulation & model distribution |
+| **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
 | **Storage** | `Store`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
 | **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
@@ -1471,7 +1464,7 @@ classDiagram
 1. **Config → Training**: `TrainConfig` holds `model_fn`, `dataset`, `optimizer_fn`, `scheduler_fn`, `parallel_mode`, `executor_kwargs`
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
-4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor` / `FSDP2Executor`
+4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
 5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `KVCache` + `SamplingPipeline`
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data`

assets/docs/params.md

@@ -84,7 +84,7 @@ Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`f
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--nprocs` | Number of GPUs / processes | 1 |
-| `--parallel_mode` | Parallel strategy (`none`, `ddp`, `fsdp`, or `fsdp2`) | none |
+| `--parallel_mode` | Parallel strategy (`none`, `ddp`, `fsdp`) | fsdp |
 | `--device_type` | Device type | cuda |
 | `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |
 | `--backend` | Distributed training backend | nccl |

astrai/parallel/__init__.py

@@ -4,7 +4,6 @@ from astrai.parallel.executor import (
     BaseExecutor,
     DDPExecutor,
     ExecutorFactory,
-    FSDP2Executor,
     FSDPExecutor,
     GradientState,
     NoneExecutor,
@@ -36,5 +35,4 @@ __all__ = [
     "NoneExecutor",
     "DDPExecutor",
     "FSDPExecutor",
-    "FSDP2Executor",
 ]

astrai/parallel/executor.py

@@ -11,11 +11,8 @@ import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.fsdp import (
     FSDPModule,
-    FullStateDictConfig,
-    StateDictType,
     fully_shard,
 )
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.tensor import DTensor
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
@@ -95,11 +92,14 @@ class BaseExecutor:
         optimizer_fn: Optional[Callable[[nn.Module], Optimizer]] = None,
         scheduler_fn: Optional[Callable[[Optimizer], LRScheduler]] = None,
         before_wrap: Optional[Callable[[nn.Module], nn.Module]] = None,
+        after_wrap: Optional[Callable[[nn.Module], nn.Module]] = None,
     ) -> Tuple[nn.Module, Optional[Optimizer], Optional[LRScheduler]]:
         model = model_fn()
         if before_wrap is not None:
             model = before_wrap(model)
         model = self._prepare_model(model)
+        if after_wrap is not None:
+            model = after_wrap(model)
         optimizer = None
         scheduler = None
         if optimizer_fn is not None:
@@ -238,88 +238,11 @@ class DDPExecutor(BaseExecutor):
 
 @ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
-    def __init__(
+    """FSDP executor using `torch.distributed.fsdp.fully_shard` (per-module API).
-        self,
-        grad_accum_steps: int = 1,
-        process_group=None,
-        sharding_strategy=None,
-        cpu_offload=None,
-        auto_wrap_policy=None,
-        backward_prefetch=None,
-        mixed_precision=None,
-        ignored_modules=None,
-        param_init_fn=None,
-        sync_module_states: bool = False,
-        forward_prefetch: bool = False,
-        limit_all_gathers: bool = True,
-        ignored_states=None,
-        device_mesh=None,
-    ):
-        super().__init__(grad_accum_steps=grad_accum_steps)
-        self._fsdp_kwargs = {
-            k: v
-            for k, v in dict(
-                process_group=process_group,
-                sharding_strategy=sharding_strategy,
-                cpu_offload=cpu_offload,
-                auto_wrap_policy=auto_wrap_policy,
-                backward_prefetch=backward_prefetch,
-                mixed_precision=mixed_precision,
-                ignored_modules=ignored_modules,
-                param_init_fn=param_init_fn,
-                sync_module_states=sync_module_states,
-                forward_prefetch=forward_prefetch,
-                limit_all_gathers=limit_all_gathers,
-                use_orig_params=True,
-                ignored_states=ignored_states,
-                device_mesh=device_mesh,
-            ).items()
-            if v is not None
-        }
-        self._original_model: Optional[nn.Module] = None
-
-    def _prepare_model(self, model: nn.Module) -> nn.Module:
-        if not self.use_distributed:
-            logger.warning("FSDP backend selected but world_size=1, model not wrapped")
-            return model
-        self._original_model = model
-        device_id = torch.device("cuda", get_rank())
-        model = FSDP(model, device_id=device_id, **self._fsdp_kwargs)
-        logger.info("Model wrapped with FSDP (world_size=%d)", get_world_size())
-        return model
-
-    def _no_sync(self, model: nn.Module):
-        if isinstance(model, FSDP):
-            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(
-                model,
-                StateDictType.FULL_STATE_DICT,
-                FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-            ):
-                return model.state_dict()
-
-        return model.state_dict()
-
-
-@ExecutorFactory.register("fsdp2")
-class FSDP2Executor(BaseExecutor):
-    """FSDP2 executor using `torch.distributed.fsdp.fully_shard` (per-module API).
 
     Wraps each child module individually via ``fully_shard``.
     Skips the root model because ``ABC + Generic[T]`` in the MRO makes
-    FSDP2's dynamic ``__class__`` assignment fail at the CPython level.
+    ``fully_shard``'s dynamic ``__class__`` assignment fail at the CPython level.
     Original ``Parameter`` objects are preserved (as DTensors) — no
     ``FlatParameter``, no ``use_orig_params=True`` hack.
     """
@@ -338,7 +261,7 @@ class FSDP2Executor(BaseExecutor):
 
     def _prepare_model(self, model: nn.Module) -> nn.Module:
         if not self.use_distributed:
-            logger.warning("FSDP2 backend selected but world_size=1, model not wrapped")
+            logger.warning("FSDP backend selected but world_size=1, model not wrapped")
             return model
 
         kwargs = dict(
@@ -356,7 +279,7 @@ class FSDP2Executor(BaseExecutor):
                 fully_shard(child, **kwargs)
 
         logger.info(
-            "FSDP2 wrapping applied to %d direct children (root skipped for ABC compat)",
+            "FSDP wrapping applied to %d direct children (root skipped for ABC compat)",
             len(list(model.children())),
         )
         return model

astrai/trainer/train_context.py

@@ -127,6 +127,9 @@ class TrainContextBuilder:
                 )
             if preloaded_state_dict is not None:
                 m.load_state_dict(preloaded_state_dict, strict=False)
+            return m
+
+        def _after_wrap(m):
             if cfg.compile_mode is not None:
                 logger.info("torch.compile enabled (mode=%s)", cfg.compile_mode)
                 m = torch.compile(m, mode=cfg.compile_mode)
@@ -148,6 +151,7 @@ class TrainContextBuilder:
             cfg.optimizer_fn,
             cfg.scheduler_fn,
             before_wrap=_before_wrap,
+            after_wrap=_after_wrap,
         )
 
         train_dataset = cfg.dataset

scripts/tools/train.py

@@ -115,7 +115,7 @@ def _merge_yaml_into_kwargs(config_path: str, passed_kwargs: dict) -> dict:
 
 
 _TRAIN_TYPE = ["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"]
-_PARALLEL = ["none", "ddp", "fsdp", "fsdp2"]
+_PARALLEL = ["none", "ddp", "fsdp"]
 _SCHEDULES = ["cosine", "sgdr", "wsd"]
 _BACKENDS = ["nccl", "gloo"]
 _START_METHODS = ["spawn", "fork", "forkserver"]
@@ -247,7 +247,7 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
 @click.option(
     "--parallel_mode",
     type=click.Choice(_PARALLEL),
-    default="none",
+    default="fsdp",
     help="Parallel strategy.",
 )
 @click.option("--device_type", type=str, default="cuda", help="Device type.")
@@ -418,9 +418,7 @@ def train(
     if not os.path.exists(param_path):
         raise FileNotFoundError(f"Model directory not found: {param_path}")
     if nprocs > 1 and parallel_mode == "none":
-        raise ValueError(
+        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
-            "--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'"
-        )
 
     # Load config
     config_path = os.path.join(param_path, "config.json")