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

[fx] provide a stable but not accurate enough version of profiler. (#1547)

Browse Source 
* [fx] compute memory stat and flop count for MetaInfoProp.

* [fx] modify node attribute.

* [fx] modify ckpt_chen.

* [fx] fix compatibility.

* [fx] fix import error.

* [fx] skip test for MetaInfoProp.

* [fx] skip test for MetaInfoProp.

* [fx] skip test for MetaInfoProp.

* [fx] skip test for MetaInfoProp.

* [fx] skip if torch 1.11.0.

* [fx] recover MetaInfoProp support for PyTorch 1.11.

* [fx] provide a stable but not accurate enough version of profiler.

* [fx] provide a stable but not accurate enough version of profiler.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix compatibility in tests.

* [fx] fix import error.
This commit is contained in:
Super Daniel
2022-09-07 11:21:04 +08:00
committed by GitHub
parent 7d49e7b2db
commit 4f59693207
38 changed files with 776 additions and 263 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/test_fx/test_meta_info_prop.py
View File

@@ -5,6 +5,8 @@ import colossalai.nn as col_nn
from torch.fx import symbolic_trace
from colossalai.fx.passes.meta_info_prop import MetaInfoProp, TensorMetadata
import pytest
BATCH_SIZE = 2
DIM_IN = 4
DIM_OUT = 16
@@ -13,7 +15,6 @@ DIM_OUT = 16
def meta_check(meta_info_spec: TensorMetadata, orig_tensor: torch.Tensor):
assert meta_info_spec.shape == orig_tensor.shape
assert meta_info_spec.dtype == orig_tensor.dtype
assert meta_info_spec.requires_grad == orig_tensor.requires_grad
assert meta_info_spec.stride == orig_tensor.stride()
assert meta_info_spec.numel == orig_tensor.numel()
@@ -23,29 +24,12 @@ def test_meta_info_prop():
input_sample = torch.rand(BATCH_SIZE, DIM_IN, device='meta')
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__'), 'The attribute Node.__param__ should not exist before MetaInfoProp procedure'
assert not hasattr(
node, '__activation__'), 'The attribute Node.__activation__ should not exist before MetaInfoProp procedure'
assert not hasattr(node,
'__flops__'), 'The attribute Node.__flops__ should not exist before MetaInfoProp procedure'
assert not hasattr(node,
'__macs__'), 'The attribute Node.__macs__ 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__'), 'The attribute Node.__param__ should exist after MetaInfoProp procedure'
assert hasattr(node,
'__activation__'), 'The attribute Node.__activation__ should exist after MetaInfoProp procedure'
assert hasattr(node, '__flops__'), 'The attribute Node.__flops__ should exist after MetaInfoProp procedure'
assert hasattr(node, '__macs__'), 'The attribute Node.__macs__ should exist after MetaInfoProp procedure'
if __name__ == '__main__':