feat(benchmarks): GPU benchmark for Beta-VAE train step + encoder inference#253
Conversation
|
Claude finished @zachtheyek's task in 3m 43s —— View job Claude Code review — PR #253Nice self-contained addition. The training step faithfully mirrors The comments I left are minor — nothing that changes correctness of the reported numbers. Inline comments posted
Things I checked and liked
Not blocking, just noting
--- · Branch |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| "--find-max", | ||
| action="store_true", | ||
| help="Double the batch size until the GPU OOMs (the default when --batch-sizes is unset).", | ||
| ) |
There was a problem hiding this comment.
--find-max is decorative — it's declared but never consulted.
The sweep is actually gated by whether --batch-sizes is set (see _batch_schedule, benchmarks/bench_gpu.py:140-149). args.find_max is never read, so --find-max alone does nothing that omitting --batch-sizes doesn't already do, and --find-max --batch-sizes 128,256 silently ignores the flag rather than sweeping past 256.
Two consistent options:
- Drop the flag and let
--batch-sizesbe the sole switch (README already says the sweep is the default), OR - Actually use it:
if args.batch_sizes and not args.find_max:(or make the two mutually exclusive viaadd_mutually_exclusive_group) so--find-maxis meaningful and mis-combinations surface as an error rather than a silent no-op.
| f"bench_gpu_{args.mode}", | ||
| { | ||
| "mode": args.mode, | ||
| "obs_shape": list(_OBS_SHAPE), |
There was a problem hiding this comment.
obs_shape in the JSON params is misleading in --mode train.
You always write _OBS_SHAPE (the encode-mode shape) even when the driving input is a cadence (6, 16, 512). A downstream reader looking at the JSON for a train run will see obs_shape=[16, 512, 1] and misinterpret what was actually measured.
Suggest logging the shape you actually drove — either
"example_shape": list(_CADENCE_SHAPE if args.mode == "train" else _OBS_SHAPE),or record both (cadence_shape + obs_shape) unconditionally. Bonus: it'd match the example_shape= label already printed on stdout at line 189.
| try: | ||
| for _ in range(warmup): | ||
| out = step_fn(*inputs) | ||
| out.numpy() # sync after warmup (also absorbs tf.function tracing) before timing | ||
| _reset_vram_peak() | ||
| start = time.perf_counter() | ||
| for _ in range(steps): | ||
| out = step_fn(*inputs) | ||
| out.numpy() # force the queued steps to finish before stopping the clock | ||
| elapsed = time.perf_counter() - start |
There was a problem hiding this comment.
out.numpy() explodes on --warmup 0.
If a user runs with --warmup 0, the first out.numpy() at line 115 raises NameError: name 'out' is not defined — the timing loop never executes. Not a common invocation, but the arg accepts 0 without complaint. Cheapest guard is either
parser.add_argument("--warmup", type=int, default=3, ...) # already there
# then:
if args.warmup < 1:
parser.error("--warmup must be >= 1")or initialise out = None and skip the sync when it's None.
| if mode == "train": | ||
| # main, true, false, target — each a batch of cadences (B, 6, 16, 512). | ||
| inputs = tuple(_rand((batch, *_CADENCE_SHAPE)) for _ in range(4)) | ||
| else: | ||
| # encode: a batch of single observations (B, 16, 512, 1). | ||
| inputs = (_rand((batch, *_OBS_SHAPE)),) |
There was a problem hiding this comment.
_rand re-seeds default_rng(0) on every call, so all four training inputs are byte-identical tensors (main == true == false == target). Fine for GPU timing (compute cost doesn't depend on values), but it's worth an explicit one-liner in the docstring — a future reader inspecting the JSON or debugging weird loss values will assume they're independent draws. Alternatively, hoist the rng to a module-level singleton so each call advances state.
|
|
||
|
|
||
| def _peak_vram_gb() -> float: | ||
| return tf.config.experimental.get_memory_info(_DEVICE)["peak"] / 1e9 |
There was a problem hiding this comment.
Unit nit: /1e9 is SI GB (10⁹ bytes), not GiB (2³⁰). Nvidia tooling (nvidia-smi, driver docs) mixes the two — a 96 GB Blackwell is ~89.4 GiB — so a reader cross-checking the README's "76.52 GB peak" at batch 4096 against nvidia-smi's "MiB" may not know which convention you're using. Either add a comment noting SI GB here, or divide by 2**30 and label as GiB in the tables. Not a bug; just prevents confusion the first time somebody tries to reconcile the numbers.
| def _batch_schedule(args) -> list[int]: | ||
| if args.batch_sizes: | ||
| return [int(x) for x in args.batch_sizes.split(",") if x.strip()] | ||
| # --find-max (default): powers of two from --start-batch up to the cap; the sweep stops at | ||
| # the first OOM, so the last success is the largest power of two that fits. | ||
| sizes, b = [], args.start_batch | ||
| while b <= args.max_batch: | ||
| sizes.append(b) | ||
| b *= 2 | ||
| return sizes |
There was a problem hiding this comment.
Minor input hardening: --batch-sizes " 128, foo" raises ValueError: invalid literal for int() from int(x) rather than a friendly argparse error, and --start-batch 0 / --max-batch 0 produces an infinite loop (since b=0, b*=2 stays 0 and is <= any positive cap). Neither is a real-world path, but a cheap fix is parser.error(...) in main() when the parsed values fall outside the sane range.
|
This branch evolved substantially since the first pass, and is now validated on-cluster (bla0, 6× A4000). Current state (
On-cluster validation (bla0, all paths clean): train single 795 / 6-GPU 3,087 (~3.9×) / accumulation K=4 4,272 cad/s @ 9.91 GB; encode single 53k / 6-GPU 214k obs/s (~4.1×). The multi-GPU + accumulation paths (the ones that couldn't be tested locally — no GPU, TF deadlocks on the dev box) are what the smoke exercised and fixed. Fresh eyes welcome on the accumulator pattern in |
| def _apply_and_reset(): | ||
| grads = [acc / accumulation_steps for acc in accumulators] | ||
| clipped, _ = tf.clip_by_global_norm(grads, _CLIP_NORM) | ||
| model.optimizer.apply_gradients(zip(clipped, model.trainable_variables, strict=False)) | ||
| for acc in accumulators: | ||
| acc.assign(tf.zeros_like(acc)) | ||
|
|
||
| @tf.function | ||
| def run(dist_inputs): | ||
| # Accumulate grads over accumulation_steps micro-batches into the per-replica | ||
| # accumulators, then apply once. The apply runs INSIDE strategy.run, so MirroredStrategy | ||
| # all-reduces on apply — the correct multi-GPU accumulate-then-apply cadence, and it | ||
| # sidesteps the cross-device error from applying pre-reduced grads in cross-replica context. | ||
| for _ in range(accumulation_steps): | ||
| strategy.run(_accumulate, args=(dist_inputs,)) | ||
| strategy.run(_apply_and_reset) |
There was a problem hiding this comment.
Correct pattern, and the "apply inside strategy.run" comment nails the right justification. One subtle semantic drift from production worth surfacing (does not affect the throughput numbers, so this is a note, not a change request):
Production _apply_gradients (src/aetherscan/train.py:2093-2098) does reduce → clip → apply in cross-replica context: strategy.reduce(MEAN) in _distributed_train_step runs first, then clip_by_global_norm(1.0) is computed once on the mean-across-replicas gradient, then optimizer.apply_gradients broadcasts to mirrors.
_apply_and_reset here does per-replica clip → apply-with-implicit-reduce: each replica clips its own acc / K independently against 1.0, and the all-reduce (mean) happens inside optimizer.apply_gradients on the clipped per-replica tensors. So the effective grad is mean_r(clip(acc_r / K)) rather than clip(mean_r(acc_r) / K) — not identical when at least one replica's global norm exceeds the threshold. Compute cost is essentially the same either way (one tree-reduce over params), so throughput/VRAM are faithful, but the header claim "exactly as _train_epoch" is a shade strong — the cadence (K micro-batches per apply) is exact; the reduce/clip ordering differs.
If you want production-exact clip semantics without giving up the per-replica accumulator (needed for the multi-GPU fix), you can do the merge explicitly before clipping — something like:
def _apply_and_reset():
ctx = tf.distribute.get_replica_context()
reduced = ctx.all_reduce(tf.distribute.ReduceOp.MEAN, [acc / accumulation_steps for acc in accumulators])
clipped, _ = tf.clip_by_global_norm(reduced, _CLIP_NORM)
model.optimizer.apply_gradients(zip(clipped, model.trainable_variables, strict=False))
for acc in accumulators:
acc.assign(tf.zeros_like(acc))That gets the same reduced-then-clipped-once semantics as _apply_gradients while keeping the per-replica accumulator + apply-inside-strategy.run structure that fixes the cross-device crash. But given the stated purpose is throughput + VRAM, leaving this as-is with a one-line comment noting the divergence is also fine.
| if mode == "train": | ||
| # Per-replica gradient accumulators (created in scope so they mirror across replicas). | ||
| with strategy.scope(): | ||
| accumulators = [ | ||
| tf.Variable(tf.zeros_like(v), trainable=False) for v in model.trainable_variables | ||
| ] |
There was a problem hiding this comment.
Second-order, doesn't affect the reported numbers: at --accumulation-steps 1 you still walk the accumulator path (one strategy.run(_accumulate) + one strategy.run(_apply_and_reset) per step, plus the assign_add + zero-out on every parameter). The whole thing is inside one @tf.function, so TF can fuse most of it — the bla0 numbers (795 → 3,087 cad/s for 1→6 GPUs, ~3.9×) suggest overhead is small — but a fused K=1 fast-path (single tape → clip → apply, no accumulator I/O) would be a slightly cleaner baseline for the "plain apply-every-step" row. Not a blocker, and probably not worth the extra code path unless you see K=1 numbers you don't trust.
| def sync(out) -> None: | ||
| # Force queued device work to finish before the clock is read. Train applies mutate the | ||
| # model variables, so reading one is a sufficient barrier; encode syncs on its output. | ||
| if mode == "train": | ||
| model.trainable_variables[0].numpy() | ||
| elif out is not None: | ||
| strategy.experimental_local_results(out)[0].numpy() | ||
|
|
There was a problem hiding this comment.
The model.trainable_variables[0].numpy() barrier is correct but the reason is non-obvious enough to be worth spelling out in the comment: it works because every apply ends in a collective all-reduce, so any replica whose stream has drained past step N implies all replicas completed the all-reduce for step N. Without that guarantee, reading one MirroredVariable would only sync the primary replica. A future reader debugging weird timing numbers might add a "properly" cross-replica sync assuming this one is sloppy, so a one-liner locks in the reasoning.
Claude Code review — PR #253 (second pass, focused on the accumulator pattern)Fresh read of the accumulator rewrite in The on-cluster validation (bla0, 6× A4000, 1→6-GPU 3.9× scaling on train, 4.1× on encode) is the right kind of evidence — the crash was masked by single-GPU, and the fix is only meaningfully verifiable under MirroredStrategy. Three notes inline — none blocking, none change the numbers:
Things I checked and liked
|
…cumulation mechanism Addresses the #253 second-pass review notes (all comment-only; no logic change, so the on-cluster baselines still stand): - Soften the module docstring's "exactly as _train_epoch" to cadence-exact, and note in `_apply_and_reset` that production reduces-then-clips while the benchmark clips-per-replica-then- all-reduces-on-apply — mean_r(clip(acc_r/K)) vs clip(mean_r(acc_r)/K), which differ only when a replica's global norm exceeds the clip threshold and are compute-identical, so throughput/VRAM stay faithful. - Spell out why reading one mutated MirroredVariable is a sufficient *cross-replica* sync barrier (every apply ends in a collective all-reduce). - README: name the mechanism behind the --accumulation-steps throughput bump (one apply + clip + all-reduce per K micro-batches amortizes the fixed per-apply overhead). Refs #252
|
Thanks — all three notes addressed in 1. Reduce/clip ordering ( I deliberately did not adopt the explicit 2. K=1 walks the accumulator path ( 3. Primary-replica sync barrier ( README. Named the mechanism behind the
|
Claude Code review — PR #253 (third pass, closing)Re-read Reduce/clip ordering ( K=1 accumulator path ( Cross-replica sync ( README amortization sentence. "One apply + clip + one cross-replica all-reduce per K micro-batches ... amortizing that fixed per-apply overhead over 4× as many cadences" — matches the mechanism (each apply's fixed cost is the collectives + the params-scale clip tree-reduce; K=4 pays it 4× less often). Makes the 3,087 → 4,272 cad/s bump self-explanatory. Since LGTM — leaving for maintainer merge, and per the loop rule not re-tagging. --- · Branch `feature/gpu-benchmarks` @ `63b54ec` |
… inference The benchmarks/ suite covered only hot CPU kernels; the Beta-VAE's GPU work (training step + encoder inference) was unmeasured. Add bench_gpu.py: builds the real model on one GPU, drives synthetic batches of the true pipeline shapes (train: cadence (B,6,16,512) = 18B enc + 6B dec passes/step; encode: obs (B,16,512,1)), reports throughput + peak VRAM, and --find-max sweeps the per-replica batch to the largest that fits. Documents blpc3 (RTX PRO 6000) baselines in benchmarks/README.md. GPU-only (runs in the NGC container), consistent with the framework-free suite. --find-max is capped at 4096 to avoid a TF int32 launch-config abort at batch >= 8192. Closes #252
Add --num-gpus: build the Beta-VAE under a MirroredStrategy over the first N GPUs and report aggregate throughput + per-GPU peak VRAM, covering the cross-replica gradient all-reduce and multi-GPU scaling (the single-GPU version measured neither). Also resolves the first-pass review nits: - --find-max / --batch-sizes are now a mutually-exclusive group (was decorative) - JSON records example_shape (was mislabeled obs_shape in train mode) - --warmup validated >= 1 and the sync guarded (was a NameError at --warmup 0) - module-level RNG so main/true/false/target inputs aren't byte-identical - VRAM label clarified to SI GB (1e9 B) vs nvidia-smi's MiB/GiB - argparse hardening (_positive_int / _batch_list reject bad tokens; no infinite --find-max loop at start/max 0) Refs #252
The train step applied grads every micro-batch, but the real _train_epoch accumulates the all-reduced grads over accumulation_steps micro-batches and applies once with the global-norm clip. Split the per-replica step into grad-compute only, and move the reduce/accumulate/clip/apply into one optimizer step in _measure — mirroring _distributed_train_step + _apply_gradients. --accumulation-steps K (default 1 = apply every step) now measures the once-per-K apply cadence and the persistent-accumulator VRAM; throughput is reported over all micro-batches. Encode mode is unchanged. Refs #252
The second-stage RF (sklearn, CPU) was unmeasured. bench_rf.py times prepare_latent_features (the per-cadence latent reshape), RandomForestClassifier.fit, and predict_proba on synthetic (n_cadences x 48) binary data at production hyperparameters (1000 trees, max_features=sqrt, n_jobs=-1). Framework-free (numpy + sklearn), so it runs alongside the other CPU benchmarks and needs no GPU. Also documents --num-gpus / --accumulation-steps in benchmarks/README.md and adds a MacBook bench_rf baseline (fit is a random-label upper bound). Multi-GPU + bla0 GPU baselines and cluster RF numbers land with the System-Requirements update (#183). Refs #252
…pattern) The accumulation refactor applied pre-reduced grads in cross-replica context, which core-dumps under MirroredStrategy (an Adam slot on GPU:0 accessed from GPU:1 — apply_gradients with a single reduced grad doesn't broadcast to the replicas' mirrored slots). Single-GPU masked it. Switch to the canonical per-replica accumulator pattern: accumulate grads into mirrored accumulator variables inside strategy.run over accumulation_steps micro-batches, then apply once INSIDE strategy.run so MirroredStrategy all-reduces on apply. _make_step_fn is now encode-only (the train path builds its accumulating step in _measure). Validated on bla0 (6x A4000): train single 795 / multi 3087 / accumulation K=4 4272 cad/s; encode single 53k / multi 214k obs/s — all num-gpus / accumulation / encode paths run clean. Refs #252
… accumulation Captures the bla0 (6× RTX A4000, the release-training host) single-GPU train + encode baselines, the multi-GPU MirroredStrategy scaling (~3.9–4.1× over 6 replicas), and the --accumulation-steps throughput/VRAM signature. Replaces the pending-baselines placeholder. Refs #252
…cumulation mechanism Addresses the #253 second-pass review notes (all comment-only; no logic change, so the on-cluster baselines still stand): - Soften the module docstring's "exactly as _train_epoch" to cadence-exact, and note in `_apply_and_reset` that production reduces-then-clips while the benchmark clips-per-replica-then- all-reduces-on-apply — mean_r(clip(acc_r/K)) vs clip(mean_r(acc_r)/K), which differ only when a replica's global norm exceeds the clip threshold and are compute-identical, so throughput/VRAM stay faithful. - Spell out why reading one mutated MirroredVariable is a sufficient *cross-replica* sync barrier (every apply ends in a collective all-reduce). - README: name the mechanism behind the --accumulation-steps throughput bump (one apply + clip + all-reduce per K micro-batches amortizes the fixed per-apply overhead). Refs #252
63b54ec to
836bd17
Compare
|
Release note draft — New Two new benchmarks let you profile the deep-learning stages on your own hardware: Draft for the changelog; edit or drop as needed. |
docs(benchmarking): reflect bench_gpu.py + bench_rf.py added in PR #253
1 participant
Closes #252
Adds
benchmarks/bench_gpu.py— the first GPU benchmark in the suite (the existing scripts timeonly hot CPU kernels). It builds the real Beta-VAE (
create_beta_vae_model) on a single GPU anddrives synthetic batches of the true pipeline shapes to measure per-GPU throughput and peak VRAM,
sweeping the per-replica batch to find the largest that fits.
What it measures
--mode train— a full training step:compute_total_loss(reconstruction + KL + bothclustering heads) →
tape.gradient→ global-norm clip (1.0, matching_apply_gradients) → Adamapply. A training example is a cadence
(B, 6, 16, 512), so one step runs18*Bencoder +6*Bdecoder passes — exactly like
_distributed_train_step.--mode encode— encoder forward on observations(B, 16, 512, 1), i.e. the inference path.--find-maxdoubles the per-replica batch until OOM and reports the largest power of two thatfits;
--batch-sizes a,b,cmeasures explicit sizes.Peak VRAM is per GPU (each replica holds a full model copy → combined ≈ ×replicas). Single-GPU +
memory growth (mirrors
main.py) so the VRAM figures track real usage; model built inside aOneDeviceStrategyscope per the repo's strategy-scope rule.Key finding
The Beta-VAE is compute-bound, not VRAM-bound. On blpc3 (RTX PRO 6000, 96 GB) training
throughput peaks at per-replica batch ~128 (2.81 GB) and does not improve toward the 76 GB VRAM
ceiling; inference peaks near batch ~1024 using <2 GB. The defaults (train 128, inference 2048) are
already near-optimal — a large-VRAM card's spare memory can't be converted into throughput for this
compact model. Full baseline tables are in
benchmarks/README.md.Notes / caveats
gpu/clusterintegration tests, it is notpart of the CI CPU surface, and (consistent with the deliberately framework-free
benchmarks/suite) it is not pytest-collected; it prints a table and writes a JSON result.
--find-maxis capped at--max-batch 4096: an encoder forward whose conv feature maps exceed2^31 elements (batch ≳ 8192) trips an uncatchable TF int32 launch-config abort rather than a
clean OOM. 4096 covers the training VRAM ceiling and a generous inference range.
Verification
ruff check/ruff formatclean; pre-commit clean.--mode train --find-max(max batch 4096 @ 76.5 GB) and--mode encode --batch-sizes 256..4096; numbers are the tables in the README.🤖 Generated with Claude Code