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

[embedding] rename FreqAwareEmbedding -> CachedEmbedding (#1699)

Browse Source
This commit is contained in:
Jiarui Fang
2022-10-13 22:22:27 +08:00
committed by GitHub
parent 0e52f3d3d5
commit 21962e1593
8 changed files with 77 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/nn/parallel/layers/cache_embedding/cache_mgr.py
View File

@@ -352,7 +352,8 @@ class CachedParamMgr(torch.nn.Module):
# move sure the cuda rows will not be evicted!
with record_function("(cache) prepare_rows_on_cuda"):
self._prepare_rows_on_cuda(comm_cpu_row_idxs)
with self.timer("prepare_rows_on_cuda") as timer:
self._prepare_rows_on_cuda(comm_cpu_row_idxs)
self.evict_backlist = torch.tensor([], device=cpu_row_idxs.device, dtype=cpu_row_idxs.dtype)