[PyTorch] Finalize quantization updates at logical backward boundaries - #34
Draft
pggPL wants to merge 9 commits into
Draft
[PyTorch] Finalize quantization updates at logical backward boundaries#34pggPL wants to merge 9 commits into
pggPL wants to merge 9 commits into
Conversation
…ization state update The backward amax reduction for delayed scaling was triggered from the backward of the 'first FP8 module' of an autocast, assuming its backward runs last. That assumption breaks with activation recompute (duplicate forward/backward updates per step) and with branched autograd graphs (ownership can land in a frame whose backward never runs). Instead, module backwards now only request the update (schedule_backward_quantization_update); it is flushed exactly once at the next top-level autocast entry, after the backward pass is complete. Forward updates at autocast exit are skipped during the recompute phase. Removes the is_first_fp8_module save/restore bookkeeping from checkpoint contexts and modules; renames reduce_and_update_fp8_tensors to reduce_and_update_quantization_state (old name kept as alias). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
test_recipe and test_fusible_ops asserted backward scales immediately after backward(); with the deferred update they flush explicitly. Drop raw torch.utils.checkpoint mode from new tests (unsupported with FP8 weight caching independently of this change) and relax exact amax history comparison for nested checkpoints (inner checkpoint re-records amaxes during outer recompute; scales unaffected). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…factor # Conflicts: # transformer_engine/pytorch/distributed.py # transformer_engine/pytorch/ops/fuser.py
…tization_update The method only idempotently marks the update as pending, so 'request' is more precise. Context booleans renamed to should_request_backward_quantization_update accordingly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Keeps them in a file already enumerated by the L0 QA suite instead of adding a new one. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The activation-checkpoint duplicate-update problem and its regression test scenarios were first diagnosed and addressed in NVIDIA#3213; this PR supersedes that approach by removing first-module ownership entirely. Co-authored-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
FP8GlobalStateManager now only orchestrates when updates run; how a recipe updates its process-global state is a RecipeState classmethod (reduce_and_update_global_state), no-op by default and overridden by DelayedScalingRecipeState with the amax reduction + scale recompute. The recipe-to-state-class mapping is extracted from RecipeState.create into class_for_recipe and reused for dispatch. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The reduction no longer runs in module backwards, so the old per-module NVTX markers had nothing left to wrap; mark it at its new single execution site in flush_backward_quantization_update instead. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Schedule one update after the outermost autograd task and expose an explicit scope for logical backwards made of multiple autograd calls or delayed weight-gradient work. Also make nested activation-recompute bookkeeping stack-safe and preserve metadata for modules that remain in eval mode. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
changed the title
6 tasks
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
1 participant
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
Replace first-FP8-module ownership with an explicit logical-backward boundary for stateful quantization updates.
The old ownership rule assumed that the first FP8 module in forward executes last in backward. That does not hold for branched autograd graphs, unused checkpoint branches, or activation recompute. It can therefore update delayed-scaling state too early, more than once, or not at all.
Module backwards now only request an update. With the public
torch.autograd.graph.queue_callbackAPI from pggPL/pytorch#3, Transformer Engine queues one callback on the outermost autograd task and flushes after the complete top-level backward. Runtimes without that API retain the next-top-level-autocast fallback.The PR also adds
transformer_engine.pytorch.backward_quantization_update_scope(). Frameworks whose logical backward spans multiple independent autograd calls, or schedules work after autograd returns, can use the scope to defer one combined update until all work is complete. Delayedmodule.backward_dw()calls must be inside this scope.Activation-recompute bookkeeping is made stack-safe as well. An inner checkpoint forward executed during an outer replay restores the parent snapshot while reserving a snapshot for its own replay. Metadata stashing no longer depends on
module.training, so an eval-mode module can legally participate in an autograd checkpoint. A replay without a matching stash now raises a descriptive error.Supersedes the ownership approach in NVIDIA#3213 and covers the practical eval-to-eval failure discussed in NVIDIA#3394.
Type of change
Changes
backward_dw().Checklist
git diff --checkis clean.Prepared with Codex.