feat(model): detect header metadata and nvfp4 in the model probe

detect_quant now reads _quantization_metadata as the authoritative quant
source when present, and resolves marker-file formats from the marked
layers' stored weight dtypes instead of file-wide dtype voting, which
mislabeled fp8 files carrying extra uint8 tensors and had no nvfp4
mapping at all. Schema bump so cached probe entries refresh.
This commit is contained in:
CalamitousFelicitousness
2026-07-11 17:32:46 +01:00
parent 0997526394
commit 5e1a0719ab
2 changed files with 97 additions and 10 deletions
+73
View File
@@ -637,6 +637,68 @@ def test_detect_comfy_nvfp4_convrot_raises():
assert 'convrot' in str(e)
# ============================================================
# model_probe quant detection (header-only, mirrors the loader's coverage)
# ============================================================
def test_probe_detects_header_quant_metadata():
import json
from modules import model_probe
header = {
'__metadata__': {'_quantization_metadata': json.dumps({'layers': {'blocks.0.wq': {'format': 'float8_e4m3fn'}}})},
'blocks.0.wq.weight': {'dtype': 'F8_E4M3', 'shape': [4, 4]},
'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
}
quant = model_probe.analyze_header(header)['quant']
assert quant == {'scheme': 'comfy_quant', 'format': 'float8_e4m3fn', 'marked_layers': 1, 'source': 'header'}
def test_probe_resolves_nvfp4_from_weight_dtype():
from modules import model_probe
header = {
'model.diffusion_model.blocks.0.wq.weight': {'dtype': 'U8', 'shape': [32, 16]},
'model.diffusion_model.blocks.0.wq.weight_scale': {'dtype': 'F8_E4M3', 'shape': [128, 4]},
'model.diffusion_model.blocks.0.wq.weight_scale_2': {'dtype': 'F32', 'shape': []},
'model.diffusion_model.blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [60]},
}
quant = model_probe.analyze_header(header)['quant']
assert quant['format'] == 'nvfp4'
assert quant['source'] == 'weight-dtype'
def test_probe_fp8_markers_despite_u8_noise():
"""The marked layer's own weight dtype decides the format; unrelated U8
tensors must not outvote it."""
from modules import model_probe
header = {
'blocks.0.wq.weight': {'dtype': 'F8_E4M3', 'shape': [8, 8]},
'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
'blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [27]},
'extra.blob1': {'dtype': 'U8', 'shape': [64]},
'extra.blob2': {'dtype': 'U8', 'shape': [64]},
}
quant = model_probe.analyze_header(header)['quant']
assert quant['format'] == 'float8_e4m3fn'
def test_probe_int8_markers():
from modules import model_probe
header = {
'blocks.0.wq.weight': {'dtype': 'I8', 'shape': [8, 8]},
'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
'blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [29]},
}
quant = model_probe.analyze_header(header)['quant']
assert quant['format'] == 'int8_tensorwise'
assert quant['marked_layers'] == 1
def test_probe_plain_file_no_quant():
from modules import model_probe
header = {'blocks.0.wq.weight': {'dtype': 'BF16', 'shape': [8, 8]}}
assert model_probe.analyze_header(header)['quant']['scheme'] is None
# ============================================================
# is_noop_converter
# ============================================================
@@ -2184,6 +2246,17 @@ def run_all():
]:
run_test(cat, fn)
log.warning('=== model_probe quant detection ===')
cat = category('probe')
for fn in [
test_probe_detects_header_quant_metadata,
test_probe_resolves_nvfp4_from_weight_dtype,
test_probe_fp8_markers_despite_u8_noise,
test_probe_int8_markers,
test_probe_plain_file_no_quant,
]:
run_test(cat, fn)
log.warning('=== noop_converter detection ===')
cat = category('noop')
for fn in [