통합 개발 환경 (IDE)의 영역은 오랫동안 자동 완성 기능의 힘과 정밀성에 의해 강화되어 왔습니다. 그러나 새로운 경쟁자가 등장했습니다. 바로 인공지능 (AI)입니다. AI가 IDE에 통합되는 것은 점점 더 보편화되고 있으며, 기존 도구에 대한 재평가를 촉발하고 있습니다. 이러한 이전 방법이 여전히 작동하지만, 개발자가 기대하는 최첨단 지원을 더 이상 제공하지 못할 수 있습니다.
이 기사는 JetBrains의 블로그 게시물을 각색한 것으로, 클라우드 기반 자동 완성의 기초가 되는 모델인 Mellum의 훈련 및 기능에 대해 자세히 설명합니다. 프로그래머를 위해 맞춤 제작된 IDE 및 코드 편집기로 유명한 JetBrains는 빠르고 효율적인 코드 자동 완성을 제공하기 위해 Mellum을 개발했습니다. Mellum은 로컬 장치에서 효과적으로 작동하도록 설계되었으며, 동급에서 가장 빠르고 컴팩트한 모델 중 하나입니다. 개념적으로 유사한 프로젝트는 Microsoft Phi입니다.
분명한 과제 해결
JetBrains 엔지니어는 AI 기반 자동 완성을 추구하면서 몇 가지 중요한 과제에 직면했습니다.
- 속도 및 비용: 기존 채팅 모델은 높은 계산 비용과 느린 응답 시간으로 인해 비실용적인 것으로 판명되었습니다. 또한 이러한 모델은 중간 채우기 (FIM) 또는 토큰 복구와 같은 코드 특정 기술에 대한 인식이 부족했습니다.
- 출력 형식: 플래그십 채팅 모델은 종종 일관성 없는 형식으로 데이터를 생성하여 응답을 구문 분석하고 편집기에 원활하게 통합하기가 어려웠습니다.
- 데이터 출처: 훈련 데이터의 출처를 결정하고 잠재적인 저작권 침해 문제를 완화하는 것은 중요한 장애물이었습니다.
Mellum: 개요
JetBrains의 개발 팀은 자체 모델을 만드는 것이 최적의 접근 방식임을 깨달았습니다. 그들의 목표는 품질, 추론 비용 및 대기 시간의 균형을 맞추면서 명확한 출처를 가진 데이터로 훈련된 모델을 설계하는 것이었습니다. 예비 연구에 따르면 약 40억 개의 매개 변수를 가진 모델이 광범위한 시나리오 및 사용자를 위한 유능한 자동 완성 기능을 제공할 수 있다고 합니다. 또한, 모델을 코드만으로 훈련함으로써 관련 없는 데이터가 없는 특수 토큰 어휘를 설정할 수 있었습니다.
Mellum 모델의 훈련 프로세스는 세 가지 뚜렷한 단계로 구성되며, 각 단계는 새로운 지식을 제공하고 생성된 코드의 품질을 향상시킵니다. 초기 단계는 개별 파일의 대규모 코퍼스에 대한 기본 사전 훈련을 포함합니다. 두 번째 단계는 더 작은 특수 예제 세트로 모델을 개선하는 것으로 구성됩니다. AI 피드백을 사용한 강화 학습 (RLAIF)은 세 번째 단계에서 모델을 IDE 특정 특성에 맞게 조정하고 원치 않는 출력을 제거하는 데 사용됩니다.
사전 훈련
모호한 데이터 출처와 관련된 잠재적인 문제를 피하기 위해 모델은 처음부터 훈련되었으며, 수많은 언어, 프로그래밍 구문, 패턴 및 핵심 개념에 대한 포괄적인 소개가 필요했습니다.
데이터 세트
사전 훈련의 주요 데이터 소스는 TheStack이었습니다. 이 데이터 세트는 데이터가 법적으로 건전하고 실질적으로 유익하도록 보장합니다.
사전 훈련 프로세스
사전 훈련 동안 데이터 세트는 약 3조 개의 토큰에 도달하기 위해 여러 번 샘플링되었습니다. 8192개의 토큰의 컨텍스트 창이 사용되었으며 데이터 세트는 이 크기의 조각으로 분할되었습니다. 중간 채우기 (FIM) 변환은 각 조각의 파일 절반에 적용되어 모델이 이전 및 이후 코드를 모두 고려하도록 장려했습니다. 이 기술은 실제 코드 생성 시나리오를 밀접하게 모방합니다.
사전 훈련 단계는 각각 8개의 H100 GPU가 장착된 16개의 노드 클러스터에서 수행되었습니다. 이 단계는 완료하는 데 약 15일이 걸렸으며, 결과적으로 40억 개의 매개 변수 Mellum-base 모델이 생성되었습니다.
사전 훈련은 많은 프로그래밍 언어에 대한 광범위한 지식을 갖춘 범용 코드 자동 완성 모델을 만듭니다. 그러나 이 단계에서 모델은 임의로 선택된 파일 세그먼트에서 다음 토큰을 예측하도록 훈련됩니다. 추가 컨텍스트가 없으면 모델은 코드 구조에 대한 인식이 부족하고 코드 생성을 중지할 시기를 결정하는 메커니즘이 없습니다.
미세 조정 단계는 이러한 제한 사항을 해결하도록 설계되었습니다.
컨텍스트 인식 미세 조정
향상된 중간 채우기
예측을 위해 코드 조각이 임의로 선택되는 사전 훈련과 달리 미세 조정은 코드를 보다 의미 있는 방식으로 분할하는 데 집중하여 모델이 ‘야생’에서 발생하는 코드 조각을 추출하도록 가르칩니다.
전문화된 예제
실제로 코드 자동 완성을 위해서는 주변 파일과 더 넓은 컨텍스트 (전체 프로젝트 포함 가능)를 이해해야 합니다.
데이터 전처리를 위해 회사는 전체 프로젝트 인덱싱 없이 일반 파일에서 직접 컨텍스트를 구축하기 위해 개발된 크로스 플랫폼 SDK 및 콘솔 유틸리티 세트인 Code Engine이라는 내부 프로젝트를 시작했습니다. 이 SDK는 내부 MapReduce 클러스터에 배포되어 합리적인 시간 내에 훈련에 유용한 많은 예제를 생성하기 위해 수천 개의 공용 리포지토리를 처리하는 데 사용되었습니다.
올바른 알고리즘을 찾는 데는 시행착오가 필요했습니다.
특정 언어에 대한 튜닝
작은 모델은 특정 언어에 대한 전문화를 통해 큰 이점을 얻을 수 있습니다. 기본 모델은 80개 이상의 언어로 훈련되지만 대부분의 사용자는 일반적으로 1~2개의 언어만 사용합니다. 이를 해결하기 위해 JetBrains는 여러 개의 전문화된 모델을 만들었습니다.
- mellum-all: JetBrains IDE에서 사용할 수 있는 대부분의 언어 및 방언을 지원하지만 자동 완성 품질은 전문화된 모델보다 낮습니다.
- mellum-python: Python 및 Jupyter를 전문으로 합니다.
- mellum-kotlin: Java 및 Kotlin을 전문으로 합니다.
- mellum-web: 웹 기술을 전문으로 합니다.
마지막 단계: RLAIF
마지막으로 학습 목표가 사용자 기대와 일치하지 않는 경우를 해결해야 합니다. 추가 훈련 단계인 RLAIF (AI 피드백을 사용한 강화 학습)는 이러한 문제를 해결하는 데 사용됩니다.
모델은 사용자 상호 작용에서 학습하고 사용자 선호도를 더 잘 반영하는 방법을 이해합니다.
이 접근 방식은 전체 품질 점수를 향상시킬 뿐만 아니라 성가신 생성 아티팩트의 수를 줄입니다.
Mellum은 얼마나 좋은가?
이 모델은 크기에 비해 매우 뛰어난 성능을 보입니다. 평가 방법은 다음과 같습니다.
- 먼저 모델은 ‘JetBrains BigCode’라는 내부 벤치마크에서 평가되었습니다.
- 그런 다음 SAFIM과 같은 잘 알려진 공용 벤치마크에서 테스트되었습니다.
- 마지막으로 기능에 대한 사용 통계가 수집되고 사용자 메트릭이 계산되었습니다.
오프라인 평가
데이터를 수집하는 것은 복잡한 작업이지만 신경망에서 제안한 새로운 제안과 원래 제안을 비교하는 좋은 메트릭을 만드는 것은 훨씬 더 어렵습니다. 우리는 작은 연구를 수행했고 궁극적으로 두 가지 주요 메트릭의 조합에 정착했습니다.
EM:
- Exact Match는 매우 인기있는 아이디어입니다.
- 완성의 첫 번째 줄이 최소한의 전처리를 통해 원래의 첫 번째 줄과 일치하면 예측이 좋은 것으로 간주됩니다.
KK:
- 이 메트릭은 작성자의 이름을 따서 명명되었습니다.
- 원래에서 제안된 줄 수를 제안된 완료의 줄 수로 나눈 값입니다.
JetBrains BigCode
이 모델은 내부 JetBrains BigCode 도구를 사용하여 얻은 벤치마크 데이터 세트에 대해 평가되었습니다.
공용 벤치마크에 의존하는 대신 데이터 세트에 대한 완전한 제어력을 유지함으로써 다양한 코딩 스타일 및 관행에 대한 모델 품질을 안정적으로 평가할 수 있습니다.
JetBrains BigCode 평가 결과는 인기 모델과 동등한 품질을 보여주지만 Mellum은 더 작고 효율적입니다.
단일 줄 제안의 품질 (EM 메트릭)
공용 벤치마크
이 모델은 내부 데이터 세트뿐만 아니라 다국어 벤치마크 SAFIM (구문 인식 중간 채우기)과 같은 다양한 공용 벤치마크에서도 평가되었습니다.
온라인 평가
주요 메트릭은 완료된 코드 비율 (RoCC)이라고합니다. 이는 코드 자동 완성을 사용하여 작성된 코드 문자 수를 편집기의 총 코드 양으로 나눈 비율로 정의됩니다.
또 다른 중요한 메트릭은 수락률 (AR)이며, 이는 표시된 모든 제안 수로 나눈 수락된 제안 수로 계산됩니다.
이것은 복잡한 여정이었지만 JetBrains 전문가들은 존엄성을 가지고 완료했습니다. 결국 하나의 일반 및 여러 개의 전문화된 모델을 얻었으며, 이는 JetBrains AI 플랫폼을 통해 사용할 수 있습니다. 그들은 이제 JetBrains AI Assistant에서 성공적으로 일하고 있습니다.
다음은 무엇일까요?
- JetBrains 엔지니어는 현재 웹 개발 언어용 모델을 개발 중입니다. 조만간 공개적으로 제공될 수 있습니다.
- 매개 변수 수와 데이터 다양성을 동시에 늘릴 계획이 있습니다. 코딩에는 다양한 작업이 있습니다. Mellum도 이를 수행 할 수 있습니다. 서비스 성능은 여전히 핵심 메트릭이므로 모델 확장은 합리적인 범위 내에서 이루어집니다.