QinLuo

with `MoeHybridParallelPlugin` I met another problem: https://github.com/hpcaitech/ColossalAI/issues/5426 ``` grad = grad.to(master_moe_param.dtype).to(master_moe_param.device) AttributeError: 'NoneType' object has no attribute 'to' ```

without `EPMixtralSparseMoeBlock` in mixtral_policy, I met another problem: hangs at the following point: 

> > with `MoeHybridParallelPlugin` I met another problem: #5426 > > ``` > > grad = grad.to(master_moe_param.dtype).to(master_moe_param.device) > > AttributeError: 'NoneType' object has no attribute 'to' > > ``` >...

GeminiPlugin can run correctly now.

Could someone in colossalai review it? Thx. cc @flybird11111

> Could you add unit tests to check this feature? I've tested in my local experiments  But how to test it in unittest: colossalai/tests

> > > Could you add unit tests to check this feature? > > > > > > I've tested in my local experiments  > > But how to...

@ver217 ``` pytest tests/test_checkpoint_io/test_torch_fsdp_checkpoint_io.py ================================================================= test session starts ================================================================== platform linux -- Python 3.10.11, pytest-8.0.0, pluggy-1.4.0 rootdir: ~/ColossalAI configfile: pytest.ini plugins: hypothesis-6.75.2, anyio-4.0.0 collected 1 item tests/test_checkpoint_io/test_torch_fsdp_checkpoint_io.py . [100%] ==================================================================...

Any ColossalAI-er could take a look?

I see the MR, that's awesome, when can we use it?