[RL] Restore model shard layouts after HF conversion

.github/workflows/integration_test_8gpu_rl.yaml

@@ -106,6 +106,9 @@ jobs:
           --assert-equal \
           --options='--trainer.parallelism.tensor_parallel_degree 4 --generator.parallelism.tensor_parallel_degree 4 --num_generators 1'
 
+        # HF adapter layout regression uses two CPU/Gloo ranks.
+        python -m pytest tests/unit_tests/rl/test_vllm_wrapper.py -v
+
         # Bitwise trainer/generator parity tests (batch-invariant). Each TP=2 case
         # runs in a separate torchrun invocation so it gets a fresh process group
         # and vLLM engine.

tests/unit_tests/rl/test_vllm_wrapper.py

@@ -4,18 +4,32 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from datetime import timedelta
+from pathlib import Path
+
 import spmd_types as spmd
 import torch
-from torchtitan.models.common.attention import QKVLinear
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import distribute_tensor, Replicate, Shard
 
+from torchtitan.distributed.parallel_dims import ParallelDims
+from torchtitan.models.common.attention import QKVLinear
 from torchtitan.models.common.decoder_sharding import dense_param_placement
 from torchtitan.models.common.feed_forward import FeedForward
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.moe import GroupedExperts
+from torchtitan.models.qwen3_5 import model_registry
+from torchtitan.models.qwen3_5.model import Qwen35Model
+from torchtitan.models.qwen3_5.state_dict_adapter import Qwen35StateDictAdapter
 from torchtitan.overrides.fused_swiglu import fused_grouped_experts
 from torchtitan.protocols.sharding import ShardingConfig
 
-from torchtitan.rl.model.vllm_wrapper import VLLMModelWrapper
+from torchtitan.rl.model.vllm_wrapper import (
+    PlainToDTensorStateDictAdapter,
+    VLLMModelWrapper,
+)
 
 
 def test_state_dict_layouts_include_native_feed_forward_weight():
@@ -103,3 +117,72 @@ def test_state_dict_layouts_include_split_expert_weights():
     assert layouts["experts.w1_EFD"] is colwise
     assert layouts["experts.w3_EFD"] is colwise
     assert layouts["experts.w2_EDF"] is rowwise
+
+
+def _check_hf_adapter_restores_local_shards(rank: int, rendezvous: str) -> None:
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "gloo",
+        init_method=rendezvous,
+        rank=rank,
+        world_size=2,
+        timeout=timedelta(seconds=60),
+    )
+    try:
+        mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("tp",))
+        model_config = model_registry("0.8B", seq_len=256, attn_backend="varlen").model
+        assert isinstance(model_config, Qwen35Model.Config)
+        model_adapter = Qwen35StateDictAdapter(model_config, hf_assets_path=None)
+        state_dict, expected, layouts = {}, {}, {}
+        for pattern, shape in (
+            ("layers.0.attn.in_proj_{}.weight", (2048, 1024)),
+            ("layers.0.attn.conv_{}.weight", (2048, 1, 4)),
+            ("vision_encoder.layers.0.attn.w{}.weight", (768, 768)),
+            ("vision_encoder.layers.0.attn.w{}.bias", (768,)),
+        ):
+            for index, part in enumerate(("q", "k", "v")):
+                key = pattern.format(part)
+                full = torch.arange(torch.Size(shape).numel()).reshape(shape).float()
+                full += index * 100
+                state_dict[key] = distribute_tensor(full, mesh, [Shard(0)])
+                expected[key] = full.chunk(2, dim=0)[rank].clone()
+                layouts[key] = dense_param_placement(tp=spmd.S(0))
+        # Check unchanged row-sharded, replicated, and plain values as well.
+        for key, placement, shape in (
+            ("layers.3.attn.wo.weight", Shard(1), (8, 8)),
+            ("norm.weight", Replicate(), (8,)),
+        ):
+            full = torch.arange(torch.Size(shape).numel()).reshape(shape).float()
+            state_dict[key] = distribute_tensor(full, mesh, [placement])
+            expected[key] = (
+                full.chunk(2, dim=1)[rank].clone()
+                if isinstance(placement, Shard)
+                else full
+            )
+            layouts[key] = dense_param_placement(
+                tp=spmd.S(1) if isinstance(placement, Shard) else spmd.R
+            )
+        state_dict["lm_head.weight"] = expected["lm_head.weight"] = torch.ones(1)
+        adapter = PlainToDTensorStateDictAdapter(
+            model_adapter,
+            layouts,
+            ParallelDims(
+                dp_replicate=1, dp_shard=1, cp=1, tp=2, pp=1, ep=1, world_size=2
+            ),
+        )
+        restored = adapter.from_hf(model_adapter.to_hf(state_dict))
+        assert restored.keys() == expected.keys()
+        for key in expected:
+            assert type(restored[key]) is torch.Tensor
+            torch.testing.assert_close(restored[key], expected[key], rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+def test_hf_adapter_restores_local_shards(tmp_path: Path) -> None:
+    mp.spawn(
+        _check_hf_adapter_restores_local_shards,
+        args=(f"file://{tmp_path / 'rendezvous'}",),
+        nprocs=2,
+        join=True,
+    )

torchtitan/rl/model/vllm_wrapper.py

@@ -41,6 +41,7 @@
 from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.protocols.model_spec import ModelSpec
 from torchtitan.protocols.module import Module
+from torchtitan.protocols.sharding import resolve_placements
 from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter
 from torchtitan.rl.distributed.parallelism import InferenceParallelismConfig
 from vllm.compilation.decorators import support_torch_compile
@@ -143,7 +144,19 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
         )
 
     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))
+        state_dict = self.adapter.from_hf(hf_state_dict)
+        # TODO(@andrewor14): Wrap the generator model with FSDP and revisit this
+        # explicit layout restoration once weights load into DTensor parameters.
+        for name, value in state_dict.items():
+            if isinstance(value, DTensor):
+                # Format conversions can reshard tensors, e.g. fused QKV splits.
+                # Restore the model's layout before discarding DTensor metadata.
+                state_dict[name] = value.redistribute(
+                    placements=resolve_placements(
+                        self.state_dict_layouts[name], value.device_mesh
+                    )
+                )
+        return dtensor_to_plain_tensor_state_dict(state_dict)
 
     def get_hf_storage_reader(
         self,