test_dtensor.py (float8 TP compile): aten.abs.default NotImplemented on Float8TrainingTensor after removing allow-in-graph autograd wrappers

[github-actions]

Commits affected: ca7488145beeda7416ceee0db43eb6f5124b2ebc
Run: https://github.com/pytorch/ao/actions/runs/33666967452

Summary

_test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile time with NotImplementedError: attempting to run aten.abs.default, this is not supported, raised from Float8TrainingTensor.__torch_dispatch__. Routing the Float8 autograd.Function classes through the standard Dynamo AutogradFunctionApply HOP (this commit) causes the backward to be retraced via make_fx, and torch.abs in the amax computation now dispatches onto a Float8TrainingTensor, which has no abs handler.

Failure details

• Failed job(s) (both fail identically, deterministically):
  • test (H100, mt-l-x86iamx-88-900-h100-4, --pre torch torchvision --index-url .../nightly/cu126)
  • test (H100 2.13, mt-l-x86iamx-88-900-h100-4, torch==2.13.0 torchvision==0.28.0 --index-url .../cu126)
• Test step: ./test/float8/test_everything_multi_gpu.sh → test/float8/test_dtensor.py FAILED (torchrun ChildFailedError, exit code 1). Preceding subtests (✅ All Tests Passed ✅) pass; the failure is isolated to the compile TP dtensor test.
• Representative error:

  Test _test_fp8_mlp_tensor_parallelism_compile failed with error: backend='inductor' raised:
  NotImplementedError: attempting to run aten.abs.default, this is not supported
  
  While executing %abs_1 : [num_users=1] = call_function[target=torch.abs](args = (%grad_output_reshaped,), kwargs = {})
  

• Dispatch chain (from the original traceback):
  • torchao/float8/float8_linear.py:112 backward → hp_tensor_to_float8_dynamic
  • torchao/float8/float8_scaling_utils.py:56 → tensor_to_scale
  • torchao/float8/float8_utils.py:106 tensor_to_scale → tensor_to_amax
  • torchao/float8/float8_utils.py:65 tensor_to_amax → amax = torch.max(torch.abs(x))
  • raised at torchao/float8/float8_training_tensor.py:373 __torch_dispatch__: raise NotImplementedError(f"attempting to run {func}, this is not supported")
• Compile-time framing: the backward runs under torch._functorch.autograd_function.backward → torch.fx.Interpreter(bwd).run(...) during AOTAutograd make_fx tracing, i.e. the op executes against the Float8 subclass while retracing, not during normal eager backward.

Root cause

Commit ca74881 ("Remove Float8 allow-in-graph custom autograd wrappers") routes the four Float8 autograd.Function classes through the standard Dynamo AutogradFunctionApply higher-order op instead of the previous allow-in-graph trampoline. Under torch.compile + DTensor tensor parallelism, the backward is now captured and retraced by make_fx/FX Interpreter. During that retrace, tensor_to_amax's torch.abs(grad_output_reshaped) is dispatched onto a Float8TrainingTensor, whose __torch_dispatch__ only services ops present in FLOAT8_OPS_TABLE; aten.abs.default is not registered there, so it raises NotImplementedError and inductor compilation fails.

Evidence this is a regression introduced by this commit (not flaky/infra):

• The immediately preceding 4xH100 run on main (parent commit a195aaf8a2da5b8496018d39e5261e93d01fb0ca, run 33632094558) and the prior 7 runs of this workflow were all green.
• The failing code path — Float8 autograd backward tracing through the HOP — is exactly what this commit changed.
• The failure is deterministic and reproduces identically on both the torch-nightly and torch-2.13.0 matrix legs, ruling out a per-run flake or a single-nightly toolchain change.

Classification

real regression — a deterministic compile-time failure in the exact code path modified by the triggering commit, green on the parent commit and reproducing on both matrix legs. Not intermittent (so not a flaky test) and not an environment/toolchain issue (so not infra).

Suggested fix

• In Float8TrainingTensor's op table (torchao/float8/float8_ops.py / float8_training_tensor.py), the abs/amax path taken during HOP-based backward retracing needs to resolve against the underlying high-precision tensor rather than the subclass. Options:
  • Have tensor_to_amax/tensor_to_scale operate on the plain (dequantized/underlying) tensor data so torch.abs/torch.max never dispatch onto Float8TrainingTensor during retrace; or
  • Register aten.abs.default (and any other amax-path ops surfaced) in FLOAT8_OPS_TABLE to delegate to the underlying data.
• Add a regression test that compiles the Float8 backward through the AutogradFunctionApply HOP path (the TP-compile dtensor scenario) so this is caught pre-merge.
• If a fix is not immediately available, consider reverting ca74881's autograd-wrapper change or gating it until the amax dispatch under retrace is handled.

[github-actions]

⚠️ Run Regression Tests on ROCm also failed on this commit (ca7488145beeda7416ceee0db43eb6f5124b2ebc) — https://github.com/pytorch/ao/actions/runs/33666967935

[github-actions]

Recurrence on main — same root cause as this issue's diagnosis.

• Run: https://github.com/pytorch/ao/actions/runs/33902500930 (Run 4xH100 tests)
• Commit: b7ac3aacf9b6f1cb0c6bb29143329f601dfcfcb1 ("Fix cvt_packfloat compatibility with nvidia-cutlass-dsl 4.6.0+", Fix cvt_packfloat compatibility with nvidia-cutlass-dsl 4.6.0+ #4766 — unrelated to float8; the ca74881 autograd-wrapper regression remains unfixed on main and keeps breaking this workflow).

Matches the existing diagnosis exactly:

• Test: _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py (via ./test/float8/test_everything_multi_gpu.sh); preceding subtests pass (✅ All Tests Passed ✅), then torchrun ChildFailedError, exit code 1.
• Error: backend='inductor' raised: NotImplementedError: attempting to run aten.abs.default, this is not supported, While executing %abs_1 = call_function[target=torch.abs](args=(%grad_output_reshaped,)).
• Same dispatch chain: torch/_functorch/autograd_function.py:874 backward → torch.fx.Interpreter(bwd).run(...) (HOP AutogradFunctionApply backward retrace under make_fx) → raised at torchao/float8/float8_training_tensor.py:373 __torch_dispatch__.
• Reproduces identically and deterministically on both matrix legs: --pre torch nightly cu126 and torch==2.13.0 cu126 — confirming it is not a per-run flake or single-nightly toolchain change.

No new information vs. the existing analysis; classification remains a real regression. Still needs the abs/amax path to resolve against the underlying high-precision tensor during HOP-based backward retrace (or aten.abs.default registered in FLOAT8_OPS_TABLE), or a revert/gate of ca74881.

[github-actions]

Recurs on main at commit 3efa9eda3bd1347fa2485d16641bbc5bc2124d53 — Run 4xH100 tests, https://github.com/pytorch/ao/actions/runs/34224961671.

Same root cause and signature as the existing diagnosis (no divergence):

• _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile time with NotImplementedError: attempting to run aten.abs.default, this is not supported.
• Raised from torchao/float8/float8_training_tensor.py:373 __torch_dispatch__, reached via the AutogradFunctionApply HOP backward retrace: torch/_functorch/autograd_function.py backward → torch.fx.Interpreter(bwd).run(...), executing %abs_1 = call_function[target=torch.abs](args=(%grad_output_reshaped,)) during AOTAutograd make_fx tracing.
• Fails identically and deterministically on both matrix legs (torch==2.13.0 and --pre torch nightly, both cu126); torchrun exits with ChildFailedError → exit code 1. The earlier float8 subtests still pass (✅ All Tests Passed ✅).

The triggering commit 3efa9ed ("Reject unsupported wrap_with_set_grad_enabled subgraphs in PT2E prepare APIs") does not touch float8 — this is the unfixed regression from ca74881 still present on main. Confirms it remains a real regression, not flaky/infra.

[github-actions]

Recurs on main at commit 6792133e5eb7c607d27f69756870e44e7ffed6af (run https://github.com/pytorch/ao/actions/runs/34280478407, "Run 4xH100 tests").

Same root cause and signature as this issue's diagnosis — no new information:

• _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile time with NotImplementedError: attempting to run aten.abs.default, this is not supported, raised from Float8TrainingTensor.__torch_dispatch__ (torchao/float8/float8_training_tensor.py:373).
• Dispatch path is identical: AOTAutograd make_fx retraces the backward through the AutogradFunctionApply HOP → torch._functorch.autograd_function.backward → torch.fx.Interpreter(bwd).run(...), so torch.abs in the amax computation dispatches onto the Float8 subclass.
• Fails deterministically on BOTH matrix legs (torch==2.13.0 ... cu126 and --pre torch ... nightly/cu126); the preceding test_everything_multi_gpu.sh subtests pass (✅ All Tests Passed ✅), then torchrun exits 1 via ChildFailedError on test/float8/test_dtensor.py.

The triggering commit 6792133 (fix(_is_mslk_available): guard against missing CUDA/ROCm runtime, #4544) does not touch float8 code — this is the same persistent regression from ca74881 still unfixed on main, not a new cause. Classification unchanged: real regression.

[github-actions]

⚠️ Run Regression Tests on ROCm also failed on this commit (6792133e5eb7c607d27f69756870e44e7ffed6af) — https://github.com/pytorch/ao/actions/runs/34280478457

[github-actions]

Recurs on main at commit 6f7bb3874302ebfc7bf7c15449a9c1d93d0ca549 — Run 4xH100 tests, run https://github.com/pytorch/ao/actions/runs/34342903087.

Same root cause and signature as this issue's diagnosis (no change): _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile with NotImplementedError: attempting to run aten.abs.default, this is not supported, raised while retracing the Float8 backward:

• torch.max(torch.abs(x)) in tensor_to_amax (torchao/float8/float8_utils.py:65)
• dispatched onto Float8TrainingTensor → __torch_dispatch__ (torchao/float8/float8_training_tensor.py:373)
• While executing %abs_1 : ... call_function[target=torch.abs](args = (%grad_output_reshaped,))

Both matrix legs fail identically and deterministically (torch==2.13.0 cu126 job 102437611971, and --pre torch nightly cu126 job 102437612207), ending in ChildFailedError / test/float8/test_dtensor.py FAILED, exit code 1.

The triggering commit 6f7bb38 ("Refactor MXTensor: align the swizzle execution...") does not touch float8 code; this is the same unresolved regression from ca74881 still present on main, not a new cause. Still a real regression — the amax torch.abs dispatch onto Float8TrainingTensor during HOP-based backward retrace remains unhandled.

[github-actions]

⚠️ Run Regression Tests on ROCm also failed on this commit (6f7bb3874302ebfc7bf7c15449a9c1d93d0ca549) — https://github.com/pytorch/ao/actions/runs/34342903178

[github-actions]

Recurs on main at commit 6ded493033ebba47e2631b087cb17f01992d51f0 — run https://github.com/pytorch/ao/actions/runs/34356567625 (Run 4xH100 tests). Same root cause and diagnosis as this issue, still unfixed.

_test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile time with NotImplementedError: attempting to run aten.abs.default, this is not supported, raised from Float8TrainingTensor.__torch_dispatch__. The dispatch chain is identical to the recorded one:

• torchao/float8/float8_linear.py:112 backward → hp_tensor_to_float8_dynamic
• torchao/float8/float8_scaling_utils.py:56 → tensor_to_scale
• torchao/float8/float8_utils.py:106 tensor_to_scale → tensor_to_amax
• torchao/float8/float8_utils.py:65 tensor_to_amax → amax = torch.max(torch.abs(x))
• raised at torchao/float8/float8_training_tensor.py:373 __torch_dispatch__

The backward is retraced via torch._functorch.autograd_function.backward → torch.fx.Interpreter(bwd).run(...) under AOTAutograd make_fx, so torch.abs dispatches onto Float8TrainingTensor during retrace — matching the existing diagnosis (HOP-based AutogradFunctionApply backward retrace).

Both matrix legs fail identically and deterministically:

• test (H100, mt-l-x86iamx-88-900-h100-4, --pre torch torchvision --index-url .../nightly/cu126)
• test (H100 2.13, mt-l-x86iamx-88-900-h100-4, torch==2.13.0 torchvision==0.28.0 --index-url .../cu126)

torchrun exits with ChildFailedError; preceding float8 multi-GPU subtests pass (✅ All Tests Passed ✅), failure isolated to the compile TP dtensor test. Confirms the regression is still present on main at this newer commit.

[github-actions]

Still failing on main at commit 22afe13b49e1dcc841ec89c616b5bfb020c357f1 — run https://github.com/pytorch/ao/actions/runs/34381027616 (Run 4xH100 tests).

Matches the existing diagnosis exactly. Both matrix legs fail identically and deterministically:

• test (H100, ... --pre torch ... nightly/cu126)
• test (H100 2.13, ... torch==2.13.0 ... cu126)

test/float8/test_dtensor.py::_test_fp8_mlp_tensor_parallelism_compile raises backend='inductor' raised: NotImplementedError: attempting to run aten.abs.default, this is not supported (torchrun ChildFailedError, exit 1). Preceding subtests print ✅ All Tests Passed ✅; the failure is isolated to the compile TP dtensor test. Same dispatch chain as reported: float8_linear.py:112 backward → float8_scaling_utils.py:56 hp_tensor_to_float8_dynamic → float8_utils.py:106 tensor_to_scale → float8_utils.py:65 tensor_to_amax → amax = torch.max(torch.abs(x)), dispatched onto Float8TrainingTensor during HOP-based backward retrace.

The triggering commit 22afe13 ("Make quantize_ skip already-quantized weights") does not touch float8 code; this is the same underlying regression from ca74881 continuing to fail on main, not a new cause. No divergence from the existing root-cause analysis.

[github-actions]

Recurred on main at commit ee0d7ce304cc341fc6894df1f3b32038db0da1c3 — run https://github.com/pytorch/ao/actions/runs/34399534209 ("Run 4xH100 tests").

Matches the existing diagnosis exactly. _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile time with:

NotImplementedError: attempting to run aten.abs.default, this is not supported
While executing %abs_1 : [num_users=1] = call_function[target=torch.abs](args = (%grad_output_reshaped,), kwargs = {})

Same dispatch chain as before: float8_linear.py:112 backward → hp_tensor_to_float8_dynamic → float8_scaling_utils.py:56 → tensor_to_scale → float8_utils.py:106 tensor_to_amax → float8_utils.py:65 amax = torch.max(torch.abs(x)), raised at float8_training_tensor.py:373 __torch_dispatch__. Backward is retraced through the AutogradFunctionApply HOP via torch._functorch/autograd_function.py:874 → torch.fx.Interpreter(bwd).run(...) during AOTAutograd make_fx, so torch.abs dispatches onto Float8TrainingTensor, which has no aten.abs.default handler.

Both matrix legs fail identically (nightly cu126 and torch==2.13.0), with a ChildFailedError / exit code 1 from torchrun; preceding subtests pass ("✅ All Tests Passed ✅"). This is the same real regression from ca74881 still present on main — the triggering commit ee0d7ce is a ruff/lint-only change (fix invalid escape sequences, W605) that does not touch torchao/float8, so it did not introduce the failure; it just re-surfaces the unfixed root cause. No new information vs. the existing diagnosis.

[github-actions]

Recurred on main at commit 3005bc1d2407646103ff759f38beb5c59a453728 ("[CI] Enable low_bit_optim CI on intel GPU").
Run: https://github.com/pytorch/ao/actions/runs/34564632115 (Run 4xH100 tests)

Same root cause and signature as the existing diagnosis — no new information:

• Both matrix legs fail identically and deterministically: test (H100 2.13, torch==2.13.0 ...) and test (H100, --pre torch ... nightly/cu126).
• test/float8/test_dtensor.py → _test_fp8_mlp_tensor_parallelism_compile fails with backend='inductor' raised: NotImplementedError: attempting to run aten.abs.default, this is not supported, via tensor_to_amax → amax = torch.max(torch.abs(x)) at torchao/float8/float8_utils.py:65, raised from Float8TrainingTensor.__torch_dispatch__. Both ranks fail; torchrun exits 1 with ChildFailedError. Preceding float8 subtests still print ✅ All Tests Passed ✅, so the failure remains isolated to the TP-compile dtensor test.

The triggering commit is a CI-config change and does not touch float8; this is the same regression introduced by ca74881 continuing to break main, not a new cause. Classification unchanged: real regression.

[github-actions]

Recurs on main at commit 5a33bcf6efb8e78acb858d93eb7a9b7c57718693 — Run 4xH100 tests, run https://github.com/pytorch/ao/actions/runs/34637189681.

Same root cause and signature as this issue's existing diagnosis (no divergence):

• Test: _test_fp8_mlp_tensor_parallelism_compile (test/float8/test_dtensor.py), via ./test/float8/test_everything_multi_gpu.sh. torchrun exits 1; preceding subtests (✅ All Tests Passed ✅) pass. Failure reproduces on the torch==2.13.0 matrix leg.
• Error: backend='inductor' raised: NotImplementedError: attempting to run aten.abs.default, this is not supported, While executing %abs_1 : [num_users=1] = call_function[target=torch.abs](args = (%grad_output_reshaped,), kwargs = {}).
• Dispatch chain unchanged: float8_linear.py:112 backward → hp_tensor_to_float8_dynamic (float8_scaling_utils.py:56) → tensor_to_scale (float8_utils.py:106) → tensor_to_amax (float8_utils.py:65, amax = torch.max(torch.abs(x))) → raised at float8_training_tensor.py:373 __torch_dispatch__. Backward runs under torch._functorch.autograd_function.backward → fx.Interpreter during AOTAutograd make_fx retrace.

Note: the triggering commit 5a33bcf ("Add code owners for low-bit CPU kernels", #4874) is a CODEOWNERS-only change and does not touch float8 code — this is the pre-existing regression (from ca74881, removing the allow-in-graph autograd wrappers) still unfixed on main, not a new cause. Classification remains: real regression.

[github-actions]

Recurs on main at commit 9e5ea7fde7132df0c09b2acb12c2defad90aecb9 (run https://github.com/pytorch/ao/actions/runs/34637555688, "Run 4xH100 tests").

Same root cause and signature as this issue's diagnosis — no new information:

• _test_fp8_mlp_tensor_parallelism_compile in test/float8/test_dtensor.py fails at inductor compile with NotImplementedError: attempting to run aten.abs.default, this is not supported.
• Raised from torchao/float8/float8_training_tensor.py:373 __torch_dispatch__, reached via torch.fx.Interpreter(bwd).run(...) (torch/_functorch/autograd_function.py:874) while AOTAutograd make_fx retraces the Float8 autograd backward through the AutogradFunctionApply HOP — exactly the path changed by ca74881.
• Both matrix legs fail identically and deterministically: test (H100 2.13, torch==2.13.0) and test (H100, --pre torch nightly cu126). Preceding float8 subtests print ✅ All Tests Passed ✅; only the compile TP dtensor test fails, then torchrun exits 1 (ChildFailedError → test/float8/test_dtensor.py FAILED).

The triggering commit 9e5ea7f ("Fix MoE FP8 Triton autotune key") does not touch float8/dtensor code, so this is a continued recurrence of the pre-existing regression on main, not a new cause. Still unresolved; classification remains real regression.

[github-actions]

⚠️ Run Regression Tests on ROCm also failed on this commit (9e5ea7fde7132df0c09b2acb12c2defad90aecb9) — https://github.com/pytorch/ao/actions/runs/34637555679