v0.44.0: New AdEMAMix optimizer, Embeddings quantization, and more!

[Titus-von-Koeller]

New optimizer: AdEMAMix

The AdEMAMix optimizer is a modification to AdamW which proposes tracking two EMAs to better leverage past gradients. This allows for faster convergence with less training data and improved resistance to forgetting.

We've implemented 8bit and paged variations: AdEMAMix, AdEMAMix8bit, PagedAdEMAMix, and PagedAdEMAMix8bit. These can be used with a similar API to existing optimizers.

import bitsandbytes as bnb

optimizer = bnb.optim.PagedAdEMAMix8bit(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.999, 0.9999),
    alpha=5.0,
    eps=1e-8,
    weight_decay=1e-2,
    alpha=5.0,
)

8-bit Optimizers Update

The block size for all 8-bit optimizers has been reduced from 2048 to 256 in this release. This is a change from the original implementation proposed in the paper which improves accuracy.

CUDA Graphs support

A fix to enable CUDA Graphs capture of kernel functions was made in #1330. This allows for performance improvements with inference frameworks like vLLM. Thanks @jeejeelee!

Quantization for Embeddings

The trend of LLMs to use larger vocabularies continues. The embeddings can take up a significant portion of a quantized model's footprint. We now have an implementation of Embedding4bit and Embedding8bit thanks to @galqiwi!

Example usage:

import torch
import torch.nn as nn

from bitsandbytes.nn import Embedding4bit

fp16_module = nn.Embedding(128, 64)
quantized_module = Embedding4bit(128, 64)

quantized_module.load_state_dict(fp16_module.state_dict())

quantized_module = quantized_module.to(0)

Continuous Builds

We are now building binary wheels for each change on main. These builds can be used to preview upcoming changes.

🚤 Continuous Build

What's Changed

• Embedding4bit and Embedding8bit implementation by @galqiwi in Embedding4bit and Embedding8bit implementation #1292
• Bugfix: Load correct nocublaslt library variant when BNB_CUDA_VERSION override is set by @matthewdouglas in Bugfix: Load correct nocublaslt library variant when BNB_CUDA_VERSION override is set #1318
• Enable certain CUDA kernels to accept specified cuda stream by @jeejeelee in Enable certain CUDA kernels to accept specified cuda stream #1330
• Initial support for ppc64le by @mgiessing in Initial support for ppc64le #1316
• Cuda source cleanup , refactor and fixes by @abhilash1910 in Cuda source cleanup , refactor and fixes #1328
• Update for VS2022 17.11 compatibility with CUDA < 12.4 by @matthewdouglas in Update for VS2022 17.11 compatibility with CUDA < 12.4 #1341
• Bump the minor-patch group with 3 updates by @dependabot in Bump the minor-patch group with 3 updates #1362
• Update matplotlib requirement from ~=3.9.1 to ~=3.9.2 in the major group by @dependabot in Update matplotlib requirement from ~=3.9.1 to ~=3.9.2 in the major group #1361
• docs: add internal reference to multi-backend guide by @Titus-von-Koeller in docs: add internal reference to multi-backend guide #1352
• Add move_to_device kwarg to the optimizer's load_state_dict by @koute in Add move_to_device kwarg to the optimizer's load_state_dict #1344
• Add AdEMAMix optimizer by @matthewdouglas in Add AdEMAMix optimizer #1360
• Change 8bit optimizer blocksize 2048->256; additional bf16 support by @matthewdouglas in Change 8bit optimizer blocksize 2048->256; additional bf16 support #1365

New Contributors

• @jeejeelee made their first contribution in Enable certain CUDA kernels to accept specified cuda stream #1330
• @mgiessing made their first contribution in Initial support for ppc64le #1316
• @abhilash1910 made their first contribution in Cuda source cleanup , refactor and fixes #1328
• @koute made their first contribution in Add move_to_device kwarg to the optimizer's load_state_dict #1344

Full Changelog: 0.43.3...v0.44.0

---

This discussion was created from the release v0.44.0: New AdEMAMix optimizer, Embeddings quantization, and more!.