mojolearn.lowbit#

Low-bit weight storage for the inference classes, with float32 arithmetic.

A model’s 2-D weight matrices can be stored as bfloat16 bits or as int8 codes with one power-of-two exponent per row, half or a quarter of the float32 size on disk and in memory. An inference class built from packed weights first converts them back to float32 exactly (bfloat16 by widening, int8 as codes * 2**exponent) and then runs its normal float32 path, so it computes the same bits as the float32 block given the converted weights. Packed weights are rounded weights: the model is not the original float32 model, and the arithmetic is not lower precision.

pack converts every 2-D tensor of a weight dict, selected by rank, not by name (so a Mamba-1 A_log of shape (d_inner, 16) is packed too). 1-D and 3-D tensors (norm weights, biases, convolution taps, D) stay float32. Use pack_one to pack selected tensors by name. unpack (also materialize) returns the float32 dict and its format; unpack(pack(unpack(p)[0], fmt)) equals unpack(p).

Conversion runs on the GPU on a GPU install and on the host binding on a CPU-only install; an install without the linalg extension raises an error. NumPy is not required: packed tensors are mojolearn.Array objects and inputs can be any buffer.

Examples

>>> import mojolearn
>>> packed = mojolearn.lowbit.pack(weights, "bfloat16")        # or "int8"
>>> blk = mojolearn.TransformerBlock(packed, n_heads=8)
>>> f32, fmt = mojolearn.lowbit.unpack(packed)

Functions

pack(weights, fmt)

Pack every 2-D tensor of a weight dict to fmt; other tensors pass through unchanged.

unpack(weights[, what])

Convert every packed tensor of a weight dict to float32, exactly.

materialize(weights[, what])

Convert every packed tensor of a weight dict to float32, exactly.

format_of(weights)

The format shared by the packed tensors of weights, or 'float32' when none is packed.

is_packed(value)

Whether value is a packed weight (BF16Weight or Int8Weight).

pack_one(value, fmt[, name])

Pack one 2-D float32 tensor to fmt.

materialize_one(value[, name])

A packed tensor converted to float32 exactly, or a float32 tensor returned as is.

widen_bf16(bits)

Convert bfloat16 bits of any rank (a uint16 buffer) to float32, exactly.

Classes

BF16Weight(bits)

A 2-D weight stored as bfloat16 bits.

Int8Weight(codes, exponents)

A 2-D weight stored as int8 codes with one int32 exponent per row: row = codes * 2**exponent.
class mojolearn.lowbit.BF16Weight(bits)[source]#

Bases: object

A 2-D weight stored as bfloat16 bits.

Parameters:

bits (2-D uint16 buffer) – bfloat16 bits, for example from mojolearn.linalg.to_bf16. Anything else raises TypeError.

Variables:
class mojolearn.lowbit.Int8Weight(codes, exponents)[source]#

Bases: object

A 2-D weight stored as int8 codes with one int32 exponent per row: row = codes * 2**exponent.

Parameters:

  • codes (2-D int8 buffer)

  • exponents (int32 buffer of shape (rows,)) – Wrong dtypes or shapes raise TypeError.

Variables:
mojolearn.lowbit.pack(weights, fmt)[source]#

Pack every 2-D tensor of a weight dict to fmt; other tensors pass through unchanged.

Parameters:

  • weights (dict) – Tensors by name. Anything but a dict raises TypeError.

  • fmt ({‘float32’, ‘bfloat16’, ‘int8’}) – Target format; unknown formats raise ValueError.

Returns:

dict – A new dict.

mojolearn.lowbit.unpack(weights, what='weights')[source]#

Convert every packed tensor of a weight dict to float32, exactly.

Parameters:

  • weights (dict) – Tensors by name; other tensors pass through unchanged.

  • what (str, default ‘weights’) – Name used in error messages.

Returns:

(dict, str) – The float32 dict and the format the weights were stored in.

mojolearn.lowbit.materialize(weights, what='weights')#

Convert every packed tensor of a weight dict to float32, exactly.

Parameters:

  • weights (dict) – Tensors by name; other tensors pass through unchanged.

  • what (str, default ‘weights’) – Name used in error messages.

Returns:

(dict, str) – The float32 dict and the format the weights were stored in.

mojolearn.lowbit.format_of(weights)[source]#

The format shared by the packed tensors of weights, or 'float32' when none is packed.

Raises ValueError when the dict mixes formats.

mojolearn.lowbit.is_packed(value)[source]#

Whether value is a packed weight (BF16Weight or Int8Weight).

mojolearn.lowbit.pack_one(value, fmt, name='weight')[source]#

Pack one 2-D float32 tensor to fmt.

Parameters:

  • value (2-D float32 buffer)

  • fmt ({‘float32’, ‘bfloat16’, ‘int8’}) – 'float32' returns a float32 copy. Unknown formats and non-2-D tensors raise ValueError.

  • name (str, default ‘weight’) – Name used in error messages.

Returns:

BF16Weight, Int8Weight or mojolearn.Array

mojolearn.lowbit.materialize_one(value, name='weight')[source]#

A packed tensor converted to float32 exactly, or a float32 tensor returned as is.

mojolearn.lowbit.widen_bf16(bits)[source]#

Convert bfloat16 bits of any rank (a uint16 buffer) to float32, exactly.