from mindspeed.fsdp.memory.swap_optimizer.swap_optimizer import AdamWSwap