GitLab

* fix reuse_fp16_shard

* disable test stm

* polish code

* add tensor placement policies

* polish comments

* polish comments

* update moe unit tests

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

* refactor memstats collector

* fix disposable

* polish code

* [WIP] stateful tensor manager

* add eviction strategy

* polish code

* polish code

* polish comment

* add unit test

* fix sampler bug

* polish code

* fix max sampling cnt resetting bug

* fix sampler bug

* polish code

* fix bug

* fix unit test
Co-authored-by: jiaruifang <fangjiarui123@gmail.com>

* adapt model weight initialization for methods in Pytorch nn.init

* remove hybrid adam in test_moe_zero_optim

* fix activation checkpointing and its unitest

* add cpu adam counter for all cpu adam

* fixed updating error in adam kernel

* support existing sharded and unsharded parameters in zero

* add unitest for moe-zero model init

* polish moe gradient handler

* hijack p.grad in sharded model

* polish comments

* polish comments

only process module's own parameters in Zero context

add zero hooks for all modules that contrain parameters

gather parameters only belonging to module itself