deepspeedai / deepspeedai/DeepSpeed
[REQUEST] NPU: implement graph operations for short kernel sequences (follow-up to #4318)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Is your feature request related to a problem? Please describe.
#4318 introduced accelerator graph operations (create_graph / capture_to_graph / replay_graph) to capture short kernel sequences — gradient norm, gradient clipping, and the BF16 optimizer's hp→lp gradient write-back, enabled via the graph_harvesting config and driven by graph_process in deepspeed/runtime/utils.py. CUDA (#4318), HPU (#4912), MLU (#6472), and SUPA (#8054) have all implemented the trio on top of their runtime graph APIs.
The NPU accelerator still carries the stub that #4318 originally left: the three methods return None / a no-op context / nothing. The feature is therefore dead on NPU — graph_process caches the graph and "replays" it every step, but the replay does nothing, so the captured computation silently stops running (verified on real hardware: from the second step on, the reported gradient norm stops tracking the actual value, clipping stops being applied, and the hp→lp write-back stops, with no error raised).
Describe the solution you'd like
Implement the three methods on NPU_Accelerator (accelerator/npu_accelerator.py) with the graph API torch_npu ships, mirroring the HPU/MLU/SUPA templates:
def create_graph(self):
return torch.npu.NPUGraph()
def capture_to_graph(self, graph, pool=None, stream=None):
return torch.npu.graph(graph, pool, stream)
def replay_graph(self, graph):
graph.replay()
Add hardware-independent contract tests (tests/unit/accelerator/test_npu_accelerator.py) that stub torch.npu and run on CPU-only CI, following the pattern of tests/unit/accelerator/test_mps_accelerator.py. Real-device validation (numerical parity and the graph_harvesting end-to-end path, with performance data) is reported in the PR, as done in #8532.
Describe alternatives you've considered
- Leave the stub as-is. The silent no-op replay trap remains for anyone who enables
graph_harvesting, and the launch-bound speedup stays unrealized on Ascend. - Make
graph_processraise when the accelerator cannot capture. Prevents the silent skip but provides none of the performance benefit, and still needs this implementation to ever be removed. - Wait for the torch_npu graph API to stabilize. The interface shape used here (graph object +
graph()context manager +replay()) is exactly what the abstract interface already assumes and what the other vendor runtimes were integrated with at the time.
Additional context
- Config chain:
ds_config["graph_harvesting"]: true→engine.graph_harvesting→BF16_Optimizer(graph_harvesting=True)→graph_process. The chain works on NPU once the trio is implemented, verified end-to-end on Ascend 910B4: the three recorded functions are captured and replay matches eager execution. - Validation environment: CANN 8.5.0, torch 2.9.0+cpu, torch_npu 2.9.0.post7, Ascend 910B4 (aarch64).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read accelerator/npu_accelerator.py and compare its graph methods with the HPU, MLU, and SUPA implementations. Follow tests/unit/accelerator/test_mps_accelerator.py to add CPU-only contract coverage using a stubbed torch.npu; done means the three methods support graph capture and replay, with real-device graph_harvesting validation reported separately.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100