Tiny Recursive Model (TRM)を試しに走らすだけのメモ

Oct. 6th 2025にarXivに投稿された次のpaper。

Less is More: Recursive Reasoning with Tiny Networks
Alexia Jolicoeur-Martineau
https://arxiv.org/abs/2510.04871

ソースコードが公開されていますので、試しました。
https://github.com/SamsungSAILMontreal/TinyRecursiveModels

gitコマンドでダウンロードして、venv環境をつくり、そこに入ります。

$ git clone https://github.com/SamsungSAILMontreal/TinyRecursiveModels.git
$ cd TinyRecursiveModels
$ python3 -m venv ./venv
$ source ./venv/bin/activate

あとはreadmeの通り requirementsをいれます。

$ pip install -r requirements.txt # install requirements
$ pip install --no-cache-dir --no-build-isolation adam-atan2 

データセットをダウンロードします。今回はテストなので一番軽いSudoku-Extremeで。

$ python dataset/build_sudoku_dataset.py --output-dir data/sudoku-extreme-1k-aug-1000  --subsample-size 1000 --num-aug 1000  # 1000 examples, 1000 augments

で次のコマンドで学習を実行しようとしますがエラーがでます。

$ run_name="pretrain_mlp_t_sudoku"
python pretrain.py \
arch=trm \
data_paths="[data/sudoku-extreme-1k-aug-1000]" \
evaluators="[]" \
epochs=50000 eval_interval=5000 \
lr=1e-4 puzzle_emb_lr=1e-4 weight_decay=1.0 puzzle_emb_weight_decay=1.0 \
arch.mlp_t=True arch.pos_encodings=none \
arch.L_layers=2 \
arch.H_cycles=3 arch.L_cycles=6 \
+run_name=${run_name} ema=True

ModuleNotFoundError: No module named 'adam_atan2_backend'

HRMのほうのissuesに解決法がありました。
https://github.com/sapientinc/HRM/issues/25

代わりにadam-atan2-pytorchを使うようです。

$ pip install adam-atan2-pytorch
$ sed -i 's/adam_atan2/adam_atan2_pytorch/g' pretrain.py
$ sed -i 's/AdamATan2/AdamAtan2/g' pretrain.py
$ sed -i 's/lr=0,/lr=0.0001,/g' pretrain.py

sedでpretain.pyのadam_atan2を書き換えます。学習率の初期値は0のままだと別のエラーがでました。ので学習率の初期値も適宜設定します。

これでもう一度学習を実行すると走ります。

(もし、fatal error: Python.h: No such file or directoryという別のエラーが出た場合は、sudo apt install python3-devをしてPython.hを入れてください。)

(wandbを使っていませんが毎回どうするか聞かれます。コード削ればよいのですが)

wandb: (1) Create a W&B account
wandb: (2) Use an existing W&B account
wandb: (3) Don't visualize my results
wandb: Enter your choice: 3

L40S GPU 1つでRuntime: < 36 hoursだそうです。

うちだとRTX5090 1つで、6時間くらいかかっています。arch.mlp_t=Trueで。

先程の実行コードでは50000 epochの学習で5000 epochごとにevalが入ります。

VRAMは 32 GBギリギリまで使っていて、負荷もフル575Wです。gpuは効率よく利用できているようです。

mlpが入らないほうだとVRAMは21 GB使用して、4h30min程度で走りそうです。4.00 it/s

21 GBだとRTX4090でもそのまま走ります(mlp入りはout of memoryになった)。こちらもやってみると、2.76it/sで6h30minかかるということです。RTX5090のほうが44%速いです。

RTX5090+RTX4090両方を使うと、5.25 it/sでした。RTX5090の速度からすると2枚使う割に1.3倍にしかなっておらず遅い気がしますが、RTX4090で律速すると考えると5.25/2.76 = 1.9なので、ほぼ4090 2枚分の速度ということで悪くはないです。RTX4090が遅いということで…(1~2年前までは爆速のイメージでしたが)。