Recurrent-TinyLlama-3T-train-recurrence-4-single-phase

Recurrent-TinyLlama-3T-train-recurrence-4-single-phase

Recurrent-TinyLlama-3T-train-recurrence-4-single-phase is part of the Retrofitting Recurrence set of models. A set of depth recurrent models trained by taking layers from pretrained feedforward language models (link to paper).

Downloading and Using the Model

Load the model like this:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("smcleish/Recurrent-TinyLlama-3T-train-recurrence-4-single-phase", torch_dtype=torch.float32, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("smcleish/Recurrent-TinyLlama-3T-train-recurrence-4-single-phase")

Modifying the Model's Depth at Test Time:

By providing the argument num_steps, the model will execute a forward pass with that amount of compute:

input_ids = tokenizer.encode("The capital of Westphalia is", return_tensors="pt", add_special_tokens=True).to(device)
model.eval()
model.to(device)
model(input_ids, num_steps=32)

Depth Recurrent Features

The modelling file is based on tomg-group-umd/huginn-0125, and therefore inherits all of the cool recurrent features demonstrated in the README.md for that model.

Training

We train using https://github.com/mcleish7/retrofitting-recurrence using AMD MI300A GPUs on Tuolumne at Lawrence Livermore National Laboratory.

Data

Train and validation data is taken from non-overlapping subsets of raw text data. As such it is not an instruction model.

Licence

This model is released under the apache-2.0 licence.

Contact

Please, feel free to contact us with any questions, or open a discussion thread.

Citation

@article{mcleish2025teaching,
    title={Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence}, 
    author={Sean McLeish and Ang Li and John Kirchenbauer and Dayal Singh Kalra and Brian R. Bartoldson and Bhavya Kailkhura and Avi Schwarzschild and Jonas Geiping and Tom Goldstein and Micah Goldblum},
    journal={arXiv preprint arXiv:2511.07384},
    year={2025}
}