deepspeedai / deepspeedai/DeepSpeed

[REQUEST] NPU: implement graph operations for short kernel sequences (follow-up to #4318)

Open
#8,604 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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_process raise 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"]: trueengine.graph_harvestingBF16_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.