github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-04-28 19:55:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] fix zero comm buffer init (#6154)

Browse Source
This commit is contained in:
Hongxin Liu 2024-12-10 16:46:15 +08:00 committed by GitHub
parent 8d826a336e
commit de3d371f65
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/low_level/low_level_optim.py
View File

@ -371,7 +371,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
for i, sz in enumerate(bucket_store.sizes): for i, sz in enumerate(bucket_store.sizes):
grp = bucket_store.torch_pg if len(bucket_store.sizes) == 1 else bucket_store.torch_pg[i] grp = bucket_store.torch_pg if len(bucket_store.sizes) == 1 else bucket_store.torch_pg[i]
flat_grads_list = list(cur_flat_grads.split(len(cur_flat_grads) // sz)) flat_grads_list = list(cur_flat_grads.split(len(cur_flat_grads) // sz))
received_grad = torch.zeros_like(flat_grads_list[0]) received_grad = torch.empty_like(flat_grads_list[0])
if self._fp8_communication: if self._fp8_communication:
reduce_scatter_fp8( reduce_scatter_fp8(
received_grad, received_grad,