[Axolotl] 플러그인에 scored rollout 디스패치, 외부 플러그인 경로 확장, vLLM 에러 처리 개선

PR 링크: axolotl-ai-cloud/axolotl#3549 상태: Merged | 변경: +116 / -13

들어가며

이 PR은 세 가지 독립적 개선을 포함합니다:

1. GRPO 트레이너에서 rollout scoring 완료 후 플러그인에 데이터를 전달하는 on_rollouts_scored 훅
2. pkgutil.extend_path로 외부 패키지의 플러그인을 자동 발견
3. vLLM worker의 에러 처리 강화 및 /reset_prefix_cache endpoint의 fire-and-forget 변환

핵심 코드 분석

1. on_rollouts_scored 플러그인 훅

# integrations/base.py - BasePlugin
def on_rollouts_scored(self, cfg, trainer, prompts, completions, rewards, advantages):
    """Called after rollouts are scored during online RL (GRPO/PPO).
    
    Args:
        prompts: List of prompt texts (one per sample).
        completions: List of completion texts (one per sample).
        rewards: Dict mapping reward function name to list of reward values.
        advantages: List of advantage values (one per sample).
    """

플러그인이 scored rollout 데이터(프롬프트, 완성, 보상, 어드밴티지)에 접근할 수 있게 됩니다. 로깅, 분석, 트레이스 저장 등에 활용할 수 있습니다.

2. 외부 플러그인 경로 확장

Before: src/axolotl/integrations/__init__.py가 비어 있음

After:

import pkgutil
__path__ = pkgutil.extend_path(__path__, __name__)

pkgutil.extend_path는 같은 패키지 이름을 가진 외부 패키지의 경로를 __path__에 추가합니다. 이를 통해 pip install로 설치한 외부 플러그인이 axolotl.integrations 네임스페이스에 자동으로 포함됩니다.

3. vLLM reset_prefix_cache fire-and-forget 변환

Before:

@app.post("/reset_prefix_cache/")
async def reset_prefix_cache():
    for conn in connections:
        conn.send({"type": "call", "method": "reset_prefix_cache"})
    results = await asyncio.gather(
        *(loop.run_in_executor(None, conn.recv) for conn in connections)
    )
    return {"message": f"Reset prefix cache: {all(results)}"}

After:

@app.post("/reset_prefix_cache/")
async def reset_prefix_cache():
    for conn in connections:
        conn.send({"type": "fire_and_forget", "method": "reset_prefix_cache"})
    return {"message": "Reset prefix cache received"}

conn.recv()를 기다리지 않는 fire-and-forget 방식으로 변경합니다. worker가 응답을 보내지 않으므로 파이프에 응답이 남아 다음 recv() 호출을 오염시키는 문제가 해결됩니다.

왜 이게 좋은가

• 확장성: 플러그인 훅으로 코어 코드 수정 없이 custom 로깅/분석을 추가할 수 있습니다.
• namespace package: 외부 개발자가 axolotl-integrations-myplugin 같은 별도 패키지로 플러그인을 배포할 수 있습니다.
• 안정성: vLLM worker 에러가 전체 트레이닝을 중단시키지 않고 경고로 처리됩니다.

정리

플러그인 생태계 확장, GRPO 훅 추가, vLLM 통신 안정성 개선이라는 세 가지 축으로 Axolotl의 확장성과 안정성을 동시에 향상시킨 PR입니다.

참고 자료

• pkgutil.extend_path 문서
• Axolotl Plugin 가이드

---

이 포스트는 AI가 작성하였으며, 사실과 다를 수 있습니다. 정확한 정보는 원본 PR을 참고해 주세요.