Calibrate SM103 quantile-map replicas in the Adam 8-bit optimizer

[heiheiha798]

Problem

The CUDA two-state 8-bit optimizer kernel copies each 256-entry quantile map
into two padded shared-memory rows:

smem_quantiles[2][257]
lane_id = threadIdx.x % 2

ADAM then repeatedly reads these maps while dequantizing both optimizer
states, running the data-dependent quantize_2D search, and checking the
quantized sign. The replica count and 257-float skew have been unchanged since
the original 2021 kernel. A different replica count may alter shared-memory
broadcast/bank-conflict behavior on B300, but more copies also increase map
initialization and shared-memory use. This issue is a measurement-gated
calibration, not a claim that more or fewer replicas are faster.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Build isolated experimental variants with QMAP_COPIES={1,2,4,8} from the
  same source, compiler, flags, and allocation. Copy count 2 is the baseline.
• Parameterize only the private CUDA
  kOptimizerStatic8bit2StateBlockwise quantile-map copy count. Keep its
  256-thread CTA, 256-element block, one value per thread, grid geometry, CUB
  operations, reductions, optimizer arithmetic, quantile values, packed state
  format, and public interface unchanged.
• If and only if a single copy count passes the calibration, correctness, and
  held-out performance gates below, select it only for exact runtime compute
  capability 10.3 and compile-time optimizer ADAM.
• ADEMAMIX, all one-state optimizers, HIP, and every non-SM103 device must keep
  the exact two-copy path. Use a narrow cached per-current-device runtime check
  and include its overhead in eager timing; do not introduce a general
  architecture-dispatch framework.
• Use official CUDA 13 x64 compatible target input
  75;80;86;89;90;100;120 as the primary deployed-code path. Report native
  103 builds only as separately labeled corroboration.
• If upstream PR Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels bitsandbytes-foundation/bitsandbytes#2040 lands first, rebase and preserve its weight-decay
  ordering without absorbing that correctness change.

Expected production changes are limited to the private declaration,
definition, selected ADAM instantiations, narrow host launch choice, and
focused tests in csrc/kernels.cuh, csrc/kernels.cu, csrc/ops.cu, and the
nearest optimizer tests. Experimental forcing may remain in validation
scaffolding but must not become a public API.

Correctness oracle

Generate valid deterministic Adam8bit states with a fixed baseline warm-up,
then clone parameter, gradient, state1, state2, absmax1, absmax2, and both
quantile maps before invoking each forced variant.

• Require raw-bit equality of parameter, state1, state2, absmax1, and absmax2
  after one and repeated updates. Require the gradient and quantile maps to
  remain raw-bit unchanged. Raw-byte/integer-view comparison is required for
  NaN controls because torch.equal reports matching NaNs as unequal.
• Cover FP16, BF16, and FP32 parameter/gradient storage; steps 1, 2, 10, and
  1000; fresh-zero and warmed/clumped states; random finite, zero-heavy,
  mixed-magnitude, NaN, and Inf inputs; weight decay 0/nonzero; and
  skip_zeros false/true.
• Cover n=255,256,257,4095,4096,4097,37887,37888,37889,64K,1M and realistic
  large aligned tensors.
• Dispatch controls must prove ADEMAMIX, MOMENTUM/RMSPROP/ADAGRAD/LION,
  non-SM103 devices, and HIP retain copy count 2. The known schedule-sensitive
  ADEMAMIX unaligned-tail result is not part of the candidate oracle; it must
  never enter the candidate path.
• Run focused public Adam8bit/AdamW8bit reference coverage, relevant
  tests/test_optim.py, the official-compatible CUDA build, and full
  pre-commit. All CUDA compilation and GPU execution must run through Slurm.

Fixed calibration and holdout plan

Use one B300/SM103 allocation and interleave all isolated variants. Record the
exact source identity, node/GPU, driver, compiler/toolkit, PyTorch and
bitsandbytes versions, resolved fatbin targets, commands/environment,
warmups/iterations, compiler resource output, build-time/library-size deltas,
raw JSONL/log paths, every sample, median, and p10/p90.

Stage A selects at most one candidate before any held-out result is inspected:

• Compare copies 1, 4, and 8 against baseline 2 using FP16 at
  n=37888, 3*37888, 1M, and 16M, with both fresh and warmed states. This is
  a fixed eight-cell calibration set.
• For each copy count, compute its eight baseline/candidate median ratios.
  It is eligible only if the median of those ratios exceeds 1.02, at least six
  of eight cells favor it, no calibration cell regresses by more than 5%, and
  interleaved dispersion is below the observed median effect.
• Choose the eligible copy count with the highest median ratio; break a tie by
  choosing the smaller shared-memory footprint. If none is eligible, record
  NO-GO. Do not choose by dtype, size, state distribution, profiler result, or
  any Stage-B measurement.

Stage B validates only the frozen Stage-A winner on disjoint held-out cells:

• BF16 and FP32, plus FP16 controls; n=4096,64K,4M,64M or the largest safe
  size, 4096*4096, and 8192*28672; steps 1, 10, and 1000; realistic warmed
  states. Use direct preallocated
  torch.ops.bitsandbytes.optimizer_update_8bit_blockwise calls, at least 20
  warmups, and at least seven interleaved CUDA-event batches with size-scaled
  iterations. Report elements/s and clearly defined effective bytes/s.
• If available in the same allocation, profile one predeclared held-out cell
  with Nsight Compute for shared-load conflict/replay and occupancy evidence.
  Profiling may explain timing but cannot select the copy count or replace it.
• Measure eager Adam8bit and AdamW8bit optimizer.step() over a fixed,
  reproducibly recorded transformer-like resident parameter collection with
  quantized tensors above min_8bit_size=4096 and unchanged small 32-bit
  parameters. Preallocate states and gradients and time only steps. Paged
  controls may be reported separately but are not acceptance evidence.

Proceed only if the frozen winner is bit-exact, improves at least two realistic
held-out direct cells by at least 5% with dispersion below the gain, improves
one representative eager step repeatably by at least 3%, and causes no
greater-than-5% regression in any adjacent held-out dtype/shape. The same
single copy count must pass the official-compatible build and public-path
tests. Native-only, Stage-A-only, direct-only, or profiler-only gains are
NO-GO and must not produce a draft PR.

Non-goals

No ADEMAMIX change; no one-state optimizer change; no persistent grid or CTA
size change; no known tail-layout fix; no bias-correction change; no optimizer
math, quantization algorithm, state format, public API, or CUDA wheel-target
change; no other architecture/backend generalization; and no dependency on
prior fork branches.

Same-address broadcast may already make two copies optimal, and a different
copy count can increase initialization work or lower occupancy. Those are
expected negative outcomes for this bounded experiment, not reasons to widen
its scope.

[heiheiha798]

B300 calibration result: NO-GO

The fixed Stage A gate did not select a quantile-map replica count, so this experiment stops without a production change or draft PR.

Environment and method

• Source: exact baseline 95f9af309d4d5793847169c39288dcd3fcbdf564 on branch b300/sm103-adam8bit-qmap-replicas; tracked worktree clean.
• Slurm job: 4948 on dev-slurm, one NVIDIA B300 SXM6 AC, runtime CC 10.3, 148 SMs, driver 580.126.09.
• Software: Python 3.12.3, PyTorch 2.13.0+cu130, CUDA toolkit 13.0 (nvcc 13.0.88), bitsandbytes 0.50.2.dev0.
• Primary deployed-code build input: 75;80;86;89;90;100;120; every isolated build resolved to 75-real;80-real;86-real;89-real;90-real;100-real;120.
• Variants: isolated QMAP_COPIES={1,2,4,8} libraries with identical Release flags; 2 copies is the baseline. Build times were 60.82, 60.52, 60.21, and 60.58 seconds respectively. Library sizes were 3,918,064; 3,922,160; 3,926,256; and 3,926,256 bytes.
• Timing: 20 warmups, seven interleaved CUDA-event batches, and size-scaled repetitions (100 at 37,888; 60 at 113,664 and 1M; 20 at 16M). Every sample and round median is in the JSONL evidence.
• Command:

sbatch benchmark-evidence/cycle7/slurm_sm103_adam8bit_qmap_replicas.sh

Fixed Stage A results

Values below are baseline qmap2 median / variant median; values above 1 favor the variant.

Count	State	qmap1	qmap4	qmap8
37,888	fresh	1.01366	1.03086	1.00000
37,888	warmed	1.00599	1.05000	0.99408
113,664	fresh	0.98834	0.99706	1.01043
113,664	warmed	0.98534	0.99556	1.01205
1,048,576	fresh	1.00000	1.00000	1.00000
1,048,576	warmed	1.00000	1.00000	1.00000
16,000,000	fresh	0.97409	0.99471	0.97925
16,000,000	warmed	0.97917	0.99434	0.97913

The predeclared eligibility summaries were:

Candidate	Median of 8 ratios	Favoring cells	Worst ratio	Median round spread	Eligible
1	0.99417	2/8	0.97409	0.01089	no
4	0.99853	2/8	0.99434	0.00976	no
8	1.00000	2/8	0.97913	0.00845	no

Each measured Stage A cell was raw-bit identical to qmap2 for parameter, packed state1/state2, and absmax1/absmax2, with gradient and both qmaps unchanged. This establishes the launch-only variants' exactness on the selection cells; it is not presented as the full held-out correctness matrix.

No candidate met the required median ratio greater than 1.02 or the 6-of-8 consistency rule. The only local result above 2% was qmap4 at 37,888 elements, while qmap4 was neutral or slower at all larger fixed calibration sizes. Per the preregistered protocol, Stage B, native-SM103 corroboration, profiler collection, full correctness, and eager Adam8bit/AdamW8bit timing were therefore correctly skipped rather than used to choose a candidate after the fact.

Evidence and disposition

• Raw Slurm log: benchmark-evidence/cycle7/slurm-4948-adam8-qmap.out
• Raw JSONL: benchmark-evidence/cycle7/sm103-adam8-qmap-job-4948.jsonl
• Terminal markers: SM103_ADAM8BIT_QMAP_REPLICAS=NO_GO_STAGE_A and SM103_ADAM8BIT_QMAP_JOB=PASS.

salvage_hint=NONE: no replica count has a stable original-threshold benefit across the fixed calibration region. The isolated small-cell qmap4 observation does not pass the original selection gate and has no held-out or end-to-end support, so narrowing around it would be post-hoc specialization rather than a bounded salvage candidate.

No tracked source change, commit, push, or PR was created. The issue remains open as a record of the negative B300 result.