LLM/LVM을 SFT, DPO, Pretrain으로 학습하고 전처리하는 코드를 정리한 repo. 자주 사용하는 학습 방식을 코드로 굳혀둔 것이고, 매번 만들던 1회용 코드를 최대한 줄이고 GPU를 낭비 없이 효율적으로 사용하는데 초점이 맞춰져 있다.
- LLM 학습은 너무 비효율적이다. GPU 자원을 최대로 쓰면서 처리량을 올리는 방법들을 구현.
- 내가 자주 하는 학습 스타일을 코드로 정리해 굳혀둔다.
- 1회용 작업을 최대한 줄인다.
{method}/ 아래는 항상 이렇게 둔다.
main.py(필수): args +def main+if "__main__" in __name__preprocessor.py: 데이터 전처리- 그 외(callback 등)는
callbacks.py같은 파일로만 분리
예전엔
args.py,collator.py로 쪼갰더니 파일과 VSCode 창만 늘고 한눈에 안 들어왔다. Python 모듈 철학과는 어긋나지만main.py통합을 택했다. 파일 수를 최소화하는 게 목적이다.
위에서부터 args → def main → def train / valid / predict 순서로 둔다.
train/valid 로직을 main 안에 박지 않고 함수로 뺀 이유는, 박아두면 코드 분리가 안 돼서 한눈에 안 보이기 때문이다.
if "__main__" in __name__: 블록은 logging, setproctitle, set_seed, args 파싱처럼 한 줄로 끝나는 잡설정 전용 장소다.
args를 model/data/train으로 나누지 않고 하나로 상속해 합쳤다.
나눠두면 wandb엔 train args만 기록돼서 model, dataset 설정이 재현성 추적에서 누락된다. 실험 관리를 wandb로 하다 보니 전부 한 곳에 기록되게 합쳤다.
학습마다 별도 전처리 코드를 짜는 대신, 자주 쓰는 기능을 args 맵으로 정규화했다.
dataset_name_map: subset(name) 매핑dataset_prefix: split → 역할(train/valid/test) 매핑dataset_truncate_map: 샘플 수 조절 / 분할dataset_files_map: 로컬 파일 매핑
여러 dataset을 섞어 쓰면 split, 분할, 이상치 처리 같은 1회용 작업이 매번 따라붙는데, 이걸 맵으로 빼서 args만 바꿔 끝낼 수 있게 했다.
권장 입력 컬럼:
- SFT:
conversations,prompt,answer,images - Pretrain:
corpus,sentence_ls
대부분의 args는 HF TrainingArguments를 그대로 따른다. SFT는 추가로 TRL SFTConfig, ModelConfig, DPO는 DPOConfig, ModelConfig를 상속. 내가 추가한 args 기능은 아래와 같다.
데이터
dataset_repo_ls: 학습에 쓸 dataset repo 목록dataset_name_map/dataset_prefix/dataset_truncate_map/dataset_files_map: 위 데이터 전처리 참고dataset_batch_size: 전처리 배치 크기
모델/토크나이저
config_kwargs: AutoConfig에 넘길 추가 인자tokenizer_kwargs(SFT) /processor_kwargs(DPO): 토크나이저, 프로세서 추가 인자chat_template_path(DPO): chat template 지정lora_kwargs: LoraConfig에 넘길 추가 인자 (peft_kwargs로 LoRA 외 PEFT 타입 확장도 가능)
학습/평가
dataset_type(SFT): sft / pretrain / dpo 선택packing/packing_strategy/eval_packing: packing 제어eval_harness_tasks: 학습 중 lm-eval-harness로 평가할 태스크 목록gpu_mem_check: forward/backward/optimizer 구간별 GPU peak memory 로깅
참고로 wandb로 report하면 코드 스냅샷을 아티팩트로 같이 저장해 재현성을 남긴다.
콜백은 src/callbacks.py에 모여 있고, SFT, DPO가 공용한다. 해당 args가 켜질 때만 붙는다.
EvalHarnessCallBack(eval_harness_tasks): 학습 중 eval 시점에 lm-eval-harness를 같이 돌린다. 예전엔 학습이 끝난 뒤 checkpoint를 하나씩 lm eval 돌렸는데, 너무 번거롭고 시간도 많이 먹었다. trainer에 eval/predict loop가 있는 이유가 자동화해 사람 공수를 줄이는 것이라서 lm eval도 그 안으로 포함시킴.WandbCodeArtifactCallback(report_to에 wandb): base code 대비 뭐가 바뀌었고 실험이 어떻게 되는지를 artifact로 남긴다. 재현성, 버전 관리를 wandb로 하기 때문이다.GpuMemoryCallback(gpu_mem_check): 단순 memory 디버깅용. 구간별 GPU peak memory를 찍는다.
SFT 스크립트 하나로 sft와 pretrain을 둘 다 돌린다. dataset_type만 sft / pretrain으로 바꾸면 된다.
전처리(assistant-only 라벨링 vs 통문장 토큰화)와 collator의 라벨 처리(labels vs input_ids)가 이 값에 따라 갈릴 뿐, packing, 캐시, 데이터 맵 같은 나머지 파이프라인은 그대로 공유한다. 그래서 같은 데이터 구성 위에서 pretrain 후 sft로 이어가는 식의 실험을 args만 바꿔서 할 수 있다.
전처리 캐시는 최초 1회 실행에서만 생성된다. 최초 1회는 main_process_first 컨텍스트로 main process에서만 만들고, 이후 실행은 datasets가 그 캐시를 그대로 로드한다.
캐시 위치는 원본 데이터와 분리하고, 상태(모델/repo/run/max_length)별로 나눠 저장한다.
datasets는 arrow 캐시를 원본과 같은 폴더에 쌓는데, 캐시를 정리하다 원본까지 지우는 사고가 잦았다. 그래서 분리했다.
git clone https://github.com/jp1924/LLM.git
docker compose up -d # 도커 환경 구성
docker compose stop # 도커 환경 종료
Docker/Compose: 27.3.1-build ce12230, v2.29.7 패키지 버전이 빨리 바뀌어서 실행 안 되면 Dockerfile 수정이 필요할 수 있다.
학습 실행은 scripts/ 아래 스크립트를 쓴다.
bash scripts/run_zero3_sft.sh
bash scripts/run_zero3_dpo.sh
pypi로 받으면 컴파일에 오래 걸리지만, flash-attn repo가 주는 whl은 컴파일이 끝나 있어 빠르다. 먼저 환경 버전을 확인한다.
python -c "import torch; print(torch.__version__, torch.version.cuda, torch.compiled_with_cxx11_abi())"
이후 flash-attn repo에서 환경에 맞는 whl을 받아 설치하면 된다.
내가 주로 작업하는 방식. 그냥 참고용임.
- docker 환경
/root/workspace안에서 작업한다. 초기 파이썬 세팅은uv sync로 맞춘다. - 주로 VSCode + tmux로 작업.
.vscode에 자주 쓰는 vscode extensions과 debug에 필요한 설정은launch.json에 위치. - tmux, bashrc, vi 같은 개발환경 구성 파일은
env폴더에 - 학습 로그는
tee로logs/{method}.log에 남긴다.
실험 관리는 wandb로 한다. 그래서 config/{method}.yaml을 실험마다 복사하거나 이름을 바꿔 차이를 설명하지 않고, WANDB_NOTES env로 실험별 특징을 적는 걸 선호한다.
로그도 {method}.log 하나로 둔다. 복사본을 만들면 나중에 diff, 버전 관리가 어려워지기 때문이다.