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

[FCE] update interface for frequency statistics in FreqCacheEmbedding (#1462)

Browse Source
This commit is contained in:
Geng Zhang
2022-08-23 17:38:24 +08:00
committed by GitHub
parent ede326298b
commit 0aad53c62b
4 changed files with 30 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_layers/test_cache_embedding.py
View File

@@ -44,7 +44,7 @@ def synthesize_1d_sparse_feature(
def test_cachemgr():
model = torch.nn.EmbeddingBag(10000, 128)
# 10 chunks, 5 in cuda
mgr = CachedParamMgr(model.weight, 5)
mgr = CachedParamMgr(model.weight.detach(), 5)
assert mgr.cuda_row_num == 5
mgr._admit(1)
@@ -74,8 +74,8 @@ def test_reorder_with_freq():
chunk_size = 1
num_chunk = 5
idx_map = np.random.randint(10000, size=(num_embed,))
sorted_idx = np.flipud(np.argsort(idx_map)).tolist()
idx_map = torch.randint(10000, size=(num_embed,))
sorted_idx = torch.argsort(idx_map, descending=True).tolist()
chunkid, offset_in_chunk = [], []
for i in range(num_embed):
idx = sorted_idx.index(i)
@@ -231,6 +231,6 @@ def test_parallel_freq_aware_embed(world_size):
if __name__ == '__main__':
# test_cachemgr()
test_cachemgr()
# test_freq_aware_embed()
test_parallel_freq_aware_embed(2)
# test_parallel_freq_aware_embed(2)