| # LOOM: Looping Beyond Twice |
|
|
| Code for **Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts** |
| (Di He, Pengxiang Li, Da Chang, Qingyan Meng, Lu Yin, Shiwei Liu). [arXiv:2610.01153](https://arxiv.org/abs/2610.01153) |
|
|
| A looped Transformer applies the same blocks several times, which adds depth without adding parameters. |
| In MoE language models, the gain from extra loops usually stops after about two loops. |
| LOOM is a training recipe for looped MoE models: each loop should add new computation, and the recurrent state should stay stable. |
| With LOOM, models from 100M to 1.7B parameters train stably with 9 to 12 loops. |
|
|
| <p align="center"> |
| <img src="assets/weave.gif" width="100%" alt="Nine loops weave through 30 experts; one token uses 14 distinct experts, and the first 48 tokens of the sentence form a cloth."> |
| </p> |
|
|
| Each loop uses different experts. The 30 vertical threads are the routed experts of layer 8 in the 1.7B LOOM model with 9 loops, and each row is one loop. |
| Where the loop's router picks an expert, the row passes over that thread (coloured); everywhere else, it passes under. |
| Over 9 loops, the token "guitar" uses 14 distinct experts, where loops that reused the same experts would use 6. |
| Averaged over 16,384 held-out tokens and all 15 layers, the count is 14.6, and two consecutive loops share about 3 of their 6 experts. |
| At the end, the view pulls back to the same weave for the first 48 tokens of the sentence. |
| The routing is measured by `tools/animations/collect_routing.py`, and `tools/animations/weave.py` draws the animation. |
|
|
| | Component | What it does | Config | |
| |---|---|---| |
| | Residual scaling | scales each loop's residual update so hidden-state variance stays bounded as loops are added | `arch.cycle_residual_scale`, `arch.cycle_scale_lambda` | |
| | Embedding re-injection | mixes the input embedding back into the state at the start of every loop | `arch.embed_inject_mode=loop_start_mix_soft`, `arch.embed_inject_*` | |
| | Per-loop routers | a separate router for each loop, so different loops can use different experts | `arch.moe.router_per_loop`, `arch.moe.router_num_loops` | |
| | Looping Residual | carries earlier loops' outputs forward into later loops | `arch.attn_res=true`, `arch.attn_res_scope=dual_axis`, `arch.dual_axis_*` | |
| | MoE output RMSNorm | RMSNorm on the MoE branch output before the residual add | `arch.ffn_branch_rms=affine` | |
| | Segmented backprop | backward and optimizer step every K loops (K = 3) | `arch.bp_segment_len` | |
|
|
| All of these are set in `scripts/train.sh`. |
|
|
|
|
| ## Install |
|
|
| ```bash |
| # Python 3.12, CUDA 12.x. Install torch and flash-attn for your CUDA version first. |
| pip install -r requirements.txt |
| ``` |
|
|
| ## Data |
|
|
| Training reads a pre-tokenized, document-packed corpus. To build one from FineWeb-Edu with the SmolLM tokenizer (vocabulary 49,152): |
|
|
| ```bash |
| python tools/pack_document_data.py \ |
| --parquet 'fineweb-edu/sample/10BT/*.parquet' \ |
| --tokenizer HuggingFaceTB/SmolLM-360M \ |
| --out data/fineweb-edu-1024 \ |
| --seq-len 1024 --epochs 9 --holdout-tokens 30000000 --workers 32 |
| ``` |
|
|
| `epoch_0` to `epoch_8` are different shuffles of the training documents, and `epoch_9` is a held-out set of documents that never appear in training. |
| `pretrain.py` evaluates on `epoch_9`, or on `epoch_0` of a separate pack given as `data.val_path` (`VAL_PATH`). |
| `--jsonl` reads `.jsonl` / `.jsonl.zst` shards (for example Dolma), and `tools/merge_document_packs.py` merges packs that use the same tokenizer. |
|
|
| ## Train |
|
|
| `H` is the number of loops. The two presets reproduce the paper's two scales: |
|
|
| ```bash |
| # 329M total parameters: L=10, hidden 512, 30 routed + 2 shared experts, top-8; 10k steps of 1M tokens |
| H=9 DATA_PATH=data/fineweb-edu-1024 bash scripts/loom_329m.sh |
| |
| # 1.7B total parameters: L=15, hidden 1280, same expert layout; 57k steps of 1M tokens (about 60B tokens) |
| H=9 DATA_PATH=data/<corpus>-1024 VAL_PATH=data/<val>-1024 bash scripts/loom_1p7b.sh |
| ``` |
|
|
| - The global batch is 1024 sequences of 1024 tokens. The world size must divide 1024, and samples per rank must be a multiple of `MICRO_BATCH_SAMPLES`. |
| - MoE capacity is computed per micro-batch, so the micro-batch size changes the computation. The presets use the paper values (16 for 329M, 8 for 1.7B). |
| - Multi-node: run the same command on every node with `NNODES`, `NODE_RANK`, `MASTER_ADDR` and `MASTER_PORT` set. |
| - Quick test on one GPU: `H=3 DATA_PATH=... TOTAL_STEPS=20 GLOBAL_BATCH_SIZE=16 MICRO_BATCH_SAMPLES=8 bash scripts/loom_329m.sh`. |
|
|
| Ablations turn off one component of the recipe: |
|
|
| | Ablation | Setting | |
| |---|---| |
| | no residual scaling | `CYCLE_RESIDUAL_SCALE=false` | |
| | no embedding re-injection | `EMBED_INJECT_MODE=none` | |
| | no MoE output RMSNorm | `FFN_BRANCH_RMS=none` | |
| | no segmented backprop (one backward through all loops) | `BP_K=0` | |
|
|
| ### Checkpoints |
|
|
| Runs are written to `checkpoints/<project>/<run>/` (`CHECKPOINT_PATH` to change). |
|
|
| - `fsdp2_epoch_N/` is the live checkpoint. It is overwritten every `CKPT_EVERY` steps (default 500). |
| - `snapshots/step_NNNNNNN/` are full copies, saved every `SNAPSHOT_EVERY_STEPS` steps. Only the newest `SNAPSHOT_KEEP_LAST` are kept. |
| - To continue a run, set `RESUME_FROM=<run dir or snapshot dir>`. |
| - Each checkpoint stores model and Adam state together in torch DCP format. `tools/export_model_only_dcp.py` writes a model-only copy, about one third of the size, and checks it bit for bit. |
|
|
| Training logs go to TensorBoard under the run directory, and every scalar is also saved to `tensorboard_scalars.json`. |
|
|
| ## Evaluate |
|
|
| ```bash |
| bash eval/run.sh checkpoints/<project>/<run> |
| ``` |
|
|
| This runs the seven 0-shot tasks from the paper (ARC-E, ARC-C, HellaSwag, OpenBookQA, PIQA, Winogrande, SIQA) with lm-eval 0.4.11. |
| Before scoring, it checks that the loaded weights reproduce the `eval/loss` logged during training. |
| See [eval/README.md](eval/README.md). |
|
|
| ## Layout |
|
|
| ```text |
| pretrain.py training loop (FSDP2, WSD schedule, segmented backprop, spike guard) |
| models/baselines/looped_transformer.py looped MoE Transformer |
| models/moe.py MoE layer (shared + routed experts, sigmoid router, per-loop routers) |
| models/dual_axis_carry.py Looping Residual |
| models/attn_res.py Attention Residuals (arXiv:2603.15031); attn_res_scope=dual_axis uses dual_axis_carry.py instead |
| config/ Hydra configs (arch/size: T_moe_L10 = 329M, T_moe_L15_E32 = 1.7B) |
| dataset_document.py document-packed dataset |
| scripts/ launchers |
| tools/ data packing, model-only export |
| eval/ 0-shot evaluation |
| ``` |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{he2026loom, |
| title = {Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts}, |
| author = {He, Di and Li, Pengxiang and Chang, Da and Meng, Qingyan and Yin, Lu and Liu, Shiwei}, |
| journal = {arXiv preprint arXiv:2610.01153}, |
| year = {2026} |
| } |
| ``` |
|
|
| ## License |
|
|
| Apache-2.0. See [LICENSE](LICENSE). |
|
|
