Metadata-Version: 2.4
Name: fireworks-prequant
Version: 0.2.0
Summary: Prequantization utilities (MXFP8 / bw128fp8_deepgemm / NVFP4) for Fireworks AI hot-load / incremental snapshot pipelines
Author-email: Fireworks AI <info@fireworks.ai>
License: The MIT License
        
        Copyright (c) Fireworks (https://fireworks.ai)
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in
        all copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
        THE SOFTWARE.
Project-URL: Homepage, https://fireworks.ai
Keywords: quantization,mxfp8,bw128fp8,deepgemm,nvfp4,safetensors,model-weights,fireworks
Classifier: Development Status :: 4 - Beta
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch<3,>=2.1
Requires-Dist: safetensors<1,>=0.4
Provides-Extra: mxfp8
Requires-Dist: flashinfer-python>=0.2; extra == "mxfp8"
Provides-Extra: nvfp4
Requires-Dist: triton>=3.0; extra == "nvfp4"
Provides-Extra: all
Requires-Dist: flashinfer-python>=0.2; extra == "all"
Requires-Dist: triton>=3.0; extra == "all"
Dynamic: license-file

# 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

```bash
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:

```bash
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):

```python
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):

```python
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       | Extra needed |
| ------------- | ------------------------------------- | ----------------------- | ------------ |
| `mxfp8_bw128` | 128x128 block FP8, fp32 scales        | `moe=bw128fp8_deepgemm` | none (pure torch) |
| `mxfp8`       | 1x32 block FP8, E8M0 scales           | `moe=mxfp8`             | `[mxfp8]` (flashinfer) |
| `nvfp4`       | packed FP4 + FP8 block scales         | FP4                     | `[nvfp4]` (triton) |

`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](LICENSE).
