Make clearer about zero_init requirements

src/transformers/integrations/deepspeed.py

@@ -17,6 +17,7 @@
 import copy
 import importlib.metadata as importlib_metadata
 import importlib.util
+import os
 import weakref
 from functools import partialmethod
 
@@ -287,6 +288,14 @@ def is_deepspeed_zero3_enabled():
         return _hf_deepspeed_config_weak_ref().is_zero3()
     else:
         return False
+
+def is_deepspeed_zero3_enabled_accelerate():
+    "Same as `is_deepspeed_zero3_enabled`, but checks only for accelerate setups"
+    accelerate_zero_stage = int(os.environ.get("ACCELERATE_DEEPSPEED_ZERO_STAGE", -1))
+    accelerate_zero_init = os.environ.get("ACCELERATE_DEEPSPEED_ZERO3_INIT", "0")
+    # This only gets triggered passively if the user launches code with a configured
+    # `accelerate launch` without making `TrainingArguments`
+    return accelerate_zero_stage != -1 and accelerate_zero_init != "0"
 
 
 def deepspeed_config():

src/transformers/modeling_utils.py

@@ -42,7 +42,7 @@
 from .configuration_utils import PretrainedConfig
 from .dynamic_module_utils import custom_object_save
 from .generation import GenerationConfig, GenerationMixin
-from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled
+from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled, is_deepspeed_zero3_enabled_accelerate
 from .pytorch_utils import (  # noqa: F401
     Conv1D,
     apply_chunking_to_forward,
@@ -1313,6 +1313,12 @@ def _from_config(cls, config, **kwargs):
         )
 
         if is_deepspeed_zero3_enabled():
+            if not is_deepspeed_zero3_enabled_accelerate():
+                raise ValueError(
+                    "Detected that you want to use `zero-3` Init, but the environment "
+                    "has not been setup yet. Please create `TrainingArguments` before "
+                    "initializing the model."
+                )
             import deepspeed
 
             logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
@@ -3387,6 +3393,12 @@ def from_pretrained(
         init_contexts = [no_init_weights(_enable=_fast_init)]
 
         if is_deepspeed_zero3_enabled() and not is_quantized:
+            if not is_deepspeed_zero3_enabled_accelerate():
+                raise ValueError(
+                    "Detected that you want to use `zero-3` Init, but the environment "
+                    "has not been setup yet. Please create `TrainingArguments` before "
+                    "initializing the model."
+                )
             import deepspeed
 
             logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")

src/transformers/training_args.py

@@ -503,6 +503,11 @@ class TrainingArguments:
             evolve in the future. The value is either the location of DeepSpeed json config file (e.g.,
             `ds_config.json`) or an already loaded json file as a `dict`"
 
+            <Tip warning={true}>
+                If enabling any Zero-init, make sure that your model is not initialized until
+                *after* initializing the `TrainingArguments`, else it will not be applied.
+            </Tip>
+
         accelerator_config (`str`, `dict`, or `AcceleratorConfig`, *optional*):
             Config to be used with the internal `Accelerator` implementation. The value is either a location of
             accelerator json config file (e.g., `accelerator_config.json`), an already loaded json file as `dict`,