Conversation
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
marked this pull request as ready for review
September 20, 2026 07:06
VenusTZZ
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 20, 2026 07:06
delock
self-requested a review
September 20, 2026 07:21
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #8604
Follow-up to #4318 (graph operations for short kernel sequences).
Summary
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):
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
Known limitation