diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1d7bb3b622ad..3ef3ad264d18 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -980,6 +980,27 @@ def convert_and_load_state_dict_in_model( We build a mapping from the keys obtained by renaming each of the checkpoint keys according to the weight_mapping rules. Then we load the tensors into the model, applying any conversion operations as needed. + # Early return for FSDP CPU RAM efficient loading + # When FSDP is enabled with CPU RAM efficient loading, only rank 0 needs to load the weights + # for FULL_SHARD and SHARD_GRAD_OP strategies. Non-rank 0 processes will receive the sharded + # weights from rank 0 via FSDP's broadcast. This prevents all ranks from temporarily allocating + # model-sized CPU RAM. + # + # For HYBRID_SHARD and HYBRID_SHARD_ZERO2 strategies, all ranks need to load weights as there + # are multiple data parallel groups. + from .integrations import should_skip_non_rank0_weight_loading + + if should_skip_non_rank0_weight_loading() and load_config.hf_quantizer is None: + # Return empty loading info - the weights will be synced from rank 0 via FSDP + loading_info = LoadStateDictInfo( + missing_keys=set(model.state_dict().keys()), + unexpected_keys=set(), + mismatched_keys=set(), + conversion_errors={}, + error_msgs=[], + ) + return loading_info, disk_offload_index + The `param_name_to_load` will look like this: { "model.layers.0.attention.q.weight": # Notice here there is only the first key of the target keys @@ -1105,6 +1126,7 @@ def convert_and_load_state_dict_in_model( pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) + for original_key, tensor in state_dict: # 1. Rename the key according to all renaming pattern and optional weight converter patterns renamed_key, source_pattern = rename_source_key( diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py index 240268b75124..0954775b58ca 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -51,3 +51,39 @@ def is_fsdp_enabled(): ) return False + + +def should_skip_non_rank0_weight_loading(): + """ + Determine if non-rank0 processes should skip weight loading in FSDP with CPU RAM efficient loading. + + This function checks the FSDP sharding strategy to decide if it's safe to skip weight loading + on non-rank0 processes: + + - FULL_SHARD and SHARD_GRAD_OP: Safe to skip, as only rank 0 needs to load weights + - HYBRID_SHARD and HYBRID_SHARD_ZERO2: Not safe to skip, as multiple data parallel groups + may need weights + + Returns: + bool: True if non-rank0 should skip weight loading, False otherwise + """ + if not is_fsdp_enabled(): + return False + + if is_torch_available(): + import torch + + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return False + + # Check if we're on rank 0 - if so, we should always load weights + if int(os.environ.get("LOCAL_RANK", "-1")) == 0: + return False + + # Check sharding strategy from environment if available + # HYBRID_SHARD (4) and HYBRID_SHARD_ZERO2 (5) require all ranks to load weights + sharding_strategy = os.environ.get("FSDP_SHARDING_STRATEGY", "1") # Default to FULL_SHARD (1) + if sharding_strategy in ["4", "5"]: # HYBRID_SHARD or HYBRID_SHARD_ZERO2 + return False + + return True diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 87e8568b7fd1..0400bbf3c421 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -472,10 +472,36 @@ def remove_tied_weights_from_state_dict( def _load_parameter_into_model(model: "PreTrainedModel", param_name: str, tensor: torch.Tensor): - """Cast a single parameter or buffer `param_name` into the `model`, with value `tensor`.""" + """Cast a single parameter or buffer `param_name` into the `model`, with value `tensor`. + + Note on `_is_hf_initialized` flag: + This flag is used to mark parameters that have already been initialized or loaded with proper weights, + so they won't be re-initialized by `_initialize_missing_keys()`. This is particularly important in + FSDP (Fully Sharded Data Parallel) scenarios: + + 1. During normal model loading, parameters loaded from state_dict are automatically marked as initialized. + 2. For missing keys, parameters are created with empty tensors and should be initialized with proper + weight distributions via `_initialize_missing_keys()`. + 3. In FSDP with CPU RAM efficient loading (for FULL_SHARD and SHARD_GRAD_OP strategies), non-rank0 + processes create empty CPU tensors that will be synced from rank0. These should be marked as initialized + to avoid expensive random initialization that would be immediately overwritten by FSDP's weight sync. + + The `_is_hf_initialized` flag is checked in: + - `_initialize_missing_keys()`: Only initializes parameters without this flag + - `_move_missing_keys_from_meta_to_device()`: Moves missing keys (without the flag) to their devices + - `_tie_weights()`: Marks tied weights as initialized to avoid re-initialization + """ parent, param_type = get_module_from_name(model, param_name) if param_type in parent._parameters and not isinstance(tensor, nn.Parameter): tensor = nn.Parameter(tensor, requires_grad=tensor.is_floating_point()) + # In FSDP CPU RAM efficient loading mode (for FULL_SHARD and SHARD_GRAD_OP strategies), + # non-rank0 processes create empty tensors that will be synced from rank0. Mark these as + # initialized to avoid expensive random initialization that would be immediately overwritten + # during FSDP's weight synchronization. + from .integrations import should_skip_non_rank0_weight_loading + + if should_skip_non_rank0_weight_loading(): + tensor._is_hf_initialized = True # We need to use setattr here, as we set non-persistent buffers as well with this function (`load_state_dict` # does not allow to do it) setattr(parent, param_type, tensor)