vime ROCm 지원 AMD Instinct GPU용

vLLM은 vime에 ROCm 지원을 통합하여, 대규모 RL pós-트레이닝을 AMD Instinct MI300X 및 MI355X GPU에서 네이티브로 실행할 수 있게 합니다. 이 통합은 사전 빌드된 컨테이너를 포함하여 소스에서 빌드할 필요를 제거하는 검증된 엔드-투-엔드 파이프라인을 제공합니다.

vime 아키텍처 및 설계

vime는 slime 프레임워크를 기반으로 한 세 단계, 분리된 훈련-추론 설계를 활용합니다. 이는 RL 프로세스를 세 가지 구별된 구성 요소로 분리합니다:

• Training (Megatron): 메인 훈련 루프, 매개변수 업데이트, 그리고 롤아웃 측으로의 가중치 동기화를 관리합니다.
• Rollout (vLLM + Router): 보상 또는 검증기 신호를 기반으로 훈련 샘플을 생성하기 위한 추론 샘플링을 처리합니다.
• Data Buffer: 훈련과 롤아웃 사이의 다리 역할을 하며, 사용자 정의 롤아웃 로직과 프롬프트 주입을 관리합니다.

AMD Instinct GPU의 RL에서의 하드웨어 장점

RL pós-트레이닝은 훈련 측 가중치(Megatron 형식)와 추론 측 KV 캐시(vLLM 롤아웃)를 동시에 저장해야 하므로 메모리 집약적입니다. AMD Instinct GPU는 이 프로파일에 최적화되어 있습니다:

• 높은 HBM 용량: MI300X는 192 GB의 HBM3를 제공하며, MI355X는 288 GB를 제공합니다. 이 큰 용량은 agresive 텐서 병렬주의 필요성을 줄여 토폴로지를 단순화하고 클러스터 활용도를 높입니다.
• 우수한 메모리 대역폭: MI300X는 5 TB/s 이상, MI355X는 8 TB/s의 총 대역폭을 제공합니다. RL 롤아웃은 메모리 대역폭에 제한되므로, 이 높은 처리량은 자율 회귀 토큰 생성 단계에서의 스텝 지연을 줄입니다.
• 네이티브 생태계 통합: ROCm의 오픈소스 특성 덕분에 vime은 별도의 코드 경로 없이 기존 vLLM 롤아웃 스택을 상속할 수 있으며, ROCm에 대한 네이티브 PyTorch 및 vLLM 지원을 활용합니다.

ROCm에서의 기술적 구현

AMD 하드웨어로 vime를 가져오기 위해 몇 가지 핵심 소프트웨어 구성 요소를 통합했습니다:

Megatron-LM 백엔드

vime는 ROCm 호환 Megatron-LM 포크를 사용합니다. 호환성을 보장하기 위해 non-CUDA 빌드에서 CUDA fused-kernel 초기화를 보호하는 패치가 구현되었습니다. 이 파이프라인은 단일 GPU에서 HuggingFace에서 torch_dist로의 체크포인트 변환을 지원하며, 그래디언트 누적을 위해 네이티브 PyTorch 경로를 활용합니다.

colocated 가중치 동기화

colocated 모드에서는 Megatron과 vLLM이 동일한 GPU 풀을 공유합니다. vime는 각 옵티마이저 단계 후에 Megatron에서 vLLM으로 업데이트된 가중치를 동기화하기 위해 프로세스 간 통신(IPC)을 사용합니다. 이는 ROCm에서 torch.cuda.get_device_properties(i).uuid 인터페이스를 통해 활성화되며, 이는 라우팅을 위한 안정적이고 프로세스 일관된 디바이스 UUID를 제공합니다.

리소스 관리 및 Ray 통합

GPU 할당은 HIP_VISIBLE_DEVICES를 통해 관리됩니다. vime는 CUDA_VISIBLE_DEVICES와 함께 이를 설정하여 Megatron 훈련 액터와 vLLM 서브프로세스 간의 일관성을 보장합니다. Ray의 AMD GPU 관리자는 이러한 마스크를 존중하도록 구성되며, 시스템은 규모에 맞게 액터 워커를 생성하기 위해 파일 디스크립터 제한을 높여야 합니다(--ulimit nofile=1048576:1048576).

MI355X에서의 성능 벤치마크

MI355X GPU에서 Qwen3-8B 모델을 테스트한 결과, 100 훈련 단계 동안 다음과 같은 결과를 보여주었습니다:

• 처리량: 약 4,100 tokens_per_gpu_per_second를 지속적으로 달성했습니다. 정책이 더 예측 가능한 출력을 학습함에 따라 처리량이 상승하여 vLLM이 더 효율적으로 배치할 수 있게 되었습니다.
• 로그프롭 안정성: train_rollout_logprob_abs_diff가 약 0.012 주변에서 안정적으로 유지되었습니다. 훈련과 롤아웃 로그 확률 사이의 낮은 차이는 성공적인 가중치 동기화를 나타내며, NVIDIA 하드웨어에서의 결과와 유사합니다.
• 보상 수렴: dapo-math-17k 데이터셋을 사용하여 raw_reward가 0 근처에서 약 0.5–0.6으로 상승하여, 정책이 검증기에 의해 보상받는 추론 패턴을 선호하도록 성공적으로 학습했음을 나타냅니다.

지원 기능 및 로드맵

현재 지원

• GRPO 훈련
• colocated 및 비동기(서로 다른 GPU 풀) 훈련 및 롤아웃
• Megatron-LM 훈련 백엔드 및 vLLM 롤아웃 백엔드
• Qwen3 Dense 및 MoE 모델 지원

미래 로드맵

• 전체 vLLM 라우터 및 PD (Prefill-Decode) 분리 지원
• FP8 파이프라인 최적화
• AMD MoE 워크로드를 위한 R3 (Rollout Routing Replay)
• 비동기 훈련에 대한 성능 최적화로 메모리 누수 및 로그프롭 차이 해결
• 멀티 에이전트 설정 및 멀티 턴 툴 호출을 위한 Agentic RL