[RL] Restore model shard layouts after HF conversion

[drisspg]

Summary

Restore the declared model sharding after Hugging Face format conversion, before extracting plain local tensors.

Qwen3.5's QKV fuse/split operations can change DTensor placements: a TP2 projection expected as Shard(0) became Shard(1), producing local [2048, 512] instead of [1024, 1024]. Convolution and vision weights were affected too. The RL checkpoint adapter now redistributes to each parameter's existing layout before to_local(); it does not reshape incorrect shards or change CUDA-graph behavior.

Adds a two-rank CPU/Gloo regression with independent expected local slices and runs it in RL CI. Naming-only follow-up #4786 is the lower stack layer.

Validation

• New regression reproduces the original failure on both ranks and passes with the fix.
• Six focused wrapper/GDN metadata tests passed on this stack; scoped pre-commit passed. Scoped Pyrefly retains five verified preexisting errors in unchanged all-reduce code; no new errors.
• Before stacking, the same adapter fix passed regular Qwen3.5-0.8B TP2 eager/FULL inference on two GB200s: exact rank-local weights, 96 tokens/logprobs, 44 state records, and 43 native FULL replays per rank.
• GPU-run caveat: both runs completed inference artifacts but hit the 60-second deadline during shutdown. Clean shutdown and batch-invariant TP2 remain unvalidated. Runtime used Torch/vLLM Sep13 nightlies, FA4 b31, and Gym commit 4ac49088a, not a clean release-pin install.

.venv/bin/python -m pytest -v -ra tests/unit_tests/rl/test_vllm_wrapper.py tests/unit_tests/rl/test_gdn_graph_metadata.py

[tianyu-l]

@claude review this PR

Is the issue caused by such code?

• https://github.com/pytorch/torchtitan/blob/main/torchtitan/rl/model/vllm_wrapper.py#L584
• https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/attention.py#L742-L743

These code exists because

• QKVLinear is defined in a fused way
• the state dict is stored in an unfused way

We are going to change this soon, so that both are defined in the fused way, i.e. fusion becomes native in torchtitan (but not in HF checkpoints).

Do you agree that after the change, we no longer need the fix in this PR?

cc @anijain2305

[claude]

Claude finished @tianyu-l's task in 3m 9s —— View job

---

Review

I read the fix, the new regression test, and traced the resharding through Qwen35StateDictAdapter.from_hf and the two code locations you referenced.

Short answer to your question

No — I don't think making fusion native in TorchTitan removes the need for this fix. The two locations you cite are real fused/unfused glue, but they're a different fused/unfused boundary than the one that actually reshards here.

Where the resharding actually comes from

In this RL path the checkpoint is loaded as DTensors, and the resharding happens inside adapter.from_hf(...) — before dtensor_to_plain_tensor_state_dict extracts the local slice. The culprits are the HF-format split/chunk/reshape ops on DTensors in qwen3_5/state_dict_adapter.py:

• deltanet in_proj_qkv -> value.split([kd,kd,vd], dim=0) (L318)
• deltanet conv1d -> value.split(..., dim=0) (L333)
• vision qkv -> value.chunk(3, dim=0) (L345)
• MoE gate_up_proj -> value.chunk(2, dim=-2) (L290)
• conv3d patch embed -> value.reshape(...) (L367)

Splitting/chunking a Shard(0) DTensor along the sharded dim can't stay Shard(0), so DTensor emits a different placement (Shard(1)), yielding a wrong local slice ([2048, 512] instead of [1024, 1024]).

This matches the PR's own example and the regression test, which exercises layers.0.attn.in_proj_{q,k,v}, conv_{q,k,v}, and vision_encoder...w{q,k,v} — i.e. the GatedDeltaNet / vision / conv conversions, not the QKVLinear (GQAttention) hooks at attention.py:742-743.

Why the referenced code isn't the cause

• attention.py:742-743 (_split_qkv_on_save / _merge_qkv_on_load) are nn.Module state-dict hooks. In this path they run on plain local tensors (after dtensor_to_plain_tensor_state_dict), so they do no DTensor resharding here.
• vllm_wrapper.py:584 (FeedForward w1/w3 glue) and the sibling QKVLinear glue at L595 only populate the layout map the fix consumes — they don't perform the resharding.

Making TT store QKV/FFN fused natively would let you delete those hooks + glue. But the HF checkpoint stays unfused, so from_hf must still fuse (cat) / split HF tensors into TT's layout — and cat/split/reshape on sharded DTensors can still change placements. The deltanet/vision/conv conversions above have no relation to the internal fused-QKV representation and remain either way.

So the fix is best understood as a general invariant at the DTensor -> local boundary: whatever a format conversion does to placements, restore the model's declared layout before discarding DTensor metadata. That's not a speculative defensive check — it enforces a concrete, demonstrated contract.

Code review of the diff itself

The fix is correct and minimal. Two small notes:

1. KeyError if a converted DTensor key lacks a declared layout (vllm_wrapper.py:157, self.state_dict_layouts[name]). Today get_state_dict_layouts covers the exercised keys (the GPU run's 44 state records passed), but a future model param without a sharding config would raise a bare KeyError here rather than a clear message. Consider a targeted error (f"no declared layout for {name}") so a coverage gap fails legibly. Minor — loud failure is acceptable per repo norms.

2. Test quality is good — independent per-rank expected slices, and it deliberately covers Shard(0), Shard(1), Replicate, and a plain tensor, so it guards against the redistribute being applied where it shouldn't. Nice that it reproduces the original failure on both ranks.

Suggestion

Since the real trigger is the HF adapter's split/chunk/reshape (not the QKV fusion glue), it'd be worth tightening the code comment at vllm_wrapper.py:152-153 to name the actual conversions (deltanet/vision/conv splits), not just "fused QKV splits" — that phrasing is what invites the "won't native fusion fix this?" reading.

---

cc @tianyu-l @anijain2305 — happy to dig further if you have a specific checkpoint/parallelism combo in mind.
· branch drisspg/rl-hf-shard-layout-fix

[tianyu-l]

@claude

Question:

Why another callsite of from_hf doesn't need this treatment? Or is there a bug over there too?
https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpointer/dcp.py#L394

Also, could you give a simple and concrete example of the issue?

[claude]

Claude finished @tianyu-l's task in 2m 11s —— View job

---

Answers

1. Why doesn't dcp.py:394's from_hf need this treatment? Is there a bug there too?

No bug there. The difference is what happens to the resharded DTensor after from_hf returns.

In dcp.py:

state_dict = self.sd_adapter.from_hf(hf_state_dict)
states[MODEL].load_state_dict(state_dict)   # <-- the fixup happens here

load_state_dict -> _load_from_state_dict copies each incoming tensor into the live model parameter with param.copy_(input_param). When both sides are DTensors, DTensor.copy_ redistributes the source to the destination param's placement. The shape check upstream compares global logical shapes (identical for Shard(0) vs Shard(1)), so it passes, and the copy silently corrects the placement. The model param keeps its declared Shard(0) layout regardless of what from_hf produced.

The RL path never calls load_state_dict. In vllm_wrapper.py, from_hf's result goes straight into dtensor_to_plain_tensor_state_dict, which does a raw value.to_local() and hands the local shard to vLLM. There is no destination parameter to copy into, so nothing corrects a wrong placement. This PR's redistribute(...) before to_local() plays exactly the role that copy_ plays in the dcp path.

So the two callsites aren't inconsistent — the dcp path already gets the fixup for free via copy_; the RL path has to do it explicitly.

2. A simple, concrete example (using the PR's own numbers)

Take GatedDeltaNet's fused linear_attn.in_proj_qkv.weight under TP=2, with q/k/v each of size 2048, so the fused weight is [6144, 1024]. TT declares each split output (in_proj_q/k/v) as Shard(0) on the TP axis -> each rank should hold [1024, 1024].

How the fused tensor is laid out (Shard(0) over TP2):

in_proj_qkv  full [6144, 1024]
  rank0: rows    0 .. 3071   (local [3072, 1024])
  rank1: rows 3072 .. 6143   (local [3072, 1024])

Now value.split([2048, 2048, 2048], dim=0) — splitting along the sharded dim:

q piece: rows    0 .. 2047   -> entirely on rank0
k piece: rows 2048 .. 4095   -> STRADDLES the shard boundary at 3072!
v piece: rows 4096 .. 6143   -> entirely on rank1

The split boundaries (2048, 4096) don't align with the shard boundary (3072), so the output cannot be represented as Shard(0). DTensor falls back to Shard(1) (sharding dim 1 always aligns evenly under TP2):

in_proj_q  full [2048, 1024], placement Shard(1)
  each rank local: [2048, 512]   <-- WRONG
expected under Shard(0):
  each rank local: [1024, 1024]  <-- what the model wants

That is exactly the PR's reported [2048, 512] instead of [1024, 1024]. to_local() then bakes in that wrong slice, and vLLM gets garbage weights.

The fix redistributes in_proj_q back to its declared Shard(0) before to_local(), restoring the [1024, 1024] local slice. The same misalignment applies to the conv1d split, the vision qkv.chunk(3), and the MoE gate_up_proj.chunk(2, dim=-2) — any split/chunk/reshape whose boundary crosses the shard boundary.

---

_{· branch drisspg/rl-hf-shard-layout-fix}

[tianyu-l]

I think we won't have this problem if let FSDP wrap the generator model.

Please add a TODO for @andrewor14 to resolve