HuggingFace 실습(PEFT) : 1. Data

허깅페이스를 이용해 PEFT 튜닝 실습을 해보자!

 

허깅페이스에서 제공하는 PEFT 라이브러리를 사용하여 GPT-2 모델을 prefix tuning 해볼 것이다.

prefix tuning 의 개념은 이전에 블로그에서 설명한 글을 참조하자!

https://mari970.tistory.com/53

 

참고할 코드는 아래와 같다.

https://huggingface.co/docs/peft/task_guides/seq2seq-prefix-tuning

Prefix tuning for conditional generation

 

PEFT 를 이용하여 튜닝을 하는 실습 코드에 대해 포스팅하지만

일반 실습을 하는 사람들도 충분히 도움이 될 수 있는 글이라고 생각한다!

---

 

from transformers import GPT2Tokenizer, AutoTokenizer, GPT2Model, AutoModelForCausalLM, get_linear_schedule_with_warmup, default_data_collator
from peft import PrefixTuningConfig, get_peft_model, TaskType # 따로 peft install 해야함
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch
import os
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', '-e', default=20, type=int,
                        dest='epochs', help='training epoch')
    parser.add_argument('--learning-rate', '-lr', default=1e-2, type=float,
                        dest='lr', help='training learning rate')
    parser.add_argument('--batch-size', '-bs', default=8, type=int,
                        dest='batch_size', help='training batch size')
    parser.add_argument('--max_length', '-ml', default=1004, type=int, # 1024 인데 prefix length=20 이라서,
                        dest='max_length', help='maximum sequence length')
    parser.add_argument('--seed', type=int, default=None)
    parser.add_argument('-save_mode', type=str, choices=['all', 'best'], default='best')
    parser.add_argument('--model_name_or_path', default= 'gpt2-large',
                        dest ='model_name_or_path', help='base model')
    parser.add_argument('--result_dir', default='output',
                        dest = 'result_dir', help='experiment result save directory')

    parser.add_argument('--data_preprocess', default='concat', choices = ['def_clm', 'concat'],
                        dest = 'data', help='data preprocess method for Causal LM')
    args = parser.parse_args()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    peft_config = PrefixTuningConfig(
        task_type=TaskType.CAUSAL_LM, # TaskType.SEQ_2_SEQ_LM, 
        inference_mode=False, 
        num_virtual_tokens=20
    )

    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path,
                                              pad_token='<pad>')

    model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
    # model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
    model.resize_token_embeddings(len(tokenizer))

    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    dataset = load_dataset("bigscience/P3", name="xsum_summarize_this_DOC_summary")

    if args.data == 'def_clm':
        def preprocess_func(examples):
            inputs = examples['inputs_pretokenized']
            targets = examples['targets_pretokenized']
            model_inputs = tokenizer(inputs, padding='max_length', truncation=True, max_length= args.max_length, return_tensors='pt') # is_split_into_words = True,
            labels = tokenizer(targets, padding='max_length', truncation=True, max_length= args.max_length, return_tensors='pt')
            labels = labels["input_ids"]
            # labels[labels == tokenizer.pad_token_id] = -100 # ?
            model_inputs["labels"] = labels
            return model_inputs


        tokenized_dataset = dataset.map(
            preprocess_func,
            batched=True,
            num_proc = 1,
            remove_columns=dataset["train"].column_names
            )

    elif args.data == 'concat':
        # Concat inputs and targets for CLM training
        dataset = dataset.map(
        lambda x : {'sent_forclm' : [x['inputs_pretokenized'][i] + x['targets_pretokenized'][i].lstrip() for i in range(len(x['targets_pretokenized']))]},
        batched= True,
        remove_columns=dataset["train"].column_names,
        num_proc = 1)

        tokenized_dataset = dataset.map(
            lambda examples : tokenizer(examples['sent_forclm'], padding='max_length', max_length=args.max_length, truncation=True, return_tensors="pt"),
            batched=True,
            num_proc = 1
            )

    train_dataset = tokenized_dataset['train']
    eval_dataset = tokenized_dataset['validation']

    train_loader = DataLoader(train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.batch_size, pin_memory=True)
    eval_loader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=args.batch_size, pin_memory=True)


    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    lr_scheduler = get_linear_schedule_with_warmup(
        optimizer = optimizer,
        num_warmup_steps = 8e4,
        num_training_steps = (len(train_loader)*args.epochs)
    )

    train(model, train_loader, eval_loader, optimizer, lr_scheduler, device, args)



if __name__ == '__main__':
    main()

위는 참고 사이트에서 조금 바꾼 코드이다. 이를 레퍼런스로 데이터를 로딩하는 것 부터 토크나이즈 하는 전처리 과정까지 알아보자!

 

(train() 은 다음 블로그에서 소개할 것이다.)

 

1. Data Loading

dataset = load_dataset("bigscience/P3", name="xsum_summarize_this_DOC_summary")

P3 데이터셋, 그 중에 XSUM (summarization 요약) 데이터를 사용할 것이다.

허깅페이스에서 제공하는 데이터셋을 링크로 받아와서 load_dataset 함수를 이용해 사용할 수 있다.

원하는 데이터셋이 있는 허깅페이스 사이트에 들어가면 아래와 같은 화면을 볼 수 있는데 "</> Use in dataset library" 라는 버튼을 누르면 위 처럼 써야하는 코드를 바로 확인할 수 있다!

https://huggingface.co/datasets/bigscience/P3

 

** P3 데이터셋은 Public Pool of Prompts 로, 일반 데이터셋에 프롬프트 문장을 추가한 데이터셋이다.

 

2. Tokenizer

2-1. 선언

tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, pad_token='<pad>')

위와 같이 GPT-2 의 학습된 tokenizer 를 불러오기 위해 from_pretrained 함수를 사용하는데, 

pre-trained 된 tokenizer 는 PAD 토큰이 학습되어있지 않아서 위와 같이 pad_token='<pad>' 를 추가해주어야 한다. 

 

** 기존에 학습된 토크나이저를 사용하다 보면 패딩해야 하는데 막 pad 토큰이 없어서 따로 지정을 해줘야 한다는 오류가 나기도 한다.

 

ValueError:
Asking to pad, but the tokenizer does not have a padding token. Please select a token to use as 'pad_token' '(tokenizer.pad_token = tokenizer.eos_token e.g.)' or add a new pad token via 'tokenizer.add_special_tokens({'pad_token': '[PAD]'})'.

Error 메세지에서 친절하게 어떻게 해결하는지 방법도 알려준다. 

필자가 위에 기재한 코드는 tokenizer 초기부터 add 해주는 방법이고, 위와 같이 이후에 add_special_tokens() 을 사용할 수도 있다!

 

그렇지만 이렇게 vocab 을 1개 추가하면 vocab size 가 커져서 나중에 모델안에서 embedding layer 를 거칠 때 (예를들어 GPT-2 의 wte layer) Index out of range 문제가 뜬다.

이를 해결하기 위해서

model.resize_token_embeddings(len(tokenizer))

를 사용한다.

 

2-2.  전처리 

tokenized_dataset = dataset.map(
            lambda examples : tokenizer(examples['sent_forclm'], padding='max_length', max_length=args.max_length, truncation=True, return_tensors="pt"),
            batched=True,
            num_proc = 1)

Huggingface dataset .map() 함수 를 사용하는데, batched =True 를 이용하여 batch 단위로 map() 함수를 사용하여 데이터 처리 속도를 빠르게 할 수 있다.

https://huggingface.co/docs/datasets/about_map_batch

Batch mapping

 

2-3. tokenizer 의 사용

padding

boolean 으로, True 를 하면 batch 안에서 제일 긴 길이에 대해 <pad> 토큰으로 패딩을 해주고, ‘max_length’ 로 하면 다른 변수인 max_length 에서 설정한 최대 길이로 패딩을 해준다.

 

truncation

도 또한 또한 패딩과 관련있는 변수인데, 최대 길이를 넘어가는 데이터가 있을 때 잘라주는 역할을 한다.

 

return_tensors

'tf' 는 tensorflow, 'pt' 는 Pytorch, 'np' 는 numpy 로, 각 라이브러리에 맞는 자료형으로 출력해준다.

 

return_attention_mask

attention_mask는 tokenizer 를 통과하면 토큰화된 출력은 [input_ids’] 과 [’attention_mask’] 로 이루어져 있는데 이때 input_ids 는 토큰화된 문장이고, ‘’attention_mask’ 는 실제 토큰이 존재하는 부분과 <pad> 토큰으로 이루어진 부분을 1 과 0 으로 표현하여 실제 문장 부분을 표시한다.

 

3. DataLoader 선언

train_loader = DataLoader(train_dataset, shuffle=True, collate_fn=default_data_collator, 
						  batch_size=args.batch_size, pin_memory=True)

3-1. DataLoader DataCollator 차이점

Data loader: 프로그램을 실행하기 위해 하드디스크에 있는 파일을 RAM 에 올리는 과정이 필요한데 그 과정을 한다.

Data collator : 데이터셋 리스트에서 batch 를 만들어내는 역할을 한다.

 

DataLoader 는 보통 Pytorch 에서 구현한 기능이고

DataCollator 는 HuggingFace 라이브러리에서 구현된 함수이다.

 

DataLoader 에 있는 collate_fn 이라는 함수(? 아니면 변수) 를 이용해 batch 로 묶을 수 있다. → dataset 이 variable length 를 가지면 (padding 을 하기 위해) collate_fn 을 꼭 사용해야 한다고 한다.

 

** 요즘 보통 data loader 가 아니라 data collator 를 사용하는 예제가 많이 보인다. 아직 이 둘이 쓰임새 차이를 확인하지 못했다 ㅜㅠ..

 

3-2. DataLoader 의 pin_memory 변수

학습을 시키기 위해 GPU 에 데이터를 로드해줘야 하는데 이때 빠르게 데이터를 옮기기 위해 사용하는 것이 pin memory 이다.

즉, 원래 지피유에 데이터를 올리기 위해서는 DRAM(pageable memory) 에서 pinned memory 로 옮겨지고 이후 VRAM(GPU 전용 메모리) 에 옮기는 과정이 필요한데,

옮기는 과정이 매우 많기 때문에 느려지게 된다.

이를 보완하기 위해 Pytorch 의 DataLoader 에서는 pin_memory 라는 옵션을 이용하여 DRAM 을 거치지 않고 바로 VRAM 으로 할당시켜주는 역할을 한다.

https://mopipe.tistory.com/191

[pytorch, 딥러닝] pin_memory 란 무엇인가?

 

** Subset 사용하기

디버깅을 할 때 큰 데이터셋으로 작성한 코드가 잘 돌아가는지 확인하는 것은 시간이 너무 많이 걸린다. 데이터셋을 작게 줄여서 사용하거나, 원하는 데이터의 인덱스만 뽑아서 사용하고 싶을 때 torch 의 Subset 을 사용할 수 있다.

from torch.utils.data import DataLoader, Subset

num_train_idxs = list(range(8))
train_dataset = Subset(dataset['train'], num_train_idxs) 
eval_dataset = Subset(dataset['validation'], num_train_idxs)

train_loader = DataLoader(train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.batch_size, pin_memory=True)
eval_loader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=args.batch_size, pin_memory=True)

기존에 load_dataset() 으로 로드한 dataset 이 있을 때, 다음과 같이 index 리스트를 만들어 Dataset class 를 샘플링할 수 있다.

하지만 실수해서는 안되는 부분이 이렇게 Subset 한 train_이나 eval_dataset 의 class 는 Dataset 이 아니라 Subset 이다.

Subset 은 이후에 DataLoader 를 이용해 로드를 해야 비로소 샘플링된 상태가 된다.