diff --git a/assets/docs/architecture.md b/assets/docs/architecture.md index 07b5921..be8d2bc 100644 --- a/assets/docs/architecture.md +++ b/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` diff --git a/assets/docs/params.md b/assets/docs/params.md index 3d359e8..c4c08bd 100644 --- a/assets/docs/params.md +++ b/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 | diff --git a/astrai/parallel/__init__.py b/astrai/parallel/__init__.py index bfe96af..45dc97b 100644 --- a/astrai/parallel/__init__.py +++ b/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", ] diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index 5c006db..13db3cc 100644 --- a/astrai/parallel/executor.py +++ b/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__( - 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 diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 8b7c51a..5d3d26f 100644 --- a/astrai/trainer/train_context.py +++ b/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 diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 7d9a2d6..8a238f7 100644 --- a/scripts/tools/train.py +++ b/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( - "--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")