Handle empty FSDP shards in sharded loading: crash and hang fix

[qgallouedec]

Fixes loading models whose FSDP sharding leaves a rank an empty local shard, which today crashes some ranks and silently hangs the rest.

Uneven sharding makes this easy to hit: GLM-4.6 has 160 experts; ep=8 leaves 20 per EP rank, and chunking 20 over fsdp=8 gives [3,3,3,3,3,3,2,0] (every fsdp-index-7 rank stores zero expert rows). Then:

1. MergeModulelist receives zero pieces and torch.stack([]) raises stack expects a non-empty TensorList (fatal at the end of loading on those ranks).
2. Their params are never marked _is_hf_initialized, so _initialize_missing_keys runs _init_weights on them, whose first DTensor RNG op is a mesh-wide collective the fully-loaded ranks never join.

Verified live at 64 ranks: the stuck set was ranks 56–63 (fsdp index 7).

Real-checkpoint repro: Mixtral-8x7B on 5 GPUs

Modern MoEs have too many experts to hit this at small scale (which is why it went unnoticed until 64-rank meshes), but Mixtral-8x7B has only 8: with DistributedConfig(fsdp_size=5) the chunking is [2,2,2,2,0]. On main, loading hangs silently until the NCCL watchdog aborts the job (Watchdog caught collective operation timeout ... SeqNum=1); with this PR it loads cleanly.

Minimal repro (4 GPUs, 0.2M params, ~30 s)

A 2-expert Qwen3-MoE toy with fsdp_size=4 leaves ranks 2–3 empty. On main: the empty ranks raise (torch.cat(): expected a non-empty list of Tensors: Qwen's gate/up fusion hits Concatenate first; other models hit torch.stack in MergeModulelist) while the loaded ranks die on the NCCL watchdog. With this PR: loads cleanly.

# torchrun --nproc_per_node 4 repro.py  (any tiny MoE with num_experts < fsdp_size works)
import os, torch
from datetime import timedelta
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.distributed import DistributedConfig

if int(os.environ["RANK"]) == 0 and not os.path.isdir("/tmp/tiny_2experts"):
    cfg = AutoConfig.for_model(model_type="qwen3_moe", num_hidden_layers=2, hidden_size=64,
        intermediate_size=128, moe_intermediate_size=64, num_experts=2, num_experts_per_tok=2,
        num_attention_heads=4, num_key_value_heads=2, head_dim=16, vocab_size=1000,
        decoder_sparse_step=1, mlp_only_layers=[])
    AutoModelForCausalLM.from_config(cfg).save_pretrained("/tmp/tiny_2experts")
torch.distributed.init_process_group(backend="nccl", timeout=timedelta(minutes=3))
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
torch.distributed.barrier()
AutoModelForCausalLM.from_pretrained("/tmp/tiny_2experts", dtype=torch.bfloat16,
    distributed_config=DistributedConfig(fsdp_size=4))
torch.distributed.barrier()
print("LOADED OK")

The fix

• Conversion pipeline: when every piece of a parameter was dropped by the sharding operation, skip the mapping
• _initialize_missing_keys: mark DTensor params with empty locals as initialized

Validated: the previously-freezing configuration (GLM-4.6, ep=8 × fsdp=8, 64 ranks) loads to completion with the fix.

Found while training 100B–753B MoEs with FSDP2 × EP (#48204, which carries the same commit).

[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: 32874850998:1
Result: failure | Jobs: 11 | Tests: 179,951 | Failures: 0 | Duration: 6h 9m

[ArthurZucker]

It does make sense to have, I don't think it fixed your issue at core tho, we should expect a better expert sharding to leverage all GPUs no? (at least putting less expert on other and reducing idle?)

[ArthurZucker]

the better fix is to go in https://github.com/huggingface/transformers/blob/qwen3_vl_moe_tp_plan/src/transformers/core_model_loading.py#L1388-L1388 and if a tensor is empty + tp + number 0 (its not missing so you should go into the set param) -> set the flag