[rl][spmd_types] DTensor state dict adapter for vLLM HF weights

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,

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:

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.