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 kernels), so a prequantized checkpoint serves identically to the bf16 checkpoint + on-the-fly conversion.

Install

pip install fireworks-prequant          # torch + safetensors — enough for mxfp8_bw128

That plain install fully supports the mxfp8_bw128 (bw128fp8_deepgemm) format, which is pure torch. The other two formats each need one extra GPU kernel:

pip install fireworks-prequant[mxfp8]   # + flashinfer  (for fmt="mxfp8")
pip install fireworks-prequant[nvfp4]   # + triton      (for fmt="nvfp4")
pip install fireworks-prequant[all]     # + both        (every format)

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

mxfp8_bw128 emits the deepseek/DeepGEMM blockwise-fp8 layout ([128,128] fp8 e4m3 weights + [ceil(M/128), ceil(K/128)] fp32 block scales) — the same layout as a natively-fp8 [128,128] HF checkpoint — and is byte-identical to the runtime fireworks::fp8_blockwise_weight_quant conversion.

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.