XYZ-AI-Lab/axrl

해결하는 문제

AxisRL은 여러 번의 상호작용을 거쳐 환경과 소통하고, 도구를 호출하며 장기적인 결과에 기반한 보상을 받는 에이전트형 강화학습(RL) 후처리 학습을 위해 설계되었습니다. 복수 타ーン 로울아웃, 환경 상태, 도구 호출, 가중치 동기화를 관리하는 조율 문제를 해결하며, 학습 과정이 안정적이고 관측 가능하도록 보장합니다.

작동 방식

AxisRL은 고스루풋 로울아웃 엔진(SGLang)과 대규모 분산 학습 엔진(Megatron)을 연결하는 시스템 계층으로 작동합니다. 로울아웃 에이전트가 에이전트 워크플로우를 실행하고, SGLang 워커가 텍스트를 생성하며, 환경이 보상을 생성하고, Megatron 워커가 정책 최적화(PPO 또는 GRPO 등)를 수행하는 루프를 관리합니다.

주요 기술적 메커니즘은 다음과 같습니다:

• 화이트박스 및 블랙박스 통합: 내부 에이전트 루프(화이트박스) 또는 OpenAI 호환 프록시를 통해 외부 하드웨어의 상호작용을 캡처하는 블랙박스를 지원합니다.
• 부분 로울아웃: 완료된 샘플을 트레이너에 즉시 전송하여 긴 꼬리 트래잭터리 지연으로 인한 대기 시간을 줄입니다.
• 핸들 기반 데이터 플레인: MoE 라우팅 데이터와 같은 무거운 데이터를 중앙 드라이버를 거치지 않고 핸들로 전달하여 버퍼링을 방지합니다.
• 컨텍스트 관리: 반복적인 어텐션 계산을 최적화하기 위해 프리픽스 트리 병합과 MagiAttention을 사용합니다.
• 일관성 도구: 로울아웃과 트레이너 실행 경로 간의 편차를 디버깅하기 위해 불일치 분석 및 "스파이크 리플레이"를 포함합니다.

대상 사용자

수백억 파라미터 이상의 대규모 에이전트형 LLM을 훈련하는 연구자 및 엔지니어를 위한 것입니다. 복잡한 복수 타ーン 상호작용과 고스루풋 RL 후처리 학습이 필요한 경우에 적합합니다.

주요 특징

• 로울아웃에 SGLang, 분산 학습에 Megatron 기반 구현.
• PPO, GRPO/GRPO2, GSPO, TOPR, TIS 등 다양한 최적화 목표 지원.
• 300턴 이상의 트래잭터리 처리 가능.
• 학습 중 MoE 전문가 라우팅을 안정화하는 Rollout Routing Replay(R3) 제공.
• 로스 스파이크 및 토큰 수준 불일치에 대한 재현 가능한 디버깅 도구 세트 제공.