github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] fixed compatiblity issue with torch 1.10 (#1331)

Browse Source
This commit is contained in:
Frank Lee
2022-07-18 11:41:27 +08:00
committed by GitHub
parent 069d6fdc84
commit 75abc75c15
7 changed files with 32 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
colossalai/fx/passes/split_module.py
View File

@@ -2,6 +2,7 @@ import torch
from torch.fx.graph_module import GraphModule
from typing import Callable, List, Dict, Any, Optional
from torch.fx._compatibility import compatibility
from packaging import version
import inspect
@@ -233,10 +234,13 @@ def split_module(
base_mod_attrs: Dict[str, torch.fx.graph_module.GraphModule] = {}
for node in m.graph.nodes:
if node.op == 'placeholder':
default_value = node.args[0] if len(node.args) > 0 else inspect.Signature.empty
base_mod_env[node.name] = base_mod_graph.placeholder(node.name,
type_expr=node.type,
default_value=default_value)
if version.parse(torch.__version__) < version.parse('1.11.0'):
base_mod_env[node.name] = base_mod_graph.placeholder(node.name, type_expr=node.type)
else:
default_value = node.args[0] if len(node.args) > 0 else inspect.Signature.empty
base_mod_env[node.name] = base_mod_graph.placeholder(node.name,
type_expr=node.type,
default_value=default_value)
base_mod_env[node.name].meta = node.meta.copy()
# Do some things iterating over the partitions in topological order again:

1
colossalai/fx/tracer/meta_patch/patched_function/arithmetic.py
View File

@@ -3,6 +3,7 @@ from ..registry import meta_patched_function
@meta_patched_function.register(torch.matmul)
@meta_patched_function.register('matmul') # for built-in op @
def torch_matmul(input, other, *, out=None):
# copied from huggingface.utils.fx
d1 = input.dim()

3
colossalai/fx/tracer/tracer.py
View File

@@ -96,6 +96,9 @@ class ColoTracer(Tracer):
# fetch patched function
if meta_patched_function.has(target):
meta_target = meta_patched_function.get(target)
elif meta_patched_function.has(target.__name__):
# use name for some builtin op like @ (matmul)
meta_target = meta_patched_function.get(target.__name__)
else:
meta_target = target