docs/bench: architecture gaps cleanup + HBM OOM finding + cold bench numbers - #47
Merged
Merged
Conversation
…h.sh architecture.md Known Gaps was stale since v0.4.0: - syrk and trsm (left-side) now have NKI kernels; only symm and trmm are PyTorch-only, and neither is in the DF-MP2 hot path - nki_mp2_energy dispatch overhead bullet superseded by batched-pair energy (#43/#46): warm 3.6×/5.2× faster than torch at small/medium shape - FP64/double-double bullet updated: decision made 2026-04-18, #10 closed "not needed", #22 deferred indefinitely scripts/run_bench.sh: runs df_mp2.py --bench via SSM on trnblas-ci-trn1. Supports --shape large/medium (default: both). Follows base64-SSM pattern from run_pyscf_tests.sh; polls up to 120 min for cold NEFF compile.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
…ll for large Three findings from the 2026-04-21 medium bench run: 1. EBS-warm cold (all NEFFs loaded from disk, not compiled): 137.2 s total - Half-transform 5.1 s (was 103.5 s compile-cold; 20× faster from EBS) - Energy 101.0 s (64 energy NEFFs still load serially at ~1.3 s/NEFF ≈ 83 s) - Adds a distinct "EBS-warm" column to the end-to-end table 2. Corrects CHANGELOG: prior 4.784 s warm figure is in-process HBM-warm, not separate-process EBS-warm (which gives ~137 s as now confirmed). HBM-warm cannot be reproduced (OOM after cold pass at medium shape). 3. Large-shape cold failed: LLVM ERROR IO failure (No space left on device) during neuronxcc compilation. Documented; needs disk investigation. Fix: add NEURON_RT_LOG_LEVEL=WARNING to run_bench.sh to suppress per-NEFF INFO messages that ate the SSM 24 KB stdout budget and truncated warm output. Also add df -h /var/tmp before each shape for disk diagnostics.
Root cause of large-shape disk-full: the 100G EBS fills to 99% (95G
used) from Neuron SDK packages + medium NEFF cache alone, leaving only
1.1G for large-shape compilation artifacts.
Changes:
- infra/terraform/main.tf: root_block_device volume_size 100 → 200
- scripts/run_bench.sh:
- Redirect bench output to /tmp log file; grep timing line to stdout.
NEURON_RT_LOG_LEVEL=WARNING doesn't suppress the [INFO] NEFF messages
(they come from libnrt.so); file redirect + grep is the only reliable
approach to staying under SSM's 24 KB stdout budget.
- Add growpart/resize2fs at startup to expand filesystem after terraform
resizes the EBS volume (idempotent: no-op if already full size).
- df -h / before each pass for disk diagnostics.
Next step: run `terraform apply` in infra/terraform/ to resize the EBS,
then re-run `AWS_PROFILE=aws ./scripts/run_bench.sh`.
4 tasks
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
1 participant
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
docs/architecture.mdKnown Gaps — corrected stale text that predated v0.5.4: syrk and trsm now have NKI kernels (only symm/trmm are PyTorch-only); batched-pair energy solved the dispatch overhead (3.6×/5.2× vs torch); FP32 decision recorded inline (Double-double FP64 emulation for Tensor Engine #10 closed, Phase 2 — double-double FP64 GEMM for chemistry workloads #22 deferred)HBM OOM finding confirmed (2026-04-21) —
benchmarks.mdincorrectly stated the medium-shape warm pass "logs a warning but computation succeeds." It fails. After the cold pass, 64 energy NEFFs remain resident in HBM (12.6 GB DMA spill + 900 MB code = 15.9 GB), leaving no room for the 1.5 GB tensor allocation needed in the warm pass.Medium cold timing recorded — fresh
run_bench.shrun: chol 29.7 s, half 103.5 s, metric 4.0 s, energy 101.3 s = 238.5 s total (partially-warm EBS cache; energy kernel compiled fresh)df_mp2.py --passes cold|warm|both— new flag separates cold and warm into distinct process invocations.--passes both(default) preserves existing behaviour for small shapes; medium/large must use separate invocations to avoid HBM OOM.scripts/run_bench.sh— new script runsdf_mp2.py --bench --batched-pair-energyontrnblas-ci-trn1via SSM, automatically doing cold and warm as separate processes. Supports--shape medium|large(default: both). Follows base64-SSM pattern fromrun_pyscf_tests.sh.Test plan
examples/df_mp2.py --bench --shape small --batched-pair-energystill runs both passes (backward compat)examples/df_mp2.py --bench --shape small --batched-pair-energy --passes coldruns only coldexamples/df_mp2.py --bench --shape small --batched-pair-energy --passes warmruns only warmAWS_PROFILE=aws ./scripts/run_bench.sh --shape mediumon trn1 completes without HBM OOM