fireworks-prequant 0.1.0

Prequantization utilities (MXFP8 / NVFP4) for Fireworks AI hot-load / incremental snapshot pipelines

fireworks-prequant

Prequantization utilities used by the Fireworks AI serving/training stack to produce already-quantized model checkpoints (MXFP8 / NVFP4) so that both the upload and the deployment-side download are smaller, and serving loads the packed weights directly instead of converting bf16 on the fly.

The quantization is byte-identical to the Fireworks runtime's on-the-fly --convert-to-precision (it uses the same flashinfer MXFP8 kernel), so a prequantized checkpoint serves identically to the bf16 checkpoint + on-the-fly conversion.

Install

pip install fireworks-prequant[mxfp8]   # torch + safetensors + flashinfer

torch and flashinfer are already present in the Fireworks trainer/learner image, so this is effectively just adding the pure-Python package there.

Quick start

In-memory (the hook the stage-sharded hot-load weight publisher calls on each PP-stage owner's gathered HF state):

from fireworks_prequant import quantize_state_dict

# hf_state: dict[str, torch.Tensor] of one PP stage (bf16 on CPU)
packed = quantize_state_dict(hf_state, fmt="mxfp8", device="cuda")
# expert weights -> {..}.weight (float8_e4m3fn) + {..}.weight_scale_inv (uint8 E8M0)
# non-expert tensors pass through unchanged

Offline (directory of safetensors shards):

from fireworks_prequant import quantize_checkpoint_dir

quantize_checkpoint_dir(
    "/path/to/bf16-checkpoint",
    "/path/to/mxfp8-checkpoint",
    fmt="mxfp8",
    num_workers=8,
)  # rebuilds the weight index + stamps quantization_config into config.json

Formats

fmt	Target	Serving precision
mxfp8	1x32 block FP8, E8M0 scales	moe=mxfp8
nvfp4	packed FP4 + FP8 block scales	FP4

Only the routed MoE expert projections (model.layers.*.mlp.experts.*.{gate,up,down}_proj.weight, prefix-tolerant) are quantized; dense/attention/gate/shared-expert weights pass through in their original dtype — matching the runtime precision policy.

License

MIT — see LICENSE.