Investigate SM103 bias-correction broadcast in 32-bit optimizers

[heiheiha798]

Problem

The CUDA 32-bit two-state optimizer kernels recompute CTA-invariant bias-correction scalars in every thread.

• kOptimizer32bit2State launches 1,024 threads per CTA and each thread evaluates the same powf(beta1, step), powf(beta2, step), sqrtf, and derived step size before processing a 4,096-element tile.
• The ADAM instantiation of kPreconditionOptimizer32bit2State, used when max_unorm > 0, launches 512 threads and each thread evaluates the same two reciprocal powf corrections.

All operands are CTA-uniform. Computing each correction once in thread 0, publishing the float bit patterns through shared memory, and consuming them after an already-required first CTA barrier could remove redundant special-function work on B300. However, the compiler or scheduler may already hide that cost, and shared publication can change resource use. This issue is a measurement-gated hypothesis, not a performance claim.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Build isolated baseline and scalar-broadcast variants from the same source/toolchain. The broadcast variant changes only the bias-correction scalars in kOptimizer32bit2State and the two-state ADAM precondition kernel.
• Reuse the kernels' existing first CTA barrier to publish the shared scalars if safe; do not add another synchronization merely for this optimization.
• If and only if B300 evidence clears the gates below, select private baseline/broadcast kernel variants through a cached per-device runtime check for exact compute capability 10.3. Every non-SM103 device and HIP must keep the current expression path.
• Validate the official CUDA 13 x64 compatible target input 75;80;86;89;90;100;120 as the primary deployed-code path, and report native SM103 builds separately. Do not rely solely on __CUDA_ARCH__ == 1030: official wheels do not currently include an SM103 image.
• Keep the existing C/Python API and optimizer arithmetic unchanged. Open upstream PR Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels bitsandbytes-foundation/bitsandbytes#2040 changes Adam/AdEMAMix weight-decay ordering in the same source file; rebase and preserve that ordering if it lands, but do not absorb the correctness fix here.

The expected production surface is limited to csrc/kernels.cu, the narrow host launch selection in csrc/ops.cu, and focused validation/benchmark coverage. No broad architecture-dispatch helper is required.

Correctness oracle

For identical deterministic inputs, require bitwise-identical parameter, state1, state2, and unorm buffers between baseline and broadcast variants after one and multiple steps. Thread 0 must evaluate the same float expressions and shared publication must preserve their bits; approximate equality is not an acceptable replacement for this launch-only optimization.

• Cover Adam32bit/AdamW32bit and AdEMAMix32bit with max_unorm=0, plus LAMB32bit or the corresponding two-state ADAM path with max_unorm>0 to exercise both precondition and update kernels.
• Cover FP16, BF16, and FP32 parameter/gradient storage where registered; steps 1, 2, 10, and 1000; weight_decay=0 plus one nonzero control; and skip_zeros false/true where existing semantics support it.
• Include n=1, 255/256/257, 4095/4096/4097, 8191/8192/8193, 64K, 1M, and realistic parameter counts. Include zero, random finite, and mixed large/small finite gradients.
• Keep the repository's public optimizer-reference tolerances as a secondary oracle and run the focused tests/test_optim.py matrix, CUDA build, and full pre-commit suite. All CUDA compilation and GPU tests run only through Slurm.

Preserve whichever weight-decay ordering is present after rebasing. Do not change max-unorm, skip_zeros, NaN/Inf, paging, or state-layout semantics.

B300 benchmark and acceptance

Use one B300/SM103 allocation and record the exact commit/build identity, node/GPU, driver, CUDA compiler, PyTorch/bitsandbytes versions, target list, commands, environment, and raw JSONL/log paths. Interleave baseline and broadcast variants built with identical flags.

• Use the official-compatible target list as the primary matrix so the runtime SM103 selection is exercised in the same form as distributed CUDA 13 x64 binaries. Use a native 103 build as a separately labeled corroborating control, not as a substitute.
• Measure direct torch.ops.bitsandbytes.optimizer_update_32bit calls with preallocated buffers at n=256, 4096, 64K, 1M, 16M, 64M or the largest memory-safe size, 4096x4096, and 8192x28672; FP16/BF16/FP32; Adam, AdEMAMix, and a max_unorm/LAMB control; steps 1, 10, and 1000.
• Measure end-to-end optimizer.step() for AdamW32bit, AdEMAMix32bit, and LAMB32bit on representative large tensors and a model-like collection containing many sub-4096 parameters. Exclude allocation and state construction from timing.
• Use at least 20 warmups where practical and at least 7 interleaved batches per cell. Scale CUDA-event iterations by problem size so each batch is stable. Report all round samples, median, p10/p90, elements/s, and a clearly defined effective-bandwidth estimate; never report best-only values.

Proceed only if the bit-exact oracle passes, the official-compatible broadcast path produces at least a 5% median direct-op gain on at least two realistic Adam/AdEMAMix shapes with dispersion below the gain, and at least one representative end-to-end optimizer step improves repeatably by 3% or more without a greater-than-5% regression in any adjacent large-shape/dtype cell. If only the native build or LAMB benefits, the compiler already eliminates the cost, or gains disappear end to end, record a no-go.

Non-goals

No optimizer-math or weight-decay-ordering change; no broadcast of unrelated invariant expressions; no CTA-size/grid tuning; no persistent optimizer grid; no 8-bit optimizer or state-layout work; no max-unorm/skip_zeros semantic fix; no public API; no CUDA wheel target change; no other architecture/backend generalization; no dependency on prior fork cycles.

[heiheiha798]

B300 result: NO-GO

The required bit-exact correctness gate did not pass, so this experiment stopped before performance measurement. No production source was changed and no draft PR was created.

Environment and method

• Baseline: 95f9af309d4d5793847169c39288dcd3fcbdf564, clean branch b300/sm103-32bit-optimizer-bias-broadcast
• GPU: NVIDIA B300 SXM6 AC, compute capability 10.3, 148 SMs
• Node: dev-slurm; PyTorch 2.13.0+cu130; CUDA toolkit/runtime 13.0; driver 580.126.09
• Primary build targets: 75;80;86;89;90;100;120, producing 75-real;80-real;86-real;89-real;90-real;100-real;120
• Separate native control target: 103
• All four libraries were built from isolated source trees with the same Release flags in each Slurm allocation. Correctness ran before any timing cell.
• Oracle: torch.equal for parameter, state1, state2, and unorm after deterministic steps 1, 2, 10, and 1000. Approximate equality was not accepted.

The exact calibration command was:

python benchmark-evidence/cycle5/sm103_optimizer_bias_broadcast.py \
  --output benchmark-evidence/cycle5/sm103-opt32-broadcast-job-${SLURM_JOB_ID}.jsonl \
  --build-root ${SLURM_TMPDIR}/bnb-sm103-opt32-broadcast-${SLURM_JOB_ID}/benchmark-build \
  --warmup 20 --repetitions 100 --rounds 7

Job 4904: both proposed broadcast sites

The first diagnostic broadcasted the correction scalars in both kOptimizer32bit2State and kPreconditionOptimizer32bit2State while reusing their existing barriers.

• 94 official-compatible correctness records completed successfully before the first mismatch. These include the complete Adam and AdEMAMix matrix across FP16/BF16/FP32, all configured sizes/steps, nonzero weight decay, and skip_zeros, plus LAMB/FP16 through n=8193.
• At LAMB/FP16, n=65536, step 1, the broadcast result differed from baseline by one byte in unorm; parameter, state1, and state2 remained bitwise identical.
• Slurm state: FAILED (NonZeroExitCode), runtime 2:39.
• Raw evidence: benchmark-evidence/cycle5/slurm-4904-opt32-broadcast.out and benchmark-evidence/cycle5/sm103-opt32-broadcast-job-4904.jsonl.

The precondition kernel accumulates per-CTA values through a cross-CTA atomicAdd. Removing redundant work before that reduction changes CTA arrival timing and can change the floating-point accumulation order, violating the strict unorm oracle.

Job 4908: update-kernel-only narrowing

The scope was then narrowed to broadcast only correction1, correction2, and the derived step_size in kOptimizer32bit2State. The complete kPreconditionOptimizer32bit2State source region was verified byte-for-byte unchanged between the isolated baseline and broadcast source trees. No reduction, launch, arithmetic, or synchronization was changed.

• The same 94 official-compatible correctness records completed successfully before the first mismatch.
• At LAMB/FP16, n=65536, step 10, the two executions again differed by one byte in unorm; parameter, state1, and state2 remained bitwise identical.
• Because the precondition source was unchanged, this demonstrates that separate executions of the existing cross-CTA atomic reduction do not reliably produce bit-identical unorm at a realistic multi-CTA size.
• Slurm state: FAILED (NonZeroExitCode), runtime 2:37.
• Raw evidence: benchmark-evidence/cycle5/slurm-4908-opt32-broadcast.out and benchmark-evidence/cycle5/sm103-opt32-broadcast-job-4908.jsonl.

Decision

Neither job reached a performance record. The issue requires bit-identical unorm, and weakening that oracle is not appropriate for a launch-only optimization. Making the cross-CTA reduction deterministic would be a separate optimizer-math/reduction change outside this issue's scope. Therefore there is no correctness- or performance-backed basis for an SM103 production dispatch change, and this investigation is recorded as NO-GO.