Jamba-v01 Model + Deepspeed Zero3 lead to "RuntimeError: Detected mismatch between collectives on ranks."
Author: KaiQiangSong
Pull request: no
Labels: name: DeepSpeed

name: bug

### System Info

- `transformers` version: 4.39.0
- Platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.29.2
- Accelerate config:    not found
- Deepspeed version: 0.14.1
- PyTorch version (GPU?): 2.1.0a0+32f93b1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

deepspeed config:
```json
{
    "zero_optimization": {
        "stage": 3, 
        "offload_optimizer": {
            "device": "cpu", 
            "pin_memory": true
        }, 
        "offload_param": {
            "device": "cpu", 
            "pin_memory": true
        }, 
        "overlap_comm": true, 
        "contiguous_gradients": true, 
        "sub_group_size": 0, 
        "reduce_bucket_size": 1.677722e+07, 
        "stage3_prefetch_bucket_size": 1.509949e+07, 
        "stage3_param_persistence_threshold": 4.096000e+04, 
        "stage3_max_live_parameters": 1.000000e+09, 
        "stage3_max_reuse_distance": 1.000000e+09, 
        "stage3_gather_16bit_weights_on_model_save": true
    }, 
    "bf16": {
        "enabled": true, 
        "auto_cast": false, 
        "loss_scale": 0, 
        "initial_scale_power": 32, 
        "loss_scale_window": 1000, 
        "hysteresis": 2, 
        "min_loss_scale": 1
    }, 
    "train_batch_size": 256, 
    "gradient_accumulation_steps": 8, 
    "train_micro_batch_size_per_gpu": 2, 
    "wall_clock_breakdown": false, 
    "steps_per_print": inf, 
    "fp16": {
        "enabled": false
    }, 
    "zero_allow_untested_optimizer": true
}
```
training with data-parallel and deepspeed zeor3 off-loading.


### Who can help?

@pacman100 


### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

SFT training using Jamba-v0.1 model with Accerate Trainer and Deepspeed Zero3 offload.

### Expected behavior

After a few iterations, there's an error met.

Error Message:
```shell
  File "/share5/users/kqsong/code/FastChat/fastchat/train/train_better_preprocessing.py", line 306, in <module>
    train()
  File "/share5/users/kqsong/code/FastChat/fastchat/train/train_better_preprocessing.py", line 300, in train
    trainer.train()
  File "/root/.local/lib/python3.10/site-packages/transformers/trainer.py", line 1780, in train
    return inner_training_loop(
  File "/root/.local/lib/python3.10/site-packages/transformers/trainer.py", line 2118, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/root/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3036, in training_step
    loss = self.compute_loss(model, inputs)
  File "/root/.local/lib/python3.10/site-packages/transformers/trainer.py", line 3059, in compute_loss
    outputs = model(**inputs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1855, in forward
    loss = self.module(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_jamba.py", line 1849, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_jamba.py", line 1715, in forward
    layer_outputs = self._gradient_checkpointing_func(
  File "/usr/local/lib/python3.10/dist-packages/torch/_compile.py", line 24, in inner
    return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py", line 450, in checkpoint
    return CheckpointFunction.apply(function, preserve, *args)
  File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py", line 230, in forward
    outputs = run_function(*args)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_jamba.py", line 1361, in forward
    hidden_states, router_logits = self.moe(hidden_states)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1568, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/root/.cache/huggingface/modules/transformers_modules/modeling_jamba.py", line 1211, in forward
    current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None]
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1557, in _call_impl
    args_result = hook(self, args)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 278, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 452, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module, forward=True)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 385, in fetch_sub_module
    self.__all_gather_params(params_to_prefetch, forward)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 434, in __all_gather_params
    self.__all_gather_params_(nonquantized_params, forward, quantize=self.zero_quantized_weights)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 463, in __all_gather_params_
    handle = param_group[0].all_gather_coalesced(param_group, quantize=quantize)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1259, in all_gather_coalesced
    handles.append(_all_gather_dtype(dtype, params, world_size, rank_in_group, ds_process_group))
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1147, in _all_gather_dtype
    handle = _dist_allgather_fn(partitions[rank_in_group], flat_tensor, ds_process_group)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 95, in _dist_allgather_fn
    return instrument_w_nvtx(dist.allgather_fn)(output_tensor, input_tensor, group=group, async_op=True)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 320, in allgather_fn
    return all_gather_into_tensor(output_tensor, input_tensor, group=group, async_op=async_op, debug=debug)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 117, in log_wrapper
    return func(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/comm/comm.py", line 305, in all_gather_into_tensor
    return cdb.all_gather_into_tensor(output_tensor=output_tensor, input_tensor=tensor, group=group, async_op=async_op)
  File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/root/.local/lib/python3.10/site-packages/deepspeed/comm/torch.py", line 199, in all_gather_into_tensor
    return self.all_gather_function(output_tensor=output_tensor,
  File "/usr/local/lib/python3.10/dist-packages/torch/distributed/c10d_logger.py", line 47, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py", line 2886, in all_gather_into_tensor
    work = group._allgather_base(output_tensor, input_tensor)
RuntimeError: Detected mismatch between collectives on ranks. Rank 0 is running collective: CollectiveFingerPrint(SequenceNumber=1843749, OpType=_ALLGATHER_BASE, TensorShape=[2228224], TensorDtypes=BFloat16, TensorDeviceTypes=TensorOptions(dtype=float (def
ault), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))), but Rank 10 is running collective: CollectiveFingerPrint(SequenceNumber=1843749, OpType=_ALLGATHER_BASE, TensorShape=[131
072], TensorDtypes=BFloat16, TensorDeviceTypes=TensorOptions(dtype=float (default), device=cuda, layout=Strided (default), requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))).Collectives differ in the following aspects: 
  Tensor Tensor shapes: 2228224vs 131072
```

Comments:
Gentle ping @pacman100 

cc @muellerzr @SunMarc 

same error

Another ping @muellerzr @SunMarc 

I've encountered this error myself, and I'd like to share my experience, 

It happened when I set some grad updates to False in a model with DeepSpeed ZeRO3 applied.
Once I realised this, I changed the require_grad to True and the model worked fine.

I'd have to go through the code to find out exactly why, but I suspect it was due to ZeRO3 being applied in the from_pretrained step.

In the ZeRO3 environment, there is a section where from_pretrained changes the context manager to DeepSpeed's Init to load the model, depending on whether it is ZeRO3 or not. 
With the model already sharded to each device, I set require_grad to False, and I'm able to get the

I'm guessing that the error is caused by the difference between the number of parameters that deepspeed needs to update and the number of parameters that need to be updated when training.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @KaiQiangSong @wwj1013 @jp1924 
Can you provide me with a repro script for this issue? I'd be happy to take a look at it. 🤗 


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

I encounter the same issue when I train InternVL with GRPO, the DeepSpeed stage is set to 2. After I turn off DeepSpeed (stage=0), the issue was gone.


The same problem occurs when I train a VLM using mixture training data of pure-text and image-text. Pure-text data won't update vision layer parameters but image-text data will, which causes TensorShape mismatch. This happends on deepspeed>=0.18.2 with deepspeed-zero2/3.