fix(dist): keep profiler record-function ops out of SAC replay accounting - #3133
Merged
Conversation
…ting torch 2.13's FSDP2 runs its pre/post-forward hooks under torch.autograd.profiler.record_function, which emits dispatchable torch.ops.profiler ops. When an FSDP module boundary sits inside a selective-activation-checkpointed region (MoE experts sharded separately inside a checkpointed decoder block), the hooks fire a different number of times during the backward recompute than during the forward. SAC replays the forward op stream by per-op invocation index, so the extra profiler op shifts the stream and training fails with: RuntimeError: profiler._record_function_enter_new.default invocation index 3 encountered during backward but not found in storage. Seen on qwen3_6_35b_medpix_ep8cp2_4k with the pytorch:26.06-py3 base (torch 2.13.0a0+8145d630e8); upstream SAC_IGNORED_OPS does not cover profiler ops as of pytorch/pytorch@main. Range ops carry no tensors SAC could cache, so adding them to SAC_IGNORED_OPS only removes them from replay accounting - they still execute. Applied from both SAC context factories (activation_checkpointing and the MoE parallelizer). CPU regression test enters a record_function range only during the recompute pass and asserts backward succeeds. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
nemo-ci validation on the pytorch:26.06-py3 base (torch 2.13.0a0+8145d630e8), 20 steps, eos: qwen3_6_35b_medpix_ep8cp2_4k — PASSED with this fix ✅
Known remaining torch-2.13 issue (separate from this PR): 🤖 Generated with Claude Code |
Contributor
Author
|
/claude review |
Contributor
|
LGTM |
Contributor
Author
|
/ok to test a87b25d |
akoumpa
reviewed
Jul 22, 2026
| still execute). No-op on torch builds without ``SAC_IGNORED_OPS`` or the | ||
| profiler op namespace. | ||
| """ | ||
| sac_ignored = getattr(torch.utils.checkpoint, "SAC_IGNORED_OPS", None) |
Contributor
There was a problem hiding this comment.
@HuiyingLi would it make sense to gate on pyt version? 🙇
Contributor
|
/ok to test 0d981f4 |
Only add profiler record-function ops to SAC_IGNORED_OPS on PyTorch 2.13+ builds, including NVIDIA alpha builds that report 2.13.0a0. This keeps older unaffected torch builds on the default SAC ignored-op set. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Contributor
Author
|
/ok to test b546d90 |
akoumpa
approved these changes
Jul 24, 2026
HuiyingLi
added a commit
that referenced
this pull request
Jul 26, 2026
…ting (#3133) * fix(dist): keep profiler record-function ops out of SAC replay accounting torch 2.13's FSDP2 runs its pre/post-forward hooks under torch.autograd.profiler.record_function, which emits dispatchable torch.ops.profiler ops. When an FSDP module boundary sits inside a selective-activation-checkpointed region (MoE experts sharded separately inside a checkpointed decoder block), the hooks fire a different number of times during the backward recompute than during the forward. SAC replays the forward op stream by per-op invocation index, so the extra profiler op shifts the stream and training fails with: RuntimeError: profiler._record_function_enter_new.default invocation index 3 encountered during backward but not found in storage. Seen on qwen3_6_35b_medpix_ep8cp2_4k with the pytorch:26.06-py3 base (torch 2.13.0a0+8145d630e8); upstream SAC_IGNORED_OPS does not cover profiler ops as of pytorch/pytorch@main. Range ops carry no tensors SAC could cache, so adding them to SAC_IGNORED_OPS only removes them from replay accounting - they still execute. Applied from both SAC context factories (activation_checkpointing and the MoE parallelizer). CPU regression test enters a record_function range only during the recompute pass and asserts backward succeeds. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(dist): gate profiler SAC ignore on torch 2.13 Only add profiler record-function ops to SAC_IGNORED_OPS on PyTorch 2.13+ builds, including NVIDIA alpha builds that report 2.13.0a0. This keeps older unaffected torch builds on the default SAC ignored-op set. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> --------- Signed-off-by: HuiyingLi <willwin.lee@gmail.com> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.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
2 participants
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.
What
Adds
torch.ops.profiler._record_function_{enter,enter_new,exit}(all overloads) totorch.utils.checkpoint.SAC_IGNORED_OPS, applied from both selective-AC context factories (activation_checkpointing.make_selective_checkpoint_context_fnand the MoEparallelizer.apply_acSAC path).Why
torch 2.13 (pytorch:26.06-py3 base,
2.13.0a0+8145d630e8) runs FSDP2 pre/post-forward hooks undertorch.autograd.profiler.record_function, which emits dispatchabletorch.ops.profilerops. With an FSDP module boundary inside a SAC region (MoE expertsfully_shard-ed separately inside a checkpointed decoder block), the hooks fire a different number of times during backward recompute than during forward. SAC replays the forward op stream by per-op invocation index, so the extra profiler op desyncs the replay:Hit by
qwen3_6_35b_medpix_ep8cp2_4k(both on main and on PR branches — pre-existing, torch-2.13-only) in nemo-ci runs on the 26.06 base: identical failure on merge-base 5008ea6 and on the #2937 branch (jobs 367239212 / 367239239). The same recipe passes on the pytorch 26.04 base (torch 2.12), and upstreamSAC_IGNORED_OPSstill doesn't cover profiler ops on pytorch main.Range ops carry no tensors SAC could cache or restore — ignoring them only removes them from replay accounting; they still execute.
Testing
record_functionrange only during the backward-time recompute; without the fix SAC raises the storage error, with it backward succeeds.profiler._record_function_enter_new.default ... not found in storage; with PR fix(dist): keep profiler record-function ops out of SAC replay accounting #3133 source overlaid, all ranks pass.qwen3_6_35b_medpix_ep8cp2_4kpassed with this fix (root pipeline 58616157, job 367298896; details in PR comment).nemotron_super_v3_te_deepeppassed with this fix overlaid at runtime:gitlab-master.nvidia.com/dl/joc/nemo-ci/main/automodel:pipe.58641817ba4ec358fecfec3bc3524a3e8a56677a2815878b; PR overlay ref:a87b25d36fb579ca93f25b9ec896572b4b95d1actorch.utils.checkpoint.SAC_IGNORED_OPS,SLURM_STATE='COMPLETED', andJob succeeded.🤖 Generated with Claude Code