Implement graph operations for NPU accelerator (#8605)
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 #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 (#4912), MLU (#6472), and SUPA (#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.
---------
Signed-off-by: Zizhao Tang <1601677496@qq.com>
