Skip to content

feat(optimizer): add opt-in packed 4-bit AdamW states - #537

Merged
xsuler merged 2 commits into
mainfrom
feat/adamw4bit
Sep 1, 2026
Merged

feat(optimizer): add opt-in packed 4-bit AdamW states#537
xsuler merged 2 commits into
mainfrom
feat/adamw4bit

Conversation

@xsuler

@xsuler xsuler commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add opt-in --adam-4bit / adam_4bit optimizer selection for the CUDA backend
  • store both Adam moments as packed 4-bit nibbles with parameter-local B128 scales (about 1.06 bytes/parameter including FP32 scale metadata for large tensors)
  • use the signed dynamic-exponent first-moment map and zero-excluding second-moment map from Memory Efficient Optimizers with 4-bit States
  • add fused CUDA updates that operate directly on packed state, plus bounded block-wise CPU reference updates and versioned checkpoint loading
  • make the existing 8-bit optimizer parameter-local/block-wise so outliers cannot set one bucket-global scale and amplify the effective step at the same learning rate
  • skip only a block containing non-finite values instead of poisoning neighboring packed state
  • expose the option through typed trainer/engine configs, CLI summaries, and dashboard command construction

Validation

  • ruff check and ruff format
  • Python compileall
  • git diff --check
  • CPU tests cover packing/storage budget, quantizer maps, checkpoint round trips, disk offload, FP32 tracking, non-finite block isolation, configuration/CLI routing, and the 8-bit same-learning-rate regression
  • CUDA parity tests are included for CI; this development host has no CUDA toolchain/device, so the fused extension was not built locally

Notes

The model checkpoint format is unchanged. Only the optimizer checkpoint carries the explicit 4-bit format version and block size.

Closes #535

@xsuler xsuler self-assigned this Sep 1, 2026
@xsuler
xsuler deployed to push_check September 1, 2026 04:43 — with GitHub Actions Active
@xsuler
xsuler merged commit 0ca1640 into main Sep 1, 2026
7 checks passed
@xsuler
xsuler deleted the feat/adamw4bit branch September 3, 2026 03:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add opt-in 4-bit AdamW optimizer states

1 participant