From fced97851234728795a0ab57b0570448cffefdfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sun, 23 Aug 2026 01:52:23 +0000 Subject: [PATCH 1/3] Handle empty FSDP shards in sharded loading: crash and hang fix Uneven FSDP sharding can assign a rank an EMPTY local shard (e.g. 2 experts chunked over fsdp=4 leave the last ranks zero rows; at larger scale, GLM-4.6's 20 local experts over fsdp=8 leave fsdp rank 7 empty). Two things then go wrong while loading: 1. The conversion ops receive zero collected pieces for the parameter and raise (torch.cat / torch.stack of an empty list) - fatal at the end of loading on those ranks. Fixed by skipping the mapping when every piece was dropped by the sharding operation: the pre-sharded empty local tensor installed at init is already correct. 2. Those 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 - while fully-loaded ranks skip it: mismatched collectives, and the group hangs silently (0% GPU, all ranks in R state). Fixed by marking empty-local DTensors before the sweep - an empty shard has nothing to initialize. Reproduces on 4 GPUs with a 0.2M-param toy (2 experts, fsdp_size=4): crash on the empty ranks, watchdog abort on the rest. At scale this froze three multi-hour 357B training runs (ep=8 x fsdp=8) before being root-caused. --- src/transformers/core_model_loading.py | 7 +++++++ src/transformers/modeling_utils.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 4c4733d14ad1..de59292fb56c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -960,6 +960,13 @@ 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()): + # Every piece of this parameter was dropped by the sharding operation: this rank owns + # none of it (uneven FSDP sharding can assign a rank an EMPTY shard, e.g. 2 experts + # chunked over fsdp=4). The pre-sharded empty local tensor installed at init is already + # correct, and running the conversion ops on empty lists would raise. + raise SkipParameters() + return collected_tensors def was_used(self) -> bool: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7259a89f36f7..6458f070517c 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4793,6 +4793,21 @@ 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: + # Uneven FSDP sharding can assign a rank an EMPTY local shard (e.g. 20 local experts + # chunked over fsdp=8 leaves the last rank zero rows). The loader then has nothing to + # load into that parameter, so it is never marked initialized, and this rank would run + # `_init_weights` on it below - whose first DTensor RNG op is a mesh-wide collective - + # while fully-loaded ranks skip it: mismatched collectives, and the whole group hangs. + # An empty shard has nothing to initialize; mark it. + 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 + # 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 From 65edf2c6a473e947a026a146b25379e5146e8c6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Sun, 23 Aug 2026 19:30:24 +0000 Subject: [PATCH 2/3] Trim the comments to one line each --- src/transformers/core_model_loading.py | 5 +---- src/transformers/modeling_utils.py | 7 +------ 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index de59292fb56c..5c7509bbca12 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -961,10 +961,7 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: collected_tensors[key] = tensors if any(len(tensors) == 0 for tensors in collected_tensors.values()): - # Every piece of this parameter was dropped by the sharding operation: this rank owns - # none of it (uneven FSDP sharding can assign a rank an EMPTY shard, e.g. 2 experts - # chunked over fsdp=4). The pre-sharded empty local tensor installed at init is already - # correct, and running the conversion ops on empty lists would raise. + # 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 diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 6458f070517c..94bc7f1b277f 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4794,12 +4794,7 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None: self._is_hf_initialized = True if getattr(self, "_device_mesh", None) is not None: - # Uneven FSDP sharding can assign a rank an EMPTY local shard (e.g. 20 local experts - # chunked over fsdp=8 leaves the last rank zero rows). The loader then has nothing to - # load into that parameter, so it is never marked initialized, and this rank would run - # `_init_weights` on it below - whose first DTensor RNG op is a mesh-wide collective - - # while fully-loaded ranks skip it: mismatched collectives, and the whole group hangs. - # An empty shard has nothing to initialize; mark it. + # 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 From c64443d6131668515520a289a8a95a0edc44cf53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= Date: Mon, 7 Sep 2026 22:09:38 +0000 Subject: [PATCH 3/3] Read _device_mesh directly and hoist the itertools import `DistributedMixin` declares `_device_mesh = None` and `PreTrainedModel` inherits it, so the attribute always resolves. `modeling_utils` already imports from `itertools` at module level. --- src/transformers/modeling_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index f669ee83dd5a..bc58d139f74e 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -27,7 +27,7 @@ from contextlib import contextmanager from dataclasses import dataclass, field from functools import partial, wraps -from itertools import cycle +from itertools import chain, cycle from threading import Thread from typing import TYPE_CHECKING, Any, TypeVar, get_type_hints, overload from zipfile import is_zipfile @@ -4742,13 +4742,11 @@ 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: + if self._device_mesh 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()): + for param_or_buffer in 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