Optimizers API

AdamW8bit

Custom 8-bit AdamW optimizer with block-wise quantization.

from legionheto.optim import AdamW8bit

optimizer = AdamW8bit(
    model.parameters(),
    lr=2e-4,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
    block_size=2048,
)

Quantization Functions

quantize_blockwise

Quantize tensor with block-wise scaling.

from legionheto.optim import quantize_blockwise

quantized, norm, mins, maxs = quantize_blockwise(
    tensor,
    block_size=2048,
)

dequantize_blockwise

Dequantize block-wise quantized tensor.

from legionheto.optim import dequantize_blockwise

dequantized = dequantize_blockwise(
    quantized,
    norm,
    mins,
    maxs,
    shape,
    block_size=2048,
)