github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 01:28:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] Modify solver linearize and add corresponding test (#1531)

Browse Source 
* [fx] modify solver linearize and add test

* [fx] add torch11 test of linearize but skip it

* [fx] remove some unused imports
This commit is contained in:
Boyuan Yao
2022-09-02 10:24:41 +08:00
committed by GitHub
parent 7dc53237c3
commit 56159049e8
4 changed files with 181 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/fx/codegen/activation_checkpoint_codegen.py
View File

@@ -167,8 +167,7 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func,
use_reentrant = True
non_leaf_input = 0
for var in input_vars[label]:
input_node = [item for item in node_list if item.name == var]
input_node = input_node[0]
input_node = next(item for item in node_list if item.name == var)
if input_node.op != "placeholder":
non_leaf_input = 1
for user in input_node.users: