From 150c974eed7e560285026878d23972e843e5c90f Mon Sep 17 00:00:00 2001 From: Cristopher Yates Date: Sun, 1 Mar 2026 05:16:56 -0330 Subject: [PATCH] =?UTF-8?q?feat(phase3):=20PyTorch=E2=86=92MLX=20weight=20?= =?UTF-8?q?converter=20(#1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(phase3): PyTorch→MLX weight converter + safetensors output 365 keys mapped (367 source - 2 num_batches_tracked skipped). Conv weights transposed (O,I,H,W→O,H,W,I), refiner stem remapped, 4ch patch embed preserved. 12 conversion tests + diagnostic report. Co-Authored-By: Claude Opus 4.6 * fix(phase3): conv allowlist, unused var, test fixture - Replace _is_conv_weight heuristic with explicit CONV_WEIGHT_KEYS frozenset (15 keys) - Remove unused `skipped` list from convert_state_dict - Module-scoped pytest fixture eliminates 11 redundant checkpoint loads Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- ...eat-corridorkey-mlx-inference-port-plan.md | 6 +- ...3-01-fix-converter-review-feedback-plan.md | 59 +++++ scripts/convert_weights.py | 85 +++++++ src/corridorkey_mlx/convert/__init__.py | 16 ++ src/corridorkey_mlx/convert/converter.py | 211 +++++++++++++++++- tests/test_conversion.py | 178 ++++++++++++++- 6 files changed, 538 insertions(+), 17 deletions(-) create mode 100644 docs/plans/2026-03-01-fix-converter-review-feedback-plan.md create mode 100644 scripts/convert_weights.py diff --git a/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md b/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md index 93343d4..0d25c07 100644 --- a/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md +++ b/docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md @@ -145,9 +145,9 @@ RGB + coarse alpha hint (4ch) **Deliverables:** -- [ ] `src/corridorkey_mlx/convert/converter.py` -- key mapping + weight transforms -- [ ] Conversion diagnostic report (source key -> dest key, shapes, transform applied) -- [ ] Output as safetensors +- [x] `src/corridorkey_mlx/convert/converter.py` -- key mapping + weight transforms +- [x] Conversion diagnostic report (source key -> dest key, shapes, transform applied) +- [x] Output as safetensors **Files touched:** diff --git a/docs/plans/2026-03-01-fix-converter-review-feedback-plan.md b/docs/plans/2026-03-01-fix-converter-review-feedback-plan.md new file mode 100644 index 0000000..865b99a --- /dev/null +++ b/docs/plans/2026-03-01-fix-converter-review-feedback-plan.md @@ -0,0 +1,59 @@ +--- +title: Fix converter review feedback (items 4, 1, 2) +type: fix +date: 2026-03-01 +--- + +# Fix converter review feedback + +Three fixes from PR #1 review, ordered by simplicity. + +## Fix 4: Remove unused `skipped` list + +**File:** `src/corridorkey_mlx/convert/converter.py:105,110` + +Delete the `skipped` list variable and its `.append()` call. The skip logic already works via `continue` — the list serves no purpose. + +## Fix 1: Explicit conv key allowlist + +**File:** `src/corridorkey_mlx/convert/converter.py:60-66` + +Replace heuristic `_is_conv_weight()` with an explicit `CONV_WEIGHT_KEYS: frozenset[str]` allowlist. This matches the "no regex guessing, no silent fallbacks" philosophy. + +The 15 conv keys (from checkpoint analysis + tests): +- `encoder.model.patch_embed.proj.weight` — (112,4,7,7) +- `alpha_decoder.linear_fuse.weight` — (256,1024,1,1) +- `fg_decoder.linear_fuse.weight` — (256,1024,1,1) +- `refiner.stem.0.weight` (pre-remap) — (64,7,3,3) +- `refiner.res{1-4}.conv{1-2}.weight` — 8 keys, all (64,64,3,3) +- `refiner.final.weight` — (4,64,1,1) + +**Note:** Use pre-remap key names since `_is_conv_weight` is called before remapping in the loop. Total: 15 keys. + +Replace `_is_conv_weight(key, arr)` → simple `key in CONV_WEIGHT_KEYS` check. Remove the `arr` parameter. + +## Fix 2: Module-scoped test fixture + +**File:** `tests/test_conversion.py` + +Add a `@pytest.fixture(scope="module")` that loads + converts once, shared by all tests. This avoids 11 redundant `load_pytorch_checkpoint` + `convert_state_dict` calls. + +```python +@pytest.fixture(scope="module") +def converted_checkpoint(): + if not CHECKPOINT_PATH.exists(): + pytest.skip("Checkpoint not found") + state_dict = load_pytorch_checkpoint(CHECKPOINT_PATH) + converted, diagnostics = convert_state_dict(state_dict) + return state_dict, converted, diagnostics +``` + +Each test method takes `converted_checkpoint` as param, destructures what it needs. Remove `_skip_if_no_checkpoint()` helper. + +## Acceptance Criteria + +- [ ] `skipped` list removed from `convert_state_dict` +- [ ] `CONV_WEIGHT_KEYS` frozenset with 15 keys replaces `_is_conv_weight` heuristic +- [ ] Module-scoped fixture eliminates repeated checkpoint loading +- [ ] `uv run pytest tests/test_conversion.py -v` — 12/12 pass +- [ ] `uv run ruff check .` — clean diff --git a/scripts/convert_weights.py b/scripts/convert_weights.py new file mode 100644 index 0000000..63d1e47 --- /dev/null +++ b/scripts/convert_weights.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Convert CorridorKey PyTorch checkpoint to MLX safetensors. + +Usage: + uv run --group reference python scripts/convert_weights.py \ + --checkpoint checkpoints/CorridorKey_v1.0.pth \ + --output checkpoints/corridorkey_mlx.safetensors +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from rich.console import Console +from rich.table import Table + +from corridorkey_mlx.convert.converter import convert_checkpoint + +console = Console() + +DEFAULT_CHECKPOINT = Path("checkpoints/CorridorKey_v1.0.pth") +DEFAULT_OUTPUT = Path("checkpoints/corridorkey_mlx.safetensors") + + +def print_diagnostics(diagnostics: list) -> None: + """Print conversion diagnostic table.""" + table = Table(title="Weight Conversion Report") + table.add_column("Source Key", style="cyan", max_width=50) + table.add_column("Dest Key", style="green", max_width=50) + table.add_column("Src Shape", style="yellow") + table.add_column("Dst Shape", style="yellow") + table.add_column("Transform", style="magenta") + + for record in diagnostics: + table.add_row( + record.source_key, + record.dest_key, + str(record.source_shape), + str(record.dest_shape), + record.transform, + ) + + console.print(table) + + # Summary stats + total = len(diagnostics) + transposed = sum(1 for r in diagnostics if "conv_transpose" in r.transform) + remapped = sum(1 for r in diagnostics if r.source_key != r.dest_key) + console.print(f"\n[bold]Total keys:[/bold] {total}") + console.print(f"[bold]Conv transposed:[/bold] {transposed}") + console.print(f"[bold]Keys remapped:[/bold] {remapped}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Convert CorridorKey PyTorch → MLX weights") + parser.add_argument( + "--checkpoint", + type=Path, + default=DEFAULT_CHECKPOINT, + help="Path to PyTorch .pth checkpoint", + ) + parser.add_argument( + "--output", + type=Path, + default=DEFAULT_OUTPUT, + help="Output .safetensors path", + ) + args = parser.parse_args() + + if not args.checkpoint.exists(): + console.print(f"[red]Checkpoint not found: {args.checkpoint}[/red]") + raise SystemExit(1) + + console.print(f"[bold]Converting:[/bold] {args.checkpoint}") + console.print(f"[bold]Output:[/bold] {args.output}") + + diagnostics = convert_checkpoint(args.checkpoint, args.output) + print_diagnostics(diagnostics) + + console.print(f"\n[green]Saved to {args.output}[/green]") + + +if __name__ == "__main__": + main() diff --git a/src/corridorkey_mlx/convert/__init__.py b/src/corridorkey_mlx/convert/__init__.py index 61b980a..b664af0 100644 --- a/src/corridorkey_mlx/convert/__init__.py +++ b/src/corridorkey_mlx/convert/__init__.py @@ -1 +1,17 @@ """PyTorch to MLX weight conversion utilities.""" + +from corridorkey_mlx.convert.converter import ( + ConversionRecord, + convert_checkpoint, + convert_state_dict, + load_pytorch_checkpoint, + save_safetensors, +) + +__all__ = [ + "ConversionRecord", + "convert_checkpoint", + "convert_state_dict", + "load_pytorch_checkpoint", + "save_safetensors", +] diff --git a/src/corridorkey_mlx/convert/converter.py b/src/corridorkey_mlx/convert/converter.py index 7432d25..5c55737 100644 --- a/src/corridorkey_mlx/convert/converter.py +++ b/src/corridorkey_mlx/convert/converter.py @@ -1,9 +1,206 @@ -"""PyTorch → MLX weight converter (not yet implemented). +"""PyTorch → MLX weight converter. -Responsibilities: -- Map state_dict keys between PyTorch and MLX naming -- Transpose conv weights from NCHW → NHWC -- Handle patched 4-channel first conv -- Emit diagnostic report (src key, dst key, shapes, transform) -- Save as safetensors +Explicit key-by-key mapping from CorridorKey PyTorch checkpoint +to MLX-compatible weights. No regex guessing, no silent fallbacks. + +Transforms applied: +- Conv2d weights: (O, I, H, W) → (O, H, W, I) +- Refiner stem: Sequential indices → named attrs (stem.0 → stem_conv, stem.1 → stem_gn) +- BatchNorm num_batches_tracked: dropped (MLX doesn't use it) +- All other weights: passed through unchanged + +Output: safetensors file with transformed weights. """ + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from pathlib import Path + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +COMPILE_PREFIX = "_orig_mod." + +# Refiner stem key remapping: PyTorch nn.Sequential indices → MLX named attrs +REFINER_STEM_MAP: dict[str, str] = { + "refiner.stem.0.weight": "refiner.stem_conv.weight", + "refiner.stem.0.bias": "refiner.stem_conv.bias", + "refiner.stem.1.weight": "refiner.stem_gn.weight", + "refiner.stem.1.bias": "refiner.stem_gn.bias", +} + +# Keys to skip entirely (not needed by MLX) +SKIP_SUFFIXES = (".num_batches_tracked",) + +# Explicit set of conv weight keys that need NCHW→NHWC transpose. +# Uses pre-remap names (checked before key remapping in the conversion loop). +CONV_WEIGHT_KEYS: frozenset[str] = frozenset( + [ + "encoder.model.patch_embed.proj.weight", + "alpha_decoder.linear_fuse.weight", + "alpha_decoder.classifier.weight", + "fg_decoder.linear_fuse.weight", + "fg_decoder.classifier.weight", + "refiner.stem.0.weight", + "refiner.final.weight", + ] + + [f"refiner.res{i}.conv{j}.weight" for i in range(1, 5) for j in range(1, 3)] +) + + +# --------------------------------------------------------------------------- +# Diagnostic record +# --------------------------------------------------------------------------- +@dataclass +class ConversionRecord: + """Single key conversion diagnostic.""" + + source_key: str + dest_key: str + source_shape: tuple[int, ...] + dest_shape: tuple[int, ...] + transform: str + + +# --------------------------------------------------------------------------- +# Transform helpers +# --------------------------------------------------------------------------- +def _is_conv_weight(key: str) -> bool: + """Check if a weight is a conv kernel that needs NCHW→NHWC transpose.""" + return key in CONV_WEIGHT_KEYS + + +def _transpose_conv_weight(arr: np.ndarray) -> np.ndarray: + """PyTorch conv (O, I, H, W) → MLX conv (O, H, W, I).""" + return np.transpose(arr, (0, 2, 3, 1)) + + +def _remap_key(key: str) -> str: + """Apply key remapping rules. Returns new key.""" + if key in REFINER_STEM_MAP: + return REFINER_STEM_MAP[key] + return key + + +def _should_skip(key: str) -> bool: + """Check if key should be dropped from output.""" + return any(key.endswith(suffix) for suffix in SKIP_SUFFIXES) + + +# --------------------------------------------------------------------------- +# Core conversion +# --------------------------------------------------------------------------- +def convert_state_dict( + state_dict: dict[str, np.ndarray], +) -> tuple[OrderedDict[str, np.ndarray], list[ConversionRecord]]: + """Convert PyTorch state_dict to MLX-compatible weight dict. + + Args: + state_dict: PyTorch state_dict as numpy arrays (already stripped of _orig_mod. prefix). + + Returns: + Tuple of (converted weights dict, list of diagnostic records). + + Raises: + ValueError: If duplicate destination keys are produced. + """ + converted: OrderedDict[str, np.ndarray] = OrderedDict() + diagnostics: list[ConversionRecord] = [] + + for src_key, arr in state_dict.items(): + # Skip keys MLX doesn't need + if _should_skip(src_key): + continue + + # Remap key + dest_key = _remap_key(src_key) + + # Apply transform + src_shape = arr.shape + if _is_conv_weight(src_key): + transformed = _transpose_conv_weight(arr) + transform_name = "conv_transpose(O,I,H,W→O,H,W,I)" + else: + transformed = arr + transform_name = "passthrough" + + # Check for duplicates + if dest_key in converted: + msg = f"Duplicate dest key: {dest_key} (from {src_key})" + raise ValueError(msg) + + converted[dest_key] = transformed + diagnostics.append( + ConversionRecord( + source_key=src_key, + dest_key=dest_key, + source_shape=src_shape, + dest_shape=transformed.shape, + transform=transform_name, + ) + ) + + return converted, diagnostics + + +# --------------------------------------------------------------------------- +# Checkpoint loading (PyTorch → numpy) +# --------------------------------------------------------------------------- +def load_pytorch_checkpoint(checkpoint_path: Path) -> dict[str, np.ndarray]: + """Load PyTorch checkpoint and return state_dict as numpy arrays. + + Strips _orig_mod. prefix from torch.compile checkpoints. + + Requires torch to be installed (reference dependency group). + """ + import torch + + raw = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + state_dict = raw.get("state_dict", raw) + + result: dict[str, np.ndarray] = {} + for key, tensor in state_dict.items(): + clean_key = key.removeprefix(COMPILE_PREFIX) + result[clean_key] = tensor.numpy() + + return result + + +# --------------------------------------------------------------------------- +# Safetensors output +# --------------------------------------------------------------------------- +def save_safetensors(weights: OrderedDict[str, np.ndarray], output_path: Path) -> None: + """Save converted weights as safetensors.""" + from safetensors.numpy import save_file + + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(weights, str(output_path)) + + +# --------------------------------------------------------------------------- +# High-level API +# --------------------------------------------------------------------------- +def convert_checkpoint( + checkpoint_path: Path, + output_path: Path, +) -> list[ConversionRecord]: + """Full conversion pipeline: load PyTorch checkpoint → convert → save safetensors. + + Args: + checkpoint_path: Path to PyTorch .pth checkpoint. + output_path: Path for output .safetensors file. + + Returns: + List of conversion diagnostic records. + """ + state_dict = load_pytorch_checkpoint(checkpoint_path) + converted, diagnostics = convert_state_dict(state_dict) + save_safetensors(converted, output_path) + return diagnostics diff --git a/tests/test_conversion.py b/tests/test_conversion.py index b5e7d1f..8d04d35 100644 --- a/tests/test_conversion.py +++ b/tests/test_conversion.py @@ -3,21 +3,185 @@ Validates key mapping, shape transforms, and round-trip integrity. """ +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np import pytest -pytestmark = pytest.mark.skip(reason="Phase 3: converter not yet implemented") +if TYPE_CHECKING: + from collections import OrderedDict + +from corridorkey_mlx.convert.converter import ( + REFINER_STEM_MAP, + SKIP_SUFFIXES, + convert_state_dict, + load_pytorch_checkpoint, +) + +CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth") +SAFETENSORS_PATH = Path("checkpoints/corridorkey_mlx.safetensors") -def test_key_mapping_complete() -> None: +@pytest.fixture(scope="module") +def checkpoint_data() -> ( + tuple[dict[str, np.ndarray], OrderedDict[str, np.ndarray], list] +): + """Load and convert checkpoint once for all tests in this module.""" + if not CHECKPOINT_PATH.exists(): + pytest.skip("Checkpoint not found — need CorridorKey_v1.0.pth") + state_dict = load_pytorch_checkpoint(CHECKPOINT_PATH) + converted, diagnostics = convert_state_dict(state_dict) + return state_dict, converted, diagnostics + + +# --------------------------------------------------------------------------- +# Key mapping completeness +# --------------------------------------------------------------------------- +class TestKeyMapping: """Every PyTorch key maps to an MLX key with no orphans.""" - ... + + def test_no_orphan_source_keys(self, checkpoint_data) -> None: + """All source keys are either mapped or explicitly skipped.""" + state_dict, _, diagnostics = checkpoint_data + + mapped_src_keys = {r.source_key for r in diagnostics} + skipped_keys = {k for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES)} + accounted_for = mapped_src_keys | skipped_keys + + orphans = set(state_dict.keys()) - accounted_for + assert orphans == set(), f"Orphan source keys: {orphans}" + + def test_no_duplicate_dest_keys(self, checkpoint_data) -> None: + """No two source keys map to the same destination key.""" + _, _, diagnostics = checkpoint_data + + dest_keys = [r.dest_key for r in diagnostics] + assert len(dest_keys) == len(set(dest_keys)), "Duplicate destination keys found" + + def test_refiner_stem_remapped(self, checkpoint_data) -> None: + """Refiner stem Sequential keys correctly remapped to named attrs.""" + _, converted, _ = checkpoint_data + + for pt_key, mlx_key in REFINER_STEM_MAP.items(): + assert mlx_key in converted, f"Expected remapped key {mlx_key} not found" + assert pt_key not in converted, f"Original key {pt_key} should not be in output" + + def test_num_batches_tracked_dropped(self, checkpoint_data) -> None: + """BatchNorm num_batches_tracked keys are not in output.""" + _, converted, _ = checkpoint_data + + for key in converted: + assert "num_batches_tracked" not in key, f"Found num_batches_tracked: {key}" + + def test_expected_key_count(self, checkpoint_data) -> None: + """Output has expected number of keys (367 source - 2 skipped = 365).""" + state_dict, converted, _ = checkpoint_data + + source_count = len(state_dict) + skip_count = sum(1 for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES)) + expected = source_count - skip_count + assert len(converted) == expected, f"Expected {expected} keys, got {len(converted)}" -def test_conv_weight_transpose() -> None: +# --------------------------------------------------------------------------- +# Conv weight transpose +# --------------------------------------------------------------------------- +class TestConvWeightTranspose: """Conv weights correctly transposed from NCHW → NHWC.""" - ... + + def test_patch_embed_conv_shape(self, checkpoint_data) -> None: + """Patch embed conv transposed: (112,4,7,7) → (112,7,7,4).""" + _, converted, _ = checkpoint_data + + key = "encoder.model.patch_embed.proj.weight" + assert key in converted + assert converted[key].shape == (112, 7, 7, 4) + + def test_decoder_conv_shapes(self, checkpoint_data) -> None: + """Decoder conv weights transposed correctly.""" + _, converted, _ = checkpoint_data + + # linear_fuse: (256, 1024, 1, 1) → (256, 1, 1, 1024) + for prefix in ("alpha_decoder", "fg_decoder"): + fuse_key = f"{prefix}.linear_fuse.weight" + assert converted[fuse_key].shape == (256, 1, 1, 1024) + + def test_refiner_conv_shapes(self, checkpoint_data) -> None: + """Refiner conv weights transposed correctly.""" + _, converted, _ = checkpoint_data + + # stem_conv: (64, 7, 3, 3) → (64, 3, 3, 7) + assert converted["refiner.stem_conv.weight"].shape == (64, 3, 3, 7) + + # res block convs: (64, 64, 3, 3) → (64, 3, 3, 64) + for i in range(1, 5): + for j in range(1, 3): + key = f"refiner.res{i}.conv{j}.weight" + assert converted[key].shape == (64, 3, 3, 64), f"Wrong shape for {key}" + + # final: (4, 64, 1, 1) → (4, 1, 1, 64) + assert converted["refiner.final.weight"].shape == (4, 1, 1, 64) + + def test_linear_weights_unchanged(self, checkpoint_data) -> None: + """Linear (2D) weights are not transposed.""" + _, _, diagnostics = checkpoint_data + + passthrough_2d = [ + r for r in diagnostics if r.transform == "passthrough" and len(r.source_shape) == 2 + ] + for record in passthrough_2d: + assert record.source_shape == record.dest_shape, ( + f"{record.source_key}: shape changed despite passthrough" + ) -def test_first_conv_4ch_preserved() -> None: +# --------------------------------------------------------------------------- +# 4-channel first conv +# --------------------------------------------------------------------------- +class TestFirstConv4Channel: """Patched 4-channel first conv preserved during conversion.""" - ... + + def test_input_channels_preserved(self, checkpoint_data) -> None: + """Patch embed conv has 4 input channels (RGB + alpha hint).""" + _, converted, _ = checkpoint_data + + weight = converted["encoder.model.patch_embed.proj.weight"] + # MLX layout: (O, H, W, I) — I is last dim + input_channels = weight.shape[-1] + assert input_channels == 4, f"Expected 4 input channels, got {input_channels}" + + def test_values_preserved(self, checkpoint_data) -> None: + """Conv values are preserved (only transposed, not modified).""" + state_dict, converted, _ = checkpoint_data + + pt_weight = state_dict["encoder.model.patch_embed.proj.weight"] + mlx_weight = converted["encoder.model.patch_embed.proj.weight"] + + # Transpose back to PyTorch layout and compare + roundtrip = np.transpose(mlx_weight, (0, 3, 1, 2)) + np.testing.assert_array_equal(roundtrip, pt_weight) + + +# --------------------------------------------------------------------------- +# Safetensors output validation +# --------------------------------------------------------------------------- +class TestSafetensorsOutput: + """Validate saved safetensors file matches conversion output.""" + + def test_safetensors_loadable(self, checkpoint_data) -> None: + """Saved safetensors file can be loaded and has correct keys.""" + if not SAFETENSORS_PATH.exists(): + pytest.skip("Run convert_weights.py first") + + from safetensors.numpy import load_file + + loaded = load_file(str(SAFETENSORS_PATH)) + + _, converted, _ = checkpoint_data + + assert set(loaded.keys()) == set(converted.keys()) + for key in converted: + np.testing.assert_array_equal(loaded[key], converted[key])