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

[NFC] update chunk manager API (#2119)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-12 16:57:22 +08:00
committed by GitHub
parent e99edfcb51
commit e5aa8333e4
5 changed files with 100 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
colossalai/nn/parallel/data_parallel.py
View File

@@ -228,16 +228,16 @@ class ZeroDDP(ColoDDP):
fp32_p = ColoTensor(fp32_data, spec=ColoTensorSpec(p.process_group))
p.data = p.data.half()
dp_world_size = p.process_group.dp_world_size()
self.chunk_manager.append_tensor(tensor=p,
group_type='fp16_param',
config_key=dp_world_size,
cpu_offload=cpu_offload,
pin_memory=pin_memory)
self.chunk_manager.append_tensor(tensor=fp32_p,
group_type='fp32_param',
config_key=dp_world_size,
cpu_offload=cpu_offload,
pin_memory=pin_memory)
self.chunk_manager.register_tensor(tensor=p,
group_type='fp16_param',
config_key=dp_world_size,
cpu_offload=cpu_offload,
pin_memory=pin_memory)
self.chunk_manager.register_tensor(tensor=fp32_p,
group_type='fp32_param',
config_key=dp_world_size,
cpu_offload=cpu_offload,
pin_memory=pin_memory)
self.fp32_params.append(fp32_p)
self.grads_device[p] = self.gemini_manager.default_device
self.chunk_manager.close_all_groups()