[Hang-get]

Problem

The CPU blockwise 8-bit optimizer update creates several full-size FP32 temporaries on every step. In particular, g.float() * 1.0 copies an already-FP32 gradient, and Adam/LAMB build a separate denominator from the dequantized variance. For large parameters, these allocations substantially increase step latency and peak memory.

Cause

The CPU implementation used allocation-producing expressions even when the input already had the required dtype and scale. It also retained the updated FP32 variance while materializing another full-size denominator, although the variance only needs to be requantized before the parameter update.

Change

• Reuse FP32 gradients when no scaling is required, while keeping weight-decay operations from modifying the input gradient.
• Requantize Adam/LAMB variance before reusing its FP32 buffer as the denominator.
• Update FP32 parameters in their existing storage and skip the redundant copy back to the same tensor.

Benchmark

Environment: Intel Xeon Platinum 8352Y, one pinned CPU core, OMP_NUM_THREADS=1, MKL_NUM_THREADS=1, Python 3.11.18, PyTorch 2.15.0+cpu, native Release CPU backend.

Workload: one 20,000,000-element FP32 parameter with Adam8bit, one warmup step followed by seven timed steps. Each result below is the median of per-process medians from three fresh processes.

	Before	After	Change
Optimizer step	0.634488 s	0.507125 s	1.25x faster, -127.363 ms
Peak RSS increase over the pre-step process	444.3 MiB	215.2 MiB	-229.1 MiB (-51.6%)

Parameter and optimizer-state hashes were bitwise identical in all six benchmark processes.

A separate baseline comparison covered Adam, LAMB, Momentum, LARS, Lion, RMSprop, Adagrad, and AdEMAMix with FP32, FP16, and BF16 parameters for five steps each. All 24 configurations produced bitwise-identical parameters and optimizer states. A gnorm_scale=0.25 case was also bitwise identical and left the input gradient unchanged.

Tests

• BNB_TEST_DEVICE=cpu python -m pytest tests/test_optim.py -q -k 'test_optimizer8bit and adam8bit_blockwise' — 9 passed
• 15 selected CPU test_optimizer8bit cases for Lion, Momentum, RMSprop, AdEMAMix, and scheduled AdEMAMix across FP32, FP16, and BF16 — 15 passed
• pre-commit run --all-files — passed

Not run: accelerator tests or the full test suite. The changed implementation is CPU-only, and the repository guidance asks for targeted tests.

AI assistance (Cursor) was used for this change.

[Hang-get]

Closing this, as I'm focusing my contributions elsewhere. Sorry for the noise.