Skip to content

Training & LoRA

Trainer

mlx_diffuser.training.DiffusionTrainer

Source code in src/mlx_diffuser/training/trainer.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class DiffusionTrainer:
    def __init__(
        self,
        model: nn.Module,
        scheduler: Scheduler,
        *,
        lr: float = 1e-4,
        optimizer: optim.Optimizer | None = None,
        weight_decay: float = 0.0,
        grad_clip: float | None = None,
        ema_decay: float | None = None,
        loss_weighting: Callable[[Scheduler, mx.array], mx.array] | None = None,
        class_dropout_prob: float = 0.0,
        compile: bool = True,
        seed: int = 0,
    ):
        self.model = model
        self.scheduler = scheduler
        self.optimizer = optimizer or optim.AdamW(learning_rate=lr, weight_decay=weight_decay)
        self.grad_clip = grad_clip
        self.loss_weighting = loss_weighting
        self.class_dropout_prob = class_dropout_prob
        self.ema = EMA(model, ema_decay) if ema_decay else None
        self._key = mx.random.key(seed)

        self.conditional = getattr(model, "y_embed", None) is not None
        self._null_label = model.config.num_classes if self.conditional else 0

        self._loss_and_grad = nn.value_and_grad(self.model, self._loss)
        if compile:
            state = [self.model.state, self.optimizer.state]
            self._step_fn = mx.compile(self._raw_step, inputs=state, outputs=state)
        else:
            self._step_fn = self._raw_step

    # --- loss / step ------------------------------------------------------
    def _loss(self, x0: mx.array, y: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        xt = self.scheduler.add_noise(x0, noise, t)
        target = self.scheduler.get_target(x0, noise, t)
        pred = self.model(xt, t, y) if self.conditional else self.model(xt, t)
        weights = self.loss_weighting(self.scheduler, t) if self.loss_weighting else None
        return mse_loss(pred, target, weights)

    def _raw_step(self, x0: mx.array, y: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        loss, grads = self._loss_and_grad(x0, y, noise, t)
        if self.grad_clip is not None:
            grads, _ = optim.clip_grad_norm(grads, self.grad_clip)
        self.optimizer.update(self.model, grads)
        return loss

    def _prep_labels(self, y: mx.array, key: mx.array) -> mx.array:
        if self.class_dropout_prob > 0:
            drop = mx.random.uniform(shape=y.shape, key=key) < self.class_dropout_prob
            y = mx.where(drop, mx.array(self._null_label), y)
        return y

    def step(self, x0: mx.array, y: mx.array | None = None) -> mx.array:
        """One optimization step on a batch; returns the (scalar) loss."""
        self._key, k_noise, k_t, k_drop = mx.random.split(self._key, 4)
        noise = mx.random.normal(x0.shape, key=k_noise)
        t = self.scheduler.sample_timesteps(x0.shape[0], k_t)
        if self.conditional:
            if y is None:
                raise ValueError("Model is class-conditional; provide labels `y`.")
            y = self._prep_labels(y, k_drop)
        else:
            y = mx.zeros((x0.shape[0],), dtype=mx.int32)

        loss = self._step_fn(x0, y, noise, t)
        mx.eval(self.model.parameters(), self.optimizer.state, loss)
        if self.ema is not None:
            self.ema.update(self.model)
        return loss

    # --- driver -----------------------------------------------------------
    def fit(
        self,
        dataset: Iterable,
        *,
        steps: int | None = None,
        epochs: int = 1,
        log_every: int = 50,
    ) -> list[float]:
        """Train over ``dataset`` (an iterable of ``x0`` or ``(x0, y)`` batches)."""
        history: list[float] = []
        done = 0
        for _ in range(epochs):
            for batch in dataset:
                x0, y = batch if isinstance(batch, tuple) else (batch, None)
                loss = self.step(x0, y)
                history.append(loss.item())
                done += 1
                if done % log_every == 0:
                    logger.info("step %d  loss %.4f", done, history[-1])
                if steps is not None and done >= steps:
                    return history
        return history

fit(dataset, *, steps=None, epochs=1, log_every=50)

Train over dataset (an iterable of x0 or (x0, y) batches).

Source code in src/mlx_diffuser/training/trainer.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def fit(
    self,
    dataset: Iterable,
    *,
    steps: int | None = None,
    epochs: int = 1,
    log_every: int = 50,
) -> list[float]:
    """Train over ``dataset`` (an iterable of ``x0`` or ``(x0, y)`` batches)."""
    history: list[float] = []
    done = 0
    for _ in range(epochs):
        for batch in dataset:
            x0, y = batch if isinstance(batch, tuple) else (batch, None)
            loss = self.step(x0, y)
            history.append(loss.item())
            done += 1
            if done % log_every == 0:
                logger.info("step %d  loss %.4f", done, history[-1])
            if steps is not None and done >= steps:
                return history
    return history

step(x0, y=None)

One optimization step on a batch; returns the (scalar) loss.

Source code in src/mlx_diffuser/training/trainer.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
def step(self, x0: mx.array, y: mx.array | None = None) -> mx.array:
    """One optimization step on a batch; returns the (scalar) loss."""
    self._key, k_noise, k_t, k_drop = mx.random.split(self._key, 4)
    noise = mx.random.normal(x0.shape, key=k_noise)
    t = self.scheduler.sample_timesteps(x0.shape[0], k_t)
    if self.conditional:
        if y is None:
            raise ValueError("Model is class-conditional; provide labels `y`.")
        y = self._prep_labels(y, k_drop)
    else:
        y = mx.zeros((x0.shape[0],), dtype=mx.int32)

    loss = self._step_fn(x0, y, noise, t)
    mx.eval(self.model.parameters(), self.optimizer.state, loss)
    if self.ema is not None:
        self.ema.update(self.model)
    return loss

mlx_diffuser.training.EMA

Tracks an EMA of a model's parameters for more stable sampling weights.

Source code in src/mlx_diffuser/training/ema.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class EMA:
    """Tracks an EMA of a model's parameters for more stable sampling weights."""

    def __init__(self, model: nn.Module, decay: float = 0.999):
        self.decay = decay
        self.shadow = tree_map(lambda p: mx.array(p), model.parameters())

    def update(self, model: nn.Module) -> None:
        d = self.decay
        self.shadow = tree_map(lambda s, p: d * s + (1.0 - d) * p, self.shadow, model.parameters())
        mx.eval(self.shadow)

    def copy_to(self, model: nn.Module) -> None:
        """Overwrite the model's parameters with the EMA weights (in place)."""
        model.update(self.shadow)

copy_to(model)

Overwrite the model's parameters with the EMA weights (in place).

Source code in src/mlx_diffuser/training/ema.py 
22
23
24
def copy_to(self, model: nn.Module) -> None:
    """Overwrite the model's parameters with the EMA weights (in place)."""
    model.update(self.shadow)

mlx_diffuser.training.batch_iterator(data, batch_size, *, shuffle=True, seed=0, drop_last=True)

Yield mini-batches from one or more aligned arrays.

A single array yields array batches; a tuple of arrays yields tuples of correspondingly-indexed batches (e.g. (images, labels)).

Source code in src/mlx_diffuser/training/data.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def batch_iterator(
    data: mx.array | tuple[mx.array, ...],
    batch_size: int,
    *,
    shuffle: bool = True,
    seed: int = 0,
    drop_last: bool = True,
) -> Iterator:
    """Yield mini-batches from one or more aligned arrays.

    A single array yields array batches; a tuple of arrays yields tuples of
    correspondingly-indexed batches (e.g. ``(images, labels)``).
    """
    arrays = data if isinstance(data, tuple) else (data,)
    n = arrays[0].shape[0]
    order = mx.array(list(range(n)))
    if shuffle:
        order = mx.random.permutation(n, key=mx.random.key(seed))

    stop = (n // batch_size) * batch_size if drop_last else n
    for start in range(0, stop, batch_size):
        idx = order[start : start + batch_size]
        batch = tuple(a[idx] for a in arrays)
        yield batch if isinstance(data, tuple) else batch[0]

mlx_diffuser.training.mse_loss(pred, target, weights=None)

Mean squared error, optionally weighted per-sample.

Source code in src/mlx_diffuser/training/losses.py 
11
12
13
14
15
16
def mse_loss(pred: mx.array, target: mx.array, weights: mx.array | None = None) -> mx.array:
    """Mean squared error, optionally weighted per-sample."""
    se = (pred - target) ** 2
    if weights is not None:
        se = expand_to(weights, se.ndim) * se
    return mx.mean(se)

mlx_diffuser.training.min_snr_weights(scheduler, t, gamma=5.0)

Min-SNR-gamma loss weighting (Hang et al., 2023) for VP diffusion.

Returns per-sample weights min(SNR, gamma) / SNR for epsilon prediction (/(SNR+1) adjustment is applied for v-prediction).

Source code in src/mlx_diffuser/training/losses.py 
19
20
21
22
23
24
25
26
27
28
29
30
def min_snr_weights(scheduler: DDPMScheduler, t: mx.array, gamma: float = 5.0) -> mx.array:
    """Min-SNR-gamma loss weighting (Hang et al., 2023) for VP diffusion.

    Returns per-sample weights ``min(SNR, gamma) / SNR`` for epsilon prediction
    (``/(SNR+1)`` adjustment is applied for v-prediction).
    """
    acp = scheduler.alphas_cumprod[t]
    snr = acp / (1.0 - acp)
    weights = mx.minimum(snr, gamma) / snr
    if scheduler.config.prediction_type == "v_prediction":
        weights = mx.minimum(snr, gamma) / (snr + 1.0)
    return weights

LoRA

mlx_diffuser.lora.inject_lora(model, rank=8, alpha=16.0, dropout=0.0, targets=DEFAULT_LORA_TARGETS)

Replace target nn.Linear layers with LoRA adapters and freeze the base.

Returns the number of layers adapted. After this, trainable_parameters contains only the adapter weights.

Source code in src/mlx_diffuser/lora/lora.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def inject_lora(
    model: nn.Module,
    rank: int = 8,
    alpha: float = 16.0,
    dropout: float = 0.0,
    targets: tuple[str, ...] = DEFAULT_LORA_TARGETS,
) -> int:
    """Replace target ``nn.Linear`` layers with LoRA adapters and freeze the base.

    Returns the number of layers adapted. After this, ``trainable_parameters``
    contains only the adapter weights.
    """
    count = 0

    def predicate(key: str, value: nn.Module) -> bool:
        return isinstance(value, nn.Linear) and key in targets

    def make(linear: nn.Linear) -> LoRALinear:
        nonlocal count
        count += 1
        return LoRALinear(linear, rank=rank, alpha=alpha, dropout=dropout)

    _replace_modules(model, predicate, make)
    model.freeze()
    for lora in _iter_lora(model):
        lora.unfreeze(keys=["lora_a", "lora_b"], recurse=False)
    return count

mlx_diffuser.lora.merge_lora(model)

Fuse all adapters into dense layers in place; returns the model.

Source code in src/mlx_diffuser/lora/lora.py 
109
110
111
112
113
114
115
116
117
def merge_lora(model: nn.Module) -> nn.Module:
    """Fuse all adapters into dense layers in place; returns the model."""
    _replace_modules(
        model,
        predicate=lambda key, value: isinstance(value, LoRALinear),
        make=lambda lora: lora.fused(),
    )
    model.unfreeze()
    return model

mlx_diffuser.lora.save_lora(model, save_directory, *, rank, alpha, targets=DEFAULT_LORA_TARGETS, dropout=0.0)

Source code in src/mlx_diffuser/lora/lora.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def save_lora(
    model: nn.Module,
    save_directory: str | Path,
    *,
    rank: int,
    alpha: float,
    targets: tuple[str, ...] = DEFAULT_LORA_TARGETS,
    dropout: float = 0.0,
) -> Path:
    save_directory = Path(save_directory)
    save_directory.mkdir(parents=True, exist_ok=True)
    config = {"rank": rank, "alpha": alpha, "targets": list(targets), "dropout": dropout}
    (save_directory / ADAPTER_CONFIG_NAME).write_text(json.dumps(config, indent=2, sort_keys=True))
    mx.save_safetensors(str(save_directory / ADAPTER_WEIGHTS_NAME), lora_state_dict(model))
    return save_directory

mlx_diffuser.lora.load_lora(model, path)

Inject adapters per adapter_config.json and load their weights.

Source code in src/mlx_diffuser/lora/lora.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def load_lora(model: nn.Module, path: str | Path) -> int:
    """Inject adapters per ``adapter_config.json`` and load their weights."""
    path = Path(path)
    config = json.loads((path / ADAPTER_CONFIG_NAME).read_text())
    count = inject_lora(
        model,
        rank=config["rank"],
        alpha=config["alpha"],
        dropout=config.get("dropout", 0.0),
        targets=tuple(config["targets"]),
    )
    weights = mx.load(str(path / ADAPTER_WEIGHTS_NAME))
    assert isinstance(weights, dict)
    model.load_weights(list(weights.items()), strict=False)
    mx.eval(model.parameters())
    return count

mlx_diffuser.lora.LoRALinear

Bases: Module

Wraps a (frozen) nn.Linear with a trainable low-rank update.

Source code in src/mlx_diffuser/lora/lora.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class LoRALinear(nn.Module):
    """Wraps a (frozen) ``nn.Linear`` with a trainable low-rank update."""

    def __init__(self, linear: nn.Linear, rank: int = 8, alpha: float = 16.0, dropout: float = 0.0):
        super().__init__()
        self.linear = linear
        out_features, in_features = linear.weight.shape
        self.rank = rank
        self.scale = alpha / rank
        bound = 1.0 / math.sqrt(in_features)
        self.lora_a = mx.random.uniform(low=-bound, high=bound, shape=(rank, in_features))
        self.lora_b = mx.zeros((out_features, rank))  # zero-init => identity at start
        self.dropout = nn.Dropout(dropout) if dropout > 0 else None

    def __call__(self, x: mx.array) -> mx.array:
        base = self.linear(x)
        z = self.dropout(x) if self.dropout is not None else x
        delta = (z @ self.lora_a.T) @ self.lora_b.T
        return base + self.scale * delta

    def fused(self) -> nn.Linear:
        """Return a dense ``nn.Linear`` with the adapter merged in."""
        out_features, in_features = self.linear.weight.shape
        has_bias = "bias" in self.linear
        merged = nn.Linear(in_features, out_features, bias=has_bias)
        merged.weight = self.linear.weight + self.scale * (self.lora_b @ self.lora_a)
        if has_bias:
            merged.bias = self.linear.bias
        return merged

fused()

Return a dense nn.Linear with the adapter merged in.

Source code in src/mlx_diffuser/lora/lora.py 
45
46
47
48
49
50
51
52
53
def fused(self) -> nn.Linear:
    """Return a dense ``nn.Linear`` with the adapter merged in."""
    out_features, in_features = self.linear.weight.shape
    has_bias = "bias" in self.linear
    merged = nn.Linear(in_features, out_features, bias=has_bias)
    merged.weight = self.linear.weight + self.scale * (self.lora_b @ self.lora_a)
    if has_bias:
        merged.bias = self.linear.bias
    return merged

Performance

mlx_diffuser.perf.compile_model(model, *, shapeless=False)

Return a compiled callable for inference with this model.

Parameters are passed as implicit inputs so weight updates (e.g. after a LoRA merge) are picked up without a stale graph. Use shapeless=True only when input ranks are stable but sizes vary a lot (read MLX's shapeless-compile caveats first) to avoid recompiling per resolution.

Source code in src/mlx_diffuser/perf.py 
31
32
33
34
35
36
37
38
39
def compile_model(model: nn.Module, *, shapeless: bool = False) -> Callable:
    """Return a compiled callable for *inference* with this model.

    Parameters are passed as implicit inputs so weight updates (e.g. after a LoRA
    merge) are picked up without a stale graph. Use ``shapeless=True`` only when
    input ranks are stable but sizes vary a lot (read MLX's shapeless-compile
    caveats first) to avoid recompiling per resolution.
    """
    return mx.compile(model, inputs=model.state, shapeless=shapeless)

mlx_diffuser.perf.memory_report()

Snapshot of unified-memory usage in GB (active / peak / cache).

Source code in src/mlx_diffuser/perf.py 
42
43
44
45
46
47
48
def memory_report() -> dict[str, float]:
    """Snapshot of unified-memory usage in GB (active / peak / cache)."""
    return {
        "active_gb": bytes_to_gb(mx.get_active_memory()),
        "peak_gb": bytes_to_gb(mx.get_peak_memory()),
        "cache_gb": bytes_to_gb(mx.get_cache_memory()),
    }