github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] Add activation checkpoint solver rotor (#1496)

Browse Source 
* [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
This commit is contained in:
Boyuan Yao
2022-08-26 10:34:21 +08:00
committed by GitHub
parent 09c023bee2
commit de1e716dc4
6 changed files with 529 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/fx/__init__.py
View File

@@ -1 +1,2 @@
from .tracer import ColoTracer
from .graph_module import ColoGraphModule