github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-11 22:10:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx] added apex normalization to patched modules (#1300)

Browse Source 
* [fx] added apex normalization to patched modules

* remove unused imports
This commit is contained in:
Frank Lee
2022-07-14 14:24:13 +08:00
committed by GitHub
parent 4165eabb1e
commit 4f4d8c3656
3 changed files with 11 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_fx/test_pipeline/test_hf_model/test_t5.py
View File

@@ -2,15 +2,6 @@ import pytest
import transformers
import torch
from hf_utils import split_model_and_compare_output
from colossalai.fx.tracer.meta_patch import meta_patched_module
try:
import apex
@meta_patched_module.register(apex.normalization.FusedRMSNorm)
def apex_fused_layernorm(self, input):
return torch.empty(input.shape, device='meta')
except ImportError:
pass
BATCH_SIZE = 1
SEQ_LENGHT = 16