Fix FSDP_CPU_RAM_EFFICIENT_LOADING (#43749)

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(

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

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)