[negvet]

Description

Enables recipe updates during training, either lazily through te.autocast() or atomically across a model with te.apply_recipe().

Diff:

• Library: +2,016 / −794 / net +1,222
• Tests: +5,829 / −274 / net +5,555

TODO - callable compatibility:
Callable factories and algorithms now require explicit semantic metadata, supplied through a decorator or recipe key argument. Consider whether unkeyed callables should remain supported and, if so, what safe fallback identity should be used.

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

• Identifies recipe changes through semantic configuration keys.
• Builds and validates complete forward/backward runtimes before publishing them.
• Invalidates cached quantized weights when their recipe changes.
• Supports role-aware CustomRecipe factories and composed MHA/DPA boundaries.
• Rejects unsupported transitions before modifying active state.
• Preserves existing numerics for unchanged supported configurations.

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

[negvet]

/te-ci pytorch L1

[negvet]

/te-ci pytorch L1

[greptile-apps]

Confidence Score: 3/5

The PR is not safe to merge until legacy unkeyed CustomRecipe factories remain activatable and custom policies are checked against the hardware capabilities their actual quantizers require.

Findings

1.  Unkeyed factories no longer activate ▶
2.  Identity recipes require FP8 ▶

Summary

This PR introduces semantic recipe configurations and transactional forward/backward runtime construction so PyTorch quantization recipes can be updated lazily or model-wide. It also propagates role-aware recipe state through linear, grouped-linear, attention, CUDA-graph, and fusible-operation paths.

• Adds te.apply_recipe() and per-module runtime planning/activation.
• Rebuilds quantizers and invalidates cached quantized weights when semantic recipe configuration changes.
• Adds stable keys for callable algorithms and custom quantizer factories.
• Extends MHA/DPA boundary roles and grouped-quantization validation.
• Adds extensive recipe-update, attention, graph, distributed, and numerics coverage.
• The new callable-key and blanket CustomRecipe hardware checks introduce two compatibility failures.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Recipe requested by autocast or apply_recipe] --> B[Build semantic recipe configuration]
  B --> C[Resolve module quantizer roles]
  C --> D[Construct forward and backward candidate states]
  D --> E[Validate candidate runtime]
  E -->|Failure| F[Keep active runtime unchanged]
  E -->|Success| G[Publish candidate runtime]
  G --> H[Clear quantized weight workspaces]
  H --> I[Execute with updated quantizers]

Loading

_{Reviews (1) · Last reviewed commit: "Merge branch 'main' into mid_train_recip..."}

[greptile-apps]

Unkeyed factories no longer activate

Existing callers can construct CustomRecipe(qfactory=callable) without a semantic key, but the first enabled autocast, calibration, or model initialization now calls quantizer_config() and raises here before invoking the factory. This breaks the previously supported callable-factory contract and prevents existing unkeyed recipes from running unless callers add new metadata.

Knowledge Base Used: PyTorch FP8 and low-precision quantization

[greptile-apps]

Identity recipes require FP8

If a CUDA environment can run Transformer Engine but fails the generic FP8 availability check, an identity-only CustomRecipe is rejected here before its factory is inspected. The documented high_precision_factory returns only IdentityQuantizers and performs no FP8 work, so this blanket gate prevents a supported high-precision policy from running even though it does not need FP8 kernels.

Knowledge Base Used: PyTorch FP8 and low-precision quantization

[zhongbozhu]

Why do we even need to dynamically change recipe during training?

[timmoon10]

This might conflict with #3717.

[timmoon10]

Why doesn't this work? Can't we just let the fuser re-fuse the ops if the recipe changes?

[timmoon10]

This attr is set when entering te.autocast, so we don't expect this to be read if the function is called outside te.autocast. Consider replacing the default recipe instead, or keep the current API where the user has to pass the recipe into te.autocast.

[timmoon10]

I see we have this scheme so that we can process all the modules before actually applying the recipe change, which means we only apply the recipe change after all the modules have had some kind of validation. It adds a lot of complexity though, so skeptical if it is worth it.

[timmoon10]

I think these are redundant with the existing logic for _output_quantizer_role and _grad_input_quantizer_role. If I understand, this is intended to allow MHA to change the output/grad input quantizer and detect recipe changes. However, I don't think there's value in having separate paths for users to set output quantizers and MHA to set output quantizers, they should use the same infrastructure.

[timmoon10]

The quantizer config infrastructure seems overengineered and redundant. #3717 uses the string representation to detect recipe changes in the op fuser.