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

[NFC] polish doc style for ColoTensor (#1457)

Browse Source
This commit is contained in:
Jiarui Fang
2022-08-16 09:21:05 +08:00
committed by GitHub
parent 0dbd61c29b
commit a1476ea882
9 changed files with 197 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -99,9 +99,11 @@ class CachedParamMgr(torch.nn.Module):
@torch.no_grad()
def reorder(self, ids_freq_mapping: Optional[List[int]] = None, warmup_ratio=0.7):
"""reorder the weight according to ids' frequency in dataset before training.
"""reorder
reorder the weight according to ids' frequency in dataset before training.
Also Build the IndexMappingTable, aka index_mapping_table.
Execute only once before training.
Args:
ids_freq_mapping (List[int]): a list, idx is id number, value is freq. if None no reorder
warmup_ratio (float): the amount of chunks preloaded in cuda cache