Google이 TPU v5e 2x4와 Keras 모델 샤딩으로 7개 LLM을 동시 로드하여 LLM의 실시간 오류 수정 능력을 평가하는 대화형 테스트 환경 구축

Context

LLM 기반 어시스턴트는 첫 시도에서 완벽한 코드를 생성하지 못할 수 있으며, 사용자의 영어 피드백만으로 오류를 수정할 수 있는 능력이 실제 사용성을 결정한다. 단순한 API 호출 작업(일정 추가/삭제)조차 일부 모델은 신뢰성 있게 처리하지 못하고, 오류 지적 후 수정 요청에 대해 일관되지 않은 응답을 생성한다.

Technical Solution

• TPU v5e 2x4(8코어, 코어당 16GB RAM, 총 128GB)를 선택하여 여러 LLM을 메모리에 동시 로드: bfloat16 포맷으로 약 8B 파라미터 모델 5개 + 약 2B 파라미터 모델 3개(총 7개) 동시 실행
• JAX와 Keras를 기술 스택으로 선택: JAX의 XLA 컴파일러로 TPU 최적화, Keras의 KerasHub 라이브러리로 사전학습 LLM 로드
• Keras API를 사용한 모델 샤딩으로 8개 TPU 코어에 모델 분산 배치: 기본 레이아웃 맵으로 자동 샤딩 처리
• Gradio를 프론트엔드, HuggingFace Spaces를 호스팅 플랫폼으로 사용하여 웹 기반 대화형 인터페이스 구현
• 2개의 병렬 채팅 세션으로 모델 간 성능 비교: 한쪽 대화를 일시 중지하고 다른 모델의 오류 수정 요청 가능하도록 UI 설계

Impact

Gemma 2 9B가 테스트를 거의 완벽하게 통과했고, Llama 3.1 8B는 잘못된 API 호출 수정에서 신뢰성이 낮았으며, Codegemma 7B는 정정된 코드 옆에 사과문(apology)을 계속 생성하는 반복적 오류를 보였다. "API call only please" 프롬프트로 Gemma의 사과 문제는 해결되었으나 Llama에서는 일관되지 않았다.

Key Takeaway

LLM의 오류 수정 능력은 모델마다 매우 다르며, 단순 API 호출 작업도 일부 모델은 재현 가능성 있게 처리하지 못한다는 것은 프롬프트 엔지니어링만으로는 부족하며 모델 선택과 파인튜닝이 실제 대화형 AI 어시스턴트 구현의 핵심임을 시사한다.