Deep_Learning

SFT Train에서의 Dataset의 변환에 관한 이야기

MoonLight314 2025. 6. 7. 20:28
728x90
반응형
 
 

1. Standard & Conversational Dataset

SFT (Supervised Fine-tuning) Trainer에서 사용할 수 있는 Dataset의 종류에 대해서 먼저 알아보도록 하겠습니다.

​

 

 

1.1. Standard( Instruction / Prompt-completion ) 형식

Standard 형식은 주로 Single-turn의 질문-응답 또는 명령어-완성(instruction-completion) 쌍으로 구성이 되어 있습니다.

​

다음과 같은 것들이 Standard 형식의 예가 될 수 있습니다.

  • "질문: 프랑스의 수도는 어디인가요? 답변: 파리입니다."
  • "명령어: 다음 문장을 요약하세요: [긴 문장] 완성: [요약된 문장]"

​

이 형식은 Model이 특정 명령을 이해하고 그에 맞는 응답을 생성하는 능력을 학습시키는 데 주로 사용되며, TRL에서는 주로 {"prompt": "...", "completion": "..."} 또는 단순히 {"text": "..."}와 같은 형태로 처리될 수 있습니다.

​

 

 

1.2. Conversational (대화형) 형식

Conversational 형식은 Multi-turn 대화 형식을 말하며, 주로 사용자와 AI 어시스턴트 간의 주고받은 대화가 순서대로 기록되어 있습니다.

제가 SFT Trainer Post에서 사용한 Capybara Dataset이 대표적인 Conversational 형식이 되겠죠.

​

​

​

​

​

2. only_completion_loss

다음으로 only_completion_loss에 대해서 알아보도록 하겠습니다.

​

 

 

2.1. only_completion_loss ?

only_completion_loss는 TRL Library의 SFTTrainer나 Hugging Face transformers Library의 Trainer에 직접적으로 보여지는 Option이나 Parameter의 이름이 아니고, 응답 부분에만 loss를 계산하는 기능을 나타내는 개념의 이름입니다.

​

구체적으로 completion 부분에 해당하는 Token들은 그대로 남겨두고 prompt 부분의 labels을 -100으로 설정하는 방식으로 구현합니다.

​

​

 

2.2. only_completion_loss 구현 방법

only_completion_loss를 구현하는 방법은 대략 2가지 정도가 있을 수 있습니다.

​

​

1) DataCollatorForCompletionOnlyLM Class 사용

DataCollatorForCompletionOnlyLM Class는 DataCollator를 확장하고 있으며, Dataset의 각 Sample이 prompt와 completion이라는 두 개의 Text 필드를 가지고 있어야 합니다.

​

DataCollatorForCompletionOnlyLM은 prompt와 completion을 합친 전체 Sequence를 input_ids로 만들고, completion 부분에 해당하는 Token들만 labels로 남겨두고 prompt 부분의 labels는 -100으로 설정합니다.

 

아래는 간단한 예제 Code입니다.

 
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("your-tokenizer")
# 이 데이터 콜레이터를 사용하려면 Dataset이 'prompt'와 'completion' 필드를 가지고 있어야 합니다.
# 예: dataset = [{"prompt": "이 문장을 요약해줘: ", "completion": "요약된 내용입니다."}]
data_collator = DataCollatorForCompletionOnlyLM(tokenizer=tokenizer, mlm=False)

trainer = SFTTrainer(
    # ... (Model, Dataset, args 등)
    data_collator=data_collator, # 여기에서 지정
)

 

 

다만, 이 방법은 한계가 있는데 DataCollatorForCompletionOnlyLM은 주로 Single Turn prompt-completion 형식 Dataset에 적합하고, Multi-turn 대화(messages 필드를 가진)에는 직접적으로 사용하기 어렵습니다.

​

conversations 필드 내의 복잡한 대화 흐름을 이해하고 "어디부터 어디까지가 프롬프트(User/System)이고 어디부터 어디까지가 완성(Assistant)이다"라고 자동으로 파악하지 못하고, Assistant 턴이 여러 번 나타날 수도 있기 때문입니다.

 

​

​

2) SFTTrainer의 formatting_func 사용하기

formatting_func은 SFTTrainer의 Parameter중에 하나이며, 입력 Text를 처리하는 역할을 합니다.

​

formatting_func를 사용하면 사용자가 Dataset의 각 Sample을 Model이 학습할 Text Sequence로 직접 변환할 수 있습니다.

​

이 변환 과정에서 Loss를 계산하고 싶은 부분(예: Assistant Turn)만 남겨두고 나머지의 labels를 -100으로 설정하는 로직을 직접 구현해서 only_completion_loss 기능을 구현하는 것이죠.

​

​

​

 

2.3. 기타 Options

1) dataset_text_field

  • SFTTrainer의 Parameter로써 Train에 사용될 Text가 train_dataset 또는 eval_dataset의 어느 필드에 저장되어 있는지 SFTTrainer에게 알려주는 역할을 합니다.
  • 개발자가 formatting_func를 제공하지 않으면, SFTTrainer는 dataset_text_field에 지정된 필드의 Text를 토크나이저로 처리합니다.
  • 또한, packing=False일 때 SFTTrainer는 각 Sample을 개별적으로 처리하는데, 이때 dataset_text_field가 지정되어 있으면 해당 필드에서 Text를 가져와 max_seq_length에 맞춰 Token화하고 패딩합니다.

 

  • dataset_text_field를 사용할 때 주의할 점이 있습니다. 만약 dataset_text_field='text'로 지정했고, 'text'에 해당하는 부분이 질문과 답변이 하나로 연결된 긴 문장(예: "질문: 프랑스의 수도는 어디인가요? 답변: 파리입니다.")이라면 SFTTrainer는 'text' 필드 내에서 질문과 답변을 자동으로 구분하지 않습니다.
  • SFTTrainer는 dataset_text_field에 지정된 필드의 내용을 하나의 연속된 Text Sequence로 간주하고, 이를 통째로 Model의 입력(input_ids)으로 사용하며, 기본적으로 Sequence의 모든 Token에 대해서 손실을 계산합니다.
  • 즉, dataset_text_field는 Text를 "통째로" 가져올 필드를 지정하는 것이며, 그 필드 내의 특정 구조(질문/답변)를 자동으로 파싱하여 Loss를 구분하는 기능은 없습니다. 이를 위해서는 formatting_func와 같은 사용자 정의 로직이 필요합니다.

 

​

​

2) tokenizer.apply_chat_template

  • PreTrainedTokenizer 클래스에 속하는 메소드로써 role(역할)과 content(내용)로 구성된 메시지 리스트를, 해당 Model이 훈련되었던 특정한 대화 형식(Chat Template)을 따르는 하나의 문자열로 변환하는 역할을 합니다.
  • 좀 더 구체적으로, 보통 최신 LLM들은 Model별 고유한 대화 형식을 위한 특수 Token(Special Tokens)과 구분자(Delimiters)를 포함되어야 하는데, apply_chat_template은 이러한 작업을 자동으로 처리하여, Model이 기대하는 정확한 입력 포맷을 생성해 줍니다.
  • 만약 입력 데이터가 {"role": "user", "content": "..."}와 같은 대화형 메시지 리스트 형식이라면, SFTTrainer는 이 메시지 리스트를 tokenizer의 chat_template에 따라 하나의 문자열로 변환되고, 이 문자열에는 <system>, <user>, <assistant> 등의 특수 Token들이 적절히 삽입되어 Model이 학습할 수 있는 형태로 만들어집니다.
  • 이는 각 Model들이 제대로 Train하는데 매우 중요합니다.
  • 'add_generation_prompt' Parameter
    1. apply_chat_template에는 'add_generation_prompt'라는 Parameter가 존재하는데, 이 파라미터는 True로 설정하면 마지막 메시지 뒤에 Model이 다음 대화를 생성해야 함을 나타내는 Token(예: <|im_start|>assistant\n 또는 [INST])을 추가해 줍니다.
    2. 이는 주로 추론 시 Model이 응답을 시작하도록 프롬프트를 구성할 때 유용하지만, SFTTrainer와 함께 사용할 때는 일반적으로 add_generation_prompt=False로 설정합니다.
    3. 그 이유는 훈련 Dataset에는 이미 Model이 생성해야 할 assistant 턴의 내용이 포함되어 있기 때문에, add_generation_prompt로 시작하는 Token을 추가하면 훈련 데이터의 실제 assistant 턴과 충돌할 수 있기 때문입니다.

 

​

3) packing=True Option

  • packing Option의 목적은 여러 개의 짧은 대화 Sequence들을 하나의 긴 Sequence로 "묶어" (max_seq_length까지 채워서) GPU 활용률을 극대화하는 것이었습니다.

​

​

​

​

 

3. 확인

위의 내용들을 바탕으로 제가 이전에 올렸던 SFT Train Code를 살펴보겠습니다.

 
# 3. (선택 사항) SFT 훈련 설정 정의
training_args = SFTConfig(
    output_dir="./sft_qwen_capybara", # Model 및 체크포인트 저장 경로
    num_train_epochs=1,              # 훈련 에포크 수
    
    per_device_train_batch_size=4,   # 장치당 배치 크기

    gradient_accumulation_steps=1,   # 그래디언트 누적 스텝
    learning_rate=2e-4,              # 학습률
    logging_steps=10,                # 로깅 간격
    max_seq_length=512,              # 최대 Sequence 길이 (메모리 사용량에 영향)
    # packing=True,                  # 여러 짧은 Sequence를 묶어 효율성 증대 (선택 사항)
)

# 4. SFTTrainer 초기화
trainer = SFTTrainer(
    model=model,                     # 훈련할 Model
    tokenizer=tokenizer,             # 토크나이저
    args=training_args,              # 훈련 설정
    train_dataset=dataset,           # 훈련 Dataset
    #dataset_text_field="text",       # Dataset에서 Text 필드 지정 (packing=False 일 때)
    # formatting_func=formatting_prompts_func, # 사용자 정의 포맷팅 함수 (선택 사항)
    # peft_config=lora_config,       # PEFT 설정 (LoRA 등 사용 시)
)

 

 

보시다시피, DataCollatorForCompletionOnlyLM을 사용하지도 않았고 , formatting_func을 지정하지도 않았으며, SFTConfig의 packing=True Option도 사용하지 않았습니다.

​

그러므로, only_completion_loss 기능은 사용되지 않았습니다.

​

​

​

​

​ 4. only_completion_loss 적용

이제 실제로 only_completion_loss 기능을 구현해서 적용해보도록 하겠습니다.

​

Train Code는 기존 SFT Train Post에서 사용했던 Train Code를 Base로 하겠습니다.

 

Example of SFT(Supervised Fine-Tuning) Trainer in TRL

안녕하세요, MoonLight입니다.​지난 번 Post에서 LLM을 Fine-tuning하고 Rreinforcement learning을 적용하는 데 사용되는 도구 모음인 TRL(Transformer Reinforcement Learning)에 대해서 알아보았습니다.https://moonlight314.

moonlight314.tistory.com

 

 

 
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
import torch

# 1. 데이터셋 로드 (예: trl-lib/Capybara 데이터셋)
dataset = load_dataset("trl-lib/Capybara", split="train")

# 2. 사전 훈련된 모델 및 토크나이저 로드 (예: Qwen/Qwen2.5-0.5B)
model_name = "Qwen/Qwen2.5-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 패딩 Token 설정 (모델에 따라 필요할 수 있음)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


# ----------------- only_completion_loss 구현을 위한 formatting_func 추가 -----------------
def formatting_func_with_apply_chat_template_and_masking(example):
    """
    tokenizer.apply_chat_template을 활용하여 텍스트를 포맷팅하고,
    labels를 수동으로 마스킹하여 only_completion_loss를 구현하는 함수.
    """
    all_input_ids = []
    all_labels = []

    for i, message in enumerate(example['messages']):
        current_turn_messages = [message]

        encoded_turn_tensor = tokenizer.apply_chat_template(
            current_turn_messages,
            tokenize=True,
            add_generation_prompt=False,
            return_tensors="pt" # PyTorch 텐서로 반환
        )        

        # 텐서 자체에서 input_ids를 가져와 리스트로 변환
        turn_input_ids = encoded_turn_tensor.squeeze(0).tolist() # .input_ids 제거
        
        # 변환 확인용
        print("\n원래 Message : ",current_turn_messages)        
        print("tokenizer.apply_chat_template 변환 : ",turn_input_ids)

        # 레이블 마스킹 여부 결정
        if message['role'] == "assistant":
            turn_labels = list(turn_input_ids)
        else: # user, system, 또는 다른 역할
            turn_labels = [-100] * len(turn_input_ids)

        print("Label Masking 적용 후 : ",turn_labels)
        
        all_input_ids.extend(turn_input_ids)
        all_labels.extend(turn_labels)

    return {
        "input_ids": all_input_ids,
        "labels": all_labels,
        "attention_mask": [1] * len(all_input_ids)
    }

# 3. (선택 사항) SFT 훈련 설정 정의
training_args = SFTConfig(
    output_dir="./sft_qwen_capybara", # 모델 및 체크포인트 저장 경로
    num_train_epochs=1,              # 훈련 에포크 수
    
    per_device_train_batch_size=4,   # 장치당 배치 크기

    gradient_accumulation_steps=1,   # 그래디언트 누적 스텝
    learning_rate=2e-4,              # 학습률
    logging_steps=10,                # 로깅 간격
    max_seq_length=512,              # 최대 시퀀스 길이 (메모리 사용량에 영향)
    packing=True,                    # 여러 짧은 시퀀스를 묶어 효율성 증대 (only_completion_loss와 함께 사용 권장)
)

# ----------------- 레이블 마스킹 확인 코드 추가 -----------------
print("\n=== 레이블 마스킹 확인 시작 (tokenizer.apply_chat_template 활용) ===")

# 데이터셋에서 첫 번째 샘플 가져오기
sample_to_check = dataset[0] # 'text' 필드에 접근하지 않고 샘플 전체를 가져옴

# only_completion_loss 로직을 적용하여 처리
processed_sample = formatting_func_with_apply_chat_template_and_masking(sample_to_check)

# 결과를 PyTorch 텐서로 변환
input_ids_tensor = torch.tensor(processed_sample['input_ids'])
labels_tensor = torch.tensor(processed_sample['labels'])

print(f"\n원본 데이터셋 샘플 (messages 필드):\n{sample_to_check['messages']}\n")

print(f"**포맷팅 및 Token화된 Input IDs (일부):**\n{input_ids_tensor[:100]}...\n") # 처음 100개 Token만 출력
print(f"**생성된 Labels (마스킹 포함, 일부):**\n{labels_tensor[:100]}...\n") # 처음 100개 레이블만 출력

# Input ID와 Label을 디코딩하여 시각적으로 비교
print("\n**Input Token, Decoded Text, and Corresponding Label (상위 50개):**")
print("-" * 50)
print(f"{'Index':<5} | {'Input Token':<15} | {'Decoded Text':<20} | {'Label':<10}")
print("-" * 50)

for i in range(min(len(input_ids_tensor), 50)): # 상위 50개 Token만 확인
    input_id = input_ids_tensor[i].item()
    label_value = labels_tensor[i].item()
    decoded_token = tokenizer.decode([input_id], skip_special_tokens=True)
    
    # 특수 Token인 경우 (예: <|im_start|>) decoded_token이 빈 문자열일 수 있어 조정
    # Qwen 토크나이저는 <|im_start|> 같은 Token을 decode하면 빈 문자열을 반환하는 경우가 많습니다.
    # Token 자체를 확인하고 싶다면 convert_ids_to_tokens를 사용합니다.
    if not decoded_token.strip(): # 디코딩된 텍스트가 비어있다면
        try:
            # 특수 Token의 경우 일반적으로 convert_ids_to_tokens가 더 정확한 문자열을 제공
            decoded_token = tokenizer.convert_ids_to_tokens(input_id)
        except Exception:
            decoded_token = f"ID:{input_id}" # 변환 실패 시 ID만 표시

    print(f"{i:<5} | {input_id:<15} | {decoded_token:<20} | {label_value:<10}")

print("-" * 50)
print("\n`-100`으로 표시된 Labels는 해당 Token의 손실이 계산되지 않음을 의미합니다.")
print("=== 레이블 마스킹 확인 완료 ===")


# 4. SFTTrainer 초기화
trainer = SFTTrainer(
    model=model,                     # 훈련할 모델
    tokenizer=tokenizer,             # 토크나이저
    args=training_args,              # 훈련 설정
    train_dataset=dataset,           # 훈련 데이터셋
    formatting_func=formatting_func_with_apply_chat_template_and_masking, # only_completion_loss 구현을 위해 사용자 정의 포맷팅 함수 지정
)

# 5. 훈련 시작
trainer.train()

# 6. (선택 사항) 훈련된 모델 저장
trainer.save_model("./sft_qwen_capybara_final") # 모델과 토크나이저 설정을 함께 저장
print("모델 훈련 및 저장이 완료되었습니다.")
 
 

추가된 부분은 only_completion_loss 구현을 위해서 formatting_func로 사용할 formatting_func_with_apply_chat_template_and_masking를 추가했다는 점입니다.

​

​

​

실제 Tokenizer를 통과한 Data가 어떻게 변하는지도 확인할 수 있습니다.

​

아래 Code는 Dataset에서 하나의 Data를 가져와서 formatting_func_with_apply_chat_template_and_masking()에 입력합니다.

 

# 데이터셋에서 첫 번째 샘플 가져오기
sample_to_check = dataset[0] # 'text' 필드에 접근하지 않고 샘플 전체를 가져옴

# only_completion_loss 로직을 적용하여 처리
processed_sample = formatting_func_with_apply_chat_template_and_masking(sample_to_check)

 

​tokenizer.apply_chat_template를 Text가 통과하면 Token값으로 바뀌고, 실제 Text의 역할에 따라서 -100으로 Masking하는 Logic입니다.

encoded_turn_tensor = tokenizer.apply_chat_template(
    current_turn_messages,
    tokenize=True,
    add_generation_prompt=False,
    return_tensors="pt" # PyTorch 텐서로 반환
)        

# 텐서 자체에서 input_ids를 가져와 리스트로 변환
turn_input_ids = encoded_turn_tensor.squeeze(0).tolist() # .input_ids 제거

# 변환 확인용
print("\n원래 Message : ",current_turn_messages)        
print("tokenizer.apply_chat_template 변환 : ",turn_input_ids)

# 레이블 마스킹 여부 결정
if message['role'] == "assistant":
    turn_labels = list(turn_input_ids)
else: # user, system, 또는 다른 역할
    turn_labels = [-100] * len(turn_input_ids)

print("Label Masking 적용 후 : ",turn_labels)

 

​

아래 2개의 서로 다른 Role을 가진 Text Data가 어떻게 변화하는지 알 수 있습니다.

원래 Message : [{'content': 'Recommend a movie to watch.\n', 'role': 'user'}]
tokenizer.apply_chat_template 변환 : [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 872, 198, 67644, 264, 5700, 311, 3736, 624, 151645, 198]
Label Masking 적용 후 : [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]

​

 

원래 Message : [{'content': 'I would recommend the movie, "The Shawshank Redemption" which is a classic drama film starring Tim Robbins and Morgan Freeman. This film tells a powerful story about hope and resilience, as it follows the story of a young man who is wrongfully convicted of murder and sent to prison. Amidst the harsh realities of prison life, the protagonist forms a bond with a fellow inmate, and together they navigate the challenges of incarceration, while holding on to the hope of eventual freedom. This timeless movie is a must-watch for its moving performances, uplifting message, and unforgettable storytelling.', 'role': 'assistant'}]
tokenizer.apply_chat_template 변환 : [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 77091, 198, 40, 1035, 6934, 279, 5700, 11, 330, 785, 35185, 927, 1180, 88641, 1, 892, 374, 264, 11416, 19584, 4531, 39400, 9354, 87215, 323, 22954, 49564, 13, 1096, 4531, 10742, 264, 7988, 3364, 911, 3900, 323, 54962, 11, 438, 432, 11017, 279, 3364, 315, 264, 3908, 883, 879, 374, 4969, 3641, 23088, 315, 9901, 323, 3208, 311, 9343, 13, 88689, 267, 279, 24939, 49346, 315, 9343, 2272, 11, 279, 45584, 7586, 264, 10815, 448, 264, 12357, 66435, 11, 323, 3786, 807, 20876, 279, 11513, 315, 69052, 11, 1393, 9963, 389, 311, 279, 3900, 315, 41735, 11290, 13, 1096, 57005, 5700, 374, 264, 1969, 84296, 369, 1181, 7218, 23675, 11, 94509, 1943, 11, 323, 59998, 47829, 13, 151645, 198]
Label Masking 적용 후 : [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 77091, 198, 40, 1035, 6934, 279, 5700, 11, 330, 785, 35185, 927, 1180, 88641, 1, 892, 374, 264, 11416, 19584, 4531, 39400, 9354, 87215, 323, 22954, 49564, 13, 1096, 4531, 10742, 264, 7988, 3364, 911, 3900, 323, 54962, 11, 438, 432, 11017, 279, 3364, 315, 264, 3908, 883, 879, 374, 4969, 3641, 23088, 315, 9901, 323, 3208, 311, 9343, 13, 88689, 267, 279, 24939, 49346, 315, 9343, 2272, 11, 279, 45584, 7586, 264, 10815, 448, 264, 12357, 66435, 11, 323, 3786, 807, 20876, 279, 11513, 315, 69052, 11, 1393, 9963, 389, 311, 279, 3900, 315, 41735, 11290, 13, 1096, 57005, 5700, 374, 264, 1969, 84296, 369, 1181, 7218, 23675, 11, 94509, 1943, 11, 323, 59998, 47829, 13, 151645, 198]

​

 

위에서 보시는 바와 같이, Text가 assistant, 즉 답변 부분으로써 우리가 예측하고자 하는 값은 Label이 그대로 나오고, 그 이외에 부분은 -100으로 변하는 것을 알 수 있습니다.

​

​

도움이 되셨기를 바라면서 이만 마칠까 합니다.

​

읽어주셔서 감사합니다.

728x90
반응형
저작자표시 (새창열림)

'Deep_Learning' 카테고리의 다른 글

Example of SFT(Supervised Fine-Tuning) Trainer in TRL  (2) 2025.06.07
TRL (Transformer Reinforcement Learning)  (6) 2025.06.07
Weights & Biases (wandb)  (0) 2025.05.13
Alignment in LLM  (0) 2025.05.13
Downstream in LLM  (0) 2025.04.19

'Deep_Learning'의 다른글

  • 현재글SFT Train에서의 Dataset의 변환에 관한 이야기

관련글

티스토리툴바