Skip to content

[PyTorch] Document heterogeneous quantization recipes - #3336

Open
negvet wants to merge 3 commits into
NVIDIA:mainfrom
negvet:fine_grained_recipe_docs
Open

[PyTorch] Document heterogeneous quantization recipes#3336
negvet wants to merge 3 commits into
NVIDIA:mainfrom
negvet:fine_grained_recipe_docs

Conversation

@negvet

@negvet negvet commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Description

Document heterogeneous quantization recipes.
This is enabled after hybrid quantization support #2817 is merged

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

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

Signed-off-by: Evgeny <etsykunov@nvidia.com>
@negvet
negvet requested a review from ksivaman as a code owner August 10, 2026 12:20
@negvet negvet added the documentation Improvements or additions to documentation label Aug 10, 2026
@negvet
negvet requested review from ptrendx and timmoon10 August 10, 2026 12:21
@greptile-apps

greptile-apps Bot commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR documents heterogeneous PyTorch quantization recipes and corrects the example’s wgrad provenance description.

  • Adds conceptual guidance for role-aware quantizer factories, GEMM-direction mapping, hybrid quantizers, and high-precision slots.
  • Adds a runnable heterogeneous quantization example and supporting diagrams.
  • Expands API documentation for hybrid and identity quantizers and tensor types.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported wgrad provenance discrepancy is corrected and the current documentation matches the quantizer configuration and Linear backward path.

Important Files Changed

Filename Overview
docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py The example now accurately describes wgrad as consuming the original BF16 columnwise input and output-gradient representations.
docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst Adds consistent documentation for role selection, GEMM-direction mapping, columnwise provenance, and high-precision quantization slots.
transformer_engine/common/recipe/init.py Clarifies the CustomRecipe factory contract, including optional roles and the requirement to return a valid quantizer for every slot.
transformer_engine/pytorch/tensor/hybrid_tensor.py Marks HybridQuantizer as experimental without changing runtime behavior.
transformer_engine/pytorch/tensor/identity_tensor.py Expands IdentityTensor API documentation and marks IdentityQuantizer as experimental without changing runtime behavior.

Reviews (3): Last reviewed commit: "resolve comments" | Re-trigger Greptile

@pggPL
pggPL self-requested a review August 10, 2026 12:52
…/pytorch_fine_grained_quantization_example.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com>
@pggPL

pggPL commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

First observation: this part is done in very different style than the rest of the Features section.
Two most important things:

  1. We use boxes for code.
  2. It would be nice to have some pictures in .svg (they are quite easy to vibe code).

@pggPL

pggPL commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

First observation: when I start reading it there is a lot of defensive text about experimental API and some philosophical section Configuration readiness. I think it would be great to start, after short introduction with short example what we can do - using more than one recipe with the picture with some code snippet. Then we can elaborate on the API in more detail.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is too long - we can have something before #START_FINE_GRAINED_QUANTIZATION_EXAMPLE or in different files like model building and require_supported_hardware() . BTW see how we deal with support hardware in boxes in the other recipes docs.

@pggPL

pggPL commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Note that we can split this: short docs in Features and tutorial in Tutorials like we have done with gemm profiling tutorial.

@negvet negvet changed the title [PyTorch] Fine-grained recipe docs [PyTorch] Document heterogeneous quantization recipes Aug 17, 2026
Signed-off-by: Evgeny <etsykunov@nvidia.com>
@negvet

negvet commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator Author

@pggPL thanks for the comments, I applied all your suggestions

@negvet
negvet requested a review from pggPL August 17, 2026 12:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants