ColossalAI/colossalai/zero/sharded_model
Jiarui Fang 3af13a2c3e [zero] polish ShardedOptimV2 unittest (#385)
* place params on cpu after zero init context

* polish code

* bucketzed cpu gpu tensor transter

* find a bug in sharded optim unittest

* add offload unittest for ShardedOptimV2.

* polish code and make it more robust
2022-03-11 15:50:28 +08:00
..
__init__.py Feature/zero (#279) 2022-03-11 15:50:28 +08:00
_zero3_utils.py [zero] Update sharded model v2 using sharded param v2 (#323) 2022-03-11 15:50:28 +08:00
param_manager.py Feature/zero (#279) 2022-03-11 15:50:28 +08:00
reduce_scatter.py [zero] add sharded grad and refactor grad hooks for ShardedModel (#287) 2022-03-11 15:50:28 +08:00
sharded_grad.py [zero] add sharded grad and refactor grad hooks for ShardedModel (#287) 2022-03-11 15:50:28 +08:00
sharded_model.py [zero] Update sharded model v2 using sharded param v2 (#323) 2022-03-11 15:50:28 +08:00
sharded_model_v2.py [zero] polish ShardedOptimV2 unittest (#385) 2022-03-11 15:50:28 +08:00
utils.py [test] polish zero related unitest (#351) 2022-03-11 15:50:28 +08:00