mx.checkpoint does not bound backward memory of a chunked recurrence when per-step inputs need gradients

Aperta
  • Qualcosa non va?
#4,837 1 commento 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 3 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
44/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
cpp, machine-learning, python
Ambito
backend, machine-learning, performance

Direzione di ricerca

Start by running the provided Python reproduction with plain and ckpt modes, then trace mx.checkpoint and its gradient handling for chunk inputs that require cotangents. The report includes measurements for gradients with respect to both the initial state and per-step inputs. Done means checkpointed backward memory follows the expected O(T/C + C) scaling in both cases.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

MLX version: 0.32.4 (latest), macOS 26.4.1, Apple M1 Max 32 GB.

Summary

A recurrence is split into chunks of 32 steps, and each chunk is wrapped in mx.checkpoint. When the gradient is taken only with respect to the initial state, backward memory stays flat as the sequence grows, as expected. When the gradient is also taken with respect to the per-step inputs (q, k, v, which any training use needs), backward memory grows linearly with the number of steps again: roughly one state-sized buffer per step stays alive.

Repro

State 32x128x128 fp32 = 2 MiB. Run python repro.py T plain and python repro.py T ckpt.

import sys
import mlx.core as mx

T, mode, CHUNK = int(sys.argv[1]), sys.argv[2], 32
H, D = 32, 128
mx.random.seed(0)
q = mx.random.normal((T, H, D)); k = mx.random.normal((T, H, D)) * 0.05; v = mx.random.normal((T, H, D))
x0 = mx.zeros((H, D, D)); mx.eval(q, k, v, x0)

def run(s, qc, kc, vc):
    ys = []
    for t in range(qc.shape[0]):
        s = s * 0.95
        delta = vc[t] - (s * kc[t][:, None, :]).sum(-1)
        s = s + delta[:, :, None] * kc[t][:, None, :]
        ys.append((s * qc[t][:, None, :]).sum(-1))
    return s, mx.stack(ys)

def loss(q, k, v):
    if mode == "plain":
        return run(x0, q, k, v)[1].sum()
    s, outs = x0, []
    for lo in range(0, T, CHUNK):
        s, y = mx.checkpoint(run)(s, q[lo:lo+CHUNK], k[lo:lo+CHUNK], v[lo:lo+CHUNK]); outs.append(y)
    return mx.concatenate(outs).sum()

mx.reset_peak_memory(); base = mx.get_active_memory()
val, grads = mx.value_and_grad(loss, argnums=(0, 1, 2))(q, k, v); mx.eval(val, grads)
print(f"T={T} {mode}: backward peak {(mx.get_peak_memory() - base) / 2**20:.0f} MiB")
Results (backward peak above inputs)
T plain ckpt (chunk 32)
256 1783 MiB 1031 MiB
512 3436 MiB 1582 MiB
1024 6510 MiB 2686 MiB

Same script with the gradient taken only w.r.t. the initial state x0: ckpt 373 / 378 / 427 MiB for T = 256 / 512 / 1024 (flat); plain 763 / 1296 / 2347 MiB.

Expected

With checkpointing, backward memory should be O(T/C + C) states whether or not the per-step inputs need cotangents.

Context

This limits LoRA training of Gated DeltaNet models (Qwen3.5 / 3.6 MoE) with mlx-lm: the training path (gated_delta_ops) keeps about 4 states per token (~8 MiB/token), and chunked checkpointing only gets it to ~6-7 MiB/token.

Lingua principale
C++
Stelle
29.1k
Fork
2.4k
Merge medio
3g 20h
PR unite (30g)
103

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di ml-explore/mlx

Tutte le issue di ml-explore/mlx

Issue simili

Altre issue su C++

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.