[PyTorch] Schedule delayed-scaling updates after backward - #39
Open
pggPL wants to merge 2 commits into
Open
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 newbackward_quantization_update_scope()is only needed when an application defines a larger logical backward, for example multiple autograd calls in a 1F1B schedule orbackward_dw()executed later.The implementation intentionally uses the existing private
Variable._execution_engine.queue_callbackhook. 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
Changes
GraphTaskfor reentrant TE checkpoints.reduce_and_update_fp8_tensorsalias.Checklist:
Testing
PYTHON_ONLY=1 bash TE/qa/L0_pytorch_lint/test.shpytest -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)