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

[FAW] export FAW in _ops (#1438)

Browse Source
This commit is contained in:
Jiarui Fang
2022-08-11 13:43:24 +08:00
committed by GitHub
parent 9056677b13
commit 30b4dd17c0
9 changed files with 13 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/test_ops/test_cache_embedding.py → tests/test_layers/test_cache_embedding.py
View File

@@ -1,15 +1,17 @@
import pytest
from functools import partial
import torch
import torch.multiprocessing as mp
import numpy as np
import random
import torch
import torch.multiprocessing as mp
import colossalai
from colossalai.utils import free_port
from colossalai.testing import rerun_if_address_is_in_use
from colossalai.tensor import ColoParameter, ProcessGroup, ShardSpec, ComputePattern, ComputeSpec
from colossalai.nn._ops.cache_embedding import CachedParamMgr, FreqAwareEmbeddingBag, ParallelFreqAwareEmbeddingBag
from colossalai.nn.parallel.layers import CachedParamMgr, FreqAwareEmbeddingBag, ParallelFreqAwareEmbeddingBag
NUM_EMBED, EMBED_DIM = 10, 8
BATCH_SIZE = 8