[lhb8125]

Summary

• Enable paged stash for the non-op-fuser TEGroupedMLP GroupedTensor path.
• Mark dynamic saved tensors explicitly through the MCore Transformer Engine adapter.
• Allow HybridEP static-budget paged stash with moe_use_grouped_tensor=True without requiring the TE op fuser.
• Preserve the paged-stash marker across fused SwiGLU/QuickGeGLU view and FP8-cast operations, and restrict this path to those supported fused activations.
• Exclude selective-recompute-managed FC2 inputs from paged-stash ownership using the existing CheckpointWithoutOutput storage-level marker, which survives TE reshape/split views.

Review follow-up

Commit 16c115625 refreshes the branch onto current dev and addresses the strict-review comments:

• Replace the tensor-object recompute marker with is_checkpoint_without_output_tensor(), and add a reshape/split view regression test.
• Drive packed-bias operand marking directly from config.moe_paged_stash, so selective-recompute detach() does not silently disable it.
• Share the paged-stash scope setup and fused-activation marker propagation helpers.
• Route grouped-tensor marker detection through the Transformer Engine adapter and ignore None values when marking optional tensors.

Mirror and dependencies

This is the dev mirror of #6828. The review follow-up above has now been ported to #6828, so the two PRs are synchronized at the feature and follow-up level, modulo branch-specific main/dev differences. Because main did not yet contain the small storage-ownership feature from #7044, #6828 also directly subsumes that change and uses its storage-level predicate for the paged-stash ownership fix.

#6847, the dev port of #6000, has merged. This branch has been refreshed onto the updated dev branch and no longer carries a separate prerequisite delta.

Transformer Engine dependencies:

• Improve device-init grouped linear module with single grouped weight support  TransformerEngine#3224 for device-initiated GroupedLinear/GroupedTensor.
• Support paged stashing for GroupedLinear activations TransformerEngine#3423 for paged-stash marking of GroupedTensorStorage.

Validation

Current-head static validation:

• CHECK_ONLY=true BASE_REF=dev bash tools/autoformat.sh: Black, isort, pylint, and ruff passed. The script’s non-blocking mypy step reports the local environment’s missing torch/Transformer Engine stubs plus existing baseline typing errors.
• python tools/check_copyright.py <changed Python files> passed.
• python -m py_compile <changed Python files> passed.
• Local pytest was not available because this checkout’s environment does not contain torch; the new targeted tests are included for CI.

Existing runtime evidence:

• An earlier-head CI-faithful GB200 single-file run of tests/unit_tests/transformer/moe/test_paged_stashing.py exited 0. The CI image predates Transformer Engine ci(fix): Parse unit tests in merge-queue #3224, so the new grouped-tensor case was dependency-gated with Installed TE GroupedLinear does not expose use_grouped_tensor.
• The corresponding [Main] Support paged stash with device-initiated GroupedLinear #6828 implementation was exercised on OCI-AGA GB300 across 32 GPUs: paged-stash ON/OFF matched loss and grad norm for 3/3 steps, and profiling observed 144 paired stash copy/pop calls. That evidence covers the feature path but is not an exact-head rerun of the storage-level ownership follow-up in 16c115625.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaox12]

/claude strict-review

[claude]

Strict review summary

Findings: CRITICAL: 1, IMPORTANT: 1, SUGGESTION: 3

Scope reviewed: the paged-stash saved-tensor hook, the new _unfused_forward split in TEGroupedMLP, marker propagation through the fused SwiGLU/QuickGeGLU boundaries, the TE adapter shim, and the new __post_init__ validation. The forward/_unfused_forward split itself is a faithful move — the pre-compute padding, tokens_per_expert dtype/device normalization, and quantization_unpadding all retain their original ordering, and the paged-stash scope is placed to match _fused_forward exactly (context wraps compute; unpadding stays outside; group_commit after). The FP32 accumulation contract in _apply_packed_bias is unchanged, and the author's reasoning about which multiply operands autograd actually saves is correct.

Most impactful

1. CRITICAL — the selective-recompute ownership fix does not fire (paged_stash.py:930-934). _RECOMPUTE_MANAGED_TENSOR_ATTR is a per-tensor-object Python attribute, but TE GroupedLinear reshapes and splits its input before save_for_backward, and those views carry a fresh TensorImpl with no Python attributes. So on_save_for_backward sees an unmarked tensor and paged-stashes the activation anyway — the exact double-ownership between CheckpointWithoutOutput (resize_(0) + share_storage rebind) and the paged stash that the commit set out to eliminate. This repo already documents the hazard and already solves it storage-level: see tensor_parallel/random.py:64-66 ("the storage-level marker is visible through aliases/views") and random.py:1146-1149, which names "TE GroupedLinear's inp.reshape() + torch.split() saved for backward" as the motivating case. is_checkpoint_without_output_tensor() is already set unconditionally on every checkpoint output (random.py:1093-1094) and is already used for the identical question by fine_grained_activation_offload.py:993 — swapping to it removes the new attribute, the new marker function, its export, and the experts.py:945-946 call site. Note test_recompute_managed_tensor_bypasses_paged_stash_save_hook passes the marked object directly, so it cannot detect this; a reshape() in that test would.

2. IMPORTANT — the packed-bias marking is dead under the config its own test uses (experts.py:316-322). paged_stash_marked is inferred from permuted_probs, but with activation_recompute on, bias_act_func receives args routed through detach_variable (random.py:778) and therefore marks a detached copy; the object _apply_bias later sees is unmarked, as is the FC2 output. TestPagedStashingGroupedTensor::test_forward_backward_without_op_fuser sets recompute_modules=["moe_act"], so both new mark_grouped_tensor calls are unreachable there. Impact is a silently lost memory saving (and a stash footprint that varies with an unrelated recompute setting), not wrong numerics — but it leaves the added marking untested. self.config.moe_paged_stash already answers this unambiguously; prefer threading it in over sniffing a TE attribute on an operand.

Suggestions

• Extract the 19-line paged-stash scope block duplicated verbatim between _fused_forward:771-789 and forward:1028-1046; the two avg_num_tokens estimates must stay in agreement for buffer sizing.
• _propagate_paged_stash_marker is byte-identical in fused_bias_swiglu.py and fused_bias_geglu.py, and its hasattr(source, "grouped_tensor_scale_inv") probe hardcodes a TE-internal attribute name in four places while the write goes through the TE adapter. An is_grouped_tensor_marked() predicate in the same adapter would keep one place aware of TE's naming.
• _mark_paged_stash_tensors(*tensors: Optional[torch.Tensor]) advertises None but forwards straight into TE, which will not accept it; bias_act_func already guards permuted_probs is not None at line 879.

One additional observation (no inline comment — outside the diff)

on_save_for_backward derives hidden_size = tensor.numel() // self.max_num_tokens (paged_stash.py:950) with no divisibility assertion. That was low-risk when only TE-internal [max_num_tokens, H] buffers were marked, but this PR extends marking to MCore-side tensors (flat_probs at hidden_size=1, bias_per_token, and fused-activation views produced by input.view(-1, ori_shape[-1]), which is shared with the 3-D dense-MLP call sites in mlp.py). A wrong hidden_size there is not an error — it silently mis-sizes the stash extents. Worth an assert tensor.numel() % self.max_num_tokens == 0 at that parallelism boundary.

Risk assessment

Medium-high. Backward compatibility is clean: the new validation only tightens a combination that previously produced a silent no-op stash, no defaults or checkpoint keys move, and the op-fuser path is untouched. Performance is fine — the added work is hasattr probes and marker sets, no new syncs, and _apply_packed_bias keeps its output_size= sync avoidance. The risk is concentrated in finding 1: if the recompute-managed guard is inert as analyzed, --moe-paged-stash combined with --recompute-granularity selective --recompute-modules moe_act on this new path can produce silently wrong expert gradients, and the PR's own validation (paged-stash-on capture iteration vs. paged-stash-on steady iteration) would not distinguish that from correct. Please either move the marker to storage level or add a view-level regression test plus a paged-stash-off reference comparison in test_forward_backward_without_op_fuser before this lands.

Also worth noting for the record: per the PR description the end-to-end grouped-tensor case is dependency-gated in CI (the image predates TE #3224), and the 32-GPU ON/OFF parity run was performed on the #6828 stack rather than this head — so at merge time neither the new path nor the ownership fix has CI coverage on this branch.

[lhb8125]

/ok to test 16c1156