Skip to content

Add missing knob for reduce_scatter_with_fp32_accumulation #4410

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asolergi-nv merged 5 commits into from
Apr 25, 2026
+12 −3
Merged

Add missing knob for reduce_scatter_with_fp32_accumulation #4410

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fce4557
add missing knob reduce_scatter_with_fp32_accumulation
WanZzzzzz Apr 21, 2026
9fb4165
rebase
WanZzzzzz Apr 23, 2026
3e40da3
lint
WanZzzzzz Apr 23, 2026
5c70ea7
explicitly del old fp4 rowwise data
WanZzzzzz Apr 24, 2026
5e736e7
Merge branch 'main' into fp4_param_small_fix
asolergi-nv Apr 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions megatron/core/distributed/distributed_data_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -265,10 +265,18 @@ def __init__(
# If bucketing is explicitly disabled, then put all buckets in a buffer into a single
# bucket group.
self.bucket_groups = partition_buckets(
self.buffers, force_single_bucket_group=disable_bucketing
self.buffers,
force_single_bucket_group=disable_bucketing,
reduce_scatter_with_fp32_accumulation=(
self.ddp_config.reduce_scatter_with_fp32_accumulation
),
)
self.expert_parallel_bucket_groups = partition_buckets(
self.expert_parallel_buffers, force_single_bucket_group=disable_bucketing
self.expert_parallel_buffers,
force_single_bucket_group=disable_bucketing,
reduce_scatter_with_fp32_accumulation=(
self.ddp_config.reduce_scatter_with_fp32_accumulation
),
)

if self.ddp_config.num_distributed_optimizer_instances > 1:
Expand Down
3 changes: 2 additions & 1 deletion megatron/core/fp4_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,8 @@ def modify_nvfp4_rowwise_storage(fp4_tensor: torch.Tensor, new_rowwise_data: tor
), "Rowwise NVFP4 storage must be uint8"
# Preserve existing values and then swap storage
new_rowwise_data.detach().copy_(old_rowwise)
setattr(fp4_tensor, "_rowwise_data", new_rowwise_data)
fp4_tensor._rowwise_data = new_rowwise_data
del old_rowwise


def quantize_nvfp4_param_shard(
Expand Down
Loading