Repository navigation
Reduce DSv4 output RoPE activation memory - #8399
Conversation
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test b44423f |
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
/ok to test d6d5cd7 |
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Fixed the determinism coverage gate in |
|
/ok to test b133adb |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
Reviewed all 11 changed files and affected SBHD/THD/CP paths. Found two correctness regressions and one review-enforced style violation. Static inspection only; tests were not executed.
Preserve importability without Triton
The new unconditional Triton imports make importing CSA fail with ModuleNotFoundError when Triton is unavailable—even when both output-RoPE fusion and the fused attention backend are disabled. Previously, the MLA kernel module explicitly guarded Triton availability. Native attention and CPU/mock contract tests should remain importable without an unused optional kernel dependency.
Correction: Reuse that module's capability-guarded bindings. Replace the two direct Triton imports with the following import, and place the binding after the existing MLA symbol imports:
from megatron.core.fusions import fused_mla_yarn_rope_apply as _mla_rope
triton, tl = _mla_rope.triton, _mla_rope.tlKeep the existing explicit kernel-symbol imports. This preserves real Triton dispatch when available and the existing import-only fallback when unavailable; it must not silently execute placeholder numerical kernels. Please verify the suggested code before adoption.
Acceptance — proposed, not executed: In a fresh interpreter with an import hook rejecting triton and triton.*, importing the new helper and CSA must succeed without launching kernels. Retain native CPU/mock contract tests with fusion disabled. With Triton/CUDA available, retain the reconstruction tests: exact forward/reference-gradient comparisons, NoPE-prefix preservation, unchanged saved output and expanded upstream gradient, and bit-exact replay.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex. This was a full review. Request a full review with /review reset=true.
| """ | ||
| if out_rope is None: | ||
| return out_flat, dO_flat | ||
| o_unrot = torch.empty_like(out_flat) |
There was a problem hiding this comment.
Handle padded-head outputs before reconstruction
Failure: Supported head counts requiring FlashMLA padding now fail during backward with output RoPE enabled. The shared adapter slices the padded output without making it dense. For example, a contiguous (4, 64, 512) backend output sliced to two heads has strides (32768, 512, 1). This empty_like allocation compacts that non-dense view to (1024, 512, 1), so reconstruction raises the matching-metadata error. SBHD, packed, and fused-indexer Functions all share this path; backward should instead reconstruct O and dO successfully.
Correction: The smallest shared correction is to replace the output-slicing line in _csa_fwd_flash_mla with:
out = out[:, :actual_num_heads, :].contiguous()This makes the saved real-head output dense before rotation, restoring the source/destination stride invariant. Do not merely remove the metadata check: the reconstruction kernel uses the input strides for both pointers. Please verify this suggested replacement before adoption.
Preservation: Keep real-head output values, LSE/indexer-LSE slicing, and the unpadded full-head path unchanged. Continue reconstructing into private storage so saved-output aliases and shared/expanded upstream gradients remain untouched; no-RoPE backward must remain unchanged.
Acceptance — proposed, not executed: Extend the existing small-head adapter test with assert out.is_contiguous(), retaining its exact output/LSE/indexer-LSE value assertions. Add SBHD and THD backward cases with two actual heads, a mocked contiguous 64-head backend output, and real OutputRopeParams; require no exception, reconstructed O within the existing dtype tolerance, reference-matching dO, and unchanged saved output/upstream storage. Retain an unpadded 64-head case and no-RoPE backward.
There was a problem hiding this comment.
Confirmed and fixed in edd16e2: only the sliced real-head output is made contiguous. Output/LSE/indexer-LSE values are preserved, and reconstruction still uses private storage. The 24 CPU layout/alias cases and both adapter controls pass, including unpadded 64 heads and no RoPE. Extended the existing GPU reconstruction tests to exercise the adapter with 2/64 heads; GPU validation is pending CI.
| batch_size = None | ||
| seq_num = cu_seqlens_q.numel() - 1 | ||
| if x.stride(-1) != 1 or x.shape[-1] != nope_dim + emb_dim: | ||
| raise ValueError('Output RoPE requires contiguous head dimensions.') |
There was a problem hiding this comment.
Use double quotes for the new literals
These new literals violate the repository's explicit double-quote convention, which Black deliberately does not enforce.
Replace this exception literal with "Output RoPE requires contiguous head dimensions.", the dictionary lookup at line 145 with meta["BLOCK_H"], and the exception literal at line 205 with "Output reconstruction requires matching tensor metadata.".
Preserve the exact exception messages and dictionary key; numerical and validation behavior must not change. Acceptance is a static check that these three new literals use double quotes and their contents remain identical.
There was a problem hiding this comment.
Fixed in edd16e2. All three literals now use double quotes; the messages and dictionary key are unchanged (AST identical).
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Can we keep Triton optional in output_rope.py? The imports on lines 14–15 are unconditional, and csa.py imports this helper even when native attention is selected and RoPE fusion is disabled. Importing CSA now fails with ModuleNotFoundError on installations without Triton. The existing MLA helper already guards this dependency. Please preserve that behavior and add an import test with Triton unavailable. |
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Fixed in 3dee74b. The helper now shares the existing MLA optional-Triton bindings and raises a clear error only when fused output RoPE is requested. Isolated subprocess tests block Triton and reload the real CSA/helper modules; native fusion-disabled RoPE forward/backward passes. Both import controls and the 24 CPU layout/alias checks pass; existing GPU test logic and tolerances are unchanged. |
Signed-off-by: Ajay Balasa <abalasa@nvidia.com>
|
/ok to test 894ce2b |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40971092855 |
Resolves the import conflict in deepseek_v4_hybrid_attention.py: main adds OutputRopeParams from megatron.core.ops.attention.csa.output_rope (NVIDIA#8399) next to the LayerNormBuilder import this branch moved to megatron.core.ops.norm.torch_norm; both are kept. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: guihong-nv <guihongl@nvidia.com>
Port the output-activation memory optimization from dev #7036 onto main. Fused CSA stores its inverse-rotated output and reconstructs the raw attention output during backward, avoiding a second output activation retained from forward. Related to #5795; credits to @ghadiaravi13 for the original dev implementation.
megatron/core/ops/attention/csaand copies the non-rotary prefix in the same launch.Validation on four GB200 GPUs: 96 tests passed per rank, including output/gradient reconstruction and bit-exact replay, shared/expanded-gradient protection, and the existing native SBHD → fused SBHD → THD → CP2/CP4 comparisons (sparse/dense indexer loss and selective recompute). The two pre-existing, flaky-marked small-head SBHD tests also fail on unmodified main with comparable BF16 KV-gradient differences; no tolerances or skip markers were changed. Formatting/lint, Python syntax, and kernel determinism coverage checks passed.
Reconstructing a saved BF16 output introduces an additional rounding step; this is short numerical/integration validation, not a long-run convergence claim.