Loading: convert transposed and RoPE-permuted weights before taking the DTensor shard

[qgallouedec]

When a parameter is loaded as a DTensor (tp_plan, or FSDP2 at load time), the loader slices the checkpoint tensor to this rank's shard first and runs the conversion ops on the slice. Transpose and PermuteForRope move elements across dims, so when the sharded dim is one they touch, the shard of the source is not the shard of the converted tensor:

• Transpose: a router stored as (hidden, experts), sharded on dim 0 over 2 ranks, comes out as (experts, hidden/2) against a (experts/2, hidden) shard. Loud: mismatched-keys error.
• PermuteForRope: takes the head size from tensor.shape[0], i.e. from the shard. Silent: q/k permuted with the wrong head size, no error.

Repro (Kimi K2.5 vision q/k go through PermuteForRope; needs FSDP2 at load, #48205):

First:

import torch
from transformers import Kimi_K25ForConditionalGeneration
from tests.models.kimi_k25.test_modeling_kimi_k25 import Kimi_K25VisionText2TextModelTester

model = Kimi_K25ForConditionalGeneration(Kimi_K25VisionText2TextModelTester(None).get_config())
model.save_pretrained(path, save_original_format=True)  # q/k written un-permuted, permuted back at load
torch.save(model.state_dict(), "/tmp/kimi_tiny/plain.pt")

# torchrun --nproc_per_node 2 repro.py
import torch
from transformers import Kimi_K25ForConditionalGeneration
from transformers.distributed import DistributedConfig

model = Kimi_K25ForConditionalGeneration.from_pretrained(path, distributed_config=DistributedConfig(fsdp_size=2))
plain = torch.load("/tmp/kimi_tiny/plain.pt")
for name, param in model.named_parameters():
    if "vision" in name and ("q_proj.weight" in name or "k_proj.weight" in name):
        print(name, (param.full_tensor().cpu() - plain[name]).abs().max().item())

Before (no error raised):

model.vision_tower.layers.0.attn.q_proj.weight 0.0897
model.vision_tower.layers.0.attn.k_proj.weight 0.1171
model.vision_tower.layers.1.attn.q_proj.weight 0.0956
model.vision_tower.layers.1.attn.k_proj.weight 0.0965

After:

model.vision_tower.layers.0.attn.q_proj.weight 0.0
model.vision_tower.layers.0.attn.k_proj.weight 0.0
model.vision_tower.layers.1.attn.q_proj.weight 0.0
model.vision_tower.layers.1.attn.k_proj.weight 0.0

Fix:

• shards_after_conversion says whether a converter has one of these ops on a sharded dim;
• for those the full tensor is converted, then sharded with the same DtensorShardOperation.

Everything else keeps shard-on-read (fused experts, and transposes on unsharded dims such as Qwen3-VL-MoE's)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33835608545:1
Result: success | Jobs: 16 | Tests: 185,859 | Failures: 0 | Duration: 15h 29m

[vasqu]

cc @3outeille could you take a first look here?

[vasqu]

Just one initial design question but letting Ferdinand comment first on the logic

[vasqu]

Wondering whether we might move this as general function within converter/renaming and override where we need it

Would make it easier to scale tbh

[qgallouedec]

Yes I wasn't sure. I went for simplicity instead of generality as a first step, but happy to generalize

[vasqu]

I just fear the function exploding at some point 😢 and it makes custom extensions simpler