ColossalAI/colossalai/fx/passes/algorithms
Boyuan Yao de1e716dc4
[fx] Add activation checkpoint solver rotor (#1496)
* [fx] fix defining ckpt functions inside forward

* [fx] Modify activation checkpoint codegen and add ColoGraphModule

* [fx] some modification

* some modifications

* some modifications

* some modifications

* some modifications

* some code modifications

* [automatic_parallel] ckpt solver rotor

* [fx] add ckpt_solver_rotor

* [fx] modification

* code refactor

* code refactor
2022-08-26 10:34:21 +08:00
..
__init__.py [fx] Add activation checkpoint solver rotor (#1496) 2022-08-26 10:34:21 +08:00
ckpt_solver_chen.py [fx] add profiler for fx nodes. (#1480) 2022-08-24 16:22:44 +08:00
ckpt_solver_rotor.py [fx] Add activation checkpoint solver rotor (#1496) 2022-08-26 10:34:21 +08:00
linearize.py [fx] Add activation checkpoint solver rotor (#1496) 2022-08-26 10:34:21 +08:00
utils.py [fx] Add activation checkpoint solver rotor (#1496) 2022-08-26 10:34:21 +08:00