[BUG] Jamba (Mamba+MoE) + ZeRO3 + LoRA training hangs

[hijkzzz]

Model Link: https://huggingface.co/ai21labs/Jamba-v0.1

Reproduce script in OpenRLHF: https://github.com/OpenLLMAI/OpenRLHF/blob/main/examples/scripts/train_sft_jamba_lora.sh

Please pip install mamba-ssm causal-conv1d>=1.2.0
and set --micro_train_batch_size 1 to reproduce this bug ( IMPORTANT! )
--micro_train_batch_size 4 works well

Mixtral + ZeRO3 + LoRA works well with the same hyperparameters (`--micro_train_batch_size 1):
see https://github.com/OpenLLMAI/OpenRLHF/blob/main/examples/scripts/train_sft_mixtral_lora.sh