github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-21 13:11:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] add recompile after graph manipulatation (#1621)

Browse Source
This commit is contained in:
YuliangLiu0306 2022-09-23 11:00:33 +08:00 committed by GitHub
parent d967779a32
commit 9eae855408
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/test_auto_parallel/test_shape_consistency_pass.py
View File

@ -65,6 +65,7 @@ def check_apply(rank, world_size, port):
solution = list(ret[0])
sharding_spec_dict, origin_spec_dict = solution_annotatation_pass(gm, solution, device_mesh)
shape_consistency_pass(gm)
gm.recompile()
nodes = [node for node in gm.graph.nodes]
# TODO: wrap the gm to avoid the influence of the user training code
output = gm(input, sharding_spec_dict, origin_spec_dict)