| |
|
|
| |
|
|
| """Dequantize the Ideogram 4.0 weight-only FP8 checkpoints to bf16. |
| |
| The fp8 checkpoints store each quantized Linear as an e4m3 ``<name>.weight`` |
| (float8_e4m3fn, shape [out, in]) plus a per-output-row ``<name>.weight_scale`` |
| (float32, shape [out]); the reference dequantizes as |
| ``w ≈ weight_fp8.to(dtype) * weight_scale[:, None]`` (see the reference |
| quantized_loading.py). Non-quantized tensors are passed through (cast to bf16 if |
| floating point). This yields a plain bf16 state dict that loads into the standard |
| (un-quantized) reference modules, to serve as the bringup golden. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| _SCALE_SUFFIX = ".weight_scale" |
|
|
|
|
| def dequant_fp8_state_dict( |
| sd: dict[str, torch.Tensor], *, dtype: torch.dtype = torch.bfloat16 |
| ) -> dict[str, torch.Tensor]: |
| out: dict[str, torch.Tensor] = {} |
| consumed: set[str] = set() |
|
|
| for key in sd: |
| if not key.endswith(_SCALE_SUFFIX): |
| continue |
| base = key[: -len(_SCALE_SUFFIX)] |
| weight_key = base + ".weight" |
| w = sd[weight_key].to(torch.float32) |
| scale = sd[key].to(torch.float32) |
| out[weight_key] = (w * scale[:, None]).to(dtype) |
| consumed.add(weight_key) |
| consumed.add(key) |
|
|
| for key, val in sd.items(): |
| if key in consumed: |
| continue |
| out[key] = val.to(dtype) if val.is_floating_point() else val |
|
|
| return out |
|
|
