feat: add qwen35 text checkpoint pruner

MVP/qwen35_prune.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import math
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable
+
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+HYBRID_BLOCK = ("linear_attention", "linear_attention", "linear_attention", "full_attention")
+DEFAULT_MINIMUM = 330_000_000
+DEFAULT_MAXIMUM = 350_000_000
+
+
+@dataclass(frozen=True)
+class ParameterReport:
+    embedding_params: int
+    layer_params: tuple[int, ...]
+    layer_types: tuple[str, ...]
+    final_norm_params: int
+    all_named_params: int
+
+    @property
+    def text_backbone_params(self) -> int:
+        return self.embedding_params + sum(self.layer_params) + self.final_norm_params
+
+
+@dataclass(frozen=True)
+class PrefixChoice:
+    layer_count: int
+    parameter_count: int
+
+
+def load_live_config(path: str | Path) -> dict:
+    with Path(path).open(encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def _numel(shape: Iterable[int]) -> int:
+    return math.prod(int(dimension) for dimension in shape)
+
+
+def _layer_index(key: str) -> int | None:
+    prefix = "model.language_model.layers."
+    if not key.startswith(prefix):
+        return None
+    remainder = key[len(prefix):]
+    index_text = remainder.split(".", 1)[0]
+    return int(index_text)
+
+
+def count_parameter_groups(safetensors_path: str | Path, config: dict) -> ParameterReport:
+    text_config = config.get("text_config", config)
+    layer_types = tuple(text_config["layer_types"])
+    layer_params = [0 for _ in layer_types]
+    embedding_params = 0
+    final_norm_params = 0
+    all_named_params = 0
+
+    with safe_open(str(safetensors_path), framework="pt", device="cpu") as handle:
+        for key in handle.keys():
+            shape = handle.get_slice(key).get_shape()
+            params = _numel(shape)
+            all_named_params += params
+            if key == "model.language_model.embed_tokens.weight":
+                embedding_params += params
+            elif key == "model.language_model.norm.weight":
+                final_norm_params += params
+            else:
+                index = _layer_index(key)
+                if index is not None:
+                    if index >= len(layer_params):
+                        raise ValueError(f"tensor layer index {index} exceeds config layer count")
+                    layer_params[index] += params
+
+    return ParameterReport(
+        embedding_params=embedding_params,
+        layer_params=tuple(layer_params),
+        layer_types=layer_types,
+        final_norm_params=final_norm_params,
+        all_named_params=all_named_params,
+    )
+
+
+def _is_complete_prefix(layer_types: tuple[str, ...], layer_count: int) -> bool:
+    if layer_count == 0 or layer_count % len(HYBRID_BLOCK) != 0:
+        return False
+    return layer_types[:layer_count] == HYBRID_BLOCK * (layer_count // len(HYBRID_BLOCK))
+
+
+def choose_prefix(
+    report: ParameterReport,
+    minimum: int,
+    maximum: int,
+    requested_layers: int | None = None,
+) -> PrefixChoice:
+    if minimum > maximum:
+        raise ValueError("minimum parameter bound exceeds maximum")
+    limit = min(len(report.layer_params), len(report.layer_types))
+    candidates = [requested_layers] if requested_layers is not None else range(
+        len(HYBRID_BLOCK), limit + 1, len(HYBRID_BLOCK)
+    )
+
+    choices: list[PrefixChoice] = []
+    for layer_count in candidates:
+        if layer_count is None or layer_count > limit:
+            continue
+        if not _is_complete_prefix(report.layer_types, layer_count):
+            if requested_layers is not None:
+                raise ValueError("requested layer count is not a complete hybrid block")
+            continue
+        parameter_count = (
+            report.embedding_params
+            + report.final_norm_params
+            + sum(report.layer_params[:layer_count])
+        )
+        if minimum <= parameter_count <= maximum:
+            choices.append(PrefixChoice(layer_count, parameter_count))
+
+    if not choices:
+        raise ValueError("no complete hybrid prefix fits the parameter interval")
+    return max(choices, key=lambda choice: choice.layer_count)
+
+
+def translate_text_key(key: str) -> str | None:
+    if key.startswith("model.visual.") or key.startswith("mtp."):
+        return None
+    prefix = "model.language_model."
+    if key.startswith(prefix):
+        return "model." + key[len(prefix):]
+    if key == "lm_head.weight":
+        return None
+    return None
+
+
+def build_text_config(full_config: dict, layer_count: int) -> dict:
+    source = copy.deepcopy(full_config.get("text_config", full_config))
+    layer_types = list(source.get("layer_types", []))
+    if layer_count <= 0 or layer_count > len(layer_types):
+        raise ValueError("layer count is outside the live text config")
+    if not _is_complete_prefix(tuple(layer_types), layer_count):
+        raise ValueError("layer count is not a complete hybrid block")
+    source["model_type"] = "qwen3_5_text"
+    source["num_hidden_layers"] = layer_count
+    source["layer_types"] = layer_types[:layer_count]
+    source["tie_word_embeddings"] = bool(
+        full_config.get("tie_word_embeddings", source.get("tie_word_embeddings", False))
+    )
+    source["architectures"] = ["Qwen3_5ForCausalLM"]
+    source.pop("vision_config", None)
+    source.pop("mtp_config", None)
+    return source
+
+
+def _keep_source_key(key: str, layer_count: int) -> bool:
+    if key == "model.language_model.embed_tokens.weight":
+        return True
+    if key == "model.language_model.norm.weight":
+        return True
+    index = _layer_index(key)
+    return index is not None and index < layer_count
+
+
+def prune_checkpoint(
+    input_weights: str | Path,
+    input_config: str | Path,
+    output_dir: str | Path,
+    layer_count: int,
+) -> dict:
+    full_config = load_live_config(input_config)
+    text_config = build_text_config(full_config, layer_count)
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    tensors = {}
+
+    with safe_open(str(input_weights), framework="pt", device="cpu") as handle:
+        for key in handle.keys():
+            if not _keep_source_key(key, layer_count):
+                continue
+            output_key = translate_text_key(key)
+            if output_key is None:
+                raise ValueError(f"source tensor cannot be translated: {key}")
+            tensors[output_key] = handle.get_tensor(key)
+
+    output_weights = output_path / "model.safetensors"
+    save_file(tensors, str(output_weights), metadata={"format": "pt"})
+    output_config = output_path / "config.json"
+    output_config.write_text(json.dumps(text_config, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+    from MVP.validate_checkpoint import validate_checkpoint
+
+    report = validate_checkpoint(
+        output_config,
+        output_weights,
+        DEFAULT_MINIMUM,
+        DEFAULT_MAXIMUM,
+    )
+    return asdict(report)
+
+
+def _main() -> None:
+    parser = argparse.ArgumentParser(description="Measure and prune Qwen3.5 text backbone tensors")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    measure = subparsers.add_parser("measure")
+    measure.add_argument("--weights", required=True)
+    measure.add_argument("--config", required=True)
+
+    prune = subparsers.add_parser("prune")
+    prune.add_argument("--weights", required=True)
+    prune.add_argument("--config", required=True)
+    prune.add_argument("--output", required=True)
+    prune.add_argument("--layers", type=int, required=True)
+
+    validate = subparsers.add_parser("validate")
+    validate.add_argument("--weights", required=True)
+    validate.add_argument("--config", required=True)
+    validate.add_argument("--minimum", type=int, default=DEFAULT_MINIMUM)
+    validate.add_argument("--maximum", type=int, default=DEFAULT_MAXIMUM)
+
+    args = parser.parse_args()
+    if args.command == "measure":
+        report = count_parameter_groups(args.weights, load_live_config(args.config))
+        print(json.dumps(asdict(report) | {"text_backbone_params": report.text_backbone_params}, indent=2))
+    elif args.command == "prune":
+        print(json.dumps(prune_checkpoint(args.weights, args.config, args.output, args.layers), indent=2))
+    else:
+        from MVP.validate_checkpoint import validate_checkpoint
+
+        print(json.dumps(asdict(validate_checkpoint(args.config, args.weights, args.minimum, args.maximum)), indent=2))
+
+
+if __name__ == "__main__":
+    _main()

MVP/tests/test_pruning_contract.py

@@ -2,10 +2,12 @@ import json
 from pathlib import Path
 
 import pytest
+from safetensors.torch import save_file
 
 from MVP.qwen35_prune import (
     ParameterReport,
     build_text_config,
+    count_parameter_groups,
     choose_prefix,
     translate_text_key,
 )
@@ -116,3 +118,38 @@ def test_validator_rejects_vision_mtp_and_duplicate_tied_head():
     with pytest.raises(ValueError, match="vision|MTP|tied"):
         validate_state_dict_keys(keys, config, 330_000_000, 350_000_000)
 
+
+def test_metadata_counter_uses_tensor_shapes_without_loading_model():
+    data = load_fixture()
+    config = {
+        "text_config": {
+            "layer_types": list(data["layer_types"][:4]),
+        }
+    }
+    tensors = {
+        "model.language_model.embed_tokens.weight": (2, 3),
+        "model.language_model.layers.0.input_layernorm.weight": (2,),
+        "model.language_model.layers.1.input_layernorm.weight": (3,),
+        "model.language_model.layers.2.input_layernorm.weight": (4,),
+        "model.language_model.layers.3.input_layernorm.weight": (5,),
+        "model.language_model.norm.weight": (2,),
+        "model.visual.patch_embed.proj.weight": (7,),
+        "mtp.fc.weight": (11,),
+    }
+
+    def make_tensor(shape):
+        import torch
+
+        return torch.zeros(shape, dtype=torch.float32)
+
+    path = FIXTURE.parent / "counter-fixture.safetensors"
+    try:
+        save_file({key: make_tensor(shape) for key, shape in tensors.items()}, str(path))
+        report = count_parameter_groups(path, config)
+    finally:
+        path.unlink(missing_ok=True)
+
+    assert report.embedding_params == 6
+    assert report.layer_params == (2, 3, 4, 5)
+    assert report.final_norm_params == 2
+    assert report.all_named_params == 40

MVP/validate_checkpoint.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Mapping, Sequence
+
+from safetensors import safe_open
+
+
+@dataclass(frozen=True)
+class ValidationReport:
+    parameter_count: int
+    tensor_count: int
+    layer_count: int
+    visual_tensor_count: int
+    mtp_tensor_count: int
+    tied_lm_head_present: bool
+    valid: bool
+
+
+def _numel(shape: Sequence[int]) -> int:
+    return math.prod(int(dimension) for dimension in shape)
+
+
+def _layer_indices(keys: Sequence[str]) -> set[int]:
+    indices = set()
+    prefix = "model.layers."
+    for key in keys:
+        if key.startswith(prefix):
+            indices.add(int(key[len(prefix):].split(".", 1)[0]))
+    return indices
+
+
+def validate_state_dict_keys(
+    shapes: Mapping[str, Sequence[int]],
+    config: Mapping[str, object],
+    minimum: int,
+    maximum: int,
+) -> ValidationReport:
+    keys = tuple(shapes.keys())
+    visual_count = sum(key.startswith("model.visual.") for key in keys)
+    mtp_count = sum(key.startswith("mtp.") for key in keys)
+    tied_head = "lm_head.weight" in shapes and bool(config.get("tie_word_embeddings", False))
+    if visual_count or mtp_count or tied_head:
+        raise ValueError("vision/MTP tensors or a duplicate tied head are present")
+    if config.get("model_type") != "qwen3_5_text":
+        raise ValueError("standalone text config must use model_type qwen3_5_text")
+
+    layer_count = int(config["num_hidden_layers"])
+    expected_indices = set(range(layer_count))
+    actual_indices = _layer_indices(keys)
+    if actual_indices != expected_indices:
+        raise ValueError(f"layer indices are not contiguous: expected {expected_indices}, got {actual_indices}")
+    layer_types = tuple(config["layer_types"])
+    if len(layer_types) != layer_count:
+        raise ValueError("config layer_types length does not match num_hidden_layers")
+    parameter_count = sum(_numel(shape) for shape in shapes.values())
+    if not minimum <= parameter_count <= maximum:
+        raise ValueError(f"parameter count {parameter_count} is outside [{minimum}, {maximum}]")
+    return ValidationReport(
+        parameter_count=parameter_count,
+        tensor_count=len(keys),
+        layer_count=layer_count,
+        visual_tensor_count=visual_count,
+        mtp_tensor_count=mtp_count,
+        tied_lm_head_present=tied_head,
+        valid=True,
+    )
+
+
+def validate_checkpoint(
+    config_path: str | Path,
+    weights_path: str | Path,
+    minimum: int,
+    maximum: int,
+) -> ValidationReport:
+    config = json.loads(Path(config_path).read_text(encoding="utf-8"))
+    shapes = {}
+    with safe_open(str(weights_path), framework="pt", device="cpu") as handle:
+        for key in handle.keys():
+            shapes[key] = tuple(handle.get_slice(key).get_shape())
+    report = validate_state_dict_keys(shapes, config, minimum, maximum)
+    Path(config_path).with_name("validation.json").write_text(
+        json.dumps(asdict(report), indent=2, sort_keys=True) + "\n", encoding="utf-8"
+    )
+    return report
+