From fe0efd5012201e68115216135e0442764de4dddc Mon Sep 17 00:00:00 2001 From: simpleqt <89645338+simpleqt@users.noreply.github.com> Date: Mon, 7 Sep 2026 01:31:51 +0800 Subject: [PATCH] fix(zenflow): restore assert and remove duplicated identical branch in gradient copy The ZenFlow fork of async_inplace_copy_grad_to_fp32_buffer_from_gpu mangled the original logic into two byte-identical if/else arms. If grad_accum were ever None, the None branch would itself call grad_accum.view(-1) and raise AttributeError, defeating the check. Restored the upstream assert + single-assignment shape from runtime/zero/stage_1_and_2.py. Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com> --- deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py index 29c702d6813d..9624abb354b8 100644 --- a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py +++ b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py @@ -685,10 +685,8 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param): 0, dest_offset, num_elements) grad_accum = self.get_param_gradient_attribute(param) - if grad_accum is None: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) - else: - src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) + assert grad_accum is not None + src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements) if src_tensor.dtype != self.master_weights_and_grads_dtype: src_tensor = src_tensor.to(self.master_weights_and_grads_dtype)