diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index 78818e809..3dd63ecc9 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/python/freetoken/layers/quantization/configs/base.py @@ -40,10 +40,15 @@ def quantization_config_of(hf_config: Any) -> dict[str, Any] | None: @dataclass(frozen=True) class Stored: - """One checkpoint tensor behind a role: its suffix, and whether it holds the quant-side scale whose reciprocal the layer wants.""" + """One checkpoint tensor behind a role: its suffix, other suffixes exports of the same dialect use for it, and whether it holds the quant-side scale whose reciprocal the layer wants.""" name: str reciprocal: bool = False + aliases: tuple[str, ...] = () + + @property + def names(self) -> tuple[str, ...]: + return (self.name, *self.aliases) class QuantConfig(ABC): diff --git a/python/freetoken/layers/quantization/configs/modelopt.py b/python/freetoken/layers/quantization/configs/modelopt.py index 6183dfc96..3547f186d 100644 --- a/python/freetoken/layers/quantization/configs/modelopt.py +++ b/python/freetoken/layers/quantization/configs/modelopt.py @@ -27,7 +27,8 @@ class ModelOptConfig(QuantConfig): STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = { QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"}, QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, - QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"}, + # same e8m0 codes under either name: MiniMax-M3 exports weight_scale_inv, Qwen3.8-Flash-Next weight_scale + QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": Stored("weight_scale_inv", aliases=("weight_scale",))}, QuantKind.NVFP4: {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2", "input_scale": "input_scale"}, } diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index bf3cecf05..e7c4676dd 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -149,7 +149,7 @@ def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] if module.rpartition(".")[2] not in _LINEAR_LEAVES: return None stored = self.stored(module) - roles = {"weight": "weight"} if stored is None else {e.name: r for r, e in self.quant.storage(stored).items()} + roles = {"weight": "weight"} if stored is None else {n: r for r, e in self.quant.storage(stored).items() for n in e.names} role = roles.get(suffix) if role is None: if suffix in _DROPPED_SUFFIXES: diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index bb71da2ff..536853be3 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -1,4 +1,4 @@ -"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant. +"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 or bf16 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant. Hash windows are pure functions of ``req.input_ids`` + ``device_len`` (prefix hits, restores and COW forks need no bookkeeping); the decode input token lives device-side under overlap scheduling and is read back here. """ @@ -22,7 +22,7 @@ from .weight import ( _PLE_SCALE_SUFFIX, _PLE_SHARD_RE, - _PLE_ST_DTYPE, + _PLE_ST_DTYPES, _ple_table_files, _safetensors_header, ) @@ -50,9 +50,10 @@ class PleRowSource: row_bytes: int row_stride: int scale: float - # packed ggml rows (a GGUF table) of ``row_values`` values each; None for one fp8 byte per value + # packed ggml rows (a GGUF table) of ``row_values`` values each; None for float rows of ``row_dtype`` ggml_type: int | None = None row_values: int | None = None + row_dtype: torch.dtype = torch.float8_e4m3fn @property def total_rows(self) -> int: @@ -62,6 +63,7 @@ def total_rows(self) -> int: def source_from_safetensors(folder: str) -> PleRowSource: """Map the checkpoint's ``ngram_embedding.shard_`` tensors in place: one extent per shard, no copy.""" rows = cols = 0 + dtype: torch.dtype | None = None scale: torch.Tensor | None = None paths: list[str] = [] path_idx: dict[str, int] = {} @@ -78,8 +80,10 @@ def source_from_safetensors(folder: str) -> PleRowSource: match = _PLE_SHARD_RE.search(key) if match is None: continue - if meta["dtype"] != _PLE_ST_DTYPE: - raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}") + shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) + if shard_dtype is None or dtype not in (None, shard_dtype): + raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") + dtype = shard_dtype if rows and tuple(meta["shape"]) != (rows, cols): raise ValueError(f"PLE shard {key} is {meta['shape']}, expected {[rows, cols]}") rows, cols = meta["shape"] @@ -93,9 +97,13 @@ def source_from_safetensors(folder: str) -> PleRowSource: if sorted(shards) != list(range(len(shards))) or not shards: raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(shards)[:8]}") if scale is None: - raise ValueError("PLE table has no weight_scale") + if dtype is torch.float8_e4m3fn: + raise ValueError("PLE table has no weight_scale") + scale = torch.tensor(1.0) order = [shards[i] for i in range(len(shards))] - return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, cols, cols, float(scale)) + row_bytes = cols * dtype.itemsize + return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, row_bytes, row_bytes, float(scale), + row_values=cols, row_dtype=dtype) def source_from_gguf(model_path: str) -> PleRowSource: @@ -137,11 +145,11 @@ def __init__( from freetoken.layers.quantization.ggml_types import float_dtype self.num_rows = source.total_rows - self.head_dim = source.row_values or source.row_bytes # fp8: one byte per element + self.head_dim = source.row_values or source.row_bytes self._ggml_type = source.ggml_type self._row_bytes = source.row_bytes # float rows (fp8, or F32 / F16 / BF16 in a GGUF) are only cast; quantized ggml rows go through ggml - self._row_dtype = torch.float8_e4m3fn if source.ggml_type is None else float_dtype(source.ggml_type) + self._row_dtype = source.row_dtype if source.ggml_type is None else float_dtype(source.ggml_type) self.dtype = dtype self.heads = int(hash_constants["num_ngram_heads"]) self.scale = source.scale diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 691651d4c..6c62d141c 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -1,10 +1,10 @@ -"""Qwen3.8-Flash-Next checkpoint reader (the NVFP4 and the official block-fp8 releases). +"""Qwen3.8-Flash-Next checkpoint reader (bf16, block-fp8, ModelOpt and llm-compressor exports). Three separate paths, because the checkpoint's three weight classes live in different places: -* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseFuser``. -* :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`. -* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named, for the offload cache's expert reader. +* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseReader``. +* :func:`load_ple_table` -- the n-gram table (47.7 GiB fp8 or 95.4 GiB bf16), 128 checkpoint shards concatenated into one pinned :class:`HostBank`. +* :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named under the checkpoint's dialect, for the offload cache's expert reader. Dropped: ``mtp.*`` (speculative head, including its stacked ``mtp.layers.0.mlp.experts.*``); ``model.visual.*`` is kept only when the model built the tower. """ @@ -28,29 +28,35 @@ from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantConfig, QuantKind, QuantScheme, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into from freetoken.utils import cached_load_hf_config, download_hf_weight from freetoken.utils.progress import byte_bar from tqdm import tqdm -# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's -# stacked ``mtp.layers.N.mlp.experts.*`` tensors. -_EXPERT_KEY_RE = re.compile( +# Routed NVFP4 experts: per-expert, un-fused. Matched against the RAW weight_map key in nvfp4_banks. +# The ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*`` tensors. +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.8-Flash-Next NVFP4 experts", -) -# Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# role -> the expert bank reader's canonical (ModelOpt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} + +# leaves the model builds as quant-aware Linear layers: their tensors are read under the QuantConfig +_LINEAR_LEAVES = frozenset({ + "q_proj", "k_proj", "v_proj", "o_proj", "index_qk_proj", + "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj", + "gate_proj", "up_proj", "down_proj", + "input_mix_weight_down", "block_inject_weight", "input_mix_weight_up", + "key_proj", "value_proj", "lm_head", +}) +# activation scales of modules whose scheme carries no input_scale role +_DROPPED_SUFFIXES = frozenset({"input_scale", "input_global_scale"}) +_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn, "e2m1": torch.uint8} +_QUANT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2, torch.uint8, torch.int8) # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." @@ -79,9 +85,6 @@ # The top-level hyper_connection_mixer has no injection and never fuses. _PAD_TO = {"input_mix_weight_down_block_inject": 16} _HC_WITH_INJECT = (".attn_hyper_connection", ".mlp_hyper_connection") -_KIND_SUFFIXES = (".weight_scale_inv", ".weight") -_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) -_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn} def _rename(raw_name: str) -> str | None: @@ -92,100 +95,157 @@ def _rename(raw_name: str) -> str | None: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks - if raw_name.endswith(_SCALE_SUFFIXES): + # static KV-cache scales of the quantizers; the KV cache runs in the engine's dtype + if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): return None return rename_vl_prefix(raw_name) -def _split_kind(name: str) -> tuple[str, str]: - """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale.""" - for suffix in _KIND_SUFFIXES: - if name.endswith(suffix): - return name[: -len(suffix)], suffix - return name, "" +def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + """Per-tensor scalar or per-channel ``[rows, 1]`` scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows.""" + flat = scale.reshape(-1).to(torch.float32) + if flat.numel() == 1: + return flat.expand(rows).contiguous() + if flat.numel() != rows: + raise ValueError(f"scale has {flat.numel()} elements for a weight with {rows} output rows (shape {tuple(scale.shape)}); expected 1 or {rows}") + return flat.contiguous() -class _DenseFuser: - """Concatenates checkpoint projection parts into the model's merged buffers, per kind (weight / block scale). +class _DenseReader: + """Routes each Linear tensor to the buffer its module's scheme declares; packed projections are concatenated per role once every part is in. - The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from. + The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and, through its dialect's storage names, which checkpoint tensor fills which role. """ - def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: + def __init__(self, quant: QuantConfig | None, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: self.quant = quant self.groups = {fused: parts for fused, parts in packed if fused != "experts"} # experts: bank reader self.by_part: dict[str, list[tuple[str, int]]] = {} for fused, parts in self.groups.items(): for idx, part in enumerate(parts): self.by_part.setdefault(part, []).append((fused, idx)) - self.buf: dict[tuple[str, str], dict[int, torch.Tensor]] = {} + # target module -> (part count, {part: {role: tensor}}, {part: the roles its module stores}, scheme) + self.pending: dict[str, tuple[int, dict[int, dict[str, torch.Tensor]], dict[int, set[str]], QuantScheme | None]] = {} - def scheme(self, module: str): + def scheme(self, module: str) -> QuantScheme | None: return None if self.quant is None else self.quant.scheme_for(module) - def _target(self, parent: str, leaf: str) -> tuple[str, int] | None: + def target(self, module: str) -> tuple[str, int, int]: + """``(fused module, part index, part count)``; a standalone linear is its own single-part target.""" + parent, _, leaf = module.rpartition(".") candidates = self.by_part.get(leaf) if not candidates: - return None + return module, 0, 1 if len(candidates) > 1: - # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py + # GDN: quantized checkpoints split qkv|z from b|a; same test as gdn.py split = self.scheme(f"{parent}.in_proj_qkvz") is not None keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} candidates = [c for c in candidates if c[0] in keep] - if not candidates: - raise ValueError(f"{parent}.{leaf}: no merged projection for the {'split' if split else 'fused'} GDN layout") fused, idx = candidates[0] if fused in _PAD_TO and not parent.endswith(_HC_WITH_INJECT): - return None - return f"{parent}.{fused}", idx + return module, 0, 1 + return f"{parent}.{fused}", idx, len(self.groups[fused]) - def check(self, module: str, name: str, tensor: torch.Tensor) -> None: - """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from.""" + def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: + """Take one tensor; the emitted ``[(name, tensor)]`` once its target module is complete, ``[]`` before, None if ``name`` is not a Linear's tensor.""" + module, _, suffix = name.rpartition(".") + if module.rpartition(".")[2] not in _LINEAR_LEAVES: + return None scheme = self.scheme(module) - if name.endswith(".weight_scale_inv"): - if scheme is None or not scheme.has("weight_scale_inv"): - raise ValueError(f"{name}: {module} has no block scale in the checkpoint's quant config ({scheme})") - return - is_fp8 = tensor.dtype in _FP8_DTYPES + roles = {"weight": "weight"} if scheme is None else {n: r for r, e in self.quant.storage(scheme).items() for n in e.names} + role = roles.get(suffix) + if role is None: + if suffix in _DROPPED_SUFFIXES: + return [] + raise ValueError(f"{name}: the checkpoint's quant config declares {module} {scheme or 'unquantized'}, stored as {sorted(roles)}") + if role == "weight": + self._check_weight(name, module, scheme, tensor) + target, idx, count = self.target(module) + _, parts, expected, _ = self.pending.setdefault(target, (count, {}, {}, scheme)) + part = parts.setdefault(idx, {}) + if role in part: + raise ValueError(f"{name}: {module} already has its {role} from another checkpoint tensor") + part[role] = tensor + expected[idx] = set(roles.values()) + if scheme is not None and set(part) == expected[idx]: + parts[idx] = self._finish(module, scheme, part) + if len(parts) < count or any(set(parts[i]) != expected[i] for i in parts): + return [] + del self.pending[target] + return self._emit(target, [parts[i] for i in range(count)], scheme) + + def missing(self) -> list[str]: + """One line per incomplete module: the roles its parts still lack.""" + lines = [] + for target, (count, parts, expected, _scheme) in sorted(self.pending.items()): + lacking = sorted(set().union(*(expected[i] - set(parts[i]) for i in parts))) + if len(parts) < count: + lacking.append(f"{count - len(parts)} of {count} fused parts") + lines.append(f"{target}: missing {lacking}") + return lines + + @staticmethod + def _check_weight(name: str, module: str, scheme: QuantScheme | None, tensor: torch.Tensor) -> None: if scheme is None: - if is_fp8: + if tensor.dtype in _QUANT_DTYPES: raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized") return - expected = _ELEM_DTYPES.get(scheme.weight.elem) - if expected is not None and tensor.dtype is not expected: + elem = _ELEM_DTYPES.get(scheme.weight.elem) + if elem is None: + raise ValueError(f"{name}: qwen4_exp has no dense reader for {scheme}") + if tensor.dtype is not elem: raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} {scheme}") - rows, cols = (scheme.weight.group or (1, 1)) - if rows > 1 and tensor.shape[0] % rows or cols > 1 and tensor.shape[1] % cols: - raise ValueError(f"{name}: {tuple(tensor.shape)} is not a multiple of the {rows}x{cols} scale block of {module}") - - def check_unfused(self, name: str, tensor: torch.Tensor) -> None: - module, kind = _split_kind(name) - if kind == ".weight_scale_inv" or (kind == ".weight" and tensor.dtype in _FP8_DTYPES): - self.check(module, name, tensor) - - def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None: - """Buffer a part; return the merged ``[(name, tensor)]`` once its kind is complete, ``[]`` while incomplete, ``None`` if ``name`` is not a part.""" - module, kind = _split_kind(name) - if not kind: - return None - parent, _, leaf = module.rpartition(".") - hit = self._target(parent, leaf) - if hit is None: - return None - fused, idx = hit - self.check(fused, name, tensor) - slots = self.buf.setdefault((fused, kind), {}) - slots[idx] = tensor - parts = self.groups[fused.rpartition(".")[2]] - if len(slots) < len(parts): - return [] - del self.buf[(fused, kind)] - rows = [slots[i] for i in range(len(parts))] - pad_to = _PAD_TO.get(fused.rpartition(".")[2], 0) if kind == ".weight" else 0 - pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0 - if pad: - rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device)) - return [(fused + kind, torch.cat(rows, dim=0))] + rows, cols = tensor.shape[0], tensor.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1) + block_rows, block_cols = scheme.weight.group or (1, 1) + if block_rows > 1 and rows % block_rows or block_cols > 1 and cols % block_cols: + raise ValueError(f"{name}: {rows}x{cols} is not a multiple of the {block_rows}x{block_cols} scale block of {module}") + + def _finish(self, module: str, scheme: QuantScheme, part: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Validate one complete part against ``scheme`` and put its scales in the layer's form.""" + storage = self.quant.storage(scheme) + part = {role: 1.0 / t.to(torch.float32) if storage[role].reciprocal else t for role, t in part.items()} + weight = part["weight"] + rows, cols = weight.shape[0], weight.shape[1] * (2 if scheme.weight.elem == "e2m1" else 1) + block_rows, block_cols = scheme.weight.group or (1, 1) + scale_role = "weight_scale_inv" if "weight_scale_inv" in part else "weight_scale" + out = dict(part) + if block_cols < 0: + try: + out[scale_role] = _per_row_scale(part[scale_role], rows) + except ValueError as exc: + raise ValueError(f"{module}.{scale_role}: {exc}") from None + else: + want = (rows // block_rows, cols // block_cols) + if tuple(part[scale_role].shape) != want: + raise ValueError(f"{module}: {scale_role} is {tuple(part[scale_role].shape)}, expected {want} for {scheme}") + if scheme.weight.scale == "e4m3" and part[scale_role].dtype is not torch.float8_e4m3fn: + raise ValueError(f"{module}: {scale_role} is {part[scale_role].dtype} but {scheme} stores e4m3 scales") + if "weight_global" in part: + g = part["weight_global"].reshape(-1).to(torch.float32) + if g.numel() != 1: + raise ValueError(f"{module}: weight_global has {g.numel()} elements, expected one per-tensor scale") + out["weight_global"] = g.expand(rows).contiguous() + if "input_scale" in part: + out["input_scale"] = part["input_scale"].reshape(()).to(torch.float32) + return out + + @staticmethod + def _emit(target: str, parts: list[dict[str, torch.Tensor]], scheme: QuantScheme | None) -> list[tuple[str, torch.Tensor]]: + pad_to = _PAD_TO.get(target.rpartition(".")[2], 0) + if pad_to and scheme is not None: + raise ValueError(f"{target}: qwen4_exp pads this fusion to {pad_to} rows and reads it bf16 only, the checkpoint's quant config declares {scheme}") + out = [] + for role in parts[0]: + tensors = [part[role] for part in parts] + if role == "input_scale": + # fused parts read the same activation, so one range covers them: max is exact then and safe if they drift + out.append((f"{target}.{role}", torch.stack(tensors).max())) + continue + pad = (-sum(t.shape[0] for t in tensors)) % pad_to if pad_to else 0 + if pad: + tensors.append(torch.zeros(pad, *tensors[0].shape[1:], dtype=tensors[0].dtype, device=tensors[0].device)) + out.append((f"{target}.{role}", tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0))) + return out def iter_weights( @@ -199,8 +259,8 @@ def iter_weights( """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers. Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts. - A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections. - Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. + Each Linear is read under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, per-channel / per-tensor fp8, block-fp8, MXFP8 and NVFP4 projections in any mix land as the model's state dict. + Fusions, per role: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader. """ if get_tp_info().size > 1: @@ -210,7 +270,7 @@ def iter_weights( hf_config = cached_load_hf_config(model_path) spec = get_model_spec(hf_config.architectures[0]) - fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping) + reader = _DenseReader(get_quant_config(), spec.packed_modules_mapping) for file in tqdm( iter_weight_files(model_path), desc="Loading weights", @@ -224,14 +284,16 @@ def iter_weights( if not include_vision and name.startswith(VISION_KEY_PREFIXES): continue tensor = f.get_tensor(raw_name) - fused = fuser.fuse(name, tensor) - if fused is None: - fuser.check_unfused(name, tensor) + emitted = None if name.startswith(VISION_KEY_PREFIXES) else reader.add(name, tensor) + if emitted is None: yield name, tensor else: - yield from fused + yield from emitted - assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}" + if reader.pending: + lines = reader.missing() + shown = "\n ".join(lines[:8]) + (f"\n ... {len(lines) - 8} more" if len(lines) > 8 else "") + raise ValueError(f"checkpoint is missing tensors the quant config declares for {len(lines)} modules:\n {shown}") def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]: @@ -251,18 +313,18 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple @dataclass(frozen=True) class PleTable: - """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor FP8 scale.""" + """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor scale (1 for a bf16 table).""" bank: HostBank weight_scale: torch.Tensor # scalar, checkpoint dtype (bf16) @property def tensor(self) -> torch.Tensor: - """``[total_rows, ngram_head_dim]`` float8_e4m3fn view of the bank.""" + """``[total_rows, ngram_head_dim]`` view of the bank in the stored dtype.""" return self.bank.tensor -_PLE_ST_DTYPE = "F8_E4M3" +_PLE_ST_DTYPES = {"F8_E4M3": torch.float8_e4m3fn, "BF16": torch.bfloat16} def _safetensors_header(path: str) -> tuple[dict, int]: @@ -322,12 +384,14 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, The checkpoint splits the table into ``split_ngram_parts`` equal row blocks named by shard index and scattered over the ``model-plefp8-*`` shards in header (lexicographic) order, so the bank is filled shard by shard at ``shard_index * rows_per_shard``. Each read is O_DIRECT: the - table is ~47.7 GiB and must not also sit in the page cache while the bank holds the same bytes. + table is ~47.7 GiB (fp8) and must not also sit in the page cache while the bank holds the same bytes. + An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled. """ folder = download_hf_weight(model_path) parts: dict[int, tuple[str, int, int]] = {} # shard index -> (path, file offset, bytes) scale: torch.Tensor | None = None rows = cols = 0 + dtype: torch.dtype | None = None for path in _ple_table_files(folder): header, base = _safetensors_header(path) for key, meta in header.items(): @@ -340,8 +404,10 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, match = _PLE_SHARD_RE.search(key) if match is None: continue - if meta["dtype"] != _PLE_ST_DTYPE: - raise ValueError(f"PLE table shard {key} has unsupported dtype {meta['dtype']}") + shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) + if shard_dtype is None or dtype not in (None, shard_dtype): + raise ValueError(f"PLE table shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") + dtype = shard_dtype shape = meta["shape"] if rows and tuple(shape) != (rows, cols): raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}") @@ -357,10 +423,12 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, if cols != qwen4_args.ngram_head_dim: raise ValueError(f"PLE table row is {cols} wide, config says {qwen4_args.ngram_head_dim}") if scale is None: - raise ValueError("PLE table has no weight_scale") + if dtype is torch.float8_e4m3fn: + raise ValueError("PLE table has no weight_scale") + scale = torch.tensor(1.0, dtype=torch.bfloat16) - bank = HostBank((expected * rows, cols), torch.float8_e4m3fn) - shard_bytes = rows * cols + bank = HostBank((expected * rows, cols), dtype) + shard_bytes = rows * cols * dtype.itemsize bar = byte_bar(expected * shard_bytes, "Loading PLE table") try: buf = bank.memoryview() @@ -382,8 +450,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index ba9cd98c6..838172074 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -68,25 +68,31 @@ def _fill(store, args, window, tokens): def _write_checkpoint(tmp_path, table, n_shards): + """uint8 ``table`` bytes are stored as a scaled fp8 table, any other dtype as an unscaled table of that dtype.""" from safetensors.torch import save_file + fp8 = table.dtype is torch.uint8 per = table.shape[0] // n_shards tensors = { - f"{_KEY_PREFIX}.shard_{i}.weight": table[i * per : (i + 1) * per].view(torch.float8_e4m3fn) + f"{_KEY_PREFIX}.shard_{i}.weight": table[i * per : (i + 1) * per].view(torch.float8_e4m3fn) if fp8 else table[i * per : (i + 1) * per].clone() for i in range(n_shards) } - tensors[f"{_KEY_PREFIX}.weight_scale"] = torch.tensor(0.03125, dtype=torch.bfloat16) + if fp8: + tensors[f"{_KEY_PREFIX}.weight_scale"] = torch.tensor(0.03125, dtype=torch.bfloat16) save_file(tensors, str(tmp_path / "model.safetensors")) -def _make_table(tmp_path): +def _make_table(tmp_path, dtype=torch.float8_e4m3fn): from freetoken.models.qwen4_exp.ple_disk import DiskRowTable, source_from_safetensors args = parse_config(toy_hf_config()).qwen4_args multipliers, sizes, offsets = hash_constants(args) total_rows = int(offsets[-1] + sizes[-1]) gen = torch.Generator().manual_seed(9) - table = torch.randint(0, 256, (total_rows, args.ngram_head_dim), dtype=torch.uint8, generator=gen) + if dtype is torch.float8_e4m3fn: + table = torch.randint(0, 256, (total_rows, args.ngram_head_dim), dtype=torch.uint8, generator=gen) + else: + table = torch.randn((total_rows, args.ngram_head_dim), generator=gen).to(dtype) n_shards = next(k for k in (4, 2, 1) if total_rows % k == 0) _write_checkpoint(tmp_path, table, n_shards) constants = { @@ -97,7 +103,7 @@ def _make_table(tmp_path): "eos_token_id": EOS, } disk = DiskRowTable(source_from_safetensors(str(tmp_path)), constants) - oracle = GpuResidentTable(table.cuda().view(torch.float8_e4m3fn), scale=0.03125) + oracle = GpuResidentTable(table.cuda().view(torch.float8_e4m3fn), scale=0.03125) if dtype is torch.float8_e4m3fn else GpuResidentTable(table.cuda()) return disk, oracle, args @@ -255,8 +261,9 @@ def run(store): @requires_cuda -def test_disk_table_matches_oracle(tmp_path): - disk, oracle, args = _make_table(tmp_path) +@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16], ids=["fp8", "bf16"]) +def test_disk_table_matches_oracle(tmp_path, dtype): + disk, oracle, args = _make_table(tmp_path, dtype) emb = _embedding() # prefill: two segments, one fresh and one mid-sequence window diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1a9a100ed..e066acbc2 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -18,9 +18,10 @@ from freetoken.kernel.aot_models import SUPPORTED_MODELS, expert_bank_row_bytes from freetoken.models.qwen4_exp.weight import ( _ZERO_CENTERED_NORM_SUFFIXES, - _DenseFuser, + _DenseReader, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into @@ -64,10 +65,12 @@ def _fp8_scale(weight: torch.Tensor) -> torch.Tensor: return torch.rand(weight.shape[0] // BLOCK, weight.shape[1] // BLOCK) + 0.5 -def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]: +def _raw_checkpoint(dense_fp8: bool = False, channel_fp8: bool = False, mxfp8_scale: str | None = None) -> dict[str, torch.Tensor]: """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise. ``dense_fp8`` stores the attention and GDN qkv|z / out projections as 128x128 block-fp8 (e4m3 ``.weight`` + fp32 ``.weight_scale_inv``) like the community NVFP4-FP8 requants. + ``channel_fp8`` stores those and the shared experts as per-channel fp8 (e4m3 ``.weight`` + bf16 ``[rows, 1]`` ``.weight_scale``) like the llm-compressor requants. + ``mxfp8_scale`` stores the attention and GDN qkv|z / out projections as MXFP8 (e4m3 ``.weight`` + uint8 e8m0 codes per 32 inputs) under that scale suffix. """ lm = "model.language_model" raw: dict[str, torch.Tensor] = { @@ -146,12 +149,63 @@ def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]: weight = raw[f"{module}.weight"] raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn) raw[f"{module}.weight_scale_inv"] = _fp8_scale(weight) + if channel_fp8: + for module in _channel_fp8_modules(gdn_layers=(0,), attn_layers=(1,), mlp_layers=(0, 1)): + weight = raw[f"{module}.weight"] + raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn) + raw[f"{module}.weight_scale"] = (torch.rand(weight.shape[0], 1) + 0.5).to(torch.bfloat16) + if mxfp8_scale is not None: + for module in (f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z", f"{gdn}.out_proj", *(f"{attn}.{p}_proj" for p in "qkvo")): + weight = raw[f"{module}.weight"] + raw[f"{module}.weight"] = weight.to(torch.float8_e4m3fn) + raw[f"{module}.{mxfp8_scale}"] = torch.randint(120, 135, (weight.shape[0], weight.shape[1] // 32), dtype=torch.uint8) return raw FP8_DENSE_QUANT = mixed_precision_quant(gdn_layers=(0,), attn_layers=(1,), moe_layers=(0, 1)) +# ModelOpt MIXED_PRECISION with MXFP8 attention / GDN projections, as in nvidia/Qwen3.8-Flash-Next-FP4-Local +MXFP8_DENSE_QUANT = { + **FP8_DENSE_QUANT, + "quantized_layers": { + k: ({"quant_algo": "MXFP8", "group_size": 32} if v["quant_algo"] == "FP8_PB_WO" else v) + for k, v in FP8_DENSE_QUANT["quantized_layers"].items() + }, +} + + +def _channel_fp8_modules(gdn_layers, attn_layers, mlp_layers) -> list[str]: + return ( + [f"{LM}.layers.{i}.linear_attn.{p}" for i in gdn_layers for p in ("in_proj_qkv", "in_proj_z", "out_proj")] + + [f"{LM}.layers.{i}.self_attn.{p}_proj" for i in attn_layers for p in "qkvo"] + + [f"{LM}.layers.{i}.mlp.shared_expert.{p}_proj" for i in mlp_layers for p in ("gate", "up", "down")] + ) + + +# llm-compressor mixed-precision as in orcarouter/Qwen3.8-Flash-Next-Uncensored-NVFP4: per-channel fp8 dense, weight-only NVFP4 experts +CT_CHANNEL_FP8_QUANT = { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "quantization_status": "compressed", + "config_groups": { + "group_0": { + "format": "naive-quantized", + "targets": _channel_fp8_modules(gdn_layers=(0,), attn_layers=(1,), mlp_layers=(0, 1)), + "weights": {"num_bits": 8, "type": "float", "strategy": "channel", "symmetric": True, "dynamic": False}, + "input_activations": None, + }, + "group_1": { + "format": "nvfp4-pack-quantized", + "targets": [r"re:.*mlp\.experts\..*proj$"], + "weights": {"num_bits": 4, "type": "float", "strategy": "tensor_group", "group_size": 16, "symmetric": True, "dynamic": False}, + "input_activations": None, + }, + }, + "ignore": ["lm_head"], +} + + def _config_json(quantization_config) -> dict: cfg = hf_config( num_layers=2, head_dim=AHD, num_q=QH, num_kv=KVH, index_head_dim=IHD, index_heads=2, @@ -227,6 +281,14 @@ def loaded_fp8(checkpoint_fp8) -> dict[str, torch.Tensor]: return _load(checkpoint_fp8[0]) +@pytest.fixture(scope="module") +def checkpoint_fp8_channel(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: + torch.manual_seed(3) + return _write_checkpoint( + tmp_path_factory.mktemp("qwen4_exp_fp8_channel_ckpt"), _raw_checkpoint(channel_fp8=True), CT_CHANNEL_FP8_QUANT + ) + + def test_tower_keys_come_out_under_the_prefix_load_weight_filters(loaded): assert {n for n in loaded if "visual" in n} == {"visual.blocks.0.attn.qkv.weight", "visual.merger.norm.weight"} @@ -348,6 +410,15 @@ def test_load_ple_table_concatenates_shards_in_index_order(checkpoint): assert float(table.weight_scale) == 0.125 +def test_load_ple_table_reads_an_unscaled_bf16_table(tmp_path): + prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" + shards = {f"{prefix}.shard_{i}.weight": _bf16(NGRAM_ROWS, NGRAM_DIM) for i in range(NGRAM_SHARDS)} + save_file(shards, str(tmp_path / "model.safetensors")) + table = load_ple_table(str(tmp_path), SimpleNamespace(split_ngram_parts=NGRAM_SHARDS, ngram_head_dim=NGRAM_DIM), pin=False) + assert torch.equal(table.tensor, torch.cat([shards[f"{prefix}.shard_{i}.weight"] for i in range(NGRAM_SHARDS)])) + assert float(table.weight_scale) == 1.0 + + def test_load_ple_table_rejects_a_shard_count_mismatch(checkpoint): folder, _raw = checkpoint args = SimpleNamespace(split_ngram_parts=NGRAM_SHARDS + 1, ngram_head_dim=NGRAM_DIM) @@ -429,11 +500,11 @@ def test_every_registry_architecture_is_claimed_by_an_aot_entry(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs cuda") def test_fusion_pad_rides_the_tensor_device(): """safetensors loads straight to cuda; a cpu-allocated pad row would break torch.cat.""" - fuser = _DenseFuser(None, get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping) + reader = _DenseReader(None, get_model_spec("Qwen4ExpForConditionalGeneration").packed_modules_mapping) down = torch.randn(320, 64, device="cuda", dtype=torch.bfloat16) inject = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) - assert fuser.fuse("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down) == [] - [(key, fused)] = fuser.fuse("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject) + assert reader.add("model.layers.0.attn_hyper_connection.input_mix_weight_down.weight", down) == [] + [(key, fused)] = reader.add("model.layers.0.attn_hyper_connection.block_inject_weight.weight", inject) assert key == "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight" assert fused.device.type == "cuda" and fused.shape[0] == 336 assert torch.equal(fused[324:], torch.zeros(12, 64, device="cuda", dtype=torch.bfloat16)) @@ -457,20 +528,55 @@ def checkpoint_nvfp4(tmp_path_factory) -> tuple[str, dict[str, torch.Tensor]]: ) -@pytest.mark.parametrize("fixture", ["checkpoint", "checkpoint_nvfp4", "checkpoint_fp8"]) +@pytest.mark.parametrize("fixture", ["checkpoint", "checkpoint_nvfp4", "checkpoint_fp8", "checkpoint_fp8_channel"]) def test_emitted_keys_are_the_model_state_dict(fixture, request): - """The reader fills exactly the buffers the engine builds from the same config, block-fp8 ones with the stored dtypes.""" + """The reader fills exactly the buffers the engine builds from the same config, fp8 ones with the stored weight dtype and fp32 scales.""" folder, _raw = request.getfixturevalue(fixture) loaded, state = _load(folder, vision=False), meta_state_dict(folder) assert set(loaded) == set(state) - if fixture != "checkpoint_fp8": + if fixture in ("checkpoint", "checkpoint_nvfp4"): assert loaded["model.layers.0.linear_attn.in_proj.weight"].dtype is torch.bfloat16 return + scale = ".weight_scale_inv" if fixture == "checkpoint_fp8" else ".weight_scale" for module in FP8_MODULES: - for kind in (".weight", ".weight_scale_inv"): + for kind in (".weight", scale): assert loaded[module + kind].shape == state[module + kind].shape, module + kind assert loaded[module + ".weight"].dtype is state[module + ".weight"].dtype is torch.float8_e4m3fn - assert loaded[module + ".weight_scale_inv"].dtype is torch.float32 # the engine casts it to the bf16 buffer at load + assert loaded[module + scale].dtype is torch.float32 # the engine casts it to the buffer dtype at load + + +def test_channel_fp8_scales_fuse_as_fp32_rows(checkpoint_fp8_channel): + folder, raw = checkpoint_fp8_channel + loaded = _load(folder) + gdn, shared = f"{LM}.layers.0.linear_attn", f"{LM}.layers.1.mlp.shared_expert" + for fused, parts in ( + ("model.layers.0.linear_attn.in_proj_qkvz", [f"{gdn}.in_proj_qkv", f"{gdn}.in_proj_z"]), + ("model.layers.1.mlp.shared_expert.gate_up_proj", [f"{shared}.gate_proj", f"{shared}.up_proj"]), + ): + assert torch.equal(loaded[f"{fused}.weight_scale"], torch.cat([raw[f"{p}.weight_scale"].reshape(-1).float() for p in parts])) + assert torch.equal(loaded[f"{fused}.weight"].view(torch.uint8), torch.cat([raw[f"{p}.weight"].view(torch.uint8) for p in parts])) + + +@pytest.mark.parametrize("scale_name", ["weight_scale_inv", "weight_scale"]) +def test_mxfp8_scales_load_under_either_modelopt_name(tmp_path, scale_name): + folder, raw = _write_checkpoint(tmp_path, _raw_checkpoint(mxfp8_scale=scale_name), MXFP8_DENSE_QUANT) + loaded, state = _load(folder, vision=False), meta_state_dict(folder) + assert set(loaded) == set(state) + gdn = f"{LM}.layers.0.linear_attn" + assert torch.equal( + loaded["model.layers.0.linear_attn.in_proj_qkvz.weight_scale_inv"], + torch.cat([raw[f"{gdn}.{p}.{scale_name}"] for p in ("in_proj_qkv", "in_proj_z")]), + ) + + +def test_expert_spec_reads_the_llm_compressor_names(checkpoint_fp8_channel): + folder, _raw = checkpoint_fp8_channel + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + base = f"{LM}.layers.1.mlp.experts.2.down_proj" + kinds = {spec.kind_map[spec.key_pattern.match(f"{base}.{k}")["kind"]] for k in ("weight_packed", "weight_scale", "weight_global_scale")} + assert kinds == {"weight", "weight_scale", "weight_scale_2"} + assert spec.global_reciprocal def _assert_fused_per_kind(loaded, raw, fused: str, parts: list[str]) -> None: diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py index 4395a8639..1c073792a 100644 --- a/tests/models/qwen4_exp/test_weight_ckpt.py +++ b/tests/models/qwen4_exp/test_weight_ckpt.py @@ -22,10 +22,10 @@ from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.weight import ( - _NVFP4_SOURCE_SPEC, _ZERO_CENTERED_NORM_SUFFIXES, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.moe.expert_banks import build_expert_banks from freetoken.moe.host_banks import HostResidency @@ -170,7 +170,7 @@ def dense_pass() -> tuple[list[str], dict[str, torch.Tensor]]: sampled: dict[str, torch.Tensor] = {} install_quant_config(MODEL_PATH) for name, tensor in iter_weights( - MODEL_PATH, torch.device("cpu"), include_moe_experts=True, include_non_moe=True + MODEL_PATH, torch.device("cpu"), include_moe_experts=True, include_non_moe=True, include_vision=False ): names.append(name) if name in wanted or name.endswith(_ZERO_CENTERED_NORM_SUFFIXES): @@ -271,8 +271,9 @@ def layer0_expert_banks(): """The real NVFP4 source-bank loader, restricted to layer 0 (1.32 GiB instead of 63.5).""" if not torch.cuda.is_available(): pytest.skip("expert bank pinning needs CUDA") + install_quant_config(MODEL_PATH) spec = dataclasses.replace( - _NVFP4_SOURCE_SPEC, layer_to_bank=lambda layer, config: 0 if layer == 0 else None + nvfp4_expert_spec(MODEL_PATH, None), layer_to_bank=lambda layer, config: 0 if layer == 0 else None ) config = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, num_moe_layers=1) diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py index 9c18b7490..6b9ecd810 100644 --- a/tests/models/test_quant_config.py +++ b/tests/models/test_quant_config.py @@ -569,9 +569,11 @@ def shape(suffix): if tensors: if dtype("weight") == "U8": return {QuantKind.MXFP4} + if dtype("weight_scale") == "U8": + return {QuantKind.MXFP8} sh = shape("weight_scale") or [] return {QuantKind.FP8_BLOCK} if len(sh) == 2 and sh[0] > 1 and sh[1] > 1 else {QuantKind.FP8_TENSOR} - return {QuantKind.FP8_TENSOR, QuantKind.FP8_BLOCK, QuantKind.MXFP4} + return {QuantKind.FP8_TENSOR, QuantKind.FP8_BLOCK, QuantKind.MXFP4, QuantKind.MXFP8} return {QuantKind.NONE}