Skip to content

convert: support compressed-tensors mixed-precision NVFP4 checkpoints #195

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
danielhanchen wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

convert: support compressed-tensors mixed-precision NVFP4 checkpoints #195

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 31 additions & 5 deletions conversion/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -489,7 +489,7 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T
quant_format == "nvfp4-pack-quantized"
or quant_format == "mixed-precision"
and bool(groups)
and all(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict))
and any(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict))
)

if len(groups) > 1 and not nvfp4_compressed_tensors:
Expand Down Expand Up @@ -538,8 +538,27 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T
if (base_name + "_zero_point") in self.model_tensors:
tensors_to_remove.append(base_name + "_zero_point")
elif nvfp4_compressed_tensors:
# Don't error from compressed-tensors, we'll handle them in _generate_nvfp4_tensors
pass
# NVFP4 tensors were already repacked by _generate_nvfp4_tensors and removed
# from model_tensors. For a "mixed-precision" checkpoint whatever weight_scale
# entries are left belong to the non-NVFP4 config group (FP8 per-channel);
# dequantize them exactly like the float-quantized branch above.
for name in self.model_tensors.keys():
if name.endswith(".weight_scale"):
weight_name = name.removesuffix("_scale")
if weight_name not in self.model_tensors:
tensors_to_remove.append(name)
continue
w = self.model_tensors[weight_name]
s = self.model_tensors[name]
is_fp8 = False
if self._fp8_as_q8:
is_fp8 = w().dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), None)
tensors_to_remove.append(name)
if is_fp8:
self._fp8_dequantized.add(weight_name)
elif name.endswith((".input_scale", ".k_scale", ".v_scale", ".weight_scale_2")):
tensors_to_remove.append(name)
else:
raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported")
elif quant_method == "modelopt":
Expand Down Expand Up @@ -751,9 +770,16 @@ def _generate_nvfp4_tensors(self):
weight = LazyTorchTensor.to_eager(self.model_tensors[name]())
scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())

# Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales)
# Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales).
# In a compressed-tensors "mixed-precision" checkpoint the FP8 group also has a
# 2D weight_scale of shape [out, 1], so shape alone is not enough: an NVFP4
# tensor is nibble-packed uint8 with an E4M3 scale, one per 16 values.
if scale.ndim < 2:
continue
if weight.dtype != torch.uint8 or scale.dtype != torch.float8_e4m3fn:
continue
if scale.shape[-1] * 16 != weight.shape[-1] * 2:
continue

scale2 = LazyTorchTensor.to_eager(self.model_tensors.get(scale2_name, lambda: torch.tensor(1.0))())
input_scale = LazyTorchTensor.to_eager(self.model_tensors.get(input_scale_name, lambda: torch.tensor(1.0))())
Expand Down Expand Up @@ -858,7 +884,7 @@ def prepare_tensors(self):
quant_format == "nvfp4-pack-quantized"
or quant_format == "mixed-precision"
and bool(quant_groups)
and all(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict))
and any(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict))
)
if quant_algo != "NVFP4":
if nvfp4_compressed_tensors:
Expand Down