diff --git a/torchtitan/distributed/spmd_types.py b/torchtitan/distributed/spmd_types.py index bfbce3aef7..e140680a6a 100644 --- a/torchtitan/distributed/spmd_types.py +++ b/torchtitan/distributed/spmd_types.py @@ -9,13 +9,14 @@ from __future__ import annotations import contextlib -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from threading import local from typing import Any, TYPE_CHECKING import spmd_types as spmd import torch from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor from torchtitan.distributed.utils import get_spmd_backend @@ -30,6 +31,9 @@ __all__ = [ "annotate_input_spmd_types", "current_spmd_mesh", + "dtensor_to_plain_tensor_state_dict", + "maybe_set_sparse_mesh", + "plain_tensor_to_dtensor_state_dict", "spmd_dense_mesh", "spmd_sparse_mesh", "spmd_mesh_size", @@ -38,13 +42,55 @@ "spmd_validate_redistributions", "set_current_spmd_mesh", "set_spmd_meshes", - "maybe_set_sparse_mesh", ] _MESH_TLS = local() +def plain_tensor_to_dtensor_state_dict( + state_dict: dict[str, Any], + *, + state_dict_layouts: Mapping[str, "SpmdLayout"], + parallel_dims: "ParallelDims", +) -> dict[str, Any]: + """Represent plain local state tensors as DTensors for state transfer.""" + from torchtitan.distributed.parallel_dims import unfold_dp_axes + from torchtitan.protocols.sharding import resolve_placements + + dtensor_state_dict = dict(state_dict) + with torch.no_grad(): + for name, target in state_dict.items(): + if not isinstance(target, torch.Tensor) or isinstance(target, DTensor): + continue + + layout = state_dict_layouts.get(name) + if layout is None: + raise KeyError(f"{name} is missing SPMD layout metadata") + + mesh = parallel_dims.get_activated_mesh(unfold_dp_axes(layout.axes())) + if mesh is None: + continue + + dtensor_state_dict[name] = DTensor.from_local( + target, + mesh, + resolve_placements(layout, mesh), + run_check=False, + ) + return dtensor_state_dict + + +def dtensor_to_plain_tensor_state_dict( + state_dict: dict[str, Any], +) -> dict[str, Any]: + """Replace DTensor state-dict entries with their plain local tensors.""" + return { + name: value.to_local() if isinstance(value, DTensor) else value + for name, value in state_dict.items() + } + + def set_spmd_meshes( *, dense_mesh: DeviceMesh, diff --git a/torchtitan/experiments/rl/actors/generator.py b/torchtitan/experiments/rl/actors/generator.py index 17267f40a2..1515f3fce4 100644 --- a/torchtitan/experiments/rl/actors/generator.py +++ b/torchtitan/experiments/rl/actors/generator.py @@ -27,10 +27,12 @@ Port, PortReceiver, ) -from torch.distributed.tensor import DTensor from torchtitan.components.checkpointer import CheckpointManager from torchtitan.config import CompileConfig, Configurable, DebugConfig, OverrideConfig -from torchtitan.distributed.parallel_dims import unfold_dp_axes +from torchtitan.distributed.spmd_types import ( + dtensor_to_plain_tensor_state_dict, + plain_tensor_to_dtensor_state_dict, +) from torchtitan.distributed.utils import get_spmd_backend, set_batch_invariance from torchtitan.experiments.rl.batch_invariance import ( force_logprobs_fn_for_batch_invariance, @@ -47,14 +49,9 @@ IntraGeneratorRouter, ) from torchtitan.experiments.rl.types import Completion -from torchtitan.models.common.attention import ( - FlexAttention, - FusedQKVLinear, - VarlenAttention, -) +from torchtitan.models.common.attention import FlexAttention, VarlenAttention from torchtitan.observability import structured_logger as sl from torchtitan.protocols.model_spec import ModelSpec -from torchtitan.protocols.sharding import resolve_placements, SpmdLayout from torchtitan.tools.logging import init_logger from torchtitan.tools.utils import has_cuda_capability from vllm import EngineArgs, LLMEngine, SamplingParams @@ -1333,82 +1330,11 @@ async def _get_spmd_state_dict(self, model_sd: dict, *, model) -> None: state-dict path, then put the local tensors back before load_state_dict. """ - def _fqn_to_spmd_layout(model: torch.nn.Module) -> dict[str, SpmdLayout]: - layouts: dict[str, SpmdLayout] = {} - - for module_fqn, module in model.named_modules(): - sharding_config = getattr(module, "_sharding_config", None) - if sharding_config is not None: - for state_name, layout in sharding_config.state_shardings.items(): - fqn = f"{module_fqn}.{state_name}" if module_fqn else state_name - layouts[fqn] = layout - - # FusedSwiGLU keeps its sharding on the fused w13 parameter, - # but its state dict exposes split w1.weight/w3.weight - # (_split_w13_on_save). Mirror w13's layout onto the split - # keys -- slicing the gate/up dim of an S(0) w13 yields S(0) - # w1/w3, which is what the DTensor path gets implicitly. - w13_layout = sharding_config.state_shardings.get("w13") - if w13_layout is not None: - for proj_name in ("w1", "w3"): - layouts[f"{module_fqn}.{proj_name}.weight"] = w13_layout - - if isinstance(module, FusedQKVLinear): - # FusedQKVLinear exposes split wq/wk/wv state-dict keys while - # the sharding layout lives on the fused wqkv parameter. - # TODO: This assumes fused and split QKV layouts stay - # equivalent. The load hook all-gathers anyway, so replace - # this with a less fragile fused-QKV state-dict path. - wqkv_sharding_config = getattr( - module.wqkv, "_sharding_config", None - ) - if wqkv_sharding_config is None: - continue - for ( - state_name, - layout, - ) in wqkv_sharding_config.state_shardings.items(): - for proj_name in ("wq", "wk", "wv"): - layouts[f"{module_fqn}.{proj_name}.{state_name}"] = layout - - return layouts - - layouts = _fqn_to_spmd_layout(model.model) - - dtensor_model_sd = dict(model_sd) - with torch.no_grad(): - for name, target in model_sd.items(): - if not isinstance(target, torch.Tensor): - continue - - layout = layouts.get(name) - if layout is None: - if name.endswith( - ( - ".vllm_attn._k_scale", - ".vllm_attn._prob_scale", - ".vllm_attn._q_scale", - ".vllm_attn._v_scale", - ) - ): - # vLLM attention scale buffers are backend-owned plain - # replicated state with no TorchTitan ShardingConfig. - continue - raise KeyError(f"{name} is missing SPMD layout metadata") - - if ( - mesh := model.parallel_dims.get_activated_mesh( - unfold_dp_axes(layout.axes()) - ) - ) is None: - continue - - dtensor_model_sd[name] = DTensor.from_local( - target, - mesh, - resolve_placements(layout, mesh), - run_check=False, - ) + dtensor_model_sd = plain_tensor_to_dtensor_state_dict( + model_sd, + state_dict_layouts=model.get_state_dict_layouts(), + parallel_dims=model.parallel_dims, + ) await ts.get_state_dict( "model_state_dict", @@ -1417,10 +1343,7 @@ def _fqn_to_spmd_layout(model: torch.nn.Module) -> dict[str, SpmdLayout]: direct_rdma=False, ) - with torch.no_grad(): - for name, value in dtensor_model_sd.items(): - if isinstance(value, DTensor): - model_sd[name] = value.to_local() + model_sd.update(dtensor_to_plain_tensor_state_dict(dtensor_model_sd)) @concurrent_endpoint async def close(self) -> None: diff --git a/torchtitan/experiments/rl/models/vllm_wrapper.py b/torchtitan/experiments/rl/models/vllm_wrapper.py index bf1a2c89a0..181313e9b0 100644 --- a/torchtitan/experiments/rl/models/vllm_wrapper.py +++ b/torchtitan/experiments/rl/models/vllm_wrapper.py @@ -14,11 +14,13 @@ import dataclasses from dataclasses import dataclass from functools import partial +from typing import Any import spmd_types as spmd import torch import torch.distributed as dist +from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.tensor import DTensor, Replicate, Shard from torchtitan.components.checkpointer import CheckpointManager from torchtitan.config import ( @@ -30,11 +32,18 @@ ) from torchtitan.distributed import utils as dist_utils from torchtitan.distributed.parallel_dims import ParallelDims -from torchtitan.distributed.spmd_types import current_spmd_mesh +from torchtitan.distributed.spmd_types import ( + current_spmd_mesh, + dtensor_to_plain_tensor_state_dict, + plain_tensor_to_dtensor_state_dict, +) from torchtitan.distributed.utils import is_in_batch_invariant_mode from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig +from torchtitan.models.common.attention import FusedQKVLinear from torchtitan.protocols.model_spec import ModelSpec from torchtitan.protocols.module import Module +from torchtitan.protocols.sharding import SpmdLayout +from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_reduce @@ -110,6 +119,41 @@ def _replace_vllm_layer_configs(model_config): return dataclasses.replace(model_config, layers=new_layers) +class PlainToDTensorStateDictAdapter(BaseStateDictAdapter): + """Add plain local tensor handling to a model-format state-dict adapter.""" + + def __init__( + self, + adapter: BaseStateDictAdapter, + state_dict_layouts: dict[str, SpmdLayout], + parallel_dims: ParallelDims, + ) -> None: + self.adapter = adapter + self.state_dict_layouts = state_dict_layouts + self.parallel_dims = parallel_dims + self.fqn_to_index_mapping = adapter.fqn_to_index_mapping + self.hf_assets_path = adapter.hf_assets_path + + def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: + return self.adapter.to_hf( + plain_tensor_to_dtensor_state_dict( + state_dict, + state_dict_layouts=self.state_dict_layouts, + parallel_dims=self.parallel_dims, + ) + ) + + def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: + return dtensor_to_plain_tensor_state_dict(self.adapter.from_hf(hf_state_dict)) + + def get_hf_storage_reader( + self, + path: str, + from_quantized: bool = False, + ) -> HuggingFaceStorageReader: + return self.adapter.get_hf_storage_reader(path, from_quantized) + + # NOTE: Monkeypatch vLLM's weak_ref_tensor to handle DTensor # This is because piecewise CUDA-graph capture calls weak_ref_tensor() # on every subgraphoutput (see vllm/compilation/cuda_graph.py). @@ -514,6 +558,12 @@ def _maybe_initial_load_weights(self) -> None: model_config=self.config, hf_assets_path=cfg.initial_load_path, ) + if self.parallel_dims.spmd_backend == "spmd_types": + sd_adapter = PlainToDTensorStateDictAdapter( + sd_adapter, + self.get_state_dict_layouts(), + self.parallel_dims, + ) # Model-only CheckpointManager: initial_load_model_only=True (default) # ensures only MODEL state is loaded, so None optimizer/lr_scheduler @@ -533,6 +583,56 @@ def _maybe_initial_load_weights(self) -> None: # the live weights fit. torch.cuda.empty_cache() + def get_state_dict_layouts(self) -> dict[str, SpmdLayout]: + """Return SPMD layouts keyed by the model's exposed state-dict names. + + TODO(pianpwk): Remove the fused QKV state-dict glue code. + """ + layouts: dict[str, SpmdLayout] = {} + + for module_fqn, module in self.model.named_modules(): + module_prefix = f"{module_fqn}." if module_fqn else "" + sharding_config = getattr(module, "_sharding_config", None) + if sharding_config is not None: + for state_name, layout in sharding_config.state_shardings.items(): + layouts[f"{module_prefix}{state_name}"] = layout + + # FusedSwiGLU exposes split w1/w3 state-dict keys while the + # layout is declared on the fused w13 parameter. + w13_layout = sharding_config.state_shardings.get("w13") + if w13_layout is not None: + for proj_name in ("w1", "w3"): + layouts[f"{module_prefix}{proj_name}.weight"] = w13_layout + + if isinstance(module, FusedQKVLinear): + # FusedQKVLinear exposes split wq/wk/wv state-dict keys while + # the layout is declared on the fused wqkv parameter. + wqkv_sharding_config = getattr( + module.wqkv, + "_sharding_config", + None, + ) + if wqkv_sharding_config is None: + continue + for ( + state_name, + layout, + ) in wqkv_sharding_config.state_shardings.items(): + for proj_name in ("wq", "wk", "wv"): + layouts[f"{module_prefix}{proj_name}.{state_name}"] = layout + + if module_fqn.rsplit(".", 1)[-1] == "vllm_attn": + for buffer_name, _ in module.named_buffers(recurse=False): + if buffer_name in { + "_k_scale", + "_prob_scale", + "_q_scale", + "_v_scale", + }: + layouts[f"{module_prefix}{buffer_name}"] = SpmdLayout({}) + + return layouts + def load_weights(self, weights_iter): """ vLLM required API. diff --git a/torchtitan/experiments/rl/tests/test_bitwise_parity.py b/torchtitan/experiments/rl/tests/test_bitwise_parity.py index 2511b75506..20bdf6c3ef 100644 --- a/torchtitan/experiments/rl/tests/test_bitwise_parity.py +++ b/torchtitan/experiments/rl/tests/test_bitwise_parity.py @@ -97,7 +97,7 @@ # TODO: directly testing against PolicyTrainer with debug model to avoid OOM def build_trainer_model( config: Controller.Config, -) -> tuple[torch.nn.Module, torch.device]: +) -> tuple[torch.nn.Module, torch.device, dist_utils.SpmdContext]: """Build, parallelize, and load weights for the trainer model. Mirrors PolicyTrainer._build_model() without the Monarch actor framework. @@ -120,6 +120,10 @@ def build_trainer_model( world_size=dist.get_world_size(), spmd_backend=parallelism.spmd_backend, ) + train_context = dist_utils.get_spmd_context( + parallel_dims=parallel_dims, + spmd_typechecking=False, + ) dist_utils.set_determinism( parallel_dims, @@ -173,7 +177,7 @@ def build_trainer_model( ) model.eval() - return model, device + return model, device, train_context # TODO: directly testing against VLLMGenerator with debug model to avoid OOM @@ -607,8 +611,6 @@ def setUpClass(cls): ) config = cls.config_fn() - config.trainer.parallelism.spmd_backend = "partial_dtensor" - config.generator.parallelism.spmd_backend = "partial_dtensor" hf_path = os.environ.get(cls.hf_assets_env_var) if hf_path: config.hf_assets_path = hf_path @@ -665,7 +667,7 @@ def setUpClass(cls): # GPU memory for vLLM to leave room for the trainer model. config.generator.gpu_memory_limit = 0.5 - cls.model, cls.device = build_trainer_model(config) + cls.model, cls.device, cls.train_context = build_trainer_model(config) cls.engine = build_inference_engine(config) if cls.sync_weights_from_trainer: _sync_trainer_weights_to_vllm(cls.model, cls.engine) @@ -735,7 +737,7 @@ def test_batch_invariance(self): n = len(self.prompt_ids) mid = max(1, n // 2) - with torch.no_grad(): + with type(self).train_context(), torch.no_grad(): lps_partial = compute_trainer_prefill_logprobs( model, self.prompt_ids[:mid], @@ -769,7 +771,7 @@ def test_trainer_vs_vllm_prefill(self): model = self.model engine = self.engine - with torch.no_grad(): + with type(self).train_context(), torch.no_grad(): trainer_lps = compute_trainer_prefill_logprobs( model, self.prompt_ids, self.device, attn_backend=self.attn_backend )