Fix multi-model training with deepspeed nvme-offloading
Author: xu-song
Pull request: yes
Labels: name: DeepSpeed

# What does this PR do?

This PR fix "multi-model training with deepspeed nvme offloading"



## 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?

## To Reproduce 

It gets error `/home/nvme/param/zero_stage_3/float16params/rank0/0_param.tensor.swp: buffer nbytes != file bytes` when multi-model training with deepspeed  nvme offloading

Here is an example deepspeed config
```yaml
# ds_config.json 
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/home/nvme/optim",   # different models share single nvme_path, which cause error
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/home/nvme/param",   # different models share single nvme_path, which cause error
        },
    }
```
<!--  
RLHF: actor_model, reference_model, reward_model(一般比较小，不用nvme)。 
DPO: actor_model, reference_ model
 -->
If you are training with multiple models (such as rlhf, dpo, multimodal) like this
```py
# train.py
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)   # it calls deepspeed.zero.Init, then swap_out param to nvme_path 
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR)   # it calls deepspeed.zero.Init again
```
Launch deepspeed training with `ds_config.json`
```
deepspeed --hostfile=hostfile train.py --deepspeed --deepspeed_config ds_config.json 
```


`deepspeed.zero.Init` first remove all files in `nvme_path`, which causes the swap files of `model_1` be overwritten by `model_2`. 

<details> 

<summary>details of deepspeed.zero.Init</summary>

```py
# https://github.com/microsoft/DeepSpeed/blob/v0.14.4/deepspeed/runtime/zero/partition_parameters.py#L1010
class Init:   
    def __init__(
        ....
        if self.remote_device == OffloadDeviceEnum.nvme:
            self.param_swapper = param_swapper or AsyncPartitionedParameterSwapper(_ds_config, self.dtype)

# https://github.com/microsoft/DeepSpeed/blob/v0.14.4/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py#L89
class AsyncPartitionedParameterSwapper(object):
    def __init__(self, ds_config, model_dtype):
         ..
        shutil.rmtree(self.swap_folder, ignore_errors=True)    
```
</details>


## After this PR

You can initialize two models with different `nvme_path`
```py
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)   # with default nvme_path
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR, disable_zero_init=True)
model_2_engine, *_ = deepspeed.initialize(model=model_2,
                                      config=another_ds_config)  # with a different nvme_path
```

## Who can review?


- deepspeed: HF Trainer/Accelerate: @muellerzr
- trainer: @muellerzr and @SunMarc




Comments:
Gentle ping @SunMarc @muellerzr 

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

@xu-song can you rebase from main? This should fix the failing tests (and sorry for taking a bit to respond)

This commit has been rebased from main.

Great, looks like the failing tests are unrelated then. cc @amyeroberts for final review :)

Hi @xu-song, thanks for opening this PR!

I'm a bit concerned about the solution needed here to enable this behaviour - it's pretty hacky to need to have one (undocumented) argument to enable the init of two models and is an indication of the wider problem with how deepspeed is being handled through the models. In particular - why should one be able to instantiate one model correctly and then for all subsequent models have to force no init and then call deepspeed to initialiaze? I realise this problem isn't your design, but I don't think adding arguments like this is a sustainable way to manage our deepspeed integration. 

At the moment the pattern for three models is: 
```py
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)   # with default nvme_path

model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR, disable_zero_init=True)
model_2_engine, *_ = deepspeed.initialize(model=model_2, config=another_ds_config_2)

model_3 = AutoModelForCausalLM.from_pretrained(MODEL_3_DIR, disable_zero_init=True)
model_3_engine, *_ = deepspeed.initialize(model=model_3, config=another_ds_config_3)
```

We should be able to do something like this: 

```py
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR)
model_3 = AutoModelForCausalLM.from_pretrained(MODEL_3_DIR)

model_1_engine, *_ = deepspeed.initialize(model=model_1, config=another_ds_config_1)
model_2_engine, *_ = deepspeed.initialize(model=model_2, config=another_ds_config_2)
model_3_engine, *_ = deepspeed.initialize(model=model_3, config=another_ds_config_3)
```

or this: 

```
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR, deepspeed_config=another_ds_config_1)
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR, deepspeed_config=another_ds_config_2)
model_3 = AutoModelForCausalLM.from_pretrained(MODEL_3_DIR, deepspeed_config=another_ds_config_3)
```

or this: 

```
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR)
model_3 = AutoModelForCausalLM.from_pretrained(MODEL_3_DIR)

model_engines = initialize_deepspeed(another_ds_config_1, another_ds_config_2, another_ds_config_3)
```

There are many different patterns - but the initialization for the models should be consistent. 

In your prefered examples, I think you can set default value `disable_zero_init=True`
```py
# train.py
model_1 = AutoModelForCausalLM.from_pretrained(MODEL_1_DIR)   # will not init deepspeed 
model_2 = AutoModelForCausalLM.from_pretrained(MODEL_2_DIR)   # will not init deepspeed 
model_3 = AutoModelForCausalLM.from_pretrained(MODEL_3_DIR)   # will not init deepspeed 

model_1_engine, *_ = deepspeed.initialize(model=model_1, config=another_ds_config_1)
model_2_engine, *_ = deepspeed.initialize(model=model_2, config=another_ds_config_2)
model_3_engine, *_ = deepspeed.initialize(model=model_3, config=another_ds_config_3)
```






@xu-song Right - we can't do this immediately though, as it would break things for current users who expect DS initialization to be happening automagically when they intialize their (single) model. 

In any case - this should be documented somewhere so people know how to use this logic in the multi model case. 

As a side note - I'd much rather we added an argument which defined a positive rather than negative action as it's clearer i.e. `zero_init=True` rather than `disable_zero_init=False` @muellerzr WDYT? We could have it unset and defaults to True if certain settings occur e.g. `is_deepspeed_zero3_enabled` is `True`

@amyeroberts I agree (we also do not support multi-modal DS yet FYI)