Skip to content

Implement graph operations for NPU accelerator - #8605

Open
VenusTZZ wants to merge 6 commits into
deepspeedai:masterfrom
VenusTZZ:npu-graph-operations
Open

VenusTZZ wants to merge 6 commits into
deepspeedai:masterfrom
VenusTZZ:npu-graph-operations

Conversation

@VenusTZZ

@VenusTZZ VenusTZZ commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

Fixes #8604

Follow-up to #4318 (graph operations for short kernel sequences).

Summary

  • Implement create_graph / capture_to_graph / replay_graph on NPU_Accelerator via the graph API torch_npu ships: torch.npu.NPUGraph() creates the graph object, torch.npu.graph(graph, pool, stream) is the capture context manager, and graph.replay() replays. Previously NPU_Accelerator carried the stub that Capture short kernel sequences to graph #4318 originally left (None / a no-op context / nothing), so graph_process cached an empty graph and every subsequent "replay" did nothing: with graph_harvesting enabled, gradient norms silently went stale, clipping silently stopped, and the hp→lp gradient write-back silently stopped — reproduced on hardware (see below).
  • Mirrors the HPU (Add Habana Labs HPU accelerator support #4912), MLU ([Accelerator] Cambricon MLU support #6472), and SUPA (Add Biren SUPA accelerator support #8054) implementations, which wrap their runtime graph APIs the same way.
  • Adds tests/unit/accelerator/test_npu_accelerator.py with three hardware-independent contract tests that stub torch.npu and run on CPU-only CI, following the pattern of tests/unit/accelerator/test_mps_accelerator.py.

Implementation notes

The torch_npu graph API. NPUGraph wraps the runtime's graph object with capture_begin / capture_end / replay / reset / pool. torch.npu.graph(graph, pool, stream) is the context manager that performs the synchronization, cache cleanup, side-stream switching, and capture begin/end around the with block. pool and stream pass through positionally and default to None; graph_process captures without pool/stream, so the defaults are part of the contract. The graph APIs are core torch_npu APIs available across the torch versions DeepSpeed supports, so the methods call them directly with no existence checks — matching the CUDA/HPU/MLU/SUPA implementations.

What it unblocks. The whole graph_harvesting chain already existed and is accelerator-agnostic: ds_config["graph_harvesting"]: true → engine.graph_harvesting → BF16_Optimizer(graph_harvesting=True) → graph_process (gradient norm, clipping, update_hp_grads). Implementing the trio is the only missing piece on NPU.

Performance on Ascend 910B4
Whole-block timing with torch.npu.synchronize() on both sides of the block, 3 rounds, best round (shared-box methodology), CANN 8.5.0, torch 2.9.0+cpu, torch_npu 2.9.0.post7, Ascend 910B4 (aarch64):

Workload eager (us/iter) graph replay (us/iter) speedup
200 independent tiny ops (norm-style) 6456 910 7.1x
50-step chained small ops (clip-style) 3067 135 22.7x
20 ops on a 1M-element tensor 1271 307 4.1x

End-to-end: with graph_harvesting: true on Ascend 910B4, _norm_tensors, clip_tensors, and _update_hp_grads_func are all captured into graph_cache, and 5 training steps (BF16_Optimizer, bf16 + fp32 gradient accumulation) are bitwise-identical to eager execution (equal losses, torch.equal parameters).

Before/after: the same graph_process scenario run against the stub vs. this implementation — with the stub the reported norm froze at its first-step value, clipping stopped being applied, and the hp accumulation stopped, with no error raised; with this change all three track eager exactly.

Test plan

  • pre-commit run --files on touched paths — yapf, flake8, check-torchdist, check-torchcuda, check-license, codespell all pass
  • pytest tests/unit/accelerator/ — 5 passed (3 new NPU contract tests + 2 existing MPS tests)
  • pytest tests/unit/v1/pin_memory/ — 45 passed (no regression)
  • Real-device on Ascend 910B4: capture → replay parity across two consecutive replays and explicit pool/stream arguments; graph_harvesting end-to-end chain; stub-vs-fix corruption demonstration

Known limitation

  • Pre-existing and unrelated to NPU: engine.step() calls zero_grad() with the torch default set_to_none=True, which trips the graph_harvesting assert in deepspeed/runtime/bf16_optimizer.py on every accelerator. The end-to-end validation above calls optimizer.step() with zero_grad(set_to_none=False); worth a separate issue.

Replace the deepspeedai#4318-era stub (create_graph returning None, no-op capture,
empty replay) with torch_npu's graph API: torch.npu.NPUGraph(),
torch.npu.graph(graph, pool, stream), and graph.replay(), mirroring the
HPU/MLU/SUPA implementations. The stub made graph_process a silent no-op
on every replay after the first call, so enabling graph_harvesting on NPU
silently froze gradient norms, clipping, and the hp->lp write-back.

Add hardware-independent contract tests (tests/unit/accelerator/
test_npu_accelerator.py) that stub torch.npu so they run on CPU-only CI.

Signed-off-by: Zizhao Tang <1601677496@qq.com>
@VenusTZZ VenusTZZ changed the title Npu graph operations Implement graph operations for NPU accelerator Sep 20, 2026
@VenusTZZ
VenusTZZ marked this pull request as ready for review September 20, 2026 07:06
@delock
delock self-requested a review September 20, 2026 07:21

This branch has not been deployed

No deployments
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.

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

1 participant