[Feature]: HyperParallel-FSDP MS后端支持CPUOffload-hyper-parallel-AtomGit

已开启

[Feature]: HyperParallel-FSDP MS后端支持CPUOffload #249

MengXY107创建于 6月25日

MengXY107成员

6月25日创建

Checklist

1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues)
2. I have read the relevant documentation.
3. I have created a minimal reproduction case that clearly demonstrates the issue, including a complete code example and the error message with full traceback and error logs.

🐛 Describe the bug

HyperParallel-FSDP在使用CPUOffload时，报错参数切片不在CPU上。然而期望参数切片（即sharded_param存在与CPU上）
问题在于当前Mindspore框架默认行为是CPU上的Tensor通过计算算子，会隐式的搬运到NPU上。经过和同事线下讨论。该问题依赖框架机制，因此该问题通过该需求跟踪，待讨论清楚框架机制后在MS后端支持CPUOffload特性。

Expected behavior

期望行为是：CPUOffload入参用来控制是否将参数，梯度，优化器状态的分片放置在CPU上，从而进一步减少Device上的显存占用。最后优化器更新会在CPU上进行.
fully_shard内部控制参数切片，梯度切片卸载到CPU上。包括对梯度累加场景是对梯度累加的+=都放置在CPU上进行。优化器动量通过在CPUDeveiceContext下创建，动量也可以放到CPU上。最终优化器.step在CPU上进行。