ByteDance-Seed/Modular-TTT

ByteDance-Seed/Modular-TTT

Language: Python

License: Apache-2.0

Stars: 12

Forks: 0

Open issues: 0

Created: 2026-06-22T06:36:17Z

Pushed: 2026-08-10T05:19:19Z

Default branch: main

Fork: no

Archived: no

README:

You can get to know us better through the following channels👇

!seed logo

Modular TTT

This repository contains the research code for Modular TTT: Rethinking Test-Time Training as Composable Modules.

Modular TTT treats test-time training (TTT) layers as composable inner-learning graphs rather than as separately hard-coded model variants. A TTT memory is represented as a directed acyclic graph whose nodes are primitive operations such as Linear, Act, Norm, Add, Mul, and loss functions. Given local train-view, train-view backward, and query-view rules for each primitive, the framework composes the full TTT computation automatically.

This code release includes the Modular TTT package, custom operators, model configs, training scripts, evaluation scripts, and the LLaMA/Gated DeltaNet baselines used for comparison.


Main Features

• Graph-structured TTT memory modules built from composable primitives.
• Local train-view, backward-transport, and query-view rules for TTT updates.
• Configurable loss functions, inner learning rates, decay, normalization, and

activation choices.

• Hugging Face-style TTTConfig, TTTModel, and TTTForCausalLM interfaces.
• Custom operators for efficient fast-weight update and query-view computation.
• Flame-based training, checkpoint conversion, and lm-eval-harness

evaluation scripts.

• RULER and per-token loss evaluation entry points.

Repository Layout

modular_ttt/ Core Modular TTT Python package.
xopes/ Custom operators used by Modular TTT and xmixers.
xmixers/ Token mixer implementations and baseline modules.
flame/ Training framework, model configs, and trainer code.
flame_script/env/ Environment setup script for this repository.
flame_script/flame_task/ Training, checkpoint conversion, and lm-eval scripts.
flame_script/ruler/ RULER evaluation scripts.
flame_script/per_token_loss/
Notes for per-token loss evaluation.
debug_ttt.sh Paper-related training and evaluation commands.

Installation

Use a Linux environment with a working PyTorch installation. From the repository root:

git clone https://github.com/ByteDance-Seed/Modular-TTT.git
cd Modular-TTT

bash flame_script/env/setup_env.sh

The setup script installs the local packages in this repository and the external Python dependencies used by the training and evaluation scripts, including flash-linear-attention, bitsandbytes, lm_eval, datasets, wandb, torchtitan, and torchdata.

If you run the setup script from a different working directory, set repo_dir explicitly:

repo_dir=/path/to/modular-ttt bash /path/to/modular-ttt/flame_script/env/setup_env.sh

Data Setup

Training data is not included in this repository. Set the data paths before launching training:

export DATA_10B_PATH=/path/to/your/10b_training_data
export DATA_100B_PATH=/path/to/your/100b_training_data

The scripts pass these paths to the Flame data loader. The optional DATASET_NAME variable defaults to default:

export DATASET_NAME=default

By default, experiment outputs, checkpoints, and caches are written under the repository directory. To keep outputs elsewhere:

export home_dir=/path/to/your/workdir

Quick Start

The main training entry point is flame_script/flame_task/script.sh. It trains a model, saves checkpoints, converts the target checkpoint to Hugging Face format, and runs lm-eval-harness.

The following command runs a short end-to-end workflow with 2 training steps and wikitext evaluation:

cd /path/to/modular-ttt
export DATA_10B_PATH=/path/to/your/10b_training_data

config_dir=modular_ttt/ttt_linear_graph_mse_no_norm_sd_sinlr_official_init

bash flame_script/flame_task/script.sh \
1 false \
${config_dir} \
ttt_linear_graph_mse_no_norm_sd_sinlr_90m_official_init \
2 1 128 false \
ttt_linear_mse_no_norm_sd \
2 1 3e-4 1 10b -1 none wikitext

Use the paper-style commands in debug_ttt.sh for the full experiment settings.

Reproducing Paper-Style Runs

debug_ttt.sh collects the paper-related commands for:

• Modular TTT scale-up variants.
• Modular TTT loss, decay, nonlinearity, and deep-memory ablations.
• LLaMA baseline runs.
• Gated DeltaNet baseline runs.

Set the data paths first, then open debug_ttt.sh and uncomment the command you want to run:

cd /path/to/modular-ttt
export DATA_10B_PATH=/path/to/your/10b_training_data
export DATA_100B_PATH=/path/to/your/100b_training_data

bash debug_ttt.sh

The script uses historical config names such as 90m, 310m, and 1_2b. These names refer roughly to the non-vocabulary backbone size: transformer blocks, token mixer, channel mixer, and norms. The paper reports total parameter counts, including token embeddings and the untied language-model head. Therefore the mapping used in the paper is:

| Config name | Paper scale | | --- | --- | | 90m | 160M | | 310m | 410M | | 1_2b | 1.45B |

Evaluation

The default flame_script/flame_task/script.sh path evaluates the converted checkpoint with lm-eval-harness. By default it covers:

• Perplexity: wikitext, lambada_openai
• Multiple choice: boolq, piqa, social_iqa, hellaswag, winogrande,

arc_easy, arc_challenge, openbookqa

• Containment-style tasks: fda, swde, squad_completion

You can override the task list with the final argument to script.sh, for example:

bash flame_script/flame_task/script.sh ... wikitext,lambada_openai

RULER commands are listed in:

flame_script/ruler/script.sh

The minimal per-token loss evaluator is in:

flame_script/per_token_loss/readme.md

Core Package Usage

After installing the repository dependencies:

from modular_ttt import TTTConfig, TTTForCausalLM

config = TTTConfig()
model = TTTForCausalLM(config)

For the standalone package details, see:

modular_ttt/README.md

Notes

• This repository releases code and configs, not training data or model

checkpoints.

• The primary implementation is modular_ttt, but efficient execution...