[Day10] LLM 스터디 1기 - 다중 GPU Llama3 파인튜닝 #1

한 권으로 끝내는 실전 LLM 파인튜닝 - 강다솔 지음/위키북스

10일 차는 이전 Gemma와 달리 파인튜닝 절차를 한 회 차에 모두 진행해서 내용이 좀 많다. 참고로 허깅페이스의 Llama 모델에 접근하려면 미리 요청해서 승인을 받아야 한다. 대략 10분 이내에 승인이 나는 듯하다.

03. 전체 파인튜닝

3.5 다중 GPU를 활용한 Llama3.1-8B-instruct 파인튜닝

지난 시간까지 단일 GPU를 이용해서 Gemma-2B-it 모델을 파인튜닝해 봤다. 이번(10일 차)에는 더 큰 규모의 모델인 Llama3.1-8B-instruct 모델을 대상으로 여러 대의 GPU를 활용한 파인튜닝을 진행한다.

런팟 환경 설정

Llama3.1-8B-instruct 모델을 전체 파인튜닝하며 주피터 노트북 환경이 아닌 터미널에서 직접 스크립트 파일 실행하는 방식으로 진행한다. 원활한 실습을 위해 최소 100GB의 VRAM, GPU는 2~4개가 필요하다. 런팟에서 아래와 같이 선택하고 Deply 버튼을 클릭했다.

출처: https://www.runpod.io/

실습에 필요한 파일들을 얻기 위해 아래 경로의 repository를 clone 받는다.

https://github.com/wikibook/llm-finetuning/

GitHub - wikibook/llm-finetuning: 《한 권으로 끝내는 실전 LLM 파인튜닝》 예제 코드

chapter3/3.5 폴더에 아래 파일들이 있는 것을 확인한다.

• 0_full_fine_tuning_config.yaml
• 1_train_full_fine_tuning
• 2_inference_notebook.ipynb
• 3_test.py
• 4_openai_test.py
• 5_score_notebook.ipynb

Llama 3.1 학습 파라미터 설정

먼저 살펴 볼 0_full_fine_tuning_config.yaml 파일은 파인튜닝에 필요한 설정 정보를 담고 있다. 전체 파일 내용은 아래와 같다.

### 3.5.2. Llama3 학습 파라미터 설정
model_name: "meta-llama/Meta-Llama-3.1-8B-Instruct"
dataset_path: "."
max_seq_length: 512
output_dir: "./llama-3.1-korean-8b-hf-20-epoch"
report_to: "wandb"
learning_rate: 0.00005
lr_scheduler_type: "linear"
num_train_epochs: 5
per_device_train_batch_size: 2
per_device_eval_batch_size: 2
gradient_accumulation_steps: 4
optim: "adamw_torch_fused"
logging_steps: 10
save_strategy: "epoch"
weight_decay: 0.01
max_grad_norm: 0.5
warmup_ratio: 0.03
bf16: true
tf32: true
gradient_checkpointing: true
fsdp: "full_shard auto_wrap"
fsdp_config:
  backward_prefetch: "backward_pre"
  forward_prefetch: "false"
  use_orig_params: "false"

dataset_path는 학습에 사용될 데이터셋이 위치한 경로인데 현재 설정에서는 현재 작업 중인 폴더(.)를 사용한다. max_seq_length는 모델이 처리할 수 있는 최대 시퀀스 길이로 한 번에 처리할 수 있는 입력 텍스트의 최대 길이를 의미한다. 모델의 성능과 메모리 사용량에 직접적인 영향을 미치므로 적절히 조정해야 한다. report_to는 학습 과정에서 생성되는 로그와 성능 지표를 어떤 플랫폼에 기록할지 지정하는 옵션으로 여기서는 wandb를 사용한다. learning_rate는 신경망의 가중치를 얼마나 크게 조정할지 결정하는 중요한 하이파라미터다. 일반적인 파라미터에서는 0.0001~0.0005 사이의 값을 사용하는데 이번 실습에서는 0.00005를 사용한다. 나머지 파라미터들에 대한 자세한 설명은 생략한다.

데이터셋 준비

이어지는 내용은 대부분 1_train_full_fine_tuning.py 파일에 해당하는 내용이다. 이 파일은 크게 '데이터셋 준비', '학습 과정 구현', '실행 부분'으로 구성돼 있는데 '데이터셋 준비'의 주요 코드는 아래와 같다.

system_prompt = "당신은 다양한 분야의 전문가들이 제공한 지식과 정보를 바탕으로 만들어진 AI 어시스턴트입니다. 사용자들의 질문에 대해 정확하고 유용한 답변을 제공하는 것이 당신의 주요 목표입니다. 복잡한 주제에 대해서도 이해하기 쉽게 설명할 수 있으며, 필요한 경우 추가 정보나 관련 예시를 제공할 수 있습니다. 항상 객관적이고 중립적인 입장을 유지하면서, 최신 정보를 반영하여 답변해 주세요. 사용자의 질문이 불분명한 경우 추가 설명을 요청하고, 당신이 확실하지 않은 정보에 대해서는 솔직히 모른다고 말해주세요."
 
train_dataset = dataset.map(
    lambda sample: 
    { 'messages' : [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": sample['instruction']},
        {"role": "assistant", "content": sample['output']}]
                   },
)

실습에서 사용하는 데이터셋은 이준범(beomi)님이 만든 KoAlpaca-v1.1a 데이터셋으로 네이버 지식인의 베스트 질문들을 크롤링해서 수집했다고 한다. system_prompt 변수에는 AI 어시스턴트의 역할과 행동 지침이 자세히 설명돼 있다. 이를 map() 함수를 이용해서 각 샘플을 시스템 프롬프트, 사용자의 지시시항, AI 응답으로 이뤄진 chat_template 형식으로 변환한다.

Llama 3.1 모델 파라미터 설정

다음의 ScriptArguments 클래스 스크립트 실행 시 전달받는 매개변수(파라미터)를 관리하는 클래스다.

@dataclass
class ScriptArguments:
    dataset_path: str = field(
        default=None,
        metadata={
            "help": "데이터셋 파일 경로"
        },
    )
    model_name: str = field(
    default=None, metadata={"help": "SFT 학습에 사용할 모델 ID"}
    )
    max_seq_length: int = field(
        default=512, metadata={"help": "SFT Trainer에 사용할 최대 시퀀스 길이"}
    )
    question_key: str = field(
    default=None, metadata={"help": "지시사항 데이터셋의 질문 키"}
    )
    answer_key: str = field(
    default=None, metadata={"help": "지시사항 데이터셋의 답변 키"}
    )

Llama 3.1 모델 학습 코드 살펴보기

이어지는 training_function 함수는 실제 학습을 수행하는 핵심 함수다.

def training_function(script_args, training_args):    
    # 데이터셋 불러오기 
    train_dataset = load_dataset(
        "json",
        data_files=os.path.join(script_args.dataset_path, "train_dataset.json"),
        split="train",
    )
    test_dataset = load_dataset(
        "json",
        data_files=os.path.join(script_args.dataset_path, "test_dataset.json"),
        split="train",
    )

    # 토크나이저 및 데이터셋 chat_template으로 변경하기      
    # 중간 코드 생략..
    def template_dataset(examples):
        return{"text":  tokenizer.apply_chat_template(examples["messages"], tokenize=False)}
    
    train_dataset = train_dataset.map(template_dataset, remove_columns=["messages"])
    test_dataset = test_dataset.map(template_dataset, remove_columns=["messages"])
    
    # 데이터가 변화되었는지 확인하기 위해 2개만 출력하기 
    with training_args.main_process_first(
        desc="Log a few random samples from the processed training set"
    ):
        for index in random.sample(range(len(train_dataset)), 2):
            print(train_dataset[index]["text"])

먼저 json 형식의 데이터셋을 로딩하고, 지정된 모델에 맞는 토크나이저도 불러(오는 부분의 코드는 생략..)온다. training_args.main_process_first는 GPU를 여러 개 사용하는 분산 학습 환경에서 사용되는데 여러 프로세스 중 메인 프로세스를 지정해 준다. 메인 프로세스가 특정 작업(여기서는 샘플 로깅)을 마칠 때까지 다른 프로세스는 대기하게 된다. 아래는 이어지는 training_function 함수의 나머지 부분이다.

    # Model 및 파라미터 설정하기 
    model = AutoModelForCausalLM.from_pretrained(
        script_args.model_name,
        attn_implementation="sdpa", 
        torch_dtype=torch.bfloat16,
        use_cache=False if training_args.gradient_checkpointing else True,  
    )
    
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    # Train 설정 
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        dataset_text_field="text",
        eval_dataset=test_dataset,
        max_seq_length=script_args.max_seq_length,
        tokenizer=tokenizer,
        packing=True,
        dataset_kwargs={
            "add_special_tokens": False,  
            "append_concat_token": False, 
        },
    )
    
    # 중간 코드 생략..

    if trainer.is_fsdp_enabled:
        trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
    trainer.save_model()

상단의 Model 및 파라미터를 설정하는 코드에서  attn_implementation은 트랜스포머 모델의 기본이 되는 어텐션 연산으로 여기서는 '쿼리, 키, 값'이 세 가지 입력을 사용하는 SDPA(Scaled Dot-Product Attention)을 사용한다. SDPA는 큰 시퀀스를 처리할 때 메모리 사용량과 계산 시간이 많이 소요되는 단점이 있다고 한다. 이 한계를 극복하기 위해 Flash Attention 2가 개발됐는데 연산을 작은 블록 단위로 나눠 계산하는 방식이다. Train 설정 부분의 STFTrainer(Supervised Fine-Tuning)은 지도 학습 방식의 파인튜닝으로 모델을 특정 작업이나 도메인에 맞게 파인튜닝하는데 주로 사용된다. training_function 함수는 여기까지다.

if __name__ == "__main__":

    parser = TrlParser((ScriptArguments, TrainingArguments))
    script_args, training_args = parser.parse_args_and_config()    
    
    if training_args.gradient_checkpointing:
        training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
    
    # set seed
    set_seed(training_args.seed)
  
    # launch training
    training_function(script_args, training_args)

위 코드는 모델 학습을 시작하기 위한 '메인 블록'이다. TrlParser로 스크립트 인자와 학습 인자를 파싱한다. 쭉 넘어가서 시드를 설정해 실험이 재현성을 보장한 뒤, training_function을 호출해서 실제 학습을 시작한다.

Llama 3.1 모델 학습 실행

ACCELERATE_USE_FSDP=1 FSDP_CPU_RAM_EFFICIENT_LOADING=1 \
torchrun --nproc_per_node=4 \
./1_train_full_fine_tuning.py \
--config ./0_full_fine_tuning_config.yaml

터미널에서 위의 명령을 실행하면 학습이 진행된다.

Wandb 설정과 사용

이어서 아래와 wandb 관련 설정을 입력하라는 창이 뜨는데 적절히 입력한다.

아래는 학습 완료된 상태의 터미널 화면이다. 완료까지 대략 1시간 정도가 소요됐다.

wandb 설정 입력 시 표시된 페이지 들어가 보면 다음과 같이 실험(학습) 기록이 표시된다.

대부분 책에 실린 이미지와 비슷한 모양인데 learning_rate(학습률)는 모양이 많이 다르다. 책에는 학습률을 고정해서 진행했기 때문에 직선으로 그려져 있다. 테스트 결과에 문제가 있으면 다시 확인해 보기로 하고 일단 넘어간다.

학습한  Llama 3.1 모델 테스트

model_name = "./llama-3.1-korean-8b-hf-20-epoch/checkpoint-4740"

// 중간 생략..

random_index = randint(0, len(test_dataset))
messages = test_dataset[random_index]["messages"][:2]

terminators = [
    tokenizer.eos_token_id,
]

# Test on sample 
input_ids = tokenizer.apply_chat_template(messages,
                                          add_generation_prompt=True,
                                          return_tensors="pt").to(model.device)

outputs = model.generate(
    input_ids,
    max_new_tokens=512,
    eos_token_id=terminators,
    do_sample=True,
    temperature=0.7,
    top_p=0.95,
)
response = outputs[0][input_ids.shape[-1]:]
print(f"질문:\n{test_dataset[random_index]['messages'][1]['content']}")
print(f"정답:\n{test_dataset[random_index]['messages'][2]['content']}")
print(f"생성:\n{tokenizer.decode(response,skip_special_tokens=True)}")

학습이 완료된 가장 마지막 체크포인트 폴더를 읽어와서 테스트를 진행하는 로직의 주요 코드이다. 훈련용 데이터에서 질문&정답을 가져와서 모델이 나머지 대화를 어떻게 이어가는지 테스트하는 방식이다. 위 코드에 medel.generate 함수가 모델이 실제 응답을 생성하는 부분이다. "do_sample=True, temperature=0.7, top_p=0.95,"은 생성의 다양성과 품질의 균형을 맞추는 설정이라고 한다. 아래는 테스트 결과물이다. 제법 그럴듯하게 나오는 것 같다.

질문:
남극에서 북극성은 볼 수 있을까요? 
오늘 학원에서 이 질문이 생각나서 궁금해졌습니다. 북극성은 북극에 위치한 별인데, 자전축이 직각이 아니어도 보일 수 있는 걸까요?
정답:
북극성은 천구상 북극에 위치한 별이기 때문에, 북극에서만 천정에서 관찰 가능합니다. 23.5도 기울어진 지구의 자전축에 맞춰진 북극에 위치하고 있기 때문입니다. 북극에서 남쪽으로 내려갈수록 북극성의 고도가 낮아지며, 적도 지방에서는 지평선에서 볼 수 있습니다. 따라서 남극에서는 북극성을 관찰할 수 없습니다.
생성:
1. 북극성은 북극에 위치한 별입니다. 
2. 남극에서는 북극성을 볼 수 없습니다. 
3. 북극성은 지평선에 있지 않기 때문에 저녁에 볼 수 있습니다. 
4. 또한, 북극성의 고도는 그 날의 날씨에 따라 다르기 때문에 항상 높게 보이는 것은 아닙니다.


질문:
남이 재채기 할 때 'God bless you'라고 말하게된 원인은?
정답:
옛날 유럽인들은 코를 통해 영혼이 빠져나가기 때문에, 심한 재채기 때 신이 영혼을 데려갈 수도 있기 때문에 "Bless you"라는 표현을 사용해 왔습니다. "God bless you"이라고 말하는 것은 그런 원리를 더욱 강조한 것입니다. 이러한 관습은 새로운 시기에는 전염병 예방 등의 목적으로 사용될 수도 있었습니다. 지금은 예의상 옆에서 재채기를 해도 누구나 이런 표현을 사용하고 있습니다.
생성:
16세기 영국에서 죽은 자를 매장할 때 무덤 앞에서 재채기를 했다가 'God bless you'라는 표현을 사용하게 되었습니다. 당시 죽은 자의 영혼이 무덤에서 나와서 살갗을 건드리지 않기 위해 사용된 표현이었습니다. 이후 이 표현은 일반적인 일상 대화에서도 사용되며, 지금까지 이어져오고 있습니다.

 

생성된 텍스트 데이터 OpenAI로 평가하기

이번에는 생성된 텍스트의 품질을 보다 구체적으로 테스트하기 위해 Open API를 활용해서 평가를 진행한다. 일단 1,000개의 질문/답변 결과 데이터를 ./test/model_genenration_result.txt 경로에 파일로 저장한다. 아래는 Open API 인스턴스를 생성하고 평가 기준을 설정하는 코드다.

def get_openai_client():
    return OpenAI(api_key="Your_OpenAI_API_KEY")

class Criterion(BaseModel):
    score: int
    explanation: str

class Evaluation(BaseModel):
    relevance: Criterion
    accuracy: Criterion
    completeness: Criterion
    clarity: Criterion
    similarity: Criterion
    average_score: float

Evaluation 클래스에 관련성(relevance), 정확성(accuracy), 완전성(completeness), 명확성(clarity), 유사성(similarity) 등의 다양한 평가 기준을 설정한다. 마지막으로 평균 점수(average_score)도 포함시킨다. 이 부분은 커스터마이징 하면서 평가하고 싶은 부분을 수정할 수 있다.

def evaluate_qa_pair(idx, qa_pairs):
    client = get_openai_client()  # 각 프로세스마다 새로운 클라이언트 생성
    save_path = f"./qa_evaluation_results/result_{idx}.json"
    if Path(save_path).exists():
        print(f"인덱스 {idx}에 대한 결과가 이미 존재합니다.")
        return None
    
    question, reference_answer, model_answer = qa_pairs[idx]
    
    prompt = f"""
질문: {question}
참조 답변: {reference_answer}
모델 생성 답변: {model_answer}

위의 질문에 대한 두 답변을 비교 평가해주세요. 다음 기준에 따라 1-10점 사이의 점수를 매겨주세요:
1. 관련성: 모델의 답변이 질문과 얼마나 관련이 있는가?
2. 정확성: 모델이 제공한 정보가 참조 답변과 비교하여 얼마나 정확한가?
3. 완전성: 모델의 답변이 질문에 대해 얼마나 포괄적인가?
4. 명확성: 모델의 답변이 얼마나 명확하고 이해하기 쉬운가?
5. 유사성: 모델의 답변이 참조 답변과 얼마나 유사한가?

각 기준에 대한 점수와 간단한 설명을 제공해주세요. 마지막으로 전체 평균 점수를 계산해주세요.
"""

    try:
        completion = client.beta.chat.completions.parse(
            model="gpt-4o-mini",  # 또는 사용 가능한 최신 모델
            messages=[
                {"role": "system", "content": "QA 모델 응답을 평가하는 임무를 맡은 AI 어시스턴트입니다."},
                {"role": "user", "content": prompt}
            ],
            response_format=Evaluation
        )

위 코드는 프롬프트와 평가한 파일을 저장하는 evaluate_qa_pair 함수의 주요 코드다. 주어진 인덱스의 질문, 참조 답변(정답), 모델 생성 답변을 추출하고, 평가를 위한 프롬프트와 함께 입력 데이터로 만든다. 생략됐지만 결과는 "./qa_evaluation_results/result_{idx}.json" 경로에 다수의 json 파일로 저장된다.

채점 점수 구하기

마지막으로 OpenAI로 채점한 결과를 확인한다. 아래 코드와 같이 개별 아이템의 점수(score) 값들로 직접 계산한 평균과 각 json 파일에 저장된 평균 점수 값(average_score)의 평균을 얻는다.

path = "./qa_evaluation_results"

path_list = glob(f"{path}/*")

calculated_average_list = [] 
model_average_score = [] 
for path in path_list:
    with open(path, "r") as file:
        data = json.load(file)
    
    data = data["choices"][0]["message"]["parsed"]
    scores = [item['score'] for item in data.values() if isinstance(item, dict)]
    calculated_average = sum(scores) / len(scores)
    calculated_average_list.append(calculated_average)
    model_average_score.append(data["average_score"])

mean_calculated_average= sum(calculated_average_list) / len(calculated_average_list)
mean_model_average_score = sum(model_average_score) / len(model_average_score)
mean_calculated_average, mean_model_average_score

두 점수를 함께 비교하는 이유는 average_score 값도 모델이 생성한 결과물이기 때문에 오류 가능성이 있기 때문이라고 한다. 출력된 점수는 (5.3294000000000095, 5.318070000000009)로 10점 만점에 약 5.32점인데 5 에폭만 돌린 것 치고는 양호한 수준이다. 추후에 에폭을 더 늘려서 학습시켜 볼 예정이다.