Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
qgallouedec wants to merge 7 commits into
base: main
Choose a base branch
from
+14 −0
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
fced978
Handle empty FSDP shards in sharded loading: crash and hang fix
qgallouedec Aug 23, 2026
385944d
Merge branch 'main' into fix-empty-fsdp-shards
qgallouedec Aug 23, 2026
65edf2c
Trim the comments to one line each
qgallouedec Aug 23, 2026
2ed891d
Merge branch 'main' into fix-empty-fsdp-shards
qgallouedec Aug 25, 2026
6f78144
Merge branch 'main' into fix-empty-fsdp-shards
qgallouedec Aug 27, 2026
ff5570e
Merge branch 'main' into fix-empty-fsdp-shards
qgallouedec Sep 3, 2026
8aa5433
Merge branch 'main' into fix-empty-fsdp-shards
qgallouedec Sep 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/transformers/core_model_loading.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -976,6 +976,10 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
# Add them to the new dictionary
collected_tensors[key] = tensors

if any(len(tensors) == 0 for tensors in collected_tensors.values()):
# Uneven FSDP sharding left this rank an empty shard: nothing to load, and its pre-sharded empty local tensor is already correct
raise SkipParameters()

return collected_tensors

def was_used(self) -> bool:
Expand Down
10 changes: 10 additions & 0 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4742,6 +4742,16 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None:
pass # may happen when handling pre-quantized weights
self._is_hf_initialized = True

if getattr(self, "_device_mesh", None) is not None:
# Empty local shards have nothing to initialize; without the mark, running _init_weights on them issues collectives the other ranks never join (hang)
import itertools

from torch.distributed.tensor import DTensor

for param_or_buffer in itertools.chain(self.parameters(), self.buffers()):
if isinstance(param_or_buffer, DTensor) and param_or_buffer._local_tensor.numel() == 0:
param_or_buffer._is_hf_initialized = True

Comment on lines +4745 to +4754

@ArthurZucker ArthurZucker Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@qgallouedec qgallouedec Sep 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ArthurZucker

Tried this and couldn't make it work yet. Routing the empty case into set_param_for_module means reconstructing the target names inside convert(), where the collected keys are still source patterns (mlp.experts.*.gate_proj.weight), so the substring trick the normal path uses doesn't apply.
compiled_sources gets you the prefix/suffix, but the same short-circuit is then needed in both WeightRenaming.convert and WeightConverter.convert, since materialize_tensors is shared.

Still, I implemented both and it still hangs on the repro, while the version in the PR loads (main hangs / PR loads / restructure hangs), same 2-expert + fsdp_size=4 fixture, same run.

torchrun --nproc_per_node 4 repro_empty_shard.py

# repro_empty_shard.py
import os
from datetime import timedelta

import torch

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

world_size = int(os.environ["WORLD_SIZE"])
torch.distributed.init_process_group("nccl", timeout=timedelta(minutes=3))
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

model = AutoModelForCausalLM.from_pretrained(
    "katuni4ka/tiny-random-qwen3moe",  # 4 experts, ~1 MB
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(fsdp_size=world_size),
)
torch.distributed.barrier()
if int(os.environ["RANK"]) == 0:
    shard = model.model.layers[1].mlp.experts.gate_up_proj
    print(f"LOADED OK (rank0 local expert shard: {tuple(shard._local_tensor.shape)})", flush=True)

Am I missing something?

# This will only initialize submodules that are not marked as initialized by the line above.
if is_deepspeed_zero3_enabled() and not is_quantized:
import deepspeed
Expand Down
Loading