Cort33x

Building and training small language models on Apple silicon

By Arpan Bhandari

Source: https://arpan.sh/projects/cort33x

Building a local language-model training workflow and working through its data, memory, and recovery problems.

I built Cort33x to work through language-model training myself, including the tokenizer, model, training loop, and inference. Running it locally meant memory use and sustained throughput were constraints throughout the project.

The work included PyTorch and native MLX implementations, a 24,576-token BPE vocabulary, and training with sequence lengths from 2K to 8K tokens.


Preserving context in the data

Long documents needed to fit into training windows without losing all context at each break. I kept document boundaries in the tokenized dataset and used overlapping windows, with a default overlap of 256 tokens.

Short documents were packed with end-of-document separators. This let their tokens contribute to training without filling the remaining space with padding targets. Keeping the dataset size separate from the planned training-token budget also made it clear what a run would actually cover.


Finding the memory bottleneck

One gradient-accumulation run reached roughly 86 GB of active memory. Smaller microbatches were still retaining computation graphs, so accumulation was not giving the memory savings I expected.

I separated the compiled single-batch path from accumulation and explicitly evaluated the microbatch gradients. The experience changed what I measured: active memory, cache use, and throughput over a sustained interval, rather than just the first few steps after a restart.

Resume exposed a different issue. The compiled training function had to be created after loading the checkpoint; otherwise, it retained references to the old state structure. Restoring weights alone was not enough. Optimizer state and the data sampler's random state had to be restored too.

Cort33x reached a working training and interactive inference workflow. It remained an experiment, without a published evaluation of model quality. The work gave me experience tracing training failures through the data pipeline, memory behavior, and execution engine.