From 9acf2392d486c632c216c487ace8c7a6193c50ae Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 9 Jun 2026 16:04:42 +0000 Subject: [PATCH 1/3] Consistent fp32 grads flow --- deepspeed/runtime/zero/stage3.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 7ebb42905456..c5adc50b73c9 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1584,8 +1584,8 @@ def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor, if self.postscale_gradients and self.gradient_predivide_factor != world_sz: buffer_to_reduce = buffer_to_reduce.mul(self.gradient_predivide_factor) - if communication_data_type != self.dtype: - buffer_to_reduce = buffer_to_reduce.to(self.dtype) + if communication_data_type != self.gradient_accumulation_dtype: + buffer_to_reduce = buffer_to_reduce.to(self.gradient_accumulation_dtype) grad_partitions = [] grad_offset_in_buffer = 0 @@ -1636,8 +1636,8 @@ def __avg_scatter_grads(self, params_to_reduce: List[Parameter], self.dp_process_group): grad_partitions_for_rank = [g.mul(self.gradient_predivide_factor) for g in grad_partitions_for_rank] - if communication_data_type != self.dtype: - grad_partitions_for_rank = [g.to(self.dtype) for g in grad_partitions_for_rank] + if communication_data_type != self.gradient_accumulation_dtype: + grad_partitions_for_rank = [g.to(self.gradient_accumulation_dtype) for g in grad_partitions_for_rank] return grad_partitions_for_rank From 07a84412d98f07fcc49a71c06177f8ca8e824863 Mon Sep 17 00:00:00 2001 From: Olatunji Ruwase Date: Tue, 9 Jun 2026 16:51:35 +0000 Subject: [PATCH 2/3] Coverage --- deepspeed/runtime/zero/stage3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index c5adc50b73c9..a9a96973c526 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1540,7 +1540,7 @@ def _apply_distributed_muon_update(self, communication_data_type: torch.dtype, b start_offset = rank * chunk_sz end_offset = start_offset + chunk_sz if end_offset > param.grad.numel(): - buffer_to_update = torch.zeros(chunk_sz, device=param.grad.device, dtype=param.grad.dtype) + buffer_to_update = torch.zeros(chunk_sz, device=param.grad.device, dtype=self.gradient_accumulation_dtype) buffer_to_update[:param.grad.numel() - start_offset] = gathered_momentum.view(-1).data[start_offset:param.grad.numel()] else: @@ -1599,7 +1599,7 @@ def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor, partition = buffer_to_reduce[start_offset:end_offset] if param.partition_numel() != partition.numel(): - padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=grad.dtype) + padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=self.gradient_accumulation_dtype) if partition.numel() > 0: padded_partition[:partition.numel()] = partition grad_partitions.append(padded_partition) From ae20d9d2c06100c2537b148d5696ab7d63a72f53 Mon Sep 17 00:00:00 2001 From: tunji-ruwase_snow Date: Tue, 9 Jun 2026 18:30:37 +0000 Subject: [PATCH 3/3] Format Signed-off-by: tunji-ruwase_snow --- deepspeed/runtime/zero/stage3.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index a9a96973c526..fc81d8bc414f 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1540,7 +1540,9 @@ def _apply_distributed_muon_update(self, communication_data_type: torch.dtype, b start_offset = rank * chunk_sz end_offset = start_offset + chunk_sz if end_offset > param.grad.numel(): - buffer_to_update = torch.zeros(chunk_sz, device=param.grad.device, dtype=self.gradient_accumulation_dtype) + buffer_to_update = torch.zeros(chunk_sz, + device=param.grad.device, + dtype=self.gradient_accumulation_dtype) buffer_to_update[:param.grad.numel() - start_offset] = gathered_momentum.view(-1).data[start_offset:param.grad.numel()] else: @@ -1599,7 +1601,9 @@ def __avg_scatter_contiguous_grads(self, buffer_to_reduce: Tensor, partition = buffer_to_reduce[start_offset:end_offset] if param.partition_numel() != partition.numel(): - padded_partition = torch.zeros(param.partition_numel(), device=grad.device, dtype=self.gradient_accumulation_dtype) + padded_partition = torch.zeros(param.partition_numel(), + device=grad.device, + dtype=self.gradient_accumulation_dtype) if partition.numel() > 0: padded_partition[:partition.numel()] = partition grad_partitions.append(padded_partition)