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

[autoparallel] use metainfo in handler (#2149)

Browse Source
This commit is contained in:
YuliangLiu0306
2022-12-20 10:31:22 +08:00
committed by GitHub
parent 9b39170a5c
commit 1cce6e36ca
11 changed files with 105 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py
View File

@@ -30,7 +30,7 @@ def avgpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem,
Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs
"""
input_tensor = next(filter(lambda x: x.type == OperationDataType.ARG, args)).data
input_tensor = args[0].data
output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data
# construct forward args for flop mapping