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

[fx] support module with bias addition (#1780)

Browse Source 
* [autoparallel] refactor tracer to fix bias addition issue

* [fx] support module with bias addition

* create bias_addition_module

* refactor file structure

* polish code

* fix unit test
This commit is contained in:
YuliangLiu0306
2022-11-01 22:53:51 +08:00
committed by GitHub
parent f3f19a5c47
commit e859380bf7
41 changed files with 624 additions and 259 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_fx/test_tracer/test_torchaudio_model/torchaudio_utils.py
View File

@@ -1,7 +1,8 @@
from colossalai.fx import ColoTracer
import torch
from torch.fx import GraphModule, Tracer
from colossalai.fx import ColoTracer
def trace_and_compare(model, data_gen, need_meta=False, need_concrete=False, kwargs_transform=False):
data = data_gen()
@@ -24,8 +25,9 @@ def trace_and_compare(model, data_gen, need_meta=False, need_concrete=False, kwa
fx_out = gm(**data)
if isinstance(fx_out, tuple):
for non_fx, fx in zip(non_fx_out, fx_out):
assert torch.allclose(non_fx,
fx), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'
assert torch.allclose(
non_fx, fx, atol=1e-5), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'
else:
assert torch.allclose(
fx_out, non_fx_out), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'
fx_out, non_fx_out,
atol=1e-5), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'