CodeParrot 처음부터 학습시키기

Hugging Face는 Python 코드 자동 완성 기능을 제공하기 위해 처음부터 학습된 대규모 GPT-2 모델인 CodeParrot를 개발했습니다. 이 프로젝트는 간소화된 학습 파이프라인과 선별된 소스 코드 데이터셋을 사용하여 GitHub Copilot의 기반이 되는 기술과 유사한 코드 생성 모델을 구축하는 방법을 보여줍니다.

데이터셋 큐레이션 및 정제

CodeParrot는 Google의 BigQuery에서 사용할 수 있는 GitHub 덤프에서 파생된 데이터셋을 사용하여 학습되었으며, 특히 Python 파일만 필터링되었습니다. 초기 원시 데이터셋은 총 180 GB에 달하는 2,000만 개의 파일로 구성되었습니다.

데이터 정제는 모델 성능에 매우 중요했습니다. Hugging Face는 중복 데이터가 결과에 심각한 영향을 미친다는 것을 발견했습니다. 분석 결과 다음과 같이 중복 콘텐츠가 고도로 집중되어 있음이 밝혀졌습니다:

0.1%의 고유 파일이 전체 파일의 15%를 차지함.
1%의 고유 파일이 전체 파일의 35%를 차지함.
10%의 고유 파일이 전체 파일의 66%를 차지함.

이를 해결하기 위해 중복을 제거하고 Codex 논문에서 사용된 정제 휴리스틱을 적용하여, 50 GB의 정제된 데이터셋(codeparrot-clean)을 생성했습니다.

모델 아키텍처 및 토큰화

사용자 정의 토크나이저

코드 토큰이 효율적으로 분할되도록 하기 위해, Python 데이터셋에 특화되어 학습된 새로운 토크나이저를 학습시켰습니다. 이는 GPT-2 토크나이저를 가져와 train_new_from_iterator() 메서드를 사용하여 소스 코드 분포에 맞게 조정함으로써 달성되었습니다.

모델 구성

CodeParrot는 15억 개의 파라미터를 가진 GPT-2 large의 하이퍼파라미터를 활용합니다. 수치적 안정성과 새로운 토크나이저와의 호환성을 유지하기 위해 다음과 같은 조정이 이루어졌습니다:

Embedding Layer: 사용자 정의 토크나이저에 맞게 조정됨.
Attention Scaling: 레이어 ID에 따라 어텐션을 스케일링하기 위해 scale_attn_by_layer_idx 플래그를 활성화함.
Precision: 수치 문제를 피하기 위해 reorder_and_upcast_attn 플래그를 사용하여 어텐션을 전체 정밀도로 계산함.

학습 구현

학습은 🤗 Accelerate 라이브러리를 사용하여 구현되었으며, 이를 통해 파이프라인을 코드 변경 없이 단일 노트북에서 멀티 GPU 환경으로 확장할 수 있습니다.

기술적 학습 세부 사항

Hardware: 모델은 16 x A100 GPU 머신에서 학습되었습니다.
Training Duration: 110M 파라미터 모델은 학습에 하루가 소요되었고, 1.5B 파라미터 모델은 일주일이 소요되었습니다.
Memory Optimization: GPU 메모리 사용량을 줄이기 위해 Gradient checkpointing을 활성화했습니다.
Data Handling: 50GB 데이터셋을 전체로 다운로드하는 대신 IterableDataset을 사용하여 스트리밍했습니다. 토큰 사용을 극대화하기 위해, 여러 예시를 EOS 토큰과 함께 연결한 후 고정된 컨텍스트 크기로 분할했습니다.
Distribution: 각 GPU 워커가 모델의 복사본을 유지하고 가중치를 업데이트하기 위해 그래디언트를 집계하는 DistributedDataParallel (DDP) 설정을 사용했습니다.

평가 및 성능

CodeParrot는 거의 200개의 코딩 과제를 테스트하는 OpenAI의 HumanEval 벤치마크를 사용하여 평가되었습니다. 성능은 pass@k 메트릭으로 측정되며, 이는 k개의 후보 생성물 중 적어도 하나가 주어진 문제의 유닛 테스트를 통과할 확률을 나타냅니다.

GPT-neo(300 billion) 또는 Codex(400 billion total)와 비교했을 때 훨씬 적은 토큰(약 25-30 billion)으로 학습되었음에도 불구하고, CodeParrot는 경쟁력 있는 다운스트림 성능을을 보여주며 학습 효율성 측면에서 높은