19분 분량

TRM mac에서 학습시켜보기 프로젝트 for ARC-AGI

읽은 깊이
시작 전

TRM mac에서 학습시켜보기 프로젝트 for ARC-AGI

목표: “TRM을 내 맥북에서 돌려볼 수 있나?”에 어디까지 되고 어디부터 안 되는지 숫자로 답하고, 안 되는 부분을 우회하는 실행 계획을 세우기.

조사일 2026-10-08. 장비는 MacBook Pro(M1 Max, 10코어, 32GB), 디스크 여유 약 62GB. 배경은 GPU와 CUDA 입문, 원격 GPU 서버 사용 입문, 로컬 LLM 파일 형식과 정밀도 입문, 벤치마크는 ARC-AGI - Abstraction and Reasoning Corpus ⭐에 있다.

한 줄 결론

질문답
처음부터 사전학습안 된다. 이 맥에서 어림 수 년이 걸린다.
공개 체크포인트로 추론·평가된다. 모델이 7M이라 메모리 문제가 없다.
체크포인트에서 시작하는 미세조정작은 규모는 된다. 대회식 12,500 step은 어림 한 달이다.
모델 규모TRM은 약 7M 고정(2층, hidden 512). 키울 이유가 없다.
LLM으로 증류한 체크포인트찾지 못했다. 그런 논문도 공개물도 없다.
학습 코드공개돼 있다. MIT, 단 CUDA 전용이고 2026-04-01에 아카이브됐다.

1. TRM이 무엇인가

TRM(Tiny Recursive Model)은 Samsung SAIL Montréal의 Alexia Jolicoeur-Martineau가 낸 모델이다. 논문은 “Less is More: Recursive Reasoning with Tiny Networks”(arXiv 2510.04871)다. 2층짜리 작은 네트워크 하나를 같은 가중치로 여러 번 되풀이해서 답을 다듬는다. 논문이 보고한 점수는 ARC-AGI-1 약 45%, ARC-AGI-2 약 8%다.

비유하면 작은 연필 한 자루로 같은 답안지를 여러 번 고쳐 쓰는 방식이다. 큰 LLM처럼 한 번에 글을 이어 쓰지 않는다. 그리드 전체(30×30 = 900칸)를 한 번에 입력받고 한 번에 답 그리드를 낸다. 이 점이 아래 계산에서 중요하다.

한 번의 학습 step 안에서 일어나는 일은 코드(models/recursive_reasoning/trm.py)에서 확인했다.

  • H_cycles(바깥 반복) 번 도는데, 마지막 한 번만 기울기를 계산하고 나머지는 torch.no_grad() 안에서 앞으로만 계산한다.
  • 한 바깥 반복마다 L_cycles(안쪽 반복)에 1을 더한 횟수만큼 2층 네트워크를 통과한다.
  • README 명령어 기준 H_cycles=3, L_cycles=4이므로 한 step에 네트워크를 15번 지나간다. 이 중 5번에 역전파가 붙는다.
  • 설정 파일 config/arch/trm.yaml의 기본값은 L_cycles: 6이다. README 명령어는 4로 덮어쓴다. 값이 서로 다르니 실행 전에 확인한다.

2. 내 맥북 사양과 실측

항목값
칩Apple M1 Max, 10코어(성능 8 + 효율 2)
메모리32GB 통합 메모리
GPU 백엔드PyTorch MPS (macOS 26)
디스크약 62GB 여유

TRM과 같은 크기(2층, hidden 512, SwiGLU, 파라미터 6.82M)의 블록을 만들어 MPS에서 속도를 쟀다. 입력은 배치 16 × 900토큰이고 앞으로 + 뒤로 계산이다.

정밀도1 iteration유효 연산 속도처리량
fp32461 ms약 1.62 TFLOPS약 31,200 토큰/초
fp16403 ms약 1.85 TFLOPS약 35,800 토큰/초
bf16446 ms약 1.67 TFLOPS약 32,300 토큰/초

정밀도를 바꿔도 속도가 거의 같다. 이 맥에서는 bf16·fp16으로 얻는 속도 이득이 없다. 코드는 기본이 forward_dtype: bfloat16이지만 맥에서는 fp32로 돌려도 손해가 크지 않다.

3. 사전학습: 안 된다

근거 1. 원 저자의 학습 규모. README는 ARC-AGI 학습을 “H100 4장으로 약 3일”이라고 쓴다. 다른 논문(McGovern, arXiv 2511.02886)은 같은 사전학습을 “H100 SXM 4장, 48시간, 옵티마이저 step 70만 회 이상”이라고 적는다. 체크포인트 이름 step_518071(ARC Prize 공개본, ARC-AGI-1)도 같은 규모다.

근거 2. 이 맥으로 환산한 어림 계산.

  • 한 sample을 한 step 학습할 때 실측으로 약 0.24초가 든다. 계산 방식: 기울기 있는 5번 통과는 5 × (900토큰 ÷ 31,200토큰/초) ≈ 0.144초, 기울기 없는 10번 통과는 앞으로만이라 약 1/3 비용으로 ≈ 0.096초.
  • 전역 배치 768로 518,071 step이면 sample-step이 약 3.98억 개다.
  • 3.98억 × 0.24초 ≈ 9.5×10⁷초 ≈ 약 3년이다.

이 숫자는 어림이다. 적응형 종료(ACT)가 평균 몇 step에서 멈추는지, 데이터 로딩 비용, torch.compile 이득을 반영하지 않았다. 그래도 자릿수(년 단위)는 바뀌지 않는다.

결론. 사전학습은 클라우드 GPU 몫이다. 맥은 사전학습 과정을 작게 흉내 내 보는 용도로만 쓴다.

4. 미세조정: 작은 규모는 된다

4-1. 미세조정이 실제로 어떻게 쓰이는가

ARC Prize 2025 대회 참가자(McGovern, arXiv 2511.02886)는 사전학습된 TRM에서 시작해 대회 과제로 12,500 gradient step을 더 학습했다. 방식은 전체 가중치 미세조정(full fine-tuning) 이다. 논문 원문 표현은 “full fine-tuning of the tiny model, not LoRA fine-tuning or fine-tuning of task embeddings alone”다. 대회 연산 한도(L4 4장, 약 12시간)가 이유다. 세미프라이빗 점수는 6.67%였다.

따라서 TRM은 LoRA 같은 부분 학습보다 통째로 학습하는 쪽이 이미 검증된 길이다. 모델이 7M이라 통째로 학습해도 메모리 부담이 작다.

4-2. 맥에서 되는 범위

단계내용이 맥에서
A공개 체크포인트로 추론·평가가능
B과제 수십수백 개, 배치 3264, 수천 step가능. 어림 수 시간~수 일
C대회식 12,500 step, 배치 768비현실적. 어림 약 27일, 메모리도 모자랄 것
D처음부터 사전학습불가
  • C의 27일은 12,500 × 768 × 0.24초 ÷ 86,400로 계산한 어림이다.
  • 배치 768은 활성값(역전파용 중간 결과) 때문에 32GB로 모자랄 가능성이 크다. 이 부분은 직접 돌려 보기 전에는 모른다. 우선 배치 32~64에서 시작해 torch.mps.current_allocated_memory()로 재 보는 것이 안전하다.
  • 단계 B는 “ARC-AGI-2 전체를 푸는 모델”이 아니라 특정 과제 묶음에서 점수가 오르는지 보는 실험이다.

4-3. 쓸 수 있는 파라미터 규모

TRM은 7M(2층, hidden 512, 헤드 8, expansion 4)이 논문과 모든 공개 체크포인트의 공통 크기다. 크기를 키우면 공개 체크포인트에서 시작할 수 없다. 맥의 병목은 모델 크기가 아니라 배치 × 900토큰 × 재귀 횟수에서 나오는 활성값과 연산량이다. 그래서 “더 큰 모델”이 아니라 “더 작은 배치·더 짧은 재귀”를 조절한다.

참고로 같은 ARC를 LLM으로 풀려는 방식(예: Llama 3 8B QLoRA)과의 비교에서, TRM은 비자기회귀 구조 덕에 처리량이 훨씬 높고 메모리가 훨씬 적다는 보고가 있다(arXiv 2512.11847). 맥북에 맞는 쪽은 LLM이 아니라 TRM이다.

5. 공개 체크포인트

체크포인트대상보고 점수크기비고
arcprize/trm_arc_prize_verificationarc_v1_public(ARC-AGI-1), arc_v2_public(ARC-AGI-2)ARC-1 40%, ARC-2 6.2%각 1.8GB, 2.5GBMIT. ARC Prize가 검증용으로 공개. 8×H100, 실행당 20~30시간
alphaXiv/trm-model-arc-agi-1ARC-AGI-1독립 재현 약 43% (검색 요약만 확인, 원문 미확인)각 1.8GBMIT
Trelis/TRM-pretrain_arc2_rearc_200kARC-AGI-2 + Re-ARC공개 평가 약 15% (커뮤니티 보고)step당 약 1.75GB, 데이터셋 포함MIT. 이어서 학습 가능
seconds-0/trm-arc2-8gpuARC-AGI-2공개 평가 pass@1 약 0.83%미확인재현·분석용 (검색 요약만 확인)

점수를 읽을 때 주의할 점.

  • Trelis 15%는 학습 데이터에 ARC-AGI-2 평가용 120개 과제의 예시 쌍이 들어 있다. 같은 평가셋에서 점수를 내는 방식이라 논문의 8%, ARC Prize 검증본의 6.2%와 같은 선에서 비교하면 안 된다.
  • 체크포인트 하나가 1.8~2.5GB인 이유는 확인하지 못했다. 모델 본체는 7M(fp32로 약 27MB)이라 나머지는 과제별 임베딩(puzzle embedding) 테이블이 대부분일 것으로 어림한다. 이것은 추정이다.
  • 체크포인트는 과제 ID에 강하게 묶여 있다. 올바른 퍼즐 ID를 빈 토큰이나 무작위 토큰으로 바꾸면 정확도가 0이 된다(arXiv 2512.11847). 새 과제를 풀게 하려면 그 과제용 임베딩부터 학습해야 한다.
  • 성능의 상당 부분은 모델이 아니라 테스트 때 증강 1,000개를 만들어 다수결을 내는 데서 나온다(같은 논문, Pass@1 약 +11%p). 맥에서 평가할 때도 이 단계가 시간을 가장 먹는다.

6. 학습 코드 공개 현황

SamsungSAILMontreal/TinyRecursiveModels

  • 라이선스 MIT. 별 약 6,500개, 포크 약 1,040개.
  • 2026-04-01에 아카이브(읽기 전용)됐다. 새 수정은 받지 않는다. 직접 포크해서 고쳐야 한다.
  • README의 사전학습 체크포인트 링크는 없다. 체크포인트는 위 표의 제3자 저장소에서 구한다.
  • ARC-AGI 데이터셋 구축은 python -m dataset.build_arc_dataset, 학습은 pretrain.py(hydra 설정)다.
  • ARC-AGI-2 학습 데이터에는 ARC-AGI-1 평가 데이터가 일부 섞여 있다. README는 두 버전을 함께 학습하고 둘 다 평가하면 안 된다고 경고한다.
  • 연구용 포크: TrelisResearch/TinyRecursiveModels(Trelis 카드가 가리키는 저장소, 미확인).

맥에서 막히는 부분

공식 코드는 NVIDIA GPU 전제다. pretrain.py와 의존성을 직접 읽어 확인한 결과다. 맥 포팅(MPS 포크)은 검색과 GitHub 검색 모두에서 찾지 못했다.

위치문제해결 방향
pretrain.py의 torch.device("cuda"), .cuda(), map_location="cuda"장치가 cuda로 박혀 있다mps 장치 선택 함수로 교체
from adam_atan2 import AdamATan2CUDA 확장을 빌드하는 패키지다순수 PyTorch로 atan2 업데이트를 직접 구현
requirements.txt의 triton, ninja맥에서 설치가 안 되거나 쓸 수 없다제거
torch.compile(model)MPS에서 불안정하다환경변수 DISABLE_COMPILE를 설정하면 코드가 건너뛴다
dist.init_process_group(backend="nccl")torchrun으로 띄울 때만 실행된다그냥 python pretrain.py로 단일 프로세스 실행
models/sparse_embedding.py의 dist.all_gather_into_tensor단일 프로세스에서 초기화 안 된 dist를 부를 수 있다분산이 아닐 때 건너뛰도록 분기 추가

adam-atan2 교체가 가장 큰 작업이다. 나머지는 줄 수정이다. README에 flash-attn 언급은 없고 models/layers.py에서도 주석 처리돼 있어 이것은 걸림돌이 아니다.

환경변수 PYTORCH_ENABLE_MPS_FALLBACK=1을 켜 두면 MPS가 지원하지 않는 연산이 CPU로 넘어가서 멈추지 않는다. 대신 그 연산은 느려진다.

7. LLM으로 증류한 체크포인트: 없다

찾은 범위에서는 LLM 교사로 TRM을 증류한 공개 논문도 체크포인트도 없다. 이유를 구조로 보면 이렇다. TRM은 단어가 아니라 0~9 색 숫자 그리드 토큰을 다룬다. LLM의 출력 분포와 어휘가 맞지 않아서 로짓을 그대로 가르칠 수 없다. 이 문장은 내 추론이고 문헌이 말한 것은 아니다.

가장 가까운 것은 TRM-Planner(ACL 2026 Findings, Baek 등)다. 교사가 LLM이 아니라 얼려 둔 TRM 체크포인트 자신이다.

  • 교사 TRM을 여러 번 돌려 토큰·로짓·step·품질 정보를 캐시에 저장한다(오프라인 단계).
  • 학생 TRM을 기본 손실에 증류 손실을 더해 학습한다.
  • 보고 결과: ARC-1 pass@2가 43.1%에서 48.1%, ARC-2가 6.7%에서 9.2%로 올랐다.
  • 이 정보는 검색 요약 두 번으로만 확인했다. 논문 전문과 코드 공개 여부는 못 읽었다(ACL 미리보기 링크는 404).

TRM의 수치를 올린 다른 갈래는 LLM 증류가 아니라 ARC 학습용 합성 데이터(Re-ARC 등)와 강한 증강이다. 맥에서 LLM 증류를 시험해 보고 싶다면 TRM-Planner의 교사-캐시 구조에서 교사를 LLM으로 바꾸는 방향을 직접 설계해야 한다.

8. 실행 계획

Phase 0. 환경. python3 -m venv로 가상환경을 만들고 torch만 설치한다. 이미 확인한 대로 pip install torch가 맥에서 MPS 지원 빌드를 받는다. 저장소는 포크해서 clone한다.

Phase 1. 포팅. 위 걸림돌 표를 순서대로 고친다. 목표는 DISABLE_COMPILE=1 PYTORCH_ENABLE_MPS_FALLBACK=1 python pretrain.py ...가 에러 없이 한 step 도는 것이다. 먼저 Sudoku-Extreme(시퀀스 81칸)나 아주 작은 ARC 데이터로 확인한다. 이 데이터셋은 시퀀스가 900칸의 약 1/11이라 샘플당 비용이 훨씬 작다(어림).

Phase 2. 공개 체크포인트로 추론·평가. arcprize/trm_arc_prize_verification의 arc_v1_public을 받아서 평가 과제 몇 개를 풀게 한다. 보고 점수(ARC-1 40%)가 소수 과제에서도 대략 재현되는지 본다. 증강 개수를 줄여(예: 1,000개 대신 32개) 시간을 먼저 쟀다.

Phase 3. 소규모 미세조정. 같은 체크포인트에서 ARC 과제 일부로 배치 32~64, 수천 step을 학습한다. 비교 대상은 같은 과제를 미세조정 전 체크포인트로 푼 점수다. 기록할 값은 step당 시간, 메모리 최대치, 과제별 정답률이다.

Phase 4(선택). 증류 실험. TRM-Planner식 교사 캐시를 만들어 본다. 교사는 공개 체크포인트(arc_v1_public)로 둔다.

각 단계 끝에 이 노트의 숫자(0.24초/sample-step, 약 3년, 약 27일)를 실측으로 갈아 끼운다.

9. 아직 확인하지 못한 것

  • 실제 학습 step 속도와 메모리. 위의 모든 시간은 블록 하나의 실측에서 곱해 낸 어림이다. 데이터 로더, ACT 평균 step 수, 평가 비용을 넣지 않았다.
  • 배치 768이 32GB에서 되는지.
  • alphaXiv 독립 재현(약 43%), seconds-0/trm-arc2-8gpu, Trelis 포크의 세부는 검색 요약만 봤다. 페이지 원문은 열지 않았다.
  • TRM-Planner의 코드 공개 여부와 학습 비용.
  • McGovern 논문(2511.02886)의 코드 공개 여부. 초록 페이지에는 코드 링크가 없었다.
  • 체크포인트 1.8GB의 내부 구성.

출처