Skip to content

[PyTorch] Schedule delayed-scaling updates after backward - #39

Open
pggPL wants to merge 2 commits into
recipe_state_global_update_refactorfrom
backward_quantization_update_scheduler
Open

[PyTorch] Schedule delayed-scaling updates after backward#39
pggPL wants to merge 2 commits into
recipe_state_global_update_refactorfrom
backward_quantization_update_scheduler

Conversation

@pggPL

@pggPL pggPL commented Aug 20, 2026

Copy link
Copy Markdown
Owner

Description

Stacked on #38.

Delayed-scaling state updates are currently owned by the first FP8 module seen in forward and run from that module's backward. This is sensitive to backward traversal order, unused branches, reentrant checkpointing, and schedules that delay weight-gradient computation.

This PR replaces forward-order ownership with a backward-boundary scheduler. Every participating module requests an update, and TE queues one callback for the enclosing autograd GraphTask. Reentrant TE checkpointing pre-registers the callback on the outer task before entering its nested backward.

Ordinary .backward() remains automatic. The new backward_quantization_update_scope() is only needed when an application defines a larger logical backward, for example multiple autograd calls in a 1F1B schedule or backward_dw() executed later.

The implementation intentionally uses the existing private Variable._execution_engine.queue_callback hook. It does not depend on an unmerged PyTorch API and has no deferred/fallback update path.

This builds on the investigation and test cases from NVIDIA#3213.

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

  • Queue one delayed/custom recipe-state update at the end of each autograd task.
  • Remove backward update ownership from forward module order.
  • Pre-register on the outer GraphTask for reentrant TE checkpoints.
  • Add an optional public scope for custom logical-backward boundaries.
  • Preserve the CUDA graph path and the legacy reduce_and_update_fp8_tensors alias.
  • Cover plain, reentrant, non-reentrant, nested, branched, unused-branch, multi-backward, and delayed-wgrad 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
  • 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

Testing

  • PYTHON_ONLY=1 bash TE/qa/L0_pytorch_lint/test.sh
  • pytest -q TE/tests/pytorch/test_recipe.py TE/tests/pytorch/test_backward_override.py TE/tests/pytorch/test_fusible_ops.py::TestFuser::test_fp8_scale_update (1380 passed, 9864 skipped)
  • Focused scheduler suite (10 passed)

pggPL and others added 2 commits August 20, 2026 17:54
Queue one quantization state update at the autograd boundary instead of assigning it to the first FP8 module seen in forward. Add an optional logical-backward scope for multi-backward schedules and delayed weight-gradient computation.

Co-authored-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Document that graphs produced under one autocast need an explicit logical-backward scope when their backward calls should share one delayed-scaling update.

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