diff --git a/optimized/mlx/README.md b/optimized/mlx/README.md index 1d7b615..b928d48 100644 --- a/optimized/mlx/README.md +++ b/optimized/mlx/README.md @@ -161,6 +161,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..4ed5619 --- /dev/null +++ b/optimized/mlx/models/defs/audio_encoding.py @@ -0,0 +1,287 @@ +"""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, +): + # 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 + 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..f7313e6 --- /dev/null +++ b/optimized/mlx/models/defs/training.py @@ -0,0 +1,228 @@ +"""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) + 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( + 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..f3d928c --- /dev/null +++ b/tests/test_mlx_audio_encoding.py @@ -0,0 +1,122 @@ +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_for_even_and_odd_overlaps(): + rng = np.random.default_rng(21) + audio = mx.array( + rng.standard_normal((2, 2, SAMPLES_PER_LATENT * 11)).astype(np.float32) + ) + encoder = GroupingEncoder() + + unchunked = encode_audio(encoder, audio, pad_modulo=16) + 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), + ) + + +@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..da9eb3a --- /dev/null +++ b/tests/test_mlx_lora.py @@ -0,0 +1,407 @@ +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 import dit_mlx, dit_mlx_medium +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))) + + +@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__() + 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..25a3e25 --- /dev/null +++ b/tests/test_mlx_training.py @@ -0,0 +1,90 @@ +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) + + +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)