Skip to content

Clear the gradient ZenFlow just offloaded, not param.grad - #8601

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-zenflow-offload-clear-20260919
Open

vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-zenflow-offload-clear-20260919

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

With ZenFlow, ZeRO-1 and a gradient accumulation dtype different from the parameter dtype, each window adds to the gradient the previous one already consumed, so gradients grow step after step.

ZenFlowZeroOptimizerParallel.async_inplace_copy_grad_to_fp32_buffer_from_gpu overrides the base method only to write into the double-buffered overlap_grad. It reads the source through get_param_gradient_attribute, which returns param.grad_accum when use_grad_accum_attribute is set, but clears param.grad, which _fill_param_grad_accum_attribute already set to None. grad_accum survives, and that function adds into it when it is not None. The base clears through clear_grad_attribute. The override also guards grad_accum is None with two identical branches, both calling .view() on it. Mirror the base class: assert the gradient is present, and clear the attribute that was read.

Test: DS_ACCELERATOR=cpu python -m pytest tests/unit/runtime/zenflow/ -q. Two regressions fail before: the grad_accum case leaves the consumed gradient in place, and the missing-gradient case raises "'NoneType' object has no attribute 'view'". Both pass after; 15 CPU tests pass across the ZenFlow directory. They drive the real method against a minimal stand-in, so no engine is needed. yapf 0.43.0 and flake8 pass on both files. Executed on Apple M2 Pro CPU; no GPU or distributed run. The offload path needs a GPU, so upstream CI covers it end to end.

Mirror the base class: assert the gradient is present, and clear the attribute that was read.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant