추론 모델을 구현해보자!

AI/LLM

추론 모델을 구현해보자!

Tech코알라 2025. 3. 22. 13:50

앞선 포스팅에서 EXAONE-Deep 관련 내용을 다루면서 추론 모델들이 굉장히 큰 인기를 끌면서 많은 추론 모델들이 출시되었다고 했는데요. 저는 Deepseek R1과 같은 모델들을 사용만 했지 실제로 어떻게 구현되는지에 대해서는 무관심했던 것 같아서 이번 포스팅에서는 실제로 어떻게 구현되는지에 대해서 작성해보고자 합니다.

이번 포스팅은 SK DEVOCEAN 사이트에 올라와 있는 "생각하는 AI? 추론 모델 빠르게 구현해 보기 (ft. S1)" 이라는 포스팅을 참고하여 작성되었습니다. 해당 포스팅에서는 "스탠포드 대학교의 연구진은 S1 모델을 공개하며 상대적으로 적은 비용으로도 높은 추론 성능을 달성할 수 있는 Test-Time Scaling 기법을 제안했습니다." 라는 말을 하면서 Test-Time Scailing 기법을 활용하여 추론을 구현하는 방법을 소개하고 있습니다.

우선, 추론이 무엇인지 이야기를 해볼텐데요. 사실 이 부분이 가장 헷갈리는 부분 일 수 있습니다. 왜냐하면 추론은 두가지를 의미합니다. Inference 와 Reasoning 두가지를 모두 추론이라고 해석하는데요. 위의 포스팅에서는 그래서 Reasoning에 대해서 '사고' 라는 표현을 대신 사용했습니다. 저도 사실 사고 라는 표현이 조금 더 쉽게 다가가기 편하다는 생각을 가지고 있습니다. 그래서 저도 사고 라는 표현을 사용하도록 하겠습니다.

Inference (추론):
주어진 정보나 데이터를 바탕으로 결론을 도출하는 과정입니다. 즉, 관찰된 사실에서 논리적으로 결론을 이끌어내는 행위라고 할 수 있습니다.
Reasoning (사고 과정):
문제를 해결하거나 의사결정을 내리기 위해 여러 정보를 종합하고 논리적 관계를 파악하는 전반적인 사고 과정입니다. 이는 단순한 결론 도출을 넘어서, 다양한 가능성을 고려하고 그 과정에서 스스로 반성하며 수정하는 복잡한 인지 활동을 포함합니다.

간단히 말해, Inference는 결론을 도출하는 구체적인 행위이고, Reasoning은 그러한 결론에 도달하기 위한 전반적인 사고의 흐름이라 볼 수 있습니다.

LLM 모델을 만드는 과정은 크게 두 단계로 나뉩니다.

첫 번째 단계인 사전 학습(Pre-Training)에서는 방대한 양의 텍스트 데이터를 통해 모델이 언어의 패턴과 구조를 학습합니다. 이 과정에서 모델은 단순히 문장을 암기하는 것이 아니라, 확률적 언어 모델링을 통해 문맥을 이해하고 다음 단어를 예측하는 능력을 갖추게 됩니다.

두 번째 단계는 후처리학습(Post-Training)으로, SFT, RLHF, DPO와 같은 기법을 활용해 모델이 우리가 원하는 방식으로 출력을 하도록 미세 조정하는 과정입니다.

두 단계 모두 모델의 성능을 높이기 위해 더 많은 데이터와 연산 자원이 필요하기 때문에, 이러한 방식은 학습 시점에서 자원과 연산량을 증가시켜 성능을 확장하는 'Train-Time Scaling'이라고 부릅니다. 그러나 이 방식의 한계는 모델이 학습한 데이터에만 의존하여 답변을 생성하고, 출력이 한 번에 이루어지기 때문에 잘못된 답변이 발생해도 추론 과정 중에 수정할 수 없다는 점입니다. 잘못된 출력을 수정하려면 해당 오류를 개선하기 위한 추가 튜닝과 재학습이 필요하며, 이로 인해 지속적인 자원과 비용이 발생하는 문제가 있습니다.

최근 주목받고 있는 테스트 시점 확장(Test-Time Scaling)은 기존의 Train-Time Scaling처럼 학습 시에 추가 컴퓨팅 자원을 투입하는 대신, 추론 과정에서 컴퓨팅 자원을 활용해 모델의 성능을 개선하는 방법입니다. 예를 들어, 작년에 OpenAI에서 공개한 o1 모델은 사용자의 질문에 대해 즉각적인 답변을 제공하는 대신, 여러 차례의 추론과 자기 검증 과정을 거쳐 “깊이 생각”하도록 설계되었습니다.

이는 CoT(Chain of Thought) 프롬프팅과 비슷해 보이나, CoT는 중간 단계에 대한 검증이 어려운 반면, Test-Time Scaling은 모델이 생성한 추론 결과를 반복적으로 검증하며 오류를 수정해 나간다는 점에서 차별화됩니다. 다만, 이 방법의 한계는 모델이 충분한 내재적 지식과 문제 해결 능력을 갖추고 있어야 효과를 발휘할 수 있다는 점입니다. 예를 들어, 초등학교 수준의 수학 지식을 가진 모델이 대학 수준의 문제를 풀 경우, 아무리 오래 생각해도 정답에 도달하기 어렵다는 한계가 존재합니다.

그렇다면 CoT(Chain of Thought) 프롬프팅 부터 예시를 들면서 설명해도록 하겠습니다. CoT(Chain of Thought) 프롬프팅은 모델이 단순히 최종 답변만을 내놓는 것이 아니라, 문제 해결 과정을 여러 단계의 사고 과정을 통해 진행하도록 유도하는 기법입니다. 이를 통해 모델은 중간에 자신의 논리적 추론 과정을 드러내며, 복잡한 문제에 대해서도 보다 체계적이고 신뢰성 있는 결과를 도출할 수 있습니다.

아래는 예시로 CoT 프롬프팅을 구현한 간단한 코드입니다. (참고 포스팅 및 S1 논문에서는 Qwen2.5-7B를 활용하였으나, 저는 Exaone3.5-7.8B를 활용하였습니다)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Choose your prompt
prompt = """
문제:
한 직사각형의 너비를 x라고 하고, 길이가 너비의 3배라고 합시다. 이 직사각형의 둘레가 64cm일 때, 직사각형의 넓이를 구하세요.

요청:
문제를 풀 때, 먼저 너비와 길이의 관계를 파악하고, 둘레의 공식을 이용해 
𝑥
x의 값을 찾은 다음, 넓이를 계산하는 과정을 단계별로 상세하게 설명해 주세요.
"""# Korean example

messages = [
    {"role": "system", 
     "content": "You are EXAONE model from LG AI Research, a helpful assistant."},
    {"role": "user", "content": prompt}
]
input_ids = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
)

output = model.generate(
    input_ids.to("cuda"),
    eos_token_id=tokenizer.eos_token_id,
    max_new_tokens=128,
    do_sample=False,
)
print(tokenizer.decode(output[0]))

이런 식으로, 답변 과정에서 최종 답변을 도출하기 위한 과정을 설명합니다. 그러나, 이는 중간에 잘못된 과정으로 가더라도 CoT 프롬프팅으로는 수정을 할 수 없습니다. 이렇게 풀이 과정과 답이 틀리는 경우는 이러한 패턴에 대한 데이터셋을 다시 만들어 학습을 시켜야 제대로 된 사고 과정과 답변을 만들어 낼 수 있습니다. Reasoning을 구현하기 위해서, S1 에서 사용한 방법을 활용하고자 합니다.

우선 S1에서 사고 과정에 Budget Forcing이라는 개념을 활용하였습니다. Budget Forcing은 모델의 추론 과정을 디코딩 단계에서 제어하기 위한 기법으로, 두 가지 매커니즘을 통해 작동합니다.

첫째, 최대 사고 토큰 수 제한을 두어 모델이 정해진 토큰 수를 넘어서 추론할 경우 강제로 “end-of-thinking” 토큰과 “Final Answer:”를 삽입해 사고를 종료시키고 최종 답변을 생성하도록 합니다.

둘째, 최소 사고 토큰 수 보장을 통해 모델이 충분한 추론 없이 답변을 너무 빨리 종료하려 할 때 “end-of-thinking” 토큰 생성을 억제하고 “Wait”이라는 단어를 삽입해 모델이 더 깊게 생각하고 검증하는 과정을 반복하도록 유도합니다. 이러한 과정을 통해 모델은 반복적으로 사고를 진행하며, 최종 출력에서 정답에 도달할 수 있게 됩니다.

예를 들어, 논문에서는 라즈베리에 포함된 “r”의 개수를 묻는 질문에서 모델이 처음에 “2”라는 잘못된 답을 제시하자, 최소 사고 토큰 보장을 통해 출력에 “Wait”이라는 단어가 추가되어 다시 추론하도록 유도합니다. 그 결과, 모델은 “Wait” 이후의 추론 과정을 거쳐 최종 출력에 “Final Answer:”와 함께 올바른 정답을 도출하게 됩니다.

이제 아래와 같이 학습을 할 수 있습니다. (예시 코드 : https://github.com/superdom/blog/tree/main/030925-reasoning-model)

# import library
from datasets import load_dataset, Dataset
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import prepare_model_for_kbit_training, LoraConfig
from trl import SFTConfig, SFTTrainer
from typing import Dict

모델 로드 및 tokenizer 구분자 추가(Reasoning 과정)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda:0")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = "<|fim_pad|>"
tokenizer.pad_token_id = 151662
tokenizer.padding_side = 'left'

dataset 로드

dataset = load_dataset("werty1248/s1k-1.1-Ko-ReGenerated-Formatted", split='train')

# 데이터 셋을 원하는 형식으로 변경 

QUERY_TEMPLATE_NOANSWER = """{Question}""".strip()

# reasoning 과정 구분자 추가
def process_cot_example(example: Dict, tokenizer):
    thinking_trajectory = example["reasoning_ko"]
    question = example["question_ko"]
    answer = example["answer_ko"]

    prompt = QUERY_TEMPLATE_NOANSWER.format(Question=question)
    answer = "Answer: " + answer if "Answer:" not in answer else answer
    text = tokenizer.apply_chat_template([
        {"role": "user", "content": prompt},
        {
            "role": "assistant", 
            "content": "<|im_start|>think\n" + thinking_trajectory.strip() + 
                       "\n<|im_start|>answer\n" + answer.strip()
        }
    ], tokenize=False)
    return dict(text=text)
    
dataset = dataset.map(lambda example: process_cot_example(example, tokenizer))
dataset = dataset.remove_columns([col for col in dataset.column_names if col != "text"])

LoRA 설정 및 학습

lora_config = LoraConfig(
    task_type="CAUSAL_LM",
    r=32,            
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    target_modules= ['k_proj', 'o_proj','q_proj', 'v_proj', 'up_proj', 'down_proj', 'gate_proj'],
)

training_arguments = SFTConfig(
    output_dir="./output",
    optim="adamw_8bit",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=32,
    gradient_checkpointing=True,
    log_level="debug",
    save_strategy="epoch",
    logging_steps=2,
    learning_rate=5e-5,
    bf16 = True,
    num_train_epochs=6,
    weight_decay=1e-4,
    warmup_ratio=0.1,
    lr_scheduler_type="linear",
    dataset_text_field="text",
    max_seq_length=20000,
    report_to='none'
  )
  
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
peft_config=lora_config,
args=training_arguments,
)

model.config.use_cache = False
trainer.train()

추론 모델 테스트

위와 같이 추론 모델로 학습된 모델을 vllm을 활용하여 인퍼런스 해보고자 합니다.

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

# 반복 사고 횟수 및 최대 출력 길이 제한
MAX_TOKENS_THINKING = 20000
NUM_IGNORE = 2

# 첫 번째 추론 단계 설정
prompt = "<|im_start|>system\nYou are Exaone, You are a helpful assistant.<|im_end|>\n"
prompt += "<|im_start|>user\n" + prompts[0] + "<|im_end|>\n<|im_start|>assistant\n"

stop_token_ids = tok("<|im_start|><|im_end|>")["input_ids"]

sampling_params = SamplingParams(
    max_tokens=MAX_TOKENS_THINKING,
    min_tokens=0,
    stop_token_ids=stop_token_ids,
    skip_special_tokens=False,
    temperature=0,
)

prompt += "<|im_start|>think"

o = model.generate(
    prompt,
    sampling_params=sampling_params,
#    lora_request=s1_adapter
)

print(prompt + o[0].outputs[0].text)


# 반복 사고 단계
ignore_str = "잠깐,"
max_tokens_thinking_tmp = MAX_TOKENS_THINKING

for i in range(NUM_IGNORE):
    max_tokens_thinking_tmp -= len(o[0].outputs[0].token_ids)
    prompt += o[0].outputs[0].text + ignore_str
    sampling_params = SamplingParams(
        max_tokens=max_tokens_thinking_tmp,
        min_tokens=0,
        stop_token_ids=stop_token_ids,
        skip_special_tokens=False,
        temperature=0.0,
    )
    o = model.generate(
        prompt,
        sampling_params=sampling_params
    )
print(prompt + o[0].outputs[0].text)

이와 같이 추론을 진행할 수 있습니다. 단 위의 학습은 실습을 위해서 7.8B의 작은 모델과 적은 에폭 사이즈 그리고 LoRA를 활용하였기 때문에 만족하실 만한 성능을 얻기 힘드실 수도 있습니다. 이를 해결하기 위해서는 더 큰 모델과 더 많은 데이터셋 그리고 FFT가 필요할 것 같습니다.