rocm/gfx1151: grad_clip_norm lane + flash_attn-backward perf ratchet baseline - #355
Merged
Merged
Conversation
…baseline The two follow-ups from PR #354's inventory §9. grad_clip_norm (§5.5) — global gradient-norm clipping g*min(1, max_norm/||g||) as a composed lane on BOTH x86 and rocm: the L2 norm's global sum-of-squares runs on the device reduce kernel (the FLOP-heavy O(n) part, reusing _device_reduce_sum_rows), host does sqrt + the clip scale; norm_type=inf uses max|g|. Single-tensor form (the atomic unit of the pytree clip_grad_norm). New rocm_grad_clip_compiled / x86_grad_clip_compiled lanes + execution-matrix rows + manifest entries + fixtures. Matches optim.clip_grad_norm within f32 tolerance (7/7 gfx1151, 4/4 x86). Perf ladder — extend the gfx1151 hot-path ratchet baseline to cover the flash_attn BACKWARD lane (rocm_flash_attn_bwd_compiled). Measured on the box: ~44 ms / 87 ms for 1x8x512x64 / 1x16x1024x128, ~15x the forward — honestly a correctness-first regression floor, NOT an MFU claim (Decision #26). The per-row ratchet (test_rocm_perf_ratchet) now re-times the backward rows too. Verified: 163-test regression sweep green; drift gate in sync; ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4eb62cc410
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…h_attn_bwd Two PR #355 review fixes (both P2): - grad_clip_norm: only L2 (norm_type==2) and inf are implemented — the norms optim.clip_grad_norm actually distinguishes (it computes tree_l2_norm for every finite norm_type). Reject other finite p-norms with a clear diagnostic instead of silently clipping by the L2 norm the caller did not ask for (Decision #21). + reject tests. - perf ratchet: the flash_attn_bwd baseline rows were recorded but never re-timed live (the live checks only covered {"wmma"}/{"flash_attn"}), so a backward regression could not fail CI. Add test_live_flash_attn_bwd_within_ ratchet (same flash lane gate) that re-times the backward rows against the committed baseline. Verified on gfx1151: 25 grad_clip+ratchet tests + the slow live bwd ratchet green; ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 (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
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.
The two follow-ups called out in PR #354's inventory §9.
grad_clip_norm (§5.5)
Global gradient-norm clipping
g * min(1, max_norm/||g||)as a composed lane on both x86 and rocm: the L2 norm's global sum-of-squares runs on the device reduce kernel (the FLOP-heavy O(n) part, reusing_device_reduce_sum_rows), host doessqrt+ the clip scale;norm_type=infusesmax|g|. Single-tensor form — the atomic unit of the pytreeclip_grad_norm.rocm_grad_clip_compiled/x86_grad_clip_compiledlanes + execution-matrix rows + manifest entries + fixtures.optim.clip_grad_normwithin f32 tolerance (the reference accumulates in f64).Perf ladder — flash_attn backward ratchet
Extends the gfx1151 hot-path ratchet baseline (
rocm_gfx1151_hot_paths.json) to cover therocm_flash_attn_bwd_compiledlane. Measured on the box: ~44 ms / 87 ms for1x8x512x64/1x16x1024x128— ≈15× the forward, honestly a correctness-first regression floor, not an MFU claim (Decision #26). The per-row ratchet (test_rocm_perf_ratchet) now re-times the backward rows too.This is the honest reading of "perf ladders / MFU sign-off": the compiled lanes now carry measured baselines, but stay correctness-first — a real MFU ladder (register/LDS blocking) is where ROCm's lead-performance-target status (Decision #28) gets earned, and is future work.
Verification
🤖 Generated with Claude Code