feat(phase3): PyTorch→MLX weight converter (#1)

* 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-01 05:16:56 -03:30 · 2026-03-01 05:16:56 -03:30 · 150c974eed
commit 150c974eed
parent 091261ccff
6 changed files with 538 additions and 17 deletions
--- 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:**

--- a/docs/plans/2026-03-01-fix-converter-review-feedback-plan.md
+++ 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
--- a/scripts/convert_weights.py
+++ 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()
--- 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",
+]
--- 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
--- 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])