monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: stage-boundary GC between backbone/decoder/refiner

Browse Source 
mx.eval + gc + clear_cache at backbone→decoder and decoder→refiner
boundaries in eager mode. _compiled flag skips GC under mx.compile.
Compilation tests updated to use compile_model() API.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates 2026-03-09 17:02:38 -02:30
parent ce30eddcff
commit 6787a7fd21
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
3 changed files with 27 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/corridorkey_mlx/inference/pipeline.py
View File

@ -68,6 +68,7 @@ def compile_model(model: GreenFormer, shapeless: bool = False) -> GreenFormer:
same resolution. Shapeless compile is experimental — the backbone uses same resolution. Shapeless compile is experimental — the backbone uses
shape-dependent reshapes that may trigger recompilation. shape-dependent reshapes that may trigger recompilation.
""" """
model._compiled = True
model.__call__ = mx.compile(model.__call__, shapeless=shapeless) # type: ignore[method-assign] model.__call__ = mx.compile(model.__call__, shapeless=shapeless) # type: ignore[method-assign]
return model return model

16
src/corridorkey_mlx/model/corridorkey.py
View File

@ -6,6 +6,7 @@ All internal operations use NHWC layout.
from __future__ import annotations from __future__ import annotations
import gc
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
import mlx.core as mx import mlx.core as mx
@ -42,6 +43,7 @@ class GreenFormer(nn.Module):
self._compute_dtype = dtype self._compute_dtype = dtype
self._fused_decode = fused_decode self._fused_decode = fused_decode
self._slim = slim self._slim = slim
self._compiled = False
self.backbone = HieraBackbone(img_size=img_size) self.backbone = HieraBackbone(img_size=img_size)
self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1) self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1)
self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3) self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3)
@ -69,6 +71,13 @@ class GreenFormer(nn.Module):
# Backbone always runs in fp32 # Backbone always runs in fp32
features = self.backbone(x) features = self.backbone(x)
# Materialize backbone output so MLX can free intermediate graph nodes.
# NOTE: mx.eval is MLX array materialization, not Python eval()
if not self._compiled:
mx.eval(features) # noqa: S307
gc.collect()
mx.clear_cache()
# Cast features to compute dtype for decoders (bf16 saves memory) # Cast features to compute dtype for decoders (bf16 saves memory)
if self._compute_dtype != mx.float32: if self._compute_dtype != mx.float32:
features = [f.astype(self._compute_dtype) for f in features] features = [f.astype(self._compute_dtype) for f in features]
@ -93,6 +102,13 @@ class GreenFormer(nn.Module):
alpha_coarse = mx.sigmoid(alpha_logits_up) alpha_coarse = mx.sigmoid(alpha_logits_up)
fg_coarse = mx.sigmoid(fg_logits_up) fg_coarse = mx.sigmoid(fg_logits_up)
# Materialize decoder output so MLX can free decoder graph nodes.
# NOTE: mx.eval is MLX array materialization, not Python eval()
if not self._compiled:
mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up) # noqa: S307
gc.collect()
mx.clear_cache()
# Refiner receives fp32 coarse predictions # Refiner receives fp32 coarse predictions
rgb = x[:, :, :, :3] # (B, H, W, 3) rgb = x[:, :, :, :3] # (B, H, W, 3)
coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1) # (B, H, W, 4) coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1) # (B, H, W, 4)

19
tests/test_compilation.py
View File

@ -5,6 +5,7 @@ from __future__ import annotations
import mlx.core as mx import mlx.core as mx
import pytest import pytest
from corridorkey_mlx.inference.pipeline import compile_model
from corridorkey_mlx.model.corridorkey import GreenFormer from corridorkey_mlx.model.corridorkey import GreenFormer
IMG_SIZE = 256 IMG_SIZE = 256
@ -32,11 +33,11 @@ def dummy_input() -> mx.array:
def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> None: def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> None:
"""Fixed-shape compiled output matches eager output within tolerance.""" """Fixed-shape compiled output matches eager output within tolerance."""
eager_out = model(dummy_input) eager_out = model(dummy_input)
mx.eval(eager_out) # noqa: S307 mx.eval(eager_out) # noqa: S307 # NOTE: MLX materialization, not Python eval
compiled_fn = mx.compile(model.__call__) compiled_model = compile_model(model)
compiled_out = compiled_fn(dummy_input) compiled_out = compiled_model(dummy_input)
mx.eval(compiled_out) # noqa: S307 mx.eval(compiled_out) # noqa: S307 # NOTE: MLX materialization
for key in OUTPUT_KEYS: for key in OUTPUT_KEYS:
diff = float(mx.max(mx.abs(eager_out[key] - compiled_out[key]))) diff = float(mx.max(mx.abs(eager_out[key] - compiled_out[key])))
@ -45,12 +46,12 @@ def test_compiled_matches_eager(model: GreenFormer, dummy_input: mx.array) -> No
def test_compiled_deterministic(model: GreenFormer, dummy_input: mx.array) -> None: def test_compiled_deterministic(model: GreenFormer, dummy_input: mx.array) -> None:
"""Compiled model produces identical results across consecutive calls.""" """Compiled model produces identical results across consecutive calls."""
compiled_fn = mx.compile(model.__call__) compiled_model = compile_model(model)
out1 = compiled_fn(dummy_input) out1 = compiled_model(dummy_input)
mx.eval(out1) # noqa: S307 mx.eval(out1) # noqa: S307 # NOTE: MLX materialization
out2 = compiled_fn(dummy_input) out2 = compiled_model(dummy_input)
mx.eval(out2) # noqa: S307 mx.eval(out2) # noqa: S307 # NOTE: MLX materialization
for key in OUTPUT_KEYS: for key in OUTPUT_KEYS:
diff = float(mx.max(mx.abs(out1[key] - out2[key]))) diff = float(mx.max(mx.abs(out1[key] - out2[key])))