Unable to use `MimiModel` with DeepSpeed ZeRO-3
Author: anferico
Pull request: yes
Labels: name: run-slow

# What does this PR do?
Allow using `MimiModel` with DeepSpeed ZeRO-3.

Fixes the following error:
```
[rank0]: Traceback (most recent call last):                                                                                                                                                                                                  
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/./mimi_mre.py", line 10, in <module>                                                                                                                                         
[rank0]:     main()                                                                                                                                                                                                                          
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/./mimi_mre.py", line 7, in main                                                                                                                                              
[rank0]:     model = MimiModel.from_pretrained("kyutai/mimi", config=config)                                                                                                                                                                 
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/modeling_utils.py", line 4110, in from_pretrained                        
[rank0]:     model = cls(config, *model_args, **model_kwargs)                                                                                             
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)                                                                                                                   
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1515, in __init__                    
[rank0]:     self.quantizer = MimiSplitResidualVectorQuantizer(config)                                                                                    
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)                                                                                                                               
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1340, in __init__                    
[rank0]:     self.semantic_residual_vector_quantizer = MimiResidualVectorQuantizer(config, self.num_semantic_quantizers)                                  
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)                                                                                                                   
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1282, in __init__                                
[rank0]:     self.layers = nn.ModuleList([MimiVectorQuantization(config) for _ in range(self.num_quantizers)])                                            
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1282, in <listcomp>                  
[rank0]:     self.layers = nn.ModuleList([MimiVectorQuantization(config) for _ in range(self.num_quantizers)])                                            
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)                                                          
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1261, in __init__                                                            
[rank0]:     self.codebook = MimiEuclideanCodebook(config)                                       
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper                           
[rank0]:     f(module, *args, **kwargs)                                                          
[rank0]:   File "/net/tscratch/people/plgfcariaggi/transformers/src/transformers/models/mimi/modeling_mimi.py", line 1217, in __init__                                                            
[rank0]:     self.register_buffer("initialized", torch.Tensor([True]))                           
[rank0]:   File "/net/tscratch/people/plgfcariaggi/envs/transformers/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 255, in new_tensor
[rank0]:     tensor = _orig_torch_empty(0, device=device).new_empty(*args, **kwargs)                                                                                                                                                         
[rank0]: TypeError: new_empty(): argument 'size' (position 1) must be tuple of ints, but found element of type bool at pos 0 
```

**Explanation**: to begin with, the use of `torch.Tensor(...)` is discouraged/deprecated (cf. [https://pytorch.org/docs/stable/tensors.html#torch.Tensor](https://pytorch.org/docs/stable/tensors.html#torch.Tensor)) and should be dropped in favor of `torch.tensor(...)`. Secondly, here is an excerpt of DeepSpeed's source code (`deepspeed/runtime/zero/partition_parameters.py`) that explains why we get this error:

```python
def zero_wrapper_for_fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype) -> Callable:

    def wrapped_fn(*args, **kwargs) -> Tensor:
        if kwargs.get("device", None) is None:
            kwargs['device'] = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"]))
        tensor: Tensor = fn(*args, **kwargs)
        if tensor.is_floating_point():
            tensor.data = tensor.data.to(target_fp_dtype)

        return tensor

    return wrapped_fn


def get_new_tensor_fn_for_dtype(dtype: torch.dtype) -> Callable:

    def new_tensor(cls, *args, **kwargs) -> Tensor:
        device = torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"]))
        if not args:
            args = (0, )
        tensor = _orig_torch_empty(0, device=device).new_empty(*args, **kwargs)
        if tensor.is_floating_point():
            tensor = tensor.to(dtype)

        return tensor

    return new_tensor

[...]

    def _add_tensor_creation_wrappers(self):
        torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(self.dtype)
        torch.tensor = zero_wrapper_for_fp_tensor_constructor(_orig_torch_tensor, self.dtype)
        torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, self.dtype)
        torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, self.dtype)
        torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, self.dtype)
        torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, self.dtype)
        torch.arange = zero_wrapper_for_fp_tensor_constructor(_orig_torch_arange, self.dtype)
        torch.eye = zero_wrapper_for_fp_tensor_constructor(_orig_torch_eye, self.dtype)
        torch.randn = zero_wrapper_for_fp_tensor_constructor(_orig_torch_randn, self.dtype)
```

as can be seen, DeepSpeed ZeRO-3 patches `Tensor.__new__` with `get_new_tensor_fn_for_dtype(self.dtype)` which, if called with `[True]` as an argument, raises an error because such argument is passed directly to `torch.Tensor.new_empty`, which expects a `size` as its first argument. On the contrary, `torch.tensor` is patched with the correct function (`zero_wrapper_for_fp_tensor_constructor(_orig_torch_tensor, self.dtype)`) and won't throw any error when passed `[True]` as an argument. 

## Reproducing the error
`mimi_mre.sh`:
```sh
OUTPUT_DIR=$HOME/mimi_mre

deepspeed \
    --num_gpus 1 \
    --master_port 60000 \
    ./mimi_mre.py \
    --output_dir $OUTPUT_DIR \
    --deepspeed zero3.json

```

`mimi_mre.py`:
```python

from transformers import AutoConfig, MimiModel, TrainingArguments, HfArgumentParser

def main():
    parser = HfArgumentParser(TrainingArguments)
    training_args = parser.parse_args_into_dataclasses()[0]
    config = AutoConfig.from_pretrained("kyutai/mimi")
    model = MimiModel.from_pretrained("kyutai/mimi", config=config)

if __name__ == "__main__":
    main()
```

`zero3.json`:
```json
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "bf16": {
        "enabled": "auto"
    },
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```


## 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.
- [x] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

@ylacombe, @eustlb


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

@ylacombe all done!

@ylacombe 1 test is failing (`tests.models.mimi.test_modeling_mimi.MimiIntegrationTest.test_integration`), which is kinda weird cause I don't think my change has any effect on the model's behavior. Any thoughts?

@ylacombe ping

@eustlb ? @ylacombe ? Any idea why some tests are failing?

Hey @anferico, the failing slow test is unrelated indeed (see https://github.com/huggingface/transformers/pull/35696).
I am of course okay with the provided fix 😊
Does it fix the originally reported issue? Can you please update the PR comment with a brief explanation if it does? 🙏
(from my understanding, it is a dtype issue entailed by `torch.Tensor` that fails in getting the right dtype - torch.bool -, contrary to `torch.tensor`, can you confirm?)

> Hey @anferico, the failing slow test is unrelated indeed (see #35696). I am of course okay with the provided fix 😊 Does it fix the originally reported issue? Can you please update the PR comment with a brief explanation if it does? 🙏 (from my understanding, it is a dtype issue entailed by `torch.Tensor` that fails in getting the right dtype - torch.bool -, contrary to `torch.tensor`, can you confirm?)

@eustlb yes I can confirm it fixes the issue, meaning I _can_ use DeepSpeed ZeRO-3 in conjunction with MimiModel with the changes introduced in this PR

@eustlb just updated the PR description too!

Great! thanks a lot 🤗

Reverted because it caused a failing test on `Moshi` that was not catched by CI, looking into it

Found the issue, we need to keep the registered buffer to be a `torch.float32`, sorry for missing that!
Would you like to open a new PR by simply adding this one as a reference and changing:
```diff
- self.register_buffer("initialized", torch.Tensor([True]))
+ self.register_buffer("initialized", torch.tensor([True], dtype=torch.float32))
```
?

@eustlb done! #35759