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

[autoparallel] support origin activation ckpt on autoprallel system (#2468)

Browse Source
This commit is contained in:
YuliangLiu0306
2023-01-16 16:25:13 +08:00
committed by GitHub
parent 3a21485ead
commit 67e1912b59
4 changed files with 111 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/auto_parallel/passes/runtime_preparation_pass.py
View File

@@ -179,6 +179,8 @@ def _size_value_converting(gm: torch.fx.GraphModule, device_mesh: DeviceMesh):
# It will be used to replace the original node with processing node in slice object
node_pairs[node] = size_processing_node
size_processing_node._meta_data = node._meta_data
if 'activation_checkpoint' in node.meta:
size_processing_node.meta['activation_checkpoint'] = node.meta['activation_checkpoint']
user_list = list(node.users.keys())
for user in user_list: