#jax

agentic RL 학습의 병목은 모델 계산보다 환경 대기 시간인 경우가 많다. Google Tunix는 asynchronous rollout과 producer-consumer pipeline으로 TPU idle time을 줄이는 JAX-native post-training library를 전면에 세웠다.

Google Cloud는 대규모 production model을 TensorFlow에서 JAX로 옮기는 작업에 specialized multi-agent AI를 투입해 이전 속도를 6x 높였다고 밝혔다. 핵심은 문법 변환이 아니라 stateful code를 JAX의 함수형·TPU 최적화 구조로 다시 짜는 문제다.

Hacker News에 오른 Nanocode는 tokenizer training, pretraining, synthetic data generation, agentic SFT, DPO를 pure JAX와 TPU workflow로 묶어 Claude Code 스타일 coding model을 재현하려는 end-to-end open project다.