ROCKET: AMD MI300X를 위한 자율 성능 최적화 에이전트

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

기존의 GPU 성능 최적화 과정은 프로파일링, 가설 수립, 코드 수정, 검증의 반복으로 많은 시간이 소요되는 수동 작업이었습니다. ROCKET은 이러한 루프를 자동화하기 위해 구축된 자율 에이전트로, AMD MI300X 하드웨어와 ROCm 7.0 환경에 최적화되어 있습니다. 에이전트는 Planner, Implementer, Validator 등 4개의 모듈로 구성되어 모델을 분석하고 5가지 핵심 최적화 도구 중 가장 효과적인 것을 선택해 적용합니다. 특히 자유로운 코드 생성 대신 정의된 도구 상자 내에서 선택하는 방식을 채택하여 코드의 유효성과 수치적 정확성을 보장합니다. Qwen2.5-7B 모델 테스트 결과, 초당 토큰 생성 속도가 62.6에서 183.5로 약 2.93배 향상되는 성과를 거두었습니다. 이는 전문 엔지니어의 판단 과정을 LLM이 성공적으로 모방하여 실무적인 성능 개선을 이끌어낼 수 있음을 보여줍니다.

의미 / 영향

AI 에이전트가 하드웨어 특화 최적화라는 전문 영역에서 인간 엔지니어의 반복적인 튜닝 작업을 대체할 수 있음을 입증했습니다. 특히 AMD 하드웨어 생태계에서 개발자 경험(DX)을 개선하고 모델 배포 효율성을 극대화하는 실질적인 도구로 활용될 가능성이 높습니다.

빠른 이해

요약 브리프

ROCKET은 AMD MI300X GPU에서 PyTorch 모델의 성능을 자동으로 튜닝하는 자율 에이전트입니다. LLM이 프로파일링 데이터를 분석하고 최적의 기법을 선택·검증하여 Qwen2.5 모델 기준 약 3배의 속도 향상을 이끌어냈습니다.

새로운 점

자유로운 코드 생성 대신 검증된 최적화 도구 상자를 활용해 24시간 내에 실무 적용 가능한 수준의 성능 향상과 정확도 검증을 동시에 달성했습니다.

핵심 메커니즘

모델 입력 → torch.profiler 분석 → Qwen2.5 Planner가 최적화 도구(bf16, compile 등) 선택 → 코드 적용 → 정확도 및 속도 재측정 → 최종 최적화 모델 출력

핵심 수치

End-to-end Speedup: 2.93x- Qwen2.5-7B-Instruct 모델 기준
Throughput Improvement: 62.6 → 183.5 tok/s- 단일 MI300X, Batch 8 환경
Optimization Tools: 5 types- dtype_cast, torch_compile, SDPA 등

섹션별 상세

ROCKET의 아키텍처와 작동 원리

ROCKET은 네 개의 전문 에이전트가 긴밀하게 협력하는 루프 구조로 설계되었습니다. Profiler가 모델의 병목 지점을 요약하면, Planner(Qwen2.5-7B 기반)가 최적화 도구를 선택하고, Implementer가 이를 실제 코드로 구현합니다. 마지막으로 Validator가 성능 변화와 수치적 정확도를 검증하여 변경 사항의 유지 여부를 결정합니다. 이러한 폐쇄형 루프 시스템은 인간의 개입 없이도 최적의 성능 지점을 찾아낼 때까지 반복 수행됩니다.

제한된 도구 상자(Bounded Toolbox) 전략

에이전트에게 자유로운 코드 생성 권한을 주는 대신, 검증된 5가지 최적화 기법(dtype_cast, torch_compile, sdpa_attention, input_padding, kv_cache_config) 내에서 선택하도록 제한했습니다. 이는 생성된 코드가 컴파일에 실패하거나 수치적 오차를 발생시키는 문제를 방지하기 위한 설계입니다. 성능 엔지니어링의 핵심은 새로운 커널을 만드는 것보다 기존의 효과적인 기법을 올바른 순서로 적용하는 데 있다는 통찰을 바탕으로 합니다. 이를 통해 24시간이라는 짧은 개발 시간 내에 신뢰할 수 있는 성능 향상 결과를 도출할 수 있었습니다.

Qwen2.5-7B 모델 최적화 결과 및 검증

단일 AMD MI300X 노드에서 Qwen2.5-7B-Instruct 모델을 대상으로 테스트한 결과, 2.93배의 속도 향상을 기록했습니다. 에이전트는 5개의 도구 중 dtype_cast(bf16 변환)가 가장 효과적임을 파악하여 적용했고, 나머지 4개 기법은 성능 향상이 미미하거나 기준 미달로 판단하여 거부했습니다. 모든 최적화 단계는 실제 벽시계 시간(Wall-clock time) 측정과 로짓(Logit) 비교를 통한 정확도 검증을 거쳤습니다. 이는 에이전트가 단순히 이론적인 최적화가 아닌, 실제 하드웨어에서 검증된 '정직한' 성능 향상을 제공함을 의미합니다.

실무 Takeaway

AMD MI300X 환경에서 ROCKET 에이전트를 사용해 Qwen2.5-7B 모델의 추론 속도를 62.6 tok/s에서 183.5 tok/s로 2.93배 향상시킬 수 있다.
LLM 에이전트에게 자유로운 코드 작성을 맡기기보다 검증된 최적화 도구 세트(Toolbox)를 제공하고 선택하게 하는 것이 실무적인 성능 튜닝에 더 효과적이다.
모든 최적화 단계에서 수치적 허용 오차(Numerical tolerance) 내의 정확도 검증을 자동화하여 성능 향상 시 발생할 수 있는 모델 오동작을 방지한다.

언급된 리소스

GitHubROCKET GitHub Repository

DemoROCKET Live Demo (HF Space)

문서원문 링크