deepspeedai / deepspeedai/DeepSpeed

[BUG][Deepcompile] reduce_grad returns undefined tensor -> Inductor compilation fails (expected a proper tensor but got None)

Open
#7,682 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug training
Dominant language
Python
Stars
43.1k
Forks
5k
Avg merge
4d 15h
Merged PRs (30d)
112

Description

Describe the bug
During AOTAutograd backward compilation, DeepSpeed’s reduce_grad op returns an undefined tensor, but the graph rewrite pass rewires all downstream gradient usages to this output.
As a result, Inductor/FakeTensor sees None as input to ops like aten.sum or reshape, causing compilation failure.

Error

torch._inductor.exc.InductorError: RuntimeError:
Expected a proper Tensor but got None (or an undefined Tensor in C++) for argument #0 'self'

Trigger path

  1. Backward graph: each parameter-grad node is rewritten to torch.ops.dc.reduce_grad.default(grad)
  2. All uses of the original grad are replaced by the output of this op
  3. Fx trace shows downstream ops (e.g., aten.sum(...,[0,1]), reshape) consuming the output of reduce_grad.
  4. c++ implementation returns at::Tensor() (undefined) in both:
  • reduce_grad()
  • reduce_grad_meta()
    This breaks FakeTensor propagation and Inductor lowering.

Root Cause
reduce_grad is treated as a functional node in the graph, but its c++ kernel and meta kernel return a undefined tensor, which cannot be consumed by downstream ops.

Since the compiler rewrites all gradient uses to this output, the output must be a valid Tensor.

Question for maintainers
In DeepSpeed/csrc/compile/deepcompile.cpp, both reduce_grad(...) and reduce_grad_meta(...) currently return an undefined tensor (at::Tensor()).
Given that the graph rewrite redirects all downstream gradient uses to the output of this op, should these two functions instead return the input grad_tensor?

This would allow downstream ops (e.g., aten.sum, reshape) to receive a valid tensor and avoid FakeTensor/Inductor errors during compilation. Is returning grad_tensor the correct fix here, or is the intended semantics different?

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

Start in DeepSpeed/csrc/compile/deepcompile.cpp and inspect reduce_grad(...) and reduce_grad_meta(...), then trace the AOTAutograd rewrite described in the issue. Reproduce the failure through the Deepcompile/Inductor path and determine the intended output semantics. Done means downstream gradient users receive a valid tensor and the reported aten.sum or reshape compilation failure no longer occurs.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.