Zihan Wang

[BUG]: AssertionError: Cannot compare difference for two sharding specs with different length.

1

### 🐛 Describe the bug **I use gemini plugin and set tp>1, the error occurs when saving optimizer.** Traceback (most recent call last): File "/mnt/lustre/tangyang2/wangzihan/ColossalAI/applications/Colossal-LLaMA-2/pretrain_np.py", line 480, in main() File...

### 🐛 Describe the bug File "/usr/local/lib/miniconda3/envs/colo/lib/python3.10/site-packages/colossalai/tensor/colo_tensor.py", line 81, in __torch_function__ File "/usr/local/lib/miniconda3/envs/colo/lib/python3.10/site-packages/colossalai/tensor/colo_tensor.py", line 81, in __torch_function__ return backward_tensor.backward(**tensor_kwargs) File "/usr/local/lib/miniconda3/envs/colo/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward result = torch_func_method(public_api, types, args,...