PithTrain: 1만 줄 코드로 만든 에이전트 친화 MoE 학습 프레임워크, 메가트론과 속도 동률

핵심 메시지

• 발표자들은 학습 프레임워크를 평가하는 축이 지금까지 처리량 같은 '학습 효율' 하나뿐이었다고 지적하며, 코딩 에이전트가 같은 작업을 얼마나 적은 턴과 토큰으로 끝내는지를 재는 '에이전트 작업 효율'이라는 두 번째 축을 제안했다.
• PithTrain은 컴팩트한 코드베이스, 파이썬 네이티브, 암묵적 간접 참조 최소화, 에이전트 스킬 동봉이라는 네 가지 설계 원칙 위에 만들어졌고 전체가 약 1만 1천 줄로, 10만 줄이 넘는 기존 프레임워크보다 한 자릿수 작다.
• H100 4노드에서 파이프라인 병렬 4·전문가 병렬 8로 수십억 토큰을 학습한 결과 손실 곡선이 메가트론과 겹쳤고, H100·B200에서 잰 처리량도 사실상 동률(한 사례에서 12만 4천 대 12만 6천)이었다.
• 직접 만든 AT-Bench 측정에서 질의응답 과제는 메가트론 대비 에이전트 턴이 최대 67% 줄었고, 환경 구축 과제는 88턴·40.5분에서 26턴·7분 미만으로 줄었다.
• 절차를 문서가 아닌 실행 가능한 스킬로 묶어두자 GPU 사용 시간은 거의 그대로인데 에이전트 턴이 114회에서 34회로, 출력 토큰이 3만에서 1만 1천으로 급감했다.

쉽게 이해하기

GPU MODE 강연에 나온 루이항 라이와 하오 캉은 최근 프런티어 모델이 대부분 전문가 혼합(MoE) 구조로 수렴했다는 관찰에서 출발한다. 토큰마다 파라미터의 일부만 활성화되기 때문에 같은 연산량으로 더 큰 용량을 쓸 수 있고, 추론 비용이 싸다는 점이 추론 중심 시대에 특히 중요해졌다는 것이다. 그래서 팀은 MoE 학습에만 집중하기로 결정했다고 밝혔다.

문제는 MoE를 학습시키려면 사실상 메가트론이나 딥스피드 같은 대형 프레임워크에 의존해야 한다는 점이다. 이들은 검증됐고 빠르지만 코드가 10만 줄을 훌쩍 넘고 C++ 확장이 얽혀 있어, 새 논문의 기법 하나를 넣거나 새 클러스터에 환경을 세우는 일마다 큰 비용이 든다. 발표자들은 2026년 현재 이 작업의 상당 부분을 사람이 아니라 코딩 에이전트가 하고 있는데도, 프레임워크는 여전히 사람이 읽는 것을 전제로 설계돼 있다고 지적한다.

그래서 이들은 프레임워크를 평가하는 두 번째 축을 제안한다. 같은 과제를 서로 다른 프레임워크 위에서 같은 에이전트에게 시켜보고, 걸린 시간·에이전트 턴 수·출력 토큰 수를 재는 방식이다. 하나의 대표 지표로 압축되지는 않지만 모두 구체적이고 측정 가능한 값이며, 이를 위해 AT-Bench라는 벤치마크를 별도로 만들었다. 일반적인 코딩 벤치마크가 프레임워크를 고정하고 에이전트를 바꾸는 것과 정확히 반대 구조다.

PithTrain의 네 원칙은 단순하다. 코드를 작게 유지하고, 파이썬으로만 쓰고(커널은 Triton 같은 파이썬 DSL로 작성해 속도를 지킨다), 다른 파일을 열어봐야 알 수 있는 암묵적 간접 참조를 줄이고, 사람이 알고 있지만 코드에는 적히지 않는 절차를 에이전트 스킬로 만들어 저장소에 함께 넣는다. 특히 스킬은 범위가 좁고 명확할 것, 전제 조건을 명시할 것, 성공·실패를 스크립트로 판정할 수 있을 것이라는 세 조건을 강조했다.

성능을 포기하지 않았다는 근거로는 딥시크의 듀얼파이프 설계를 기반으로 한 5단계 계층 분해가 제시됐다. MoE 계층을 통신 경계에서 쪼개 계산과 통신을 서로 다른 마이크로배치끼리 겹치게 하고, 여기에 그래프 전체 컴파일과 가중치 기울기 지연, 토큰 중복 제거, 융합 커널을 얹었다. 발표자들은 개별 기법 중 새로운 것은 없지만, 이를 모아 순수 파이썬 프레임워크를 처리량 동률까지 끌어올린 것이 요점이라고 정리했다.

주요 인사이트

• 코드베이스 크기가 성능이 아니라 '작업 비용'의 변수로 다뤄진다. 코드가 작으면 에이전트가 방향을 잡기 위해 읽어야 할 파일과 추적할 의존성이 줄어들고, 그 차이가 턴 수와 토큰 사용량으로 그대로 나타난다.
• 실패가 어디서 드러나는지도 언어 선택에 좌우된다. 파이썬 네이티브면 에이전트가 정확한 줄을 가리키는 트레이스백을 받지만, 컴파일된 확장이 끼면 원인이 불투명해 에이전트가 엉뚱한 수정안을 내놓고 재실행을 반복하게 된다.
• 실제 실험에서 메가트론의 재실행은 상당수가 암묵적 상호작용에서 나왔다. 에이전트가 붙인 커맨드라인 플래그가 코드 깊숙한 곳에서 자동 생성되는 설정과 충돌했고, 이 연결은 처음 코드를 읽을 때 보이지 않았다.
• 스킬 유무 실험은 '토큰을 줄이는 것'과 'GPU를 덜 쓰는 것'이 별개임을 보여준다. 정해진 학습 실행 횟수가 있는 과제라 GPU 시간은 20.8분 대 22.5분으로 비슷했지만, 절차를 다시 추론할 필요가 없어진 에이전트의 턴과 토큰은 3분의 1 이하로 떨어졌다.
• 발표자들은 스스로 한계를 짚었다. 산업계는 토큰을 아끼는 데 큰 관심이 없을 수 있고, 모델·기능 커버리지는 대형 프레임워크에 못 미치며, 간접 참조를 얼마나 줄였는지 정량화하는 방법도 아직 없다고 인정했다.

자주 묻는 질문

관련 AI 소식