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

3 lines
87 B
Python
Raw Normal View History

[fx] support meta tracing for aten level computation graphs like functorch. (#1536) * [fx] support meta tracing for aten level computation graphs like functorch. * [fx] support meta tracing for aten level computation graphs like functorch. * [fx] remove redundant import. * [fx] add docstring.
2022-09-05 04:10:09 +00:00
from .tracer import ColoTracer, meta_trace
[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 .graph_module import ColoGraphModule