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

[embedding] rollback for better FAW performance (#1625)

Browse Source
This commit is contained in:
Jiarui Fang
2022-09-22 11:16:25 +08:00
committed by GitHub
parent d925122020
commit 38c68b5b9a
2 changed files with 50 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_layers/test_cache_embedding.py
View File

@@ -90,7 +90,7 @@ def test_reorder_with_freq():
offset_in_chunk = torch.tensor(offset_in_chunk, dtype=torch.long, device=dev)
weight = torch.rand(num_embed, 2)
mgr = CachedParamMgr(weight, num_chunk, use_cpu_caching=dev.type == 'cpu')
mgr = CachedParamMgr(weight, num_chunk)
mgr.reorder(idx_map)