github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
ColossalAI/colossalai/fx/passes/algorithms/__init__.py

4 lines
122 B
Python
Raw Normal View History

[fx] add rules to linearize computation graphs for searching. (#1461) * [fx] polish ckpt_test. * [fx] add rules to linearize computation graphs for searching. * [fx] remove chen_sqrt for sake of simplicity * [fx] fix inconsistencies.
2022-08-17 06:47:12 +00:00
from .ckpt_solver_chen import chen_greedy
[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 02:34:21 +00:00
from .linearize import linearize
from .ckpt_solver_rotor import solver_rotor