fix(mp): FLOP-weighted calibration budget, per-row qk MP, drop alpha/beta - #6
Open
heroarmor wants to merge 6 commits into
Open
Conversation
…beta - calibrate_mp_thresholds.py: budget is now sum(rows * macs * stoc_len) instead of sum(rows * stoc_len), so an "avgX" target means X in real compute, not row count. macs_per_row (out*in for linears, N*head_dim for qk/av) is threaded through add() -> the global lambda solve + iso-budget check. qk recording is switched to per-row. Validated: realized FLOP-weighted avg = 47.97 for target 48. - sc_attention.py: _sc_qk MP is now per-query-row (||Q_row||_inf), matching the per_row quant scale; removed the per-head classifier and static per-head path. - quant_sc_main.py: removed all alpha/beta remnants -- --mp_alpha/--mp_beta and the 14 per-op variants, the crashing AdaptiveMPConfig(alpha=,beta=, operator_params=) construction, and the run-name tag. --adaptive_mp now requires --adaptive_mp_table. - sc_attention.py + sc_mlp.py: repaired the adaptive_classify_rows signature drift vs kernel 9358239 (6 call sites passed the old arg order -> TypeError, so the adaptive path could not run at all). Invalidates existing calibration threshold tables (new cost model -> new allocations); re-calibration required. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Run artifacts (slurm-*.out, figures/) are regenerated, not source. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
tools/: compare_old_vs_new{,_isolated,_thorough}.py, compare_vs_cpu_reference.py, sc_call_counter{,_full}.py, smoke_test_e2e.py
scripts/: global-PR calib + MP sbatch (calib_globalpr{,_flop}{,_halve}, mp_globalpr_{auto,fill,flop_halve_auto}, uniform{,_halve}_auto, opfreeze_probe) + eval/ CPU eval
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eanup Submodule now at local branch local/trace+scramble-cache (6c4539a): PR#25 trace subsystem + cherry-picked PR#28 enable-table scramble cache-key fix + mp/config doc cleanup. NOTE: 6c4539a is LOCAL-ONLY (not yet pushed to the heroarmor fork); a fresh 'git submodule update' cannot resolve it until pushed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ion; mse_rel metric The halve MP line calibrated and ran levels above the halve-mode maximum (2^(sc_prec-1)=128): the 256/192 levels are unrealizable on halve hardware and silently inflated attention accuracy (up to 99% of qk rows ran 256/192- cycle streams). Fixes, one per layer: - scmp_kernels bump -> c967d3a: sc_matmul now raises on explicit stoc_len/rng_levels above 2^(sc_prec-1) under halve+bipolar (upstream kernels PR #30), so stale 256-max tables crash instead of overspending. - calibrate_mp_thresholds.py: hard-validate mp_levels and budget_ref_stoc_len <= 2^(sc_prec-1) under --sc_halve; add --metric mse_rel (squared relative L2 = cosine's quadratic curvature + gain-error sensitivity; a pure 10% gain error scores 0.0 under cosine, 0.01 under mse_rel). - sbatch_calib_globalpr_flop_halve.sb / sbatch_calib_globalpr_halve.sb: 128-capped level grid, budget_ref 128, BUDGET_RATIO computed from AVG in-script (avg/128 — passing the old avg/256 ratios silently halved every budget), optional METRIC arg with suffixed OUTDIR, and calibration now runs at the deployment CFG (teacher_cfg_scale 1.5, batch 16 so the CFG-doubled batch keeps the old 32-row memory footprint) instead of cfg=0. - sbatch_mp_globalpr_flop_halve_auto.sb: [METRIC] and [NUM_FID] args (metric-suffixed calib/output dirs; NUM_FID for fast probes), points at the 128-capped tables; old 256-grid tables/arms preserved in the un-suffixed dirs for comparison. Validated: avg48 recalibration hits expected_avg_stoc_len=48.00; runtime realized per-op distribution matches the table within 1-3pp per level; on the noise-aligned index prefix the fixed adaptive arms (cosine and mse_rel) beat uniform48 on fidelity-to-FP (2.5% vs 5% trajectory divergence). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Small single-purpose probes developed while root-causing the halve-line
regressions; kept because each verifies an invariant worth re-checking:
- eval/sbatch_eval_uniform_halve.sb: KID + FID/IS/Prec/Recall for one
uniform-halve arm (companion to sbatch_eval_mp_flop_halve.sb).
- eval/sbatch_eval_ab_kid100.sb: matched-index KID across probe arms + an FP
anchor on the same indices (the anchor exposes the class-composition floor
that dominates absolute KID at small n — FP scores the same ~30x1e-3 as SC
arms on 10 classes).
- eval/sbatch_probe_realized_dist{,_old}.sb: validate-path run that dumps
debug_mp_distribution.csv — compares the RUNTIME-realized per-op level
fractions against the calibration table's intent (A4 check).
- eval/sbatch_kernel_level_err.sb: kernel-level rel-err vs stoc_len under
halve on qk/av/mlp shapes (catches level-grid pathologies without running
the diffusion model).
- sbatch_uniform48_rerun100.sb: regenerate uniform48 idx 0-99 into a fresh
dir — reproducibility check. NOTE: latent noise pairs by position in each
worker's index list, not by global index, so PSNR comparisons against runs
with a different GPU split are only valid on the aligned worker-0 prefix.
- sbatch_only1op_sc32.sb: single-op SC isolation (everything FP except one op
at stoc_len 32 + halve) via per-op config JSONs; the legacy layerwise flags
are ignored whenever --sc_config supplies a precision map.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Pushed two follow-up commits on top of e061384/ca5eab0: 1b74fab — fix(halve): enforce max stoc_len 128 end-to-end Validated: avg48 recalibration hits expected_avg_stoc_len=48.00 exactly; runtime-realized per-op level distributions match the table within 1–3pp; on the noise-aligned index prefix the fixed adaptive arms (cosine and mse_rel) beat uniform48 on fidelity-to-FP (2.5% vs 5% trajectory divergence). c328d36 — tools(probe): verification sbatch scripts 🤖 Generated with Claude Code |
There was a problem hiding this comment.
Pull request overview
This PR updates the mixed-precision (MP) calibration and runtime integration for SC-based diffusion sampling: it switches calibration to a FLOP-weighted budget, implements per-query-row QK mixed precision (aligned with per-row quant scaling), and removes the deprecated adaptive MP alpha/beta closed-form threshold path. It also adds a suite of verification/probing tools and Slurm scripts, plus a small .gitignore update for run artifacts.
Changes:
- Reworked SC attention/MLP integration to use per-row QK MP (and updated adaptive MP row-classification call sites).
- Updated MP calibration to enforce halve-mode constraints, add an additional error metric option, and use FLOP-weighted budgeting with per-op MAC accounting.
- Simplified adaptive MP runtime configuration to require calibrated threshold tables; added new tooling + sbatch/eval scripts and ignored additional run artifacts.
Reviewed changes
Copilot reviewed 31 out of 32 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/smoke_test_e2e.py | Adds an end-to-end GPU smoke test for kernel/API/import sanity and basic correctness checks. |
| tools/sc_call_counter.py | Adds a wrapper to count sc_matmul calls during quant_sc_main.py sampling. |
| tools/sc_call_counter_full.py | Adds a “full SC” variant call counter run (all ops enabled, uniform stoc_len). |
| tools/compare_vs_cpu_reference.py | Adds GPU-vs-CPU reference comparison against scmp_llm CPU implementation. |
| tools/compare_old_vs_new.py | Adds numerical comparison between legacy scmp_llm Triton path and new scmp_kernels dispatcher. |
| tools/compare_old_vs_new_thorough.py | Adds an exhaustive bit-identical verification sweep across modes/shapes/seeds/levels. |
| tools/compare_old_vs_new_isolated.py | Adds an isolator script to pinpoint the source of bipolar per-tensor diffs. |
| scripts/sbatch_uniform48_rerun100.sb | Adds a one-off reproducibility rerun script for uniform48 halve-mode. |
| scripts/sbatch_uniform_halve_auto.sb | Adds timeout-safe chained uniform (halve) generation script. |
| scripts/sbatch_uniform_auto.sb | Adds timeout-safe chained uniform generation script (non-halve). |
| scripts/sbatch_opfreeze_probe.sb | Adds an operator-freeze sensitivity probe job script for adaptive MP. |
| scripts/sbatch_only1op_sc32.sb | Adds a single-op SC isolation probe script (uniform stoc_len=32 + halve). |
| scripts/sbatch_mp_globalpr_flop_halve_auto.sb | Adds timeout-safe chained FLOP+halve adaptive MP generation using capped level grid. |
| scripts/sbatch_mp_globalpr_fill.sb | Adds an auto-resume “fill” variant for global+per-row adaptive MP generation. |
| scripts/sbatch_mp_globalpr_auto.sb | Adds an auto-resume script for global+per-row adaptive MP generation. |
| scripts/sbatch_calib_globalpr.sb | Adds a per-row global calibration sbatch wrapper. |
| scripts/sbatch_calib_globalpr_halve.sb | Adds a halve-mode per-row global calibration sbatch wrapper with ref/level checks. |
| scripts/sbatch_calib_globalpr_flop.sb | Adds a FLOP-weighted per-row global calibration sbatch wrapper. |
| scripts/sbatch_calib_globalpr_flop_halve.sb | Adds a FLOP+halve calibration sbatch wrapper with metric selection and ref checks. |
| scripts/quant_sc_main.py | Removes alpha/beta adaptive MP fallback; requires calibrated threshold tables and updates output tagging. |
| scripts/eval/sbatch_probe_realized_dist.sb | Adds an instrumented probe to compare realized MP distributions old vs new tables. |
| scripts/eval/sbatch_probe_realized_dist_old.sb | Adds a control probe for realized MP distribution using old tables without halve. |
| scripts/eval/sbatch_kernel_level_err.sb | Adds a kernel-level error-vs-level probe under halve and non-halve modes. |
| scripts/eval/sbatch_eval_uniform_halve.sb | Adds a GPU eval script for uniform(halve) configs (KID + FID/IS/etc). |
| scripts/eval/sbatch_eval_mp_flop_halve.sb | Adds a GPU eval script for FLOP+halve adaptive MP configs (KID + FID/IS/etc). |
| scripts/eval/sbatch_eval_ab_kid100.sb | Adds a one-off A/B probe to compute KID on idx 0–99 across multiple arms. |
| scripts/eval/eval_mp_flop_halve_cpu.sh | Adds a CPU-only eval wrapper for FLOP+halve MP results packaging + evaluation. |
| scripts/calibrate_mp_thresholds.py | Implements FLOP-weighted budget accounting, per-row QK calibration recording, halve-mode enforcement, and an additional metric option. |
| qdit/sc_integration/sc_mlp.py | Updates adaptive MP row classification call sites to the newer argument style. |
| qdit/sc_integration/sc_attention.py | Reworks QK to per-query-row MP, aligning MP granularity with per-row quant scaling. |
| .gitignore | Ignores Slurm output logs and generated figures/ artifacts. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+53
to
+55
| if not torch.cuda.is_available(): | ||
| print(" no CUDA — skipping kernel execution"); sys.exit(2) | ||
| print(f" device: {torch.cuda.get_device_name(0)}", flush=True) |
Comment on lines
+10
to
+11
| ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') | ||
| sys.path.insert(0, str(ROOT)) |
Comment on lines
+6
to
+7
| ROOT = Path('/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_diffusion') | ||
| sys.path.insert(0, str(ROOT)) |
Comment on lines
+20
to
+24
| import sys, time | ||
| from pathlib import Path | ||
|
|
||
| SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") | ||
| sys.path.insert(0, str(SCMP_LLM_SC)) |
| from pathlib import Path | ||
|
|
||
| # Old impl — bare imports relative to SC/ folder | ||
| SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") |
Comment on lines
+8
to
+12
| import sys | ||
| from pathlib import Path | ||
|
|
||
| SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") | ||
| sys.path.insert(0, str(SCMP_LLM_SC)) |
Comment on lines
+13
to
+17
| import sys | ||
| from pathlib import Path | ||
|
|
||
| # Old impl | ||
| SCMP_LLM_SC = Path("/gpfs/accounts/nbleier_owned_root/nbleier_owned1/zhkangqi/scmp_llm/SC") |
| #SBATCH --mem-per-gpu=12G | ||
| #SBATCH --time=1-12:00:00 | ||
| # | ||
| # global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: |
| #SBATCH --mem-per-gpu=12G | ||
| #SBATCH --time=1-12:00:00 | ||
| # | ||
| # global+per-row generation that AUTO-RESUMES to NUM_FID=10000, TIMEOUT-safe: |
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.
Summary
FLOP-weighted MP calibration and per-row qk mixed precision, plus supporting tooling and the kernels submodule bump.
Commits:
e061384fix(mp): FLOP-weighted calibration budget, per-row qk MP, drop alpha/beta — the core change. Reworksqdit/sc_integration/sc_attention.py(per-row qk MP),sc_mlp.py,scripts/calibrate_mp_thresholds.py(FLOP-weighted budget), and simplifiesscripts/quant_sc_main.py; drops the alpha/beta closed-form threshold path. (+156 / -209 across 4 files.)0b66c15tools+scripts — SC correctness/verification tools (tools/compare_old_vs_new{,_isolated,_thorough}.py,compare_vs_cpu_reference.py,sc_call_counter{,_full}.py,smoke_test_e2e.py) and MP/calibration sbatch scripts (scripts/sbatch_{calib_globalpr*,mp_globalpr_*,uniform*,opfreeze_probe},scripts/eval/).7a32695chore — gitignoreslurm-*.outandfigures/run artifacts.ca5eab0kernels bump — pin thescmp_kernelssubmodule to6c4539a.Note on the submodule pin
ca5eab0pinsscmp_kernelsat6c4539a, which lives on theheroarmor/scmp_kernelsfork branchlocal/trace+scramble-cache. It bundles kernels changes still in review upstream (trace = #25 / #26, enable-table scramble cache-key = #28). This matches the existing vendored-kernels workflow — the submodule previously pinned9358239, also not onscmp_kernels:main.Testing
e061384exercised through the calibration + E2E FID sweep scripts.clear_rng_cache()).🤖 Generated with Claude Code