Make non-paged 8-bit optimizers safe for FSDP2 local shards

[heiheiha798]

Problem

Upstream issue #1633 reports a deterministic illegal CUDA memory access on the first AdamW8bit.step() under PyTorch FSDP2. The issue is open, labeled Contributions Welcome, FSDP, and Optimizers, and currently has no linked branch or pull request.

The current optimizer runtime assumes every parameter, gradient, and state object passed to the native kernels is a plain tensor. Under FSDP2, the optimizer key is a DTensor: state allocation through zeros_like(p) preserves the wrapper, while update_step() ultimately passes wrapper-backed pointers and a local gradient element count into the existing CUDA optimizer kernel. The detailed upstream diagnosis shows that both plain local state allocation and passing the parameter's local tensor to the kernel are required; either change alone still fails.

This is a public distributed-training correctness gap, not a kernel-tuning hypothesis. The goal is to reuse the existing optimizer kernels safely on each FSDP2 local shard while preserving the global DTensor parameter as the optimizer key.

Bounded Scope

Investigate a Python-only runtime fix in bitsandbytes/optim/optimizer.py plus focused distributed tests:

• detect an FSDP2 DTensor through public, lazily imported PyTorch APIs;
• keep the global parameter object as the param-group and state-dict key;
• keep optimizer-mode and min_8bit_size selection based on the global parameter size;
• use to_local() to obtain the contiguous, nonempty local parameter and gradient views at the native-kernel boundary;
• allocate state1/state2 and 256-element absmax arrays as plain tensors sized for that local shard; and
• leave the existing CUDA kernels, arithmetic, state formats, and plain-tensor path unchanged.

The supported claim is deliberately narrow: ordinary contiguous FSDP2 Shard(0) parameters on the non-paged Optimizer1State and Optimizer2State runtime paths with max_unorm == 0, validated through AdamW8bit and Lion8bit, including their small-parameter FP32-state fallback.

Explicit Non-Goals

• paged/UVM optimizers;
• LAMB/LARS or any max_unorm > 0 path requiring distributed-global norm semantics;
• scheduled AdEMAMix, which owns separate state/update overrides;
• arbitrary hand-created DTensor placements, noncontiguous or empty local shards;
• FSDP2 distributed checkpoint/save-load support beyond preserving existing behavior;
• changes to CUDA kernels, optimizer math, block size, quantization maps, FSDP collectives, Linear4bit, or the separate issue FSDP2 fully_shard breaks Linear4bit forward via NaN canonicalization of packed-NF4 bf16 storage bitsandbytes-foundation/bitsandbytes#1945.

Unsupported DTensor combinations must not be passed as raw wrapper pointers to native kernels. If encountered during implementation, fail clearly rather than silently copy/update detached storage or expand this issue.

Correctness Plan

1. On the exact upstream baseline, reproduce issue illegal memory access with FSDP2 and AdamW8bit bitsandbytes-foundation/bitsandbytes#1633 with torchrun on at least two B300 GPUs and CUDA_LAUNCH_BLOCKING=1; the first optimizer step must reach the documented native illegal access.
2. Validate patched AdamW8bit and Lion8bit at world sizes 1 and 2; world size 4 is useful when available but not required for the minimum patch.
3. Use fixed deterministic gradients and compare every local shard update to an explicit plain-local-tensor bitsandbytes reference from an identical snapshot. Require exact parameter, state, and absmax equality where the shard/block partition is identical.
4. Cover FP16, BF16, and FP32 parameter/gradient storage; global sizes below and above min_8bit_size; aligned shards; local tails around 255/256/257 and 4095/4096/4097; and multiple steps.
5. Assert that optimizer state remains keyed by the global parameter, while native state tensors are plain tensors with local numel, device, dtype, and ceil(local_numel / 256) absmax sizing.
6. Run plain-tensor optimizer regression and state-dict round-trip coverage, the focused optimizer suite, full pre-commit, and a small compute-sanitizer check. Distributed FSDP2 checkpoint support is reported separately and is not implied by a passing runtime step.

B300 Validation and System Evidence

Use one multi-B300 Slurm allocation with recorded node/GPU/driver/CUDA/PyTorch/bitsandbytes identities and exact commands. Do not run GPU work on the login node.

• Run a synthetic fully trainable model and a small public FSDP2 training reproduction with AdamW8bit; add Lion8bit as the one-state coverage path.
• Record per-rank optimizer-state bytes, peak allocated/reserved memory, median/p10/p90 step latency, and finite training loss over at least 100 timed steps after warmup.
• Compare AdamW8bit state storage with torch.optim.AdamW under the same FSDP2 topology. The expected two-moment storage is approximately 2.03125 bytes per parameter for uint8 states plus FP32 block scales versus 8 bytes per parameter for two FP32 moments; measure rather than infer the actual per-rank and peak-memory effect.
• Trace one step to confirm that the existing optimizer kernels receive local n, and that the adapter adds no host-device copies, implicit all-gather, or extra per-parameter collective.
• Confirm that the unchanged plain-tensor bitsandbytes path has no material latency regression.

This is a correctness and distributed-memory-support contribution. Do not claim a B300 speedup unless the measurements independently justify one.

Acceptance

• the exact-baseline first-step failure is reproduced;
• patched world-size-2 FSDP2 steps complete without invalid access;
• local parameter/state/absmax results satisfy the exact reference oracle across the aligned and tail matrix;
• AdamW8bit and Lion8bit cover both the 8-bit and small-parameter FP32-state paths;
• state is genuinely local and the measured Adam moment-state reduction is consistent with the expected approximately 74.6% reduction;
• the trace shows no adapter-induced all-gather, transfer, or collective; and
• plain-tensor tests and latency remain healthy.

If public to_local() views do not alias the updatable shard safely, if exact local parity fails, or if correct support requires paged/global-norm/checkpoint redesign, record a bounded no-go rather than expanding the patch.

Baseline: 95f9af309d4d5793847169c39288dcd3fcbdf564.

[heiheiha798]

Cycle result: aborted before any GPU workload

The required exact-baseline WS2 reproduction could not be submitted under the available Slurm policy, so this implementation cycle is being recorded as aborted for environment capacity/policy, not as a correctness or performance no-go.

Mechanical evidence

• Baseline and branch head: 95f9af309d4d5793847169c39288dcd3fcbdf564

• Branch: b300/fsdp2-local-optimizer-runtime

• Required request: partition gpu, two GPUs, one node, 20-minute limit

• Submission result, before a job ID was created:

  sbatch: error: QOSMaxGRESPerUser
  sbatch: error: Batch job submission failed: Job violates accounting/QOS policy (job submit limit, user's size and/or time limits)
  

• squeue -u 12921178 was empty immediately after the rejection, ruling out an active job from this user consuming the allowance.

• Cluster discovery exposes the default gpu partition with a one-hour limit and a node advertising eight GPUs, while prior accepted single-GPU allocations in this workspace resolved to gpu1_qos. The two-GPU request is nevertheless rejected by the effective per-user GRES policy.

The ignored, syntax-checked reproducer and wrapper remain at:

• benchmark-evidence/cycle16/fsdp2_baseline_repro.py
• benchmark-evidence/cycle16/slurm_baseline_repro.sh

They were prepared to extract the exact baseline, build the official CUDA 13 x64-compatible target list 75;80;86;89;90;100;120, assert two B300/CC10.3 devices and loaded-library provenance, and run torchrun --standalone --nproc_per_node=2 with CUDA_LAUNCH_BLOCKING=1.

Disposition

The issue's acceptance contract requires a real two-B300/WS2 FSDP2 run. I did not guess another account or QoS, substitute a one-GPU or oversubscribed-process run, or weaken that gate. No Slurm workload or CUDA build ran, so there is no baseline reproduction, correctness result, memory measurement, performance result, or technical conclusion about the proposed local-view/state design.

No production or tracked file was changed, no commit was created, no branch was pushed, and no draft PR was opened. The issue remains open for a future environment that permits the required multi-GPU allocation.