Conversation
The ZenFlow optimizer module in engine_stage3.py hardcodes device='cuda' in 4 locations (grad_sum_buffer, selected_indices, sum_buffer, and critic_sum_buffer tensor creation). This causes a crash on non-CUDA accelerators (NPU, XPU, HPU, etc.) with 'Torch not compiled with CUDA enabled'. The DeepSpeedZeroOptimizer_Stage3 already has a 'device' attribute set via get_accelerator().current_device_name(), which correctly reports the active accelerator. Replace all hardcoded 'cuda' with optimizer_z3.device. DeepSpeed already has a device-agnostic accelerator abstraction layer (deepspeed.accelerator.get_accelerator()) used throughout the codebase, so this change is consistent with the project's own design. Signed-off-by: li-lizhe <147392333@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 90950325ba
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| total_chunk_size += param_max_chunk_size | ||
|
|
||
| optimizer_z3.grad_sum_buffer = torch.zeros(total_chunk_size, dtype=optimizer_z3.dtype, device='cuda') | ||
| optimizer_z3.grad_sum_buffer = torch.zeros(total_chunk_size, dtype=optimizer_z3.dtype, device=optimizer_z3.device) |
There was a problem hiding this comment.
When ZeRO optimizer offload is enabled, DeepSpeedZeroOptimizer_Stage3.device is deliberately set to CPU (stage3.py:329), including in the existing ZenFlow integration configuration (tests/unit/runtime/zenflow/test_zf.py:48-52). This therefore moves grad_sum_buffer to CPU, but line 227 immediately passes it to the accelerator process group (NCCL on CUDA, HCCL on NPU), which cannot collect CPU tensors; the same problem affects the newly moved auto-update buffers at lines 257-260. Use the active accelerator device independently of the optimizer-state device, such as get_accelerator().current_device_name(), so non-CUDA support does not regress optimizer-offloaded training.
AGENTS.md reference: AGENTS.md:L35-L36
Useful? React with 👍 / 👎.
1 participant
Summary
The ZenFlow optimizer module in
engine_stage3.pyhardcodesdevice='cuda'in 4 locations when creating tensors (grad_sum_buffer,selected_indices,sum_buffer, andcritic_sum_buffer). This causes a crash on non-CUDA accelerators (NPU, XPU, HPU, etc.) withTorch not compiled with CUDA enabled.Root cause
The
DeepSpeedZeroOptimizer_Stage3already has adeviceattribute set viaget_accelerator().current_device_name(), which correctly reports the active accelerator. DeepSpeed has a device-agnostic accelerator abstraction (deepspeed.accelerator.get_accelerator()) used throughout the codebase.Fix
Replace all 4 hardcoded
device='cuda'withdevice=optimizer_z3.device.Testing
On a non-CUDA accelerator (Ascend NPU), the hardcoded
device='cuda'raises:RuntimeError: Torch not compiled with CUDA enabledWith this fix, tensor creation uses
optimizer_z3.devicewhich resolves to the active accelerator, and ZenFlow selection logic runs correctly.On CUDA,
optimizer_z3.deviceiscuda(viaget_accelerator().current_device_name()), so behavior is unchanged.Related
Consistent with DeepSpeed's existing device-agnostic design via
get_accelerator().