[Fine-tuning] Llama2 파인튜닝 코드해석 2탄

https://wiz-tech.tistory.com/78에 이어서 진행해보겠습니다.

[Fine-tuning] Llama2 파인튜닝 코드해석

 

9. 모델 로드부터 저장까지

# Load dataset (you can process it here)
dataset = load_dataset(dataset_name, split="train")

# Load tokenizer and model with QLoRA configuration
compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=use_4bit,
    bnb_4bit_quant_type=bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=use_nested_quant,
)

# Check GPU compatibility with bfloat16
if compute_dtype == torch.float16 and use_4bit:
    major, _ = torch.cuda.get_device_capability()
    if major >= 8:
        print("=" * 80)
        print("Your GPU supports bfloat16: accelerate training with bf16=True")
        print("=" * 80)

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map=device_map
)
model.config.use_cache = False
model.config.pretraining_tp = 1

# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training

# Load LoRA configuration
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

# Set training parameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=num_train_epochs,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    save_steps=save_steps,
    logging_steps=logging_steps,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    fp16=fp16,
    bf16=bf16,
    max_grad_norm=max_grad_norm,
    max_steps=max_steps,
    warmup_ratio=warmup_ratio,
    group_by_length=group_by_length,
    lr_scheduler_type=lr_scheduler_type,
    report_to="tensorboard"
)

# Set supervised fine-tuning parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    args=training_arguments,
    packing=packing,
)

# Train model
trainer.train()

# Save trained model
trainer.model.save_pretrained(new_model)

1. 데이터셋 로드 : load_dataset 함수는 Hugging Face의 데이터셋 라이브러리에서 지정된 데이터셋을 로드합니다. 여기서는 train 분할을 사용하여 트레이닝 데이터를 로드합니다.
2. 양자화 설정 로드: BitsAndBytesConfig는 모델의 가중치를 4비트로 양자화하는 설정을 정의합니다. 이는 메모리 사용량을 줄이는 데 도움을 줍니다.
3. GPU 호환성 확인: 이 코드는 GPU가 bfloat16을 지원하는지 확인합니다. bfloat16은 계산 효율성을 높이는 데 사용됩니다.
4. 기본 모델 로드: AutoModelForCausalLM.from_pretrained는 사전 트레이닝된 모델을 로드하는 함수입니다. 여기서는 bnb_config 양자화 설정을 적용하여 모델을 로드합니다.
5. 토크나이저 로드: AutoTokenizer.from_pretrained는 모델에 맞는 토크나이저를 로드합니다. 이 토크나이저는 입력 텍스트를 모델이 이해할 수 있는 형태로 변환하는 데 사용됩니다.
6. LoRA 구성 로드: LoraConfig는 LoRA(저랭크 적응) 설정을 정의합니다. 이 설정은 모델의 파인튜닝 시 파라미터 수를 줄이면서도 성능을 유지하는 데 도움을 줍니다.
7. 트레이닝 파라미터 설정: TrainingArguments는 트레이닝에 필요한 다양한 매개변수를 설정합니다.
8. 감독된 파인튜닝 매개변수 설정: SFTTrainer는 모델, 트레이닝 데이터셋, LoRA 구성 등을 설정하여 감독된 방식으로 모델을 파인튜닝하는 데 사용됩니다.
9. 모델 트레이닝: trainer.train()은 모델을 트레이닝하는 함수입니다.
10. 트레이닝된 모델 저장: trainer.model.save_pretrained()는 트레이닝된 모델을 저장하는 데 사용됩니다.

이 코드는 전반적으로 Llama 2 모델을 데이터셋에 맞게 파인튜닝하고, 트레이닝된 모델을 저장하는 과정을 구현합니다. 이를 통해 모델의 성능을 향상시키고, 특정 작업에 맞게 조정할 수 있습니다.

 

10. Fine-tuning된 모델 테스트

# Ignore warnings
logging.set_verbosity(logging.CRITICAL)

# Run text generation pipeline with our next model
prompt = "What is a large language model?"
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
result = pipe(f"<s>[INST] {prompt} [/INST]")
print(result[0]['generated_text'])

 

output :

<s>[INST] What is a large language model? [/INST] A large language model is a type of artificial intelligence (AI) model that is trained on a large dataset of text to generate human-like language outputs. It is designed to be able to understand and generate text in a way that is similar to human language, and can be used for a variety of applications such as chatbots, language translation, and text summarization.

Large language models are typically trained using deep learning techniques, such as recurrent neural networks (RNNs) or transformer models, and are often based on pre-trained models such as BERT or RoBERTa. These models are trained on large datasets of text, such as books, articles, or websites, and are designed to learn the patterns and structures of language.

Some examples of large language models include:

* BERT (Bidirectional Encoder Representations from Transformers):

 

11. VRAM 비워주기

# Empty VRAM
del model
del pipe
del trainer
import gc
gc.collect()
gc.collect()

 

11. Fine-tuning된 모델 다시 load하고, LoRA  weights와 병합하는 과정

# Reload model in FP16 and merge it with LoRA weights
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map=device_map,
)
model = PeftModel.from_pretrained(base_model, new_model)
model = model.merge_and_unload()

# Reload tokenizer to save it
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

1. 모델 재로드(FP16 형식):
   •   AutoModelForCausalLM.from_pretrained 함수는 사전 트레이닝된 모델을 재로드합니다. 여기서는 모델을 FP16(16비트 부동 소수점) 형식으로 로드하고, CPU 메모리 사용량을 줄이기 위한 설정(low_cpu_mem_usage=True)을 적용합니다.
   •   torch_dtype=torch.float16는 모델이 16비트 부동 소수점 형식으로 로드되도록 지정합니다. 이는 모델의 메모리 사용량을 줄이고, 효율성을 높이는 데 도움을 줍니다.
2. PeftModel 사용하여 LoRA 가중치 병합:
   •   PeftModel.from_pretrained는 사전 트레이닝된 모델에 LoRA 가중치를 적용하고 로드합니다. 이를 통해 파인튜닝된 LoRA 설정이 모델에 적용됩니다.
   •   model.merge_and_unload()는 LoRA 가중치를 모델과 병합하고, 더 이상 필요하지 않은 가중치를 메모리에서 제거하여 효율을 높입니다.
3. 토크나이저 재로드 및 저장:
   •   AutoTokenizer.from_pretrained는 사전 트레이닝된 모델에 해당하는 토크나이저를 재로드합니다. 이는 텍스트를 모델이 이해할 수 있는 형식으로 변환하는 데 사용됩니다.
   •   tokenizer.pad_token = tokenizer.eos_token와 tokenizer.padding_side = "right"는 토크나이저의 패딩 설정을 조정합니다. 이는 모델이 텍스트를 처리할 때 발생할 수 있는 문제를 해결하는 데 도움이 됩니다.

12. 모델 HuggingFace에 업로드

import locale
locale.getpreferredencoding = lambda: "UTF-8".

!huggingface-cli login

model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)

• 모델과 토크나이저 업로드, huggingface에서 write token을 발급받아서 진행하면 된다.

 

이제, 이 과정을 활용해서 ko-llm 을 fine-tuning해보고 openllm-leaderboad에 업로드 해보겠습니다.