🏗️ 에이전트 기반 정렬(Agentic Alignment)의 시스템 병목 현상

정적인 인간 피드백 기반 강화학습 (RLHF)에서 능동적인 에이전트 기반 강화학습 (RL)으로의 전환은 대규모 언어 모델 (LLM)이 정렬되는 방식의 중대한 변화를 의미합니다. 에이전트 기반 RL에서 언어 모델은 단순히 프롬프트에 대해 단일 응답을 출력하는 것에 그치지 않습니다. 모델은 환경 내에서 에이전트로 동작하며, 다단계 추론을 수행하고, 외부 API를 호출하며, 샌드박스 환경에서 코드를 실행하고, 환경의 피드백에 동적으로 반응합니다.

이러한 에이전트 루프는 모델의 문제 해결 능력을 향상시키지만, 심각한 시스템 엔지니어링 병목 현상을 초래합니다. 전통적인 RL 훈련 파이프라인은 동기식 (synchronous)입니다. 모델이 일련의 롤아웃 (rollouts) 배치를 생성하면, 환경이 해당 액션들을 실행하고, 보상 (rewards)이 계산되며, 모델이 가중치를 업데이트합니다. 만약 환경에 Python 코드 실행, 웹 페이지 로딩 대기, 또는 외부 데이터베이스 쿼리와 같이 호스트 측의 느린 작업이 포함된다면, 값비싼 TPU (Tensor Processing Units)는 완전히 유휴 상태로 방치됩니다. 이러한 스텝 락 (step-lock) 문제는 낮은 모델 플롭스 활용도 (MFU, Model Flops Utilization)를 초래하며, 종종 가속기 효율을 10% 미만으로 떨어뜨립니다.

이 병목 현상을 해결하기 위해 Google은 LLM 사후 훈련 (post-training) 및 고처리량 RL 오케스트레이션을 위해 특별히 설계된 JAX 네이티브 라이브러리인 Tunix를 개발했습니다. 저는 Tunix가 궤적 생성 (trajectory generation)을 정책 최적화 (policy optimization)로부터 분리 (decoupling)함으로써 어떻게 에이전트 기반 RL을 확장하는지 분석했습니다. 이 글에서 저는 Tunix의 비동기 아키텍처를 살펴보고, 동적 워크로드 하에서 TPU 메모리와 JAX 컴파일을 어떻게 오케스트레이션하는지 상세히 설명하며, 오프-정책 발산 (off-policy divergence)의 수학적 과제를 다루고, 에이전트 기반 훈련 파이프라인을 확장하고자 하는 엔지니어링 리더들을 위한 구체적인 구현 청사진을 제공합니다.

JAX 네이티브 오케스트레이션 (JAX-native orchestration)을 사용하여 궤적 생성 (trajectory generation)을 정책 최적화 (policy optimization)로부터 분리함으로써, Google의 Tunix 라이브러리가 에이전트 기반 RL (agentic RL)의 스텝-락 (step-lock) 병목 현상을 어떻게 해결하는지에 대한 심층적인 기술 분석입니다.

🏗️ 분리된 궤적 생성의 아키텍처 (The Architecture of Decoupled Trajectory Generation)

표준적인 동기식 액터-러너 (synchronous actor-learner) 아키텍처에서는 학습 루프가 가장 느린 구성 요소에 의해 제한됩니다. 만약 에이전트가 도구 호출 (tool call)을 실행하고 응답을 받는 데 5초가 걸린다면, 전체 TPU 클러스터는 그 5초 동안 차단(blocked)된 상태로 유지됩니다. 높은 처리량 (throughput)을 달성하기 위해서는 이러한 동기식 실행 모델을 깨뜨려야 합니다. Tunix는 궤적 생성 (the actors)을 정책 최적화 (the learner)로부터 완전히 분리함으로써 이를 달성합니다.

Tunix 프레임워크에서 액터 (actors)와 러너 (learners)는 종종 이기종 하드웨어 (heterogeneous hardware) 상에서 독립적인 비동기 프로세스로 실행됩니다. 액터—정책 모델을 실행하여 행동을 생성하고 환경과 상호작용하는 역할—는 더 저렴한 GPU 인스턴스나 네트워크 및 I/O 작업에 최적화된 CPU 풀 (CPU pools)로 분산될 수 있습니다. 러너—모델 가중치를 업데이트하기 위해 계산 집약적인 순전파 (forward pass) 및 역전파 (backward pass)를 수행하는 역할—는 전용 TPU Pod 슬라이스 (TPU Pod slice)에서 지속적으로 실행됩니다.

이 두 구성 요소는 중앙 집중식의 고처리량 궤적 버퍼 (high-throughput trajectory buffer)를 통해 통신합니다. 이 분리된 아키텍처에서 궤적의 라이프사이클 (lifecycle)은 다음과 같이 진행됩니다:

• 비동기 롤아웃 (Asynchronous Rollout): 액터 (actor)는 (러너 (learner)로부터 주기적으로 동기화된) 현재 정책 파라미터 (policy parameters)를 쿼리하여 액션 (actions)을 생성합니다. 이들은 병렬 환경 (parallel environments)에서 이러한 액션을 실행하고, 전이 (transitions: 상태 (state), 액션 (action), 로그 확률 (log probabilities), 보상 (reward), 다음 상태 (next state))를 수집하여 표준화된 궤적 텐서 (trajectory tensors)로 패킹합니다.
• 버퍼 인제스션 (Buffer Ingestion): 이러한 궤적들은 Tunix가 관리하는 분산 인메모리 리플레이 버퍼 (distributed, in-memory replay buffer)로 비동기적으로 스트리밍됩니다. 버퍼는 중복 제거 (deduplication), 시퀀스 패킹 (sequence packing), 그리고 우선순위 지정 (prioritization)을 처리합니다.
• 연속 학습 (Continuous Learning): 액터 상태와 완전히 독립적으로, TPU 기반의 러너는 버퍼에서 궤적 배치 (batches of trajectories)를 지속적으로 가져와 경사 하강법 (gradient descent)을 수행하고 글로벌 정책 파라미터를 업데이트합니다.
• 파라미터 동기화 (Parameter Synchronization): 정기적이고 설정 가능한 간격으로, 업데이트된 파라미터가 액터들에게 다시 브로드캐스트 (broadcast)됩니다.

이러한 루프들을 디커플링 (decoupling)함으로써, 러너는 환경 스텝 (environment steps)을 위해 기다리지 않습니다. 환경 지연 시간 (latency)이 급증하더라도, 러너는 (안전한 신선도 임계값 (safe staleness threshold) 내에서) 버퍼링된 궤적을 사용하여 계속 학습을 진행하며, 100%에 가까운 TPU 활용률을 유지합니다. 반대로 러너가 액터보다 빠를 경우, 러너 코드를 수정하거나 TPU 클러스터 토폴로지 (topology)를 변경하지 않고도 액터 풀 (actor pool)을 수평적으로 확장할 수 있습니다.

TPU 오케스트레이션 (Orchestration) 및 JAX 네이티브 메모리 관리

JAX에서 이러한 디커플링된 아키텍처를 오케스트레이션하려면 몇 가지 독특한 컴파일러 및 메모리 과제를 극복해야 합니다. JAX는 XLA (Accelerated Linear Algebra)를 통한 사전 컴파일 (Ahead-Of-Time (AOT) compilation)에 의존합니다. XLA는 정적 텐서 모양 (static tensor shapes)을 기반으로 계산 그래프 (computation graphs)를 컴파일합니다. 그러나 에이전트 기반 궤적 (agentic trajectories)은 본질적으로 동적입니다. 어떤 에피소드는 두 번의 도구 호출 (tool calls) 후에 종료될 수 있는 반면, 다른 에피소드는 30단계 동안 실행될 수도 있습니다.

가변 길이 궤적 (variable-length trajectories)이 JAX 학습 루프에 직접 입력되면 지속적인 재컴파일 (recompilations)이 트리거되어, 가속기 (accelerators)의 성능 이점이 상쇄됩니다. Tunix는 데이터가 TPU 고대역폭 메모리 (HBM, High Bandwidth Memory)에 도달하기 전에 엄격한 시퀀스 패킹 (sequence packing) 및 버케팅 (bucketing) 전략을 구현함으로써 이 문제를 해결합니다.

동적 시퀀스 패킹 및 버케팅 (Dynamic Sequence Packing and Bucketing)

모든 궤적을 가능한 최대 에피소드 길이로 패딩 (padding) 하는 대신—이는 패딩 토큰에 메모리와 연산 자원을 낭비하게 만듭니다—Tunix는 유사한 길이의 궤적들을 버킷 (buckets) 단위로 그룹화합니다. 그런 다음 여러 개의 짧은 궤적을 하나의 고정된 크기의 시퀀스 윈도우 (sequence window, 예: 4096 또는 8192 토큰)에 패킹하며, 어텐션 마스킹 (attention masking)을 활용하여 궤적 간의 간섭 (cross-trajectory contamination)을 방지합니다. 이를 통해 입력 형태 (input shape)를 일정하게 유지하여 XLA 재컴파일을 방지하는 동시에 배치 (batch)당 토큰 밀도를 극대화합니다.

분산 메모리 레이아웃 및 샤딩 (Distributed Memory Layout and Sharding)

거대한 에이전트 모델 (예: 70B+ 파라미터)을 확장하기 위해서는 모델 가중치 (model weights)와 유입되는 궤적 배치 (trajectory batches) 모두 TPU 토폴로지 (topology) 전체에 걸쳐 샤딩 (sharded)되어야 합니다. Tunix는 JAX의 네이티브 jax.experimental.shard_map (ShardMap) 및 SPMD (Single Program, Multiple Data) 파티션 스펙 (partition specs)을 활용하여 이러한 메모리 레이아웃을 관리합니다.

대규모 에이전트 RL을 위해 Tunix를 구성할 때 다음과 같은 하이브리드 병렬화 (hybrid parallelization) 전략을 권장합니다:

• 텐서 병렬화 (Tensor Parallelism, TP): 순전파 (forward pass) 및 역전파 (backward pass) 중 지연 시간 (latency)을 최소화하기 위해 단일 노드 내 인접한 TPU들에 모델의 어텐션 헤드 (attention heads)와 MLP 레이어를 샤딩합니다.
• 파이프라인 병렬화 (Pipeline Parallelism, PP): 단일 노드의 HBM 용량을 초과하는 모델을 수용하기 위해 모델 레이어를 서로 다른 TPU 스테이지 (stages)에 분할합니다.
• 데이터 병렬화 (Data Parallelism, DP): 궤적 배치를 서로 다른 TPU 노드들에 샤딩합니다. Tunix의 데이터 로더 (data loaders)는 궤적 버퍼 (trajectory buffer)의 별도 청크 (chunks)를 해당 DP 샤드로 직접 스트리밍하도록 설계되어, 불필요한 호스트-장치 간 전송 (host-to-device transfers)을 방지합니다.

이러한 의도적인 메모리 오케스트레이션 (memory orchestration)은 비동기 액터 (asynchronous actors)로부터 들어오는 고처리량 (high-throughput) 데이터 스트림이 호스트 측 오버헤드를 최소화하면서 TPU의 로컬 HBM으로 직접 유입되도록 보장하며, 텐서 코어 (tensor cores)를 포화 상태로 유지합니다.

오프-폴리시 발산 (Off-Policy Divergence) 및 스테일 그래디언트 (Stale Gradients) 완화

액터 (actors)를 러너 (learner)로부터 분리하는 것이 처리량 병목 현상을 해결해주기는 하지만, 근본적인 알고리즘적 과제인 스테일 그래디언트 (stale gradients) 문제를 야기합니다. 액터들이 비동기적으로 실행되고 매개변수 (parameters)를 주기적으로만 동기화하기 때문에, 궤적 (trajectory)을 생성하는 데 사용된 정책 (policy)은 러너에서 현재 업데이트 중인 정책보다 다소 오래된 경우가 많기 때문입니다.

만약 정책이 너무 빠르게 업데이트되면, 버퍼 (buffer) 내의 궤적들은 매우 오프-폴리시 (off-policy) 상태가 됩니다. PPO (Proximal Policy Optimization)와 같은 표준 온-폴리시 (on-policy) RL 알고리즘은 이러한 조건에서 급격히 성능이 저하되어, 학습 불안정성, 정책 붕괴 (policy collapse) 또는 발산 (divergence)으로 이어질 수 있습니다. 수학적 정확성과 학습 안정성을 유지하기 위해, Tunix는 여러 완화 메커니즘을 통합합니다.

1. 중요도 샘플링 보정 (Importance Sampling Correction)

행동 정책 (behavior policy)과 타겟 정책 (target policy) 사이의 불일치를 보정하기 위해, Tunix는 정책 그래디언트 손실 (policy gradient loss)에 중요도 샘플링 비율 (importance sampling ratios)을 적용합니다. 표준 목적 함수 (objective)는 중요도 비율을 포함하도록 수정되며, 이 비율은 현재 정책 하에서의 행동 확률을 행동 정책 하에서의 확률로 나눈 값으로 계산됩니다. Tunix는 각 궤적을 생성하는 데 사용된 정확한 매개변수 버전을 추적하고, 전이 데이터 (transition data)와 함께 행동 로그 확률 (action log-probabilities)을 저장합니다. 러너의 순전파 (forward pass) 과정에서, 정확한 비율을 계산하기 위해 현재 매개변수 하에서의 로그 확률을 계산합니다.

2. 엄격한 스테일니스 경계 (Strict Staleness Bounds)

학습기 (learner)가 완전히 구식인 데이터로 학습하는 것을 방지하기 위해, Tunix는 동적 동기화 장벽 (dynamic synchronization barrier)을 구현합니다. 저는 이를 최대 스테일니스 임계값 (maximum staleness threshold)을 사용하여 설정합니다. 학습기의 현재 스텝과 궤적 (trajectory)이 생성된 시점의 스텝 사이의 차이가 이 임계값을 초과하면, 해당 궤적은 버퍼 (buffer)에서 폐기됩니다. 만약 버퍼의 평균 스테일니스 (average staleness)가 임계값을 초과하면, 학습기는 액터 (actors)들이 따라잡아 더 신선한 궤적들로 버퍼를 채울 때까지 실행 속도를 조절 (throttle)합니다.

3. Kullback-Leibler (KL) 발산 제약 (Constraints)

중요도 비율 (importance ratio)을 클리핑 (clipping)하는 것 외에도, Tunix는 최적화 과정에서 정책 (policy)이 너무 멀리 벗어나는 것을 방지하기 위해 현재 정책과 참조 모델 (reference model, 일반적으로 초기 SFT 모델) 사이의 KL 발산 (KL divergence) 페널티를 강제합니다. 이는 단 한 번의 잘못된 업데이트가 모델을 무한 루프에 빠뜨리거나 쓰레기 도구 호출 (garbage tool calls)을 생성하게 하여, 고품질 궤적을 수집하는 능력을 영구적으로 손상시킬 수 있는 에이전트 환경 (agentic environments)에서 특히 중요합니다.

⚙️ 구현 청사진 (Implementation Blueprint): Tunix 파이프라인 구성하기

이러한 개념들이 실제 프로덕션 코드에서 어떻게 변환되는지 설명하기 위해, Tunix의 아키텍처 패턴을 사용하는 실질적인 JAX 구현 과정을 살펴보겠습니다. 다음 코드 블록은 비동기 궤적 버퍼 (asynchronous trajectory buffer)를 설정하고, shard_map을 사용하여 샤딩된 학습기 스텝 (sharded learner step)을 정의하며, 비동기 학습 루프를 관리하는 방법을 보여줍니다.

import jax
import jax.numpy as jnp
from jax.experimental import shard_map
...

이 청사진은 Tunix의 설계 원칙을 활용하면 고도로 최적화된 분산 강화학습 (distributed RL) 코드를 작성하는 것이 얼마나 간단한지를 보여줍니다. 배치 차원 (batch dimension)을 메쉬 (mesh)의 data 축에 매핑하고 shard_map을 사용함으로써, tensor 차원을 통해 모델 가중치를 동기화하는 동시에 TPU 클러스터가 들어오는 궤적들을 최대의 병렬 효율성으로 처리하도록 보장할 수 있습니다.

🏗️ 아키텍처 비교: 동기식(Synchronous) vs. 디커플링된(Decoupled) RL

귀하의 엔지니어링 팀이 Tunix와 같은 디커플링된 (Decoupled) 아키텍처로 전환하는 것이 적절한지 평가하는 데 도움이 되도록, 동기식 (Synchronous) 에이전트 기반 RL 시스템과 디커플링된 시스템 간의 운영상의 트레이드오프 (Trade-offs) 비교를 정리했습니다.

🎯 운영상의 다음 단계

LLM (Large Language Models)이 수동적인 텍스트 생성기에서 능동적인 에이전트 (Agents)로 전환됨에 따라, 훈련 인프라는 상호작용형 환경의 고유한 병목 현상을 처리할 수 있도록 진화해야 합니다. 표준 동기식 훈련 파이프라인은 더 이상 실행 가능하지 않습니다. 호스트 측의 도구 실행, 샌드박스 내 코드 컴파일, API 응답을 기다리는 동안 귀중한 가속기 사이클을 낭비하기 때문입니다.

Google의 Tunix 라이브러리는 궤적 생성 (Trajectory generation)을 정책 최적화 (Policy optimization)로부터 분리함으로써 이 병목 현상을 해결합니다. 비동기식 액터-러너 (Actor-learner) 아키텍처를 활용하고, XLA 재컴파일을 방지하기 위한 엄격한 시퀀스 패킹 (Sequence packing)을 구현하며, 수학적으로 타당한 오프-폴리시 교정 (Off-policy corrections)을 적용함으로써, Tunix는 엔지니어링 팀이 TPU 투자를 극대화하고 에이전트 기반 RL을 수백억 개의 파라미터를 가진 모델로 확장할 수 있도록 지원합니다.