Skip to content

Commit 573ce48

Browse files
Skip zero-sized parameters in HP fragment mapping (#8298)
## Summary - exclude zero-sized parameters from ZeRO-1/2 HP-fragment overlap mapping - avoid constructing an empty fragment when an empty parameter's offset lies inside a partition Fixes #8297. ## Testing - current `master` (`edaa722`), 2x A6000: minimized ZeRO-1 reproducer fails before the change with `fragment start 8 should be < fragment_end 8` and completes after the change - 2-process CPU/Gloo: minimized reproducer completes with both ZeRO-1 and ZeRO-2 after the change - the trailing zero-sized-parameter reproducer from #8279 still completes - `python3 -m py_compile deepspeed/utils/mixed_precision_linkage.py` Signed-off-by: Cao Yuhang <caoyuhang@fwerkor.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
1 parent f76ab88 commit 573ce48

1 file changed

Lines changed: 5 additions & 2 deletions

File tree

deepspeed/utils/mixed_precision_linkage.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,14 @@ def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_gr
4141
# 1) current_offset < partition_end,
4242
# 2) current_offset + lp_param.numel() >= partition_start
4343
lp_param_end = current_offset + lp_param.numel()
44-
if current_offset < partition_end and lp_param_end > partition_start:
45-
param_and_offset_list.append((lp_param, current_offset))
44+
in_partition = (partition_start <= current_offset < partition_end
45+
or current_offset < partition_start < lp_param_end)
46+
if in_partition:
4647
lp_param._index_in_param_group = index_in_param_group
4748
# Indices for params in this partition/GPU
4849
index_in_param_group += 1
50+
if lp_param.numel() > 0:
51+
param_and_offset_list.append((lp_param, current_offset))
4952
current_offset += lp_param.numel()
5053

5154
return param_and_offset_list

0 commit comments

Comments
 (0)