Skip to content

[PyTorch] Delegate delayed-scaling updates to RecipeState - #38

Open
pggPL wants to merge 1 commit into
mainfrom
recipe_state_global_update_refactor
Open

[PyTorch] Delegate delayed-scaling updates to RecipeState#38
pggPL wants to merge 1 commit into
mainfrom
recipe_state_global_update_refactor

Conversation

@pggPL

@pggPL pggPL commented Aug 20, 2026

Copy link
Copy Markdown
Owner

Description

Extract the non-behavioral recipe-state refactor from #34. FP8GlobalStateManager continues to own update timing and iteration over the existing delayed-scaling global buffers, while the corresponding RecipeState subclass owns the update algorithm.

This does not change the global buffer registry, update cadence, or make non-delayed recipes stateful.

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

  • Extract recipe-to-RecipeState resolution into RecipeState.class_for_recipe.
  • Reuse that mapping for state construction and global update dispatch.
  • Move the existing delayed-scaling concat/reduce/scale-update implementation into DelayedScalingRecipeState.reduce_and_update_global_state.

Test plan

  • pytest -q tests/pytorch/test_recipe.py plus targeted delayed-scaling CustomRecipe and fuser tests: 58 passed, 92 skipped.
  • python -m pylint transformer_engine/pytorch/quantization.py: 10.00/10.
  • Forced unfused module tests reproduce the same 12 failures on an unmodified main; the fuser variants pass in both cases.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • No documentation changes are required for this internal refactor
  • My changes generate no new warnings
  • Existing tests cover the refactored behavior
  • New and existing relevant unit tests pass locally with my changes

Keep FP8GlobalStateManager responsible for iterating registered delayed-scaling buckets, while RecipeState subclasses define the corresponding update algorithm. Extract recipe-to-state dispatch so construction and global updates share one mapping.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant