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

[fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages (#1425)

Browse Source 
* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages

* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages

* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages
This commit is contained in:
Super Daniel
2022-08-10 16:36:35 +08:00
committed by GitHub
parent 89c434a0a6
commit f20cb4e893
2 changed files with 28 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/test_fx/test_meta_info_prop.py
View File

@@ -23,12 +23,24 @@ def test_meta_info_prop():
input_sample = torch.rand(BATCH_SIZE, DIM_IN)
orig_output = model(input_sample)
gm = symbolic_trace(model)
for node in gm.graph.nodes:
assert not hasattr(node,
'node_size'), 'The attribute Node.node_size should not exist before MetaInfoProp procedure'
assert not hasattr(node,
'param_size'), 'The attribute Node.param_size should not exist before MetaInfoProp procedure'
assert not hasattr(
node,
'activation_size'), 'The attribute Node.activation_size should not exist before MetaInfoProp procedure'
MetaInfoProp(gm).run(input_sample)
for node in gm.graph.nodes:
if node.op == 'placeholder':
meta_check(node.meta['tensor_meta'], input_sample)
if node.op == 'output':
meta_check(node.meta['tensor_meta'], orig_output)
assert hasattr(node, 'node_size'), 'The attribute Node.node_size should exist after MetaInfoProp procedure'
assert hasattr(node, 'param_size'), 'The attribute Node.param_size should exist after MetaInfoProp procedure'
assert hasattr(
node, 'activation_size'), 'The attribute Node.activation_size should exist after MetaInfoProp procedure'
if __name__ == '__main__':