[BUG] Mixtral + QLoRA + zero3 occured error

Hello everyone.

When I trained Mixtral-MOE with QLoRA + Zero3, it occurs error like below.

_TypeErrorwork = group._allgather_base(output_tensor, input_tensor, opts)
TypeError: output tensor must have the same type as input tensor.

It seems odd because after several iterations, about 10 to 20, it works fine, but then suddenly the error occurs.
Also, 'Mixtral-MOE + QLoRA + Zero2' and 'llama-7B + QLoRA + Zero3' works well.

Any solutions?

My environment settings:

• transformers 4.36.2
• torch 2.2.0a0+81ea7a4
• deepspeed 0.13.1
• accelerate 0.26.1

deepspeed config is

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "bf16": {"enabled": true},
  "optimizer": {
      "type": "AdamW",
      "params": {
          "lr": "auto",
          "betas": "auto",
          "eps": "auto",
          "weight_decay": "auto"
      }
  },
"scheduler": {
  "type": "WarmupDecayLR",
  "params": {
    "last_batch_iteration": -1,
    "total_num_steps": "auto",
    "warmup_min_lr": 5e-7,
    "warmup_max_lr": "auto",
    "warmup_num_steps": "auto"
  }
},
  "zero_optimization": {
      "stage": 3,
      "offload_param": {"device": "cpu", "pin_memory": true},
      "offload_optimizer": {"device": "cpu", "pin_memory": true},
      "allgather_partitions": true,
      "allgather_bucket_size": 5e8,
      "contiguous_gradients": true,
      "overlap_comm": true,
      "reduce_scatter": true,
      "reduce_bucket_size": 5e8,

      "zero_quantized_weights": false,
      "zero_hpz_partition_size": 8,
      "zero_quantized_gradients": false
  },
  "gradient_accumulation_steps": 2
}