StarCoder로 코딩 어시스턴트 만들기 – 요약

개요

  • Goal: StarCoder(16 B 파라미터 코드 생성 모델)를 채팅형 코딩 어시스턴트인 StarChat‑α 로 변환합니다.
  • 핵심 단계:
    1. 특수 채팅 토큰 추가 (<|system|>, <|assistant|>, <|user|>, "\n").
    2. 대화 데이터셋 준비: Open‑Assistant oasst1 데이터셋을 영어 대화만 남기도록 필터링하고 새로운 토큰을 추가합니다.
    3. 사용자 라벨 마스킹: 손실이 어시스턴트 출력에만 계산되도록 합니다.
    4. Fine‑tune: 모델이 단일 GPU에 맞지 않으므로 DeepSpeed ZeRO‑3(또는 LoRA/FSDP와 같은 대안)으로 파인튜닝합니다.

토크나이저 및 마스킹

special_tokens = {
    "additional_special_tokens": ["<|system|>", "<|assistant|>", "<|user|>", "\n"]
}
tokenizer.add_special_tokens(special_tokens)
  • <|assistant|>가 단일 토큰 ID(49153)와 매핑되는지 확인합니다.
  • mask_user_labels는 사용자 턴 토큰 ID와 다음 어시스턴트 토큰까지의 모든 토큰을 -100으로 교체하여 손실 계산에서 무시합니다.

DeepSpeed ZeRO‑3를 이용한 학습

git clone https://github.com/bigcode-project/starcoder.git
cd starcoder/chat
conda create -n starchat python=3.10 && conda activate starchat
# install pytorch (see pytorch.org) then:
pip install -r requirements.txt
huggingface-cli login
sudo apt-get install git-lfs

torchrun --nproc_per_node=8 train.py config.yaml \
    --deepspeed deepspeed_z3_config_bf16.json
  • 8개의 GPU(A100, 80 GB) 환경을 사용하며, 학습은 약 45분에 완료됩니다.
  • config.yaml은 데이터셋 경로, 모델 이름, 학습 하이퍼파라미터 등을 정의합니다.

샘플 출력 (코딩 작업)

프롬프트 모델 출력
Bar plot – dict에서 키와 나이를 플롯 완전한 pandas/seaborn 스크립트를 생성하고, 두 개의 서브플롯을 포함하며, 그림을 표시합니다.
World map – 독일과 스페인을 빨간색으로 색칠 geopandas 스크립트를 생성하여 GeoJSON을 로드하고, 두 국가를 분리한 뒤 요청된 색상으로 플롯합니다.
Basketball stats – 득점과 리바운드의 산점도 seaborn 산점도를 반환하며, 각 선수에 대한 주석과 적절한 제목/라벨을 포함합니다.

평가

  • 표준 벤치마크 (ARC, HellaSwag, MMLU, TruthfulQA) ...
    모델 ARC HellaSwag MMLU TruthfulQA
    StarCoderBase 0.30 0.46 0.33 0.40
    StarChat‑α 0.33 0.49 0.34 0.44
  • Human‑in‑the‑loop / AI‑in‑the‑loop: 115개의 파이썬 프롬프트(ChatGPT가 생성한 시드)로 구성된 데이터셋을 사용하고, GPT‑4와 ChatGPT가 모델 응답을 1‑8점 척도로 평가했습니다. StarChat‑α는 기본 모델보다 약 96 %의 비율로 우수했습니다.

제한 사항 및 위험

  • 환각 및 안전하지 않은 콘텐츠가 여전히 발생할 수 있습니다(RLHF 또는 안전 필터링 없음).
  • 인구통계적 편향은 대부분의 학습 데이터가 출처인 GitHub 커뮤니티를 반영합니다.
  • 벤치마크는 대화 품질을 포착하지 못하므로 인간 또는 LLM 기반 평가가 필요합니다.

향후 방향

  • 비용 효율적인 파인튜닝을 위해 LoRA 또는 기타 PEFT 방법을 탐색합니다.
  • 코드 중심 데이터와 보다 자연스러운 언어 대화를 결합하여 대화 능력을 향상시킵니다.
  • 보다 넓은 커뮤니티 사용을 위해 오픈소스 어시스턴트로 배포합니다.

리소스

인용

@article{Tunstall2023starchat-alpha,
  author = {Tunstall, Lewis and Lambert, Nathan and Rajani, Nazneen and Beeching, Edward and Le Scao, Teven and von Werra, Leandro and Han, Sheon and Schmid, Philipp and Rush, Alexander},
  title = {Creating a Coding Assistant with StarCoder},
  journal = {Hugging Face Blog},
  year = {2023},
  note = {https://huggingface.co/blog/starchat-alpha},
}

Sources