[muutot]

Description

QuantizerRole, QuantizerRequest, and DelayedScalingRequest are plain dataclasses, but they were defined in quantization.py. Since HybridQuantizer.__init__ needs QuantizerRequest, the tensor package had to import it lazily from inside a method, and quantization.py compensated with five deferred quantizer imports in RecipeState.make_quantizers() (Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer). Three of those carry TODO(ksivamani) notes asking for a better design. The lazy imports date back to TE 2.0, where the tensor package genuinely imported fp8; that back-edge is long gone, and this PR removes what it left behind.

Fixes: no tracking issue — resolves the TODO(ksivamani) notes in transformer_engine/pytorch/quantization.py.

Type of change

• Code refactoring

Changes

• Add transformer_engine/pytorch/quantizer_types.py, a leaf module containing QuantizerRole, QuantizerRequest, and DelayedScalingRequest (definitions unchanged). It imports nothing from the rest of transformer_engine.pytorch.
• quantization.py: import and re-export the three dataclasses from the new module (still listed in __all__), hoist the five deferred quantizer imports to module level, and drop the now-unused import dataclasses.
• tensor/hybrid_tensor.py, custom_recipes/quantizer_factories.py, custom_recipes/quantizer_factory_zoo.py: import the dataclasses from the leaf module; no more local imports of quantization.
• Backward compatible: transformer_engine.pytorch.quantization.QuantizerRole and friends still resolve, so user code, docs (autoapi), and old pickles are unaffected.
• The tensor-package imports in quantization.py are placed before from .utils import ..., with a comment: .utils pulls in quantized_tensor (via the debug tooling), which imports a helper from inside tensor/, so the tensor package must be initializing for that import to succeed. This removes the reliance on the current import order within the package.

No behavior change; pure import restructuring.

Testing

• black (24.4.2, repo settings) and pylint (3.3.1, repo pylintrc) on all touched files: no new warnings.
• Static import-graph check: transformer_engine.pytorch.tensor no longer reaches quantization at module level, and no new cycles are introduced.
• python -c "import transformer_engine.pytorch"
• python -c "import transformer_engine.pytorch.quantization"
• pytest tests/pytorch/test_custom_recipe.py

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

[Low impact] The PR appears safe to merge; no actionable issues were found.

Summary

Moves QuantizerRole, QuantizerRequest, and DelayedScalingRequest into quantizer_types.py to remove the tensor package’s dependency on quantization.py.

• Tensor and recipe code share quantizer types without importing them from quantization.py.

No actionable issues found. Runtime tests were not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Q["quantization.py"] --> T["quantizer_types.py"]
  Q --> P["tensor package"]
  H["tensor/hybrid_tensor.py"] --> T
  F["custom recipe factories"] --> T
  A["Existing public type imports"] --> Q

Loading

_{Reviews (1) · Last reviewed commit: "[PyTorch] Break circular import between ..." · Reviewed by Greptile}

[muutot]

@ksivaman , Please review this PR and let me know if any adjustments are needed. Thanks for your time!