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
2026-07-29 23:09:37 +08:00
parent 8150ab6c32
commit 646b1b0f46
6 changed files with 18 additions and 102 deletions
@@ -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`
@@ -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 |
@@ -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",
 ]
@@ -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__(
-        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).
+    """FSDP 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
@@ -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
@@ -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(
-            "--nprocs > 1 requires --parallel_mode to be 'ddp', 'fsdp', or 'fsdp2'"
-        )
+        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")

    # Load config
    config_path = os.path.join(param_path, "config.json")