Repository: huggingface/transformers
Pull request: fix zero3 init config

# What does this PR do?

Supersedes https://github.com/huggingface/transformers/pull/43847

When using zero3 + from_config, the model was incorrectly initialized as we were not gathering the params. Added a test also. 

cc @tohtana 

Changed files:
filename: src/transformers/integrations/deepspeed.py
patch: @@ -290,6 +290,41 @@ def deepspeed_config():
         return None
 
 
+def initialize_weights_zero3(model):
+    """
+    DeepSpeed ZeRO-3 variant of `PreTrainedModel.initialize_weights`. Mirrors the `smart_apply`
+    dispatch logic but gathers each module's partitioned parameters before calling
+    `_initialize_weights`, so initialization operates on full tensors instead of empty shards.
+    Only rank 0 performs the actual init.
+    """
+    import deepspeed
+    import torch
+
+    from ..initialization import guard_torch_init_functions
+    from ..modeling_utils import PreTrainedModel
+
+    is_remote_code = model.is_remote_code()
+
+    def _apply_zero3(model_or_module, fn):
+        for child in model_or_module.children():
+            if isinstance(child, PreTrainedModel):
+                _apply_zero3(child, child._initialize_weights)
+            else:
+                _apply_zero3(child, fn)
+
+        params = list(model_or_module.parameters(recurse=False))
+        if params:
+            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
+                if deepspeed.comm.get_rank() == 0:
+                    fn(model_or_module, is_remote_code)
+        else:
+            fn(model_or_module, is_remote_code)
+
+    with torch.no_grad():
+        with guard_torch_init_functions():
+            _apply_zero3(model, model._initialize_weights)
+
+
 def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     """
     Apply weight conversions (renaming and merging/splitting operations) to a state dict.

filename: src/transformers/modeling_utils.py
patch: @@ -1462,18 +1462,34 @@ def _from_config(cls, config, **kwargs):
         if dtype is not None:
             init_contexts.append(local_torch_dtype(dtype, cls.__name__))
 
-        if is_deepspeed_zero3_enabled() and not _is_quantized and not _is_ds_init_called:
+        needs_zero3_init = is_deepspeed_zero3_enabled() and not _is_quantized and not _is_ds_init_called
+        if needs_zero3_init:
             logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
             # this immediately partitions the model across all gpus, to avoid the overhead in time
             # and memory copying it on CPU or each GPU first
             import deepspeed
 
-            init_contexts.extend([deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()])
+            init_contexts.extend(
+                [
+                    init.no_init_weights(),
+                    deepspeed.zero.Init(config_dict_or_path=deepspeed_config()),
+                    set_zero3_state(),
+                ]
+            )
 
         # Instantiate the model
         with ContextManagers(init_contexts):
             model = cls(config, **kwargs)
 
+        # Under ZeRO-3, parameters were partitioned into empty tensors during construction,
+        # so weight init was suppressed. Re-initialize using the ZeRO-3 variant which gathers
+        # each module's parameters before init to avoid OOM on large models.
+        if needs_zero3_init:
+            from .integrations.deepspeed import initialize_weights_zero3
+
+            initialize_weights_zero3(model)
+            model.tie_weights()
+
         return model
 
     @property

filename: tests/deepspeed/test_deepspeed.py
patch: @@ -274,6 +274,59 @@ def test_init_zero3_fp16(self):
                     AutoModel.from_pretrained(T5_TINY)
         self.assertNotIn("Detected DeepSpeed ZeRO-3", cl.out)
 
+    @require_torch_accelerator
+    def test_from_config_zero3_weight_init(self):
+        # test that from_config() correctly initializes weights under zero3
+        # (regression test: without the fix, _init_weights runs on partitioned empty tensors
+        # and custom initialization is silently skipped)
+        import deepspeed
+        import torch
+
+        from transformers import AutoConfig, AutoModelForCausalLM
+
+        ds_config = {
+            "train_batch_size": 1,
+            "zero_optimization": {
+                "stage": 3,
+            },
+            "bf16": {
+                "enabled": True,
+            },
+        }
+
+        config = AutoConfig.from_pretrained(GPT2_TINY)
+
+        # 1. Baseline: from_config without DeepSpeed, in bf16 to match ZeRO-3 dtype
+        torch.manual_seed(42)
+        baseline_model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16)
+        baseline_std = baseline_model.transformer.h[0].attn.c_attn.weight.data.float().std().item()
+        del baseline_model
+
+        # 2. ZeRO-3: from_config with DeepSpeed
+        torch.manual_seed(42)
+        HfDeepSpeedConfig(ds_config)
+
+        with mockenv_context(**self.dist_env_1_gpu):
+            model = AutoModelForCausalLM.from_config(config)
+
+        param = model.transformer.h[0].attn.c_attn.weight
+        with deepspeed.zero.GatheredParameters([param]):
+            zero3_std = param.data.float().std().item()
+
+        # Weight std should be in the same ballpark between baseline and ZeRO-3.
+        # The seeds produce slightly different values due to different init ordering,
+        # but delta=0.3 catches the buggy case (ratio ~0.58, i.e. 42% off).
+        ratio = zero3_std / baseline_std
+        self.assertAlmostEqual(
+            ratio,
+            1.0,
+            delta=0.3,
+            msg=(
+                f"ZeRO-3 from_config() weight init diverges from baseline: "
+                f"baseline_std={baseline_std:.6f}, zero3_std={zero3_std:.6f}, ratio={ratio:.4f}"
+            ),
+        )
+
     def test_init_zero3_missing_params(self):
         # test that zero.Init() for missing parameters works correctly under zero3
         import deepspeed