Unable to use MimiModel with DeepSpeed ZeRO-3

[anferico]

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 

Reproducing the error

mimi_mre.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:

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:

{
    "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).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ylacombe, @eustlb