From 9362cc4e07811ea3c80c942963bcaaccc83b23a8 Mon Sep 17 00:00:00 2001 From: drisspg Date: Fri, 18 Sep 2026 12:18:54 -0700 Subject: [PATCH 1/2] [RL] Restore model shard layouts after HF conversion HF format conversions can change DTensor placements. Restore the declared model layout before extracting each local tensor, rather than loading a shard with the wrong shape or ownership. --- .../workflows/integration_test_8gpu_rl.yaml | 3 + tests/unit_tests/rl/test_vllm_wrapper.py | 87 ++++++++++++++++++- torchtitan/rl/model/vllm_wrapper.py | 13 ++- 3 files changed, 100 insertions(+), 3 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_rl.yaml b/.github/workflows/integration_test_8gpu_rl.yaml index 14fe54ee97b..2c7796c985e 100644 --- a/.github/workflows/integration_test_8gpu_rl.yaml +++ b/.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. diff --git a/tests/unit_tests/rl/test_vllm_wrapper.py b/tests/unit_tests/rl/test_vllm_wrapper.py index aaa70bc507b..543e3ac32b9 100644 --- a/tests/unit_tests/rl/test_vllm_wrapper.py +++ b/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, + ) diff --git a/torchtitan/rl/model/vllm_wrapper.py b/torchtitan/rl/model/vllm_wrapper.py index ca68b4cf060..4f558b605dd 100644 --- a/torchtitan/rl/model/vllm_wrapper.py +++ b/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,17 @@ 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) + 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, From d9965ce27c55e202cf16f8d459d6f9728ef2e8c5 Mon Sep 17 00:00:00 2001 From: drisspg Date: Mon, 21 Sep 2026 15:59:23 -0700 Subject: [PATCH 2/2] [RL] Add TODO for FSDP generator layout handling --- torchtitan/rl/model/vllm_wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torchtitan/rl/model/vllm_wrapper.py b/torchtitan/rl/model/vllm_wrapper.py index 4f558b605dd..52ac02ffcc1 100644 --- a/torchtitan/rl/model/vllm_wrapper.py +++ b/torchtitan/rl/model/vllm_wrapper.py @@ -145,6 +145,8 @@ 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]: 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.