MLX-LM training memory and gradient control
Parent: MLX-LM training and adapter lifecycle · Published reference · snapshot 2026-10-23 · viability 4.65 / 5 (research)
↓ Facts as markdownall context files
Source-anchored research for MLX-LM training memory and gradient control. Run status: COMPLETED.
These notes link each claim to its source. Some citations identify archived research reports that are no longer hosted on this site. A published reference means the content is available; it does not certify independent review or accuracy.Read the editorial policy and follow the sources before relying on a claim.
Structure and components
- MLX-LM's guide also recommends reducing `--batch-size`, `--num-layers`, or example length to lower memory; it cautions that training fewer layers may reduce quality on data-heavy runs, while research treats microbatching and checkpointing as separate techniques with workload-dependent memory and compute tradeoffs. (medium confidence) [source]
- --grad-checkpoint enables training gradient checkpointing. --batch-size reduces examples in each minibatch, while --num-layers reduces the last-layer subset selected for training. For one worker, batch_size B and --grad-accumulation-steps A combine A minibatches before an optimizer update, giving an effective example batch size B*A. (medium confidence) (verified as of 2026-10-23) [source]
How it works
- train() rejects grad_accumulation_steps below one with ValueError("grad_accumulation_steps must be at least 1"). It sums successive microbatch gradients; at an update it calls average_gradients, divides each accumulated gradient by grad_accumulation_steps when greater than one, then calls optimizer.update. (low confidence) (verified as of 2026-10-23) [source]
Problems, failure modes and limitations
- A June 2025 user report recorded Metal `Insufficient Memory` at a reported 34.108 GB peak with batch size 1, `--grad-checkpoint`, and a 5,000-token maximum sequence; after reinstalling MLX-LM from source, the follow-up log reached iteration 100, so the report does not isolate a memory-control limit or give a portable capacity estimate. (low confidence) [source]
- Activation memory grows with batch size and sequence length, so the MLX-LM guide recommends lowering `--batch-size` and shortening long examples when memory is constrained; it notes that smaller batches may slow training. (high confidence) (verified as of 2026-10-23) [source]
- iterate_batches truncates examples longer than max_seq_length, emits a warning, and updates the recorded lengths to the truncated lengths. If the dataset contains fewer examples than batch_size, it raises ValueError stating that the dataset must have at least batch_size examples. Pre-split long examples and choose a batch size no larger than the dataset. (low confidence) (verified as of 2026-10-23) [source]
Comparisons and alternatives
- MLX-LM's `--grad-checkpoint` trades extra computation for lower activation memory by recomputing intermediate values needed during backpropagation; its guide says the option is more helpful with larger batches or sequences. Framework documentation describes the same recomputation tradeoff and notes that checkpoint placement affects peak-memory savings. (high confidence) (verified as of 2026-10-23) [source]
- As of 2026-10-23, open MLX-LM issue #1206 contains reports of Qwen3.5 training OOM despite `--grad-checkpoint` on M5 Max and M3 Pro systems, alongside a base-M5 report that the OOM was not reproducible and training completed normally through iteration 5; treat this as configuration-specific evidence rather than a general memory limit. (low confidence) (verified as of 2026-10-23) [source]
Facts and statements
- MLX-LM's memory guide says `--grad-accumulation-steps N` accumulates gradients across N batches before updating, so users can keep each microbatch small while increasing effective batch size without increasing batch memory; Hugging Face's training documentation describes the same effective-batch relationship. (high confidence) (verified as of 2026-10-23) [source]
Children
- No children recorded.