GitLab

[colotensor] add megatron initialization for gpt2

* process group supports getting ranks in group

* chunk mgr receives a process group

* update unit test

* fix unit tests

* sharded optim supports loading local state dict

* polish code

* add unit test

* zero optim supports loading local state dict

* polish code

* add unit test

* [test] ignore 8 gpu test

* polish code

* polish workflow

* polish workflow

* refactor StatefulTensor, tensor utilities

* add unitest for GeminiMemoryManager

* [test] refactored with the new rerun decorator

* polish test case

* fix reuse_fp16_shard

* disable test stm

* polish code

* add tensor placement policies

* polish comments

* polish comments

* update moe unit tests