diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 580f46849c3c..ea9e8b958a65 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -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: diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 61f338e82057..bc428d03ec54 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 @@ -4766,6 +4766,14 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None: pass # may happen when handling pre-quantized weights self._is_hf_initialized = True + 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) + from torch.distributed.tensor import DTensor + + 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 + # 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