perf(optimizer): stream factored AdamW4bit state

[xsuler]

Summary

• replace matrix-sized AdamW4bit second moments with Adafactor-style FP32 row/column factors for tensors with rank >= 2
• retain packed signed 4-bit first moments and packed B=128 zero-excluding 4-bit second moments for one-dimensional tensors
• add fused bounded CUDA passes for factored gradient-square statistics and parameter updates without parameter-sized FP32 moment temporaries
• stream microbatches into BF16 DP gradient shards and perform shard-aware norm/clipping only for --adam-4bit
• update parameters in parameter order and release completed gradient/state buckets promptly
• preserve AdamW8bit and FP32 AdamW's existing FP32 accumulation path and collective schedule
• cover checkpoint format v3, clearing, CPU/disk offload, distributed factor statistics, and memory metrics

Algorithm

The packed 4-bit codebooks follow Memory Efficient Optimizers with 4-bit States. Matrix second moments follow the row/column estimator from Adafactor: each TP-local tensor is viewed as [shape[0], -1], DP ranks sum partial squared-gradient row/column statistics, and the update reconstructs row[i] * column[j] / mean(row).

This is deliberately a factored AdamW approximation, not the prior rank-1 quantization design. It removes the per-element packed second moment for matrices. Vectors keep the 4-bit B128 second moment because row/column factorization is not applicable.

Adam4-only scope

AdamW4bit opts into stream_gradient_shards=True and BF16 gradient shards. The trainer gates the streamed path on that optimizer capability. The shared FP32-master base defaults remain False and FP32, so AdamW8bit and FP32 AdamW continue through the pre-existing full FP32 main_grad accumulation and step-time synchronization path.

Memory impact

For a large matrix, persistent optimizer moments approach 0.5 byte/parameter for packed momentum plus B128 scale metadata and 4 * (rows + columns) factored variance bytes. The Adam4 DP gradient shard is BF16 instead of FP32 and the trainer does not retain a full-model FP32 main_grad for this mode.

Peak savings still depend on DP/TP topology, activation memory, temporary collectives, parameter gathering, and state offload. This PR does not claim a fixed end-to-end GiB reduction without GPU profiling.

Compatibility

• model checkpoint format is unchanged
• Adam4 optimizer state is versioned as format v3
• previous experimental Adam4 optimizer-state formats are intentionally rejected
• no AdamW8bit/FP32 optimizer-state or runtime behavior change is intended
• CUDA only; MLX is unchanged

Validation performed in this checkout

• ruff check on all changed Python modules and the AdamW4bit test module: passed
• ruff format --check on all changed Python modules and the AdamW4bit test module: passed
• python3 -m py_compile on changed Python modules and tests: passed
• git diff --check: passed

The current local Python environment does not provide PyTorch, and this host does not provide a CUDA build/runtime, so the revised CPU/Gloo test suite and fused CUDA build were not executed locally. Those results are intentionally not claimed here; PR CI and a CUDA runner must validate them.

Follow-up validation required

• targeted CPU/Gloo optimizer tests
• CUDA extension build and fused-kernel parity tests
• representative 7B multi-GPU peak-memory and throughput comparison against AdamW8bit and FP32 AdamW
• convergence/reward comparison across multiple seeds

Addresses #542