Skip to content
fast-axolotl_
GitHub
intermediate 10 minutes

Fix Axolotl RAM OOM with Rust streaming

Axolotl issue #2975 documents a large-dataset preload spiking RAM. This guide switches the reader to fast-axolotl streaming so memory is bounded by batch size, not dataset size.

  1. 1

    Install and import the shim

    Install fast-axolotl and import it before axolotl so the streaming reader is available.

    uv add fast-axolotl
  2. 2

    Enable Rust streaming in your YAML

    Turn on streaming in the Axolotl config. This routes dataset reads through the Rust streaming_dataset_reader.

    dataset_use_rust_streaming: true
sequence_len: 32768
  3. 3

    Run training and watch RAM

    Start your training run. Peak RAM should now be bounded by the batch size rather than climbing with dataset size.

    accelerate launch -m axolotl.cli.train config.yml
  4. 4

    Confirm the streaming path is active

    If memory still spikes, assert the extension linked; a False here means the pure-Python reader is still in use.

    python -c "import fast_axolotl; print(fast_axolotl.is_available())"

Why this works

  • →Streaming reads fixed-size batches on demand, so the full dataset is never held in memory.
  • →The README reports a 77x speedup on Parquet streaming at 50,000 rows on 16-core Linux.
  • →Parquet, Arrow, JSON, JSONL, CSV, and text are supported with transparent ZSTD/Gzip decompression.

More guides in the index, or read the architecture to see what the shim is doing.