From 91e3a2832d85230699cb751eee1d9938f615890d Mon Sep 17 00:00:00 2001 From: betweentwomidnights Date: Sun, 14 Jun 2026 16:03:33 -0700 Subject: [PATCH 01/28] Add MLX LoRA training and audio encoding primitives Add pure-MLX LoRA, DoRA, BoRA, and XS adapter injection, checkpoint interoperability, fixed-strength inference support, waveform-to-SAME-latent encoding, SA3 timestep sampling, distribution shifting, rectified-flow loss, and focused parity tests. --- optimized/mlx/README.md | 52 ++ optimized/mlx/models/defs/audio_encoding.py | 285 ++++++ optimized/mlx/models/defs/lora.py | 970 ++++++++++++++++++++ optimized/mlx/models/defs/training.py | 224 +++++ tests/test_mlx_audio_encoding.py | 121 +++ tests/test_mlx_lora.py | 368 ++++++++ tests/test_mlx_training.py | 69 ++ 7 files changed, 2089 insertions(+) create mode 100644 optimized/mlx/models/defs/audio_encoding.py create mode 100644 optimized/mlx/models/defs/lora.py create mode 100644 optimized/mlx/models/defs/training.py create mode 100644 tests/test_mlx_audio_encoding.py create mode 100644 tests/test_mlx_lora.py create mode 100644 tests/test_mlx_training.py diff --git a/optimized/mlx/README.md b/optimized/mlx/README.md index 70c4069..7bcc9b9 100644 --- a/optimized/mlx/README.md +++ b/optimized/mlx/README.md @@ -191,6 +191,58 @@ Sample run on **M4 Pro / 48 GB**: └───────────┴─────────┴─────────┴───────────┴───────────┴────────────┘ ``` +## LoRA training primitives + +The optimized runtime includes MLX counterparts to the adapter family already +supported by the official PyTorch implementation: standard LoRA, DoRA, BoRA, +and LoRA-XS variants. + +```python +from models.defs.lora import ( + apply_lora_checkpoint, + inject_trainable_lora, + save_lora_checkpoint, +) +from models.defs.training import ( + rectified_flow_loss, + sample_training_timesteps, + shift_training_timesteps, +) +from models.defs.audio_encoding import encode_audio +``` + +`inject_trainable_lora()` freezes the base model and leaves only adapter +parameters trainable through `mlx.nn.value_and_grad`. `save_lora_checkpoint()` +writes the same safetensors keys and `lora_config` metadata used by the +PyTorch trainer. `apply_lora_checkpoint()` loads that format through MLX +without adding PyTorch or safetensors as runtime dependencies. Application +materializes a fixed strength into the loaded model's in-memory weights; it +does not modify the base checkpoint on disk. Reload the base model before +applying a different strength rather than applying repeatedly to the same +instance. + +`encode_audio()` provides the waveform-to-latent bridge used by training. It +accepts a batch of already-decoded 44.1 kHz stereo waveforms, applies SAME's +patched pretransform, pads to the encoder's required alignment, optionally +encodes long clips in overlapping chunks, and returns both latents and a +ceiling-scaled padding mask: + +```python +encoded = encode_audio( + same_l_encoder, + audio, # [batch, 2, samples] + valid_sample_lengths=[samples_a, samples_b], + pad_modulo=16, # 32 for SAME-S + chunked=True, +) +latents = encoded.latents +loss_mask = encoded.padding_mask +``` + +These are model-level primitives rather than a dataset or training CLI. +Callers remain responsible for file decoding and resampling, dataset +discovery and persistence, conditioning, optimization, and checkpoint cadence. + ## Flag reference | Flag | Default | Notes | diff --git a/optimized/mlx/models/defs/audio_encoding.py b/optimized/mlx/models/defs/audio_encoding.py new file mode 100644 index 0000000..c636670 --- /dev/null +++ b/optimized/mlx/models/defs/audio_encoding.py @@ -0,0 +1,285 @@ +"""Waveform-to-latent encoding utilities for the standalone MLX models.""" + +from __future__ import annotations + +import math +import typing as tp +from dataclasses import dataclass + +import mlx.core as mx +import numpy as np + + +PATCH_SIZE = 256 +ENCODER_STRIDE = 16 +SAMPLES_PER_LATENT = PATCH_SIZE * ENCODER_STRIDE +SAME_ENCODER_PAD_MODULO = { + "same-s": 32, + "same-l": 16, +} + + +@dataclass(frozen=True) +class EncodedAudio: + """Latents and validity metadata produced from a padded waveform batch.""" + + latents: mx.array + padding_mask: mx.array + valid_latent_lengths: tuple[int, ...] + source_samples: int + padded_samples: int + + +def patch_audio(audio, *, patch_size: int = PATCH_SIZE): + """Apply SAME's patched pretransform to a ``[B, C, T]`` waveform batch.""" + + if patch_size <= 0: + raise ValueError("patch_size must be positive.") + if audio.ndim != 3: + raise ValueError(f"Expected audio shaped [B, C, T], got {audio.shape}.") + batch, channels, samples = (int(value) for value in audio.shape) + if samples % patch_size != 0: + raise ValueError( + f"Audio length {samples} must be divisible by patch_size={patch_size}." + ) + patch_count = samples // patch_size + return ( + audio.reshape( + batch, + channels, + patch_count, + patch_size, + ) + .transpose(0, 1, 3, 2) + .reshape( + batch, + channels * patch_size, + patch_count, + ) + ) + + +def encode_audio( + encoder, + audio, + *, + valid_sample_lengths: tp.Sequence[int] | np.ndarray | None = None, + pad_modulo: int = SAME_ENCODER_PAD_MODULO["same-l"], + patch_size: int = PATCH_SIZE, + encoder_stride: int = ENCODER_STRIDE, + chunked: bool = False, + chunk_size: int = 128, + overlap: int = 32, + chunk_batch_size: int = 1, +) -> EncodedAudio: + """Encode a waveform batch with a SAME-S or SAME-L MLX encoder. + + The caller owns file decoding, resampling, and channel conversion. ``audio`` + must be a batch of equal-length waveforms shaped ``[B, C, T]``. The current + SAME encoders expect 44.1 kHz stereo input. + """ + + audio = mx.array(audio) + if audio.ndim != 3: + raise ValueError(f"Expected audio shaped [B, C, T], got {audio.shape}.") + batch_size, channels, source_samples = (int(value) for value in audio.shape) + if batch_size < 1 or source_samples < 1: + raise ValueError("Audio batches and waveforms must be non-empty.") + if channels != 2: + raise ValueError(f"SAME encoders expect stereo audio, got {channels} channels.") + if patch_size <= 0 or encoder_stride <= 0: + raise ValueError("patch_size and encoder_stride must be positive.") + if pad_modulo <= 0 or pad_modulo % encoder_stride != 0: + raise ValueError("pad_modulo must be a positive multiple of encoder_stride.") + + valid_lengths = _normalize_valid_lengths( + valid_sample_lengths, + batch_size=batch_size, + source_samples=source_samples, + ) + sample_alignment = patch_size * pad_modulo + padded_samples = math.ceil(source_samples / sample_alignment) * sample_alignment + if padded_samples != source_samples: + audio = mx.pad( + audio, + ((0, 0), (0, 0), (0, padded_samples - source_samples)), + ) + + patches = patch_audio(audio, patch_size=patch_size) + latents = _encode_patches( + encoder, + patches, + encoder_stride=encoder_stride, + pad_modulo=pad_modulo, + chunked=chunked, + chunk_size=chunk_size, + overlap=overlap, + chunk_batch_size=chunk_batch_size, + ) + + samples_per_latent = patch_size * encoder_stride + valid_latent_lengths = tuple( + min( + int(latents.shape[-1]), + math.ceil(length / samples_per_latent), + ) + for length in valid_lengths + ) + positions = mx.arange(int(latents.shape[-1]))[None, :] + padding_mask = positions < mx.array(valid_latent_lengths)[:, None] + return EncodedAudio( + latents=latents, + padding_mask=padding_mask, + valid_latent_lengths=valid_latent_lengths, + source_samples=source_samples, + padded_samples=padded_samples, + ) + + +def _normalize_valid_lengths( + valid_sample_lengths: tp.Sequence[int] | np.ndarray | None, + *, + batch_size: int, + source_samples: int, +) -> tuple[int, ...]: + if valid_sample_lengths is None: + return (source_samples,) * batch_size + values = tuple(int(value) for value in valid_sample_lengths) + if len(values) != batch_size: + raise ValueError( + "valid_sample_lengths must contain one value per audio batch item." + ) + if any(value < 0 or value > source_samples for value in values): + raise ValueError( + f"Valid sample lengths must be between 0 and {source_samples}." + ) + return values + + +def _encode_patches( + encoder, + patches, + *, + encoder_stride: int, + pad_modulo: int, + chunked: bool, + chunk_size: int, + overlap: int, + chunk_batch_size: int, +): + total_patches = int(patches.shape[-1]) + if total_patches % encoder_stride != 0: + raise ValueError( + f"Patch length {total_patches} must be divisible by " + f"encoder_stride={encoder_stride}." + ) + if total_patches % pad_modulo != 0: + raise ValueError( + f"Patch length {total_patches} must be divisible by " + f"pad_modulo={pad_modulo}." + ) + + total_latents = total_patches // encoder_stride + if not chunked or total_latents <= chunk_size: + return encoder(patches) + if chunk_size < 1: + raise ValueError("chunk_size must be positive.") + if overlap < 0 or overlap >= chunk_size: + raise ValueError("overlap must be non-negative and smaller than chunk_size.") + if chunk_batch_size < 1: + raise ValueError("chunk_batch_size must be positive.") + chunk_patches = chunk_size * encoder_stride + if chunk_patches % pad_modulo != 0: + raise ValueError( + "chunk_size produces a patch length incompatible with pad_modulo." + ) + + hop_latents = chunk_size - overlap + chunk_starts = list(range(0, total_latents - chunk_size + 1, hop_latents)) + final_start = total_latents - chunk_size + if chunk_starts[-1] != final_start: + chunk_starts.append(final_start) + + batch_size = int(patches.shape[0]) + encoded_chunks = [] + for offset in range(0, len(chunk_starts), chunk_batch_size): + batch_starts = chunk_starts[offset : offset + chunk_batch_size] + chunk_inputs = mx.concatenate( + [ + patches[ + ..., + start * encoder_stride : (start + chunk_size) * encoder_stride, + ] + for start in batch_starts + ], + axis=0, + ) + encoded_batch = encoder(chunk_inputs) + mx.eval(encoded_batch) + encoded_chunks.extend( + encoded_batch[index * batch_size : (index + 1) * batch_size] + for index in range(len(batch_starts)) + ) + + return _stitch_encoded_chunks( + encoded_chunks, + chunk_starts=chunk_starts, + total_latents=total_latents, + chunk_size=chunk_size, + overlap=overlap, + ) + + +def _stitch_encoded_chunks( + encoded_chunks, + *, + chunk_starts: list[int], + total_latents: int, + chunk_size: int, + overlap: int, +): + half_overlap = overlap // 2 + intervals = [] + last_index = len(chunk_starts) - 1 + for index, (start, chunk) in enumerate( + zip(chunk_starts, encoded_chunks, strict=True) + ): + is_first = index == 0 + is_last = index == last_index + output_start = total_latents - chunk_size if is_last else start + left = 0 if is_first else half_overlap + right = chunk_size if is_last else chunk_size - half_overlap + intervals.append((output_start + left, output_start + right, left, chunk)) + + pieces = [] + cursor = 0 + output_shape = tuple(int(value) for value in encoded_chunks[0].shape[:-1]) + for index, (target_start, target_end, left, chunk) in enumerate(intervals): + next_start = ( + intervals[index + 1][0] if index + 1 < len(intervals) else target_end + ) + target_end = min(target_end, next_start) + clipped_start = max(target_start, cursor) + if clipped_start > cursor: + pieces.append( + mx.zeros( + (*output_shape, clipped_start - cursor), + dtype=encoded_chunks[0].dtype, + ) + ) + cursor = clipped_start + if target_end <= clipped_start: + continue + source_start = left + clipped_start - target_start + source_end = source_start + target_end - clipped_start + pieces.append(chunk[..., source_start:source_end]) + cursor = target_end + + if cursor < total_latents: + pieces.append( + mx.zeros( + (*output_shape, total_latents - cursor), + dtype=encoded_chunks[0].dtype, + ) + ) + return mx.concatenate(pieces, axis=-1) diff --git a/optimized/mlx/models/defs/lora.py b/optimized/mlx/models/defs/lora.py new file mode 100644 index 0000000..8628fa2 --- /dev/null +++ b/optimized/mlx/models/defs/lora.py @@ -0,0 +1,970 @@ +"""MLX LoRA-family adapters compatible with Stable Audio 3 checkpoints. + +This module intentionally depends only on MLX and NumPy so it can be used by +the standalone optimized MLX runtime without pulling in PyTorch or safetensors. +""" + +from __future__ import annotations + +import json +import math +import re +import typing as tp +from dataclasses import dataclass +from itertools import product +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +from mlx.utils import tree_flatten, tree_unflatten + + +_LORA_KEY_RE = re.compile( + r"^(?P.+)\.parametrizations\.weight\.(?P\d+)\." + r"(?Plora_A|lora_B|M_xs|magnitude|magnitude_r|magnitude_c|U|V)$" +) +_XS_ADAPTER_TYPES = { + "lora-xs", + "dora-rows-xs", + "dora-cols-xs", + "bora-xs", +} +_SUPPORTED_ADAPTER_TYPES = { + "lora", + "dora-rows", + "dora-cols", + "bora", + *_XS_ADAPTER_TYPES, +} +_FULL_WEIGHT_ADAPTER_TYPES = _SUPPORTED_ADAPTER_TYPES - {"lora"} + + +@dataclass(frozen=True) +class LoRAInjectionReport: + layer_names: tuple[str, ...] + trainable_parameters: int + adapter_type: str + + @property + def layer_count(self) -> int: + return len(self.layer_names) + + +@dataclass(frozen=True) +class LoRAApplyReport: + path: str + adapter_type: str + loaded_layers: int + applied_layers: int + missing_targets: tuple[str, ...] = () + skipped_layers: tuple[str, ...] = () + + +class LoRALinear(nn.Module): + """Trainable LoRA-family wrapper for an MLX Linear layer.""" + + def __init__( + self, + base: nn.Linear, + *, + rank: int, + alpha: float, + source_name: str, + adapter_type: str = "lora", + ): + super().__init__() + if rank <= 0: + raise ValueError(f"LoRA rank must be positive, got {rank}.") + + self.base = base + self.base.freeze() + self.rank = int(rank) + self.alpha = float(alpha) + self.scaling = self.alpha / self.rank + self.source_name = str(source_name) + self.checkpoint_name = _checkpoint_layer_name(self.source_name) + self.adapter_type = canonical_adapter_type(adapter_type) + + fan_out, fan_in = (int(value) for value in base.weight.shape) + source_weight = _linear_source_weight_2d(base.weight) + _validate_rank( + self.rank, + fan_out=fan_out, + fan_in=fan_in, + source_name=self.source_name, + ) + _initialize_adapter(self, source_weight, fan_out=fan_out, fan_in=fan_in) + + def __call__(self, x): + if self.adapter_type in _FULL_WEIGHT_ADAPTER_TYPES: + adapted_weight = _adapted_weight_2d( + _linear_source_weight_2d(self.base.weight), + adapter_type=self.adapter_type, + layer=self, + ) + output = x.astype(mx.float32) @ adapted_weight.T + bias = getattr(self.base, "bias", None) + if bias is not None: + output = output + bias.astype(mx.float32) + return output.astype(x.dtype) + + base_output = self.base(x) + adapter_output = (x.astype(mx.float32) @ self.lora_A.T) @ self.lora_B.T + return base_output + (adapter_output * self.scaling).astype(base_output.dtype) + + +class LoRAConv1d(nn.Module): + """Trainable LoRA-family wrapper for an MLX Conv1d layer.""" + + def __init__( + self, + base: nn.Conv1d, + *, + rank: int, + alpha: float, + source_name: str, + adapter_type: str = "lora", + ): + super().__init__() + if rank <= 0: + raise ValueError(f"LoRA rank must be positive, got {rank}.") + + self.base = base + self.base.freeze() + self.rank = int(rank) + self.alpha = float(alpha) + self.scaling = self.alpha / self.rank + self.source_name = str(source_name) + self.checkpoint_name = _checkpoint_layer_name(self.source_name) + self.adapter_type = canonical_adapter_type(adapter_type) + + fan_out, kernel_size, fan_in_per_group = ( + int(value) for value in base.weight.shape + ) + fan_in = fan_in_per_group * kernel_size + source_weight = _conv1d_source_weight_2d(base.weight) + _validate_rank( + self.rank, + fan_out=fan_out, + fan_in=fan_in, + source_name=self.source_name, + ) + _initialize_adapter(self, source_weight, fan_out=fan_out, fan_in=fan_in) + + def __call__(self, x): + fan_out, kernel_size, fan_in_per_group = ( + int(value) for value in self.base.weight.shape + ) + + if self.adapter_type in _FULL_WEIGHT_ADAPTER_TYPES: + adapted_source = _adapted_weight_2d( + _conv1d_source_weight_2d(self.base.weight), + adapter_type=self.adapter_type, + layer=self, + ) + adapted_weight = _conv1d_weight_from_source_2d( + adapted_source, + fan_out=fan_out, + fan_in_per_group=fan_in_per_group, + kernel_size=kernel_size, + ) + output = mx.conv1d( + x.astype(mx.float32), + adapted_weight, + self.base.stride, + self.base.padding, + self.base.dilation, + self.base.groups, + ) + bias = getattr(self.base, "bias", None) + if bias is not None: + output = output + bias.astype(mx.float32) + return output.astype(x.dtype) + + base_output = self.base(x) + delta_weight = _conv1d_weight_from_source_2d( + self.lora_B @ self.lora_A, + fan_out=fan_out, + fan_in_per_group=fan_in_per_group, + kernel_size=kernel_size, + ) + adapter_output = mx.conv1d( + x.astype(mx.float32), + delta_weight, + self.base.stride, + self.base.padding, + self.base.dilation, + self.base.groups, + ) + return base_output + (adapter_output * self.scaling).astype(base_output.dtype) + + +TrainableLoRALayer = LoRALinear | LoRAConv1d + + +def inject_trainable_lora( + model: nn.Module, + *, + rank: int = 16, + alpha: float | None = None, + include: tp.Sequence[str] | None = None, + exclude: tp.Sequence[str] | None = None, + adapter_type: str = "lora", +) -> LoRAInjectionReport: + """Freeze an MLX model and replace selected Linear/Conv1d layers.""" + + alpha = float(rank if alpha is None else alpha) + adapter_type = canonical_adapter_type(adapter_type) + model.freeze() + + replacements: list[tuple[str, TrainableLoRALayer]] = [] + for name, layer in model.named_modules(): + if not name or not _name_is_selected(name, include=include, exclude=exclude): + continue + if isinstance(layer, nn.Linear): + replacement = LoRALinear( + layer, + rank=rank, + alpha=alpha, + source_name=name, + adapter_type=adapter_type, + ) + elif isinstance(layer, nn.Conv1d): + replacement = LoRAConv1d( + layer, + rank=rank, + alpha=alpha, + source_name=name, + adapter_type=adapter_type, + ) + else: + continue + replacements.append((name, replacement)) + + if not replacements: + raise ValueError("No MLX Linear or Conv1d layers matched the LoRA filters.") + + model.update_modules(tree_unflatten(replacements)) + trainable_parameters = sum( + int(value.size) for _, value in tree_flatten(model.trainable_parameters()) + ) + return LoRAInjectionReport( + layer_names=tuple(name for name, _ in replacements), + trainable_parameters=trainable_parameters, + adapter_type=adapter_type, + ) + + +def iter_trainable_lora_layers( + model: nn.Module, +) -> tp.Iterator[TrainableLoRALayer]: + for _, layer in model.named_modules(): + if isinstance(layer, (LoRALinear, LoRAConv1d)): + yield layer + + +def save_lora_checkpoint( + model: nn.Module, + path: str | Path, + *, + include: tp.Sequence[str] | None = None, + exclude: tp.Sequence[str] | None = None, + extra_config: dict[str, tp.Any] | None = None, +) -> Path: + """Save trainable adapters using the official SA3 safetensors contract.""" + + layers = list(iter_trainable_lora_layers(model)) + if not layers: + raise ValueError("The model has no trainable MLX LoRA layers to save.") + + ranks = {layer.rank for layer in layers} + alphas = {layer.alpha for layer in layers} + adapter_types = {layer.adapter_type for layer in layers} + if len(ranks) != 1 or len(alphas) != 1 or len(adapter_types) != 1: + raise ValueError("A checkpoint must use one rank, alpha, and adapter type.") + + rank = next(iter(ranks)) + alpha = next(iter(alphas)) + adapter_type = next(iter(adapter_types)) + state_dict: dict[str, mx.array] = {} + for layer in layers: + prefix = f"{layer.checkpoint_name}.parametrizations.weight.0" + if adapter_type in _XS_ADAPTER_TYPES: + state_dict[f"{prefix}.M_xs"] = layer.M_xs.astype(mx.float16) + else: + state_dict[f"{prefix}.lora_A"] = layer.lora_A.astype(mx.float16) + state_dict[f"{prefix}.lora_B"] = layer.lora_B.astype(mx.float16) + + if adapter_type in { + "dora-rows", + "dora-cols", + "dora-rows-xs", + "dora-cols-xs", + }: + state_dict[f"{prefix}.magnitude"] = layer.magnitude.astype(mx.float16) + elif adapter_type in {"bora", "bora-xs"}: + state_dict[f"{prefix}.magnitude_r"] = layer.magnitude_r.astype(mx.float16) + state_dict[f"{prefix}.magnitude_c"] = layer.magnitude_c.astype(mx.float16) + + config: dict[str, tp.Any] = { + "rank": rank, + "alpha": alpha, + "adapter_type": adapter_type, + "include": list(include) if include else None, + "exclude": list(exclude) if exclude else None, + } + if extra_config: + protected = {"rank", "alpha", "adapter_type", "include", "exclude"} + overlap = protected.intersection(extra_config) + if overlap: + raise ValueError( + "extra_config cannot override checkpoint fields: " + + ", ".join(sorted(overlap)) + ) + config.update(extra_config) + + output_path = Path(path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + mx.save_safetensors( + str(output_path), + state_dict, + metadata={"lora_config": json.dumps(config)}, + ) + return output_path + + +def load_lora_checkpoint( + path: str | Path, +) -> tuple[dict[str, mx.array], dict[str, tp.Any]]: + """Load an SA3 LoRA safetensors checkpoint without PyTorch.""" + + checkpoint_path = Path(path).expanduser().resolve() + if checkpoint_path.suffix != ".safetensors": + raise ValueError("The standalone MLX runtime supports .safetensors LoRAs.") + state_dict, metadata = mx.load(str(checkpoint_path), return_metadata=True) + config = {} + if metadata and metadata.get("lora_config"): + config = json.loads(metadata["lora_config"]) + return dict(state_dict), config + + +def apply_lora_checkpoint( + model: nn.Module, + path: str | Path, + *, + strength: float = 1.0, +) -> LoRAApplyReport: + """Materialize one checkpoint into a loaded MLX model at a fixed strength. + + Only the model's in-memory weights are changed; the base checkpoint on disk + is untouched. The operation is not reversible because the original target + weights are not retained. Reload the base model before applying a different + strength instead of calling this function repeatedly on the same instance. + """ + + state_dict, config = load_lora_checkpoint(path) + adapter_type = _adapter_type_from_state( + config.get("adapter_type", "lora"), + state_dict, + ) + if adapter_type not in _SUPPORTED_ADAPTER_TYPES: + raise ValueError(f"Unsupported MLX LoRA adapter type: {adapter_type!r}") + + grouped = _group_lora_state_dict(state_dict) + target_params = dict(tree_flatten(model.parameters())) + target_keys = tuple(target_params) + missing_targets: list[str] = [] + skipped_layers: list[str] = [] + applied_layers = 0 + + global_rank = int(config.get("rank") or _infer_global_rank(grouped) or 0) + alpha_value = config.get("alpha", config.get("lora_alpha")) + alpha = float(alpha_value if alpha_value is not None else (global_rank or 1)) + + for source_name, params in grouped.items(): + target_key = _resolve_target_key(f"{source_name}.weight", target_keys) + if target_key is None: + missing_targets.append(source_name) + continue + try: + adapted = _apply_checkpoint_layer( + target_params[target_key], + params, + adapter_type=adapter_type, + alpha=alpha, + strength=float(strength), + ) + except ValueError as exc: + skipped_layers.append(f"{source_name}: {exc}") + continue + + target_dtype = target_params[target_key].dtype + updated = mx.array(adapted) + if updated.dtype != target_dtype: + updated = updated.astype(target_dtype) + model.update(tree_unflatten([(target_key, updated)])) + target_params[target_key] = updated + applied_layers += 1 + + if applied_layers: + mx.eval(model.parameters()) + + return LoRAApplyReport( + path=str(Path(path).expanduser().resolve()), + adapter_type=adapter_type, + loaded_layers=len(grouped), + applied_layers=applied_layers, + missing_targets=tuple(sorted(set(missing_targets))), + skipped_layers=tuple(skipped_layers), + ) + + +def apply_lora_checkpoints( + model: nn.Module, + paths: tp.Sequence[str | Path], + *, + strengths: float | tp.Sequence[float] = 1.0, +) -> tuple[LoRAApplyReport, ...]: + """Materialize an ordered checkpoint stack into a loaded MLX model. + + This is a fixed-strength, in-place operation. In particular, DoRA and BoRA + composition is order-dependent. Callers that need mutable strengths should + retain canonical base weights and rebuild the complete ordered stack from + those values rather than applying updates cumulatively. + """ + + if isinstance(strengths, (int, float)): + values = [float(strengths)] * len(paths) + else: + values = [float(value) for value in strengths] + if len(values) == 1: + values *= len(paths) + if len(values) != len(paths): + raise ValueError( + f"Expected 1 or {len(paths)} strengths, got {len(values)}." + ) + + return tuple( + apply_lora_checkpoint(model, path, strength=strength) + for path, strength in zip(paths, values, strict=True) + ) + + +def canonical_adapter_type(adapter_type: str) -> str: + adapter_type = str(adapter_type or "lora").strip().lower() + aliases = { + "dora": "dora-rows", + "dora-xs": "dora-rows-xs", + "xs": "lora-xs", + } + adapter_type = aliases.get(adapter_type, adapter_type) + if adapter_type not in _SUPPORTED_ADAPTER_TYPES: + supported = ", ".join(sorted(_SUPPORTED_ADAPTER_TYPES)) + raise ValueError( + f"Unsupported MLX adapter type {adapter_type!r}. Expected one of: " + f"{supported}." + ) + return adapter_type + + +def _initialize_adapter(layer, source_weight, *, fan_out: int, fan_in: int) -> None: + if layer.adapter_type in _XS_ADAPTER_TYPES: + layer.U, layer.V = _svd_bases(source_weight, layer.rank) + layer.M_xs = mx.zeros((layer.rank, layer.rank), dtype=mx.float32) + layer.freeze(keys=["U", "V"], recurse=False) + else: + init_scale = 1.0 / math.sqrt(fan_in) + layer.lora_A = mx.random.uniform( + low=-init_scale, + high=init_scale, + shape=(layer.rank, fan_in), + dtype=mx.float32, + ) + layer.lora_B = mx.zeros((fan_out, layer.rank), dtype=mx.float32) + + if layer.adapter_type in {"dora-rows", "dora-rows-xs"}: + layer.magnitude = _row_norms(source_weight) + elif layer.adapter_type in {"dora-cols", "dora-cols-xs"}: + layer.magnitude = _column_norms(source_weight) + elif layer.adapter_type in {"bora", "bora-xs"}: + layer.magnitude_r = _row_norms(source_weight) + layer.magnitude_c = _column_norms(source_weight) + + +def _apply_checkpoint_layer( + target_weight, + params: dict[str, np.ndarray], + *, + adapter_type: str, + alpha: float, + strength: float, +) -> np.ndarray: + target = np.asarray(target_weight, dtype=np.float32) + if strength == 0: + return target + + if adapter_type in _XS_ADAPTER_TYPES: + source_shape = _source_shape_for_xs(tuple(target.shape), params) + else: + delta, _ = _lora_delta_2d(params) + source_shape = _source_shape_for_delta(tuple(target.shape), delta.shape) + + source = _target_to_source_weight(target, source_shape) + base_2d = source.reshape(source_shape[0], -1).astype(np.float32, copy=False) + if adapter_type in _XS_ADAPTER_TYPES: + delta, rank = _xs_delta_2d(params, base_2d) + else: + delta, rank = _lora_delta_2d(params) + + value = base_2d + (float(alpha) / rank) * strength * delta + if adapter_type in {"lora", "lora-xs"}: + adapted = value + elif adapter_type in {"dora-rows", "dora-rows-xs"}: + adapted = _dora_weight_2d( + value, + magnitude=_require_param(params, "magnitude").reshape(-1), + norm_dim=1, + ) + elif adapter_type in {"dora-cols", "dora-cols-xs"}: + adapted = _dora_weight_2d( + value, + magnitude=_require_param(params, "magnitude").reshape(-1), + norm_dim=0, + ) + else: + adapted = _bora_weight_2d( + value, + magnitude_r=_require_param(params, "magnitude_r").reshape(-1), + magnitude_c=_require_param(params, "magnitude_c").reshape(-1), + ) + + return _source_to_target_weight(adapted.reshape(source_shape), target.shape) + + +def _adapted_weight_2d(weight_2d, *, adapter_type: str, layer): + value = weight_2d.astype(mx.float32) + _adapter_delta_2d(layer) * float( + layer.scaling + ) + if adapter_type in {"lora", "lora-xs"}: + return value + if adapter_type in {"dora-rows", "dora-rows-xs"}: + return _dora_weight_2d(value, magnitude=layer.magnitude, norm_dim=1) + if adapter_type in {"dora-cols", "dora-cols-xs"}: + return _dora_weight_2d(value, magnitude=layer.magnitude, norm_dim=0) + return _bora_weight_2d( + value, + magnitude_r=layer.magnitude_r, + magnitude_c=layer.magnitude_c, + ) + + +def _adapter_delta_2d(layer): + if layer.adapter_type in _XS_ADAPTER_TYPES: + return layer.U @ layer.M_xs.astype(mx.float32) @ layer.V.T + return layer.lora_B @ layer.lora_A + + +def _linear_source_weight_2d(weight): + return weight.astype(mx.float32) + + +def _conv1d_source_weight_2d(weight): + fan_out, kernel_size, fan_in_per_group = (int(value) for value in weight.shape) + return ( + weight.astype(mx.float32) + .transpose(0, 2, 1) + .reshape( + fan_out, + fan_in_per_group * kernel_size, + ) + ) + + +def _conv1d_weight_from_source_2d( + source, + *, + fan_out: int, + fan_in_per_group: int, + kernel_size: int, +): + return source.reshape( + fan_out, + fan_in_per_group, + kernel_size, + ).transpose(0, 2, 1) + + +def _row_norms(weight_2d): + return mx.sqrt(mx.sum(weight_2d.astype(mx.float32) ** 2, axis=1)).astype(mx.float32) + + +def _column_norms(weight_2d): + return mx.sqrt(mx.sum(weight_2d.astype(mx.float32) ** 2, axis=0)).astype(mx.float32) + + +def _dora_weight_2d(value, *, magnitude, norm_dim: int): + if isinstance(value, np.ndarray): + norms = np.linalg.norm(value, axis=norm_dim, keepdims=True) + value_hat = value / np.maximum(norms, 1e-12) + if norm_dim == 1: + if magnitude.shape[0] != value.shape[0]: + raise ValueError("DoRA row magnitude does not match the weight.") + return value_hat * magnitude[:, None] + if magnitude.shape[0] != value.shape[1]: + raise ValueError("DoRA column magnitude does not match the weight.") + return value_hat * magnitude[None, :] + + norms = mx.sqrt(mx.sum(value**2, axis=norm_dim, keepdims=True)) + value_hat = value / mx.maximum(norms, 1e-12) + if norm_dim == 1: + return value_hat * magnitude.astype(mx.float32)[:, None] + return value_hat * magnitude.astype(mx.float32)[None, :] + + +def _bora_weight_2d(value, *, magnitude_r, magnitude_c): + if isinstance(value, np.ndarray): + if magnitude_r.shape[0] != value.shape[0]: + raise ValueError("BoRA row magnitude does not match the weight.") + if magnitude_c.shape[0] != value.shape[1]: + raise ValueError("BoRA column magnitude does not match the weight.") + row_norms = np.linalg.norm(value, axis=1, keepdims=True) + row_scaled = value / np.maximum(row_norms, 1e-12) + row_scaled *= magnitude_r[:, None] + column_norms = np.linalg.norm(row_scaled, axis=0, keepdims=True) + return (row_scaled / np.maximum(column_norms, 1e-12)) * magnitude_c[None, :] + + row_norms = mx.sqrt(mx.sum(value**2, axis=1, keepdims=True)) + row_scaled = (value / mx.maximum(row_norms, 1e-12)) * magnitude_r.astype( + mx.float32 + )[:, None] + column_norms = mx.sqrt(mx.sum(row_scaled**2, axis=0, keepdims=True)) + return (row_scaled / mx.maximum(column_norms, 1e-12)) * magnitude_c.astype( + mx.float32 + )[None, :] + + +def _group_lora_state_dict( + state_dict: dict[str, tp.Any], +) -> dict[str, dict[str, np.ndarray]]: + grouped: dict[str, dict[str, np.ndarray]] = {} + for key, value in state_dict.items(): + match = _LORA_KEY_RE.match(key) + if match is None: + continue + grouped.setdefault(match.group("prefix"), {})[match.group("param")] = ( + np.asarray(value, dtype=np.float32) + ) + return grouped + + +def _adapter_type_from_state( + adapter_type: str, + state_dict: dict[str, tp.Any], +) -> str: + raw_type = str(adapter_type or "lora").strip().lower() + keys = tuple(state_dict) + has_xs = any(key.endswith(".M_xs") for key in keys) + if has_xs: + if raw_type in {"bora", "bora-xs"} or any( + key.endswith((".magnitude_r", ".magnitude_c")) for key in keys + ): + return "bora-xs" + if raw_type in {"dora-cols", "dora-cols-xs"}: + return "dora-cols-xs" + if raw_type in {"dora", "dora-rows", "dora-rows-xs"} or any( + key.endswith(".magnitude") for key in keys + ): + return "dora-rows-xs" + return "lora-xs" + if raw_type == "lora": + if any(key.endswith((".magnitude_r", ".magnitude_c")) for key in keys): + return "bora" + if any(key.endswith(".magnitude") for key in keys): + return "dora-rows" + return canonical_adapter_type(raw_type) + + +def _infer_global_rank(grouped: dict[str, dict[str, np.ndarray]]) -> int: + ranks = { + rank for params in grouped.values() if (rank := _rank_from_params(params)) > 0 + } + if not ranks: + return 0 + if len(ranks) > 1: + raise ValueError(f"Multiple adapter ranks found: {sorted(ranks)}.") + return next(iter(ranks)) + + +def _rank_from_params(params: dict[str, np.ndarray]) -> int: + core = params.get("M_xs") + if core is not None and core.ndim == 2 and core.shape[0] == core.shape[1]: + return int(core.shape[0]) + adapter_a = params.get("lora_A") + adapter_b = params.get("lora_B") + if adapter_a is None or adapter_b is None: + return 0 + if adapter_b.shape[-1] == adapter_a.shape[0]: + return int(adapter_a.shape[0]) + if adapter_a.shape[-1] == adapter_b.shape[0]: + return int(adapter_a.shape[-1]) + return 0 + + +def _lora_delta_2d( + params: dict[str, np.ndarray], +) -> tuple[np.ndarray, int]: + adapter_a = _require_param(params, "lora_A").astype(np.float64) + adapter_b = _require_param(params, "lora_B").astype(np.float64) + if adapter_b.shape[-1] == adapter_a.shape[0]: + delta = adapter_b @ adapter_a + rank = adapter_a.shape[0] + elif adapter_a.shape[-1] == adapter_b.shape[0]: + delta = adapter_a @ adapter_b + rank = adapter_a.shape[-1] + else: + raise ValueError( + "Unable to multiply LoRA matrices with shapes " + f"A={adapter_a.shape}, B={adapter_b.shape}." + ) + if not np.isfinite(delta).all(): + raise ValueError("LoRA delta contains non-finite values.") + return delta.astype(np.float32), int(rank) + + +def _xs_delta_2d( + params: dict[str, np.ndarray], + base_2d: np.ndarray, +) -> tuple[np.ndarray, int]: + core = _require_param(params, "M_xs") + if core.ndim != 2 or core.shape[0] != core.shape[1]: + raise ValueError(f"LoRA-XS core must be square, got {core.shape}.") + rank = int(core.shape[0]) + if rank > min(base_2d.shape): + raise ValueError( + f"LoRA-XS rank {rank} exceeds base weight shape {base_2d.shape}." + ) + + u = params.get("U") + v = params.get("V") + if u is None or v is None: + u, v = _svd_bases_numpy(base_2d, rank) + delta = u.astype(np.float64) @ core.astype(np.float64) @ v.astype(np.float64).T + if not np.isfinite(delta).all(): + raise ValueError("LoRA-XS delta contains non-finite values.") + return delta.astype(np.float32), rank + + +def _require_param(params: dict[str, np.ndarray], name: str) -> np.ndarray: + value = params.get(name) + if value is None: + raise ValueError(f"Adapter layer is missing {name}.") + return value.astype(np.float32, copy=False) + + +def _resolve_target_key( + source_weight_key: str, + target_keys: tuple[str, ...], +) -> str | None: + target_key_set = set(target_keys) + candidates = _target_key_candidates(source_weight_key) + for candidate in candidates: + if candidate in target_key_set: + return candidate + + suffix_matches = { + target_key + for candidate in candidates + for target_key in target_keys + if target_key.endswith(candidate) + } + if len(suffix_matches) == 1: + return next(iter(suffix_matches)) + return None + + +def _target_key_candidates(source_weight_key: str) -> tuple[str, ...]: + prefixes = ("model.model.", "model.") + candidates = [source_weight_key] + for prefix in prefixes: + if source_weight_key.startswith(prefix): + candidates.append(source_weight_key[len(prefix) :]) + + candidates.extend( + candidate.replace("to_local_embed.0.", "to_local_embed.seq.0.").replace( + "to_local_embed.2.", "to_local_embed.seq.2." + ) + for candidate in tuple(candidates) + ) + return tuple(dict.fromkeys(candidates)) + + +def _checkpoint_layer_name(name: str) -> str: + return name.replace("to_local_embed.seq.0", "to_local_embed.0").replace( + "to_local_embed.seq.2", "to_local_embed.2" + ) + + +def _source_shape_for_delta( + target_shape: tuple[int, ...], + delta_shape: tuple[int, int], +) -> tuple[int, ...]: + if len(target_shape) == 2: + candidates = (target_shape, (target_shape[1], target_shape[0])) + elif len(target_shape) == 3: + candidates = ( + (target_shape[0], target_shape[2], target_shape[1]), + target_shape, + ) + else: + candidates = (target_shape,) + + for candidate in candidates: + if ( + candidate[0] == delta_shape[0] + and int(np.prod(candidate[1:])) == delta_shape[1] + ): + return candidate + raise ValueError( + f"Unable to map LoRA delta {delta_shape} to target shape {target_shape}." + ) + + +def _source_shape_for_xs( + target_shape: tuple[int, ...], + params: dict[str, np.ndarray], +) -> tuple[int, ...]: + u = params.get("U") + v = params.get("V") + if u is not None and v is not None: + return _source_shape_for_delta( + target_shape, + (int(u.shape[0]), int(v.shape[0])), + ) + if len(target_shape) == 3: + return (target_shape[0], target_shape[2], target_shape[1]) + return target_shape + + +def _target_to_source_weight( + target: np.ndarray, + source_shape: tuple[int, ...], +) -> np.ndarray: + if target.shape == source_shape: + return target + if target.ndim == 2 and target.T.shape == source_shape: + return target.T + if target.ndim == 3: + candidate = target.transpose(0, 2, 1) + if candidate.shape == source_shape: + return candidate + raise ValueError( + f"Unable to map target shape {target.shape} to source shape {source_shape}." + ) + + +def _source_to_target_weight( + source: np.ndarray, + target_shape: tuple[int, ...], +) -> np.ndarray: + if source.shape == target_shape: + return source + if source.ndim == 2 and source.T.shape == target_shape: + return source.T + if source.ndim == 3: + candidate = source.transpose(0, 2, 1) + if candidate.shape == target_shape: + return candidate + raise ValueError( + f"Unable to map source shape {source.shape} to target shape {target_shape}." + ) + + +def _validate_rank( + rank: int, + *, + fan_out: int, + fan_in: int, + source_name: str, +) -> None: + max_rank = min(fan_out, fan_in) + if rank > max_rank: + raise ValueError( + f"Adapter rank {rank} exceeds maximum rank {max_rank} for " + f"{source_name!r} with shape ({fan_out}, {fan_in})." + ) + + +def _svd_bases(weight_2d, rank: int): + u, v = _svd_bases_numpy(np.asarray(weight_2d, dtype=np.float32), rank) + return mx.array(u, dtype=mx.float32), mx.array(v, dtype=mx.float32) + + +def _svd_bases_numpy( + weight_2d: np.ndarray, + rank: int, +) -> tuple[np.ndarray, np.ndarray]: + u, _, vh = np.linalg.svd( + weight_2d.astype(np.float32, copy=False), + full_matrices=False, + ) + u, vh = _canonicalize_svd_signs(u, vh) + return ( + u[:, :rank].astype(np.float32, copy=False), + vh[:rank, :].T.astype(np.float32, copy=False), + ) + + +def _canonicalize_svd_signs( + u: np.ndarray, + vh: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + max_abs_indices = np.abs(u).argmax(axis=0) + signs = np.sign(u[max_abs_indices, np.arange(u.shape[1])]) + signs[signs == 0] = 1 + return u * signs[None, :], vh * signs[:, None] + + +def _name_is_selected( + name: str, + *, + include: tp.Sequence[str] | None, + exclude: tp.Sequence[str] | None, +) -> bool: + if include and not _matches_any(name, include): + return False + return not (exclude and _matches_any(name, exclude)) + + +def _matches_any(name: str, patterns: tp.Sequence[str]) -> bool: + return any( + expanded in name for pattern in patterns for expanded in _expand(pattern) + ) + + +def _expand(pattern: str) -> list[str]: + parts = re.split(r"\[(\d+)-(\d+)\]", pattern) + if len(parts) == 1: + return [pattern] + + literals = parts[0::3] + starts = parts[1::3] + ends = parts[2::3] + ranges = [] + for start, end in zip(starts, ends, strict=True): + start_value = int(start) + end_value = int(end) + step = 1 if end_value >= start_value else -1 + ranges.append( + [str(value) for value in range(start_value, end_value + step, step)] + ) + + expanded = [] + for values in product(*ranges): + pieces = [] + for index, literal in enumerate(literals): + pieces.append(literal) + if index < len(values): + pieces.append(values[index]) + expanded.append("".join(pieces)) + return expanded diff --git a/optimized/mlx/models/defs/training.py b/optimized/mlx/models/defs/training.py new file mode 100644 index 0000000..02bb313 --- /dev/null +++ b/optimized/mlx/models/defs/training.py @@ -0,0 +1,224 @@ +"""Training utilities for the standalone Stable Audio 3 MLX models.""" + +from __future__ import annotations + +import math +import typing as tp +from statistics import NormalDist + +import mlx.core as mx +import numpy as np + + +_TIMESTEP_SAMPLERS = { + "uniform", + "logit_normal", + "trunc_logit_normal", + "log_snr", + "log_snr_uniform", +} +_DISTRIBUTION_SHIFTS = {"none", "full", "flux", "logsnr"} +_STANDARD_NORMAL = NormalDist() + + +def sample_training_timesteps( + sampler: str, + batch_size: int, + *, + rng: np.random.Generator, + options: dict[str, float] | None = None, +) -> np.ndarray: + """Sample SA3 training timesteps with the PyTorch implementation's defaults.""" + + sampler = str(sampler).strip().lower() + if sampler not in _TIMESTEP_SAMPLERS: + raise ValueError(f"Unsupported timestep sampler: {sampler!r}.") + if batch_size <= 0: + raise ValueError("batch_size must be positive.") + + options = options or {} + if sampler == "uniform": + values = rng.random(batch_size) + elif sampler == "logit_normal": + values = _sigmoid(rng.standard_normal(batch_size)) + elif sampler == "trunc_logit_normal": + values = 1.0 - _truncated_logistic_normal_rescaled( + batch_size, + rng=rng, + ) + elif sampler == "log_snr": + mean = float(options.get("mean_logsnr", -1.2)) + std = float(options.get("std_logsnr", 2.0)) + logsnr = rng.standard_normal(batch_size) * std + mean + values = np.clip(_sigmoid(-logsnr), 1e-4, 1.0 - 1e-4) + else: + minimum = float(options.get("min_logsnr", -6.0)) + maximum = float(options.get("max_logsnr", 5.0)) + if maximum <= minimum: + raise ValueError("max_logsnr must be greater than min_logsnr.") + logsnr = rng.uniform(minimum, maximum, batch_size) + values = np.clip(_sigmoid(-logsnr), 1e-4, 1.0 - 1e-4) + return np.asarray(values, dtype=np.float32) + + +def shift_training_timesteps( + timesteps: tp.Sequence[float] | np.ndarray, + sequence_length: int | tp.Sequence[int] | np.ndarray, + *, + shift_type: str = "full", + options: dict[str, float | int | bool] | None = None, +) -> np.ndarray: + """Apply an SA3 training distribution shift to a timestep batch.""" + + shift_type = str(shift_type).strip().lower() + if shift_type == "identity": + shift_type = "none" + if shift_type not in _DISTRIBUTION_SHIFTS: + raise ValueError(f"Unsupported distribution shift: {shift_type!r}.") + + values = np.asarray(timesteps, dtype=np.float64).reshape(-1) + lengths = np.asarray(sequence_length, dtype=np.float64).reshape(-1) + if lengths.size == 1: + lengths = np.repeat(lengths, values.size) + if lengths.size != values.size: + raise ValueError( + "sequence_length must contain one value or match the timestep batch." + ) + if shift_type == "none": + return values.astype(np.float32) + + options = options or {} + shifted = [ + _shift_timestep( + float(timestep), + float(length), + shift_type=shift_type, + options=options, + ) + for timestep, length in zip(values, lengths, strict=True) + ] + return np.asarray(shifted, dtype=np.float32) + + +def rectified_flow_loss( + model, + clean, + timesteps, + *, + noise=None, + loss_mask=None, + model_kwargs: dict[str, tp.Any] | None = None, +): + """Compute the SA3 rectified-flow velocity loss for pre-encoded latents.""" + + if noise is None: + noise = mx.random.normal(clean.shape, dtype=clean.dtype) + timesteps = timesteps.astype(mx.float32) + alpha = (1.0 - timesteps)[:, None, None].astype(clean.dtype) + sigma = timesteps[:, None, None].astype(clean.dtype) + noised = clean * alpha + noise * sigma + target = noise - clean + prediction = model(noised, timesteps, **(model_kwargs or {})) + mse = (prediction.astype(mx.float32) - target.astype(mx.float32)) ** 2 + + if loss_mask is None: + return mx.mean(mse) + mask = loss_mask[:, None, :].astype(mx.float32) + denominator = mx.maximum(mx.sum(mask) * mse.shape[1], 1.0) + return mx.sum(mse * mask) / denominator + + +def _shift_timestep( + timestep: float, + sequence_length: float, + *, + shift_type: str, + options: dict[str, float | int | bool], +) -> float: + if timestep <= 0: + return 0.0 + if timestep >= 1: + return 1.0 + + if shift_type == "full": + minimum = float(options.get("min_length", 256)) + maximum = float(options.get("max_length", 4096)) + length = _clamp(sequence_length, minimum, maximum) + base_shift = float(options.get("base_shift", 0.5)) + max_shift = float(options.get("max_shift", 1.15)) + mu = -( + base_shift + + (max_shift - base_shift) * (length - minimum) / (maximum - minimum) + ) + exp_mu = math.exp(mu) + shifted = 1.0 - exp_mu / (exp_mu + (1.0 / (1.0 - timestep) - 1.0)) + if bool(options.get("use_sine", False)): + shifted = math.sin(shifted * math.pi / 2.0) + return shifted + + if shift_type == "flux": + minimum = float(options.get("min_length", 256)) + maximum = float(options.get("max_length", 4096)) + length = _clamp(sequence_length, minimum, maximum) + alpha_min = max(float(options.get("alpha_min", 1.0)), 1e-8) + alpha_max = max(float(options.get("alpha_max", 1.0)), 1e-8) + denominator = math.log(maximum) - math.log(minimum) + fraction = (math.log(length) - math.log(minimum)) / max( + denominator, + 1e-8, + ) + log_alpha = math.log(alpha_min) + fraction * ( + math.log(alpha_max) - math.log(alpha_min) + ) + alpha = math.exp(log_alpha) + return alpha * timestep / (1.0 + (alpha - 1.0) * timestep) + + anchor_length = float(options.get("anchor_length", 2000)) + anchor_logsnr = float(options.get("anchor_logsnr", -6.2)) + rate = float(options.get("rate", 1.0)) + logsnr_end = float(options.get("logsnr_end", 2.0)) + logsnr_start = anchor_logsnr - rate * math.log2( + max(sequence_length, 1.0) / anchor_length + ) + logsnr = logsnr_end - timestep * (logsnr_end - logsnr_start) + return 1.0 / (1.0 + math.exp(logsnr)) + + +def _sigmoid(values: np.ndarray) -> np.ndarray: + return 1.0 / (1.0 + np.exp(-np.asarray(values, dtype=np.float64))) + + +def _truncated_logistic_normal_rescaled( + size: int, + *, + rng: np.random.Generator, + left_trunc: float = 0.075, + right_trunc: float = 1.0, +) -> np.ndarray: + if not 0.0 < left_trunc < right_trunc <= 1.0: + raise ValueError("Expected 0 < left_trunc < right_trunc <= 1.") + + lower_logit = math.log(left_trunc / (1.0 - left_trunc)) + lower_cdf = _STANDARD_NORMAL.cdf(lower_logit) + upper_cdf = ( + 1.0 + if right_trunc == 1.0 + else _STANDARD_NORMAL.cdf(math.log(right_trunc / (1.0 - right_trunc))) + ) + uniforms = lower_cdf + (upper_cdf - lower_cdf) * rng.random(size) + epsilon = np.finfo(np.float64).eps + logits = np.asarray( + [ + _STANDARD_NORMAL.inv_cdf(float(np.clip(value, epsilon, 1.0 - epsilon))) + for value in uniforms + ], + dtype=np.float64, + ) + samples = _sigmoid(logits) + return (samples - left_trunc) / (right_trunc - left_trunc) + + +def _clamp(value: float, minimum: float, maximum: float) -> float: + if maximum <= minimum: + raise ValueError("max_length must be greater than min_length.") + return min(max(value, minimum), maximum) diff --git a/tests/test_mlx_audio_encoding.py b/tests/test_mlx_audio_encoding.py new file mode 100644 index 0000000..9b2bbc7 --- /dev/null +++ b/tests/test_mlx_audio_encoding.py @@ -0,0 +1,121 @@ +import numpy as np +import pytest + +pytest.importorskip("mlx.core") + +import mlx.core as mx + +from optimized.mlx.models.defs.audio_encoding import ( + SAMPLES_PER_LATENT, + encode_audio, + patch_audio, +) + + +class GroupingEncoder: + def __init__(self, stride: int = 16): + self.stride = stride + + def __call__(self, patches): + batch, channels, patch_count = patches.shape + return mx.mean( + patches.reshape( + batch, + channels, + patch_count // self.stride, + self.stride, + ), + axis=-1, + ) + + +def test_patch_audio_matches_patched_pretransform_layout(): + audio = mx.array(np.arange(2 * 8, dtype=np.float32).reshape(1, 2, 8)) + + patched = patch_audio(audio, patch_size=4) + + np.testing.assert_array_equal( + np.asarray(patched), + np.array( + [ + [ + [0, 4], + [1, 5], + [2, 6], + [3, 7], + [8, 12], + [9, 13], + [10, 14], + [11, 15], + ] + ], + dtype=np.float32, + ), + ) + + +def test_encode_audio_pads_to_codec_alignment_and_builds_mask(): + audio = mx.zeros((2, 2, SAMPLES_PER_LATENT + 1)) + + encoded = encode_audio( + GroupingEncoder(), + audio, + valid_sample_lengths=[SAMPLES_PER_LATENT, SAMPLES_PER_LATENT + 1], + pad_modulo=32, + ) + + assert encoded.source_samples == SAMPLES_PER_LATENT + 1 + assert encoded.padded_samples == SAMPLES_PER_LATENT * 2 + assert encoded.latents.shape == (2, 512, 2) + assert encoded.valid_latent_lengths == (1, 2) + np.testing.assert_array_equal( + np.asarray(encoded.padding_mask), + np.array([[True, False], [True, True]]), + ) + + +def test_chunked_encoding_matches_unchunked_encoding(): + rng = np.random.default_rng(21) + audio = mx.array( + rng.standard_normal((2, 2, SAMPLES_PER_LATENT * 10)).astype(np.float32) + ) + encoder = GroupingEncoder() + + unchunked = encode_audio(encoder, audio, pad_modulo=16) + chunked = encode_audio( + encoder, + audio, + pad_modulo=16, + chunked=True, + chunk_size=4, + overlap=2, + chunk_batch_size=2, + ) + + np.testing.assert_allclose( + np.asarray(chunked.latents), + np.asarray(unchunked.latents), + atol=1e-6, + ) + np.testing.assert_array_equal( + np.asarray(chunked.padding_mask), + np.asarray(unchunked.padding_mask), + ) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"valid_sample_lengths": [1, 2]}, "one value per audio batch item"), + ({"pad_modulo": 17}, "positive multiple of encoder_stride"), + ( + {"chunked": True, "chunk_size": 3, "overlap": 1, "pad_modulo": 32}, + "incompatible with pad_modulo", + ), + ], +) +def test_encode_audio_rejects_invalid_contracts(kwargs, message): + audio = mx.zeros((1, 2, SAMPLES_PER_LATENT * 4)) + + with pytest.raises(ValueError, match=message): + encode_audio(GroupingEncoder(), audio, **kwargs) diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py new file mode 100644 index 0000000..5b05883 --- /dev/null +++ b/tests/test_mlx_lora.py @@ -0,0 +1,368 @@ +from functools import partial +from pathlib import Path + +import numpy as np +import pytest +import torch + +pytest.importorskip("mlx.core") + +import mlx.core as mx +import mlx.nn as nn +import mlx.optimizers as optim +from mlx.utils import tree_flatten + +from optimized.mlx.models.defs.lora import ( + apply_lora_checkpoint, + apply_lora_checkpoints, + inject_trainable_lora, + load_lora_checkpoint, + save_lora_checkpoint, +) +from stable_audio_3.models.lora import ( + LoRAParametrization, + add_lora, + get_lora_state_dict, + load_lora_checkpoint as load_torch_lora_checkpoint, + save_lora_safetensors, +) + + +class TinyMLXLinear(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(3, 2, bias=False) + self.layer.weight = mx.array( + [[1.0, -2.0, 0.5], [-0.5, 1.5, 2.0]], + dtype=mx.float32, + ) + + def __call__(self, x): + return self.layer(x) + + +class TinyMLXRegressor(nn.Module): + def __init__(self): + super().__init__() + self.input = nn.Linear(3, 4, bias=False) + self.output = nn.Linear(4, 2, bias=False) + self.output.weight = mx.zeros_like(self.output.weight) + + def __call__(self, x): + return self.output(nn.silu(self.input(x))) + + +class TinyMLXConv1d(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Conv1d(2, 3, kernel_size=3, bias=False) + source_weight = np.arange(18, dtype=np.float32).reshape(3, 2, 3) / 20 + self.layer.weight = mx.array(source_weight.transpose(0, 2, 1)) + + +@pytest.mark.parametrize( + ("model", "inputs"), + [ + (TinyMLXLinear(), mx.ones((1, 3))), + (TinyMLXConv1d(), mx.ones((1, 5, 2))), + ], +) +def test_trainable_dora_supports_bias_free_mlx_layers(model, inputs): + inject_trainable_lora( + model, + rank=1, + alpha=1, + adapter_type="dora", + ) + + output = model.layer(inputs) + mx.eval(output) + + assert bool(mx.all(mx.isfinite(output))) + + +class TinyTorchLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(3, 2, bias=False) + + +class TinyTorchConv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.layer = torch.nn.Conv1d(2, 3, kernel_size=3, bias=False) + + +def test_trainable_lora_updates_only_adapter_parameters(): + mx.random.seed(7) + model = TinyMLXRegressor() + report = inject_trainable_lora( + model, + rank=2, + alpha=2, + include=["output"], + ) + base_before = mx.array(model.output.base.weight) + inputs = mx.array( + [[1.0, -2.0, 0.5], [-1.0, 0.5, 2.0]], + dtype=mx.float32, + ) + target = mx.array( + [[0.5, -1.0], [-0.25, 0.75]], + dtype=mx.float32, + ) + + def loss_fn(local_model, values, expected): + return mx.mean((local_model(values) - expected) ** 2) + + loss_and_grad = nn.value_and_grad(model, loss_fn) + optimizer = optim.AdamW(learning_rate=0.1) + initial_loss = float(loss_fn(model, inputs, target)) + for _ in range(40): + loss, grads = loss_and_grad(model, inputs, target) + optimizer.update(model, grads) + mx.eval(model.parameters(), optimizer.state, loss) + + assert report.layer_names == ("output",) + assert report.trainable_parameters == 12 + assert [name for name, _ in tree_flatten(model.trainable_parameters())] == [ + "output.lora_A", + "output.lora_B", + ] + assert mx.array_equal(model.output.base.weight, base_before) + assert float(loss_fn(model, inputs, target)) < initial_loss * 0.1 + + +def test_mlx_checkpoint_round_trips_through_official_torch_loader( + tmp_path: Path, +): + model = TinyMLXLinear() + inject_trainable_lora(model, rank=1, alpha=1, adapter_type="dora") + model.layer.lora_A = mx.array([[0.25, -0.5, 1.0]]) + model.layer.lora_B = mx.array([[0.5], [-0.25]]) + model.layer.magnitude = mx.array([3.0, 2.0]) + + checkpoint = save_lora_checkpoint( + model, + tmp_path / "mlx-dora.safetensors", + extra_config={"step": 20}, + ) + mlx_state, mlx_config = load_lora_checkpoint(checkpoint) + torch_state, torch_config = load_torch_lora_checkpoint(checkpoint) + + assert sorted(mlx_state) == sorted(torch_state) + assert mlx_config == torch_config + assert mlx_config["adapter_type"] == "dora-rows" + assert mlx_config["step"] == 20 + + +@pytest.mark.parametrize( + "adapter_type", + [ + "lora", + "dora-rows", + "dora-cols", + "bora", + "lora-xs", + "dora-rows-xs", + "dora-cols-xs", + "bora-xs", + ], +) +def test_mlx_inference_matches_official_torch_adapter_math( + tmp_path: Path, + adapter_type: str, +): + base_weight = torch.tensor( + [[1.0, -2.0, 0.5], [-0.5, 1.5, 2.0]], + dtype=torch.float32, + ) + torch_model = TinyTorchLinear() + torch_model.layer.weight.data.copy_(base_weight) + config = { + torch.nn.Linear: { + "weight": partial( + LoRAParametrization.from_linear, + rank=1, + lora_alpha=1, + adapter_type=adapter_type, + ) + } + } + add_lora(torch_model, config) + adapter = torch_model.layer.parametrizations.weight[0] + if adapter_type.endswith("-xs"): + adapter.M_xs.data.fill_(0.5) + else: + adapter.lora_A.data.copy_(torch.tensor([[0.25, -0.5, 1.0]])) + adapter.lora_B.data.copy_(torch.tensor([[0.5], [-0.25]])) + + if adapter_type in {"dora-rows", "dora-rows-xs"}: + adapter.magnitude.data.copy_(torch.tensor([3.0, 2.0])) + elif adapter_type in {"dora-cols", "dora-cols-xs"}: + adapter.magnitude.data.copy_(torch.tensor([1.5, 2.5, 3.5])) + elif adapter_type in {"bora", "bora-xs"}: + adapter.magnitude_r.data.copy_(torch.tensor([3.0, 2.0])) + adapter.magnitude_c.data.copy_(torch.tensor([1.5, 2.5, 3.5])) + + checkpoint = tmp_path / f"{adapter_type}.safetensors" + save_lora_safetensors( + get_lora_state_dict(torch_model), + {"rank": 1, "alpha": 1, "adapter_type": adapter_type}, + checkpoint, + ) + + mlx_model = TinyMLXLinear() + report = apply_lora_checkpoint(mlx_model, checkpoint) + expected = torch_model.layer.weight.detach().numpy() + + assert report.adapter_type == adapter_type + assert report.applied_layers == 1 + assert report.missing_targets == () + assert report.skipped_layers == () + assert np.allclose(np.asarray(mlx_model.layer.weight), expected, atol=2e-3) + + +def test_multiple_lora_checkpoints_apply_with_independent_strengths( + tmp_path: Path, +): + base_weight = np.array( + [[1.0, -2.0, 0.5], [-0.5, 1.5, 2.0]], + dtype=np.float32, + ) + checkpoints = [] + deltas = [] + for index, (lora_a, lora_b) in enumerate( + ( + ( + [[0.25, -0.5, 1.0]], + [[0.5], [-0.25]], + ), + ( + [[-0.75, 0.5, 0.25]], + [[0.2], [0.4]], + ), + ) + ): + checkpoint = tmp_path / f"lora-{index}.safetensors" + mx.save_safetensors( + str(checkpoint), + { + "layer.parametrizations.weight.0.lora_A": mx.array(lora_a), + "layer.parametrizations.weight.0.lora_B": mx.array(lora_b), + }, + metadata={"lora_config": '{"rank": 1, "alpha": 1, "adapter_type": "lora"}'}, + ) + checkpoints.append(checkpoint) + deltas.append(np.asarray(lora_b, dtype=np.float32) @ np.asarray(lora_a)) + + model = TinyMLXLinear() + strengths = (0.25, 0.75) + reports = apply_lora_checkpoints( + model, + checkpoints, + strengths=strengths, + ) + expected = base_weight + strengths[0] * deltas[0] + strengths[1] * deltas[1] + + assert [report.applied_layers for report in reports] == [1, 1] + assert np.allclose(np.asarray(model.layer.weight), expected, atol=2e-3) + + +def test_checkpoint_names_map_to_optimized_local_embed_layout(tmp_path: Path): + class LocalEmbed(nn.Module): + def __init__(self): + super().__init__() + self.to_local_embed = type("LocalEmbedSeq", (nn.Module,), {})() + self.to_local_embed.seq = [ + nn.Linear(3, 2, bias=False), + None, + nn.Linear(2, 2, bias=False), + ] + + model = LocalEmbed() + original = np.asarray(model.to_local_embed.seq[0].weight).copy() + checkpoint = tmp_path / "local-embed.safetensors" + mx.save_safetensors( + str(checkpoint), + { + "to_local_embed.0.parametrizations.weight.0.lora_A": mx.array( + [[1.0, 0.0, 0.0]] + ), + "to_local_embed.0.parametrizations.weight.0.lora_B": mx.array( + [[0.5], [-0.25]] + ), + }, + metadata={"lora_config": ('{"rank": 1, "alpha": 1, "adapter_type": "lora"}')}, + ) + + report = apply_lora_checkpoint(model, checkpoint) + + assert report.applied_layers == 1 + assert not np.array_equal( + np.asarray(model.to_local_embed.seq[0].weight), + original, + ) + + +def test_saved_local_embed_name_maps_back_to_pytorch_layout(tmp_path: Path): + class LocalEmbed(nn.Module): + def __init__(self): + super().__init__() + self.to_local_embed = type("LocalEmbedSeq", (nn.Module,), {})() + self.to_local_embed.seq = [ + nn.Linear(3, 2, bias=False), + None, + nn.Linear(2, 2, bias=False), + ] + + model = LocalEmbed() + inject_trainable_lora( + model, + rank=1, + include=["to_local_embed.seq.0"], + ) + checkpoint = save_lora_checkpoint( + model, + tmp_path / "local-embed-save.safetensors", + ) + state_dict, _ = load_torch_lora_checkpoint(checkpoint) + + assert sorted(state_dict) == [ + "to_local_embed.0.parametrizations.weight.0.lora_A", + "to_local_embed.0.parametrizations.weight.0.lora_B", + ] + + +def test_conv1d_checkpoint_maps_pytorch_weight_layout_to_mlx(tmp_path: Path): + torch_model = TinyTorchConv1d() + torch_model.layer.weight.data.copy_( + torch.arange(18, dtype=torch.float32).reshape(3, 2, 3) / 20 + ) + config = { + torch.nn.Conv1d: { + "weight": partial( + LoRAParametrization.from_conv1d, + rank=1, + lora_alpha=1, + adapter_type="lora", + ) + } + } + add_lora(torch_model, config) + adapter = torch_model.layer.parametrizations.weight[0] + adapter.lora_A.data.copy_(torch.tensor([[0.1, -0.2, 0.3, -0.4, 0.5, -0.6]])) + adapter.lora_B.data.copy_(torch.tensor([[0.5], [-0.25], [0.75]])) + checkpoint = tmp_path / "conv1d.safetensors" + save_lora_safetensors( + get_lora_state_dict(torch_model), + {"rank": 1, "alpha": 1, "adapter_type": "lora"}, + checkpoint, + ) + + mlx_model = TinyMLXConv1d() + report = apply_lora_checkpoint(mlx_model, checkpoint) + expected = torch_model.layer.weight.detach().numpy().transpose(0, 2, 1) + + assert report.applied_layers == 1 + assert np.allclose(np.asarray(mlx_model.layer.weight), expected, atol=2e-3) diff --git a/tests/test_mlx_training.py b/tests/test_mlx_training.py new file mode 100644 index 0000000..ce939ca --- /dev/null +++ b/tests/test_mlx_training.py @@ -0,0 +1,69 @@ +import numpy as np +import pytest + +pytest.importorskip("mlx.core") + +import mlx.core as mx + +from optimized.mlx.models.defs.training import ( + rectified_flow_loss, + sample_training_timesteps, + shift_training_timesteps, +) + + +def test_truncated_logit_normal_sampler_matches_sa3_distribution(): + values = sample_training_timesteps( + "trunc_logit_normal", + 20_000, + rng=np.random.default_rng(17), + ) + + assert values.dtype == np.float32 + assert np.all((values >= 0.0) & (values <= 1.0)) + assert 0.52 < float(values.mean()) < 0.55 + assert 0.52 < float(np.median(values)) < 0.56 + + +def test_full_training_shift_matches_sa3_defaults(): + shifted = shift_training_timesteps( + [0.5], + 507, + shift_type="full", + ) + + assert shifted[0] == pytest.approx(0.6323907626) + + +def test_rectified_flow_loss_uses_velocity_target_and_mask(): + clean = mx.array([[[1.0, 2.0, 3.0]]], dtype=mx.float32) + noise = mx.array([[[4.0, 5.0, 6.0]]], dtype=mx.float32) + timesteps = mx.array([0.25], dtype=mx.float32) + mask = mx.array([[True, False, True]]) + + def perfect_model(noised, timestep): + del noised, timestep + return noise - clean + + def zero_model(noised, timestep): + del timestep + return mx.zeros_like(noised) + + assert float( + rectified_flow_loss( + perfect_model, + clean, + timesteps, + noise=noise, + loss_mask=mask, + ) + ) == pytest.approx(0.0) + assert float( + rectified_flow_loss( + zero_model, + clean, + timesteps, + noise=noise, + loss_mask=mask, + ) + ) == pytest.approx(9.0) From e5af3b42de256a63b2e768f4e835c4eda25ff4cc Mon Sep 17 00:00:00 2001 From: betweentwomidnights Date: Mon, 15 Jun 2026 11:48:00 -0700 Subject: [PATCH 02/28] Clarify odd-overlap chunk stitching --- optimized/mlx/models/defs/audio_encoding.py | 2 + tests/test_mlx_audio_encoding.py | 41 +++++++++++---------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/optimized/mlx/models/defs/audio_encoding.py b/optimized/mlx/models/defs/audio_encoding.py index c636670..4ed5619 100644 --- a/optimized/mlx/models/defs/audio_encoding.py +++ b/optimized/mlx/models/defs/audio_encoding.py @@ -238,6 +238,8 @@ def _stitch_encoded_chunks( chunk_size: int, overlap: int, ): + # Floor division intentionally gives the later chunk the extra latent when + # an overlap is odd. The interval clipping below prevents gaps or duplicates. half_overlap = overlap // 2 intervals = [] last_index = len(chunk_starts) - 1 diff --git a/tests/test_mlx_audio_encoding.py b/tests/test_mlx_audio_encoding.py index 9b2bbc7..f3d928c 100644 --- a/tests/test_mlx_audio_encoding.py +++ b/tests/test_mlx_audio_encoding.py @@ -74,33 +74,34 @@ def test_encode_audio_pads_to_codec_alignment_and_builds_mask(): ) -def test_chunked_encoding_matches_unchunked_encoding(): +def test_chunked_encoding_matches_unchunked_for_even_and_odd_overlaps(): rng = np.random.default_rng(21) audio = mx.array( - rng.standard_normal((2, 2, SAMPLES_PER_LATENT * 10)).astype(np.float32) + rng.standard_normal((2, 2, SAMPLES_PER_LATENT * 11)).astype(np.float32) ) encoder = GroupingEncoder() unchunked = encode_audio(encoder, audio, pad_modulo=16) - chunked = encode_audio( - encoder, - audio, - pad_modulo=16, - chunked=True, - chunk_size=4, - overlap=2, - chunk_batch_size=2, - ) + for overlap in (2, 1): + chunked = encode_audio( + encoder, + audio, + pad_modulo=16, + chunked=True, + chunk_size=4, + overlap=overlap, + chunk_batch_size=2, + ) - np.testing.assert_allclose( - np.asarray(chunked.latents), - np.asarray(unchunked.latents), - atol=1e-6, - ) - np.testing.assert_array_equal( - np.asarray(chunked.padding_mask), - np.asarray(unchunked.padding_mask), - ) + np.testing.assert_allclose( + np.asarray(chunked.latents), + np.asarray(unchunked.latents), + atol=1e-6, + ) + np.testing.assert_array_equal( + np.asarray(chunked.padding_mask), + np.asarray(unchunked.padding_mask), + ) @pytest.mark.parametrize( From 669618cb04125213d309ce3af405b6d96e51dbb9 Mon Sep 17 00:00:00 2001 From: betweentwomidnights Date: Fri, 26 Jun 2026 00:39:36 -0700 Subject: [PATCH 03/28] Cover optimized small and medium DiT LoRA injection --- tests/test_mlx_lora.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py index 5b05883..da9eb3a 100644 --- a/tests/test_mlx_lora.py +++ b/tests/test_mlx_lora.py @@ -12,6 +12,7 @@ import mlx.optimizers as optim from mlx.utils import tree_flatten +from optimized.mlx.models.defs import dit_mlx, dit_mlx_medium from optimized.mlx.models.defs.lora import ( apply_lora_checkpoint, apply_lora_checkpoints, @@ -81,6 +82,44 @@ def test_trainable_dora_supports_bias_free_mlx_layers(model, inputs): assert bool(mx.all(mx.isfinite(output))) +@pytest.mark.parametrize( + ("model_factory", "minimum_layer_count"), + [ + (dit_mlx.DiT, 190), + (dit_mlx_medium.DiT, 220), + ], +) +def test_trainable_lora_targets_optimized_small_and_medium_dits( + model_factory, + minimum_layer_count: int, +): + model = model_factory(T_lat=8) + target_names = [ + name + for name, layer in model.named_modules() + if isinstance(layer, (nn.Linear, nn.Conv1d)) + ] + include = [ + "transformer.project_in", + "preprocess_conv", + "transformer.layers.0.to_local_embed.seq.0", + ] + + assert len(target_names) >= minimum_layer_count + assert all(name in target_names for name in include) + + report = inject_trainable_lora( + model, + rank=1, + alpha=1, + adapter_type="lora", + include=include, + ) + + assert set(report.layer_names) == set(include) + assert report.adapter_type == "lora" + + class TinyTorchLinear(torch.nn.Module): def __init__(self): super().__init__() From c05cbbd6eb86ef8d2a0a6e2dd2eb714c594dcf63 Mon Sep 17 00:00:00 2001 From: betweentwomidnights Date: Fri, 26 Jun 2026 06:31:19 -0700 Subject: [PATCH 04/28] Match PyTorch masked loss averaging --- optimized/mlx/models/defs/training.py | 8 ++++++-- tests/test_mlx_training.py | 21 +++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/optimized/mlx/models/defs/training.py b/optimized/mlx/models/defs/training.py index 02bb313..f7313e6 100644 --- a/optimized/mlx/models/defs/training.py +++ b/optimized/mlx/models/defs/training.py @@ -124,8 +124,12 @@ def rectified_flow_loss( if loss_mask is None: return mx.mean(mse) mask = loss_mask[:, None, :].astype(mx.float32) - denominator = mx.maximum(mx.sum(mask) * mse.shape[1], 1.0) - return mx.sum(mse * mask) / denominator + per_sample_denominator = mx.maximum( + mx.sum(mask, axis=(1, 2)) * mse.shape[1], + 1.0, + ) + per_sample_loss = mx.sum(mse * mask, axis=(1, 2)) / per_sample_denominator + return mx.mean(per_sample_loss) def _shift_timestep( diff --git a/tests/test_mlx_training.py b/tests/test_mlx_training.py index ce939ca..25a3e25 100644 --- a/tests/test_mlx_training.py +++ b/tests/test_mlx_training.py @@ -67,3 +67,24 @@ def zero_model(noised, timestep): loss_mask=mask, ) ) == pytest.approx(9.0) + + +def test_rectified_flow_loss_averages_masked_loss_per_sample(): + clean = mx.zeros((2, 1, 3), dtype=mx.float32) + noise = mx.array([[[1.0, 1.0, 1.0]], [[3.0, 3.0, 3.0]]], dtype=mx.float32) + timesteps = mx.array([0.25, 0.25], dtype=mx.float32) + mask = mx.array([[True, False, False], [True, True, True]]) + + def zero_model(noised, timestep): + del noised, timestep + return mx.zeros_like(clean) + + loss = rectified_flow_loss( + zero_model, + clean, + timesteps, + noise=noise, + loss_mask=mask, + ) + + assert float(loss) == pytest.approx(5.0) From 7d59118cd8c3ba151804ac249db92c73ac24c2a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 02:22:34 -0400 Subject: [PATCH 05/28] docs: underfit training-convention inventory + port gap analysis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Complete inventory of every training convention and default in underfit (adapter config, checkpoint key roots and metadata, loss/timesteps/CFG dropout, optimizer and loop mechanics, pre-encode and dataset semantics, prompt augmentation, demos), cross-referenced to source, plus the gap analysis against the MLX primitives this branch starts from and the build order. Working reference for the branch — fold into README before the PR finalizes. --- optimized/mlx/TRAINING_CONVENTIONS.md | 155 ++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 optimized/mlx/TRAINING_CONVENTIONS.md diff --git a/optimized/mlx/TRAINING_CONVENTIONS.md b/optimized/mlx/TRAINING_CONVENTIONS.md new file mode 100644 index 0000000..53f2016 --- /dev/null +++ b/optimized/mlx/TRAINING_CONVENTIONS.md @@ -0,0 +1,155 @@ +# Underfit training conventions → MLX port reference + +Working reference for the `mlx-lora-training` branch: every training convention and +default in [underfit](https://github.com/dada-bots/underfit) (the SA3 LoRA trainer), +inventoried from its source (2026-07-14), plus the gap analysis against the MLX +primitives this branch starts from (PR #51 by @betweentwomidnights). The goal: +the MLX runtime can run as an underfit backend — same behaviors, same defaults, +same checkpoint format. Trim/fold this file into README before the PR finalizes. + +## 1. Adapter configuration + +| Convention | Underfit value | Source | +|---|---|---| +| adapter_type (config-layer default) | `"lora"`; legacy `"dora"` → `"dora-rows"` | training/lora.py:34, resolve_adapter_type | +| adapter_type (product/dashboard default) | **`dora-rows`** | dashboard index.html:1336 | +| rank (config-layer default) | 8 | training/lora.py:32 | +| rank (product/dashboard default) | **16** | README.md:157 | +| alpha default | = rank | training/lora.py:33 | +| targets | every `nn.Linear` + `nn.Conv1d` under BOTH `model.model` (DiT) and `model.conditioner`, filtered by include/exclude | training/lora.py:53-64 | +| include/exclude | substring match + `[i-j]` numeric-range expansion; dashboard default EXCLUDES all conditioner + "other" modules (embedders/projections/convs) | dashboard index.html:5741-49 | +| adapter dropout | none | — | +| trainable dtype | LoRA params forced fp32; base stays fp16 (`base_precision`) | loop.py:344-347 | + +## 2. Checkpoint format + +- Keys: `{**get_lora_state_dict(model.model), **get_lora_state_dict(model.conditioner)}` + → DiT keys come out `model.transformer.layers...` (inner `.model` on the DiT wrapper), + conditioner keys `conditioners.seconds_total.embedder.embedding.1...`; per-key shape + `.parametrizations.weight..{lora_A,lora_B,magnitude,magnitude_r,magnitude_c,M_xs,U,V}` + (index 0 for the finetune adapter). +- Tensors saved **fp16** (`save_lora_safetensors` casts `.half()`); DoRA magnitudes + squeezed to 1-D on load (`prepare_dora_state_dict`). +- Metadata `lora_config` JSON: `rank, alpha, adapter_type, include, exclude` + + injected `step` (int), `epoch` (int), `base_model` (str, from model_config). +- Filename: `{run_label}-step={step}-epoch={epoch}.safetensors` (run_label = run name + with trailing 14-digit timestamp stripped); dir `///checkpoints`. +- -xs: only `M_xs` (+ magnitudes) trained/saved; U/V bases precomputed OFFLINE + (compute_svd.py) from fp32 base weights into a separate `.pt` + (`{key: {"U": fp16, "V": fp16, "S": fp32, "shape"}}`, sign-canonicalized by max-abs + U-column entry), loaded via `svd_bases_path` from the model registry; per-layer SVD + at attach time is the fallback. +- Validation contract (lora_validate.py): `.safetensors` only, ≥1 parametrization key, + `lora_config` present with `rank` + `adapter_type` (alpha:=rank, include/exclude:=[] + back-filled), adapter type inferable from key fingerprints. + +## 3. Loss (loop.py:561-613 + loss.py) + +- Objective from the model: `rectified_flow` → `alphas, sigmas = 1−t, t`, + `noised = x·α + noise·σ`, `target = noise − x` (`"v"` objective also exists: + cos/sin schedule, `target = noise·α − x·σ`). +- MSE per-element; `loss_normalization` default `"none"` (timestep/sample/sample_channel + variance-normalization modes exist, eps 1e-6, detached variance). +- Masking: `mask_loss_weight` default 0.0 and `mask_padding_attention` default False → + **signal-only** loss; reduction is per-sample mean over valid positions×channels, + then batch mean. Inpainting further restricts the mask to the regenerated region. +- No sigma/logSNR loss weighting. Runs under fp16 autocast; no explicit casts. + +## 4. Timesteps + +- Sampler default **`uniform`** (`torch.rand`); also logit_normal, + trunc_logit_normal (left 0.075, flipped `1−t`), log_snr (mean −1.2, std 2.0), + log_snr_uniform (min −6, max 5). +- Distribution shift comes from the MODEL config: SA3 models ship + `{"type": "full", "min_length": 256, "max_length": 4096}` (base_shift 0.5, + max_shift 1.15 defaults) — applied as `t = dist_shift.shift(t, seq_len)` with + seq_len = the latent time dim (crop length) by default + (`use_effective_length_for_schedule` False; when True: `ceil(seconds_total·sr/4096)`). + +## 5. CFG dropout (dit.py:441-450) + +- `cfg_dropout_prob` default **0.1**, applied INSIDE the model forward when + `cfg_scale == 1.0`: per-sample bernoulli mask replaces the whole + `cross_attn_cond` (prompt+seconds concat) with **zeros**; `global_cond` is NOT + dropped. (Matches inference: the uncond branch zeroes cross_attn, keeps global.) + +## 6. Optimizer / schedule / loop mechanics + +| Convention | Underfit value | +|---|---| +| optimizer | AdamW, single param group (LoRA params only) | +| lr | **no default — must be supplied** (`--lr` or optimizer_configs) | +| betas/eps/weight_decay | torch defaults: (0.9, 0.999) / 1e-8 / 0.0 | +| LR schedule | none by default (InverseLR + any torch scheduler configurable) | +| grad clip | off by default (`gradient_clip_val = 0.0`); when set: clip_grad_norm_ | +| grad accumulation | none (defaults.ini key exists but inert) | +| EMA / validation / early stopping / save_top_k | none (inert keys) | +| precision | `"16-mixed"`: fp16 autocast + GradScaler; base fp16, LoRA fp32 | +| batch_size | 1 (defaults.ini) | +| seed | 42 | +| steps | step-driven; `max_steps` absolute global-step target (default ~1e10) | +| checkpoint cadence | every 1000 steps + final save; save BEFORE demos | +| resume offsets | explicit config > safetensors metadata `step/epoch` > filename `step=/epoch=` tokens > step//steps_per_epoch | +| logging | train/loss, train/lr, grad_norm + lora_magnitude (fp32, post-clip); `loss_by_timestep.bin` binary telemetry (step, t_mean, loss_mean) flushed every 10 steps | + +## 7. Data pipeline + +- **Pre-encode** (offline, once): whole files (no clip splitting), up to 600 s + (aligned down to ds_ratio=4096 samples), 44.1 kHz stereo (mono→stereo by channel + repeat, resample via torchaudio defaults, NO loudness/peak norm), encoder fp32, + `chunked=True` for >30 s. Saved per file: `.npy` latents `[D,T]` fp32 + **unscaled** (training divides by pretransform scale) + `.json` with + `seconds_total` (full duration, rounded 3dp), `seconds_start: 0`, + `audio_samples`, `latent_shape`, latent-resolution 0/1 `padding_mask`, and all tags. +- **Training dataset** (PreEncodedDataset semantics): crop/pad every item to + `latent_crop_length` — **1300** latents (sm-music/sfx, ≈120 s) / **4096** (medium, + ≈380 s); `random_crop=True` (dashboard default) choosing a start within the valid + region; short items padded with `silence.npy` latent if present else zeros, mask + extended with zeros. Fixed-length batches + padding mask (no bucketing). + **`seconds_total` stays the FULL song duration after cropping** (not recomputed). +- **Small datasets**: if `len(ds) < batch_size`, RandomSampler with replacement, + `num_samples = batch_size*100` (~100 steps/epoch) — the core underfit workflow. +- **Prompts** (prompt_templates.py, per-sample at train time): tag-based by default + (`use_tags`, tags→`Title:/Artist:/Album:/Genre:/Label:/Year:/Composer:/BPM:/Prompt:` + joined ", "); augmentation: 50% shuffle-all vs 50% random-subset (1..n tags); + optional trigger token prepended with prob **80%**; path/fixed prompt sources with + weighted balance (tags 50/paths 50/fixed 0); `lyrics=""` always. Tag priority: + JSON sidecar > `.txt` sidecar (whole content = `prompt`) > embedded ID3 + (`title,artist,album,genre,label,date,composer,bpm`). +- **Demos**: `demo_every` 1000 (template; loop default 0=off) + baseline at step 0; + `demo_steps` 50 (ARC 8), cfg [7] (ARC 1), per-entry seed/duration/lora_strength/ + `lora_interval_max` (σ-interval demo path uses an explicit Euler loop); + peak-normalized int16 → mp3 (`ffmpeg -q:a 0`), trimmed to seconds_total+5 s, + `demo__.mp3` + json sidecar, idempotent per step. + +## 8. Gap analysis: PR #51 primitives vs the above + +Already aligned: adapter math for all 9 types (torch-parity tested), rf loss target ++ masked per-sample-then-mean reduction, the 5 timestep samplers + full/flux/logsnr +shifts with matching defaults, fp16 checkpoint tensors + `lora_config` metadata, +`.safetensors`-only trust boundary, fp32 adapter params over a fp16 base, SVD sign +canonicalization, `[i-j]` include/exclude expansion, whole-file chunked encoding +with validity masks. + +To build (roughly in order): +1. **lora.py**: save with underfit key roots (`model.` + `conditioners.`) instead of + bare DiT keys; adapter for the baked seconds-conditioner Linear + (`cond.seconds_total_weight` ↔ `conditioners.seconds_total.embedder.embedding.1`); + match include/exclude against checkpoint-convention names so dashboard filter + strings work verbatim; config-layer defaults rank 8 / alpha=rank / lora; + metadata `step/epoch/base_model` injection helpers; accept `svd_bases_path` + `.pt`-equivalent (or keep recompute fallback, which underfit also has). +2. **conditioning + dropout**: training-time conditioning builder (T5Gemma frozen, + prompt padding, seconds token; MLX runtime pieces exist) + per-sample + cross_attn→zeros dropout at prob 0.1 (keep global_cond). +3. **trainer loop** (`lora_train`-equivalent CLI): AdamW (torch-default betas/eps/wd, + lr required), no schedule, no clip, batch 1, seed 42, step-driven with + checkpoint-every-1000 + final save, underfit filename convention, resume-offset + ladder, loss/lr/grad_norm logging, `loss_by_timestep.bin`. +4. **dataset**: PreEncodedDataset-equivalent (npy+json pairs, latent_crop_length + 1300/4096 defaults, random_crop, silence/zero pad, mask handling, seconds_total + NOT recomputed, oversample-with-replacement ×100 for tiny sets) + + prompt_templates port (50/50 shuffle/subset, trigger 80%, display-name map). +5. **pre-encode CLI** on top of `audio_encoding.py` (600 s cap, fp32, unscaled fp32 + npy + json sidecar with the exact metadata fields). +6. **demos** (optional, last): pipeline-based demo step with underfit cadence/format. From 06fe643986a0a02e68cda3994ebed3bdaf028469 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 13:46:56 -0400 Subject: [PATCH 06/28] =?UTF-8?q?mlx=20lora:=20underfit=20checkpoint=20con?= =?UTF-8?q?ventions=20=E2=80=94=20key=20roots,=20conditioner=20adapter,=20?= =?UTF-8?q?defaults?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Checkpoint keys now use underfit's roots (model. prefix for DiT layers via a checkpoint_prefix on inject; conditioners.seconds_total.embedder. embedding.1 for the seconds conditioner) — verified key-for-key and shape-for-shape against a real underfit checkpoint. include/exclude filters match both bare runtime names and checkpoint names so dashboard filter strings work verbatim. New: underfit_lora_config() (dora-rows, rank 16, alpha=rank, the dashboard exclude list), inject_from_lora_config (config-layer fallbacks rank 8/alpha=rank/lora), TrainableSecondsEmbedder (bit-exact vs the inference embedder at init), load_trainable_lora_state for resume (strict=False, DoRA magnitude squeeze). 24 tests. --- optimized/mlx/models/defs/lora.py | 263 +++++++++++++++++++++++++++- tests/conftest.py | 22 ++- tests/test_mlx_lora.py | 276 +++++++++++++++++++++++++++++- 3 files changed, 543 insertions(+), 18 deletions(-) diff --git a/optimized/mlx/models/defs/lora.py b/optimized/mlx/models/defs/lora.py index 8628fa2..fdfb324 100644 --- a/optimized/mlx/models/defs/lora.py +++ b/optimized/mlx/models/defs/lora.py @@ -72,6 +72,7 @@ def __init__( alpha: float, source_name: str, adapter_type: str = "lora", + checkpoint_prefix: str = "model.", ): super().__init__() if rank <= 0: @@ -83,7 +84,7 @@ def __init__( self.alpha = float(alpha) self.scaling = self.alpha / self.rank self.source_name = str(source_name) - self.checkpoint_name = _checkpoint_layer_name(self.source_name) + self.checkpoint_name = _checkpoint_name(source_name, checkpoint_prefix) self.adapter_type = canonical_adapter_type(adapter_type) fan_out, fan_in = (int(value) for value in base.weight.shape) @@ -125,6 +126,7 @@ def __init__( alpha: float, source_name: str, adapter_type: str = "lora", + checkpoint_prefix: str = "model.", ): super().__init__() if rank <= 0: @@ -136,7 +138,7 @@ def __init__( self.alpha = float(alpha) self.scaling = self.alpha / self.rank self.source_name = str(source_name) - self.checkpoint_name = _checkpoint_layer_name(self.source_name) + self.checkpoint_name = _checkpoint_name(source_name, checkpoint_prefix) self.adapter_type = canonical_adapter_type(adapter_type) fan_out, kernel_size, fan_in_per_group = ( @@ -203,6 +205,95 @@ def __call__(self, x): TrainableLoRALayer = LoRALinear | LoRAConv1d +class _ExpoFourierFeatures(nn.Module): + """MLX port of stable_audio_tools ExpoFourierFeatures (fp32, no params). + + Occupies ``embedder.embedding.0`` in TrainableSecondsEmbedder so the + trainable Linear lands at ``embedder.embedding.1``, matching underfit's + torch Sequential(ExpoFourierFeatures, Linear) checkpoint layout. The math + mirrors sa3_pipeline.expo_fourier_features exactly. + """ + + def __init__( + self, + dim: int = 256, + min_freq: float = 0.5, + max_freq: float = 10000.0, + ): + super().__init__() + self.dim = int(dim) + self.min_freq = float(min_freq) + self.max_freq = float(max_freq) + + def __call__(self, t: mx.array) -> mx.array: + t = t.astype(mx.float32).reshape(-1, 1) + half = self.dim // 2 + ramp = mx.arange(half, dtype=mx.float32) / max(half - 1, 1) + freqs = mx.exp( + ramp * (math.log(self.max_freq) - math.log(self.min_freq)) + + math.log(self.min_freq) + ) + args = t * freqs * 2 * math.pi + return mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) + + +class TrainableSecondsEmbedder(nn.Module): + """Trainable seconds_total conditioner (underfit checkpoint-compatible). + + Built from the baked conditioner weights as loaded by + sa3_pipeline.load_conditioner_from_npz: ``weight`` [768, 256] fp32 and + ``bias`` [768] fp32. Replicates SecondsTotalEmbedder's exact math (clip to + [min_val, max_val], normalize, expo Fourier features, Linear) and returns + the [B, 1, 768] seconds token. + + The module tree places the Linear at the runtime name + ``embedder.embedding.1`` so injecting with + ``checkpoint_prefix="conditioners.seconds_total."`` reproduces underfit's + checkpoint key root ``conditioners.seconds_total.embedder.embedding.1``. + When injected, the base Linear is frozen and only the adapter trains. + """ + + def __init__( + self, + weight: mx.array, + bias: mx.array, + *, + min_val: float = 0.0, + max_val: float = 384.0, + fourier_dim: int = 256, + ): + super().__init__() + weight = mx.array(weight).astype(mx.float32) + bias = mx.array(bias).astype(mx.float32) + fan_out, fan_in = (int(value) for value in weight.shape) + if fan_in != int(fourier_dim): + raise ValueError( + f"Conditioner weight fan-in {fan_in} does not match " + f"fourier_dim {fourier_dim}." + ) + linear = nn.Linear(fan_in, fan_out) + linear.weight = weight + linear.bias = bias + embedder = nn.Module() + embedder.embedding = [_ExpoFourierFeatures(dim=int(fourier_dim)), linear] + self.embedder = embedder + self.min_val = float(min_val) + self.max_val = float(max_val) + self.fourier_dim = int(fourier_dim) + + def __call__(self, seconds: float | tp.Sequence[float] | mx.array) -> mx.array: + if isinstance(seconds, (int, float)): + seconds = [float(seconds)] + if not isinstance(seconds, mx.array): + seconds = mx.array([float(value) for value in seconds], dtype=mx.float32) + values = seconds.astype(mx.float32).reshape(-1) + values = mx.clip(values, self.min_val, self.max_val) + norm = (values - self.min_val) / (self.max_val - self.min_val) + features = self.embedder.embedding[0](norm) # (B, fourier_dim) + token = self.embedder.embedding[1](features) # (B, 768) + return token[:, None, :] # (B, 1, 768) + + def inject_trainable_lora( model: nn.Module, *, @@ -211,8 +302,17 @@ def inject_trainable_lora( include: tp.Sequence[str] | None = None, exclude: tp.Sequence[str] | None = None, adapter_type: str = "lora", + checkpoint_prefix: str = "model.", ) -> LoRAInjectionReport: - """Freeze an MLX model and replace selected Linear/Conv1d layers.""" + """Freeze an MLX model and replace selected Linear/Conv1d layers. + + ``checkpoint_prefix`` is prepended to each layer's remapped runtime name + to form its checkpoint key root (underfit convention: DiT layers save as + ``model.``, the seconds conditioner as + ``conditioners.seconds_total.``). Include/exclude patterns match + against both the bare runtime name and the full checkpoint name, so + underfit dashboard filter strings work verbatim. + """ alpha = float(rank if alpha is None else alpha) adapter_type = canonical_adapter_type(adapter_type) @@ -220,7 +320,12 @@ def inject_trainable_lora( replacements: list[tuple[str, TrainableLoRALayer]] = [] for name, layer in model.named_modules(): - if not name or not _name_is_selected(name, include=include, exclude=exclude): + if not name or not _name_is_selected( + name, + _checkpoint_name(name, checkpoint_prefix), + include=include, + exclude=exclude, + ): continue if isinstance(layer, nn.Linear): replacement = LoRALinear( @@ -229,6 +334,7 @@ def inject_trainable_lora( alpha=alpha, source_name=name, adapter_type=adapter_type, + checkpoint_prefix=checkpoint_prefix, ) elif isinstance(layer, nn.Conv1d): replacement = LoRAConv1d( @@ -237,6 +343,7 @@ def inject_trainable_lora( alpha=alpha, source_name=name, adapter_type=adapter_type, + checkpoint_prefix=checkpoint_prefix, ) else: continue @@ -256,6 +363,77 @@ def inject_trainable_lora( ) +def underfit_lora_config() -> dict[str, tp.Any]: + """Effective underfit product defaults (what the dashboard writes). + + The exclude list is exactly what underfit's dashboard puts in the + ``lora_config`` metadata of trained checkpoints (confirmed against a real + underfit checkpoint): adapters go on every attention/FF linear, while the + embedders/projections/convs stay frozen. + """ + + return { + "adapter_type": "dora-rows", + "rank": 16, + "alpha": 16, + "include": None, + "exclude": [ + "to_timestep_embed", + "to_cond_embed", + "to_global_embed", + "to_local_embed", + "global_cond_embedder", + "project_in", + "project_out", + "preprocess_conv", + "postprocess_conv", + ], + } + + +def inject_from_lora_config( + model: nn.Module, + lora_config: dict[str, tp.Any] | None, + *, + checkpoint_prefix: str = "model.", +) -> tuple[LoRAInjectionReport, dict[str, tp.Any]]: + """Inject adapters from an underfit ``lora_config`` dict. + + Mirrors underfit's apply_lora_from_config config-layer fallbacks: rank + defaults to 8, alpha defaults to rank, adapter_type defaults to "lora" + (legacy "dora" resolves to "dora-rows"). Returns the injection report and + the saved-config dict ({rank, alpha, adapter_type, include, exclude}) + ready to pass to save metadata. + """ + + lora_config = dict(lora_config or {}) + rank = lora_config.get("rank") + rank = 8 if rank is None else int(rank) + alpha = lora_config.get("alpha") + alpha = rank if alpha is None else alpha + adapter_type = canonical_adapter_type(lora_config.get("adapter_type") or "lora") + include = lora_config.get("include") + exclude = lora_config.get("exclude") + + report = inject_trainable_lora( + model, + rank=rank, + alpha=float(alpha), + include=include, + exclude=exclude, + adapter_type=adapter_type, + checkpoint_prefix=checkpoint_prefix, + ) + saved_config = { + "rank": rank, + "alpha": alpha, + "adapter_type": adapter_type, + "include": include, + "exclude": exclude, + } + return report, saved_config + + def iter_trainable_lora_layers( model: nn.Module, ) -> tp.Iterator[TrainableLoRALayer]: @@ -309,7 +487,8 @@ def save_lora_checkpoint( config: dict[str, tp.Any] = { "rank": rank, - "alpha": alpha, + # Underfit writes integral alphas as JSON ints — keep metadata parity. + "alpha": int(alpha) if float(alpha).is_integer() else alpha, "adapter_type": adapter_type, "include": list(include) if include else None, "exclude": list(exclude) if exclude else None, @@ -334,6 +513,60 @@ def save_lora_checkpoint( return output_path +def load_trainable_lora_state( + model: nn.Module, + state_dict: dict[str, tp.Any], +) -> int: + """Restore trainable adapter parameters from an underfit checkpoint. + + ``state_dict`` uses underfit checkpoint keys + (``.parametrizations.weight..``). Each group is + matched to the injected trainable layer with the same checkpoint_name; + unmatched groups and unknown params are tolerated (torch strict=False + semantics). 2-D DoRA magnitudes are squeezed to 1-D like underfit's + prepare_dora_state_dict. Parameters are restored in fp32. Returns the + number of layers restored. + """ + + grouped = _group_lora_state_dict(state_dict) + layers = { + layer.checkpoint_name: layer for layer in iter_trainable_lora_layers(model) + } + restored = 0 + for checkpoint_name, params in grouped.items(): + layer = layers.get(checkpoint_name) + if layer is None: + continue + applied = False + for param_name in ( + "lora_A", + "lora_B", + "magnitude", + "magnitude_r", + "magnitude_c", + "M_xs", + ): + value = params.get(param_name) + if value is None: + continue + current = getattr(layer, param_name, None) + if current is None: + continue + array = np.asarray(value, dtype=np.float32) + if param_name == "magnitude" and array.ndim == 2: + array = array.squeeze() + if tuple(array.shape) != tuple(current.shape): + raise ValueError( + f"Checkpoint tensor {checkpoint_name}.{param_name} has " + f"shape {tuple(array.shape)}, expected " + f"{tuple(current.shape)}." + ) + setattr(layer, param_name, mx.array(array, dtype=mx.float32)) + applied = True + restored += int(applied) + return restored + + def load_lora_checkpoint( path: str | Path, ) -> tuple[dict[str, mx.array], dict[str, tp.Any]]: @@ -805,6 +1038,10 @@ def _checkpoint_layer_name(name: str) -> str: ) +def _checkpoint_name(source_name: str, checkpoint_prefix: str) -> str: + return str(checkpoint_prefix or "") + _checkpoint_layer_name(str(source_name)) + + def _source_shape_for_delta( target_shape: tuple[int, ...], delta_shape: tuple[int, int], @@ -926,14 +1163,22 @@ def _canonicalize_svd_signs( def _name_is_selected( - name: str, - *, + *names: str, include: tp.Sequence[str] | None, exclude: tp.Sequence[str] | None, ) -> bool: - if include and not _matches_any(name, include): + """Match include/exclude filters against any of the given aliases. + + Callers pass both the bare runtime name and the full checkpoint name so + underfit dashboard filter strings (written in torch/checkpoint naming) + select the same layers verbatim. + """ + + if include and not any(_matches_any(name, include) for name in names): return False - return not (exclude and _matches_any(name, exclude)) + return not ( + exclude and any(_matches_any(name, exclude) for name in names) + ) def _matches_any(name: str, patterns: tp.Sequence[str]) -> bool: diff --git a/tests/conftest.py b/tests/conftest.py index 07ac1f0..fd5e946 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,18 +1,28 @@ +from __future__ import annotations + import re from pathlib import Path import pytest -import torch -import torchaudio -from stable_audio_3 import AutoencoderModel, StableAudioModel -from stable_audio_3.model_configs import ae_models, base_models +try: # torch is absent in the torch-free MLX venv; MLX-only tests don't need it + import torch + import torchaudio + + from stable_audio_3 import AutoencoderModel, StableAudioModel + from stable_audio_3.model_configs import ae_models, base_models +except ModuleNotFoundError: + torch = None + torchaudio = None + AutoencoderModel = StableAudioModel = None + ae_models = () + base_models = () # --------------------------------------------------------------------------- # Hardware detection — used by fixtures and tests to gate GPU-only paths # --------------------------------------------------------------------------- -HAS_CUDA = torch.cuda.is_available() -HAS_MPS = torch.backends.mps.is_available() +HAS_CUDA = torch is not None and torch.cuda.is_available() +HAS_MPS = torch is not None and torch.backends.mps.is_available() HAS_ACCEL = HAS_CUDA or HAS_MPS ACCEL_DEVICE = "cuda" if HAS_CUDA else ("mps" if HAS_MPS else "cpu") diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py index da9eb3a..1d460c8 100644 --- a/tests/test_mlx_lora.py +++ b/tests/test_mlx_lora.py @@ -14,12 +14,18 @@ from optimized.mlx.models.defs import dit_mlx, dit_mlx_medium from optimized.mlx.models.defs.lora import ( + TrainableSecondsEmbedder, apply_lora_checkpoint, apply_lora_checkpoints, + inject_from_lora_config, inject_trainable_lora, + iter_trainable_lora_layers, load_lora_checkpoint, + load_trainable_lora_state, save_lora_checkpoint, + underfit_lora_config, ) +from optimized.mlx.models.defs.sa3_pipeline import SecondsTotalEmbedder from stable_audio_3.models.lora import ( LoRAParametrization, add_lora, @@ -28,6 +34,16 @@ save_lora_safetensors, ) +# Real underfit-trained checkpoint (medium DiT, dora-rows rank 16) used to +# verify byte-convention key-naming parity with underfit's saver. +UNDERFIT_REFERENCE_CHECKPOINT = Path( + "/Users/cj/clod/speed-metal/scripts/lora_bench/plini-sa3-380.safetensors" +) +needs_underfit_reference = pytest.mark.skipif( + not UNDERFIT_REFERENCE_CHECKPOINT.exists(), + reason="real underfit reference checkpoint not available", +) + class TinyMLXLinear(nn.Module): def __init__(self): @@ -184,15 +200,25 @@ def test_mlx_checkpoint_round_trips_through_official_torch_loader( checkpoint = save_lora_checkpoint( model, tmp_path / "mlx-dora.safetensors", - extra_config={"step": 20}, + extra_config={"step": 20, "epoch": 3, "base_model": "sa3-medium"}, ) mlx_state, mlx_config = load_lora_checkpoint(checkpoint) torch_state, torch_config = load_torch_lora_checkpoint(checkpoint) + # Underfit key convention: "model." + the bare runtime name. Underfit + # loads with strict=False into model.model/model.conditioner, so the + # model.-prefixed keys are what the official stack expects. + assert sorted(mlx_state) == [ + "model.layer.parametrizations.weight.0.lora_A", + "model.layer.parametrizations.weight.0.lora_B", + "model.layer.parametrizations.weight.0.magnitude", + ] assert sorted(mlx_state) == sorted(torch_state) assert mlx_config == torch_config assert mlx_config["adapter_type"] == "dora-rows" assert mlx_config["step"] == 20 + assert mlx_config["epoch"] == 3 + assert mlx_config["base_model"] == "sa3-medium" @pytest.mark.parametrize( @@ -368,8 +394,8 @@ def __init__(self): state_dict, _ = load_torch_lora_checkpoint(checkpoint) assert sorted(state_dict) == [ - "to_local_embed.0.parametrizations.weight.0.lora_A", - "to_local_embed.0.parametrizations.weight.0.lora_B", + "model.to_local_embed.0.parametrizations.weight.0.lora_A", + "model.to_local_embed.0.parametrizations.weight.0.lora_B", ] @@ -405,3 +431,247 @@ def test_conv1d_checkpoint_maps_pytorch_weight_layout_to_mlx(tmp_path: Path): assert report.applied_layers == 1 assert np.allclose(np.asarray(mlx_model.layer.weight), expected, atol=2e-3) + + +def test_underfit_lora_config_matches_dashboard_defaults(): + assert underfit_lora_config() == { + "adapter_type": "dora-rows", + "rank": 16, + "alpha": 16, + "include": None, + "exclude": [ + "to_timestep_embed", + "to_cond_embed", + "to_global_embed", + "to_local_embed", + "global_cond_embedder", + "project_in", + "project_out", + "preprocess_conv", + "postprocess_conv", + ], + } + + +class WideMLXRegressor(nn.Module): + """Layers wide enough for underfit's config-layer default rank (8).""" + + def __init__(self): + super().__init__() + self.input = nn.Linear(16, 12, bias=False) + self.output = nn.Linear(12, 8, bias=False) + + def __call__(self, x): + return self.output(nn.silu(self.input(x))) + + +def test_inject_from_lora_config_applies_underfit_config_layer_fallbacks(): + model = WideMLXRegressor() + report, saved_config = inject_from_lora_config(model, {}) + + assert report.adapter_type == "lora" + assert saved_config == { + "rank": 8, + "alpha": 8, + "adapter_type": "lora", + "include": None, + "exclude": None, + } + for layer in iter_trainable_lora_layers(model): + assert layer.rank == 8 + assert layer.alpha == 8.0 + assert layer.checkpoint_name == f"model.{layer.source_name}" + + legacy = WideMLXRegressor() + report, saved_config = inject_from_lora_config( + legacy, + {"adapter_type": "dora", "rank": 4, "include": ["output"]}, + ) + + assert report.adapter_type == "dora-rows" + assert report.layer_names == ("output",) + assert saved_config == { + "rank": 4, + "alpha": 4, + "adapter_type": "dora-rows", + "include": ["output"], + "exclude": None, + } + + +def test_include_filters_match_checkpoint_convention_names(): + model = dit_mlx.DiT(T_lat=8) + # Underfit dashboard filter strings are written in torch/checkpoint + # naming: "model." prefix and "to_local_embed.0" (no ".seq"). + report = inject_trainable_lora( + model, + rank=1, + include=[ + "model.transformer.layers.[0-1].self_attn.to_qkv", + "model.transformer.layers.0.to_local_embed.0", + ], + exclude=["model.transformer.layers.1."], + ) + + assert set(report.layer_names) == { + "transformer.layers.0.self_attn.to_qkv", + "transformer.layers.0.to_local_embed.seq.0", + } + + +def test_trainable_seconds_embedder_matches_pipeline_conditioner(): + rng = np.random.default_rng(2) + weight = mx.array(rng.standard_normal((768, 256)).astype(np.float32)) + bias = mx.array(rng.standard_normal(768).astype(np.float32)) + reference = SecondsTotalEmbedder(weight, bias) + trainable = TrainableSecondsEmbedder(weight, bias) + + for seconds in ([30.0], [0.0, 1.5, 47.25, 285.0, 384.0, 500.0]): + expected = reference(seconds) + actual = trainable(seconds) + assert actual.shape == (len(seconds), 1, 768) + assert bool(mx.all(actual == expected)) + # mx.array input path (what the trainer passes) is bit-identical too. + assert bool( + mx.all(trainable(mx.array([12.5, 380.0])) == reference([12.5, 380.0])) + ) + + report, _ = inject_from_lora_config( + trainable, + underfit_lora_config(), + checkpoint_prefix="conditioners.seconds_total.", + ) + layer = next(iter(iter_trainable_lora_layers(trainable))) + + assert report.layer_names == ("embedder.embedding.1",) + assert layer.checkpoint_name == ( + "conditioners.seconds_total.embedder.embedding.1" + ) + # Base Linear is frozen — only the adapter trains. + assert sorted( + name for name, _ in tree_flatten(trainable.trainable_parameters()) + ) == [ + "embedder.embedding.1.lora_A", + "embedder.embedding.1.lora_B", + "embedder.embedding.1.magnitude", + ] + assert trainable([30.0]).shape == (1, 1, 768) + + +def test_load_trainable_lora_state_round_trips_saved_adapters(tmp_path: Path): + def build_model(): + mx.random.seed(11) + model = TinyMLXRegressor() + inject_trainable_lora(model, rank=2, alpha=2, adapter_type="dora-rows") + return model + + trained = build_model() + rng = np.random.default_rng(5) + for layer in iter_trainable_lora_layers(trained): + for name in ("lora_A", "lora_B", "magnitude"): + shape = tuple(getattr(layer, name).shape) + # fp16-representable values so the fp16 checkpoint is lossless. + values = rng.standard_normal(shape).astype(np.float16) + setattr(layer, name, mx.array(values.astype(np.float32))) + + checkpoint = save_lora_checkpoint(trained, tmp_path / "resume.safetensors") + state_dict, _ = load_lora_checkpoint(checkpoint) + inputs = mx.array([[1.0, -2.0, 0.5], [0.25, 0.75, -1.5]], dtype=mx.float32) + + fresh = build_model() + assert not np.array_equal(np.asarray(trained(inputs)), np.asarray(fresh(inputs))) + + restored = load_trainable_lora_state(fresh, state_dict) + + assert restored == 2 + assert np.array_equal(np.asarray(trained(inputs)), np.asarray(fresh(inputs))) + for layer in iter_trainable_lora_layers(fresh): + assert layer.lora_A.dtype == mx.float32 + assert layer.magnitude.dtype == mx.float32 + + # strict=False semantics: unmatched checkpoint layers are tolerated, and + # 2-D DoRA magnitudes are squeezed like prepare_dora_state_dict. + partial = { + key: value + for key, value in state_dict.items() + if key.startswith("model.output.") + } + magnitude_key = "model.output.parametrizations.weight.0.magnitude" + partial[magnitude_key] = partial[magnitude_key].reshape(-1, 1) + partial["model.missing.parametrizations.weight.0.lora_A"] = mx.zeros((2, 3)) + + assert load_trainable_lora_state(build_model(), partial) == 1 + + +@needs_underfit_reference +def test_checkpoint_key_naming_matches_real_underfit_checkpoint(tmp_path: Path): + reference_state, reference_config = load_lora_checkpoint( + UNDERFIT_REFERENCE_CHECKPOINT + ) + config = underfit_lora_config() + + # The reference checkpoint's lora_config metadata is exactly the product + # defaults (plus the injected step/epoch fields). + assert {key: reference_config[key] for key in config} == config + + # Injecting the medium DiT with the product config selects exactly the + # layer set underfit trained, with byte-identical checkpoint keys. + model = dit_mlx_medium.DiT(T_lat=8) + report, _ = inject_from_lora_config(model, config) + produced_keys = set() + for layer in iter_trainable_lora_layers(model): + root = f"{layer.checkpoint_name}.parametrizations.weight.0" + produced_keys.update( + f"{root}.{param}" for param in ("lora_A", "lora_B", "magnitude") + ) + reference_dit_keys = { + key for key in reference_state if key.startswith("model.") + } + + assert produced_keys == reference_dit_keys + assert report.layer_count == len(reference_dit_keys) // 3 + + # Resume: the real underfit checkpoint loads into the injected model. + assert load_trainable_lora_state(model, reference_state) == report.layer_count + + # The saver reproduces a real medium DiT layer's keys AND shapes exactly. + single = dit_mlx_medium.DiT(T_lat=8) + inject_from_lora_config( + single, dict(config, include=["layers.0.self_attn.to_qkv"]) + ) + saved_state, _ = load_lora_checkpoint( + save_lora_checkpoint(single, tmp_path / "single-layer.safetensors") + ) + layer_root = "model.transformer.layers.0.self_attn.to_qkv" + expected = { + key: tuple(value.shape) + for key, value in reference_state.items() + if key.startswith(f"{layer_root}.") + } + + assert len(expected) == 3 + assert {key: tuple(value.shape) for key, value in saved_state.items()} == expected + + # The conditioner layer comes out at underfit's exact key root too. + rng = np.random.default_rng(3) + seconds_module = TrainableSecondsEmbedder( + mx.array(rng.standard_normal((768, 256)).astype(np.float32)), + mx.array(rng.standard_normal(768).astype(np.float32)), + ) + inject_from_lora_config( + seconds_module, config, checkpoint_prefix="conditioners.seconds_total." + ) + saved_cond, _ = load_lora_checkpoint( + save_lora_checkpoint(seconds_module, tmp_path / "conditioner.safetensors") + ) + expected_cond = { + key: tuple(value.shape) + for key, value in reference_state.items() + if key.startswith("conditioners.") + } + + assert len(expected_cond) == 3 + assert { + key: tuple(value.shape) for key, value in saved_cond.items() + } == expected_cond + assert load_trainable_lora_state(seconds_module, reference_state) == 1 From dbea9acc6e7c9bc8036258c34b0214ea635a2375 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 13:46:56 -0400 Subject: [PATCH 07/28] mlx training data: pre-encode CLI + latent dataset with underfit semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pre_encode_mlx.py: offline audio→latent CLI (whole files ≤600 s aligned to 4096, fp32 SAME encoders, chunked >30 s, no normalization; npy + json with seconds_total/seconds_start/audio_samples/latent_shape/padding_mask and sidecar tag extraction — underfit field parity). latent_dataset.py: PreEncodedDataset port (latent_crop_length crop with random_crop over the mask-valid region, silence/zero padding, seconds_total kept at FULL duration after cropping, min/max-length rejection redraw) + prompt_templates port (tag prompts with the 50/50 shuffle-vs-subset augmentation, trigger token at 80%, path/fixed sources, legacy fallback) + iterate_batches with the oversample-with-replacement x100 tiny-dataset behavior. 31 tests, all synthetic fixtures. --- optimized/mlx/models/defs/latent_dataset.py | 388 ++++++++++++++++++ optimized/mlx/scripts/pre_encode_mlx.py | 419 ++++++++++++++++++++ tests/test_mlx_latent_dataset.py | 327 +++++++++++++++ tests/test_mlx_pre_encode.py | 313 +++++++++++++++ 4 files changed, 1447 insertions(+) create mode 100644 optimized/mlx/models/defs/latent_dataset.py create mode 100644 optimized/mlx/scripts/pre_encode_mlx.py create mode 100644 tests/test_mlx_latent_dataset.py create mode 100644 tests/test_mlx_pre_encode.py diff --git a/optimized/mlx/models/defs/latent_dataset.py b/optimized/mlx/models/defs/latent_dataset.py new file mode 100644 index 0000000..aa36b15 --- /dev/null +++ b/optimized/mlx/models/defs/latent_dataset.py @@ -0,0 +1,388 @@ +"""Torch-free pre-encoded latent dataset for LoRA training (numpy + stdlib only). + +Faithful port of the underfit data pipeline (TRAINING_CONVENTIONS.md §7): + +- ``PreEncodedLatentDataset`` mirrors stable_audio_3.data.dataset.PreEncodedDataset + (latent cropping, silence padding, mask handling, min/max-length rejection with + random redraw) over the underfit pre-encode format: ``.npy`` latents + ``[D, T]`` float32 + ``.json`` sidecar metadata. +- ``build_prompt`` ports underfit's dataset_processing/prompt_templates.py + (tag/path/fixed sources with weighted balance, 50/50 shuffle-vs-subset tag + augmentation, trigger token at ``trigger_pct``), falling back to the legacy + songs_simple.py fixed-order tag prompt when no prompt_config is given. +- ``iterate_batches`` mirrors underfit's tiny-dataset core behavior: when + ``len(dataset) < batch_size``, sample with replacement for ``batch_size * 100`` + samples per epoch (~100 batches); otherwise a shuffled epoch, drop_last=False. + +Deliberate underfit convention: ``seconds_total`` stays the FULL stored duration +and is never recomputed after cropping. +""" + +import json +import os +import random + +import numpy as np + +__all__ = ["PreEncodedLatentDataset", "build_prompt", "iterate_batches"] + + +# --------------------------------------------------------------------------- +# Prompt building (port of underfit dataset_processing/prompt_templates.py) +# --------------------------------------------------------------------------- + +_TAG_DISPLAY = { + "title": "Title", "artist": "Artist", "album": "Album", + "genre": "Genre", "label": "Label", "date": "Year", + "composer": "Composer", "bpm": "BPM", "prompt": "Prompt", +} +_ALL_TAG_KEYS = list(_TAG_DISPLAY.keys()) + +# Legacy (no prompt_config) fixed tag order — from underfit songs_simple.py. +_LEGACY_TAG_KEYS = ["artist", "title", "date", "bpm", "genre", "label", "album", "composer"] + + +def _get(metadata, key): + val = metadata.get(key, "") + if isinstance(val, (list, tuple)): + val = val[0] if val else "" + return str(val).strip() + + +def _build_tag_prompt(metadata, pc, rng): + tag_keys = pc.get("tag_keys", _ALL_TAG_KEYS) + hide_names = pc.get("hide_tag_names", False) + hide_commas = pc.get("hide_commas", False) + split_commas = pc.get("split_commas", False) + + parts = [] + for key in tag_keys: + val = _get(metadata, key) + if not val: + continue + label = _TAG_DISPLAY.get(key, key) + if split_commas and "," in val: + for sub in val.split(","): + sub = sub.strip() + if sub: + parts.append(sub if hide_names else f"{label}: {sub}") + else: + parts.append(val if hide_names else f"{label}: {val}") + + if not parts: + return "" + + if pc.get("shuffle", True): + # 50% shuffle all, 50% random subset + if rng.random() < 0.5: + rng.shuffle(parts) + else: + parts = rng.sample(parts, rng.randint(1, len(parts))) + + return (" " if hide_commas else ", ").join(parts) + + +def _build_path_prompt(metadata, pc): + relpath = _get(metadata, "relpath") + if not relpath: + return "" + + path_opts = pc.get("path_opts", {}) + # Frontend sends camelCase keys + hide_ext = path_opts.get("hideExt", path_opts.get("hide_ext", False)) + hide_dirs = path_opts.get("hideDirs", path_opts.get("hide_dirs", False)) + hide_topmost = path_opts.get("hideTopmostDir", path_opts.get("hide_topmost_dir", False)) + + parts = relpath.replace("\\", "/").split("/") + filename = parts[-1] + dirs = parts[:-1] + + if hide_ext: + dot = filename.rfind(".") + if dot > 0: + filename = filename[:dot] + + if hide_dirs: + if hide_topmost or not dirs: + return filename + return dirs[0] + "/" + filename + + if hide_topmost and dirs: + return os.path.join(*dirs[1:], filename) if len(dirs) > 1 else filename + + return os.path.join(*dirs, filename) if dirs else filename + + +def _legacy_prompt(metadata, rng): + """No-prompt_config fallback: underfit songs_simple.py behavior.""" + properties = [] + for key in _LEGACY_TAG_KEYS: + val = _get(metadata, key) + if val: + properties.append(f"{_TAG_DISPLAY[key]}: {val}") + + if not properties: + return str(metadata.get("text", "")) + + # 50% shuffle all, 50% random subset + if rng.random() < 0.5: + rng.shuffle(properties) + else: + properties = rng.sample(properties, rng.randint(1, len(properties))) + return ", ".join(properties) + + +def build_prompt(metadata, prompt_config, rng): + """Build one training prompt from sidecar metadata. + + Port of prompt_templates.get_custom_metadata (returns the prompt string + directly; underfit's ``lyrics`` is always "" and is omitted here). ``rng`` + is a ``random.Random`` so every __getitem__ draw is fresh (per-epoch + prompt augmentation) yet reproducible from the dataset seed. + """ + pc = prompt_config + + if pc is None: + return _legacy_prompt(metadata, rng) + + use_tags = pc.get("use_tags", True) + use_paths = pc.get("use_paths", False) + use_fixed = pc.get("use_fixed", False) + fixed_text = pc.get("fixed_text", "") + balance = pc.get("balance", {}) + trigger = pc.get("trigger", "") + trigger_pct = pc.get("trigger_pct", 80) + + # Build candidate prompts for each enabled method + candidates = [] # (method_name, prompt_text) + weights = [] + + if use_tags: + candidates.append(("tags", _build_tag_prompt(metadata, pc, rng))) + weights.append(balance.get("tags", 50)) + + if use_paths: + candidates.append(("paths", _build_path_prompt(metadata, pc))) + weights.append(balance.get("paths", 50)) + + if use_fixed: + candidates.append(("fixed", fixed_text)) + weights.append(balance.get("fixed", 0)) + + # Pick one method based on balance weights + chosen_method = None + prompt = "" + if not candidates: + prompt = str(metadata.get("text", "")) + else: + total = sum(weights) + if total <= 0: + chosen_method, prompt = rng.choice(candidates) + else: + chosen_method, prompt = rng.choices(candidates, weights=weights, k=1)[0] + + # Prepend trigger token based on trigger_pct + if trigger and trigger_pct > 0 and rng.random() * 100 < trigger_pct: + if prompt: + use_comma = chosen_method == "tags" and not pc.get("hide_commas", False) + prompt = trigger + (", " if use_comma else " ") + prompt + else: + prompt = trigger + + return prompt + + +# --------------------------------------------------------------------------- +# Dataset (port of PreEncodedDataset semantics, torch-free) +# --------------------------------------------------------------------------- + +class PreEncodedLatentDataset: + """Pre-encoded latent dataset over ``.npy`` + ``.json`` pairs. + + Scans ``root_dir`` recursively for underfit pre-encode outputs: ``.npy`` + latents ``[D, T]`` float32 with a matching ``.json`` sidecar containing + ``seconds_total``, ``seconds_start``, ``audio_samples``, ``latent_shape``, + ``padding_mask`` (0/1 at latent resolution) plus arbitrary tag keys + (title/artist/genre/prompt/...). ``silence.npy`` at the root is used to + pad short items (never served as a sample). + + ``__getitem__`` returns a dict: + latents np.float32 ``[D, latent_crop_length]`` + padding_mask np.bool_ ``[latent_crop_length]`` (True = valid frame) + prompt str (freshly randomized every call — per-epoch augmentation) + seconds_total float — the FULL stored duration (never recomputed on crop) + relpath str (path of the .npy relative to root_dir) + + Items outside [min_length_sec, max_length_sec] and items that fail to load + are rejected by redrawing a different random index, like the reference. + """ + + _MAX_REDRAWS = 100 + + def __init__( + self, + root_dir, + latent_crop_length, + *, + random_crop=True, + prompt_config=None, + min_length_sec=None, + max_length_sec=None, + seed=42, + ): + self.root_dir = os.path.abspath(root_dir) + self.latent_crop_length = latent_crop_length + self.random_crop = random_crop + self.prompt_config = prompt_config + self.min_length_sec = min_length_sec + self.max_length_sec = max_length_sec + self._rng = random.Random(seed) + + # Silence latent for variable-length padding (reference stores [1, C, N]) + self.silence_latents = None + silence_path = os.path.join(self.root_dir, "silence.npy") + if os.path.exists(silence_path): + silence = np.load(silence_path) + if silence.ndim == 3: + silence = silence.squeeze(0) # [C, N] + self.silence_latents = np.asarray(silence, dtype=np.float32) + + # Recursive scan for npy+json pairs (sorted for determinism) + self.filenames = [] + for dirpath, _dirnames, files in os.walk(self.root_dir): + for fn in files: + if not fn.endswith(".npy") or fn == "silence.npy": + continue + npy_path = os.path.join(dirpath, fn) + json_path = npy_path[: -len(".npy")] + ".json" + if os.path.exists(json_path): + self.filenames.append((npy_path, json_path)) + self.filenames.sort() + + if not self.filenames: + raise ValueError(f"No .npy + .json pairs found under {self.root_dir}") + + def __len__(self): + return len(self.filenames) + + def __getitem__(self, idx): + for _ in range(self._MAX_REDRAWS): + item = self._load_item(idx) + if item is not None: + return item + # Rejected (length filter) or failed to load — redraw a random index + idx = self._rng.randrange(len(self)) + raise RuntimeError( + f"Gave up after {self._MAX_REDRAWS} redraws — every drawn item was " + f"rejected or failed to load (check min/max_length_sec and data files)" + ) + + def _load_item(self, idx): + latent_filename, md_filename = self.filenames[idx] + try: + latents = np.load(latent_filename).astype(np.float32, copy=False) # [D, T] + with open(md_filename, "r") as f: + meta = json.load(f) + + padding_mask = list(meta["padding_mask"]) + + crop = self.latent_crop_length + if crop is not None: + stored_length = latents.shape[1] + + if stored_length > crop: + # Crop: last valid index = index of the last 1 in the mask + last_ix = len(padding_mask) - 1 - padding_mask[::-1].index(1) + + if self.random_crop and last_ix > crop: + start = self._rng.randint(0, last_ix - crop) + else: + start = 0 + + latents = latents[:, start:start + crop] + padding_mask = padding_mask[start:start + crop] + + elif stored_length < crop: + # Pad with the silence latent (tiled/sliced) if present, else zeros + pad_needed = crop - stored_length + silence = self.silence_latents + + if silence is not None: + if silence.shape[1] >= pad_needed: + silence_pad = silence[:, :pad_needed] + else: + reps = (pad_needed // silence.shape[1]) + 1 + silence_pad = np.tile(silence, (1, reps))[:, :pad_needed] + latents = np.concatenate([latents, silence_pad], axis=1) + else: + latents = np.pad(latents, ((0, 0), (0, pad_needed))) + + # Valid frames from the stored mask, zeros for padding + padding_mask = padding_mask[:stored_length] + [0] * pad_needed + + # Deliberate convention: seconds_total is the FULL stored duration — + # never recomputed after cropping. + seconds_total = float(meta["seconds_total"]) + + if self.min_length_sec is not None and seconds_total < self.min_length_sec: + return None + if self.max_length_sec is not None and seconds_total > self.max_length_sec: + return None + + relpath = os.path.relpath(latent_filename, self.root_dir) + # Path prompts prefer the sidecar's own relpath (pre_encode writes it) + meta.setdefault("relpath", relpath) + + prompt = build_prompt(meta, self.prompt_config, self._rng) + + return { + "latents": np.ascontiguousarray(latents, dtype=np.float32), + "padding_mask": np.asarray(padding_mask, dtype=np.bool_), + "prompt": prompt, + "seconds_total": seconds_total, + "relpath": relpath, + } + except Exception as e: + print(f"Couldn't load file {latent_filename}: {e}") + return None + + +# --------------------------------------------------------------------------- +# Batching (underfit small-dataset oversampling + plain epochs) +# --------------------------------------------------------------------------- + +def iterate_batches(dataset, batch_size, *, shuffle=True, seed=42): + """Yield collated batches for one epoch. + + When ``len(dataset) < batch_size``: draw with replacement for + ``batch_size * 100`` samples (≈100 batches/epoch) — underfit's tiny-dataset + core behavior (RandomSampler(replacement=True, num_samples=batch_size*100)); + ``shuffle`` is irrelevant in this mode. Otherwise iterate a (optionally + shuffled) permutation of the dataset, drop_last=False. + + Each batch: {"latents": np.float32 [B, D, L], "padding_mask": np.bool_ [B, L], + "prompt": list[str], "seconds_total": list[float], "relpath": list[str]}. + """ + n = len(dataset) + if n == 0: + return + + rng = random.Random(seed) + + if n < batch_size: + # Oversample with replacement (~100 batches per epoch) + indices = [rng.randrange(n) for _ in range(batch_size * 100)] + else: + indices = list(range(n)) + if shuffle: + rng.shuffle(indices) + + for i in range(0, len(indices), batch_size): + items = [dataset[j] for j in indices[i:i + batch_size]] + yield { + "latents": np.stack([it["latents"] for it in items]).astype(np.float32, copy=False), + "padding_mask": np.stack([it["padding_mask"] for it in items]), + "prompt": [it["prompt"] for it in items], + "seconds_total": [it["seconds_total"] for it in items], + "relpath": [it["relpath"] for it in items], + } diff --git a/optimized/mlx/scripts/pre_encode_mlx.py b/optimized/mlx/scripts/pre_encode_mlx.py new file mode 100644 index 0000000..9f429fc --- /dev/null +++ b/optimized/mlx/scripts/pre_encode_mlx.py @@ -0,0 +1,419 @@ +"""Pre-encode audio files into SAME latents for LoRA finetuning (MLX, torch-free). + +Offline audio → latent CLI mirroring underfit's ``dataset_processing/pre_encode.py`` +conventions, ported to the standalone MLX runtime: single-process, soundfile-based +loading (no torch / torchaudio), and ``models.defs.audio_encoding.encode_audio`` +doing the padding / chunking / masking. + +Every audio file under ``--audio-dir`` (recursive) becomes a ``.npy`` + ``.json`` +pair at the same relative path under ``--output-dir``: + + / latents [D, T] float32 + / metadata (underfit field parity) + +Conventions matched to underfit: + * Whole files, no clip splitting; cap at ``--max-duration`` seconds (default + 600) with the sample cap aligned DOWN to a multiple of 4096 (ds_ratio). + * 44.1 kHz stereo — mono is duplicated to stereo, >2 channels keep the first + two, other sample rates are linearly resampled. NO loudness/peak norm. + * Encoder runs fp32; ``chunked=True`` for clips longer than 30 s. + * Metadata: ``path, relpath, src_relpath, seconds_total`` (rounded 3dp), + ``seconds_start: 0, audio_samples, latent_shape``, latent-resolution 0/1 + ``padding_mask``, merged with tags extracted from sidecars. + * Skip files whose outputs already exist unless ``--overwrite``. + * ``/details.json`` records the global encode params. + +Latent scaling note: the MLX SAME encoders output DiT-space (softnorm) latents +directly — the equivalent of torch's unscaled-latents-then-divide-by-scale +convention with scale 1.0 — so latents are saved as-is and the trainer uses +them directly (no pretransform-scale division at train time). + +Tag extraction (ported from underfit's ``extract_tags``): JSON sidecar +(``.json`` next to the audio, or ``/json/.json``) keeps all +non-empty string/number values; else a ``.txt`` sidecar (same dir or +``/txt/``) whose whole stripped content becomes the ``prompt`` tag. +Embedded ID3/Vorbis tags are a DELIBERATE OMISSION — they need an extra +dependency (``audio_metadata``), and sidecars are the underfit-recommended +path anyway. + +Usage: + python pre_encode_mlx.py --audio-dir ~/clips --output-dir ~/latents \ + --codec same-s [--max-duration 600] [--overwrite] + +``--codec same-s`` pairs with the sm-music / sm-sfx DiTs, ``same-l`` with medium. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parent.parent # optimized/mlx (scripts/ is one level down) +if str(REPO) not in sys.path: + sys.path.insert(0, str(REPO)) +if str(REPO / "scripts") not in sys.path: + sys.path.insert(0, str(REPO / "scripts")) + +import mlx.core as mx # noqa: E402 + +from models.defs.audio_encoding import ( # noqa: E402 + SAME_ENCODER_PAD_MODULO, + SAMPLES_PER_LATENT, + encode_audio, +) + +SAMPLE_RATE = 44100 +DS_RATIO = SAMPLES_PER_LATENT # 4096 samples per latent +CHUNKED_THRESHOLD_SECONDS = 30.0 + +AUDIO_EXTS = {".wav", ".mp3", ".flac", ".ogg", ".aif", ".aiff", ".m4a"} +_MIN_AUDIO_SIZE = 4096 # skip files smaller than this (macOS resource forks, corrupt) + + +# --------------------------------------------------------------------------- +# Audio discovery +# --------------------------------------------------------------------------- + + +def find_audio_files(root) -> list[Path]: + """Recursively find audio files under ``root``, sorted by path. + + Skips macOS resource forks (``._*``) and files too small to be real audio. + """ + files = [] + for dirpath, _, filenames in os.walk(root): + for fn in filenames: + if fn.startswith("._"): + continue + fp = Path(dirpath) / fn + if fp.suffix.lower() in AUDIO_EXTS: + try: + if fp.stat().st_size < _MIN_AUDIO_SIZE: + continue + except OSError: + continue + files.append(fp) + files.sort() + return files + + +# --------------------------------------------------------------------------- +# Audio loading (torch-free) +# --------------------------------------------------------------------------- + + +def load_audio(path) -> np.ndarray: + """Load an audio file → ``(2, T)`` float32 at 44.1 kHz. + + soundfile decodes; mono is duplicated to stereo, >2 channels keep the + first two; non-44.1 kHz input is linearly resampled at + align_corners=False positions (same pattern as sa3_gradio.read_audio_any). + No loudness/peak normalization. + """ + import soundfile as sf + + data, sr = sf.read(str(path), dtype="float32", always_2d=True) # (T, ch) + y = data.T + y = np.stack([y[0], y[0]]) if y.shape[0] == 1 else y[:2] + if sr != SAMPLE_RATE: + in_len = y.shape[-1] + new_len = int(round(in_len * SAMPLE_RATE / sr)) + scale = in_len / new_len + pos = np.clip((np.arange(new_len) + 0.5) * scale - 0.5, 0, in_len - 1) + y = np.stack([np.interp(pos, np.arange(in_len), ch) for ch in y]) + return np.ascontiguousarray(y, dtype=np.float32) + + +def max_samples_for_duration(max_duration: float) -> int: + """Sample cap for ``max_duration`` seconds, aligned DOWN to ds_ratio (4096).""" + return (int(max_duration * SAMPLE_RATE) // DS_RATIO) * DS_RATIO + + +# --------------------------------------------------------------------------- +# Sidecar tag extraction (ported from underfit pre_encode.extract_tags; +# embedded ID3/Vorbis tags deliberately omitted — see module docstring) +# --------------------------------------------------------------------------- + + +def extract_tags(filepath) -> dict: + """Extract tag fields for an audio file from sidecar files. + + Priority: + 1. JSON sidecar (``.json`` same dir, or ``/json/.json``) + — all non-empty string/number values kept (custom keys preserved). + 2. Plain ``.txt`` sidecar (same dir or ``/txt/``) — whole stripped + content becomes the ``prompt`` tag. + Returns ``{}`` when neither sidecar exists (no embedded-tag fallback). + """ + fp = Path(filepath) + stem = fp.stem + + # 1. JSON sidecar — same dir, then sibling json/ dir + sidecar_candidates = [ + fp.with_suffix(".json"), + fp.parent.parent / "json" / (stem + ".json"), + ] + for sidecar in sidecar_candidates: + if sidecar.exists(): + try: + with open(sidecar) as f: + sc = json.load(f) + tags_out = { + k: str(v) + for k, v in sc.items() + if v and isinstance(v, (str, int, float)) + } + if tags_out: + return tags_out + except Exception: + pass + + # 2. Plain .txt sidecar (SA3 convention) — same dir, then sibling txt/ dir. + txt_candidates = [ + fp.with_suffix(".txt"), + fp.parent.parent / "txt" / (stem + ".txt"), + ] + for txt in txt_candidates: + if txt.exists(): + try: + content = txt.read_text(encoding="utf-8", errors="replace").strip() + if content: + return {"prompt": content} + except Exception: + pass + + return {} + + +# --------------------------------------------------------------------------- +# Core encode +# --------------------------------------------------------------------------- + + +def encode_file( + encoder, + audio_np: np.ndarray, + *, + pad_modulo: int, + actual_samples: int | None = None, +): + """Encode one ``(2, T)`` float32 waveform → (latents, padding_mask, seconds). + + Returns: + latents_np: ``[D, T_latent]`` float32 (softnorm/DiT space, saved as-is) + padding_mask: latent-resolution 0/1 int list, length ``T_latent`` — + zeros over the alignment-padding tail + seconds_total: ``round(actual_samples / 44100, 3)`` + + ``chunked=True`` is used automatically for clips longer than 30 s. + """ + audio_np = np.asarray(audio_np) + if audio_np.ndim != 2 or audio_np.shape[0] != 2: + raise ValueError(f"Expected audio shaped (2, T), got {audio_np.shape}.") + if actual_samples is None: + actual_samples = int(audio_np.shape[-1]) + duration = actual_samples / SAMPLE_RATE + encoded = encode_audio( + encoder, + audio_np[None, ...], + valid_sample_lengths=[actual_samples], + pad_modulo=pad_modulo, + chunked=duration > CHUNKED_THRESHOLD_SECONDS, + ) + latents_np = np.asarray(encoded.latents.astype(mx.float32))[0] # [D, T] + padding_mask = [int(v) for v in np.asarray(encoded.padding_mask)[0]] + seconds_total = round(actual_samples / SAMPLE_RATE, 3) + return latents_np, padding_mask, seconds_total + + +def build_metadata( + fpath, + rel, + out_rel, + *, + actual_samples: int, + latent_shape, + padding_mask, + seconds_total, +) -> dict: + """Underfit-parity metadata dict for one encoded file (tags merged last).""" + meta = { + "path": str(fpath), + "relpath": str(out_rel), + "src_relpath": str(rel), + "seconds_total": seconds_total, + "seconds_start": 0, + "audio_samples": int(actual_samples), + "latent_shape": [int(s) for s in latent_shape], + "padding_mask": padding_mask, + } + meta.update(extract_tags(fpath)) + return meta + + +def run( + audio_dir, + output_dir, + codec: str, + encoder, + *, + max_duration: float = 600.0, + overwrite: bool = False, +) -> dict: + """Encode every audio file under ``audio_dir`` into ``output_dir``. + + Returns stats: ``{"encoded", "skipped", "errors", "count"}``. + """ + if codec not in SAME_ENCODER_PAD_MODULO: + raise ValueError(f"Unknown codec {codec!r}; expected same-s or same-l.") + pad_modulo = SAME_ENCODER_PAD_MODULO[codec] + + audio_dir = Path(audio_dir).expanduser().resolve() + output_dir = Path(output_dir).expanduser().resolve() + files = find_audio_files(audio_dir) + max_samples = max_samples_for_duration(max_duration) + + output_dir.mkdir(parents=True, exist_ok=True) + details = { + "codec": codec, + "sample_rate": SAMPLE_RATE, + "max_duration": max_duration, + "max_samples": max_samples, + "pad_modulo": pad_modulo, + "audio_dir": str(audio_dir), + "count": len(files), + } + with open(output_dir / "details.json", "w") as f: + json.dump(details, f, indent=2) + + total = len(files) + encoded = skipped = errors = 0 + t0 = time.time() + for index, fpath in enumerate(files, 1): + rel = fpath.relative_to(audio_dir) + out_rel = rel.with_suffix(".npy") + npy_path = output_dir / out_rel + json_path = npy_path.with_suffix(".json") + tag = f"[{index}/{total}]" + if npy_path.exists() and json_path.exists() and not overwrite: + skipped += 1 + print(f" {tag} SKIP {out_rel} (exists)") + continue + try: + audio = load_audio(fpath) + actual_samples = int(audio.shape[-1]) + if actual_samples > max_samples: + audio = audio[:, :max_samples] + actual_samples = max_samples + + latents_np, padding_mask, seconds_total = encode_file( + encoder, audio, pad_modulo=pad_modulo, actual_samples=actual_samples + ) + + npy_path.parent.mkdir(parents=True, exist_ok=True) + np.save(str(npy_path), latents_np) + meta = build_metadata( + fpath, + rel, + out_rel, + actual_samples=actual_samples, + latent_shape=latents_np.shape, + padding_mask=padding_mask, + seconds_total=seconds_total, + ) + with open(json_path, "w") as f: + json.dump(meta, f) + + encoded += 1 + shape_str = "x".join(str(s) for s in latents_np.shape) + print(f" {tag} {out_rel} {seconds_total:.1f}s -> [{shape_str}]") + except Exception as e: + errors += 1 + print(f" {tag} ERROR {rel}: {e}") + + elapsed = time.time() - t0 + mins, secs = divmod(int(elapsed), 60) + print( + f"\nDone in {mins}m{secs:02d}s: {encoded} encoded, {skipped} skipped, " + f"{errors} errors (of {total})" + ) + print(f"Output: {output_dir}") + return {"encoded": encoded, "skipped": skipped, "errors": errors, "count": total} + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def load_codec_encoder(codec: str): + """Load the real SAME encoder for ``codec`` in fp32. Returns (encoder, pad_modulo).""" + from sa3_mlx import load_encoder # lazy: pulls weight-download machinery + + return load_encoder(codec, mx.float32) + + +def main(): + parser = argparse.ArgumentParser( + description="Pre-encode audio into SAME latents for MLX LoRA finetuning", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--audio-dir", "-i", required=True, + help="Directory of audio files (searched recursively)", + ) + parser.add_argument( + "--output-dir", "-o", required=True, + help="Output root; each audio file becomes /.npy + .json", + ) + parser.add_argument( + "--codec", required=True, choices=sorted(SAME_ENCODER_PAD_MODULO), + help="Autoencoder: same-s pairs with sm-music/sm-sfx, same-l with medium", + ) + parser.add_argument( + "--max-duration", type=float, default=600.0, + help="Crop files longer than N seconds (default: 600; aligned down to 4096 samples)", + ) + parser.add_argument( + "--overwrite", action="store_true", + help="Re-encode files whose outputs already exist", + ) + args = parser.parse_args() + + audio_dir = Path(args.audio_dir).expanduser().resolve() + if not audio_dir.is_dir(): + print(f"Error: not a directory: {audio_dir}") + sys.exit(1) + files = find_audio_files(audio_dir) + if not files: + print(f"No audio files found in {audio_dir}") + sys.exit(1) + + max_samples = max_samples_for_duration(args.max_duration) + print(f"Codec: {args.codec} (pad_modulo={SAME_ENCODER_PAD_MODULO[args.codec]}, fp32)") + print(f"Input: {audio_dir} ({len(files)} audio files)") + print(f"Output: {Path(args.output_dir).expanduser().resolve()}") + print(f"Max: {max_samples / SAMPLE_RATE:.1f}s ({max_samples:,} samples)\n") + + print("Loading encoder...") + encoder, pad_modulo = load_codec_encoder(args.codec) + assert pad_modulo == SAME_ENCODER_PAD_MODULO[args.codec] + + run( + audio_dir, + args.output_dir, + args.codec, + encoder, + max_duration=args.max_duration, + overwrite=args.overwrite, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_mlx_latent_dataset.py b/tests/test_mlx_latent_dataset.py new file mode 100644 index 0000000..5830072 --- /dev/null +++ b/tests/test_mlx_latent_dataset.py @@ -0,0 +1,327 @@ +"""Tests for the torch-free pre-encoded latent dataset (underfit data pipeline port). + +No torch, no mlx — numpy + stdlib only, matching the module under test. +""" + +import json +import random + +import numpy as np +import pytest + +from optimized.mlx.models.defs.latent_dataset import ( + PreEncodedLatentDataset, + build_prompt, + iterate_batches, +) + +D = 8 # latent channels +CROP = 16 # latent_crop_length used throughout + +LONG_T = 64 # stored length of the long item (> CROP) +LONG_VALID = 50 # 1s in its padding mask (trailing zeros after) +LONG_SECONDS = 120.5 + +SHORT_T = 6 # stored length of the short item (< CROP) +SHORT_SECONDS = 2.786 + +SILENCE_VALUE = 7.5 + + +def _write_item(root, stem, T, seconds_total, tags, valid=None): + """Write an underfit pre-encode style .npy + .json pair. + + Latents are [D, T] with latents[d, t] == t so a crop's start offset can be + read straight off the values. + """ + valid = T if valid is None else valid + latents = np.tile(np.arange(T, dtype=np.float32), (D, 1)) + np.save(str(root / f"{stem}.npy"), latents) + meta = { + "relpath": f"{stem}.npy", + "seconds_total": seconds_total, + "seconds_start": 0, + "audio_samples": int(seconds_total * 44100), + "latent_shape": [D, T], + "padding_mask": [1] * valid + [0] * (T - valid), + } + meta.update(tags) + (root / f"{stem}.json").write_text(json.dumps(meta)) + + +def _make_root(tmp_path, with_silence): + root = tmp_path / ("latents" if with_silence else "latents_nosilence") + root.mkdir() + _write_item( + root, "long", LONG_T, LONG_SECONDS, + {"title": "Neon Skyline", "artist": "The Testers", + "genre": "Synthwave", "bpm": "120"}, + valid=LONG_VALID, + ) + # "prompt" tag = what pre_encode extracts from a .txt sidecar caption + _write_item(root, "short", SHORT_T, SHORT_SECONDS, + {"prompt": "a lofi hip hop beat"}) + if with_silence: + # Reference stores silence as [1, C, N] (it squeezes axis 0 on load) + silence = np.full((1, D, 4), SILENCE_VALUE, dtype=np.float32) + np.save(str(root / "silence.npy"), silence) + return root + + +@pytest.fixture +def data_root(tmp_path): + return _make_root(tmp_path, with_silence=True) + + +@pytest.fixture +def nosilence_root(tmp_path): + return _make_root(tmp_path, with_silence=False) + + +def _get_by_relpath(ds, relpath): + for i in range(len(ds)): + item = ds[i] + if item["relpath"] == relpath: + return item + raise AssertionError(f"{relpath} not found in dataset") + + +# --------------------------------------------------------------------------- +# Cropping / padding / masks +# --------------------------------------------------------------------------- + + +def test_crop_yields_exact_length_and_mask(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, random_crop=False) + assert len(ds) == 2 # silence.npy is not a sample + + item = _get_by_relpath(ds, "long.npy") + assert item["latents"].shape == (D, CROP) + assert item["latents"].dtype == np.float32 + assert item["padding_mask"].shape == (CROP,) + assert item["padding_mask"].dtype == np.bool_ + # random_crop=False → start at 0; first CROP frames are all valid + np.testing.assert_array_equal(item["latents"][0], np.arange(CROP, dtype=np.float32)) + assert item["padding_mask"].all() + + +def test_random_crop_false_starts_at_zero(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, random_crop=False) + for _ in range(20): + item = _get_by_relpath(ds, "long.npy") + assert item["latents"][0, 0] == 0.0 + + +def test_random_crop_stays_in_valid_region(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, random_crop=True, seed=1234) + last_ix = LONG_VALID - 1 # last 1 in the stored padding mask + starts = set() + for _ in range(200): + item = _get_by_relpath(ds, "long.npy") + start = int(item["latents"][0, 0]) + # reference: start = randint(0, last_ix - crop), inclusive + assert 0 <= start <= last_ix - CROP + # every crop is contiguous and the mask matches the stored one + np.testing.assert_array_equal( + item["latents"][0], np.arange(start, start + CROP, dtype=np.float32) + ) + np.testing.assert_array_equal( + item["padding_mask"], (start + np.arange(CROP)) < LONG_VALID + ) + starts.add(start) + assert len(starts) > 1, "random_crop=True should vary the crop start" + + +def test_short_item_pads_with_silence(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, random_crop=False) + item = _get_by_relpath(ds, "short.npy") + + assert item["latents"].shape == (D, CROP) + # stored frames untouched + np.testing.assert_array_equal( + item["latents"][:, :SHORT_T], + np.tile(np.arange(SHORT_T, dtype=np.float32), (D, 1)), + ) + # padded frames come from the tiled/sliced silence latent + np.testing.assert_array_equal( + item["latents"][:, SHORT_T:], + np.full((D, CROP - SHORT_T), SILENCE_VALUE, dtype=np.float32), + ) + # mask: valid frames then zeros for the padding + expected_mask = np.array([True] * SHORT_T + [False] * (CROP - SHORT_T)) + np.testing.assert_array_equal(item["padding_mask"], expected_mask) + + +def test_short_item_zero_pads_without_silence(nosilence_root): + ds = PreEncodedLatentDataset(nosilence_root, CROP, random_crop=False) + item = _get_by_relpath(ds, "short.npy") + + np.testing.assert_array_equal( + item["latents"][:, SHORT_T:], np.zeros((D, CROP - SHORT_T), dtype=np.float32) + ) + expected_mask = np.array([True] * SHORT_T + [False] * (CROP - SHORT_T)) + np.testing.assert_array_equal(item["padding_mask"], expected_mask) + + +def test_seconds_total_unchanged_by_crop_and_pad(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, random_crop=True) + # deliberate underfit convention: FULL stored duration, never recomputed + assert _get_by_relpath(ds, "long.npy")["seconds_total"] == LONG_SECONDS + assert _get_by_relpath(ds, "short.npy")["seconds_total"] == SHORT_SECONDS + + +# --------------------------------------------------------------------------- +# Length filters (rejection redraw) +# --------------------------------------------------------------------------- + + +def test_min_length_sec_rejection_redraws(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, min_length_sec=50.0) + # the short item (2.786 s) is always rejected → every index resolves to long + for i in range(len(ds)): + for _ in range(5): + item = ds[i] + assert item["relpath"] == "long.npy" + assert item["seconds_total"] == LONG_SECONDS + + +def test_all_items_rejected_raises(data_root): + ds = PreEncodedLatentDataset(data_root, CROP, min_length_sec=1e6) + with pytest.raises(RuntimeError): + ds[0] + + +# --------------------------------------------------------------------------- +# Batching +# --------------------------------------------------------------------------- + + +def test_oversample_epoch_yields_batch_size_x100_samples(data_root): + ds = PreEncodedLatentDataset(data_root, CROP) + batch_size = 4 # > len(ds) == 2 → oversample with replacement + batches = list(iterate_batches(ds, batch_size, seed=0)) + + assert len(batches) == 100 + total = sum(b["latents"].shape[0] for b in batches) + assert total == batch_size * 100 + + for b in batches: + assert b["latents"].shape == (batch_size, D, CROP) + assert b["latents"].dtype == np.float32 + assert b["padding_mask"].shape == (batch_size, CROP) + assert b["padding_mask"].dtype == np.bool_ + assert len(b["prompt"]) == batch_size + assert len(b["seconds_total"]) == batch_size + assert len(b["relpath"]) == batch_size + + # with replacement over 400 draws, both items must appear + seen = {rp for b in batches for rp in b["relpath"]} + assert seen == {"long.npy", "short.npy"} + + +def test_regular_epoch_drop_last_false(data_root): + ds = PreEncodedLatentDataset(data_root, CROP) + # len(ds) == 2 >= batch_size → plain epoch, no oversampling + batches = list(iterate_batches(ds, 2, seed=0)) + assert len(batches) == 1 + assert sorted(batches[0]["relpath"]) == ["long.npy", "short.npy"] + + # odd split keeps the remainder (drop_last=False) + batches = list(iterate_batches(ds, 1, seed=0)) + assert [b["latents"].shape[0] for b in batches] == [1, 1] + + +# --------------------------------------------------------------------------- +# Prompt building +# --------------------------------------------------------------------------- + + +def test_trigger_probability_approx_80pct(): + meta = {"title": "Song"} + pc = {"trigger": "zkq"} # trigger_pct defaults to 80 + rng = random.Random(123) + + prompts = [build_prompt(meta, pc, rng) for _ in range(2000)] + with_trigger = sum(p.startswith("zkq, ") for p in prompts) + frac = with_trigger / len(prompts) + assert 0.77 < frac < 0.83, f"trigger fraction {frac} not ≈ 0.8" + + # tag prompt itself is always present (single tag → augmentation is a no-op) + assert all(p.endswith("Title: Song") for p in prompts) + + +def test_shuffle_vs_subset_50_50(): + meta = {"title": "Neon", "artist": "Testers", "genre": "Synthwave", "bpm": "120"} + expected_parts = {"Title: Neon", "Artist: Testers", "Genre: Synthwave", "BPM: 120"} + pc = {} # use_tags default True, shuffle default True, no trigger + rng = random.Random(7) + + n_draws = 2000 + n_full = 0 + full_orderings = set() + for _ in range(n_draws): + prompt = build_prompt(meta, pc, rng) + parts = prompt.split(", ") + assert set(parts) <= expected_parts + assert len(parts) == len(set(parts)) + if len(parts) == 4: + n_full += 1 + full_orderings.add(prompt) + + # subset branch (p=0.5) draws size k ~ U{1..4}; k<4 w.p. 3/4 → partial ≈ 0.375 + partial_frac = 1 - n_full / n_draws + assert 0.32 < partial_frac < 0.43, f"partial fraction {partial_frac} not ≈ 0.375" + # shuffle-all branch occurred and actually shuffles + assert n_full > 0 + assert len(full_orderings) > 1 + + +def test_legacy_prompt_without_prompt_config(): + meta = {"artist": "The Testers", "title": "Neon Skyline"} + rng = random.Random(3) + lengths = set() + for _ in range(300): + prompt = build_prompt(meta, None, rng) + parts = prompt.split(", ") + assert set(parts) <= {"Artist: The Testers", "Title: Neon Skyline"} + assert len(parts) >= 1 + lengths.add(len(parts)) + assert lengths == {1, 2} # subset branch produces 1-tag prompts too + + # no tags at all → "text" fallback + assert build_prompt({"text": "raw caption"}, None, rng) == "raw caption" + + +def test_legacy_prompt_via_dataset(data_root): + ds = PreEncodedLatentDataset(data_root, CROP) # prompt_config=None → legacy + expected = {"Artist: The Testers", "Title: Neon Skyline", + "BPM: 120", "Genre: Synthwave"} + for _ in range(20): + prompt = _get_by_relpath(ds, "long.npy")["prompt"] + assert prompt + assert set(prompt.split(", ")) <= expected + + +def test_txt_derived_prompt_tag(data_root): + # the short item's caption came from a .txt sidecar → "prompt" tag key + ds = PreEncodedLatentDataset(data_root, CROP, prompt_config={"shuffle": False}) + assert _get_by_relpath(ds, "short.npy")["prompt"] == "Prompt: a lofi hip hop beat" + + ds = PreEncodedLatentDataset( + data_root, CROP, + prompt_config={"shuffle": False, "hide_tag_names": True}, + ) + assert _get_by_relpath(ds, "short.npy")["prompt"] == "a lofi hip hop beat" + + +def test_path_prompt_and_space_joined_trigger(): + meta = {"relpath": "artistX/track01.npy"} + rng = random.Random(11) + + pc = {"use_tags": False, "use_paths": True, "path_opts": {"hideExt": True}} + assert build_prompt(meta, pc, rng) == "artistX/track01" + + # non-tag method → trigger joined with a space, not ", " + pc = {"use_tags": False, "use_paths": True, + "path_opts": {"hideExt": True}, "trigger": "zkq", "trigger_pct": 100} + assert build_prompt(meta, pc, rng) == "zkq artistX/track01" diff --git a/tests/test_mlx_pre_encode.py b/tests/test_mlx_pre_encode.py new file mode 100644 index 0000000..1e0a07d --- /dev/null +++ b/tests/test_mlx_pre_encode.py @@ -0,0 +1,313 @@ +"""Unit tests for optimized/mlx/scripts/pre_encode_mlx.py (stub encoder, no weights).""" + +import json +import sys +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("mlx.core") +soundfile = pytest.importorskip("soundfile") + +import mlx.core as mx + +SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "optimized" / "mlx" / "scripts" +if str(SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPTS_DIR)) + +import pre_encode_mlx as pe + +SR = 44100 + + +class StubEncoder: + """[B, 512, P] patches → [B, 8, P/16]: 16-stride mean pool, first 8 channels. + + Stride-local like the real encoder's latent rate, so chunked stitching is + bit-exact against unchunked — no model download needed. + """ + + def __call__(self, patches): + batch, channels, patch_count = patches.shape + pooled = mx.mean( + patches.reshape(batch, channels, patch_count // 16, 16), axis=-1 + ) + return pooled[:, :8, :] + + +@pytest.fixture(scope="module") +def dataset(tmp_path_factory): + """Audio tree: sub/three_sec.wav (stereo + JSON sidecar), long45.wav (mono + .txt).""" + root = tmp_path_factory.mktemp("audio") + + t3 = np.arange(3 * SR) / SR + stereo = np.stack( + [ + 0.4 * np.sin(2 * np.pi * 440 * t3) + 0.1, + 0.4 * np.sin(2 * np.pi * 330 * t3) + 0.1, + ] + ).astype(np.float32) + (root / "sub").mkdir() + soundfile.write(root / "sub" / "three_sec.wav", stereo.T, SR) + (root / "sub" / "three_sec.json").write_text( + json.dumps( + {"artist": "Test Artist", "bpm": 128, "nested": {"x": 1}, "empty": ""} + ) + ) + + t45 = np.arange(45 * SR) / SR + mono = (0.3 * np.sin(2 * np.pi * 220 * t45) + 0.05).astype(np.float32) + soundfile.write(root / "long45.wav", mono, SR) + (root / "long45.txt").write_text("a funky bassline, live drums\n") + + return root + + +@pytest.fixture() +def encoded(dataset, tmp_path): + out = tmp_path / "latents" + stats = pe.run(dataset, out, "same-s", StubEncoder()) + return out, stats + + +# --------------------------------------------------------------------------- +# End-to-end run: outputs, relpath preservation, details.json +# --------------------------------------------------------------------------- + + +def test_run_outputs_preserve_relpaths(encoded): + out, stats = encoded + assert stats == {"encoded": 2, "skipped": 0, "errors": 0, "count": 2} + for rel in ("sub/three_sec.npy", "sub/three_sec.json", "long45.npy", "long45.json"): + assert (out / rel).exists(), rel + + details = json.loads((out / "details.json").read_text()) + assert details["codec"] == "same-s" + assert details["sample_rate"] == SR + assert details["max_duration"] == 600.0 + assert details["count"] == 2 + + +def test_three_sec_latents_and_metadata(encoded, dataset): + out, _ = encoded + latents = np.load(out / "sub" / "three_sec.npy") + # 3 s = 132300 samples → padded to 17*8192 = 139264 (same-s alignment) → 34 latents + assert latents.shape == (8, 34) + assert latents.dtype == np.float32 + + meta = json.loads((out / "sub" / "three_sec.json").read_text()) + assert meta["path"] == str(dataset / "sub" / "three_sec.wav") + assert meta["relpath"] == "sub/three_sec.npy" + assert meta["src_relpath"] == "sub/three_sec.wav" + assert meta["seconds_total"] == 3.0 + assert meta["seconds_start"] == 0 + assert meta["audio_samples"] == 3 * SR + assert meta["latent_shape"] == [8, 34] + # ceil(132300 / 4096) = 33 valid latents; final latent is alignment padding + mask = meta["padding_mask"] + assert len(mask) == 34 + assert mask[:33] == [1] * 33 + assert mask[33] == 0 + # JSON sidecar tags merged: string/number values kept, nested/empty dropped + assert meta["artist"] == "Test Artist" + assert meta["bpm"] == "128" + assert "nested" not in meta + assert "empty" not in meta + + +def test_long45_chunked_mono_and_txt_prompt(encoded, dataset): + out, _ = encoded + latents = np.load(out / "long45.npy") + # 45 s = 1984500 samples → padded to 243*8192 = 1990656 → 486 latents + assert latents.shape == (8, 486) + + meta = json.loads((out / "long45.json").read_text()) + assert meta["audio_samples"] == 45 * SR + assert meta["seconds_total"] == 45.0 + mask = meta["padding_mask"] + assert len(mask) == 486 + assert sum(mask) == 485 # ceil(1984500/4096) = 485 valid + assert mask[-1] == 0 + assert meta["prompt"] == "a funky bassline, live drums" + + # >30 s took the chunked path; stub is stride-local → must match unchunked + audio = pe.load_audio(dataset / "long45.wav") + ref = pe.encode_audio( + StubEncoder(), audio[None, ...], pad_modulo=32, chunked=False + ) + np.testing.assert_allclose(latents, np.asarray(ref.latents)[0], atol=1e-6) + assert np.any(latents != 0.0) + + +# --------------------------------------------------------------------------- +# Skip vs overwrite +# --------------------------------------------------------------------------- + + +def test_skip_then_overwrite(dataset, tmp_path): + out = tmp_path / "latents" + stub = StubEncoder() + + first = pe.run(dataset, out, "same-s", stub) + assert first["encoded"] == 2 and first["skipped"] == 0 + + second = pe.run(dataset, out, "same-s", stub) + assert second["encoded"] == 0 and second["skipped"] == 2 + + # Corrupt an output, then --overwrite must re-encode it + np.save(str(out / "long45.npy"), np.zeros((8, 486), dtype=np.float32)) + third = pe.run(dataset, out, "same-s", stub, overwrite=True) + assert third["encoded"] == 2 and third["skipped"] == 0 + assert np.any(np.load(out / "long45.npy") != 0.0) + + +# --------------------------------------------------------------------------- +# Max-duration cap aligned down to 4096 +# --------------------------------------------------------------------------- + + +def test_max_duration_aligns_down_to_4096(dataset, tmp_path): + out = tmp_path / "latents" + stats = pe.run(dataset, out, "same-s", StubEncoder(), max_duration=1.0) + assert stats["encoded"] == 2 + + # 1 s = 44100 samples → aligned down to 40960 = 10 * 4096 + assert pe.max_samples_for_duration(1.0) == 40960 + for name in ("sub/three_sec", "long45"): + meta = json.loads((out / f"{name}.json").read_text()) + assert meta["audio_samples"] == 40960 + assert meta["seconds_total"] == round(40960 / SR, 3) == 0.929 + latents = np.load(out / f"{name}.npy") + assert latents.shape == (8, 10) # 40960 is already 8192-aligned, no padding + assert meta["padding_mask"] == [1] * 10 + + +# --------------------------------------------------------------------------- +# encode_file core +# --------------------------------------------------------------------------- + + +def test_encode_file_seconds_rounding_and_mask(): + rng = np.random.default_rng(7) + audio = rng.standard_normal((2, 130000)).astype(np.float32) + + latents, mask, seconds_total = pe.encode_file( + StubEncoder(), audio, pad_modulo=32 + ) + assert seconds_total == round(130000 / SR, 3) == 2.948 + # 130000 → padded to 16*8192 = 131072 → 32 latents; ceil(130000/4096) = 32 valid + assert latents.shape == (8, 32) + assert latents.dtype == np.float32 + assert mask == [1] * 32 + + # Cropping semantics: caller-declared valid length shortens the mask + latents2, mask2, seconds2 = pe.encode_file( + StubEncoder(), audio, pad_modulo=32, actual_samples=100000 + ) + assert seconds2 == round(100000 / SR, 3) + assert len(mask2) == latents2.shape[-1] + assert sum(mask2) == 25 # ceil(100000/4096) + + +def test_encode_file_rejects_non_stereo(): + with pytest.raises(ValueError, match=r"\(2, T\)"): + pe.encode_file(StubEncoder(), np.zeros((1, 8192), dtype=np.float32), pad_modulo=32) + + +# --------------------------------------------------------------------------- +# Audio loading: mono→stereo, >2ch, resample +# --------------------------------------------------------------------------- + + +def test_load_audio_mono_duplicates_to_stereo(dataset): + audio = pe.load_audio(dataset / "long45.wav") + assert audio.shape == (2, 45 * SR) + assert audio.dtype == np.float32 + np.testing.assert_array_equal(audio[0], audio[1]) + + +def test_load_audio_takes_first_two_of_multichannel(tmp_path): + data = np.stack( + [np.full(SR, 0.1), np.full(SR, 0.2), np.full(SR, 0.3)], axis=1 + ).astype(np.float32) + path = tmp_path / "three_ch.wav" + soundfile.write(path, data, SR) + + audio = pe.load_audio(path) + assert audio.shape == (2, SR) + np.testing.assert_allclose(audio[0], 0.1, atol=1e-3) + np.testing.assert_allclose(audio[1], 0.2, atol=1e-3) + + +def test_load_audio_resamples_to_44100(tmp_path): + sr_in = 22050 + t = np.arange(sr_in) / sr_in # 1 s + mono = (0.3 * np.sin(2 * np.pi * 110 * t)).astype(np.float32) + path = tmp_path / "half_rate.wav" + soundfile.write(path, mono, sr_in) + + audio = pe.load_audio(path) + assert audio.shape == (2, SR) + # Still a ~110 Hz sine after resampling: correlate with the ideal signal + ideal = 0.3 * np.sin(2 * np.pi * 110 * np.arange(SR) / SR) + corr = np.dot(audio[0], ideal) / ( + np.linalg.norm(audio[0]) * np.linalg.norm(ideal) + ) + assert corr > 0.99 + + +# --------------------------------------------------------------------------- +# Tag extraction +# --------------------------------------------------------------------------- + + +def test_extract_tags_json_beats_txt(tmp_path): + wav = tmp_path / "clip.wav" + (tmp_path / "clip.json").write_text(json.dumps({"title": "A", "year": 2020})) + (tmp_path / "clip.txt").write_text("prompt text") + assert pe.extract_tags(wav) == {"title": "A", "year": "2020"} + + +def test_extract_tags_txt_becomes_prompt(tmp_path): + wav = tmp_path / "clip.wav" + (tmp_path / "clip.txt").write_text(" spaced out prompt \n") + assert pe.extract_tags(wav) == {"prompt": "spaced out prompt"} + + +def test_extract_tags_sibling_dirs(tmp_path): + audio_dir = tmp_path / "audio" + audio_dir.mkdir() + wav = audio_dir / "clip.wav" + + (tmp_path / "json").mkdir() + (tmp_path / "json" / "clip.json").write_text(json.dumps({"genre": "funk"})) + assert pe.extract_tags(wav) == {"genre": "funk"} + + (tmp_path / "json" / "clip.json").unlink() + (tmp_path / "txt").mkdir() + (tmp_path / "txt" / "clip.txt").write_text("sibling prompt") + assert pe.extract_tags(wav) == {"prompt": "sibling prompt"} + + +def test_extract_tags_no_sidecar_returns_empty(tmp_path): + assert pe.extract_tags(tmp_path / "clip.wav") == {} + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + + +def test_find_audio_files_filters_and_sorts(tmp_path): + rng = np.random.default_rng(3) + big = (0.1 * rng.standard_normal(SR)).astype(np.float32) # noise: flac stays >4 KB + soundfile.write(tmp_path / "b.wav", big, SR) + (tmp_path / "deep").mkdir() + soundfile.write(tmp_path / "deep" / "a.flac", big, SR) + soundfile.write(tmp_path / "._resource.wav", big, SR) # macOS fork prefix + (tmp_path / "tiny.wav").write_bytes(b"RIFF") # below min size + (tmp_path / "notes.txt").write_text("not audio") + + found = pe.find_audio_files(tmp_path) + assert found == [tmp_path / "b.wav", tmp_path / "deep" / "a.flac"] From 3d4ca37810fb6374a6dce90bb17e4f69f663f70b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 13:46:56 -0400 Subject: [PATCH 08/28] =?UTF-8?q?mlx=20trainer:=20lora=5Ftrain=5Fmlx.py=20?= =?UTF-8?q?=E2=80=94=20underfit=20loop=20conventions=20end=20to=20end?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step-driven LoRA training on Apple Silicon mirroring underfit's raw loop: rectified-flow velocity target with signal-only masked MSE (per-sample then mean), uniform timestep sampler + the SA3 'full' distribution shift (256/4096), CFG dropout 0.1 (whole cross_attn to zeros per sample inside the grad scope, global_cond kept), AdamW at torch defaults with required --lr (weight decay explicitly 0.0 — MLX's default differs), fp32 adapters over a frozen fp16 base, checkpoints every 1000 steps + final ({run}-step={s}-epoch={e}.safetensors with step/epoch/base_model metadata), resume ladder (flags > metadata > filename), per-step loss_by_timestep.bin telemetry, ///checkpoints layout. E2E verified on sm-music: pre-encode → 20 steps dora-rows r16 (140 DiT layers = the underfit per-block set + conditioner, 9.2M trainable, 1.6 it/s at T=256) → checkpoint passes underfit's lora_validate contract → resume restores 141 layers with metadata offsets → checkpoint loads directly via sa3_mlx.py --lora (141 layers merged). --- optimized/mlx/.gitignore | 3 + optimized/mlx/scripts/lora_train_mlx.py | 367 ++++++++++++++++++++++++ 2 files changed, 370 insertions(+) create mode 100644 optimized/mlx/scripts/lora_train_mlx.py diff --git a/optimized/mlx/.gitignore b/optimized/mlx/.gitignore index 1568428..5bd9379 100644 --- a/optimized/mlx/.gitignore +++ b/optimized/mlx/.gitignore @@ -14,3 +14,6 @@ output/gradio/ # Local LoRA library scanned by the gradio UI (user adapters, never committed) loras/ + +# Trainer telemetry (written to cwd, underfit convention) +loss_by_timestep.bin diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py new file mode 100644 index 0000000..0b43cc0 --- /dev/null +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -0,0 +1,367 @@ +"""SA3 LoRA training on Apple Silicon — pure MLX, underfit conventions. + +The MLX counterpart of underfit's raw-PyTorch training loop +(github.com/dada-bots/underfit, underfit/training/loop.py), built on the +trainable adapters in models/defs/lora.py. Every training convention and +default mirrors underfit (see ../TRAINING_CONVENTIONS.md): + + - rectified-flow velocity target (noise − x), signal-only masked MSE, + per-sample-then-mean reduction + - uniform timestep sampler + the SA3 models' "full" distribution shift + (min_length 256, max_length 4096) + - CFG dropout 0.1: per-sample, the whole cross-attention conditioning + (prompt + seconds token) is zeroed; global_cond is kept + - AdamW, single param group, torch defaults (betas 0.9/0.999, eps 1e-8, + weight decay 0.0), no schedule, no clipping; LR has NO default and must + be supplied (--lr) + - adapters train in fp32 over a frozen fp16 base + - step-driven loop; checkpoints every --checkpoint-every (1000) steps and + at the end, saved BEFORE anything else can fail, named + {run_label}-step={step}-epoch={epoch}.safetensors with lora_config + metadata {rank, alpha, adapter_type, include, exclude, step, epoch, + base_model} + - resume offsets: explicit flags > checkpoint metadata > filename tokens + - per-step loss_by_timestep.bin telemetry (struct "Iff": step, t_mean, + loss_mean), flushed every 10 steps + +Dataset: pre-encoded latents from scripts/pre_encode_mlx.py (npy + json +pairs), cropped to --latent-crop-length with random crop, tiny datasets +oversampled with replacement (batch_size*100 draws/epoch) — the underfit +overfitting workflow. Prompts are built per-sample with underfit's tag +augmentation. seconds_total stays the FULL source duration after cropping +(deliberate underfit convention). + +Usage: + python scripts/lora_train_mlx.py --dit sm-music \ + --latents-dir output/latents/my-set --lr 1e-4 \ + --name my-lora --save-dir output/runs --max-steps 2000 +""" +from __future__ import annotations + +import argparse +import json +import math +import re +import struct +import sys +import time +import uuid +from pathlib import Path + +import numpy as np + +SCRIPTS_DIR = Path(__file__).resolve().parent +REPO = SCRIPTS_DIR.parent +sys.path.insert(0, str(REPO)) +sys.path.insert(0, str(SCRIPTS_DIR)) + +import mlx.core as mx # noqa: E402 +import mlx.nn as nn # noqa: E402 +import mlx.optimizers as optim # noqa: E402 + +from sa3_mlx import DIT_CHOICES, load_dit # noqa: E402 +from weights import ensure_local # noqa: E402 +from sa3_mlx import T5GEMMA_NPZ_REL # noqa: E402 +from models.defs.sa3_pipeline import ( # noqa: E402 + apply_prompt_padding, load_conditioner_from_npz, +) +from models.defs.t5gemma_mlx import T5Gemma # noqa: E402 +from models.defs.training import ( # noqa: E402 + rectified_flow_loss, sample_training_timesteps, shift_training_timesteps, +) +from models.defs.lora import ( # noqa: E402 + TrainableSecondsEmbedder, inject_from_lora_config, load_lora_checkpoint, + load_trainable_lora_state, save_lora_checkpoint, underfit_lora_config, +) +from models.defs.latent_dataset import ( # noqa: E402 + PreEncodedLatentDataset, iterate_batches, +) + +# underfit registry conventions per model +BASE_MODEL_NAMES = {"sm-music": "sa3-sm-music", "sm-sfx": "sa3-sm-sfx", + "medium": "sa3-medium"} +LATENT_CROP_DEFAULTS = {"sm-music": 1300, "sm-sfx": 1300, "medium": 4096} +DIST_SHIFT_DEFAULT = {"shift_type": "full", + "options": {"min_length": 256, "max_length": 4096}} + + +class TrainBundle(nn.Module): + """DiT + (optional) trainable seconds conditioner under one module root so + a single value_and_grad covers every adapter parameter. The bundle's + trainable_parameters() are the LoRA params only (bases frozen at inject).""" + + def __init__(self, dit, secs): + super().__init__() + self.dit = dit + self.secs = secs # TrainableSecondsEmbedder or None + + +def parse_args(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--dit", required=True, choices=list(DIT_CHOICES.keys())) + ap.add_argument("--latents-dir", required=True, + help="Root of pre-encoded npy+json pairs (scripts/pre_encode_mlx.py)") + ap.add_argument("--lr", type=float, required=True, + help="Learning rate (underfit has NO default — must be supplied)") + ap.add_argument("--name", default="underfit-run", help="Run name (defaults.ini)") + ap.add_argument("--save-dir", default=str(REPO / "output" / "runs")) + ap.add_argument("--batch-size", type=int, default=1) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--max-steps", type=int, default=10_000_000_000, + help="Absolute global-step target (resume-aware)") + ap.add_argument("--checkpoint-every", type=int, default=1000) + # adapter config — underfit dashboard defaults + ap.add_argument("--adapter-type", default="dora-rows") + ap.add_argument("--rank", type=int, default=16) + ap.add_argument("--alpha", type=float, default=None, help="Default: = rank") + ap.add_argument("--include", default=None, + help="Comma-separated substring filters (underfit convention)") + ap.add_argument("--exclude", default=None, + help="Comma-separated; default = underfit's dashboard exclude list") + ap.add_argument("--no-conditioner-lora", action="store_true", + help="Skip the seconds-conditioner adapter (underfit adapts it " + "by default — e.g. the plini checkpoint carries its delta)") + # data + ap.add_argument("--latent-crop-length", type=int, default=None, + help="Default per model: 1300 (sm-*) / 4096 (medium)") + ap.add_argument("--random-crop", action=argparse.BooleanOptionalAction, default=True) + ap.add_argument("--prompt-config", default=None, + help="JSON file with underfit prompt_config (trigger, tag options…)") + # schedule / conditioning + ap.add_argument("--timestep-sampler", default="uniform", + choices=["uniform", "logit_normal", "trunc_logit_normal", + "log_snr", "log_snr_uniform"]) + ap.add_argument("--dist-shift", default="full", + choices=["none", "full", "flux", "logsnr"]) + ap.add_argument("--cfg-dropout-prob", type=float, default=0.1) + # resume + ap.add_argument("--lora-ckpt-path", default=None) + ap.add_argument("--step-offset", type=int, default=None) + ap.add_argument("--epoch-offset", type=int, default=None) + return ap.parse_args() + + +def resolve_offsets(args, ckpt_config) -> tuple[int, int]: + """Underfit's resume ladder: explicit flags > checkpoint metadata > + filename step=/epoch= tokens > zero.""" + if args.step_offset is not None or args.epoch_offset is not None: + return int(args.step_offset or 0), int(args.epoch_offset or 0) + if ckpt_config: + step = ckpt_config.get("step") + epoch = ckpt_config.get("epoch") + if step is not None: + return int(step), int(epoch or 0) + if args.lora_ckpt_path: + name = Path(args.lora_ckpt_path).name + ms = re.search(r"step=(\d+)", name) + me = re.search(r"epoch=(\d+)", name) + if ms: + return int(ms.group(1)), int(me.group(1)) if me else 0 + return 0, 0 + + +def build_conditioning(t5, padding_emb, prompts, max_len=256): + """Frozen prompt-side conditioning: T5Gemma embeddings with the learned + padding embedding applied. Returns [B, 256, 768] fp32 (the trainable + seconds token is concatenated later, inside the grad scope).""" + embeds, mask = t5.encode(list(prompts), max_len=max_len) + mx.eval(embeds, mask) + padded = apply_prompt_padding(embeds.astype(mx.float32), mask, + padding_emb.astype(mx.float32)) + return mx.stop_gradient(padded) + + +def main(): + args = parse_args() + np_rng = np.random.default_rng(args.seed) + mx.random.seed(args.seed) + + alpha = float(args.alpha) if args.alpha is not None else float(args.rank) + lora_config = underfit_lora_config() + lora_config.update({"adapter_type": args.adapter_type, "rank": args.rank, + "alpha": alpha}) + if args.include is not None: + lora_config["include"] = [s.strip() for s in args.include.split(",") if s.strip()] + if args.exclude is not None: + lora_config["exclude"] = [s.strip() for s in args.exclude.split(",") if s.strip()] + + crop_len = args.latent_crop_length or LATENT_CROP_DEFAULTS[args.dit] + base_model = BASE_MODEL_NAMES[args.dit] + + # ── run dirs (underfit layout: ///checkpoints) ──────── + run_label = re.sub(r"-\d{14}$", "", args.name) if args.name else None + session = uuid.uuid4().hex[:8] + ckpt_dir = Path(args.save_dir) / args.name / session / "checkpoints" + ckpt_dir.mkdir(parents=True, exist_ok=True) + print(f"run: {args.name} session {session}") + print(f"checkpoints → {ckpt_dir}") + + # ── dataset ────────────────────────────────────────────────────────────── + prompt_config = None + if args.prompt_config: + prompt_config = json.loads(Path(args.prompt_config).read_text()) + dataset = PreEncodedLatentDataset( + args.latents_dir, crop_len, random_crop=args.random_crop, + prompt_config=prompt_config, seed=args.seed) + print(f"dataset: {len(dataset)} latent file(s), crop {crop_len} latents" + + (" (oversampling with replacement — tiny dataset)" + if len(dataset) < args.batch_size else "")) + + # ── models ─────────────────────────────────────────────────────────────── + t0 = time.time() + dit_model, _ = load_dit(args.dit, T_lat=crop_len, dtype=mx.float16) + print(f"DiT loaded ({time.time()-t0:.1f}s, fp16 base, T_lat={crop_len})") + t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL))) # frozen + padding_emb, secs_embedder = load_conditioner_from_npz( + str(ensure_local(DIT_CHOICES[args.dit]["ckpt"])), prefix="cond.") + + # ── inject adapters ───────────────────────────────────────────────────── + report, saved_config = inject_from_lora_config( + dit_model, lora_config, checkpoint_prefix="model.") + print(f"lora: {report.layer_count} DiT layer(s), {report.adapter_type}, " + f"rank {args.rank}, alpha {alpha:g} " + f"({report.trainable_parameters/1e6:.2f}M trainable)") + + secs_module = None + if not args.no_conditioner_lora: + secs_module = TrainableSecondsEmbedder(secs_embedder.W, secs_embedder.b) + try: + cond_report, _ = inject_from_lora_config( + secs_module, lora_config, + checkpoint_prefix="conditioners.seconds_total.") + print(f"lora: conditioner seconds embedder adapted " + f"({cond_report.layer_count} layer)") + except ValueError: + secs_module = None # filtered out by include/exclude + bundle = TrainBundle(dit_model, secs_module) + + # ── resume ─────────────────────────────────────────────────────────────── + step_offset = epoch_offset = 0 + if args.lora_ckpt_path: + sd, ckpt_cfg = load_lora_checkpoint(args.lora_ckpt_path) + restored = load_trainable_lora_state(bundle, sd) + step_offset, epoch_offset = resolve_offsets(args, ckpt_cfg) + print(f"resumed {restored} adapter layer(s) from " + f"{Path(args.lora_ckpt_path).name} (step {step_offset}, " + f"epoch {epoch_offset})") + else: + step_offset, epoch_offset = resolve_offsets(args, None) + + max_new_steps = max(0, args.max_steps - step_offset) + + # ── optimizer: AdamW, torch defaults, single group, no schedule ────────── + optimizer = optim.AdamW(learning_rate=args.lr, betas=[0.9, 0.999], + eps=1e-8, weight_decay=0.0) + + # ── loss (adapters in the bundle are the only trainable params) ────────── + def loss_fn(bundle, latents, timesteps, loss_mask, prompt_cond, + seconds_totals, drop_mask): + if bundle.secs is not None: + sec_tok = bundle.secs(seconds_totals) # [B, 1, 768] fp32 + else: + sec_tok = mx.stop_gradient( + secs_embedder(list(np.asarray(seconds_totals)))) + cross = mx.concatenate([prompt_cond, sec_tok.astype(mx.float32)], axis=1) + # CFG dropout (dit.py:441): whole cross_attn → zeros per sample; + # global_cond (the seconds token) is kept. + cross = mx.where(drop_mask[:, None, None], mx.zeros_like(cross), cross) + global_cond = sec_tok[:, 0, :] + cross16 = cross.astype(mx.float16) + global16 = global_cond.astype(mx.float16) + + def model_fn(noised, t): + return bundle.dit(noised, t.astype(noised.dtype), cross16, global16) + + return rectified_flow_loss(model_fn, latents, timesteps, + loss_mask=loss_mask) + + value_and_grad = nn.value_and_grad(bundle, loss_fn) + + # ── telemetry (underfit's loss_by_timestep.bin, struct "Iff") ──────────── + tele = open("loss_by_timestep.bin", "ab") + + def checkpoint(step, epoch): + fname = (f"{run_label}-step={step}-epoch={epoch}.safetensors" + if run_label else f"step={step}-epoch={epoch}.safetensors") + path = ckpt_dir / fname + save_lora_checkpoint( + bundle, path, + include=lora_config.get("include"), + exclude=lora_config.get("exclude"), + extra_config={"step": int(step), "epoch": int(epoch), + "base_model": base_model}) + print(f"\n ✓ checkpoint {path.name}") + return path + + # ── train loop (step-driven; epoch = one dataloader pass) ──────────────── + raw_step = 0 + epoch = epoch_offset + last_saved_step = None + t_start = time.time() + print(f"training: max {max_new_steps} new step(s) " + f"(global target {args.max_steps}), lr {args.lr:g}, " + f"batch {args.batch_size}, cfg-dropout {args.cfg_dropout_prob}") + try: + while raw_step < max_new_steps: + for batch in iterate_batches(dataset, args.batch_size, + shuffle=True, + seed=args.seed + epoch): + if raw_step >= max_new_steps: + break + global_step = raw_step + step_offset + 1 + + latents = mx.array(batch["latents"]).astype(mx.float16) + loss_mask = mx.array(batch["padding_mask"]) + seconds = mx.array(np.asarray(batch["seconds_total"], + dtype=np.float32)) + + # timesteps: sampler + model-config distribution shift + t_np = sample_training_timesteps( + args.timestep_sampler, latents.shape[0], rng=np_rng) + if args.dist_shift != "none": + t_np = shift_training_timesteps( + t_np, latents.shape[-1], shift_type=args.dist_shift, + options=DIST_SHIFT_DEFAULT["options"] + if args.dist_shift == "full" else None) + timesteps = mx.array(t_np) + + prompt_cond = build_conditioning(t5, padding_emb, + batch["prompt"]) + drop = mx.array((np_rng.random(latents.shape[0]) + < args.cfg_dropout_prob)) + + loss, grads = value_and_grad(bundle, latents, timesteps, + loss_mask, prompt_cond, seconds, + drop) + optimizer.update(bundle, grads) + mx.eval(loss, bundle.trainable_parameters(), optimizer.state) + + raw_step += 1 + loss_v = float(loss) + tele.write(struct.pack("Iff", global_step, + float(t_np.mean()), loss_v)) + if raw_step % 10 == 0: + tele.flush() + rate = raw_step / max(time.time() - t_start, 1e-9) + print(f"step {global_step} train/loss {loss_v:.6f} " + f"train/lr {args.lr:.3e} epoch {epoch} " + f"({rate:.2f} it/s)", flush=True) + + if args.checkpoint_every > 0 and \ + global_step % args.checkpoint_every == 0: + checkpoint(global_step, epoch) + last_saved_step = global_step + epoch += 1 + except KeyboardInterrupt: + print("\ninterrupted — saving final checkpoint") + + global_step = raw_step + step_offset + if raw_step > 0 and global_step != last_saved_step: + checkpoint(global_step, epoch) + tele.close() + print(f"done: {raw_step} step(s) in {time.time()-t_start:.0f}s") + + +if __name__ == "__main__": + main() From db8c55c064517919662265f3a4602287a1b7d589 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 15:07:25 -0400 Subject: [PATCH 09/28] mlx trainer: base-ckpt training + the two parity-critical forward conventions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --dit-weights flag: train on the BASE (rectified-flow) checkpoint like underfit does — inference uses the ARC weights, training must not (loud warning otherwise). Conversion recipe for the HF *-base safetensors into the 441-key MLX npz documented in TRAINING_CONVENTIONS.md. Training forward now feeds the diffusion_cond_inpaint pure-generation conditioning (all-ones inpaint mask + zero context) instead of the inference path's zeros — verified against underfit's torch loop with a controlled forward: loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB (fp16-weight bound; torch MPS-vs-CPU 3e-6 rel). Using inference-style zeros trains against the wrong conditioning regime (4x loss difference). Cross-attention correctly runs over all 257 padding-embedded tokens at training time (parity-tested; masking/slicing to valid tokens is wrong). Benchmark (M4 Pro, 30 steps, identical setup, full underfit loop on MPS via the mps-training/mps-support branches): MLX 1.65 it/s vs MPS 0.668 it/s — MLX 2.47x faster. Loss regimes match; both checkpoints pass underfit validation with identical key sets. --- optimized/mlx/TRAINING_CONVENTIONS.md | 31 +++++++++++++++++++++ optimized/mlx/scripts/lora_train_mlx.py | 37 ++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/optimized/mlx/TRAINING_CONVENTIONS.md b/optimized/mlx/TRAINING_CONVENTIONS.md index 53f2016..7ad02e7 100644 --- a/optimized/mlx/TRAINING_CONVENTIONS.md +++ b/optimized/mlx/TRAINING_CONVENTIONS.md @@ -153,3 +153,34 @@ To build (roughly in order): 5. **pre-encode CLI** on top of `audio_encoding.py` (600 s cap, fp32, unscaled fp32 npy + json sidecar with the exact metadata fields). 6. **demos** (optional, last): pipeline-based demo step with underfit cadence/format. + +## 9. Training-forward conventions discovered by parity testing (2026-07-14) + +Verified by a controlled forward (identical latents crop, numpy-seeded noise, +t=0.5, fixed prompt, fp32) through underfit's torch loop on MPS vs the MLX +trainer — final agreement **loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB** +(bounded by the fp16 npz weights; torch MPS-vs-CPU agrees at 3e-6 rel): + +- **Train on the BASE checkpoint** (`stabilityai/stable-audio-3-*-base`, + rectified_flow), NOT the shipped ARC (rf_denoiser) weights that inference + uses. `lora_train_mlx.py --dit-weights` + a 441-key npz conversion + (strip `model.model.`, `.gamma`→`.weight`, conv (out,in,k)→(out,k,in), + `to_local_embed.{0,2}`→`.seq.{0,2}`, conditioner keys → `cond.*`). +- **local_add_cond during training = all-ONES inpaint mask + zero context** + (the models are the diffusion_cond_inpaint variant). The inference path's + `local_add_cond=None` (≡ zeros) is an inference-only convention — using it + in training changes the loss by ~4x (wrong conditioning regime). +- **Cross-attention runs over all 257 tokens at training time too**: the + torch conditioner substitutes the learned padding_embedding into padded + positions (exactly like the optimized runtimes) and the attention mask is + effectively pass-through. Do NOT mask/slice to valid tokens. + +## 10. MPS vs MLX benchmark (M4 Pro 48 GB, sm-music base, 30 steps, +## dora-rows r16 α16 = 141 layers / 9.2M params, batch 1, crop 256, fp16 base) + +| | MPS (underfit loop, torch 2.13) | MLX (lora_train_mlx.py) | +|---|---|---| +| steady it/s | 0.668 (1498 ms/step) | **1.65 (2.47x faster)** | +| peak memory | 3.67 GiB driver-allocated | (not instrumented; comparable class) | +| loss regime | 1–43, spikes at low t | same (0.9–35, spikes at low t) | +| checkpoint | underfit-valid, 423 keys | underfit-valid, identical key set | diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 0b43cc0..6600f2e 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -100,6 +100,13 @@ def parse_args(): ap = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--dit", required=True, choices=list(DIT_CHOICES.keys())) + ap.add_argument("--dit-weights", default=None, + help="Override the DiT weight file — REQUIRED in practice for " + "training parity with underfit: train on the BASE " + "(rectified-flow) checkpoint, not the shipped ARC " + "(rf_denoiser) weights that inference uses. Convert the " + "HF *-base model.safetensors to npz first (441-key " + "layout; see TRAINING_CONVENTIONS.md).") ap.add_argument("--latents-dir", required=True, help="Root of pre-encoded npy+json pairs (scripts/pre_encode_mlx.py)") ap.add_argument("--lr", type=float, required=True, @@ -210,11 +217,21 @@ def main(): # ── models ─────────────────────────────────────────────────────────────── t0 = time.time() - dit_model, _ = load_dit(args.dit, T_lat=crop_len, dtype=mx.float16) + if args.dit_weights: + import importlib + mod = importlib.import_module(DIT_CHOICES[args.dit]["loader"]) + dit_model = mod.load_dit(args.dit_weights, T_lat=crop_len, + dtype=mx.float16, compile_=False) + cond_src = args.dit_weights + else: + print("WARNING: training on the shipped (ARC/rf_denoiser) weights — " + "underfit convention is to train on the BASE checkpoint " + "(pass --dit-weights)") + dit_model, _ = load_dit(args.dit, T_lat=crop_len, dtype=mx.float16) + cond_src = str(ensure_local(DIT_CHOICES[args.dit]["ckpt"])) print(f"DiT loaded ({time.time()-t0:.1f}s, fp16 base, T_lat={crop_len})") t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL))) # frozen - padding_emb, secs_embedder = load_conditioner_from_npz( - str(ensure_local(DIT_CHOICES[args.dit]["ckpt"])), prefix="cond.") + padding_emb, secs_embedder = load_conditioner_from_npz(cond_src, prefix="cond.") # ── inject adapters ───────────────────────────────────────────────────── report, saved_config = inject_from_lora_config( @@ -254,6 +271,17 @@ def main(): optimizer = optim.AdamW(learning_rate=args.lr, betas=[0.9, 0.999], eps=1e-8, weight_decay=0.0) + # ── training-time local conditioning (underfit/upstream convention) ────── + # The SA3 models are the diffusion_cond_inpaint variant: the trainer feeds + # pure generation as an all-ONES inpaint mask + zero context (verified at + # 80.7 dB forward parity vs the torch loop; the inference path's + # local_add_cond=None ≡ zeros is an inference-only convention and trains + # against the wrong conditioning regime — 4x loss difference). + lac_const = mx.array( + np.concatenate([np.ones((1, 1, crop_len), np.float32), + np.zeros((1, 256, crop_len), np.float32)], + axis=1).transpose(0, 2, 1)).astype(mx.float16) + # ── loss (adapters in the bundle are the only trainable params) ────────── def loss_fn(bundle, latents, timesteps, loss_mask, prompt_cond, seconds_totals, drop_mask): @@ -271,7 +299,8 @@ def loss_fn(bundle, latents, timesteps, loss_mask, prompt_cond, global16 = global_cond.astype(mx.float16) def model_fn(noised, t): - return bundle.dit(noised, t.astype(noised.dtype), cross16, global16) + return bundle.dit(noised, t.astype(noised.dtype), cross16, global16, + local_add_cond=lac_const) return rectified_flow_loss(model_fn, latents, timesteps, loss_mask=loss_mask) From a13e9ba49c1cb2728e5cbacecc8235ba1d750a82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 16:37:14 -0400 Subject: [PATCH 10/28] =?UTF-8?q?mlx=20lora:=20reformulate=20DoRA=20traini?= =?UTF-8?q?ng=20forward=20=E2=80=94=20no=20weight=20materialization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dora-rows/cols and the -xs variants (and lora-xs) no longer build the full [out,in] adapted weight per call. dora-rows is exactly a row-scale of the (base + low-rank) output with the norm computed in closed form: y = (x@W0.T + s·(x@A.T)@B.T) ⊙ m/rownorm + bias rownorm² = Σ_row W0² (cached const) + 2s·rowsum((W0@A.T)⊙B) + s²·rowsum((B@AAᵀ)⊙B) dora-cols is the input-feature-scale dual. One fp16 read of W0 through a rank-r matmul replaces ~5 fp32 full-matrix passes + an fp32 GEMM; the base matmul runs in native dtype. bora keeps the full-weight path (its nested col-norm has no rank-r expansion without storing W0²); SA3_LORA_NAIVE_DORA=1 restores the old path for ablation. Equivalence proven: outputs ≤9e-6 rel fp32, gradients ≤2.4e-7 wrt A/B/magnitude/M_xs vs the naive path; 55 tests. Layer micro-bench (12288x1536 dora-rows fwd+bwd): 7.2x faster. Full training step: sm-music 1.60→3.15 it/s, medium 0.52→~1.2 it/s, and peak training memory 9.8→5.5 GB (small) / 26.8→12.9 GB (medium) since the fp32 weight copies are gone. --- optimized/mlx/models/defs/lora.py | 147 ++++++++++++++++++++++++-- tests/test_mlx_lora.py | 164 ++++++++++++++++++++++++++++++ 2 files changed, 300 insertions(+), 11 deletions(-) diff --git a/optimized/mlx/models/defs/lora.py b/optimized/mlx/models/defs/lora.py index fdfb324..1c8aaac 100644 --- a/optimized/mlx/models/defs/lora.py +++ b/optimized/mlx/models/defs/lora.py @@ -8,6 +8,7 @@ import json import math +import os import re import typing as tp from dataclasses import dataclass @@ -38,6 +39,18 @@ *_XS_ADAPTER_TYPES, } _FULL_WEIGHT_ADAPTER_TYPES = _SUPPORTED_ADAPTER_TYPES - {"lora"} +# bora/bora-xs stay on the full-weight-materializing path: their column norm +# is taken over the row-rescaled intermediate m_r·V/rownorm(V), which has no +# rank-r expansion without also storing W0**2 — and bora is rare. Every other +# adapted type uses the reformulated low-rank forward (see +# _reformulated_linear_forward), which never builds the [out, in] weight. +_NAIVE_ONLY_ADAPTER_TYPES = {"bora", "bora-xs"} +_DORA_ROW_ADAPTER_TYPES = {"dora-rows", "dora-rows-xs"} +_DORA_COL_ADAPTER_TYPES = {"dora-cols", "dora-cols-xs"} + +# Ablation/fallback toggle (read once at import): SA3_LORA_NAIVE_DORA=1 +# restores the original full-weight forward for the reformulated types. +_NAIVE_DORA = os.environ.get("SA3_LORA_NAIVE_DORA", "") == "1" @dataclass(frozen=True) @@ -96,27 +109,50 @@ def __init__( source_name=self.source_name, ) _initialize_adapter(self, source_weight, fan_out=fan_out, fan_in=fan_in) + # Frozen-base row/column energy Σ W0² for the reformulated DoRA norm. + # The base never trains, so this is a constant computed once. The + # leading underscore keeps it out of MLX parameters()/checkpoints. + if self.adapter_type in _DORA_ROW_ADAPTER_TYPES: + self._w0_sq = mx.sum(source_weight * source_weight, axis=1) + elif self.adapter_type in _DORA_COL_ADAPTER_TYPES: + self._w0_sq = mx.sum(source_weight * source_weight, axis=0) def __call__(self, x): if self.adapter_type in _FULL_WEIGHT_ADAPTER_TYPES: - adapted_weight = _adapted_weight_2d( - _linear_source_weight_2d(self.base.weight), - adapter_type=self.adapter_type, - layer=self, - ) - output = x.astype(mx.float32) @ adapted_weight.T - bias = getattr(self.base, "bias", None) - if bias is not None: - output = output + bias.astype(mx.float32) - return output.astype(x.dtype) + if _NAIVE_DORA or self.adapter_type in _NAIVE_ONLY_ADAPTER_TYPES: + return self._full_weight_forward(x) + return _reformulated_linear_forward(self, x) base_output = self.base(x) adapter_output = (x.astype(mx.float32) @ self.lora_A.T) @ self.lora_B.T return base_output + (adapter_output * self.scaling).astype(base_output.dtype) + def _full_weight_forward(self, x): + """Original forward: materialize the full adapted weight per call. + + Kept verbatim as the bora/bora-xs path and as the + SA3_LORA_NAIVE_DORA=1 ablation fallback for the reformulated types. + """ + + adapted_weight = _adapted_weight_2d( + _linear_source_weight_2d(self.base.weight), + adapter_type=self.adapter_type, + layer=self, + ) + output = x.astype(mx.float32) @ adapted_weight.T + bias = getattr(self.base, "bias", None) + if bias is not None: + output = output + bias.astype(mx.float32) + return output.astype(x.dtype) + class LoRAConv1d(nn.Module): - """Trainable LoRA-family wrapper for an MLX Conv1d layer.""" + """Trainable LoRA-family wrapper for an MLX Conv1d layer. + + Stays on the original (full-weight for DoRA/BoRA) paths: the underfit + product exclude list removes all convs, so this wrapper is never on the + training hot path and does not need the reformulated forward. + """ def __init__( self, @@ -798,6 +834,95 @@ def _adapter_delta_2d(layer): return layer.lora_B @ layer.lora_A +def _effective_low_rank_factors(layer): + """Return fp32 ``(A, B)`` with ``delta == B @ A``, folding -xs cores. + + For the -xs variants the effective factors are ``B̃ = U @ M_xs`` [out, r] + and ``Ã = V.T`` [r, in] (U, V frozen); B̃ is recomputed per call — it is a + cheap [out, rank] product, and gradients flow to M_xs through it. + """ + + if layer.adapter_type in _XS_ADAPTER_TYPES: + return layer.V.T, layer.U @ layer.M_xs.astype(mx.float32) + return layer.lora_A, layer.lora_B + + +def _reformulated_linear_forward(layer, x): + """Adapted-linear forward without materializing the [out, in] weight. + + Mathematically identical to ``x @ _adapted_weight_2d(...).T + bias`` for + lora-xs / dora-rows(-xs) / dora-cols(-xs), with V = W0 + s·B@A: + + * lora-xs: y = x @ W0.T + s·((x @ Ã.T) @ B̃.T) + bias + * dora-rows: W' = diag(m / rownorm(V)) · V, a row scale of the output: + y = (x @ W0.T + s·((x @ A.T) @ B.T)) ⊙ c + bias, c = m / rownorm(V) + * dora-cols: W' = V · diag(m / colnorm(V)), a scale of the input features + (x @ W'.T = (x ⊙ c) @ V.T): + y = ((x ⊙ c) @ W0.T + s·(((x ⊙ c) @ A.T) @ B.T)) + bias + + The bias is NOT parametrized (torch applies the parametrization to the + weight only: y = x @ W'.T + bias), so it is added un-scaled after the + row/column scaling. The base matmul runs in the layer's native dtype; + the low-rank term, norms, and scaling are fp32; the result is cast back + to x.dtype like the original path. + """ + + a, b = _effective_low_rank_factors(layer) + scaling = float(layer.scaling) + weight = layer.base.weight + bias = getattr(layer.base, "bias", None) + + x32 = x.astype(mx.float32) + if layer.adapter_type in _DORA_COL_ADAPTER_TYPES: + col_scale = _dora_scale_no_materialize(layer, a, b, weight, norm_dim=0) + x32 = x32 * col_scale + base_output = x32.astype(x.dtype) @ weight.T + else: + base_output = x @ weight.T + output = base_output.astype(mx.float32) + ((x32 @ a.T) @ b.T) * scaling + if layer.adapter_type in _DORA_ROW_ADAPTER_TYPES: + output = output * _dora_scale_no_materialize( + layer, a, b, weight, norm_dim=1 + ) + if bias is not None: + output = output + bias.astype(mx.float32) + return output.astype(x.dtype) + + +def _dora_scale_no_materialize(layer, a, b, weight, *, norm_dim: int): + """``m / norm(V, axis=norm_dim)`` without materializing V = W0 + s·B@A. + + Expands the squared norm so only rank-r matmuls touch W0 (read once, in + its native dtype), with the constant ΣW0² cached at init: + + rownorm² = ΣW0²(rows) + 2s·rowsum((W0 @ A.T) ⊙ B) + s²·rowsum((B@G) ⊙ B) + with G = A @ A.T [r, r] + colnorm² = ΣW0²(cols) + 2s·rowsum((W0.T @ B) ⊙ A.T) + s²·colsum((M@A) ⊙ A) + with M = B.T @ B [r, r] + """ + + scaling = float(layer.scaling) + if norm_dim == 1: + cross_lhs = mx.matmul(weight, a.T.astype(weight.dtype)).astype( + mx.float32 + ) # W0 @ A.T -> [out, r] + cross = mx.sum(cross_lhs * b, axis=1) + gram = a @ a.T + quad = mx.sum((b @ gram) * b, axis=1) + else: + cross_lhs = mx.matmul(weight.T, b.astype(weight.dtype)).astype( + mx.float32 + ) # W0.T @ B -> [in, r] + cross = mx.sum(cross_lhs * a.T, axis=1) + gram = b.T @ b + quad = mx.sum((gram @ a) * a, axis=0) + norm_sq = layer._w0_sq + 2.0 * scaling * cross + scaling * scaling * quad + # Match _dora_weight_2d semantics: clamp the NORM (not norm²) at 1e-12; + # the max(·, 0) only guards tiny negative rounding in the expansion. + norm = mx.sqrt(mx.maximum(norm_sq, 0.0)) + return layer.magnitude.astype(mx.float32) / mx.maximum(norm, 1e-12) + + def _linear_source_weight_2d(weight): return weight.astype(mx.float32) diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py index 1d460c8..6237bb0 100644 --- a/tests/test_mlx_lora.py +++ b/tests/test_mlx_lora.py @@ -1,3 +1,4 @@ +import time from functools import partial from pathlib import Path @@ -13,6 +14,7 @@ from mlx.utils import tree_flatten from optimized.mlx.models.defs import dit_mlx, dit_mlx_medium +from optimized.mlx.models.defs import lora as lora_module from optimized.mlx.models.defs.lora import ( TrainableSecondsEmbedder, apply_lora_checkpoint, @@ -603,6 +605,168 @@ def build_model(): assert load_trainable_lora_state(build_model(), partial) == 1 +# --------------------------------------------------------------------------- +# Reformulated (no-full-weight-materialization) DoRA forward +# --------------------------------------------------------------------------- + +# lora-xs is included: it moved off the full-weight path onto the cheap +# low-rank path (delta == (U @ M_xs) @ V.T needs no materialization either). +REFORMULATED_ADAPTER_TYPES = ( + "dora-rows", + "dora-cols", + "dora-rows-xs", + "dora-cols-xs", + "lora-xs", +) + + +def _build_reform_layer( + adapter_type: str, + *, + bias: bool, + dtype, + fan_in: int = 24, + fan_out: int = 18, + rank: int = 4, + alpha: float = 6.0, + seed: int = 3, +): + mx.random.seed(seed) + base = nn.Linear(fan_in, fan_out, bias=bias) + base.weight = (mx.random.normal((fan_out, fan_in)) * 0.5).astype(dtype) + if bias: + base.bias = mx.random.normal((fan_out,)).astype(dtype) + layer = lora_module.LoRALinear( + base, + rank=rank, + alpha=alpha, + source_name="layer", + adapter_type=adapter_type, + ) + # Non-trivial adapter state so both the delta and the norms move. + if adapter_type.endswith("-xs"): + layer.M_xs = mx.random.normal((rank, rank)) * 0.2 + else: + layer.lora_A = mx.random.normal((rank, fan_in)) * 0.4 + layer.lora_B = mx.random.normal((fan_out, rank)) * 0.4 + if "dora" in adapter_type: + layer.magnitude = layer.magnitude * ( + 1.0 + 0.1 * mx.random.normal(layer.magnitude.shape) + ) + return layer + + +@pytest.mark.parametrize("adapter_type", REFORMULATED_ADAPTER_TYPES) +@pytest.mark.parametrize("bias", [False, True]) +@pytest.mark.parametrize("dtype", [mx.float32, mx.float16]) +def test_reformulated_dora_forward_matches_naive( + monkeypatch, adapter_type, bias, dtype +): + layer = _build_reform_layer(adapter_type, bias=bias, dtype=dtype) + x = mx.random.normal((5, 24)).astype(dtype) + + assert not lora_module._NAIVE_DORA # reformulation is the default + reformed_output = layer(x) + assert reformed_output.dtype == dtype + reformed = np.asarray(reformed_output, dtype=np.float32) + + monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) + naive = np.asarray(layer(x), dtype=np.float32) + + tol = 1e-4 if dtype == mx.float32 else 2e-2 + np.testing.assert_allclose(reformed, naive, rtol=tol, atol=tol) + + +@pytest.mark.parametrize("adapter_type", REFORMULATED_ADAPTER_TYPES) +@pytest.mark.parametrize("dtype", [mx.float32, mx.float16]) +def test_reformulated_dora_gradients_match_naive(monkeypatch, adapter_type, dtype): + layer = _build_reform_layer(adapter_type, bias=True, dtype=dtype) + x = mx.random.normal((5, 24)).astype(dtype) + + def loss_fn(model, values): + return mx.mean(model(values).astype(mx.float32) ** 2) + + loss_and_grad = nn.value_and_grad(layer, loss_fn) + loss_new, grads_new = loss_and_grad(layer, x) + mx.eval(loss_new, grads_new) + + monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) + loss_old, grads_old = loss_and_grad(layer, x) + mx.eval(loss_old, grads_old) + + new_flat = dict(tree_flatten(grads_new)) + old_flat = dict(tree_flatten(grads_old)) + expected_params = ( + {"M_xs"} if adapter_type.endswith("-xs") else {"lora_A", "lora_B"} + ) + if "dora" in adapter_type: + expected_params = expected_params | {"magnitude"} + assert expected_params.issubset(new_flat) + assert set(new_flat) == set(old_flat) + + tol = 1e-3 if dtype == mx.float32 else 3e-2 + np.testing.assert_allclose( + float(loss_new), float(loss_old), rtol=tol, atol=tol + ) + for name in sorted(new_flat): + np.testing.assert_allclose( + np.asarray(new_flat[name], dtype=np.float32), + np.asarray(old_flat[name], dtype=np.float32), + rtol=tol, + atol=tol, + err_msg=f"gradient mismatch for {name}", + ) + + +def test_reformulated_dora_forward_backward_is_faster(monkeypatch): + """Micro-timing: one medium-sized dora-rows layer, fwd+bwd, 20 iters. + + Direction-only check (the dedicated ablation benchmarks live elsewhere): + the reformulated path must beat materializing the full 12288x1536 weight. + """ + + fan_in, fan_out = 1536, 12288 + mx.random.seed(0) + base = nn.Linear(fan_in, fan_out, bias=True) + layer = lora_module.LoRALinear( + base, + rank=16, + alpha=16.0, + source_name="layer", + adapter_type="dora-rows", + ) + layer.lora_B = mx.random.normal((fan_out, 16)) * 0.02 + x = mx.random.normal((64, fan_in)) + + def loss_fn(model, values): + return mx.mean(model(values) ** 2) + + loss_and_grad = nn.value_and_grad(layer, loss_fn) + + def average_step_seconds(iterations: int = 20) -> float: + for _ in range(3): # warmup + loss, grads = loss_and_grad(layer, x) + mx.eval(loss, grads) + start = time.perf_counter() + for _ in range(iterations): + loss, grads = loss_and_grad(layer, x) + mx.eval(loss, grads) + return (time.perf_counter() - start) / iterations + + reformed = average_step_seconds() + monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) + naive = average_step_seconds() + + print( + f"\ndora-rows 12288x1536 fwd+bwd: reformed {reformed * 1e3:.2f} ms/iter" + f" vs naive {naive * 1e3:.2f} ms/iter ({naive / reformed:.1f}x)" + ) + assert reformed < naive, ( + f"reformed {reformed * 1e3:.2f} ms/iter is not faster than naive " + f"{naive * 1e3:.2f} ms/iter" + ) + + @needs_underfit_reference def test_checkpoint_key_naming_matches_real_underfit_checkpoint(tmp_path: Path): reference_state, reference_config = load_lora_checkpoint( From b4462313020658dc81fe06f6d51647825c066582 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 16:37:14 -0400 Subject: [PATCH 11/28] mlx trainer: compiled step, T5 conditioning cache, wired limit, --grad-checkpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Train step (loss+grad+optimizer update) is mx.compile'd with [bundle.state, optimizer.state, mx.random.state] capture (--no-compile for eager; optimizer.init up front so captured state is stable; the frozen-seconds branch hoisted out of the step for purity). T5Gemma prompt conditioning cached per exact prompt string (--no-t5-cache; bit-exact, ~100% hit rate in the tiny-dataset workflow). Wired memory limit raised to the device recommendation. --grad-checkpoint ports mlx-lm's grad_checkpoint onto the DiT blocks: bit-exact losses, −19%+ peak memory, ~1.4x step cost — for big crops on small Macs. Compiled-vs-eager losses agree to 8e-7 rel; save/resume verified under compile. Combined ablation (30 same-seed steps, crop 256): sm-music 1.60→3.78 it/s and medium 0.52→1.29 it/s at ≤3e-4 max loss deviation; vs the MPS/underfit loop that is 5.7x (small) and 5.1x (medium). --- optimized/mlx/scripts/lora_train_mlx.py | 178 +++++++++++++++++++++--- 1 file changed, 158 insertions(+), 20 deletions(-) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 6600f2e..42133bd 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -31,6 +31,19 @@ augmentation. seconds_total stays the FULL source duration after cropping (deliberate underfit convention). +Performance: + - the train step (loss + grad + optimizer update) is mx.compile'd by + default (MLX docs "Compile > Training graphs" pattern; --no-compile for + the eager path). Batch shapes are constant, so it compiles once. + - T5Gemma prompt conditioning is cached per exact prompt string (per-run, + in-memory; --no-t5-cache to disable) — the tiny-dataset workflow repeats + a handful of prompts, so the hit rate is ≈100% after the first epoch. + - --grad-checkpoint recomputes each transformer block's activations in the + backward pass (mlx-lm style mx.checkpoint) — memory for big crops on + small Macs, at the cost of roughly one extra forward. + - wired memory limit is raised to the device's recommended working-set + size at startup (mx.set_wired_limit). + Usage: python scripts/lora_train_mlx.py --dit sm-music \ --latents-dir output/latents/my-set --lr 1e-4 \ @@ -46,6 +59,7 @@ import sys import time import uuid +from functools import partial from pathlib import Path import numpy as np @@ -83,6 +97,28 @@ LATENT_CROP_DEFAULTS = {"sm-music": 1300, "sm-sfx": 1300, "medium": 4096} DIST_SHIFT_DEFAULT = {"shift_type": "full", "options": {"min_length": 256, "max_length": 4096}} +T5_CACHE_CAP = 512 # prompt-conditioning cache entries (oldest evicted first) + + +def grad_checkpoint(layer): + """ + Update all instances of type(layer) to use gradient checkpointing. + + Verbatim port of mlx-lm's mlx_lm/tuner/trainer.py:grad_checkpoint — + activations of every instance of the layer's class are recomputed during + the backward pass instead of being kept alive, trading ~one extra forward + for a much smaller peak working set. + """ + fn = type(layer).__call__ + + def checkpointed_fn(model, *args, **kwargs): + def inner_fn(params, *args, **kwargs): + model.update(params) + return fn(model, *args, **kwargs) + + return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs) + + type(layer).__call__ = checkpointed_fn class TrainBundle(nn.Module): @@ -142,6 +178,20 @@ def parse_args(): ap.add_argument("--dist-shift", default="full", choices=["none", "full", "flux", "logsnr"]) ap.add_argument("--cfg-dropout-prob", type=float, default=0.1) + # performance + ap.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True, + help="mx.compile the train step (loss+grad+optimizer update; " + "MLX 'Compile > Training graphs' pattern). " + "--no-compile restores the eager path.") + ap.add_argument("--grad-checkpoint", action="store_true", + help="Gradient-checkpoint every transformer block (mlx-lm " + "style mx.checkpoint): activations are recomputed in " + "the backward pass — much lower peak memory for big " + "crops, ~one extra forward of compute") + ap.add_argument("--t5-cache", action=argparse.BooleanOptionalAction, default=True, + help="Cache the padded T5Gemma prompt conditioning per exact " + f"prompt string (per-run, in-memory, max {T5_CACHE_CAP} " + "entries)") # resume ap.add_argument("--lora-ckpt-path", default=None) ap.add_argument("--step-offset", type=int, default=None) @@ -168,15 +218,46 @@ def resolve_offsets(args, ckpt_config) -> tuple[int, int]: return 0, 0 -def build_conditioning(t5, padding_emb, prompts, max_len=256): +def build_conditioning(t5, padding_emb, prompts, max_len=256, + cache=None, cache_stats=None): """Frozen prompt-side conditioning: T5Gemma embeddings with the learned padding embedding applied. Returns [B, 256, 768] fp32 (the trainable - seconds token is concatenated later, inside the grad scope).""" - embeds, mask = t5.encode(list(prompts), max_len=max_len) - mx.eval(embeds, mask) - padded = apply_prompt_padding(embeds.astype(mx.float32), mask, - padding_emb.astype(mx.float32)) - return mx.stop_gradient(padded) + seconds token is concatenated later, inside the grad scope). + + When ``cache`` (a dict) is given, the padded embedding is cached per + EXACT prompt string. The cache is per-run only (in-memory, never + persisted) and stores evaluated [1, max_len, 768] fp32 mx arrays; it is + capped at T5_CACHE_CAP entries, evicting oldest-inserted first. The + tiny-dataset underfit workflow repeats a handful of prompts, so after + the first epoch the T5Gemma forward is skipped entirely (~100% hits).""" + prompts = list(prompts) + if cache is None: + embeds, mask = t5.encode(prompts, max_len=max_len) + mx.eval(embeds, mask) + padded = apply_prompt_padding(embeds.astype(mx.float32), mask, + padding_emb.astype(mx.float32)) + return mx.stop_gradient(padded) + + missing = [p for p in dict.fromkeys(prompts) if p not in cache] + if cache_stats is not None: + n_miss = sum(1 for p in prompts if p in set(missing)) + cache_stats["misses"] += n_miss + cache_stats["hits"] += len(prompts) - n_miss + if missing: + embeds, mask = t5.encode(missing, max_len=max_len) + padded = apply_prompt_padding(embeds.astype(mx.float32), mask, + padding_emb.astype(mx.float32)) + mx.eval(padded) + for i, p in enumerate(missing): + while len(cache) >= T5_CACHE_CAP: + cache.pop(next(iter(cache))) # evict oldest + row = padded[i:i + 1] + mx.eval(row) + cache[p] = row + print(f" t5-cache: encoded {len(missing)} new prompt(s) " + f"({len(cache)} cached)") + return mx.stop_gradient(mx.concatenate([cache[p] for p in prompts], + axis=0)) def main(): @@ -184,6 +265,13 @@ def main(): np_rng = np.random.default_rng(args.seed) mx.random.seed(args.seed) + # Let Metal wire (pin) up to the recommended working set — avoids paging + # stalls when the training working set approaches physical memory. + if hasattr(mx, "set_wired_limit") and hasattr(mx, "device_info"): + wired = mx.device_info().get("max_recommended_working_set_size") + if wired: + mx.set_wired_limit(int(wired)) + alpha = float(args.alpha) if args.alpha is not None else float(args.rank) lora_config = underfit_lora_config() lora_config.update({"adapter_type": args.adapter_type, "rank": args.rank, @@ -240,6 +328,13 @@ def main(): f"rank {args.rank}, alpha {alpha:g} " f"({report.trainable_parameters/1e6:.2f}M trainable)") + if args.grad_checkpoint: + # Patch the TransformerBlock class (both DiT defs expose the blocks at + # dit.transformer.layers) so every block recomputes in backward. + grad_checkpoint(dit_model.transformer.layers[0]) + print(f"grad-checkpoint: on " + f"({len(dit_model.transformer.layers)} transformer blocks)") + secs_module = None if not args.no_conditioner_lora: secs_module = TrainableSecondsEmbedder(secs_embedder.W, secs_embedder.b) @@ -283,13 +378,16 @@ def main(): axis=1).transpose(0, 2, 1)).astype(mx.float16) # ── loss (adapters in the bundle are the only trainable params) ────────── + # `seconds_in` is the raw seconds batch [B] fp32 when the trainable + # seconds conditioner is active; with --no-conditioner-lora it is the + # PRE-computed frozen seconds token [B, 1, 768] (built outside the step so + # the compiled step stays pure mx — no numpy inside the traced graph). def loss_fn(bundle, latents, timesteps, loss_mask, prompt_cond, - seconds_totals, drop_mask): + seconds_in, drop_mask): if bundle.secs is not None: - sec_tok = bundle.secs(seconds_totals) # [B, 1, 768] fp32 + sec_tok = bundle.secs(seconds_in) # [B, 1, 768] fp32 else: - sec_tok = mx.stop_gradient( - secs_embedder(list(np.asarray(seconds_totals)))) + sec_tok = seconds_in # frozen, precomputed cross = mx.concatenate([prompt_cond, sec_tok.astype(mx.float32)], axis=1) # CFG dropout (dit.py:441): whole cross_attn → zeros per sample; # global_cond (the seconds token) is kept. @@ -307,6 +405,27 @@ def model_fn(noised, t): value_and_grad = nn.value_and_grad(bundle, loss_fn) + # ── train step (compiled by default — MLX "Compile > Training graphs") ── + # Captured state: model params (the LoRA adapters the optimizer mutates), + # optimizer moments, and the global PRNG key — rectified_flow_loss draws + # its noise INSIDE the step. Everything np/python (timestep sampling, CFG + # dropout draw, T5 encode, telemetry) stays OUTSIDE the step; batch + # shapes/dtypes are constant (fixed crop + batch), so this traces once. + # optimizer.init() materializes the moment arrays up front so the captured + # state structure never changes (no recompile on step 2). + optimizer.init(bundle.trainable_parameters()) + state = [bundle.state, optimizer.state, mx.random.state] + + def train_step(latents, timesteps, loss_mask, prompt_cond, seconds_in, + drop): + loss, grads = value_and_grad(bundle, latents, timesteps, loss_mask, + prompt_cond, seconds_in, drop) + optimizer.update(bundle, grads) + return loss + + if args.compile: + train_step = partial(mx.compile, inputs=state, outputs=state)(train_step) + # ── telemetry (underfit's loss_by_timestep.bin, struct "Iff") ──────────── tele = open("loss_by_timestep.bin", "ab") @@ -327,10 +446,15 @@ def checkpoint(step, epoch): raw_step = 0 epoch = epoch_offset last_saved_step = None + t5_cache = {} if args.t5_cache else None + t5_stats = {"hits": 0, "misses": 0} + if hasattr(mx, "reset_peak_memory"): + mx.reset_peak_memory() # measure the training loop, not weight loading t_start = time.time() print(f"training: max {max_new_steps} new step(s) " f"(global target {args.max_steps}), lr {args.lr:g}, " - f"batch {args.batch_size}, cfg-dropout {args.cfg_dropout_prob}") + f"batch {args.batch_size}, cfg-dropout {args.cfg_dropout_prob}, " + f"compile {'on' if args.compile else 'off'}") try: while raw_step < max_new_steps: for batch in iterate_batches(dataset, args.batch_size, @@ -356,15 +480,24 @@ def checkpoint(step, epoch): timesteps = mx.array(t_np) prompt_cond = build_conditioning(t5, padding_emb, - batch["prompt"]) + batch["prompt"], + cache=t5_cache, + cache_stats=t5_stats) drop = mx.array((np_rng.random(latents.shape[0]) < args.cfg_dropout_prob)) - - loss, grads = value_and_grad(bundle, latents, timesteps, - loss_mask, prompt_cond, seconds, - drop) - optimizer.update(bundle, grads) - mx.eval(loss, bundle.trainable_parameters(), optimizer.state) + if bundle.secs is not None: + seconds_in = seconds + else: + # frozen conditioner: build the seconds token OUTSIDE the + # (possibly compiled) step — see loss_fn + seconds_in = mx.stop_gradient(secs_embedder( + [float(s) for s in batch["seconds_total"]])) + + t_step = time.time() + loss = train_step(latents, timesteps, loss_mask, prompt_cond, + seconds_in, drop) + mx.eval(state, loss) + step_ms = (time.time() - t_step) * 1000.0 raw_step += 1 loss_v = float(loss) @@ -375,7 +508,7 @@ def checkpoint(step, epoch): rate = raw_step / max(time.time() - t_start, 1e-9) print(f"step {global_step} train/loss {loss_v:.6f} " f"train/lr {args.lr:.3e} epoch {epoch} " - f"({rate:.2f} it/s)", flush=True) + f"({rate:.2f} it/s, {step_ms:.0f} ms)", flush=True) if args.checkpoint_every > 0 and \ global_step % args.checkpoint_every == 0: @@ -389,6 +522,11 @@ def checkpoint(step, epoch): if raw_step > 0 and global_step != last_saved_step: checkpoint(global_step, epoch) tele.close() + if t5_cache is not None: + print(f"t5 cache: {t5_stats['hits']} hit(s) / " + f"{t5_stats['misses']} miss(es), {len(t5_cache)} cached") + print(f"peak memory: {mx.get_peak_memory()/2**30:.2f} GB" + + (" (grad-checkpoint on)" if args.grad_checkpoint else "")) print(f"done: {raw_step} step(s) in {time.time()-t_start:.0f}s") From 577a041d8c1a4c0c4d06a2f02f6342fd05cd2a51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Tue, 14 Jul 2026 16:52:12 -0400 Subject: [PATCH 12/28] mlx lora: BoRA speed-mode reformulation + --bora-mode {speed,memory} MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BoRA/BoRA-XS training forwards join the no-materialization path: y = (((x⊙β) @ W0.T + s·(((x⊙β)@A.T)@B.T)) ⊙ α) + bias α reuses the dora-rows closed-form rownorm; β's colnorm of the row-rescaled intermediate expands into three cheap terms, the W0²-weighted one via a cached native-dtype W0⊙W0 — the speed-mode memory cost (+1 copy of adapted weights, e.g. ~2.6 GiB for a full medium bora injection; the underfit default config adapts no convs and bora is rare, so usually 0). --bora-mode memory (trainer flag, threaded through inject_*) keeps the exact old full-weight path with no cache; SA3_LORA_NAIVE_DORA=1 still forces naive everywhere for ablation. Equivalence: fp32 forward ≤1.4e-6 abs, grads ≤8.5e-5 rel (incl. magnitude_r/magnitude_c/M_xs); fp16 within the shared tolerances with no loosening. Layer micro-bench (12288x1536 bora, fwd+bwd): 5.8x vs the full-weight path. 71 tests. --- optimized/mlx/models/defs/lora.py | 158 +++++++++++++++++++++--- optimized/mlx/scripts/lora_train_mlx.py | 12 +- tests/test_mlx_lora.py | 153 +++++++++++++++++++++++ 3 files changed, 305 insertions(+), 18 deletions(-) diff --git a/optimized/mlx/models/defs/lora.py b/optimized/mlx/models/defs/lora.py index 1c8aaac..ec05b89 100644 --- a/optimized/mlx/models/defs/lora.py +++ b/optimized/mlx/models/defs/lora.py @@ -39,12 +39,17 @@ *_XS_ADAPTER_TYPES, } _FULL_WEIGHT_ADAPTER_TYPES = _SUPPORTED_ADAPTER_TYPES - {"lora"} -# bora/bora-xs stay on the full-weight-materializing path: their column norm -# is taken over the row-rescaled intermediate m_r·V/rownorm(V), which has no -# rank-r expansion without also storing W0**2 — and bora is rare. Every other -# adapted type uses the reformulated low-rank forward (see -# _reformulated_linear_forward), which never builds the [out, in] weight. -_NAIVE_ONLY_ADAPTER_TYPES = {"bora", "bora-xs"} +# bora/bora-xs have a selectable forward (LoRALinear ``bora_mode``): their +# column norm is taken over the row-rescaled intermediate diag(α)·V with +# α = m_r/rownorm(V), which has no rank-r expansion without also storing +# W0**2. "speed" (default) caches that W0² once at init (+1 weight copy per +# bora layer, native dtype) and uses the reformulated low-rank forward; +# "memory" keeps the original full-weight-materializing forward and +# allocates no cache. Every other adapted type always uses the reformulated +# forward (see _reformulated_linear_forward), which never builds the +# [out, in] weight. +_BORA_ADAPTER_TYPES = {"bora", "bora-xs"} +_BORA_MODES = {"speed", "memory"} _DORA_ROW_ADAPTER_TYPES = {"dora-rows", "dora-rows-xs"} _DORA_COL_ADAPTER_TYPES = {"dora-cols", "dora-cols-xs"} @@ -86,10 +91,16 @@ def __init__( source_name: str, adapter_type: str = "lora", checkpoint_prefix: str = "model.", + bora_mode: str = "speed", ): super().__init__() if rank <= 0: raise ValueError(f"LoRA rank must be positive, got {rank}.") + if bora_mode not in _BORA_MODES: + raise ValueError( + f"bora_mode must be one of {sorted(_BORA_MODES)}, " + f"got {bora_mode!r}." + ) self.base = base self.base.freeze() @@ -99,6 +110,7 @@ def __init__( self.source_name = str(source_name) self.checkpoint_name = _checkpoint_name(source_name, checkpoint_prefix) self.adapter_type = canonical_adapter_type(adapter_type) + self.bora_mode = str(bora_mode) fan_out, fan_in = (int(value) for value in base.weight.shape) source_weight = _linear_source_weight_2d(base.weight) @@ -109,17 +121,28 @@ def __init__( source_name=self.source_name, ) _initialize_adapter(self, source_weight, fan_out=fan_out, fan_in=fan_in) - # Frozen-base row/column energy Σ W0² for the reformulated DoRA norm. - # The base never trains, so this is a constant computed once. The - # leading underscore keeps it out of MLX parameters()/checkpoints. + # Frozen-base row/column energy Σ W0² for the reformulated DoRA/BoRA + # norms. The base never trains, so this is a constant computed once. + # The leading underscore keeps it out of MLX parameters()/checkpoints. if self.adapter_type in _DORA_ROW_ADAPTER_TYPES: self._w0_sq = mx.sum(source_weight * source_weight, axis=1) elif self.adapter_type in _DORA_COL_ADAPTER_TYPES: self._w0_sq = mx.sum(source_weight * source_weight, axis=0) + elif self.adapter_type in _BORA_ADAPTER_TYPES and self.bora_mode == "speed": + # BoRA rownorm(V) reuses the dora-rows closed form (row Σ W0²); + # its colnorm(diag(α)·V) additionally needs the FULL element-wise + # W0² (α² re-weights the rows), cached once in the weight's native + # dtype — this is speed mode's memory cost: +1 weight copy per + # bora layer. "memory" mode allocates neither. + self._w0_sq = mx.sum(source_weight * source_weight, axis=1) + self._w0_sq_full = self.base.weight * self.base.weight def __call__(self, x): if self.adapter_type in _FULL_WEIGHT_ADAPTER_TYPES: - if _NAIVE_DORA or self.adapter_type in _NAIVE_ONLY_ADAPTER_TYPES: + if _NAIVE_DORA or ( + self.adapter_type in _BORA_ADAPTER_TYPES + and self.bora_mode == "memory" + ): return self._full_weight_forward(x) return _reformulated_linear_forward(self, x) @@ -130,8 +153,9 @@ def __call__(self, x): def _full_weight_forward(self, x): """Original forward: materialize the full adapted weight per call. - Kept verbatim as the bora/bora-xs path and as the - SA3_LORA_NAIVE_DORA=1 ablation fallback for the reformulated types. + Kept verbatim as the bora/bora-xs ``bora_mode="memory"`` path and as + the SA3_LORA_NAIVE_DORA=1 ablation fallback for the reformulated + types. """ adapted_weight = _adapted_weight_2d( @@ -152,6 +176,8 @@ class LoRAConv1d(nn.Module): Stays on the original (full-weight for DoRA/BoRA) paths: the underfit product exclude list removes all convs, so this wrapper is never on the training hot path and does not need the reformulated forward. + ``bora_mode`` intentionally does not apply here — bora/bora-xs convs + always use the full-weight forward (equivalent to "memory" mode). """ def __init__( @@ -339,6 +365,7 @@ def inject_trainable_lora( exclude: tp.Sequence[str] | None = None, adapter_type: str = "lora", checkpoint_prefix: str = "model.", + bora_mode: str = "speed", ) -> LoRAInjectionReport: """Freeze an MLX model and replace selected Linear/Conv1d layers. @@ -348,6 +375,11 @@ def inject_trainable_lora( ``conditioners.seconds_total.``). Include/exclude patterns match against both the bare runtime name and the full checkpoint name, so underfit dashboard filter strings work verbatim. + + ``bora_mode`` selects the bora/bora-xs Linear forward: "speed" (default) + caches W0² per layer for the reformulated forward, "memory" keeps the + original full-weight forward with no cache. It only affects bora-family + Linear layers (Conv1d always uses the full-weight path). """ alpha = float(rank if alpha is None else alpha) @@ -371,6 +403,7 @@ def inject_trainable_lora( source_name=name, adapter_type=adapter_type, checkpoint_prefix=checkpoint_prefix, + bora_mode=bora_mode, ) elif isinstance(layer, nn.Conv1d): replacement = LoRAConv1d( @@ -432,6 +465,7 @@ def inject_from_lora_config( lora_config: dict[str, tp.Any] | None, *, checkpoint_prefix: str = "model.", + bora_mode: str = "speed", ) -> tuple[LoRAInjectionReport, dict[str, tp.Any]]: """Inject adapters from an underfit ``lora_config`` dict. @@ -439,7 +473,9 @@ def inject_from_lora_config( defaults to 8, alpha defaults to rank, adapter_type defaults to "lora" (legacy "dora" resolves to "dora-rows"). Returns the injection report and the saved-config dict ({rank, alpha, adapter_type, include, exclude}) - ready to pass to save metadata. + ready to pass to save metadata. ``bora_mode`` is a runtime speed-vs-memory + knob (see inject_trainable_lora), not checkpoint semantics, so it is not + part of the saved config. """ lora_config = dict(lora_config or {}) @@ -459,6 +495,7 @@ def inject_from_lora_config( exclude=exclude, adapter_type=adapter_type, checkpoint_prefix=checkpoint_prefix, + bora_mode=bora_mode, ) saved_config = { "rank": rank, @@ -851,7 +888,8 @@ def _reformulated_linear_forward(layer, x): """Adapted-linear forward without materializing the [out, in] weight. Mathematically identical to ``x @ _adapted_weight_2d(...).T + bias`` for - lora-xs / dora-rows(-xs) / dora-cols(-xs), with V = W0 + s·B@A: + lora-xs / dora-rows(-xs) / dora-cols(-xs) / bora(-xs), with + V = W0 + s·B@A: * lora-xs: y = x @ W0.T + s·((x @ Ã.T) @ B̃.T) + bias * dora-rows: W' = diag(m / rownorm(V)) · V, a row scale of the output: @@ -859,6 +897,8 @@ def _reformulated_linear_forward(layer, x): * dora-cols: W' = V · diag(m / colnorm(V)), a scale of the input features (x @ W'.T = (x ⊙ c) @ V.T): y = ((x ⊙ c) @ W0.T + s·(((x ⊙ c) @ A.T) @ B.T)) + bias + * bora(-xs): W' = diag(α) · V · diag(β), both scales at once (see + _bora_reformulated_forward) The bias is NOT parametrized (torch applies the parametrization to the weight only: y = x @ W'.T + bias), so it is added un-scaled after the @@ -867,6 +907,9 @@ def _reformulated_linear_forward(layer, x): to x.dtype like the original path. """ + if layer.adapter_type in _BORA_ADAPTER_TYPES: + return _bora_reformulated_forward(layer, x) + a, b = _effective_low_rank_factors(layer) scaling = float(layer.scaling) weight = layer.base.weight @@ -889,8 +932,91 @@ def _reformulated_linear_forward(layer, x): return output.astype(x.dtype) +def _bora_reformulated_forward(layer, x): + """BoRA speed-mode forward: W' = diag(α)·V·diag(β) without building V. + + With V = W0 + s·B@A, α = m_r / rownorm(V) and + β = m_c / colnorm(diag(α)·V): + + y = (((x ⊙ β) @ W0.T + s·(((x ⊙ β) @ A.T) @ B.T)) ⊙ α) + bias + + α reuses the dora-rows closed-form rownorm (_v_norm_no_materialize with + the cached row ΣW0²); β needs the α²-reweighted column norm, which is + where the cached full W0² comes in (_bora_col_scale_no_materialize). α + depends on m_r, A, B; β depends on everything including α — all + differentiable, autodiff handles it. Same dtype policy as the other + reformulated types: base matmul in the layer's native dtype, low-rank + term and scales in fp32, bias un-scaled last, result cast to x.dtype. + """ + + a, b = _effective_low_rank_factors(layer) + scaling = float(layer.scaling) + weight = layer.base.weight + bias = getattr(layer.base, "bias", None) + + row_scale = layer.magnitude_r.astype(mx.float32) / _v_norm_no_materialize( + layer, a, b, weight, norm_dim=1 + ) + col_scale = _bora_col_scale_no_materialize(layer, a, b, weight, row_scale) + + x32 = x.astype(mx.float32) * col_scale + base_output = x32.astype(x.dtype) @ weight.T + output = base_output.astype(mx.float32) + ((x32 @ a.T) @ b.T) * scaling + output = output * row_scale + if bias is not None: + output = output + bias.astype(mx.float32) + return output.astype(x.dtype) + + +def _bora_col_scale_no_materialize(layer, a, b, weight, row_scale): + """``m_c / colnorm(diag(α)·V)`` without materializing V or diag(α)·V. + + colnorm²(diag(α)·V) per column j is Σᵢ αᵢ²·Vᵢⱼ², expanded into three + terms with V = W0 + s·B@A: + + 1. Σᵢ αᵢ²·W0ᵢⱼ² = α² @ W0² (cached _w0_sq_full, native + dtype — one weight-sized read) + 2. 2s·Σᵢ αᵢ²·W0ᵢⱼ·(BA)ᵢⱼ = 2s·rowsum((W0.T @ (α²[:,None] ⊙ B)) ⊙ A.T) + (one W0 read into a rank-r + matmul) + 3. s²·Σᵢ αᵢ²·(BA)ᵢⱼ² = s²·colsum((M@A) ⊙ A), M = B.T @ (α²[:,None]⊙B) + [r, r] — tiny + + Term 1 runs the matvec in the cache's native dtype (fp16 for an fp16 + base — accepted accumulation precision, validated by the equivalence + tests) and casts the result to fp32; terms 2-3 follow the existing + _v_norm_no_materialize dtype pattern. + """ + + scaling = float(layer.scaling) + alpha_sq = row_scale * row_scale # [out] fp32 + w0_sq = layer._w0_sq_full + base_term = mx.matmul( + alpha_sq.astype(w0_sq.dtype)[None, :], w0_sq + ).astype(mx.float32)[0] # α² @ W0² -> [in] + weighted_b = alpha_sq[:, None] * b # [out, r] fp32 + cross_lhs = mx.matmul(weight.T, weighted_b.astype(weight.dtype)).astype( + mx.float32 + ) # W0.T @ (α² ⊙ B) -> [in, r] + cross = mx.sum(cross_lhs * a.T, axis=1) + gram = b.T @ weighted_b + quad = mx.sum((gram @ a) * a, axis=0) + norm_sq = base_term + 2.0 * scaling * cross + scaling * scaling * quad + # Same clamp semantics as _v_norm_no_materialize / _bora_weight_2d. + norm = mx.sqrt(mx.maximum(norm_sq, 0.0)) + return layer.magnitude_c.astype(mx.float32) / mx.maximum(norm, 1e-12) + + def _dora_scale_no_materialize(layer, a, b, weight, *, norm_dim: int): - """``m / norm(V, axis=norm_dim)`` without materializing V = W0 + s·B@A. + """``m / norm(V, axis=norm_dim)`` without materializing V = W0 + s·B@A.""" + + return layer.magnitude.astype(mx.float32) / _v_norm_no_materialize( + layer, a, b, weight, norm_dim=norm_dim + ) + + +def _v_norm_no_materialize(layer, a, b, weight, *, norm_dim: int): + """Clamped ``norm(V, axis=norm_dim)`` without materializing V = W0 + s·B@A. Expands the squared norm so only rank-r matmuls touch W0 (read once, in its native dtype), with the constant ΣW0² cached at init: @@ -920,7 +1046,7 @@ def _dora_scale_no_materialize(layer, a, b, weight, *, norm_dim: int): # Match _dora_weight_2d semantics: clamp the NORM (not norm²) at 1e-12; # the max(·, 0) only guards tiny negative rounding in the expansion. norm = mx.sqrt(mx.maximum(norm_sq, 0.0)) - return layer.magnitude.astype(mx.float32) / mx.maximum(norm, 1e-12) + return mx.maximum(norm, 1e-12) def _linear_source_weight_2d(weight): diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 42133bd..f4bd6be 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -165,6 +165,12 @@ def parse_args(): ap.add_argument("--no-conditioner-lora", action="store_true", help="Skip the seconds-conditioner adapter (underfit adapts it " "by default — e.g. the plini checkpoint carries its delta)") + ap.add_argument("--bora-mode", default="speed", choices=["speed", "memory"], + help="bora/bora-xs forward: 'speed' (default) caches W0² per " + "adapted layer (+1 weight copy) for the reformulated " + "no-materialize forward; 'memory' keeps the original " + "full-weight forward with no cache. Ignored for other " + "adapter types.") # data ap.add_argument("--latent-crop-length", type=int, default=None, help="Default per model: 1300 (sm-*) / 4096 (medium)") @@ -323,7 +329,8 @@ def main(): # ── inject adapters ───────────────────────────────────────────────────── report, saved_config = inject_from_lora_config( - dit_model, lora_config, checkpoint_prefix="model.") + dit_model, lora_config, checkpoint_prefix="model.", + bora_mode=args.bora_mode) print(f"lora: {report.layer_count} DiT layer(s), {report.adapter_type}, " f"rank {args.rank}, alpha {alpha:g} " f"({report.trainable_parameters/1e6:.2f}M trainable)") @@ -341,7 +348,8 @@ def main(): try: cond_report, _ = inject_from_lora_config( secs_module, lora_config, - checkpoint_prefix="conditioners.seconds_total.") + checkpoint_prefix="conditioners.seconds_total.", + bora_mode=args.bora_mode) print(f"lora: conditioner seconds embedder adapted " f"({cond_report.layer_count} layer)") except ValueError: diff --git a/tests/test_mlx_lora.py b/tests/test_mlx_lora.py index 6237bb0..d8e8e24 100644 --- a/tests/test_mlx_lora.py +++ b/tests/test_mlx_lora.py @@ -611,12 +611,17 @@ def build_model(): # lora-xs is included: it moved off the full-weight path onto the cheap # low-rank path (delta == (U @ M_xs) @ V.T needs no materialization either). +# bora/bora-xs are included in their default "speed" mode, which caches W0² +# once at init to expand colnorm(diag(α)·V) without materializing V +# ("memory" mode keeps the naive full-weight forward). REFORMULATED_ADAPTER_TYPES = ( "dora-rows", "dora-cols", "dora-rows-xs", "dora-cols-xs", "lora-xs", + "bora", + "bora-xs", ) @@ -653,6 +658,13 @@ def _build_reform_layer( layer.magnitude = layer.magnitude * ( 1.0 + 0.1 * mx.random.normal(layer.magnitude.shape) ) + elif "bora" in adapter_type: + layer.magnitude_r = layer.magnitude_r * ( + 1.0 + 0.1 * mx.random.normal(layer.magnitude_r.shape) + ) + layer.magnitude_c = layer.magnitude_c * ( + 1.0 + 0.1 * mx.random.normal(layer.magnitude_c.shape) + ) return layer @@ -673,6 +685,9 @@ def test_reformulated_dora_forward_matches_naive( monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) naive = np.asarray(layer(x), dtype=np.float32) + # bora needs NO looser fp16 tolerance: its colnorm term-1 matvec runs in + # the cache's native dtype (fp16 for an fp16 base), but the resulting + # error stays ~4e-3 max abs — within the shared 2e-2 fp16 tolerance. tol = 1e-4 if dtype == mx.float32 else 2e-2 np.testing.assert_allclose(reformed, naive, rtol=tol, atol=tol) @@ -701,6 +716,8 @@ def loss_fn(model, values): ) if "dora" in adapter_type: expected_params = expected_params | {"magnitude"} + elif "bora" in adapter_type: + expected_params = expected_params | {"magnitude_r", "magnitude_c"} assert expected_params.issubset(new_flat) assert set(new_flat) == set(old_flat) @@ -767,6 +784,142 @@ def average_step_seconds(iterations: int = 20) -> float: ) +@pytest.mark.parametrize("adapter_type", ["bora", "bora-xs"]) +def test_bora_memory_mode_is_naive_and_mode_threads_from_config( + monkeypatch, adapter_type +): + """bora_mode="memory" IS the naive full-weight code path (bit-identical), + allocates no W0² cache, and the mode threads through + inject_from_lora_config → inject_trainable_lora → LoRALinear.""" + + def build(bora_mode): + mx.random.seed(9) + model = WideMLXRegressor() + report, _ = inject_from_lora_config( + model, + {"adapter_type": adapter_type, "rank": 4}, + bora_mode=bora_mode, + ) + assert report.adapter_type == adapter_type + for layer in iter_trainable_lora_layers(model): + if not adapter_type.endswith("-xs"): + layer.lora_B = mx.random.normal(layer.lora_B.shape) * 0.3 + else: + layer.M_xs = mx.random.normal(layer.M_xs.shape) * 0.3 + return model + + x = mx.random.normal((3, 16)) + + memory_model = build("memory") + for layer in iter_trainable_lora_layers(memory_model): + assert layer.bora_mode == "memory" + assert not hasattr(layer, "_w0_sq_full") # no cache allocated + assert not hasattr(layer, "_w0_sq") + memory_output = memory_model(x) + + # Bit-identical to SA3_LORA_NAIVE_DORA=1 — memory mode dispatches to the + # very same _full_weight_forward, not a numerically-close reimplementation. + monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) + assert bool(mx.array_equal(memory_output, memory_model(x))) + monkeypatch.setattr(lora_module, "_NAIVE_DORA", False) + + # Default mode is "speed": cache allocated, same math within fp32 tol. + speed_model = build("speed") + for layer in iter_trainable_lora_layers(speed_model): + assert layer.bora_mode == "speed" + assert layer._w0_sq_full.dtype == layer.base.weight.dtype + assert layer._w0_sq_full.shape == layer.base.weight.shape + np.testing.assert_allclose( + np.asarray(speed_model(x), dtype=np.float32), + np.asarray(memory_output, dtype=np.float32), + rtol=1e-4, + atol=1e-4, + ) + + # The env ablation toggle still overrides bora_mode="speed" too. + monkeypatch.setattr(lora_module, "_NAIVE_DORA", True) + assert bool(mx.array_equal(speed_model(x), memory_output)) + + with pytest.raises(ValueError, match="bora_mode"): + build("fast") + + +def test_bora_w0_sq_cache_stays_out_of_parameters_and_checkpoints( + tmp_path: Path, +): + mx.random.seed(13) + model = WideMLXRegressor() + inject_trainable_lora(model, rank=4, alpha=4, adapter_type="bora") + + param_names = [name for name, _ in tree_flatten(model.parameters())] + assert any(name.endswith(".lora_A") for name in param_names) + assert not any("_w0_sq" in name for name in param_names) # incl. _w0_sq_full + assert not any( + "_w0_sq" in name for name, _ in tree_flatten(model.trainable_parameters()) + ) + + checkpoint = save_lora_checkpoint(model, tmp_path / "bora.safetensors") + state_dict, config = load_lora_checkpoint(checkpoint) + + assert config["adapter_type"] == "bora" + assert not any("_w0_sq" in key for key in state_dict) + assert sorted(key.rsplit(".", 1)[1] for key in state_dict) == sorted( + ["lora_A", "lora_B", "magnitude_r", "magnitude_c"] * 2 + ) + + +def test_reformulated_bora_forward_backward_is_faster(): + """Micro-timing: one medium-sized bora layer, fwd+bwd, 20 iters. + + Direction-only check like the dora-rows one: speed mode (reformulated, + cached W0²) must beat memory mode (materializing the full 12288x1536 + weight). Expected in the same ballpark as the dora-rows reform (7.2x) + minus the extra α² @ W0² matvec. + """ + + fan_in, fan_out = 1536, 12288 + mx.random.seed(0) + base = nn.Linear(fan_in, fan_out, bias=True) + layer = lora_module.LoRALinear( + base, + rank=16, + alpha=16.0, + source_name="layer", + adapter_type="bora", + ) + layer.lora_B = mx.random.normal((fan_out, 16)) * 0.02 + x = mx.random.normal((64, fan_in)) + + def loss_fn(model, values): + return mx.mean(model(values) ** 2) + + loss_and_grad = nn.value_and_grad(layer, loss_fn) + + def average_step_seconds(iterations: int = 20) -> float: + for _ in range(3): # warmup + loss, grads = loss_and_grad(layer, x) + mx.eval(loss, grads) + start = time.perf_counter() + for _ in range(iterations): + loss, grads = loss_and_grad(layer, x) + mx.eval(loss, grads) + return (time.perf_counter() - start) / iterations + + assert layer.bora_mode == "speed" + speed = average_step_seconds() + layer.bora_mode = "memory" # dispatch-only switch; cache is just unused + memory = average_step_seconds() + + print( + f"\nbora 12288x1536 fwd+bwd: speed {speed * 1e3:.2f} ms/iter" + f" vs memory {memory * 1e3:.2f} ms/iter ({memory / speed:.1f}x)" + ) + assert speed < memory, ( + f"speed mode {speed * 1e3:.2f} ms/iter is not faster than memory " + f"mode {memory * 1e3:.2f} ms/iter" + ) + + @needs_underfit_reference def test_checkpoint_key_naming_matches_real_underfit_checkpoint(tmp_path: Path): reference_state, reference_config = load_lora_checkpoint( From d39980ba914f1e3c20dd655eb48e21be47b6c791 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 12:06:03 -0400 Subject: [PATCH 13/28] mlx trainer: training-time demos (RF Euler + CFG) + gradio LoRA preload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add underfit-style demos to lora_train_mlx.py: baseline at step 0 then every --demo-every steps + a final render. models/defs/demo_mlx.py runs MLX inference on the base rectified_flow model + the trained LoRA — plain Euler velocity integration (x += dt*v, distinct from the distilled rf_denoiser pingpong), the same conditioning/CFG as sa3_mlx (uncond = cross->zeros, denoised-space guidance, local_add_cond=None), decode via SAME-S/SAME-L, peak-normalized int16 -> mp3 with a json sidecar, files demo__.mp3, idempotent per step. Per-entry cfg/seed/steps/duration; lora_strength and lora_interval_max (sigma gating) are applied by scaling lora_B (M_xs for -xs) pre-norm, which is exactly underfit's lora_strength semantics for lora/DoRA and touches no forward code. sa3_gradio.py: --lora PATH (repeatable, preloads a slot active-on-launch) and --port, so underfit can open the MLX gradio with a trained checkpoint loaded. --- optimized/mlx/models/defs/demo_mlx.py | 171 ++++++++++++++++++++++++ optimized/mlx/scripts/lora_train_mlx.py | 115 +++++++++++++++- optimized/mlx/scripts/sa3_gradio.py | 28 +++- 3 files changed, 305 insertions(+), 9 deletions(-) create mode 100644 optimized/mlx/models/defs/demo_mlx.py diff --git a/optimized/mlx/models/defs/demo_mlx.py b/optimized/mlx/models/defs/demo_mlx.py new file mode 100644 index 0000000..fe5072a --- /dev/null +++ b/optimized/mlx/models/defs/demo_mlx.py @@ -0,0 +1,171 @@ +"""Training-time demos for lora_train_mlx.py — MLX RF inference on the base model. + +Mirrors the RF/V path of underfit's demo step (underfit/training/demo_step.py): +plain Euler velocity integration (x += dt·v) over the sa3 sampling schedule, +optional CFG (uncond = cross→zeros, guidance done in denoised space — identical +to sa3_mlx), decode through the SAME-S/SAME-L decoder, peak-normalized int16 → +mp3 with a json sidecar, files named ``demo__.mp3`` in the demo +dir, idempotent per step (a resume never re-renders an existing demo). + +Why Euler and not the pingpong sampler sa3_mlx ships: the trainer finetunes the +rectified_flow BASE checkpoint, whose velocity field is integrated with a plain +ODE step. sa3_mlx's ``sample_flow_pingpong`` is the 8-step re-noising sampler for +the DISTILLED rf_denoiser (ARC) weights, which is a different model. Conditioning, +CFG math and ``local_add_cond=None`` (the inference convention — training uses the +all-ones inpaint mask, see TRAINING_CONVENTIONS §9) are otherwise identical. + +Per-entry ``lora_strength`` and ``lora_interval_max`` (σ-interval gating) are +applied by scaling each adapter's ``lora_B`` (``M_xs`` for -xs) — pre-norm, which +is exactly underfit's ``lora_strength`` semantics for lora/DoRA (it scales the +low-rank delta inside V = W0 + scaling·strength·B@A before the DoRA norm). This +touches no forward code, so the parity-verified training path is untouched. +""" +from __future__ import annotations + +import json +import os +import subprocess + +import mlx.core as mx +import numpy as np + +from .sa3_pipeline import build_pingpong_schedule, patched_decode + +SAMPLE_RATE = 44100 +SAMPLES_PER_LATENT = 4096 + + +# ── sampler ──────────────────────────────────────────────────────────────── + +def make_rf_model_fn(dit, cross_attn, global_cond, cfg=1.0, + null_cross_attn=None, apg=1.0): + """Velocity model_fn(x, t) for the base rectified_flow DiT. + + cfg == 1.0: a single conditioned forward. cfg != 1.0: the sa3_mlx dual pass — + batch cond+uncond, convert to denoised space (denoised = x − σ·v, σ = t), + guide there (optional APG orthogonal projection), convert back to velocity. + local_add_cond stays None (inference convention). + """ + def model_fn(x, t): + if cfg == 1.0: + return dit(x, t, cross_attn, global_cond, local_add_cond=None) + x2 = mx.concatenate([x, x], axis=0) + t2 = mx.concatenate([t, t], axis=0) + cross2 = mx.concatenate([cross_attn, null_cross_attn], axis=0) + global2 = mx.concatenate([global_cond, global_cond], axis=0) + v_batched = dit(x2, t2, cross2, global2, local_add_cond=None) + cond_v, uncond_v = mx.split(v_batched, 2, axis=0) + sigma = t.reshape(-1, 1, 1).astype(mx.float32) + cond_d = x.astype(mx.float32) - cond_v.astype(mx.float32) * sigma + uncond_d = x.astype(mx.float32) - uncond_v.astype(mx.float32) * sigma + diff = cond_d - uncond_d + if apg <= 0.0: + cfg_diff = diff + else: + flat_c = cond_d.reshape(cond_d.shape[0], -1) + flat_d = diff.reshape(diff.shape[0], -1) + unit = flat_c / (mx.linalg.norm(flat_c, axis=1, keepdims=True) + 1e-8) + proj = mx.sum(flat_d * unit, axis=1, keepdims=True) * unit + diff_orth = (flat_d - proj).reshape(diff.shape) + cfg_diff = diff_orth if apg >= 1.0 else (apg * diff_orth + (1.0 - apg) * diff) + cfg_d = cond_d + (cfg - 1.0) * cfg_diff + cfg_v = (x.astype(mx.float32) - cfg_d) / sigma + return cfg_v.astype(x.dtype) + + return model_fn + + +def rf_euler_sample(model_fn, x, sigmas, before_step=None): + """Euler ODE integration of the RF velocity field: x_{i+1} = x_i + Δt·v. + + ``sigmas`` is (steps+1,) descending from σ_max to 0. ``before_step(i, σ)``, + if given, runs right before each velocity eval (per-step LoRA gating). + """ + num_steps = sigmas.shape[0] - 1 + for i in range(num_steps): + t_curr = sigmas[i] + t_next = sigmas[i + 1] + if before_step is not None: + before_step(i, float(t_curr)) + t_vec = t_curr * mx.ones((x.shape[0],), dtype=x.dtype) + v = model_fn(x, t_vec) + x = x + (t_next - t_curr).astype(x.dtype) * v + mx.eval(x) + return x + + +# ── decode ───────────────────────────────────────────────────────────────── + +def decode_latents(decoder, chunk_fn, chunk_cfg, latents, T_lat): + """Latent [1,256,T] → audio np.float32 [2, T·4096]. Mirrors sa3_mlx's + chunk/kernel dispatch (odd T_lat routes through even-kernel chunked decode).""" + latents = latents.astype(mx.float32) + chunk, ovl = chunk_cfg + kernel = chunk + 2 * ovl + if T_lat > kernel: + patches = chunk_fn(decoder, latents, chunk, ovl) + elif T_lat % 2 == 0: + patches = decoder(latents) + elif T_lat > 6: + patches = chunk_fn(decoder, latents, 2, 2) + else: + latents_even = mx.concatenate([latents, latents[..., -1:]], axis=-1) + patches = decoder(latents_even)[..., : T_lat * 16] + audio = patched_decode(patches, patch_size=256, channels=2) # (1,2,T*4096) + mx.eval(audio) + return np.array(audio.astype(mx.float32))[0] + + +def save_demo_mp3(audio_np, demo_index, step, sample_rate, out_dir, meta=None): + """Peak-normalized int16 → mp3 + json sidecar (underfit's on-disk demo format). + + Atomic (temp wav → ffmpeg → rename); idempotent (skips if the mp3 exists). + """ + final_mp3 = os.path.join(out_dir, f"demo_{demo_index}_{step:08d}.mp3") + if os.path.exists(final_mp3): + return final_mp3 + peak = float(np.abs(audio_np).max()) + pcm = (audio_np / max(peak, 1e-8) * 32767.0).astype(np.int16) # (2, T) + tmp_wav = os.path.join(out_dir, f".tmp_demo_{demo_index}_{step:08d}.wav") + tmp_mp3 = os.path.join(out_dir, f".tmp_demo_{demo_index}_{step:08d}.mp3") + import wave + with wave.open(tmp_wav, "wb") as w: + w.setnchannels(pcm.shape[0]) + w.setsampwidth(2) + w.setframerate(sample_rate) + w.writeframes(pcm.T.tobytes()) # (T, C) interleaved + subprocess.run( + ["ffmpeg", "-i", tmp_wav, "-q:a", "0", "-y", tmp_mp3, "-loglevel", "error"], + check=True, + ) + os.replace(tmp_mp3, final_mp3) + os.remove(tmp_wav) + if meta is not None: + with open(os.path.join(out_dir, f"demo_{demo_index}_{step:08d}.json"), "w") as f: + json.dump(meta, f) + return final_mp3 + + +# ── per-entry LoRA strength / σ-interval (scale lora_B pre-norm) ───────────── + +def snapshot_adapters(layers): + """Freeze the current adapter factors so per-step scaling starts from the + trained values each time (avoids compounding across steps).""" + snap = [] + for layer in layers: + if hasattr(layer, "M_xs"): + snap.append((layer, "M_xs", layer.M_xs)) + elif hasattr(layer, "lora_B"): + snap.append((layer, "lora_B", layer.lora_B)) + return snap + + +def scale_adapters(snapshot, strength): + """Set every adapter's low-rank factor to (trained value)·strength — matches + underfit's pre-norm lora_strength for lora/DoRA. strength 1.0 restores.""" + for layer, attr, original in snapshot: + setattr(layer, attr, original * strength if strength != 1.0 else original) + + +def build_sigmas(steps): + return build_pingpong_schedule(steps) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index f4bd6be..cb76bce 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -54,6 +54,7 @@ import argparse import json import math +import os import re import struct import sys @@ -84,9 +85,11 @@ rectified_flow_loss, sample_training_timesteps, shift_training_timesteps, ) from models.defs.lora import ( # noqa: E402 - TrainableSecondsEmbedder, inject_from_lora_config, load_lora_checkpoint, - load_trainable_lora_state, save_lora_checkpoint, underfit_lora_config, + TrainableSecondsEmbedder, inject_from_lora_config, iter_trainable_lora_layers, + load_lora_checkpoint, load_trainable_lora_state, save_lora_checkpoint, + underfit_lora_config, ) +from models.defs import demo_mlx as demo # noqa: E402 from models.defs.latent_dataset import ( # noqa: E402 PreEncodedLatentDataset, iterate_batches, ) @@ -202,6 +205,20 @@ def parse_args(): ap.add_argument("--lora-ckpt-path", default=None) ap.add_argument("--step-offset", type=int, default=None) ap.add_argument("--epoch-offset", type=int, default=None) + # demos (underfit convention: RF Euler inference on the base model + trained + # LoRA, decode → mp3, baseline at step 0 then every --demo-every steps) + ap.add_argument("--demo-every", type=int, default=0, + help="Generate demos every N steps (0 = off). A baseline " + "set is also rendered at step 0.") + ap.add_argument("--demo-config", default=None, + help="JSON list of demo entries. Each: {prompt, cfg=7, " + "seed=0, steps=50, lora_strength?, lora_interval_max?}.") + ap.add_argument("--demo-decoder", default=None, + choices=["same-s", "same-l"], + help="Decoder for demos (default: the DiT's default_decoder).") + ap.add_argument("--demo-dir", default=None, + help="Where to write demo__.mp3 (default: cwd, " + "which the dashboard sets to /demos/).") return ap.parse_args() @@ -456,6 +473,93 @@ def checkpoint(step, epoch): last_saved_step = None t5_cache = {} if args.t5_cache else None t5_stats = {"hits": 0, "misses": 0} + + # ── demos (underfit RF path) ────────────────────────────────────────────── + # RF Euler inference on the base model + trained LoRA, decoded to mp3. Baseline + # at the start, then every --demo-every steps + a final render. Idempotent. + demo_entries = None + if args.demo_every > 0: + if args.demo_config: + demo_entries = json.loads(Path(args.demo_config).read_text()) + print(f"demos: {len(demo_entries)} prompt(s) every {args.demo_every} " + f"step(s) → {args.demo_dir or '.'}") + else: + print("WARNING: --demo-every set but no --demo-config — demos disabled") + demo_dir = args.demo_dir or "." + demo_decoder_name = args.demo_decoder or DIT_CHOICES[args.dit]["default_decoder"] + _demo_dec = {} # lazy: (decoder, chunk_fn, chunk_cfg) + + def run_demos(step): + if not demo_entries: + return + os.makedirs(demo_dir, exist_ok=True) + pending = [(i, e) for i, e in enumerate(demo_entries) + if not os.path.exists( + os.path.join(demo_dir, f"demo_{i}_{step:08d}.mp3"))] + if not pending: + return + if "dec" not in _demo_dec: + from sa3_mlx import load_decoder + t_d = time.time() + _demo_dec["dec"] = load_decoder(demo_decoder_name, dtype=mx.float32) + print(f" demo decoder {demo_decoder_name} loaded ({time.time()-t_d:.1f}s)") + decoder, chunk_fn, chunk_cfg = _demo_dec["dec"] + adapters = list(iter_trainable_lora_layers(bundle)) + T_lat = crop_len + seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE + for i, entry in pending: + prompt = entry.get("prompt", "") + cfg = float(entry.get("cfg", 7)) + seed = int(entry.get("seed", 0)) + steps = int(entry.get("steps", 50)) + strength = entry.get("lora_strength") + interval = entry.get("lora_interval_max") + prompt_cond = build_conditioning(t5, padding_emb, [prompt], + cache=t5_cache, cache_stats=t5_stats) + if bundle.secs is not None: + sec_tok = bundle.secs(mx.array([seconds_val], dtype=mx.float32)) + else: + sec_tok = secs_embedder([seconds_val]) + sec_tok = mx.stop_gradient(sec_tok.astype(mx.float32)) + cross = mx.concatenate([prompt_cond, sec_tok], axis=1).astype(mx.float16) + gcond = sec_tok[:, 0, :].astype(mx.float16) + null = mx.zeros_like(cross) if cfg != 1.0 else None + model_fn = demo.make_rf_model_fn(bundle.dit, cross, gcond, cfg=cfg, + null_cross_attn=null) + sigmas = demo.build_sigmas(steps) + noise = mx.random.normal((1, 256, T_lat), dtype=mx.float16, + key=mx.random.key(seed)) + snap = before = None + if interval is not None: + snap = demo.snapshot_adapters(adapters) + s_on = float(strength) if strength is not None else 1.0 + + def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): + demo.scale_adapters(snap, s_on if sigma <= interval else 0.0) + elif strength is not None and float(strength) != 1.0: + snap = demo.snapshot_adapters(adapters) + demo.scale_adapters(snap, float(strength)) + try: + latent = demo.rf_euler_sample(model_fn, noise, sigmas, + before_step=before) + audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, latent, T_lat) + meta = {"prompt": prompt, "cfg": cfg, "seed": seed, + "steps": steps, "step": int(step)} + if strength is not None: + meta["lora_strength"] = strength + if interval is not None: + meta["lora_interval_max"] = interval + demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE, demo_dir, meta) + print(f" ♪ demo {i} @ step {step}: {prompt[:48]!r} " + f"→ demo_{i}_{step:08d}.mp3", flush=True) + except Exception as e: + print(f" demo {i} failed: {e}", flush=True) + finally: + if snap is not None: + demo.scale_adapters(snap, 1.0) # restore trained factors + if hasattr(mx, "clear_cache"): + mx.clear_cache() + if hasattr(mx, "reset_peak_memory"): mx.reset_peak_memory() # measure the training loop, not weight loading t_start = time.time() @@ -463,6 +567,7 @@ def checkpoint(step, epoch): f"(global target {args.max_steps}), lr {args.lr:g}, " f"batch {args.batch_size}, cfg-dropout {args.cfg_dropout_prob}, " f"compile {'on' if args.compile else 'off'}") + run_demos(step_offset) # baseline (untrained / resumed state) try: while raw_step < max_new_steps: for batch in iterate_batches(dataset, args.batch_size, @@ -520,8 +625,10 @@ def checkpoint(step, epoch): if args.checkpoint_every > 0 and \ global_step % args.checkpoint_every == 0: - checkpoint(global_step, epoch) + checkpoint(global_step, epoch) # save BEFORE demos last_saved_step = global_step + if args.demo_every > 0 and global_step % args.demo_every == 0: + run_demos(global_step) epoch += 1 except KeyboardInterrupt: print("\ninterrupted — saving final checkpoint") @@ -529,6 +636,8 @@ def checkpoint(step, epoch): global_step = raw_step + step_offset if raw_step > 0 and global_step != last_saved_step: checkpoint(global_step, epoch) + if raw_step > 0: + run_demos(global_step) # final render tele.close() if t5_cache is not None: print(f"t5 cache: {t5_stats['hits']} hit(s) / " diff --git a/optimized/mlx/scripts/sa3_gradio.py b/optimized/mlx/scripts/sa3_gradio.py index 3e5259f..6ff207c 100644 --- a/optimized/mlx/scripts/sa3_gradio.py +++ b/optimized/mlx/scripts/sa3_gradio.py @@ -779,7 +779,8 @@ def render_queue_status(entry=None, generating=False): # ── Gradio UI ────────────────────────────────────────────────────────────── def build_ui(initial_dit: str, initial_decoder: str, *, share: bool, - default_seconds: float, default_steps: int): + default_seconds: float, default_steps: int, + preload_loras: tuple = (), server_port: int | None = None): import gradio as gr import random as _random @@ -1103,19 +1104,23 @@ def pregen(dit_name, decoder_name, prompt, negative_prompt, _dd0 = _lora_dd_choices(initial_dit) lora_inputs, lora_rows, lora_rm_btns = [], [], [] lora_dds, lora_files, lora_srows = [], [], [] + # --lora on the CLI preloads slot(s) 0..N: the trained ckpt + # is active on launch (underfit's "load the LoRA by default"). for _i in range(1, 4): - with gr.Group(visible=False) as _grp: + _prel = preload_loras[_i - 1] if _i - 1 < len(preload_loras) else None + with gr.Group(visible=_prel is not None) as _grp: with gr.Row(equal_height=True): _lf = gr.File(label=f"LoRA {_i} (.safetensors)", file_types=[".safetensors"], - type="filepath", scale=1, height=88) + type="filepath", scale=1, height=88, + value=_prel) _ld = gr.Dropdown( label=f"…or pick from loras/{LORA_DIR_NAMES[initial_dit]}/", choices=_dd0, value=_DD_NONE, scale=1, visible=len(_dd0) > 1) _rm = gr.Button("✕", size="sm", scale=0, min_width=36) - with gr.Row(visible=False) as _srow: + with gr.Row(visible=_prel is not None) as _srow: _ls = gr.Slider(label="strength", minimum=0.0, maximum=10.0, value=1.0, step=0.1, scale=2, min_width=110) @@ -1133,7 +1138,7 @@ def pregen(dit_name, decoder_name, prompt, negative_prompt, lora_srows.append(_srow) lora_inputs += [_lf, _ld, _ls, _ln, _lx] lora_add_btn = gr.Button("+ Add LoRA", size="sm") - lora_vis = gr.State([False, False, False]) + lora_vis = gr.State([k < len(preload_loras) for k in range(3)]) # Per-model LoRA memory: switching DiT saves the outgoing # model's slots and restores the incoming model's. lora_mem = gr.State({}) @@ -1324,6 +1329,7 @@ def wrapped(*args): ) demo.queue(max_size=16).launch(share=share, server_name="0.0.0.0", + server_port=server_port, allowed_paths=[str(OUTPUT_DIR)], prevent_thread_lock=False, show_error=True) @@ -1340,17 +1346,27 @@ def main(): ap.add_argument("--default-steps", type=int, default=8) ap.add_argument("--share", action=argparse.BooleanOptionalAction, default=True, help="Create a public gradio.live URL (default on)") + ap.add_argument("--lora", action="append", default=None, metavar="PATH", + help="Preload a LoRA .safetensors into a slot (active on " + "launch). Repeatable for up to 3 slots. underfit uses " + "this to open the gradio with a trained checkpoint loaded.") + ap.add_argument("--port", type=int, default=None, + help="Server port (default: gradio picks 7860+).") args = ap.parse_args() if args.decoder is None: args.decoder = DEFAULT_DECODERS[args.dit] + preload = tuple(args.lora or ())[:3] print(f"\n━━━ SA3 MLX — gradio ━━━") print(f" initial dit: {args.dit}") print(f" initial decoder: {args.decoder}") print(f" models: {', '.join(DIT_CHOICES.keys())} (runtime-switchable)") + if preload: + print(f" preloaded LoRA: {', '.join(preload)}") build_ui(args.dit, args.decoder, share=args.share, - default_seconds=args.default_seconds, default_steps=args.default_steps) + default_seconds=args.default_seconds, default_steps=args.default_steps, + preload_loras=preload, server_port=args.port) if __name__ == "__main__": From 2aac00a7032045902ecc39ac58b0da773ff4b2aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 12:47:05 -0400 Subject: [PATCH 14/28] mlx trainer: full default-config support (effective-length, optimizer betas/wd, InverseLR) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The underfit SA3 training templates set conventions the trainer didn't honor: - use_effective_length_for_schedule=True: shift timesteps by the effective latent length ceil(int(seconds_total*44100)/4096) per sample, not the crop length (--use-effective-length). - AdamW betas [0.9,0.95] + weight_decay 0.01 (--beta1/--beta2/--eps/--weight-decay; MLX AdamW's decoupled wd matches torch AdamW). - InverseLR scheduler with warmup (--lr-scheduler inverse + --inv-gamma/--lr-power/ --lr-warmup/--lr-final): lr(step)=(1-warmup^(step+1))*max(final,lr*(1+step/inv_gamma)^-power), stepped per optimizer step (last_epoch=raw_step, fresh on resume like torch). At the template's warmup=0.995 the step-0 LR is ~5e-7, not the nominal 1e-4 — omitting it diverges immediately. Set via optimizer.learning_rate each step; the compiled step reads it through inputs=state (verified: exact InverseLR values, no recompile). --- optimized/mlx/scripts/lora_train_mlx.py | 66 ++++++++++++++++++++++--- 1 file changed, 60 insertions(+), 6 deletions(-) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index cb76bce..25548c8 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -187,6 +187,32 @@ def parse_args(): ap.add_argument("--dist-shift", default="full", choices=["none", "full", "flux", "logsnr"]) ap.add_argument("--cfg-dropout-prob", type=float, default=0.1) + ap.add_argument("--use-effective-length", action=argparse.BooleanOptionalAction, + default=False, + help="Shift timesteps by the EFFECTIVE latent length " + "ceil(int(seconds_total*44100)/4096) per sample instead " + "of the crop length (underfit model flag " + "use_effective_length_for_schedule; True in the SA3 " + "templates). Only affects the 'full'/'flux'/'logsnr' shift.") + # optimizer (underfit reads these from optimizer_configs; torch AdamW defaults) + ap.add_argument("--beta1", type=float, default=0.9) + ap.add_argument("--beta2", type=float, default=0.999, + help="AdamW β2 (torch default 0.999; SA3 templates use 0.95)") + ap.add_argument("--eps", type=float, default=1e-8) + ap.add_argument("--weight-decay", type=float, default=0.0, + help="AdamW decoupled weight decay (torch default 0.0; SA3 " + "templates use 0.01)") + # LR schedule (underfit's optimizer_configs.diffusion.scheduler; InverseLR + # with warmup — at step 0 the SA3 template LR is ~5e-7, not the nominal lr) + ap.add_argument("--lr-scheduler", default="none", choices=["none", "inverse"], + help="'inverse' = SAT-dev InverseLR: " + "lr(step) = (1-warmup^(step+1)) * " + "max(final, lr*(1+step/inv_gamma)^-power)") + ap.add_argument("--inv-gamma", type=float, default=1e6) + ap.add_argument("--lr-power", type=float, default=0.5) + ap.add_argument("--lr-warmup", type=float, default=0.0, + help="InverseLR exponential-warmup base in [0,1) (template 0.995)") + ap.add_argument("--lr-final", type=float, default=0.0) # performance ap.add_argument("--compile", action=argparse.BooleanOptionalAction, default=True, help="mx.compile the train step (loss+grad+optimizer update; " @@ -387,9 +413,20 @@ def main(): max_new_steps = max(0, args.max_steps - step_offset) - # ── optimizer: AdamW, torch defaults, single group, no schedule ────────── - optimizer = optim.AdamW(learning_rate=args.lr, betas=[0.9, 0.999], - eps=1e-8, weight_decay=0.0) + # ── optimizer: AdamW, single group; betas/eps/wd from config (torch + # AdamW's decoupled weight decay — matches MLX's AdamW update exactly) ───── + def scheduled_lr(step): + """SAT-dev InverseLR at 0-based step (== torch last_epoch on a fresh run, + whose scheduler restarts even on resume). --lr-scheduler none → flat.""" + if args.lr_scheduler != "inverse": + return args.lr + warmup_factor = 1.0 - args.lr_warmup ** (step + 1) + lr_mult = (1.0 + step / args.inv_gamma) ** (-args.lr_power) + return warmup_factor * max(args.lr_final, args.lr * lr_mult) + + optimizer = optim.AdamW(learning_rate=scheduled_lr(0), + betas=[args.beta1, args.beta2], + eps=args.eps, weight_decay=args.weight_decay) # ── training-time local conditioning (underfit/upstream convention) ────── # The SA3 models are the diffusion_cond_inpaint variant: the trainer feeds @@ -582,12 +619,21 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): seconds = mx.array(np.asarray(batch["seconds_total"], dtype=np.float32)) - # timesteps: sampler + model-config distribution shift + # timesteps: sampler + model-config distribution shift. The shift + # sequence length is the EFFECTIVE latent length per sample when + # --use-effective-length (underfit use_effective_length_for_schedule, + # True in the SA3 templates), else the crop length. t_np = sample_training_timesteps( args.timestep_sampler, latents.shape[0], rng=np_rng) if args.dist_shift != "none": + if args.use_effective_length: + seq_len = np.asarray( + [int(math.ceil(int(float(s) * 44100) / 4096)) + for s in batch["seconds_total"]]) + else: + seq_len = latents.shape[-1] t_np = shift_training_timesteps( - t_np, latents.shape[-1], shift_type=args.dist_shift, + t_np, seq_len, shift_type=args.dist_shift, options=DIST_SHIFT_DEFAULT["options"] if args.dist_shift == "full" else None) timesteps = mx.array(t_np) @@ -606,6 +652,14 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): seconds_in = mx.stop_gradient(secs_embedder( [float(s) for s in batch["seconds_total"]])) + # InverseLR (or flat): the scheduler steps per optimizer step, + # last_epoch = raw_step (0-based new-step index), matching torch's + # fresh-on-resume scheduler. Setting learning_rate updates the leaf + # in optimizer.state, which the compiled step reads via inputs=state. + lr_now = scheduled_lr(raw_step) + if args.lr_scheduler != "none": + optimizer.learning_rate = lr_now + t_step = time.time() loss = train_step(latents, timesteps, loss_mask, prompt_cond, seconds_in, drop) @@ -620,7 +674,7 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): tele.flush() rate = raw_step / max(time.time() - t_start, 1e-9) print(f"step {global_step} train/loss {loss_v:.6f} " - f"train/lr {args.lr:.3e} epoch {epoch} " + f"train/lr {lr_now:.3e} epoch {epoch} " f"({rate:.2f} it/s, {step_ms:.0f} ms)", flush=True) if args.checkpoint_every > 0 and \ From 35c7dee9f59f8ccb0f1fbe934ff6b0d85cb6be32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 13:06:54 -0400 Subject: [PATCH 15/28] =?UTF-8?q?mlx:=20=C2=A79=20forward-parity=20harness?= =?UTF-8?q?=20(torch=20MPS=20vs=20MLX,=20bug-localizing)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reusable controlled-forward comparison: parity_forward_torch.py (sa3 venv) builds the fp32 base model, injects shared latents/noise/timesteps, replicates underfit's training-step forward, saves prediction/loss/cross/global + pretransform.scale; parity_forward_mlx.py (MLX venv) runs the same forward and reports three PSNRs that localize any mismatch — DiT-only (inject torch's cross/global → isolates the forward), conditioning (T5+seconds), end-to-end — plus asserts pretransform.scale==1.0. Fresh sm-music result: scale=1.0 (verified against the real model, not just assumed), DiT-only 79.0 dB, conditioning 88.5/77.6 dB, end-to-end 79.0 dB, loss Δ 6.1e-4. A convention bug would read ≪40 dB; ~79 dB is the cross-framework fp32 floor. No bug — the MLX and torch(MPS) training forwards agree to the floating-point limit. --- optimized/mlx/TRAINING_CONVENTIONS.md | 12 +++ optimized/mlx/scripts/parity_forward_mlx.py | 93 ++++++++++++++++ optimized/mlx/scripts/parity_forward_torch.py | 100 ++++++++++++++++++ 3 files changed, 205 insertions(+) create mode 100644 optimized/mlx/scripts/parity_forward_mlx.py create mode 100644 optimized/mlx/scripts/parity_forward_torch.py diff --git a/optimized/mlx/TRAINING_CONVENTIONS.md b/optimized/mlx/TRAINING_CONVENTIONS.md index 7ad02e7..c56a793 100644 --- a/optimized/mlx/TRAINING_CONVENTIONS.md +++ b/optimized/mlx/TRAINING_CONVENTIONS.md @@ -156,6 +156,18 @@ To build (roughly in order): ## 9. Training-forward conventions discovered by parity testing (2026-07-14) +**Harness** (reusable, 2026-07-15): `scripts/parity_forward_torch.py` (sa3 venv — +builds the fp32 base model, injects shared latents/noise/timesteps, saves +prediction/loss/cross/global/scale) + `scripts/parity_forward_mlx.py` (MLX venv — +same forward, reports three bug-localizing PSNRs). Fresh run (sm-music base, B=3, +t={0.25,0.5,0.75}): **pretransform.scale=1.0 asserted against the real model** +(confirms the softnorm no-scale convention — a prime hidden-bug candidate), +DiT-only forward (torch cross/global injected → isolates the DiT) **79.0 dB**, +conditioning cross/global **88.5/77.6 dB**, end-to-end **79.0 dB**, loss torch +4.716201 vs MLX 4.716814 (**Δ 6.1e-4**). A convention bug (RoPE/attn-mask/adaLN/ +local_add_cond layout/scale) reads ≪40 dB; ~79 dB is the cross-framework fp32 +reduction-order floor through 20 layers. Confirms parity, no bug. + Verified by a controlled forward (identical latents crop, numpy-seeded noise, t=0.5, fixed prompt, fp32) through underfit's torch loop on MPS vs the MLX trainer — final agreement **loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB** diff --git a/optimized/mlx/scripts/parity_forward_mlx.py b/optimized/mlx/scripts/parity_forward_mlx.py new file mode 100644 index 0000000..5ff4dc2 --- /dev/null +++ b/optimized/mlx/scripts/parity_forward_mlx.py @@ -0,0 +1,93 @@ +"""§9 forward-parity harness — MLX side + diff. + +Reads shared_inputs.npz + torch_out.npz (from parity_forward_torch.py) and runs +the SAME controlled forward through the MLX base DiT in fp32, then reports three +PSNRs that LOCALIZE any mismatch: + + (1) DiT-only — feed torch's cross/global + the shared noised into the MLX DiT. + Isolates the DiT forward convention (weights are the same fp16 source→fp32, + so ~79 dB is the cross-framework fp32 floor; a convention bug reads ≪40 dB). + (2) conditioning — MLX-native cross/global vs torch's (T5 + seconds; fp16-T5 bound). + (3) end-to-end — full MLX-native prediction + loss vs torch. + +Also asserts the torch model's pretransform.scale == 1.0 (the softnorm no-scale +convention the MLX trainer relies on). Run in the MLX venv: + + python parity_forward_mlx.py --base-npz --workdir +""" +import argparse +import importlib +import sys + +import numpy as np + + +def psnr(a, b): + a = np.asarray(a, np.float64); b = np.asarray(b, np.float64) + mse = np.mean((a - b) ** 2) + if mse == 0: + return float("inf") + return 20 * np.log10(max(np.abs(a).max(), np.abs(b).max()) / np.sqrt(mse)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--base-npz", required=True, help="MLX base DiT weights (dit_-base_f16.npz)") + ap.add_argument("--workdir", required=True) + ap.add_argument("--dit-module", default="models.defs.dit_mlx") + ap.add_argument("--prompt", default="techno") + ap.add_argument("--seconds", type=float, default=30.0) + args = ap.parse_args() + + import mlx.core as mx + from sa3_mlx import T5GEMMA_NPZ_REL + from weights import ensure_local + from models.defs.sa3_pipeline import apply_prompt_padding, load_conditioner_from_npz + from models.defs.t5gemma_mlx import T5Gemma + + tt = np.load(f"{args.workdir}/torch_out.npz") + inp = np.load(f"{args.workdir}/shared_inputs.npz", allow_pickle=True) + assert abs(float(tt["scale"]) - 1.0) < 1e-9, \ + f"pretransform.scale={float(tt['scale'])} != 1.0 — the MLX trainer's no-scale " \ + "softnorm convention is WRONG for this model (real bug)." + + latents = mx.array(inp["latents"]).astype(mx.float32) + noise = mx.array(inp["noise"]).astype(mx.float32) + t = mx.array(inp["t"]).astype(mx.float32) + B, C, T = latents.shape + dit = importlib.import_module(args.dit_module).load_dit( + args.base_npz, T_lat=T, dtype=mx.float32, compile_=False) + + noised = latents * (1.0 - t)[:, None, None] + noise * t[:, None, None] + target = noise - latents + lac = mx.concatenate([mx.ones((B, 1, T)), mx.zeros((B, C, T))], axis=1).transpose(0, 2, 1).astype(mx.float32) + + # (1) DiT-only: inject torch's cross/global + pred1 = dit(noised, t, mx.array(tt["cross"]).astype(mx.float32), + mx.array(tt["gcond"]).astype(mx.float32), local_add_cond=lac) + mx.eval(pred1) + + # (2) MLX-native conditioning + t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL))) + padding_emb, secs = load_conditioner_from_npz(args.base_npz, prefix="cond.") + embeds, mask = t5.encode([args.prompt] * B, max_len=256) + prompt_cond = apply_prompt_padding(embeds.astype(mx.float32), mask, padding_emb.astype(mx.float32)) + sec_tok = mx.stop_gradient(secs([args.seconds] * B).astype(mx.float32)) + cross_m = mx.concatenate([prompt_cond, sec_tok], axis=1) + gcond_m = sec_tok[:, 0, :] + mx.eval(cross_m, gcond_m) + + # (3) end-to-end + pred3 = dit(noised, t, cross_m, gcond_m, local_add_cond=lac) + mx.eval(pred3) + loss_m = float(mx.mean((pred3 - target) ** 2)) + + print(f"noised sanity PSNR = {psnr(np.array(noised), tt['noised']):.1f} dB (inf = bit-identical)") + print(f"(1) DiT-only forward = {psnr(np.array(pred1), tt['prediction']):6.2f} dB [conv bug ≪40; fp32 floor ~79]") + print(f"(2) conditioning = cross {psnr(np.array(cross_m), tt['cross']):.2f} dB / global {psnr(np.array(gcond_m), tt['gcond']):.2f} dB") + print(f"(3) end-to-end pred = {psnr(np.array(pred3), tt['prediction']):6.2f} dB") + print(f"loss: torch {float(tt['loss']):.6f} MLX {loss_m:.6f} |Δ| {abs(float(tt['loss'])-loss_m):.2e}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/optimized/mlx/scripts/parity_forward_torch.py b/optimized/mlx/scripts/parity_forward_torch.py new file mode 100644 index 0000000..d2c9462 --- /dev/null +++ b/optimized/mlx/scripts/parity_forward_torch.py @@ -0,0 +1,100 @@ +"""§9 forward-parity harness — torch (MPS) side. + +Controlled single forward of the SA3 base model, replicating underfit's +training-step forward (underfit/training/loop.py) on INJECTED inputs so the +result is deterministic and diffable against the MLX trainer's forward +(parity_forward_mlx.py). fp32, cfg_dropout=0. + +Run in an env with `stable_audio_3` importable (the sa3 repo venv): + + python parity_forward_torch.py \ + --base-dir # dir with model_config.json + model.safetensors + --latents # a [256,T] pre-encoded latent + --workdir # writes shared_inputs.npz + torch_out.npz + +Then run parity_forward_mlx.py (MLX venv) with the same --workdir. + +The harness localizes any mismatch: it saves cross/global/noised/prediction/ +loss + the model's pretransform.scale (the MLX side asserts scale==1.0 for the +softnorm no-scale convention; a non-1.0 value would be a real bug). Weights on +both sides come from the same fp16 source upcast to fp32, so agreement is the +cross-framework fp32 floor (~79 dB); a convention bug reads far lower. +""" +import argparse +import json +import os + +import numpy as np +import torch +from safetensors.torch import load_file +from stable_audio_3.factory import create_diffusion_cond_from_config +from stable_audio_3.loading_utils import remap_state_dict_keys + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--base-dir", required=True, + help="Base checkpoint dir (model_config.json + model.safetensors)") + ap.add_argument("--latents", required=True, help="A [256,T] pre-encoded latent .npy") + ap.add_argument("--workdir", required=True) + ap.add_argument("--crop", type=int, default=256) + ap.add_argument("--prompt", default="techno") + ap.add_argument("--seconds", type=float, default=30.0) + ap.add_argument("--noise-seed", type=int, default=12345) + ap.add_argument("--device", default=None) + args = ap.parse_args() + os.makedirs(args.workdir, exist_ok=True) + device = args.device or ("mps" if torch.backends.mps.is_available() else "cpu") + + # Shared injected inputs (numpy — both engines consume these verbatim). + clip = np.load(args.latents)[:, : args.crop] + lat = np.stack([clip, clip, clip], 0).astype(np.float32) # B=3 + noise = np.random.default_rng(args.noise_seed).standard_normal(lat.shape).astype(np.float32) + t = np.array([0.25, 0.5, 0.75], dtype=np.float32) + np.savez(f"{args.workdir}/shared_inputs.npz", latents=lat, noise=noise, t=t, + prompt=args.prompt, seconds_total=args.seconds) + + cfg = json.load(open(f"{args.base_dir}/model_config.json")) + print("building base model (fp32)...", flush=True) + model = create_diffusion_cond_from_config(cfg) + sd = load_file(f"{args.base_dir}/model.safetensors") + sd = remap_state_dict_keys(sd, model.state_dict()) + missing, unexpected = model.load_state_dict(sd, strict=False) + print(f"loaded: {len(missing)} missing / {len(unexpected)} unexpected", flush=True) + model = model.to(device).eval().requires_grad_(False) + + scale = float(getattr(model.pretransform, "scale", 1.0)) if model.pretransform is not None else 1.0 + print(f"pretransform_scale = {scale} diffusion_objective = {model.diffusion_objective}", flush=True) + + latents = torch.tensor(lat).to(device) + noise_t = torch.tensor(noise).to(device) + tt = torch.tensor(t).to(device) + B, C, T = latents.shape + diffusion_input = latents / scale if scale != 1.0 else latents + noised = diffusion_input * (1 - tt)[:, None, None] + noise_t * tt[:, None, None] + target = noise_t - diffusion_input + + metadata = [{"prompt": args.prompt, "seconds_total": args.seconds, "seconds_start": 0.0} + for _ in range(B)] + conditioning = model.conditioner(metadata, device) + conditioning["inpaint_mask"] = [torch.ones(B, 1, T, device=device)] + conditioning["inpaint_masked_input"] = [torch.zeros(B, C, T, device=device)] + + with torch.no_grad(): + output = model(noised, tt, cond=conditioning, cfg_dropout_prob=0.0, + padding_mask=torch.ones(B, T, dtype=torch.bool, device=device)) + ci = model.get_conditioning_inputs(conditioning) + + loss = ((output.float() - target.float()) ** 2).mean() # full mask → signal-only == mean + print(f"loss = {float(loss):.8f}", flush=True) + np.savez(f"{args.workdir}/torch_out.npz", + prediction=output.float().cpu().numpy(), loss=float(loss), + cross=ci["cross_attn_cond"].float().cpu().numpy(), + gcond=ci["global_cond"].float().cpu().numpy(), + noised=noised.float().cpu().numpy(), target=target.float().cpu().numpy(), + scale=scale) + print(f"wrote {args.workdir}/torch_out.npz", flush=True) + + +if __name__ == "__main__": + main() From a80e1c7b63d4eb82507c58172360ed7338714119 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 13:21:20 -0400 Subject: [PATCH 16/28] =?UTF-8?q?mlx:=20=C2=A79=20harness=20also=20checks?= =?UTF-8?q?=20the=20backward=20(d=20loss/d=20noised)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the parity harness to compare the full DiT backward sweep cross-framework: torch saves d loss/d noised (autograd w.r.t. the injected noised input, the same backward every interior/adapter gradient is built from); the MLX side computes mx.grad of the same loss (torch cross/global injected to isolate the DiT backward) and reports its PSNR. Result: 77.5 dB — the same fp32 floor as the forward (79 dB), so forward AND backward agree to the floating-point limit. The local adapter-param VJP was separately verified within-framework (DoRA reformulation, <=2.4e-7). --- optimized/mlx/TRAINING_CONVENTIONS.md | 11 ++++++++--- optimized/mlx/scripts/parity_forward_mlx.py | 14 ++++++++++++++ optimized/mlx/scripts/parity_forward_torch.py | 19 ++++++++++++------- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/optimized/mlx/TRAINING_CONVENTIONS.md b/optimized/mlx/TRAINING_CONVENTIONS.md index c56a793..0c82d61 100644 --- a/optimized/mlx/TRAINING_CONVENTIONS.md +++ b/optimized/mlx/TRAINING_CONVENTIONS.md @@ -164,9 +164,14 @@ t={0.25,0.5,0.75}): **pretransform.scale=1.0 asserted against the real model** (confirms the softnorm no-scale convention — a prime hidden-bug candidate), DiT-only forward (torch cross/global injected → isolates the DiT) **79.0 dB**, conditioning cross/global **88.5/77.6 dB**, end-to-end **79.0 dB**, loss torch -4.716201 vs MLX 4.716814 (**Δ 6.1e-4**). A convention bug (RoPE/attn-mask/adaLN/ -local_add_cond layout/scale) reads ≪40 dB; ~79 dB is the cross-framework fp32 -reduction-order floor through 20 layers. Confirms parity, no bug. +4.716201 vs MLX 4.716814 (**Δ 6.1e-4**). **Backward also checked**: ∂loss/∂noised +through the full DiT (torch cross/global injected) agrees at **77.5 dB** — the same +floor as the forward, so the backward sweep every adapter gradient is built from +matches too (the local adapter-param VJP was separately verified within-framework to +≤2.4e-7 MLX / ≤4.4e-7 torch during the DoRA reformulation). A convention/autograd bug +(RoPE/attn-mask/adaLN/local_add_cond layout/scale) reads ≪40 dB; ~78 dB is the +cross-framework fp32 reduction-order floor through 20 layers. Confirms fwd+bwd parity, +no bug. Verified by a controlled forward (identical latents crop, numpy-seeded noise, t=0.5, fixed prompt, fp32) through underfit's torch loop on MPS vs the MLX diff --git a/optimized/mlx/scripts/parity_forward_mlx.py b/optimized/mlx/scripts/parity_forward_mlx.py index 5ff4dc2..53c5ef6 100644 --- a/optimized/mlx/scripts/parity_forward_mlx.py +++ b/optimized/mlx/scripts/parity_forward_mlx.py @@ -82,10 +82,24 @@ def main(): mx.eval(pred3) loss_m = float(mx.mean((pred3 - target) ** 2)) + # (4) BACKWARD: ∂loss/∂noised through the DiT (with torch's cross/global so + # the backward is compared on identical inputs — isolates the DiT backward). + cross_t = mx.array(tt["cross"]).astype(mx.float32) + gcond_t = mx.array(tt["gcond"]).astype(mx.float32) + + def loss_of_noised(nz): + pred = dit(nz, t, cross_t, gcond_t, local_add_cond=lac) + return mx.mean((pred - target) ** 2) + + grad_m = mx.grad(loss_of_noised)(noised) + mx.eval(grad_m) + p_grad = psnr(np.array(grad_m), tt["grad_noised"]) if "grad_noised" in tt else float("nan") + print(f"noised sanity PSNR = {psnr(np.array(noised), tt['noised']):.1f} dB (inf = bit-identical)") print(f"(1) DiT-only forward = {psnr(np.array(pred1), tt['prediction']):6.2f} dB [conv bug ≪40; fp32 floor ~79]") print(f"(2) conditioning = cross {psnr(np.array(cross_m), tt['cross']):.2f} dB / global {psnr(np.array(gcond_m), tt['gcond']):.2f} dB") print(f"(3) end-to-end pred = {psnr(np.array(pred3), tt['prediction']):6.2f} dB") + print(f"(4) BACKWARD d loss/d noised PSNR = {p_grad:6.2f} dB [full DiT backward; conv bug ≪40]") print(f"loss: torch {float(tt['loss']):.6f} MLX {loss_m:.6f} |Δ| {abs(float(tt['loss'])-loss_m):.2e}") diff --git a/optimized/mlx/scripts/parity_forward_torch.py b/optimized/mlx/scripts/parity_forward_torch.py index d2c9462..c7439ea 100644 --- a/optimized/mlx/scripts/parity_forward_torch.py +++ b/optimized/mlx/scripts/parity_forward_torch.py @@ -79,20 +79,25 @@ def main(): conditioning = model.conditioner(metadata, device) conditioning["inpaint_mask"] = [torch.ones(B, 1, T, device=device)] conditioning["inpaint_masked_input"] = [torch.zeros(B, C, T, device=device)] - with torch.no_grad(): - output = model(noised, tt, cond=conditioning, cfg_dropout_prob=0.0, - padding_mask=torch.ones(B, T, dtype=torch.bool, device=device)) ci = model.get_conditioning_inputs(conditioning) + # BACKWARD: gradient of the loss w.r.t. the noised input — the full DiT + # backward sweep (every interior gradient the adapter grads are built from). + noised.requires_grad_(True) + output = model(noised, tt, cond=conditioning, cfg_dropout_prob=0.0, + padding_mask=torch.ones(B, T, dtype=torch.bool, device=device)) loss = ((output.float() - target.float()) ** 2).mean() # full mask → signal-only == mean - print(f"loss = {float(loss):.8f}", flush=True) + loss.backward() + grad_noised = noised.grad + print(f"loss = {float(loss):.8f} |grad_noised| = {float(grad_noised.abs().mean()):.6e}", flush=True) np.savez(f"{args.workdir}/torch_out.npz", - prediction=output.float().cpu().numpy(), loss=float(loss), + prediction=output.detach().float().cpu().numpy(), loss=float(loss), cross=ci["cross_attn_cond"].float().cpu().numpy(), gcond=ci["global_cond"].float().cpu().numpy(), - noised=noised.float().cpu().numpy(), target=target.float().cpu().numpy(), - scale=scale) + noised=noised.detach().float().cpu().numpy(), + target=target.float().cpu().numpy(), + grad_noised=grad_noised.float().cpu().numpy(), scale=scale) print(f"wrote {args.workdir}/torch_out.npz", flush=True) From 42ed8f6d71c398768833dc56a7a45ce76045d313 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 13:35:06 -0400 Subject: [PATCH 17/28] ci: MLX test job on Apple Silicon (macos-latest / arm64) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GitHub-hosted macos-latest runners are Apple Silicon, which MLX requires. Runs the weight-free adapter-math suite (all 9 LoRA/DoRA/BoRA types, scripts/ test_lora_merge.py) + a core-module import smoke + lora_train_mlx.py/ pre_encode_mlx.py --help. Scoped via paths filter to optimized/mlx changes so it doesn't spend macOS runner minutes on unrelated PRs. (test_all_configs.py needs the shipped npz weights and isn't CI-runnable — left out.) --- .github/workflows/mlx.yml | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 .github/workflows/mlx.yml diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml new file mode 100644 index 0000000..02d7135 --- /dev/null +++ b/.github/workflows/mlx.yml @@ -0,0 +1,32 @@ +name: MLX (Apple Silicon) + +# GitHub-hosted `macos-latest` runners are Apple Silicon (arm64), which the MLX +# runtime requires. Scoped to optimized/mlx changes so it doesn't spend macOS +# runner minutes on unrelated PRs. +on: + pull_request: + paths: + - "optimized/mlx/**" + - ".github/workflows/mlx.yml" + +jobs: + test: + runs-on: macos-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install MLX deps + run: | + python -m pip install --upgrade pip + pip install -r optimized/mlx/requirements.txt pytest + - name: Adapter math — all 9 LoRA/DoRA/BoRA types (weight-free) + run: pytest optimized/mlx/scripts/test_lora_merge.py -q + - name: Import + CLI smoke (no weights downloaded) + working-directory: optimized/mlx + run: | + python -c "import mlx.core; from models.defs import lora, training, demo_mlx, latent_dataset, lora_merge, sa3_pipeline; print('core imports OK')" + python scripts/lora_train_mlx.py --help > /dev/null + python scripts/pre_encode_mlx.py --help > /dev/null + echo "CLI smoke OK" From 98d4d2fb93b79054d4207bad71f91d9069b39d59 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 13:45:01 -0400 Subject: [PATCH 18/28] =?UTF-8?q?docs:=20README=20LoRA=20training=20sectio?= =?UTF-8?q?n=20(pre-encode=20=E2=86=92=20train=20=E2=86=92=20generate)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old 'LoRA training primitives' section predated the trainer and said there was 'no dataset or training CLI' — now stale. Replace it with the real end-to-end workflow: pre_encode_mlx.py → lora_train_mlx.py (train on the BASE ckpt via --dit-weights, underfit-default flags + the full-template flag set, demos) → generate with --lora. Keep a short 'build your own loop' note for the primitives, and list the training files + TRAINING_CONVENTIONS.md in the Files tree. --- optimized/mlx/README.md | 122 ++++++++++++++++++++++++---------------- 1 file changed, 73 insertions(+), 49 deletions(-) diff --git a/optimized/mlx/README.md b/optimized/mlx/README.md index 7bcc9b9..23529eb 100644 --- a/optimized/mlx/README.md +++ b/optimized/mlx/README.md @@ -191,57 +191,69 @@ Sample run on **M4 Pro / 48 GB**: └───────────┴─────────┴─────────┴───────────┴───────────┴────────────┘ ``` -## LoRA training primitives - -The optimized runtime includes MLX counterparts to the adapter family already -supported by the official PyTorch implementation: standard LoRA, DoRA, BoRA, -and LoRA-XS variants. - -```python -from models.defs.lora import ( - apply_lora_checkpoint, - inject_trainable_lora, - save_lora_checkpoint, -) -from models.defs.training import ( - rectified_flow_loss, - sample_training_timesteps, - shift_training_timesteps, -) -from models.defs.audio_encoding import encode_audio +## LoRA training + +Finetune SA3 on your own audio, entirely in MLX (no PyTorch) — a full training +CLI that replicates [underfit](https://github.com/dada-bots/underfit)'s +conventions, defaults, and checkpoint format (so it can run as underfit's +Apple-Silicon backend). Three steps: **pre-encode → train → generate**. See +`TRAINING_CONVENTIONS.md` for the complete convention inventory and the +torch(MPS)-vs-MLX forward+backward parity results. + +**1. Pre-encode** your audio to SAME latents (once, offline — torch-free): + +```bash +uv run python scripts/pre_encode_mlx.py \ + --audio-dir ~/my-clips --output-dir ~/my-latents --codec same-s +``` + +Each file becomes `.npy` latents `[D, T]` + a `.json` sidecar +(duration, padding mask, tags) — the exact format underfit's dataset uses. +Use `same-s` for sm-music/sm-sfx, `same-l` for medium. + +**2. Train.** Train on the **BASE** checkpoint (`stabilityai/stable-audio-3-*-base`), +*not* the shipped ARC weights that inference uses — pass its MLX `.npz` via +`--dit-weights` (base-ckpt→npz conversion recipe in `TRAINING_CONVENTIONS.md` §9; +omitting it trains on ARC and prints a warning): + +```bash +uv run python scripts/lora_train_mlx.py \ + --dit sm-music --dit-weights models/mlx/dit_sm-music-base_f16.npz \ + --latents-dir ~/my-latents --lr 1e-4 --name my-lora \ + --adapter-type dora-rows --rank 16 --max-steps 2000 ``` -`inject_trainable_lora()` freezes the base model and leaves only adapter -parameters trainable through `mlx.nn.value_and_grad`. `save_lora_checkpoint()` -writes the same safetensors keys and `lora_config` metadata used by the -PyTorch trainer. `apply_lora_checkpoint()` loads that format through MLX -without adding PyTorch or safetensors as runtime dependencies. Application -materializes a fixed strength into the loaded model's in-memory weights; it -does not modify the base checkpoint on disk. Reload the base model before -applying a different strength rather than applying repeatedly to the same -instance. - -`encode_audio()` provides the waveform-to-latent bridge used by training. It -accepts a batch of already-decoded 44.1 kHz stereo waveforms, applies SAME's -patched pretransform, pads to the encoder's required alignment, optionally -encodes long clips in overlapping chunks, and returns both latents and a -ceiling-scaled padding mask: - -```python -encoded = encode_audio( - same_l_encoder, - audio, # [batch, 2, samples] - valid_sample_lengths=[samples_a, samples_b], - pad_modulo=16, # 32 for SAME-S - chunked=True, -) -latents = encoded.latents -loss_mask = encoded.padding_mask +Defaults mirror underfit (dora-rows / rank 16, AdamW, uniform sampler + the +"full" distribution shift, CFG-dropout 0.1, signal-only masked rectified-flow +loss, checkpoint every 1000 steps). To match the full dashboard template config, +add `--timestep-sampler trunc_logit_normal --use-effective-length --beta2 0.95 +--weight-decay 0.01 --lr-scheduler inverse --lr-warmup 0.995`. Checkpoints land at +`output/runs///checkpoints/-step=S-epoch=E.safetensors` with +`lora_config` metadata; resume with `--lora-ckpt-path `. + +Optional training-time demos (RF-Euler inference → mp3, underfit's on-disk +format): `--demo-every 1000 --demo-config demos.json`, where `demos.json` is a +list of `{"prompt", "cfg", "seed", "steps", "lora_strength"?, "lora_interval_max"?}`. + +**3. Generate** with your adapter — the checkpoint applies to the shipped ARC +model at inference via `--lora` (see "Apply a LoRA finetune" above; also +`scripts/sa3_gradio.py --lora ` to open the web UI with it preloaded): + +```bash +./sa3 --dit sm-music --decoder same-s --prompt "..." --out finetuned.wav \ + --lora "output/runs/my-lora/"*"/checkpoints/my-lora-step=2000-epoch="*.safetensors ``` -These are model-level primitives rather than a dataset or training CLI. -Callers remain responsible for file decoding and resampling, dataset -discovery and persistence, conditioning, optimization, and checkpoint cadence. +### Building your own loop + +The CLI is assembled from reusable pieces: `inject_trainable_lora` / +`save_lora_checkpoint` / `apply_lora_checkpoint` (`models/defs/lora.py` — all 9 +LoRA/DoRA/BoRA/-XS types with the no-weight-materialization forward), the RF +loss + samplers + distribution shift (`models/defs/training.py`), the +pre-encoded dataset + prompt templates (`models/defs/latent_dataset.py`), and +`encode_audio` (`models/defs/audio_encoding.py`). Checkpoints use the same +safetensors keys + `lora_config` metadata as the PyTorch trainer, so adapters +are interchangeable in either direction. ## Flag reference @@ -273,6 +285,7 @@ sa3_mlx/ ├── sa3 ← shell wrapper (use this) ├── install.sh ← uv bootstrap (run once) ├── README.md +├── TRAINING_CONVENTIONS.md ← LoRA-training conventions + MPS-vs-MLX parity ├── requirements.txt ├── output/ ← default landing zone for generated WAVs ├── scripts/ @@ -281,7 +294,12 @@ sa3_mlx/ │ ├── examples.py ← shared examples block (--help + post-install) │ ├── install.py ← install.sh's Python half (bundle picker) │ ├── test_all_configs.py ← npz + CLI config sanity tests -│ └── benchmark.py ← wall-time + peak-RAM matrix across model × duration +│ ├── benchmark.py ← wall-time + peak-RAM matrix across model × duration +│ ├── lora_train_mlx.py ← LoRA training CLI (underfit conventions) +│ ├── pre_encode_mlx.py ← audio → SAME-latent pre-encode (torch-free) +│ ├── sa3_gradio.py ← web UI (invoked by ./sa3-gradio; --lora preload) +│ ├── test_lora_merge.py ← adapter-math tests (weight-free; run in CI) +│ └── parity_forward_{torch,mlx}.py ← torch(MPS)-vs-MLX fwd+bwd parity harness └── models/ ├── defs/ │ ├── sa3_pipeline.py ← sampler + conditioner + unpatch @@ -289,7 +307,13 @@ sa3_mlx/ │ ├── dit_mlx.py ← small DiT (sm-music + sm-sfx) │ ├── dit_mlx_medium.py ← medium DiT (differential attention) │ ├── same_s_{encoder,decoder}.py ← small codec - │ └── same_l_{encoder,decoder}.py ← large codec + │ ├── same_l_{encoder,decoder}.py ← large codec + │ ├── lora.py ← trainable LoRA/DoRA/BoRA adapters (9 types) + │ ├── lora_merge.py ← inference-time LoRA merge + per-step gating + │ ├── training.py ← RF loss + timestep samplers + distribution shift + │ ├── latent_dataset.py ← pre-encoded dataset + underfit prompt templates + │ ├── audio_encoding.py ← waveform → SAME latents (pre-encode bridge) + │ └── demo_mlx.py ← training-time RF-Euler demos → mp3 └── mlx/ ← .npz weights (auto-downloaded; ~8.4 GB total) ├── t5gemma_f16.npz 541 MB text encoder + tokenizer ├── dit_sm-music_f16.npz 877 MB DiT + conditioner baked in From fb2d51ceb6a77964d37f0ed22d79332c217f3290 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 14:37:09 -0400 Subject: [PATCH 19/28] mlx: auto-download base-model npz for training + export_base_npz.py converter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Training uses the BASE checkpoint (not the shipped ARC weights). Ship the base weights via HF so training works out of the box: - weights.py: TRAINING_BASE manifest entries for dit_{sm-music,sm-sfx,medium}-base_f16.npz (MLX/… in stabilityai/stable-audio-3-optimized), added to FLAT_MANIFEST for lazy ensure_local download. - lora_train_mlx.py: default to the base npz (ensure_local auto-download) when --dit-weights is omitted — drop the old 'warn + train on ARC' fallback. --dit-weights still overrides with a custom path. - export_base_npz.py: torch-free converter (safetensors.numpy) from a *-base checkpoint → dit_-base_f16.npz (DiT + baked conditioner). Validated to reproduce the known-good sm-music/medium npz within fp16 (419/441 bit-identical, rest ≤1 fp16 ULP); used to produce the sm-sfx base npz. - README: training uses the auto-downloaded base npz (no flag needed). --- optimized/mlx/README.md | 12 ++-- optimized/mlx/scripts/export_base_npz.py | 90 ++++++++++++++++++++++++ optimized/mlx/scripts/lora_train_mlx.py | 20 +++--- optimized/mlx/scripts/weights.py | 11 +++ 4 files changed, 117 insertions(+), 16 deletions(-) create mode 100644 optimized/mlx/scripts/export_base_npz.py diff --git a/optimized/mlx/README.md b/optimized/mlx/README.md index 23529eb..f9c4a7a 100644 --- a/optimized/mlx/README.md +++ b/optimized/mlx/README.md @@ -211,15 +211,15 @@ Each file becomes `.npy` latents `[D, T]` + a `.json` sidecar (duration, padding mask, tags) — the exact format underfit's dataset uses. Use `same-s` for sm-music/sm-sfx, `same-l` for medium. -**2. Train.** Train on the **BASE** checkpoint (`stabilityai/stable-audio-3-*-base`), -*not* the shipped ARC weights that inference uses — pass its MLX `.npz` via -`--dit-weights` (base-ckpt→npz conversion recipe in `TRAINING_CONVENTIONS.md` §9; -omitting it trains on ARC and prints a warning): +**2. Train.** Training uses the **BASE** checkpoint (`stabilityai/stable-audio-3-*-base`, +rectified_flow), *not* the shipped ARC weights inference uses. The base-model npz is +**auto-downloaded** from HuggingFace on first run — no flag needed (override with +`--dit-weights `; regenerate one yourself from a `-base` repo with +`scripts/export_base_npz.py`): ```bash uv run python scripts/lora_train_mlx.py \ - --dit sm-music --dit-weights models/mlx/dit_sm-music-base_f16.npz \ - --latents-dir ~/my-latents --lr 1e-4 --name my-lora \ + --dit sm-music --latents-dir ~/my-latents --lr 1e-4 --name my-lora \ --adapter-type dora-rows --rank 16 --max-steps 2000 ``` diff --git a/optimized/mlx/scripts/export_base_npz.py b/optimized/mlx/scripts/export_base_npz.py new file mode 100644 index 0000000..f551fac --- /dev/null +++ b/optimized/mlx/scripts/export_base_npz.py @@ -0,0 +1,90 @@ +"""Convert a Stable Audio 3 *-base torch checkpoint → the MLX training npz. + +Torch-free (safetensors.numpy). Produces `dit_-base_f16.npz` — the DiT +weights + baked conditioner that `lora_train_mlx.py --dit-weights` expects. +LoRA training uses the BASE checkpoint (not the shipped ARC weights inference +uses); these npz are auto-downloaded by weights.py, or reproduced here from the +HF `stabilityai/stable-audio-3-*-base` repos. + + python scripts/export_base_npz.py \ + --base-dir # dir with model.safetensors + model_config.json + --output models/mlx/dit_-base_f16.npz + [--validate-against ] # assert bit-parity with a known-good npz + +Conversion (identical to the shipped inference-npz layout): strip the DiT +wrapper prefix `model.model.`, transpose the two Conv1d weights +[out,in,k]→[out,k,in], rename RMSNorm `.gamma`→`.weight`, remap +`.to_local_embed.{0,2}`→`.seq.{0,2}`, and fold the 3 conditioner tensors under +the `cond.` prefix. Same for small (sm-music/sm-sfx) and medium. +""" +import argparse + +import mlx.core as mx +import numpy as np +from safetensors.numpy import load_file + +_DIT_PREFIX = "model.model." +_CONV_WEIGHTS = ("preprocess_conv.weight", "postprocess_conv.weight") +_COND_MAP = { + "conditioner.conditioners.prompt.padding_embedding": "cond.padding_embedding", + "conditioner.conditioners.seconds_total.embedder.embedding.1.weight": "cond.seconds_total_weight", + "conditioner.conditioners.seconds_total.embedder.embedding.1.bias": "cond.seconds_total_bias", +} + + +def convert(base_dir): + sd = load_file(f"{base_dir}/model.safetensors") + out = {} + for k, arr in sd.items(): + if k in _COND_MAP: + out[_COND_MAP[k]] = mx.array(arr.astype(np.float32)) + continue + if not k.startswith(_DIT_PREFIX): + continue # pretransform / T5 backbone — not part of the DiT npz + sk = k[len(_DIT_PREFIX):] + a = arr.astype(np.float32) + if sk in _CONV_WEIGHTS: # Conv1d [out,in,k] -> MLX [out,k,in] + out[sk] = mx.array(a.transpose(0, 2, 1)) + elif sk.endswith(".gamma"): # RMSNorm .gamma -> .weight + out[sk[: -len(".gamma")] + ".weight"] = mx.array(a) + else: + sk = sk.replace(".to_local_embed.0.", ".to_local_embed.seq.0.") + sk = sk.replace(".to_local_embed.2.", ".to_local_embed.seq.2.") + out[sk] = mx.array(a) + return out + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--base-dir", required=True) + ap.add_argument("--output", required=True) + ap.add_argument("--validate-against", default=None, + help="Assert key-set + value parity with a known-good npz.") + args = ap.parse_args() + + wd = convert(args.base_dir) + print(f"converted {len(wd)} keys ({sum(1 for k in wd if k.startswith('cond'))} conditioner)") + + if args.validate_against: + ref = dict(mx.load(args.validate_against)) + extra, missing = set(wd) - set(ref), set(ref) - set(wd) + assert not extra and not missing, f"key mismatch: extra={extra}, missing={missing}" + worst = max( + float(mx.max(mx.abs(wd[k].astype(mx.float32) - ref[k].astype(mx.float32)))) + for k in ref + ) + print(f"validate: {len(ref)} keys match, max|Δ| = {worst:.2e}") + # Tolerance is fp16-scale: the source ckpt is fp32 and both npz are fp16, + # so a handful of the specially-transformed keys (conv transpose, the + # to_local_embed rename) can differ by ~1 fp16 ULP (~4e-3) depending on + # the rounding path. A real structural bug (wrong transpose/rename) gives + # an O(0.1-1) diff or a shape/key mismatch — both caught well below this. + assert worst < 1e-2, "values diverge from the reference npz (structural bug)" + + mx.savez(args.output, **{k: v.astype(mx.float16) for k, v in wd.items()}) + print(f"wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 25548c8..4405aae 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -353,19 +353,19 @@ def main(): if len(dataset) < args.batch_size else "")) # ── models ─────────────────────────────────────────────────────────────── + # LoRA training uses the BASE checkpoint (rectified_flow), NOT the shipped + # ARC weights inference uses. Default to the base-model npz (auto-downloaded + # from HF via weights.TRAINING_BASE); --dit-weights overrides with a path. t0 = time.time() + import importlib if args.dit_weights: - import importlib - mod = importlib.import_module(DIT_CHOICES[args.dit]["loader"]) - dit_model = mod.load_dit(args.dit_weights, T_lat=crop_len, - dtype=mx.float16, compile_=False) - cond_src = args.dit_weights + base_path = args.dit_weights else: - print("WARNING: training on the shipped (ARC/rf_denoiser) weights — " - "underfit convention is to train on the BASE checkpoint " - "(pass --dit-weights)") - dit_model, _ = load_dit(args.dit, T_lat=crop_len, dtype=mx.float16) - cond_src = str(ensure_local(DIT_CHOICES[args.dit]["ckpt"])) + base_path = str(ensure_local(f"models/mlx/dit_{args.dit}-base_f16.npz")) + mod = importlib.import_module(DIT_CHOICES[args.dit]["loader"]) + dit_model = mod.load_dit(base_path, T_lat=crop_len, dtype=mx.float16, + compile_=False) + cond_src = base_path print(f"DiT loaded ({time.time()-t0:.1f}s, fp16 base, T_lat={crop_len})") t5 = T5Gemma.from_npz(str(ensure_local(T5GEMMA_NPZ_REL))) # frozen padding_emb, secs_embedder = load_conditioner_from_npz(cond_src, prefix="cond.") diff --git a/optimized/mlx/scripts/weights.py b/optimized/mlx/scripts/weights.py index 996ac42..cda37a6 100644 --- a/optimized/mlx/scripts/weights.py +++ b/optimized/mlx/scripts/weights.py @@ -46,6 +46,15 @@ ("models/mlx/t5gemma_f16.npz", "MLX/t5gemma_f16.npz"), ] +# Base-model DiT weights — needed only for LoRA *training* (lora_train_mlx.py +# trains on the BASE checkpoint, not the ARC weights inference uses). Not part +# of any install bundle; auto-downloaded lazily by the trainer via ensure_local. +TRAINING_BASE: list[tuple[str, str]] = [ + ("models/mlx/dit_sm-music-base_f16.npz", "MLX/dit_sm-music-base_f16.npz"), + ("models/mlx/dit_sm-sfx-base_f16.npz", "MLX/dit_sm-sfx-base_f16.npz"), + ("models/mlx/dit_medium-base_f16.npz", "MLX/dit_medium-base_f16.npz"), +] + # Human-friendly bundle sizes (for the install prompt). BUNDLE_SIZES = { "medium": "5.9 GB (medium DiT + SAME-L codec)", @@ -61,6 +70,8 @@ FLAT_MANIFEST[_rel] = _hf for _rel, _hf in SHARED: FLAT_MANIFEST[_rel] = _hf +for _rel, _hf in TRAINING_BASE: + FLAT_MANIFEST[_rel] = _hf def _hf_token_configured() -> bool: From cf8194c70dcdbe14d4019c5a691c389394d58837 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 16:30:09 -0400 Subject: [PATCH 20/28] mlx pre-encode: print engine + device banner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Print "Engine: MLX · device: Device(gpu, 0)" at the top of the encode run so the dashboard log (and CLI) shows unambiguously that the MLX SAME encoder on the Metal GPU is running — not the torch (MPS) pre_encode path. --- optimized/mlx/scripts/pre_encode_mlx.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/optimized/mlx/scripts/pre_encode_mlx.py b/optimized/mlx/scripts/pre_encode_mlx.py index 9f429fc..efe3f98 100644 --- a/optimized/mlx/scripts/pre_encode_mlx.py +++ b/optimized/mlx/scripts/pre_encode_mlx.py @@ -396,6 +396,8 @@ def main(): sys.exit(1) max_samples = max_samples_for_duration(args.max_duration) + import mlx.core as mx + print(f"Engine: MLX · device: {mx.default_device()}") print(f"Codec: {args.codec} (pad_modulo={SAME_ENCODER_PAD_MODULO[args.codec]}, fp32)") print(f"Input: {audio_dir} ({len(files)} audio files)") print(f"Output: {Path(args.output_dir).expanduser().resolve()}") From ba8fed8df2624d4a3022c089455751e587508d6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 16:37:03 -0400 Subject: [PATCH 21/28] mlx trainer: print engine + device banner (Engine: MLX / device: ...) --- optimized/mlx/scripts/lora_train_mlx.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 4405aae..9379e20 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -338,6 +338,7 @@ def main(): session = uuid.uuid4().hex[:8] ckpt_dir = Path(args.save_dir) / args.name / session / "checkpoints" ckpt_dir.mkdir(parents=True, exist_ok=True) + print(f"Engine: MLX · device: {mx.default_device()}") print(f"run: {args.name} session {session}") print(f"checkpoints → {ckpt_dir}") From 78ddecae08e5e19ed5528aaba9b975030d2ad720 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 16:44:06 -0400 Subject: [PATCH 22/28] mlx pre-encode: --exclude-file support (skip relpaths, torch pre_encode parity) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add --exclude-file (text file of relpaths relative to --audio-dir, one per line) — same format/semantics as dataset_processing/pre_encode.py. Filters the discovered files in both main() (banner shows N excluded) and run(); the dashboard passes its exclude.txt through so per-file exclusions work on the MLX encode path. --- optimized/mlx/scripts/pre_encode_mlx.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/optimized/mlx/scripts/pre_encode_mlx.py b/optimized/mlx/scripts/pre_encode_mlx.py index efe3f98..72d8fd5 100644 --- a/optimized/mlx/scripts/pre_encode_mlx.py +++ b/optimized/mlx/scripts/pre_encode_mlx.py @@ -265,6 +265,7 @@ def run( *, max_duration: float = 600.0, overwrite: bool = False, + exclude: set | None = None, ) -> dict: """Encode every audio file under ``audio_dir`` into ``output_dir``. @@ -277,6 +278,8 @@ def run( audio_dir = Path(audio_dir).expanduser().resolve() output_dir = Path(output_dir).expanduser().resolve() files = find_audio_files(audio_dir) + if exclude: + files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude] max_samples = max_samples_for_duration(max_duration) output_dir.mkdir(parents=True, exist_ok=True) @@ -384,6 +387,10 @@ def main(): "--overwrite", action="store_true", help="Re-encode files whose outputs already exist", ) + parser.add_argument( + "--exclude-file", type=str, default=None, + help="Text file of relpaths (one per line, relative to --audio-dir) to skip", + ) args = parser.parse_args() audio_dir = Path(args.audio_dir).expanduser().resolve() @@ -395,11 +402,23 @@ def main(): print(f"No audio files found in {audio_dir}") sys.exit(1) + exclude = set() + if args.exclude_file and Path(args.exclude_file).is_file(): + exclude = {ln.strip() for ln in Path(args.exclude_file).read_text().splitlines() + if ln.strip()} + n_before = len(files) + if exclude: + files = [f for f in files if str(f.relative_to(audio_dir)) not in exclude] + if not files: + print(f"All {n_before} file(s) excluded by {args.exclude_file}") + sys.exit(1) + max_samples = max_samples_for_duration(args.max_duration) import mlx.core as mx print(f"Engine: MLX · device: {mx.default_device()}") print(f"Codec: {args.codec} (pad_modulo={SAME_ENCODER_PAD_MODULO[args.codec]}, fp32)") - print(f"Input: {audio_dir} ({len(files)} audio files)") + print(f"Input: {audio_dir} ({len(files)} audio files" + + (f", {n_before - len(files)} excluded)" if exclude else ")")) print(f"Output: {Path(args.output_dir).expanduser().resolve()}") print(f"Max: {max_samples / SAMPLE_RATE:.1f}s ({max_samples:,} samples)\n") @@ -414,6 +433,7 @@ def main(): encoder, max_duration=args.max_duration, overwrite=args.overwrite, + exclude=exclude, ) From ccadf8aff40d373d8eb4d1b9a7e6e3188db95d3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 17:17:40 -0400 Subject: [PATCH 23/28] MLX trainer: torch-identical tqdm output + ARC demos Emit the same per-step tqdm line as underfit's torch loop (desc "Step N, Epoch E"; postfix train/loss, train/lr, train/grad_norm, train/lora_magnitude) instead of a plain print, so the dashboard collapses the progress bar and charts grad-norm / lora-magnitude with no translation layer. Both norms are one global L2 over the adapter grads (post-clip) and params (pre-update), matching loop.py _compute_grad_and_lora_norms. Add optional --gradient-clip-val (the dashboard passes 1.0) via mlx clip_grad_norm; grad_norm is post-clip like torch. Demos: entries tagged arc=true now render on the shipped rf_denoiser weights (train-on-base / demo-on-ARC). The trained LoRA is merged into a fresh copy of dit__f16.npz (auto-downloaded, or --arc-weights) and sampled with the pingpong integrator; the seconds conditioner is model-independent so the trained embedder is reused. The ARC model is loaded once per demo round and freed after. RF/base entries keep the Euler path. Documents both in TRAINING_CONVENTIONS section 11. --- optimized/mlx/TRAINING_CONVENTIONS.md | 24 +++ optimized/mlx/models/defs/demo_mlx.py | 15 ++ optimized/mlx/scripts/lora_train_mlx.py | 220 ++++++++++++++++++++---- 3 files changed, 225 insertions(+), 34 deletions(-) diff --git a/optimized/mlx/TRAINING_CONVENTIONS.md b/optimized/mlx/TRAINING_CONVENTIONS.md index 0c82d61..4d03285 100644 --- a/optimized/mlx/TRAINING_CONVENTIONS.md +++ b/optimized/mlx/TRAINING_CONVENTIONS.md @@ -201,3 +201,27 @@ trainer — final agreement **loss 3.9829 vs 3.9823, prediction PSNR 80.7 dB** | peak memory | 3.67 GiB driver-allocated | (not instrumented; comparable class) | | loss regime | 1–43, spikes at low t | same (0.9–35, spikes at low t) | | checkpoint | underfit-valid, 423 keys | underfit-valid, identical key set | + +## 11. Dashboard integration (implements §6/§7 for the underfit dashboard) + +- **Progress output is byte-identical to underfit's torch loop** (loop.py:516): + one `tqdm` per epoch, `desc="Step N, Epoch E"`, `mininterval=0, miniters=1, + file=sys.stdout`, postfix `{train/loss:.6f, train/lr:.3e, train/grad_norm:.6f, + train/lora_magnitude:.6f}`. The dashboard collapses the progress bar on tqdm's + `\r` and parses those four keys per step (server.py `_HISTORY_RE`/`_LR_RE`/ + `_GN_RE`/`_LM_RE`) — so `underfit.backends.mlx_engine.run_mlx_training` just + inherits stdout (no line translation, which would break `\r` and drop the + postfix). grad_norm/lora_magnitude are one global L2 over the adapter grads + (post-clip) / params (pre-update), matching `_compute_grad_and_lora_norms`. +- **`--gradient-clip-val`** (dashboard passes 1.0) clips the global adapter-grad + norm via `mlx.optimizers.clip_grad_norm` before the step; grad_norm is reported + post-clip like torch. Default 0 = off. +- **ARC demos** (entries `arc=true`): the trained LoRA is merged into a fresh copy + of the shipped rf_denoiser weights (`dit__f16.npz`, auto-downloaded, or + `--arc-weights`) and sampled with the pingpong integrator — underfit's + train-on-base / demo-on-ARC story (its torch loop swaps the full ARC model in + place; the MLX path loads it with the LoRA merged at load, once per demo round, + freed after). The seconds conditioner is model-independent, so the trained + `bundle.secs` is reused for both RF and ARC conditioning. RF/base entries keep + the plain-Euler path. cfg/steps default to 1/8 for ARC, `demo_cfg_scales[0]`/ + `demo_steps` for RF. diff --git a/optimized/mlx/models/defs/demo_mlx.py b/optimized/mlx/models/defs/demo_mlx.py index fe5072a..927ab88 100644 --- a/optimized/mlx/models/defs/demo_mlx.py +++ b/optimized/mlx/models/defs/demo_mlx.py @@ -94,6 +94,21 @@ def rf_euler_sample(model_fn, x, sigmas, before_step=None): return x +def pingpong_sample(model_fn, x, sigmas, seed=0, before_step=None): + """Ping-pong (rf_denoiser) sampler for ARC demos — the distilled 8-step + re-noising integrator sa3_mlx ships (``sample_flow_pingpong``). Same + ``model_fn(x, t)`` / ``x`` / ``sigmas`` interface as ``rf_euler_sample``. + + Base demos integrate the rectified_flow BASE model with plain Euler; ARC + demos run the SHIPPED rf_denoiser (ARC) weights with the trained LoRA merged + in, so they use the pingpong sampler instead (matches underfit's ARC demo, + which overrides diffusion_objective to rf_denoiser).""" + from .sa3_pipeline import sample_flow_pingpong + return sample_flow_pingpong(model_fn, x, sigmas, seed=seed, + before_step=(lambda i: before_step(i, float(sigmas[i]))) + if before_step is not None else None) + + # ── decode ───────────────────────────────────────────────────────────────── def decode_latents(decoder, chunk_fn, chunk_cfg, latents, T_lat): diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 9379e20..2d31053 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -64,6 +64,7 @@ from pathlib import Path import numpy as np +from tqdm import tqdm SCRIPTS_DIR = Path(__file__).resolve().parent REPO = SCRIPTS_DIR.parent @@ -73,6 +74,7 @@ import mlx.core as mx # noqa: E402 import mlx.nn as nn # noqa: E402 import mlx.optimizers as optim # noqa: E402 +from mlx.utils import tree_flatten # noqa: E402 from sa3_mlx import DIT_CHOICES, load_dit # noqa: E402 from weights import ensure_local # noqa: E402 @@ -103,6 +105,14 @@ T5_CACHE_CAP = 512 # prompt-conditioning cache entries (oldest evicted first) +def _rm(path): + """Remove a file, ignoring a missing one (temp-ckpt cleanup).""" + try: + os.remove(path) + except OSError: + pass + + def grad_checkpoint(layer): """ Update all instances of type(layer) to use gradient checkpointing. @@ -202,6 +212,11 @@ def parse_args(): ap.add_argument("--weight-decay", type=float, default=0.0, help="AdamW decoupled weight decay (torch default 0.0; SA3 " "templates use 0.01)") + ap.add_argument("--gradient-clip-val", type=float, default=0.0, + help="Global grad-norm clip over the adapter params before " + "the optimizer step (0 = off). underfit's dashboard " + "passes 1.0; the grad_norm metric is then post-clip, " + "matching underfit's torch loop.") # LR schedule (underfit's optimizer_configs.diffusion.scheduler; InverseLR # with warmup — at step 0 the SA3 template LR is ~5e-7, not the nominal lr) ap.add_argument("--lr-scheduler", default="none", choices=["none", "inverse"], @@ -245,6 +260,12 @@ def parse_args(): ap.add_argument("--demo-dir", default=None, help="Where to write demo__.mp3 (default: cwd, " "which the dashboard sets to /demos/).") + ap.add_argument("--arc-weights", default=None, + help="DiT weights for ARC demos (entries with arc=true): the " + "SHIPPED rf_denoiser npz, NOT the training base. Default " + "auto-downloads models/mlx/dit__f16.npz. The trained " + "LoRA is merged into a fresh copy of these weights and " + "sampled with the pingpong integrator.") return ap.parse_args() @@ -352,6 +373,13 @@ def main(): print(f"dataset: {len(dataset)} latent file(s), crop {crop_len} latents" + (" (oversampling with replacement — tiny dataset)" if len(dataset) < args.batch_size else "")) + # Batches per epoch — mirrors iterate_batches: tiny sets (< batch) oversample + # to batch_size*100 draws (→ 100 batches); else ceil(n / batch). Used as the + # tqdm total so the progress bar (and the dashboard's collapse signature) + # matches underfit's torch dataloader. + _n_idx = (args.batch_size * 100 if len(dataset) < args.batch_size + else len(dataset)) + steps_per_epoch = max(1, (_n_idx + args.batch_size - 1) // args.batch_size) # ── models ─────────────────────────────────────────────────────────────── # LoRA training uses the BASE checkpoint (rectified_flow), NOT the shipped @@ -478,13 +506,30 @@ def model_fn(noised, t): # state structure never changes (no recompile on step 2). optimizer.init(bundle.trainable_parameters()) state = [bundle.state, optimizer.state, mx.random.state] + clip_val = float(args.gradient_clip_val or 0.0) + + def _global_l2(tree): + """sqrt(Σ ‖leaf‖²) over a param/grad tree in fp32 — the single global L2 + across all adapter leaves that underfit's _compute_grad_and_lora_norms + reports (loop.py:243).""" + total = mx.zeros((), dtype=mx.float32) + for _, v in tree_flatten(tree): + total = total + (v.astype(mx.float32) ** 2).sum() + return mx.sqrt(total) def train_step(latents, timesteps, loss_mask, prompt_cond, seconds_in, drop): loss, grads = value_and_grad(bundle, latents, timesteps, loss_mask, prompt_cond, seconds_in, drop) + if clip_val > 0.0: + grads, _ = optim.clip_grad_norm(grads, clip_val) + # grad_norm on the POST-clip grads, lora_magnitude on the PRE-update + # params — the order + definition underfit's torch loop uses (clip → + # compute norms → optimizer.step, loop.py:628-637). + grad_norm = _global_l2(grads) + lora_mag = _global_l2(bundle.trainable_parameters()) optimizer.update(bundle, grads) - return loss + return loss, grad_norm, lora_mag if args.compile: train_step = partial(mx.compile, inputs=state, outputs=state)(train_step) @@ -512,14 +557,21 @@ def checkpoint(step, epoch): t5_cache = {} if args.t5_cache else None t5_stats = {"hits": 0, "misses": 0} - # ── demos (underfit RF path) ────────────────────────────────────────────── - # RF Euler inference on the base model + trained LoRA, decoded to mp3. Baseline - # at the start, then every --demo-every steps + a final render. Idempotent. + # ── demos (underfit RF + ARC paths) ────────────────────────────────────── + # RF/base entries: Euler inference on the BASE model + trained LoRA. ARC + # entries (arc=true): the trained LoRA merged into the SHIPPED rf_denoiser + # weights, sampled with the pingpong integrator — underfit's "train on base, + # demo on ARC" story (its torch loop swaps in the full ARC model; here we + # load the ARC npz with the LoRA merged at load). Both decode to mp3. + # Baseline at the start, then every --demo-every steps + a final render. + # Idempotent (a resume never re-renders an existing demo). demo_entries = None if args.demo_every > 0: if args.demo_config: demo_entries = json.loads(Path(args.demo_config).read_text()) - print(f"demos: {len(demo_entries)} prompt(s) every {args.demo_every} " + n_arc = sum(1 for e in demo_entries if e.get("arc")) + print(f"demos: {len(demo_entries)} prompt(s) " + f"({n_arc} ARC) every {args.demo_every} " f"step(s) → {args.demo_dir or '.'}") else: print("WARNING: --demo-every set but no --demo-config — demos disabled") @@ -527,6 +579,76 @@ def checkpoint(step, epoch): demo_decoder_name = args.demo_decoder or DIT_CHOICES[args.dit]["default_decoder"] _demo_dec = {} # lazy: (decoder, chunk_fn, chunk_cfg) + def _demo_conditioning(prompt, seconds_val): + """cross_attn [1,257,768] fp16 + global_cond [1,768] fp16 for one demo + prompt. Model-independent (T5Gemma + the trained seconds conditioner), + so RF (base) and ARC demos share it.""" + prompt_cond = build_conditioning(t5, padding_emb, [prompt], + cache=t5_cache, cache_stats=t5_stats) + if bundle.secs is not None: + sec_tok = bundle.secs(mx.array([seconds_val], dtype=mx.float32)) + else: + sec_tok = secs_embedder([seconds_val]) + sec_tok = mx.stop_gradient(sec_tok.astype(mx.float32)) + cross = mx.concatenate([prompt_cond, sec_tok], axis=1).astype(mx.float16) + gcond = sec_tok[:, 0, :].astype(mx.float16) + return cross, gcond + + def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, + T_lat, seconds_val): + """Merge the trained LoRA into a fresh copy of the SHIPPED rf_denoiser + weights and sample with the pingpong integrator. Loaded once for the + whole ARC set, freed after (peak = base + ARC DiT during the demo).""" + arc_path = args.arc_weights or str( + ensure_local(f"models/mlx/dit_{args.dit}_f16.npz")) + tmp_ckpt = ckpt_dir / f".arc_demo_tmp_{step:08d}.safetensors" + save_lora_checkpoint( + bundle, tmp_ckpt, include=lora_config.get("include"), + exclude=lora_config.get("exclude"), + extra_config={"step": int(step), "base_model": base_model}) + t_a = time.time() + try: + arc_dit = mod.load_dit(arc_path, T_lat=T_lat, dtype=mx.float16, + compile_=False, lora_paths=[str(tmp_ckpt)], + lora_strength=1.0, lora_log=lambda *a, **k: None) + print(f" ARC model ({args.dit}) + LoRA merged " + f"({time.time()-t_a:.1f}s)") + except Exception as e: + print(f" ARC demos skipped: could not load " + f"{Path(arc_path).name} with LoRA merged: {e}", flush=True) + _rm(tmp_ckpt) + return + try: + for i, entry in arc_pending: + prompt = entry.get("prompt", "") + cfg = float(entry.get("cfg", 1)) + seed = int(entry.get("seed", 0)) + steps = int(entry.get("steps", 8)) + cross, gcond = _demo_conditioning(prompt, seconds_val) + null = mx.zeros_like(cross) if cfg != 1.0 else None + model_fn = demo.make_rf_model_fn(arc_dit, cross, gcond, cfg=cfg, + null_cross_attn=null) + sigmas = demo.build_sigmas(steps) + noise = mx.random.normal((1, 256, T_lat), dtype=mx.float16, + key=mx.random.key(seed)) + try: + latent = demo.pingpong_sample(model_fn, noise, sigmas, seed=seed) + audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, + latent, T_lat) + meta = {"prompt": prompt, "cfg": cfg, "seed": seed, + "steps": steps, "step": int(step), "arc": True} + demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE, + demo_dir, meta) + print(f" ♪ ARC demo {i} @ step {step}: {prompt[:44]!r} " + f"→ demo_{i}_{step:08d}.mp3", flush=True) + except Exception as e: + print(f" ARC demo {i} failed: {e}", flush=True) + finally: + del arc_dit + if hasattr(mx, "clear_cache"): + mx.clear_cache() + _rm(tmp_ckpt) + def run_demos(step): if not demo_entries: return @@ -545,22 +667,18 @@ def run_demos(step): adapters = list(iter_trainable_lora_layers(bundle)) T_lat = crop_len seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE - for i, entry in pending: + rf_pending = [(i, e) for i, e in pending if not e.get("arc")] + arc_pending = [(i, e) for i, e in pending if e.get("arc")] + + # ── RF/base demos (Euler on the base model + trained LoRA) ── + for i, entry in rf_pending: prompt = entry.get("prompt", "") cfg = float(entry.get("cfg", 7)) seed = int(entry.get("seed", 0)) steps = int(entry.get("steps", 50)) strength = entry.get("lora_strength") interval = entry.get("lora_interval_max") - prompt_cond = build_conditioning(t5, padding_emb, [prompt], - cache=t5_cache, cache_stats=t5_stats) - if bundle.secs is not None: - sec_tok = bundle.secs(mx.array([seconds_val], dtype=mx.float32)) - else: - sec_tok = secs_embedder([seconds_val]) - sec_tok = mx.stop_gradient(sec_tok.astype(mx.float32)) - cross = mx.concatenate([prompt_cond, sec_tok], axis=1).astype(mx.float16) - gcond = sec_tok[:, 0, :].astype(mx.float16) + cross, gcond = _demo_conditioning(prompt, seconds_val) null = mx.zeros_like(cross) if cfg != 1.0 else None model_fn = demo.make_rf_model_fn(bundle.dit, cross, gcond, cfg=cfg, null_cross_attn=null) @@ -595,6 +713,12 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): finally: if snap is not None: demo.scale_adapters(snap, 1.0) # restore trained factors + + # ── ARC demos (pingpong on the shipped ARC weights + trained LoRA) ── + if arc_pending: + _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, + T_lat, seconds_val) + if hasattr(mx, "clear_cache"): mx.clear_cache() @@ -608,9 +732,18 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): run_demos(step_offset) # baseline (untrained / resumed state) try: while raw_step < max_new_steps: - for batch in iterate_batches(dataset, args.batch_size, - shuffle=True, - seed=args.seed + epoch): + # tqdm per epoch — byte-identical to underfit's torch loop + # (loop.py:516) so the dashboard's progress collapse, history parser, + # and grad-norm / lora-magnitude charts work unchanged. desc + # "Step N, Epoch E"; postfix train/loss,lr,grad_norm,lora_magnitude. + pbar = tqdm( + iterate_batches(dataset, args.batch_size, shuffle=True, + seed=args.seed + epoch), + total=steps_per_epoch, + desc=f"Step {raw_step + step_offset}, Epoch {epoch}", + mininterval=0, miniters=1, file=sys.stdout, + ) + for batch in pbar: if raw_step >= max_new_steps: break global_step = raw_step + step_offset + 1 @@ -661,11 +794,9 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): if args.lr_scheduler != "none": optimizer.learning_rate = lr_now - t_step = time.time() - loss = train_step(latents, timesteps, loss_mask, prompt_cond, - seconds_in, drop) - mx.eval(state, loss) - step_ms = (time.time() - t_step) * 1000.0 + loss, grad_norm, lora_mag = train_step( + latents, timesteps, loss_mask, prompt_cond, seconds_in, drop) + mx.eval(state, loss, grad_norm, lora_mag) raw_step += 1 loss_v = float(loss) @@ -673,17 +804,38 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): float(t_np.mean()), loss_v)) if raw_step % 10 == 0: tele.flush() - rate = raw_step / max(time.time() - t_start, 1e-9) - print(f"step {global_step} train/loss {loss_v:.6f} " - f"train/lr {lr_now:.3e} epoch {epoch} " - f"({rate:.2f} it/s, {step_ms:.0f} ms)", flush=True) - - if args.checkpoint_every > 0 and \ - global_step % args.checkpoint_every == 0: - checkpoint(global_step, epoch) # save BEFORE demos - last_saved_step = global_step - if args.demo_every > 0 and global_step % args.demo_every == 0: - run_demos(global_step) + + # Metrics on the tqdm postfix — the dashboard parses + # train/loss=, train/lr=, train/grad_norm=, train/lora_magnitude= + # per step (server.py _HISTORY_RE/_LR_RE/_GN_RE/_LM_RE), and the + # "Step N," prefix gives it the authoritative global step. + pbar.set_postfix({ + "train/loss": f"{loss_v:.6f}", + "train/lr": f"{lr_now:.3e}", + "train/grad_norm": f"{float(grad_norm):.6f}", + "train/lora_magnitude": f"{float(lora_mag):.6f}", + }) + pbar.set_description(f"Step {global_step}, Epoch {epoch}") + + save_now = (args.checkpoint_every > 0 + and global_step % args.checkpoint_every == 0) + demo_now = (args.demo_every > 0 + and global_step % args.demo_every == 0) + if save_now or demo_now: + # Clear + disable the bar so checkpoint/demo prints land on + # their own lines (matches the torch loop, loop.py:680). + pbar.clear() + pbar.disable = True + try: + if save_now: + checkpoint(global_step, epoch) # save BEFORE demos + last_saved_step = global_step + if demo_now: + run_demos(global_step) + finally: + pbar.disable = False + pbar.refresh() + pbar.close() epoch += 1 except KeyboardInterrupt: print("\ninterrupted — saving final checkpoint") From 71061c5364911ff19cef2ca32d96566141367877 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 18:26:12 -0400 Subject: [PATCH 24/28] MLX demos: honor per-entry length; make the DiT length-agnostic Demo entries carry a `duration` (seconds) when they aren't at the crop length (the dashboard's full-length demos set it, crop-length ones omit it). The trainer was ignoring it and rendering every demo at the crop length. Now each demo computes its own T_lat from `duration` (falling back to the crop) for both the RF/base and ARC paths. For that to work on a model loaded at the training crop length, the DiT is now length-agnostic: the local-add-cond zeros in the text-to-audio path (local_add_cond=None) are built at the INPUT length instead of the baked self.T_lat. Identical for inference (which always loads T_lat == the gen length) and unused by training (which passes the all-ones inpaint cond), so this only unlocks reusing one loaded model across demo lengths. Both DiT sizes (dit_mlx_medium, dit_mlx) updated; attention is full (no mask) and RoPE is dynamic, so length is otherwise free. Verified: crop 128 + a 256-latent demo (via duration) render at 11.9 s / 23.8 s. --- optimized/mlx/models/defs/dit_mlx.py | 4 ++- optimized/mlx/models/defs/dit_mlx_medium.py | 6 ++++- optimized/mlx/scripts/lora_train_mlx.py | 28 +++++++++++++++------ 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/optimized/mlx/models/defs/dit_mlx.py b/optimized/mlx/models/defs/dit_mlx.py index c6f825a..5a4ba81 100644 --- a/optimized/mlx/models/defs/dit_mlx.py +++ b/optimized/mlx/models/defs/dit_mlx.py @@ -255,7 +255,9 @@ def __call__(self, x, t, cross_attn_cond_raw, global_cond_raw, local_add_cond=No x_pp = self.preprocess_conv(x_lc) + x_lc if local_add_cond is None: - local = mx.broadcast_to(self._local_zeros_1, (B, self.T_lat, LOCAL_ADD_COND_DIM)) + # Zero local-add-cond at the INPUT length (not the baked self.T_lat), + # so one loaded model runs any sequence length (see dit_mlx_medium). + local = mx.zeros((B, x.shape[-1], LOCAL_ADD_COND_DIM)) else: local = local_add_cond h = self.transformer(x_pp, context, global_embed, local) diff --git a/optimized/mlx/models/defs/dit_mlx_medium.py b/optimized/mlx/models/defs/dit_mlx_medium.py index 8cfcade..5128089 100644 --- a/optimized/mlx/models/defs/dit_mlx_medium.py +++ b/optimized/mlx/models/defs/dit_mlx_medium.py @@ -338,7 +338,11 @@ def __call__(self, x, t, cross_attn_cond_raw, global_cond_raw, local_add_cond=No x_pp = self.preprocess_conv(x_lc) + x_lc if local_add_cond is None: - local = mx.broadcast_to(self._local_zeros_1, (B, self.T_lat, LOCAL_ADD_COND_DIM)) + # Zero local-add-cond at the INPUT length (not the baked self.T_lat), + # so one loaded model runs any sequence length — text-to-audio + # inference always loads T_lat == the gen length (identical result), + # but training demos reuse the crop-length model for longer clips. + local = mx.zeros((B, x.shape[-1], LOCAL_ADD_COND_DIM)) else: local = local_add_cond diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 2d31053..0580202 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -579,6 +579,18 @@ def checkpoint(step, epoch): demo_decoder_name = args.demo_decoder or DIT_CHOICES[args.dit]["default_decoder"] _demo_dec = {} # lazy: (decoder, chunk_fn, chunk_cfg) + def _demo_T_lat(entry): + """Per-demo latent length. The dashboard sets `duration` (seconds) on a + demo entry only when it isn't at the crop length (see index.html + demo_cond serialization) — full-length demos carry it, crop-length ones + omit it. The DiT is length-agnostic (dit_mlx*), so one loaded model + serves both without reloading.""" + dur = entry.get("duration") + if dur: + return max(2, int(round(float(dur) * demo.SAMPLE_RATE + / demo.SAMPLES_PER_LATENT))) + return crop_len + def _demo_conditioning(prompt, seconds_val): """cross_attn [1,257,768] fp16 + global_cond [1,768] fp16 for one demo prompt. Model-independent (T5Gemma + the trained seconds conditioner), @@ -594,11 +606,10 @@ def _demo_conditioning(prompt, seconds_val): gcond = sec_tok[:, 0, :].astype(mx.float16) return cross, gcond - def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, - T_lat, seconds_val): + def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg): """Merge the trained LoRA into a fresh copy of the SHIPPED rf_denoiser weights and sample with the pingpong integrator. Loaded once for the - whole ARC set, freed after (peak = base + ARC DiT during the demo).""" + whole ARC set (length-agnostic), freed after (peak = base + ARC DiT).""" arc_path = args.arc_weights or str( ensure_local(f"models/mlx/dit_{args.dit}_f16.npz")) tmp_ckpt = ckpt_dir / f".arc_demo_tmp_{step:08d}.safetensors" @@ -608,7 +619,7 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, extra_config={"step": int(step), "base_model": base_model}) t_a = time.time() try: - arc_dit = mod.load_dit(arc_path, T_lat=T_lat, dtype=mx.float16, + arc_dit = mod.load_dit(arc_path, T_lat=crop_len, dtype=mx.float16, compile_=False, lora_paths=[str(tmp_ckpt)], lora_strength=1.0, lora_log=lambda *a, **k: None) print(f" ARC model ({args.dit}) + LoRA merged " @@ -624,6 +635,8 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, cfg = float(entry.get("cfg", 1)) seed = int(entry.get("seed", 0)) steps = int(entry.get("steps", 8)) + T_lat = _demo_T_lat(entry) + seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE cross, gcond = _demo_conditioning(prompt, seconds_val) null = mx.zeros_like(cross) if cfg != 1.0 else None model_fn = demo.make_rf_model_fn(arc_dit, cross, gcond, cfg=cfg, @@ -665,8 +678,6 @@ def run_demos(step): print(f" demo decoder {demo_decoder_name} loaded ({time.time()-t_d:.1f}s)") decoder, chunk_fn, chunk_cfg = _demo_dec["dec"] adapters = list(iter_trainable_lora_layers(bundle)) - T_lat = crop_len - seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE rf_pending = [(i, e) for i, e in pending if not e.get("arc")] arc_pending = [(i, e) for i, e in pending if e.get("arc")] @@ -678,6 +689,8 @@ def run_demos(step): steps = int(entry.get("steps", 50)) strength = entry.get("lora_strength") interval = entry.get("lora_interval_max") + T_lat = _demo_T_lat(entry) + seconds_val = T_lat * demo.SAMPLES_PER_LATENT / demo.SAMPLE_RATE cross, gcond = _demo_conditioning(prompt, seconds_val) null = mx.zeros_like(cross) if cfg != 1.0 else None model_fn = demo.make_rf_model_fn(bundle.dit, cross, gcond, cfg=cfg, @@ -716,8 +729,7 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): # ── ARC demos (pingpong on the shipped ARC weights + trained LoRA) ── if arc_pending: - _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg, - T_lat, seconds_val) + _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg) if hasattr(mx, "clear_cache"): mx.clear_cache() From abac31bab2e3aba97b8cc26978f0e6a092e4fb19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 18:42:55 -0400 Subject: [PATCH 25/28] MLX demos: live step-progress so slow full-length demos don't look hung MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A full-length (4096-latent) demo is ~10 s/forward on MLX medium → ~8 min for a 50-step RF demo, and nothing was printed until it finished — so training looked stuck at the step-0 baseline. Now each demo prints its length + step count before running and shows a tqdm step bar (rf_euler_sample + pingpong_sample gain a `desc`). Makes clear it's progressing, not hung. --- optimized/mlx/models/defs/demo_mlx.py | 27 +++++++++++++++++++------ optimized/mlx/scripts/lora_train_mlx.py | 9 +++++++-- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/optimized/mlx/models/defs/demo_mlx.py b/optimized/mlx/models/defs/demo_mlx.py index 927ab88..6e7c77a 100644 --- a/optimized/mlx/models/defs/demo_mlx.py +++ b/optimized/mlx/models/defs/demo_mlx.py @@ -25,9 +25,11 @@ import json import os import subprocess +import sys import mlx.core as mx import numpy as np +from tqdm import tqdm from .sa3_pipeline import build_pingpong_schedule, patched_decode @@ -75,14 +77,19 @@ def model_fn(x, t): return model_fn -def rf_euler_sample(model_fn, x, sigmas, before_step=None): +def rf_euler_sample(model_fn, x, sigmas, before_step=None, desc=None): """Euler ODE integration of the RF velocity field: x_{i+1} = x_i + Δt·v. ``sigmas`` is (steps+1,) descending from σ_max to 0. ``before_step(i, σ)``, if given, runs right before each velocity eval (per-step LoRA gating). + ``desc``, if given, shows a tqdm step bar — a full-length demo is minutes of + generation on MLX, so the bar makes clear it's progressing, not hung. """ num_steps = sigmas.shape[0] - 1 - for i in range(num_steps): + steps = range(num_steps) + if desc is not None: + steps = tqdm(steps, desc=desc, file=sys.stdout, leave=False, mininterval=0.5) + for i in steps: t_curr = sigmas[i] t_next = sigmas[i + 1] if before_step is not None: @@ -94,7 +101,7 @@ def rf_euler_sample(model_fn, x, sigmas, before_step=None): return x -def pingpong_sample(model_fn, x, sigmas, seed=0, before_step=None): +def pingpong_sample(model_fn, x, sigmas, seed=0, before_step=None, desc=None): """Ping-pong (rf_denoiser) sampler for ARC demos — the distilled 8-step re-noising integrator sa3_mlx ships (``sample_flow_pingpong``). Same ``model_fn(x, t)`` / ``x`` / ``sigmas`` interface as ``rf_euler_sample``. @@ -104,9 +111,17 @@ def pingpong_sample(model_fn, x, sigmas, seed=0, before_step=None): in, so they use the pingpong sampler instead (matches underfit's ARC demo, which overrides diffusion_objective to rf_denoiser).""" from .sa3_pipeline import sample_flow_pingpong - return sample_flow_pingpong(model_fn, x, sigmas, seed=seed, - before_step=(lambda i: before_step(i, float(sigmas[i]))) - if before_step is not None else None) + bar = (tqdm(total=int(sigmas.shape[0]) - 1, desc=desc, file=sys.stdout, + leave=False, mininterval=0.5) if desc is not None else None) + try: + return sample_flow_pingpong( + model_fn, x, sigmas, seed=seed, + before_step=(lambda i: before_step(i, float(sigmas[i]))) + if before_step is not None else None, + on_step=((lambda *_: bar.update(1)) if bar is not None else None)) + finally: + if bar is not None: + bar.close() # ── decode ───────────────────────────────────────────────────────────────── diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 0580202..68cdad9 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -644,8 +644,11 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg): sigmas = demo.build_sigmas(steps) noise = mx.random.normal((1, 256, T_lat), dtype=mx.float16, key=mx.random.key(seed)) + print(f" ▸ ARC demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), " + f"{steps} steps", flush=True) try: - latent = demo.pingpong_sample(model_fn, noise, sigmas, seed=seed) + latent = demo.pingpong_sample(model_fn, noise, sigmas, seed=seed, + desc=f"ARC demo {i}") audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, latent, T_lat) meta = {"prompt": prompt, "cfg": cfg, "seed": seed, @@ -708,9 +711,11 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): elif strength is not None and float(strength) != 1.0: snap = demo.snapshot_adapters(adapters) demo.scale_adapters(snap, float(strength)) + print(f" ▸ demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), " + f"{steps} steps", flush=True) try: latent = demo.rf_euler_sample(model_fn, noise, sigmas, - before_step=before) + before_step=before, desc=f"demo {i}") audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, latent, T_lat) meta = {"prompt": prompt, "cfg": cfg, "seed": seed, "steps": steps, "step": int(step)} From 4527a75f179b78119c1ad3391732940daabe8294 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 19:01:25 -0400 Subject: [PATCH 26/28] =?UTF-8?q?MLX=20demos:=20report=20per-demo=20wall-c?= =?UTF-8?q?lock=20in=20the=20=E2=99=AA=20line?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each demo's completion line now shows how long generation took, e.g. "♪ demo 0 @ step 0: '...' → demo_0_00000000.mp3 (482s)" — useful given full-length demos run minutes each on MLX. --- optimized/mlx/scripts/lora_train_mlx.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index 68cdad9..fdf2611 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -647,6 +647,7 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg): print(f" ▸ ARC demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), " f"{steps} steps", flush=True) try: + t_gen = time.time() latent = demo.pingpong_sample(model_fn, noise, sigmas, seed=seed, desc=f"ARC demo {i}") audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, @@ -656,7 +657,7 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg): demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE, demo_dir, meta) print(f" ♪ ARC demo {i} @ step {step}: {prompt[:44]!r} " - f"→ demo_{i}_{step:08d}.mp3", flush=True) + f"→ demo_{i}_{step:08d}.mp3 ({time.time()-t_gen:.0f}s)", flush=True) except Exception as e: print(f" ARC demo {i} failed: {e}", flush=True) finally: @@ -714,6 +715,7 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): print(f" ▸ demo {i}: {T_lat} latents (~{seconds_val:.0f}s audio), " f"{steps} steps", flush=True) try: + t_gen = time.time() latent = demo.rf_euler_sample(model_fn, noise, sigmas, before_step=before, desc=f"demo {i}") audio = demo.decode_latents(decoder, chunk_fn, chunk_cfg, latent, T_lat) @@ -725,7 +727,7 @@ def before(_i, sigma, s_on=s_on, interval=interval, snap=snap): meta["lora_interval_max"] = interval demo.save_demo_mp3(audio, i, step, demo.SAMPLE_RATE, demo_dir, meta) print(f" ♪ demo {i} @ step {step}: {prompt[:48]!r} " - f"→ demo_{i}_{step:08d}.mp3", flush=True) + f"→ demo_{i}_{step:08d}.mp3 ({time.time()-t_gen:.0f}s)", flush=True) except Exception as e: print(f" demo {i} failed: {e}", flush=True) finally: From 9fba072a2e7a528cbba98f8003474b34a050c105 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 19:10:49 -0400 Subject: [PATCH 27/28] MLX trainer: note the SAME-S demo-decoder substitution in the console When --demo-decoder differs from the model's default (SA3-medium's SAME-L), print "demos: decoding with SAME-S instead of the model's SAME-L (faster demos)" so the console shows the swap. Both RF and ARC demos use the chosen decoder. --- optimized/mlx/scripts/lora_train_mlx.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index fdf2611..ff4b6f8 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -577,6 +577,10 @@ def checkpoint(step, epoch): print("WARNING: --demo-every set but no --demo-config — demos disabled") demo_dir = args.demo_dir or "." demo_decoder_name = args.demo_decoder or DIT_CHOICES[args.dit]["default_decoder"] + _model_default_decoder = DIT_CHOICES[args.dit]["default_decoder"] + if demo_entries and demo_decoder_name != _model_default_decoder: + print(f"demos: decoding with {demo_decoder_name.upper()} instead of the " + f"model's {_model_default_decoder.upper()} (faster demos)") _demo_dec = {} # lazy: (decoder, chunk_fn, chunk_cfg) def _demo_T_lat(entry): From 0f5c31ccafb1ba4bd8700f923224a2fd7776c8c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”> Date: Wed, 15 Jul 2026 19:47:31 -0400 Subject: [PATCH 28/28] MLX trainer: write the ARC-demo scratch LoRA outside the checkpoints dir The temp LoRA merged into the ARC DiT was written to the run's checkpoints/ dir (as a dotfile), so it could surface in the dashboard's checkpoint list during the demo (or leak on a hard kill). Write it to the session dir instead (ckpt_dir.parent, not the scanned checkpoints/ subdir); it's still a dotfile and still removed after the demo. --- optimized/mlx/scripts/lora_train_mlx.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/optimized/mlx/scripts/lora_train_mlx.py b/optimized/mlx/scripts/lora_train_mlx.py index ff4b6f8..73b174f 100644 --- a/optimized/mlx/scripts/lora_train_mlx.py +++ b/optimized/mlx/scripts/lora_train_mlx.py @@ -616,7 +616,10 @@ def _run_arc_demos(step, arc_pending, decoder, chunk_fn, chunk_cfg): whole ARC set (length-agnostic), freed after (peak = base + ARC DiT).""" arc_path = args.arc_weights or str( ensure_local(f"models/mlx/dit_{args.dit}_f16.npz")) - tmp_ckpt = ckpt_dir / f".arc_demo_tmp_{step:08d}.safetensors" + # Scratch LoRA for merging into the ARC DiT — write it to the SESSION + # dir (ckpt_dir.parent), NOT the scanned checkpoints/ dir, so it never + # shows up as a checkpoint (dotfile too, belt + suspenders). + tmp_ckpt = ckpt_dir.parent / f".arc_demo_tmp_{step:08d}.safetensors" save_lora_checkpoint( bundle, tmp_ckpt, include=lora_config.get("include"), exclude=lora_config.get("exclude"),