[fx] Add pofo solver (#1608)

* [fx] add pofo algorithm * [fx] Add pofo solver * [fx] code refactor * [fx] fix test_linearize import
hpcaitech · Sep 20, 2022 · 933b6c6 · 933b6c6
1 parent d32cf84
commit 933b6c6
Show file tree

Hide file tree

Showing 5 changed files with 450 additions and 13 deletions.
diff --git a/colossalai/fx/passes/algorithms/__init__.py b/colossalai/fx/passes/algorithms/__init__.py
@@ -1,3 +1,4 @@
 from .ckpt_solver_chen import chen_greedy
 from .linearize import linearize
 from .ckpt_solver_rotor import solver_rotor
+from .ckpt_solver_pofo import solver_pofo