Skip to content
ifm-aiPublic
forked from ifm-ai/xllm

About

Official code for Towards Looped Models Done Right based on xLLM

Resources

Stars

41 stars

Watchers

0 watching

Forks

 
 

Repository files navigation

xLLM-Loop

Official code for Towards Looped Models Done Right

Part I paper Part II paper (arXiv) PyTorch 2.11+ CUDA 12.8+ 49 checkpoints on Hugging Face Apache 2.0 License

Papers  |  Highlights  |  Installation  |  Quick Start  |  Reproducing the Papers  |  Checkpoints  |  Issues

xLLM-Loop is the code release of our work on looped language models, which reuse a block of layers several times per token. It adds looped Transformer and Huginn architectures with depth-controlled training to the xLLM training framework, together with the training, evaluation and checkpoint code of two papers.

Papers

Paper Scope Code Guide
Part I: Topology, Input Injection, Recurrent-State Organization
Blog post · (PDF and arXiv link to be added)
17 Dense and MoE recipes; 25 checkpoints xllm/paper_part1 release/paper-part1
Part II: Rethinking at Fixed Points
arXiv · PDF
Depth priors, input-injection variants and distilled prefill; 24 checkpoints xllm/paper_part2 release/paper-part2

Highlights

Area Capabilities
Looped architectures LoopedTransformer iterates a range of layers; Huginn runs a prelude, a recurrent block and a coda, with a single or a hierarchical (H/L) recurrent state. See looped.py.
Depth-controlled training DepthControlledHuginn trains at a fixed, sampled (Poisson-lognormal) or learned depth with truncated backpropagation, and logs training FLOPs in expectation over the depths, or at the drawn depths with the learned prior.
Input injection Diagonal, linear and orthogonal input injection, with optional normalization of the injected input or the recurrent state.
Training loop Gradient accumulation, early stopping on an unchanged schedule (stop_step), and resume that can restore per-rank RNG states.
Paper protocols Recipe launchers, evaluators and data preparation scripts for the experiments of both papers.
Released checkpoints 49 checkpoints in native xLLM format on Hugging Face; the loaders verify every file against its manifest.

Installation

xLLM-Loop installs like xLLM. Start with PyTorch >= 2.11 and CUDA >= 12.8; the release was checked with Python 3.12, PyTorch 2.11.0+cu128 and FlashAttention 3.0.0. The papers' experiments ran on NVIDIA H200 GPUs.

1. Install xLLM-Loop

git clone https://github.com/ifm-ai/xllm-loop.git
cd xllm-loop

python -m pip install -r requirements.txt
python -m pip install -e . --no-build-isolation --config-settings editable_mode=compat
python -m pip install "transformers[torch]" pyarrow

2. Set Up FlashAttention 3

The paper code runs with FlashAttention 3. Follow the upstream installation instructions, then select it for every command:

export ENABLE_FLASH_ATTENTION_3=true

The xLLM installation guide covers the native extensions and the other attention backends.

Quick Start

Download a released checkpoint and generate from it:

hf download IFM/LoopedLM-P1-dense-huginn-336b --local-dir checkpoints/dense-huginn-336b
from xllm.paper_part1.native_inference import generate_native, load_native_model

model, tokenizer, config = load_native_model("checkpoints/dense-huginn-336b")
tokens = generate_native(model, tokenizer, ["The capital of France is"],
                         max_gen_len=32, use_sampling=False)
print(tokenizer.decode(tokens[0]))

Preview the full training config of a paper recipe. Copy common-base.example.json to base.json and fill in the data and tokenizer paths first:

python train_paper_part1.py --recipe dense_huginn --target 336b \
  --base-config base.json --dump-dir runs/dense-huginn --print-config

Reproducing the Papers

I want to... Start here
Train a Part I recipe Part I guide; entry point: train_paper_part1.py.
Export a trained Part I checkpoint for evaluation Part I guide; entry point: export_paper_part1.py.
Evaluate a Part I checkpoint Part I guide; entry point: eval_paper_part1.py.
Train a Part II recipe Part II guide; entry point: train_paper_part2.py.
Export a trained Part II checkpoint for evaluation or distillation Part II guide; entry point: export_paper_part2.py.
Distill a Part II prefill student Part II guide; entry point: distill_paper_part2.py.
Evaluate a Part II checkpoint Part II guide; entry point: eval_paper_part2.py.
Prepare the data Part I training and evaluation data; Part II training and evaluation data.

Checkpoints

All 49 checkpoints are in the Hugging Face collection Towards Looped Models Done Right. They use the native xLLM format and load with this code, not with Hugging Face Transformers.

Paper Repositories Contents
Part I IFM/LoopedLM-P1-<recipe>-<336b|500b> 25 checkpoints of the 17 recipes; list
Part II IFM/LoopedLM-P2-<name> 24 checkpoints: depth priors, injection variants, D4 and distilled students; list

Repository Layout

xLLM-Loop adds these paths to xLLM:

xllm-loop/
├── xllm/
│   ├── models/
│   │   ├── looped.py              # Looped Transformer and Huginn models
│   │   ├── looped_depth.py        # Recurrent-depth distributions
│   │   ├── looped_depth_prior.py  # Learned depth prior
│   │   └── looped_flops.py        # Training FLOPs
│   ├── paper_part1/               # Part I recipes, training, evaluation
│   └── paper_part2/               # Part II recipes, training, evaluation
├── release/
│   ├── paper-part1/               # Part I guide and weight license
│   └── paper-part2/               # Part II guide and weight license
├── papers/                        # Paper PDFs
├── train_paper_part1.py           # Part I: train
├── export_paper_part1.py          # Part I: export a checkpoint
├── eval_paper_part1.py            # Part I: evaluate
├── train_paper_part2.py           # Part II: train
├── export_paper_part2.py          # Part II: export a checkpoint
├── eval_paper_part2.py            # Part II: evaluate
└── distill_paper_part2.py         # Part II: distill a prefill student

Built on xLLM

The rest of the repository is xLLM: data loading, configuration, distributed training, evaluation and checkpointing. The xLLM README covers its general usage, including data preparation, training configuration, evaluation and checkpoint conversion. xLLM-Loop's additions to the shared training loop are opt-in: xLLM models train as they do in xLLM, and checkpoints additionally record per-rank RNG states.

Citation

@misc{huang2026loopedmodels,
  title  = {Towards Looped Models Done Right. Part I: Topology, Input Injection, Recurrent-State Organization},
  author = {Benhao Huang and Chufan Shi and Junlin Chen and Shicheng Wen and Zhengzhong Liu and Eric Xing and Xuezhe Ma},
  year   = {2026},
  url    = {https://huskydoge.github.io/husky-blog/posts/recursive_models/towards-looped-models-done-right/}
}

@article{huang2026fixedpoints,
  title   = {Towards Looped Models Done Right, Part II: Rethinking at Fixed Points},
  author  = {Benhao Huang and Chufan Shi and Junlin Chen and Shicheng Wen and Zhengzhong Liu and Eric Xing and Xuezhe Ma},
  journal = {arXiv preprint arXiv:2610.06833},
  year    = {2026}
}

License

xLLM-Loop is released under the Apache 2.0 License. Third-party code in the evaluation modules is listed with its licenses in THIRD_PARTY_NOTICES.md. The released weights carry their own licenses: Part I and Part II.

About

Official code for Towards Looped Models Done Right based on xLLM

Resources

Stars

41 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages