tests/test_moe/test_moe_zero_optim.py · a9b8300d54f842c4ae840552e99f2ffadffcc172 · OpenDAS / ColossalAI

HELSON authored Apr 11, 2022

* adapt post grad hooks for not-shard parameters
* adapt optimizer for not-shard parameters
* offload gradients for not-replicated parameters

a9b8300d

test_moe_zero_optim.py 5.52 KB