Fix torch.onnx._internal imports for torch 2.9

[kevalmorabia97]

What does this PR do?

Type of change: Import fixes

Testing

Summary by CodeRabbit

• Refactor

  • Improved PyTorch compatibility for ONNX export and Diffusers attention by using forward-referenced types and version-aware imports.
  • Aligned compile/export mode checks with current PyTorch APIs for more robust behavior.

• Chores

  • Added type-checking guards to ensure stable operation across multiple PyTorch releases.

• Documentation

  • Clarified compatibility-related behavior in code comments.

• Notes

  • No functional or API changes expected; ONNX export and model execution should behave as before with improved cross-version stability.

[coderabbitai]

Walkthrough

Refactors typing and import patterns to use forward-referenced "GraphContext" and string-literal torch.Value annotations, adds TYPE_CHECKING guards and version-aware imports for torch.onnx internals, updates attention/export symbolics accordingly, and switches is_compiling checks from torch._dynamo to torch.compiler in tensor quantizer logic.

Changes

Cohort / File(s)	Summary
ONNX export typing & version-guards modelopt/torch/quantization/export_onnx.py	Replaces direct jit_utils.GraphContext/Value annotations with string-based forward refs; adds TYPE_CHECKING and conditional imports; adjusts version-aware selection of ONNX symbolics; updates many function signatures (int8/int4/fp8/fp4/mxfp8 and attention) to use "GraphContext" and "torch._C.Value" types.
Tensor quantizer compile-check modelopt/torch/quantization/nn/modules/tensor_quantizer.py	Replaces torch._dynamo.is_compiling() with torch.compiler.is_compiling() in validation/forward paths; no API or behavioral changes beyond the check source.
Diffusers FP8 SDPA symbolic typing modelopt/torch/quantization/plugins/diffusers.py	Removes explicit _onnx_symbolic binding; adds TYPE_CHECKING and conditional GraphContext imports; updates FP8SDPA.symbolic signature to string-literal forward refs for GraphContext and Value types.

Sequence Diagram(s)

sequenceDiagram
    actor U as Caller
    participant D as Diffusers.FP8SDPA.symbolic
    participant E as export_onnx.scaled_dot_product_attention
    participant O as ONNX Graph

    Note over D,E: GraphContext and Value types are forward-referenced (string) with TYPE_CHECKING guards.

    U->>D: call symbolic(g, q, k, v, attn_mask?, scale?)
    alt torch.onnx version A
        D->>E: delegate SDPA export (version-aware imports)
    else torch.onnx version B
        D->>E: delegate SDPA export (alt symbolics)
    end
    E->>O: emit attention nodes (scale/mask/causal paths)
    O-->>U: constructed subgraph

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

A nibble of types, a hop through time,
I forward-reference, ears in line.
Guards at import, versions sway—
yet graphs still build the ONNX way.
With compiler checks I softly tread,
Carrots committed, bugs unfed. 🥕

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The title "Fix torch.onnx._internal imports for torch 2.9" concisely and accurately summarizes the primary change in this PR—making torch.onnx internal imports and typing compatible with PyTorch 2.9—and directly matches the modifications in export_onnx.py and related files to use forward references and version-conditional imports.
Docstring Coverage	✅ Passed	Docstring coverage is 82.35% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch kmorabia/torch2.9-fixes

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7d5f636 and 0a4521a.

📒 Files selected for processing (3)

• modelopt/torch/quantization/export_onnx.py (18 hunks)
• modelopt/torch/quantization/nn/modules/tensor_quantizer.py (3 hunks)
• modelopt/torch/quantization/plugins/diffusers.py (3 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

modelopt/torch/quantization/utils.py (1)

• is_torch_export_mode (332-334)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)

• GitHub Check: linux
• GitHub Check: wait-checks / wait
• GitHub Check: wait-checks / wait
• GitHub Check: build-docs
• GitHub Check: code-quality

🔇 Additional comments (11)

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (3)

551-551: LGTM! API migration to new torch.compiler module.

The change from torch._dynamo.is_compiling() to torch.compiler.is_compiling() correctly migrates to the public API. The official documentation confirms that torch.compiler.is_compiling() indicates whether a graph is executed/traced as part of torch.compile() or torch.export(), and notes that torch._dynamo.external_utils.is_compiling() should be deprecated eventually.

---

917-917: LGTM! Consistent API migration.

This change maintains consistency with the earlier migration from torch._dynamo.is_compiling() to torch.compiler.is_compiling() in line 551.

---

883-884: Confirm torch.onnx._internal import pattern across supported PyTorch versions

The repo uses hasattr(...) fallbacks to torch.onnx._internal.torchscript_exporter.* (e.g. modelopt/torch/quantization/nn/modules/tensor_quantizer.py:881–884; similar patterns in export_onnx.py and plugins/diffusers.py). Verify the flat modules (torch.onnx._globals / _type_utils / symbolic_opset14, etc.) exist on older supported PyTorch and that the _internal.torchscript_exporter paths exist on newer releases — consider switching to try/except ImportError for more robust detection.

modelopt/torch/quantization/export_onnx.py (5)

106-106: LGTM! Proper TYPE_CHECKING import.

The addition of TYPE_CHECKING import enables forward reference typing without runtime overhead.

---

113-117: LGTM! Version-aware GraphContext import strategy.

The conditional import logic properly handles the movement of GraphContext from torch.onnx._internal.jit_utils to torch.onnx._internal.torchscript_exporter.jit_utils in torch >= 2.9, wrapped in TYPE_CHECKING to avoid runtime imports.

---

133-133: LGTM! Consistent forward reference typing.

The systematic replacement of concrete GraphContext type annotations with string literals "GraphContext" enables proper forward referencing while maintaining type safety. The parameter type annotations for torch._C.Value types are also consistently updated with string literals.

Also applies to: 192-192, 216-216, 244-244, 271-271, 287-294, 371-378, 513-513, 549-549, 564-564, 575-575, 593-593, 611-611, 629-629

---

298-309: LGTM! Version-aware JitScalarType import.

The conditional import pattern correctly handles the relocation of JitScalarType and symbolic functions between torch versions. This follows the same pattern as the GraphContext imports and maintains backward compatibility.

---

343-356: LGTM! Consistent import pattern for symbolic utilities.

The conditional imports for JitScalarType and symbolic functions (_attention_scale, _causal_attention_mask) maintain consistency across all functions that require these utilities, properly handling the module reorganization in torch >= 2.9.

Also applies to: 410-421, 472-485

modelopt/torch/quantization/plugins/diffusers.py (3)

21-21: LGTM! TYPE_CHECKING import for forward references.

Proper import for enabling forward reference type annotations.

---

31-35: LGTM! Consistent version-aware GraphContext import.

The conditional import pattern matches the implementation in export_onnx.py, ensuring consistency across the codebase for handling GraphContext imports between torch versions.

---

211-218: LGTM! Forward reference type annotations.

The conversion to string literal type annotations for GraphContext and torch._C.Value types maintains consistency with the patterns established in export_onnx.py while enabling proper forward referencing.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

❌ Patch coverage is 12.50000% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.46%. Comparing base (8b1cedf) to head (0a4521a).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/export_onnx.py	6.66%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #356      +/-   ##
==========================================
- Coverage   73.48%   73.46%   -0.02%     
==========================================
  Files         172      172              
  Lines       17636    17640       +4     
==========================================
  Hits        12960    12960              
- Misses       4676     4680       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.