TPU가 tool 호출을 기다리지 않게… Google Tunix의 agentic RL 병목 제거

multi-turn agent를 학습시키는 비용 문제는 모델이 생각하는 시간보다 tool 호출, web search, code 실행, 환경 step을 기다리는 시간에서 커진다. Google Tunix는 이 병목을 정면으로 겨냥한 JAX-native post-training library다. Google은 Tunix의 최신 release가 reasoning agent 학습에서 TPU가 host-side 환경 대기 때문에 멈춰 서는 문제를 asynchronous rollout과 decoupled training pipeline으로 줄인다고 설명했다.

기존 synchronous rollout 구조에서는 한 trajectory가 환경 초기화나 reward 반환을 기다리면 accelerator에도 빈 시간이 생긴다. batch 안에 느린 trajectory가 하나만 있어도 전체 latency가 그쪽에 묶이는 straggler effect도 생긴다. Tunix는 Python asyncio 기반 RolloutOrchestrator로 많은 agent-environment interaction을 동시에 관리한다. 한 agent가 tool 실행을 기다리는 동안 inference engine은 다른 trajectory의 token 생성을 처리한다. async vLLM-TPU와 SGLang-Jax 통합도 이 설계의 일부다.

두 번째 축은 rollout과 training의 분리다. naive pipeline은 batch 전체가 끝날 때까지 trainer가 기다릴 수 있다. Tunix는 producer-consumer 구조로 완료된 trajectory를 계속 queue에 넣고, AgenticRLLearner가 이를 소비한다. GRPO처럼 prompt당 여러 reasoning path가 필요한 알고리즘에서는 asynchronous trajectory를 동적으로 묶어 post-processing과 scoring을 거쳐 trainer로 흘려보낸다. 목표는 synchronous trainer를 계속 먹여 살려 end-to-end throughput을 높이는 것이다.

개발자 경험도 agentic RL 쪽으로 맞췄다. Tunix는 ModelAgent, ToolAgent, TaskEnvironment, ToolEnvironment 같은 기본 구성요소를 제공하고, SWE-bench, WebArena, game engine, custom environment를 붙일 수 있는 API 경계를 둔다. 표준 profiler가 짧은 trace에 강한 반면, Tunix는 rollout, training, weight sync, tool latency 같은 RL-specific metric을 가볍게 계속 추적한다. 이는 agent 학습이 “잘 안 돈다”는 감각을 TPU starvation, data generation, micro-batch 설정 같은 조정 지점으로 바꾸는 도구다. 다음 관전 포인트는 repo 채택 속도와, Gemma·Qwen 계열 fine-tuning recipe가 실제 agent benchmark 개선으로 이어지는지다.