// architecture

How the shim swaps hot paths

fast-axolotl is not a fork of Axolotl. It is a Python package that, on import, installs Rust implementations of four data-pipeline functions into sys.modules — so existing axolotl.utils.* calls run native code with no source edit.

import fast_axolotl   # must precede: import axolotl
        │
        ▼  installs shim into sys.modules
   ┌───────────────────────────────────────────────┐
   │  Python  axolotl.utils.data / packing / ...   │
   │        │  calls resolve to shimmed impls       │
   │  ······│···········FFI boundary················ │
   │        ▼                                        │
   │  Rust    streaming reader   ──▶ 77x          │
   │          parallel SHA256    ──▶ 1.9x         │
   │          pack_sequences     ──▶ 0.4x (small)   │
   │          pad_sequences      ──▶ 0.5x (small)   │
   └───────────────────────────────────────────────┘
        │
        ▼
   Axolotl training proceeds — unchanged config

Import order matters: fast_axolotl must be imported before axolotl so the shim is in place when Axolotl resolves its utilities.

// the four hot paths

What crosses the boundary

01 · read

win

streaming_dataset_reader

Reads a file in fixed-size batches without materializing it. Format detected by extension; ZSTD/Gzip decompressed transparently. This is the 77x win.

02 · dedup

win

parallel_hash_rows / deduplicate_indices

Fans SHA256 across cores and returns unique row indices plus hashes, optionally filtered against a seen set. This is the 1.9x win.

03 · pack

caveat

pack_sequences

Emits input_ids, labels, and attention_mask at max_length in one Rust pass. Exposed, but ~0.4x at small batches due to FFI overhead.

04 · pad

caveat

pad_sequences

Left/right padding and pad_to_multiple_of alignment. Same ~0.5x small-batch FFI caveat as packing.

Why the FFI boundary matters

Every call from Python into Rust pays a fixed crossing cost. For streaming and hashing, each call moves a large amount of work across the boundary, so the crossing is negligible and Rust's throughput dominates. For packing and padding at small batch sizes, the work per call is tiny, so the crossing cost outweighs the gain — which is exactly why the README reports ~0.4x and ~0.5x there. Those operations are expected to pay off only at larger training-scale batches where the FFI overhead is amortized.

Why nothing in Axolotl changes

Because the shim installs itself into sys.modules, the substitution is transparent to Axolotl. Your training script, YAML config, and the Axolotl source stay untouched. Read the install guide or the FAQ for the practical steps.