Yell at the user if zero-3 init wasn't performed, but expected to have been done
Author: muellerzr
Pull request: yes

# What does this PR do?

This PR adds a guard if the user has selected Zero-3 but the init was *not* called. This can easily happen when a user does not instantiate their `TrainingArguments()` before calling the model init. Will now raise a very helpful error message telling them to instantiate their `TrainingArguments` prior to init if we've found this has happened

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc @philschmid @LysandreJik 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32299). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts " why does this need to be done this way" is because we need to initialize Zero-3's init and enable it *before* creating a model. This can only happen automatically when creating the `TrainingArguments`. Otherwise there is no other way for us to catch that we want to use Zero, and this can't be done post model init. It's just how deepspeed works, and we have to work around it. This has been a thing for awhile, however no warning was ever entered so users were just frustrated when it seemed like zero-init wasn't happening

@amyeroberts another way to explain it:

There is logic in `from_pretrained` that utilizes the `with ZeroInit()`... logic when we detect that DeepSpeed has been initialized, which only happens when certain criteria are met, aka a DeepSpeed object has been created (which is persistent) in `transformers`. This gets triggered if and only if we enable it via the `TrainingArguments` (or if `accelerate launch` is used, but I do **not** want to rely on env variables that strongly). 

Think of it akin to when we init via `device="meta"` during big model loading, but because we need to know deepspeed exists and configured, this can *only* happen after we initialize the `TrainingArguments` and pass in something for the deepspeed plugin. As a result, we are stuck with this change in API (that again, has always been a thing) where we need to be careful about *when* model init happens

Thanks for taking the time to write this out @muellerzr! I'm still a bit confused about the process here and I might be missing something (including maybe some context for the normal workflow / design desicions). My understanding now is that yes, it depends on global state and that is created through TrainingArguments. 

> This gets triggered if and only if we enable it via the TrainingArguments (or if accelerate launch is used, but I do not want to rely on env variables that strongly).

Why is this the case? Couldn't we create a deepspeed object outside of TrainingArguments?  

> There is logic in from_pretrained that utilizes the with ZeroInit()... logic when we detect that DeepSpeed has been initialized

Is it necessary to use the global scope in `from_pretrained` to infer this? Couldn't we do something similar to `device_map` where we tell the model it's to be initialized for deepspeed? 

> Is it necessary to use the global scope in from_pretrained to infer this? Couldn't we do something similar to device_map where we tell the model it's to be initialized for deepspeed?

We could, yes. That's potentially a big breaking change though from something that's been a thing since the deepspeed integration was introduced 😓 

Let me bring up the LOC for it to make a bit more sense flow-wise:

In the `.from_pretrained` and `.from_config` implementation, we get to a point where we need to see if we're doing init_empty_weights or zero init:

```python
        # Instantiate model.
        init_contexts = [no_init_weights(_enable=_fast_init)]

        if is_deepspeed_zero3_enabled() and not is_quantized:
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
            init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config())] + init_contexts
        elif low_cpu_mem_usage:
            init_contexts.append(init_empty_weights())

        config = copy.deepcopy(config)  # We do not want to modify the config inplace in from_pretrained.
        config = cls._autoset_attn_implementation(
            config, use_flash_attention_2=use_flash_attention_2, torch_dtype=torch_dtype, device_map=device_map
        )

        with ContextManagers(init_contexts):
            # Let's make sure we don't run the init function of buffer modules
            model = cls(config, *model_args, **model_kwargs)
```

Because the deepspeed integration is meant to be a *zero code change* integration, where the user should only ever have to pass in a deepspeed plugin config to the `TrainingArguments`, we need to check for that `is_deepspeed_zero3_enabled` flag. 

This is done by checking a weakref to the deepspeed config, which is noted as:

```python
# keep the config object global to be able to access it anywhere during TrainingArguments life-cycle
_hf_deepspeed_config_weak_ref = None
```
This only gets set during `TrainingArguments`/made. So not even when I said "accelerate launch" works, I was incorrect there. Only with FSDP is that done since we check a global variable (again, something I *do not want to rely on*).

I don't believe we want to add a new arg or a new device_map, as it doesn't make sense. If a user is doing zero3, `transformers` should know that we are doing `zero3`. 

Re;
> Couldn't we create a deepspeed object outside of TrainingArguments

Nope, we currently don't have this capability. (nor is our API designed to have this capability). The entirety of the DS integration is designed to rely off of the `TrainingArguments` as it stands, and I think would be a substantial rewrite for us to not do so.  


@muellerzr OK, I see, thanks for writing this all out. 

Yes, if it's a substantial rewrite, and effectively breaking then it doesn't make sense to modify. Given it's how things have been working up until now, agreed keeping like this makes most sense. Also agreed on not relying on global flags (even if private). 

My main hesitation is still that this is very much hidden and non-obvious to the user (why should I have to define these objects in this order?). Having this defensive error is the best protection - I'd also like to make sure this is documented somewhere on the [deepspeed page](https://huggingface.co/docs/transformers/main/en/deepspeed#deepspeed)

@amyeroberts it's documented there: 
![image](https://github.com/user-attachments/assets/da3256be-5110-447a-a8d6-fd0d615c5856)

https://huggingface.co/docs/transformers/main/en/deepspeed#deepspeed

But I don't disagree it's hard to find :)

@muellerzr FYI this broke my DPO workflow using TRL + Deepspeed. I have to create a copy reference model which doesn't seem to play nicely with this change. I have to downgrade to 4.43.1 to avoid this change and will not be able to upgrade in the future. 

@sam-h-bean Could you open an issue detailing this, including a code snippet to reproduce? This ensures the issue is properly tracked and this comment not lost

@amyeroberts @muellerzr 

Does it have a deeper reason that the implementations for `from_config` and `from_pretrained` are different and the flag won't be set for models with a quant confg?
```python
model._transformers_zero3_init_used = is_deepspeed_zero3_enabled() and not is_quantized
```

This broke a working finetuning script for me (which worked with Zero3 before) and I wonder why this was implemented in a way that breaks old code without an obvious migration path?

@jphme because we have never done zero init if we are doing zero3 and quantization:

https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L1469-L1476

If it worked for you before, we can add it back in (likely in this case my assumption is you go onto the second config which uses `init_empty_weights()`). This logic is what I was basing the implementation off of. 



@sam-h-bean @jphme can either of you try `pip install https://github.com/huggingface/transformers@muellerzr-enable-quant` to see if this resolves the problem?

> @sam-h-bean @jphme can either of you try `pip install https://github.com/huggingface/transformers@muellerzr-enable-quant` to see if this resolves the problem?

will do so and report back here (but will only get to it in the afternoon), thank you 👍 

EDIT: Works for me again with your branch :)

Hi @muellerzr, I read the PR and the discussion above, but I'm still not quite getting why this PR is necessary. Maybe could you explain a bit more, because as the discussion above pointed out, the new changes break all the trained code. Not sure if there are any workarounds available so far?

> @sam-h-bean Could you open an issue detailing this, including a code snippet to reproduce? This ensures the issue is properly tracked and this comment not lost

I also ran into the same problem, which is actually very simple code, using deepspeed xxx.py to start DPO training, everything else is exactly the same as the examples in trl.

@paulcx did you try the workaround in the linked PR above? (Which worked for the other user)

> @paulcx did you try the workaround in the linked PR above? (Which worked for the other user)

https://github.com/huggingface/transformers@muellerzr-enable-quant this one? If so, it does not work for me.

What I did is trying to manually replace the codes in trainer.py (transformers-4.44.0):
`if is_deepspeed_zero3_enabled() and not getattr(model, "_transformers_zero3_init_used", True):`
with
`if not (hasattr(model, "hf_quantizer") and model.hf_quantizer.is_trainable) and is_deepspeed_zero3_enabled() and not getattr(model, "_transformers_zero3_init_used", True):`

This broke my process of running code which ran well before using zero3 for inference. 
I confirm that 1) I used `from_pre_trained` to initialize the model and 2) use `accelerate  launch` to run the script, which mismatches the comment:

```
        # Will reach this branch if the user has
        # 1. Used `.from_pretrained` or `.from_config` to initialize their model
        # 2. Did not configure Zero-3 via `TrainingArguments` or `accelerate launch` beforehand
        # New models init such as `MyModel()` will not hit this step
```

 A trick to make the code run well is commenting this out:
```
        if is_deepspeed_zero3_enabled() and not getattr(model, "_transformers_zero3_init_used", True):
            raise ValueError(
                "Model was not initialized with `Zero-3` despite being configured for DeepSpeed Zero-3. Please re-initialize your model via `Model.from_pretrained(...)` or `Model.from_config(...)` after creating your `TrainingArguments`!"
            )
```

Not sure what's wrong here...

I’ll go ahead and put a revert in the AM for this then, apologies everyone :/

> This broke my process of running code which ran well before using zero3 for inference. I confirm that 1) I used `from_pre_trained` to initialize the model and 2) use `accelerate launch` to run the script, which mismatches the comment:
> 
> ```
>         # Will reach this branch if the user has
>         # 1. Used `.from_pretrained` or `.from_config` to initialize their model
>         # 2. Did not configure Zero-3 via `TrainingArguments` or `accelerate launch` beforehand
>         # New models init such as `MyModel()` will not hit this step
> ```
> 
> A trick to make the code run well is commenting this out:
> 
> ```
>         if is_deepspeed_zero3_enabled() and not getattr(model, "_transformers_zero3_init_used", True):
>             raise ValueError(
>                 "Model was not initialized with `Zero-3` despite being configured for DeepSpeed Zero-3. Please re-initialize your model via `Model.from_pretrained(...)` or `Model.from_config(...)` after creating your `TrainingArguments`!"
>             )
> ```
> 
> Not sure what's wrong here...

It saved my day. You are genius！