diff --git a/docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md b/docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md index 498b4b4..46ebcf3 100644 --- a/docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md +++ b/docs/plans/2026-03-08-feat-mlx-memory-optimizations-plan.md @@ -122,9 +122,9 @@ def __call__(self, x: mx.array) -> dict[str, mx.array]: ### 1.3 Guardrails -- [ ] Verify all `nn.GroupNorm` layers retain `pytorch_compatible=True` after any decoder modifications -- [ ] Verify `nn.BatchNorm` in decoder fusion layer handles bf16 inputs correctly (running stats are fp32 in eval mode — matmul should auto-promote) -- [ ] If mx.compile + bf16 fails (Spike 0a), add `if self._compute_dtype != mx.float32: self._compiled = False` guard +- [x] Verify all `nn.GroupNorm` layers retain `pytorch_compatible=True` after any decoder modifications +- [x] Verify `nn.BatchNorm` in decoder fusion layer handles bf16 inputs correctly (running stats are fp32 in eval mode — matmul should auto-promote) +- [x] If mx.compile + bf16 fails (Spike 0a), add `if self._compute_dtype != mx.float32: self._compiled = False` guard — N/A, spike passed ### 1.4 Test Updates @@ -201,10 +201,10 @@ class FusedDecoderPair(nn.Module): ### 2.3 Constraints -- [ ] `nn.Upsample` instances stay pre-allocated in `__init__` (Phase 6 guardrail) -- [ ] Concatenation on `axis=-1` only (NHWC) -- [ ] Checkpoint loading unaffected (same `alpha_decoder.*` / `fg_decoder.*` keys) -- [ ] mx.compile shape tracing: channel dim changes from 256 to 512 in fused path — verify compile handles this +- [x] `nn.Upsample` instances stay pre-allocated in `__init__` (Phase 6 guardrail) +- [x] Concatenation on `axis=-1` only (NHWC) +- [x] Checkpoint loading unaffected (same `alpha_decoder.*` / `fg_decoder.*` keys) +- [x] mx.compile shape tracing: channel dim changes from 256 to 512 in fused path — verify compile handles this ### 2.4 Fallback Criteria diff --git a/src/corridorkey_mlx/engine.py b/src/corridorkey_mlx/engine.py index 1ef7251..2fbd4c5 100644 --- a/src/corridorkey_mlx/engine.py +++ b/src/corridorkey_mlx/engine.py @@ -6,6 +6,7 @@ mirrors the Torch engine contract expected by the main CorridorKey repository. from __future__ import annotations +import gc import logging import warnings from pathlib import Path @@ -152,6 +153,10 @@ class CorridorKeyMLXEngine: alpha_refined = outputs["alpha_final"] fg_refined = outputs["fg_final"] + # Release unused intermediate tensors (logits, delta_logits, etc.) + del outputs + gc.collect() + if not self._use_refiner or refiner_scale == 0.0: alpha_out = alpha_coarse fg_out = fg_coarse diff --git a/src/corridorkey_mlx/inference/tiling.py b/src/corridorkey_mlx/inference/tiling.py index 84be576..2a323e4 100644 --- a/src/corridorkey_mlx/inference/tiling.py +++ b/src/corridorkey_mlx/inference/tiling.py @@ -6,6 +6,7 @@ then blends results using linear ramp weights in the overlap region. from __future__ import annotations +import gc from typing import TYPE_CHECKING import mlx.core as mx @@ -157,6 +158,11 @@ def tiled_inference( fg_accum[y_start:y_end, x_start:x_end, :] += fg_tile * w3d weight_accum[y_start:y_end, x_start:x_end, :] += w3d + # Deterministic memory cleanup — release MLX graph refs between tiles + del out, tile, alpha_tile, fg_tile + gc.collect() + mx.clear_cache() + # Normalize by accumulated weights weight_accum = np.maximum(weight_accum, 1e-8) alpha_final = mx.array(alpha_accum / weight_accum)[None] # (1, H, W, 1) diff --git a/src/corridorkey_mlx/model/corridorkey.py b/src/corridorkey_mlx/model/corridorkey.py index a121ba7..b5155f3 100644 --- a/src/corridorkey_mlx/model/corridorkey.py +++ b/src/corridorkey_mlx/model/corridorkey.py @@ -13,7 +13,7 @@ import mlx.nn as nn from safetensors import safe_open from corridorkey_mlx.model.backbone import HieraBackbone -from corridorkey_mlx.model.decoder import DecoderHead +from corridorkey_mlx.model.decoder import DecoderHead, FusedDecoderPair from corridorkey_mlx.model.hiera import ENCODER_KEY_PREFIX, _interpolate_pos_embed, _prod from corridorkey_mlx.model.refiner import CNNRefinerModule @@ -31,13 +31,23 @@ class GreenFormer(nn.Module): Output: dict with coarse/final alpha and foreground maps in NHWC. """ - def __init__(self, img_size: int = 512) -> None: + def __init__( + self, + img_size: int = 512, + dtype: mx.Dtype = mx.float32, + fused_decode: bool = False, + ) -> None: super().__init__() + self._compute_dtype = dtype + self._fused_decode = fused_decode self.backbone = HieraBackbone(img_size=img_size) self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1) self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3) self.refiner = CNNRefinerModule() + if fused_decode: + self._fused_pair = FusedDecoderPair(self.alpha_decoder, self.fg_decoder) + # Decoder outputs at stride-4 (H/4, W/4); upsampler is always 4x self._logit_upsampler = nn.Upsample( scale_factor=(4.0, 4.0), mode="linear", align_corners=False @@ -54,22 +64,34 @@ class GreenFormer(nn.Module): alpha_coarse, fg_coarse, delta_logits, alpha_final, fg_final. All tensors in NHWC format. """ - # Backbone -> 4 multiscale feature maps in NHWC + # Backbone always runs in fp32 features = self.backbone(x) + # Cast features to compute dtype for decoders (bf16 saves memory) + if self._compute_dtype != mx.float32: + features = [f.astype(self._compute_dtype) for f in features] + # Decoder heads -> logits at H/4 resolution - alpha_logits = self.alpha_decoder(features) # (B, H/4, W/4, 1) - fg_logits = self.fg_decoder(features) # (B, H/4, W/4, 3) + if self._fused_decode: + alpha_logits, fg_logits = self._fused_pair(features) + else: + alpha_logits = self.alpha_decoder(features) # (B, H/4, W/4, 1) + fg_logits = self.fg_decoder(features) # (B, H/4, W/4, 3) # Upsample logits to full input resolution (4x from stride-4 decoder) alpha_logits_up = self._logit_upsampler(alpha_logits) # (B, H, W, 1) fg_logits_up = self._logit_upsampler(fg_logits) # (B, H, W, 3) - # Coarse predictions via sigmoid + # Cast back to fp32 before sigmoid (precision at saturation boundaries) + if self._compute_dtype != mx.float32: + alpha_logits_up = alpha_logits_up.astype(mx.float32) + fg_logits_up = fg_logits_up.astype(mx.float32) + + # Coarse predictions via sigmoid (always fp32) alpha_coarse = mx.sigmoid(alpha_logits_up) fg_coarse = mx.sigmoid(fg_logits_up) - # Refiner: RGB + coarse predictions -> delta logits + # Refiner receives fp32 coarse predictions rgb = x[:, :, :, :3] # (B, H, W, 3) coarse_pred = mx.concatenate([alpha_coarse, fg_coarse], axis=-1) # (B, H, W, 4) delta_logits = self.refiner(rgb, coarse_pred) # (B, H, W, 4) @@ -79,8 +101,8 @@ class GreenFormer(nn.Module): fg_final = mx.sigmoid(fg_logits_up + delta_logits[:, :, :, 1:4]) return { - "alpha_logits": alpha_logits, - "fg_logits": fg_logits, + "alpha_logits": alpha_logits.astype(mx.float32), + "fg_logits": fg_logits.astype(mx.float32), "alpha_logits_up": alpha_logits_up, "fg_logits_up": fg_logits_up, "alpha_coarse": alpha_coarse, diff --git a/src/corridorkey_mlx/model/decoder.py b/src/corridorkey_mlx/model/decoder.py index fbe4abb..1d70f78 100644 --- a/src/corridorkey_mlx/model/decoder.py +++ b/src/corridorkey_mlx/model/decoder.py @@ -8,9 +8,14 @@ Architecture mirrors nikopueringer/CorridorKey DecoderHead (SegFormer-style). from __future__ import annotations +from typing import TYPE_CHECKING + import mlx.core as mx import mlx.nn as nn +if TYPE_CHECKING: + from collections.abc import Sequence + class MLP(nn.Module): """Single linear projection: input_dim -> embed_dim.""" @@ -98,3 +103,73 @@ class DecoderHead(nn.Module): fused = nn.relu(fused) # No dropout at inference (eval mode) return self.classifier(fused) + + def _project_features(self, features: list[mx.array]) -> list[mx.array]: + """Project features without upsampling or fusion.""" + linears = [self.linear_c1, self.linear_c2, self.linear_c3, self.linear_c4] + projected = [] + for feat, linear in zip(features, linears, strict=True): + b, h, w, _c = feat.shape + x = feat.reshape(b, h * w, _c) + x = linear(x) + x = x.reshape(b, h, w, -1) + projected.append(x) + return projected + + def _fuse_and_classify(self, projected: list[mx.array]) -> mx.array: + """Fuse pre-upsampled projections and classify.""" + fused = mx.concatenate(projected[::-1], axis=-1) + fused = self.linear_fuse(fused) + fused = self.bn(fused) + fused = nn.relu(fused) + return self.classifier(fused) + + +class FusedDecoderPair(nn.Module): + """Runs two DecoderHeads with batched upsampling. + + Concatenates alpha+fg projections along channel axis before each upsample, + reducing 6 Metal dispatch calls to 3. + """ + + def __init__(self, alpha_head: DecoderHead, fg_head: DecoderHead) -> None: + super().__init__() + self.alpha_head = alpha_head + self.fg_head = fg_head + # Reuse alpha_head's pre-allocated upsamplers + self._upsamplers: Sequence[nn.Upsample | None] = [ + None, + alpha_head._upsampler_2x, + alpha_head._upsampler_4x, + alpha_head._upsampler_8x, + ] + + def __call__( + self, features: list[mx.array] + ) -> tuple[mx.array, mx.array]: + """Forward pass with batched upsampling. + + Returns: + (alpha_logits, fg_logits) both in NHWC. + """ + alpha_projs = self.alpha_head._project_features(features) + fg_projs = self.fg_head._project_features(features) + + alpha_up = [] + fg_up = [] + for a_proj, f_proj, up in zip( + alpha_projs, fg_projs, self._upsamplers, strict=True + ): + if up is not None: + # Batch upsample: concat along channel axis (NHWC) + fused = mx.concatenate([a_proj, f_proj], axis=-1) + fused = up(fused) + embed_dim = a_proj.shape[-1] + a_proj = fused[:, :, :, :embed_dim] + f_proj = fused[:, :, :, embed_dim:] + alpha_up.append(a_proj) + fg_up.append(f_proj) + + alpha_logits = self.alpha_head._fuse_and_classify(alpha_up) + fg_logits = self.fg_head._fuse_and_classify(fg_up) + return alpha_logits, fg_logits diff --git a/tests/test_model_contract.py b/tests/test_model_contract.py index 1b10002..f87c31f 100644 --- a/tests/test_model_contract.py +++ b/tests/test_model_contract.py @@ -192,3 +192,101 @@ def test_no_nan_inf() -> None: arr = np.array(out[key]) assert not np.isnan(arr).any(), f"{key} contains NaN" assert not np.isinf(arr).any(), f"{key} contains Inf" + + +# --------------------------------------------------------------------------- +# bf16 mixed precision contract +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def bf16_model_output() -> dict[str, mx.array]: + """Forward pass with bf16 compute dtype, random weights.""" + model = GreenFormer(img_size=SMALL_IMG_SIZE, dtype=mx.bfloat16) + mx.random.seed(42) + x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(x) # noqa: S307 + out = model(x) + mx.eval(out) # noqa: S307 + return out + + +def test_bf16_output_dtype_always_fp32(bf16_model_output: dict[str, mx.array]) -> None: + """All outputs must be fp32 regardless of compute dtype.""" + for key, arr in bf16_model_output.items(): + assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}" + + +def test_bf16_sigmoid_outputs_in_range(bf16_model_output: dict[str, mx.array]) -> None: + """Post-sigmoid outputs must be in [0, 1] even with bf16 compute.""" + for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"): + arr = bf16_model_output[key] + assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0" + assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1" + + +def test_bf16_no_nan_inf(bf16_model_output: dict[str, mx.array]) -> None: + """bf16 forward produces finite outputs.""" + for key in ("alpha_final", "fg_final"): + arr = np.array(bf16_model_output[key]) + assert not np.isnan(arr).any(), f"{key} contains NaN" + assert not np.isinf(arr).any(), f"{key} contains Inf" + + +def test_fused_decode_matches_unfused() -> None: + """Fused decoder pair produces bit-exact output vs independent decoders.""" + from mlx.utils import tree_flatten + + mx.random.seed(77) + x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(x) # noqa: S307 + + model_unfused = GreenFormer(img_size=SMALL_IMG_SIZE, fused_decode=False) + mx.eval(model_unfused.parameters()) # noqa: S307 + + model_fused = GreenFormer(img_size=SMALL_IMG_SIZE, fused_decode=True) + model_fused.load_weights(tree_flatten(model_unfused.parameters())) + mx.eval(model_fused.parameters()) # noqa: S307 + + out_u = model_unfused(x) + out_f = model_fused(x) + mx.eval(out_u) # noqa: S307 + mx.eval(out_f) # noqa: S307 + + for key in out_u: + np.testing.assert_array_equal( + np.array(out_u[key]), + np.array(out_f[key]), + err_msg=f"{key} differs between fused and unfused decode", + ) + + +def test_fp32_default_unchanged() -> None: + """GreenFormer(dtype=mx.float32) behaves identically to GreenFormer().""" + mx.random.seed(99) + x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(x) # noqa: S307 + + model_default = GreenFormer(img_size=SMALL_IMG_SIZE) + model_explicit = GreenFormer(img_size=SMALL_IMG_SIZE, dtype=mx.float32) + + # Same weights via flattened tree + from mlx.utils import tree_flatten + + model_explicit.load_weights(tree_flatten(model_default.parameters())) + mx.eval(model_explicit.parameters()) # noqa: S307 + + out_default = model_default(x) + out_explicit = model_explicit(x) + mx.eval(out_default) # noqa: S307 + mx.eval(out_explicit) # noqa: S307 + + for key in out_default: + np.testing.assert_array_equal( + np.array(out_default[key]), + np.array(out_explicit[key]), + err_msg=f"{key} differs between default and explicit fp32", + )