ColossalAI/colossalai/zero/sharded_model
Jiarui Fang 21dc54e019
[zero] memtracer to record cuda memory usage of model data and overall system (#395)
2022-03-14 22:05:30 +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] memtracer to record cuda memory usage of model data and overall system (#395) 2022-03-14 22:05:30 +08:00
utils.py [test] polish zero related unitest (#351) 2022-03-11 15:50:28 +08:00