fix(llama3): make q/k HF-conversion permute DTensor-safe

[Kaif10]

Llama3StateDictAdapter._permute / _reverse_permute do a head-splitting view() on the q/k projection weights during HF checkpoint conversion. In the live save/load path those weights are DTensors that FSDP shards along dim 0, the exact dim the view unflattens, so an FSDP degree that does not evenly divide the head count raises Cannot unflatten unevenly sharded tensor. For llama3-8B (n_kv_heads=8) this breaks HF export above an 8-way FSDP degree, i.e. any typical multi-node run.

Fix

muse_glimmer's adapter hit the identical failure and fixed it locally with a _dtensor_safe decorator (redistribute to Replicate, run the reshape on the full local tensor, redistribute back). This hoists that decorator into the shared torchtitan/protocols/state_dict_adapter.py as dtensor_safe, applies it to llama3's two permute helpers, and points muse_glimmer at the shared implementation instead of its own copy.

Net diff is +41/-39: llama3 gains two decorators, muse_glimmer loses its 30-line duplicate.

Testing

Two new CPU unit tests in test_state_dict_adapter.py:

• test_permute_does_not_raise_when_fsdp_degree_exceeds_head_count reproduces the production failure with a fake process group (32-way mesh, 16 heads) and asserts to_hf no longer raises.
• test_permute_roundtrip_preserves_dtensor_values_and_placement uses a real gloo process group to confirm the decorator does not alter values or placement in the common (evenly-sharded) case.

Both fail on main — the first with the exact production RuntimeError — and pass with this change. Full test_state_dict_adapter.py file passes (7 tests). pre-commit (flake8, ufmt, pydoclint, codespell, license headers) is clean on the changed files.

This is a checkpoint-conversion bug fix, not a training-numerics change, so no loss/grad_norm comparison applies.

[meta-cla]

Hi @Kaif10!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[tianyu-l]

@claude

Validate the issue mentioned in the PR. Give a simple and concrete example where things are broken.

Would either #4797 or #4787 solve this?

[claude]

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

---

Validation of the reported bug

• Understand the reported bug in _permute / _reverse_permute
• Construct a simple, concrete example where it breaks
• Check whether PR fix(llama3): make q/k HF-conversion permute DTensor-safe #4797 (this PR) fixes it
• Check what PR [RL] Restore model shard layouts after HF conversion #4787 does and whether it fixes it — see note on access below

---

The bug is real

_permute/_reverse_permute unflatten dim 0 of the q/k weight with a view(n_heads, dim1 // n_heads // 2, 2, dim2). In the live save/load path these weights arrive as DTensors that FSDP shards along dim 0 — the exact dim being unflattened. DTensor can only unflatten a sharded dim when the leading factor of the split (n_heads / n_kv_heads) can be evenly distributed across the shards. Once the FSDP degree exceeds the head count, the per-rank shard is smaller than one head group and the op raises Cannot unflatten unevenly sharded tensor.

Concrete broken example (llama3-8B wk, 16-way FSDP)

llama3-8B: dim = 4096, n_heads = 32, n_kv_heads = 8, head_dim = 128 (confirmed in torchtitan/models/llama3/__init__.py:249-251).

The wk weight is [head_dim * n_kv_heads, dim] = [1024, 4096]. On a 16-way FSDP mesh it is Shard(0) → 64 rows per rank. to_hf then calls:

# key_value_dim = head_dim * n_kv_heads = 1024
self._permute(wk, n_kv_heads=8, dim1=1024, dim2=4096)
#   -> w.view(8, 1024 // 8 // 2, 2, 4096) == w.view(8, 64, 2, 4096)

To build that view, DTensor must unflatten the Shard(0) dim (global size 1024) into a leading factor of 8. But the tensor is split 16 ways, each shard holding 64 rows while one head group is 1024 / 8 = 128 rows — the shard boundary lands in the middle of a head group. The shard can't map onto the leading 8, so:

RuntimeError: Cannot unflatten unevenly sharded tensor

So HF export/import breaks for llama3-8B (8 KV heads) at any FSDP degree > 8, i.e. essentially every multi-node run. The debugmodel repro in the new test uses the same mechanism (16 heads, [256,256] weight, 32-way mesh → 8 rows/rank < 16 rows/head).

Does PR #4797 (this PR) solve it? — Yes.

dtensor_safe (torchtitan/protocols/state_dict_adapter.py:25) detects a DTensor, redistribute(...Replicate()) to gather the full tensor, runs the view on the full local tensor (where the unflatten is always valid), rewraps with from_local, and redistributes back to the original placements. It's independent of whether the FSDP degree divides the head count, and the non-DTensor path is unchanged. The two new CPU tests exercise exactly this (fake 32-way PG for the raise, gloo PG for value/placement round-trip). This is a checkpoint-conversion fix with no training-numerics impact, so no loss comparison is needed — consistent with the PR description.

One thing worth confirming with reviewers: gathering to full Replicate() materializes the whole q/k weight on each rank during conversion. For 8B that's fine; for very large models it's a memory spike, but it matches what muse_glimmer already did, so #4797 doesn't regress anything.

Does PR #4787 solve it?

I was not able to inspect #4787 in this run — network access (gh, curl, git fetch) and WebFetch are all blocked in this sandbox, so I can only reason about the local #4797 branch. Rather than guess at its contents, here is the concrete criterion that decides it:

#4787 fixes this bug iff it removes the head-splitting view from the sharded path, by one of:

1. Replicating/gathering q/k to a full tensor before permuting (same idea as dtensor_safe, possibly centralized in the checkpointer/save path) — would fix it, and more broadly than a per-adapter decorator.
2. Reordering rows without unflattening the sharded dim (e.g. an index/gather permutation that operates on Shard(0) directly) — would fix it without an all-gather.

It does not fix this bug if it only touches an unrelated code path (routing, a different adapter, the even-division case, etc.).

If #4787 takes approach (1) at the save-path level, it would likely supersede this PR (making the per-adapter decorator redundant); if it takes (2), it's a genuinely different and lower-memory strategy worth comparing. @tianyu-l, if you can point me at #4787's diff (or approve network access for this action), I'll confirm which of the two it is and whether the two PRs overlap or should be merged into one approach.
· branch fix/llama3-dtensor-safe-permute

[tianyu-l]

@claude

I see. More questions:

1. What is the life-cycle of such all-gathered replicate tensors? My understanding is that for loading checkpoint from HF format, the gathering happens in to_hf, and the gathered tensor will be there until from_hf consumes them to restore to titan model definition & sharding.
2. Can we replace this all-gather with all-to-all to let it shard on other tensor dims, so that peak memory won't increase, at the potential cost of invoking another collective during from_hf as well.

[claude]

Claude finished @tianyu-l's task in 1m 43s —— View job

---

Answers

1. Lifecycle of the replicated tensor

The replicated tensor is transient and scoped to a single _permute/_reverse_permute call -- it does not persist from to_hf to from_hf. Looking at the decorator (torchtitan/protocols/state_dict_adapter.py:38-50):

replicated = w.redistribute(..., placements=[Replicate()] * mesh.ndim)  # all-gather
local = fn(self, replicated.to_local(), *args, **kwargs)                 # reshape full local
out = DTensor.from_local(local, mesh, [Replicate()] * mesh.ndim, ...)
return out.redistribute(device_mesh=mesh, placements=placements)         # reshard back

The last line redistributes back to the original (sharded) placements before returning, and redistribute is eager, so the full-size replicated/out locals are freed when the wrapper returns. Two consequences worth being precise about:

• to_hf and from_hf are separate operations (save vs. load, typically separate processes). Within a single conversion only one of them runs, and each of its _permute calls gathers + reshards internally. So there's no "gather in to_hf, consume in from_hf" hand-off -- the mental model in the question doesn't quite match the code.
• Peak extra memory is one fully-replicated q or k weight at a time, held only for the duration of that one helper call -- it is not accumulated across layers, and the returned state-dict entries are sharded exactly as they came in.

So for 8B this spike is tiny (wq is [4096,4096], wk is [1024,4096]); it only becomes a concern for very large hidden dims at high FSDP degree.

2. All-to-all instead of all-gather -- yes, this works and is a good optimization

The permute only unflattens dim 0 (view(n_heads, dim1//n_heads//2, 2, dim2)). It is a pure row-reordering that is identical for every column. So instead of gathering to Replicate (full tensor on every rank), we can reshard Shard(0) -> Shard(1) (an all-to-all), which keeps dim 0 fully local -- making the unflatten valid -- while sharding the innermost dim2 (untouched by the reshape). Peak memory stays at the original ~1/N shard size, no replication spike.

Result equivalence holds because permuting a column-shard and then gathering columns is the same as permuting the full tensor -- the row permutation commutes with column sharding.

A nice detail: both helpers already default dim1 = w.shape[0] and dim2 = w.shape[1]. Under Shard(1) the local tensor is [dim1, dim2/N], so if the wrapper lets those defaults take over (rather than threading the explicit global dim2), the reshape "just works" on the local shard with no extra bookkeeping.

Caveats to weigh before switching:

• Explicit dim2 for wk. to_hf/from_hf currently pass the global dim as dim2 (llama3 state_dict_adapter.py:99,144). The all-to-all variant must use the local dim2, so the decorator would drop/override that argument and rely on the local shape.
• Multi-dim meshes (FSDP + TP). q/k are Shard(0) on both the dp_shard and TP axes. Redistributing to [Shard(1)] * mesh.ndim reshards every row-sharded axis onto dim 1; that's valid here since these weights are never Partial, but the placement construction needs to handle >1D meshes rather than assuming a single FSDP axis.
• Extra collective on from_hf. As you note, this adds an all-to-all on the load path too (currently from_hf also all-gathers, so it's an all-to-all replacing an all-gather rather than a net-new collective -- generally cheaper in both bandwidth and memory).

My recommendation: the current all-gather is a correct, minimal first fix and I'd keep it as-is to unblock 8B multi-node export. The Shard(0) -> Shard(1) all-to-all is the right follow-up when the replication spike actually matters (large models / high degree). If you'd like, I can implement the all-to-all variant in this PR or draft it as a follow-up.

_{· branch fix/llama3-dtensor-safe-permute}