[Axolotl] 1. Axolotl로 언어 모델 파인튜닝하기

👋 시작하며

안녕하세요!  

 

Axolotl (아홀로틀)은 위의 그림에 있는 굉장히 귀여운 수중 생물인데요! 이 axolotl과 같은 이름을 가지는 AI 모델 학습 도구가 있습니다. 최근 이 Axolotl을 사용하여 언어 모델 학습과 최적화를 수행하고 있는데 개인적으로 굉장히 좋은 도구라고 생각해서 이를 사용하는 방법을 정리해서 블로그 글을 작성해보려고 합니다. 

 

Axolotl은 복잡한 LLM 미세조정(fine-tuning)을 훨씬 쉽고 유연하게 할 수 있도록 도와주는 오픈소스 프레임워크입니다. Hugging Face Transformers, LoRA, Quantization 등의 기술을 한데 모아 다양한 하드웨어 환경에서 자유롭게 설정할 수 있다는 점이 큰 장점입니다! 💪

 

앞으로 axolotl을 사용한 단순한 파인튜닝부터 다양한 최신 기술들까지 폭넓은 내용을 천천히 다뤄볼까합니다. 수정이 필요한 부분이나 의견이 있으시면 언제든 편하게 남겨주세요! 💬

 

우선 Axolotl 관련 링크들은 다음과 같습니다. 

• Axolotl 깃허브: https://github.com/axolotl-ai-cloud/axolotl
• Axolotl 문서: https://docs.axolotl.ai/

 

📥 Axolotl의 설치

우선 Axolotl을 설치해보겠습니다. 

저는 Ubuntu 22.04, Python 3.11 환경에서 다음과 같은 명령어를 통해 설치를 진행했습니다. 

pip install -U packaging setuptools wheel ninja
pip install --no-build-isolation axolotl[flash-attn,deepspeed]

 

그리고 효율적인 학습을 위해 아래와 같은 명령어를 통해 cut cross entropy까지 설치하겠습니다

pip install "cut-cross-entropy @ git+https://github.com/apple/ml-cross-entropy.git"

 

Cut cross entropy가 뭔지 간단하게만 살펴보겠습니다!🧐

Cut cross entropy는 Apple에서 제안한 효율적인 학습 기법으로 loss 계산 시에 전체 단어들에 대한 softmax를 계산하는 대신 상위 확률 토큰 일부와 정답 토큰만 사용하여 연산량을 줄여줍니다!🤗

 

이렇게 간단하게 설치가 완료되었습니다!! 🪄

 

Axolotl 학습을 위해서는 아래의 딱 2가지만 있으면 됩니다!

• 데이터 셋
• 학습 설정 파일

이제 위의 두가지를 어떻게 생성하고 설정하는지 차근차근 살펴보겠습니다!😆  

 

🗃️ 데이터셋 준비

데이터셋 살펴보기

다음으로 학습에 사용할 데이터 셋을 준비해보겠습니다. 

데이터를 찾다보니 아래와 같이 옛날 한글로 된 홍길동전을 현대의 한글로 변환한 데이터셋이 있었습니다. 

• 데이터셋 링크: https://huggingface.co/datasets/bebechien/HongGildongJeon 

bebechien/HongGildongJeon · Datasets at Hugging Face

 

이 데이터에서 사용한 옛날 한글과 현대 한글의 예시는 다음과 같습니다. (사실 현대 한글도 옛날 말이긴 합니다😅) 

# 옛날 한글

이날길동삼쳔젹군을거ᄂᆞ려망망ᄃᆡᄒᆡ로ᄯᅥᄀᆞ더니셩도라ᄒᆞᄂᆞᆫ도즁의이르러창고을지으며궁실을지여안돈ᄒᆞ고군ᄉᆞ로ᄒᆞ여곰농업을심쓰고각국의왕ᄂᆡᄒᆞ야믈화을통ᄒᆞ며무예을슝상ᄒᆞ야병법을ᄀᆞ르치니삼연지ᄂᆡ예군긔군량이뫼갓고군ᄉᆞ강ᄒᆞ야당젹ᄒᆞ리업슬네라

# 현대 한글
이날 길동 삼천 적군을 거느려 망망대해로 떠나더니, 성도라 하는 도중에 이르러 창고를 지으며, 궁실을 지어 안돈하고, 군사로 하여금 농업을 힘쓰고, 각국에 왕래하여 물화를 통하며, 무예를 숭상하여 병법을 가르치니, 삼년지내에 군기 군량이 산같고, 군사 강하여 당적할 이 없더라.

 

간단하게 파인튜닝하는 데이터로 사용하기 좋다고 생각되어 한번 이를 사용해서 데이터를 생성해보겠습니다.  

 

데이터셋 변환 예시

이제 다음과 같은 과정을 통해 데이터를 학습을 위한 형식으로 변환해보겠습니다! 

Axolotl을 사용하여 언어 모델을 학습하기 위한 데이터 포맷은 아래 링크 내용을 참고해주세요! 

https://docs.axolotl.ai/docs/dataset-formats/

Dataset Formats – Axolotl

 

저희는 instruction following (명령을 수행하는) 학습을 위한 형식 중 하나인 Alpaca 형식으로 데이터 셋을 변경해보겠습니다!

데이터의 형식은 아래와 같이 "instruction", "input", "output"으로 구성됩니다. 

{"instruction": "...", "input": "...", "output": "..."}

 

각각에 내용을 입력하면 실제 모델 학습을 위한 프롬프트는 아래와 같이 설정됩니다. 

Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Input:
{input}

### Response:
{output}

 

저는 이 데이터에 아래와 같이 값들을 대입하여 사용할 예정입니다. 

Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
입력의 옛 한글 내용을 현대의 한글로 변환해줘

### Input:
이날길동삼쳔젹군을거ᄂᆞ려망망ᄃᆡᄒᆡ로ᄯᅥᄀᆞ더니셩도라ᄒᆞᄂᆞᆫ도즁의이르러창고을지으며궁실을지여안돈ᄒᆞ고군ᄉᆞ로ᄒᆞ여곰농업을심쓰고각국의왕ᄂᆡᄒᆞ야믈화을통ᄒᆞ며무예을슝상ᄒᆞ야병법을ᄀᆞ르치니삼연지ᄂᆡ예군긔군량이뫼갓고군ᄉᆞ강ᄒᆞ야당젹ᄒᆞ리업슬네라

### Response:
이날 길동 삼천 적군을 거느려 망망대해로 떠나더니, 성도라 하는 도중에 이르러 창고를 지으며, 궁실을 지어 안돈하고, 군사로 하여금 농업을 힘쓰고, 각국에 왕래하여 물화를 통하며, 무예를 숭상하여 병법을 가르치니, 삼년지내에 군기 군량이 산같고, 군사 강하여 당적할 이 없더라.

 

즉, Instruction에는 명령을, Input에는 입력을, Response에는 LLM의 응답을 입력하면 됩니다! 

이제 실제 코드를 통해 변환을 수행해보겠습니다!!

 

데이터셋 변환 코드

Huggingface의 HongGildongJeon 데이터 셋을 Alpaca 형식으로 변환하기 위한 코드는 다음과 같습니다! 

from datasets import load_dataset
import json
import random

# ✅ 데이터셋 불러오기
dataset = load_dataset("bebechien/HongGildongJeon")
data = list(dataset["train"])  # list로 변환

# ✅ 섞기
random.shuffle(data)
n_total = len(data)
n_test = max(1, int(n_total * 0.05))  # 최소 1개는 test로

test_data = data[:n_test]
train_data = data[n_test:]

# ✅ JSONL 저장 함수
def save_jsonl(filename, data):
	with open(filename, "w", encoding="utf-8") as f:
		for row in data:
			original_text = row.get("original", "").strip()
			modern_translation_text = row.get("modern translation", "").strip()
			
			json_obj = {
				"instruction": "입력의 옛 한글 내용을 현대의 한글로 변환해줘",
				"input": original_text,
				"output": modern_translation_text
			}
			
			f.write(json.dumps(json_obj, ensure_ascii=False) + "\n")

# ✅ 파일 저장
save_jsonl("./honggildong_train.jsonl", train_data)
save_jsonl("./honggildong_test.jsonl", test_data)

print(f"✅ train: {len(train_data)}개, test: {len(test_data)}개 저장 완료!")

 

코드에 대한 간략한 설명은 다음과 같습니다. 

• 먼저 load_dataset을 통해 bebechien/HongGildongJeon 데이터셋을 불러옵니다. 
• 전체 데이터를 한번 섞고 5%의 데이터를 테스트 세트로 설정합니다. 
• save_jsonl 함수는 데이터를 instruction, input, output을 key로 하는 jsonl 형식으로 저장합니다. 
  • instruction: 항상 "입력의 옛 한글 내용을 현대의 한글로 변환해줘" 내용을 포함합니다. 
  • input: 데이터 중 옛 한글에 해당하는 original_text를 포함합니다. 
  • output: 데이터 중 현대어에 해당하는 modern_translation_text를 포함합니다. 
• 마지막으로 이 save_jsonl 함수에 train_data, test_data를 입력하여 각각을 jsonl 파일로 저장합니다. 

위 코드를 수행하면 아래와 같은 결과를 확인하실 수 있습니다! 

✅ train: 425개, test: 22개 저장 완료!

 

총 425개의 학습 데이터를 포함한 jsonl 파일 (honggildong_train.jsonl)과 22개의 테스트 데이터를 포함한 jsonl 파일 (honggildong_test.jsonl)이 생성되었습니다! 

일반적인 언어 모델 학습을 위해서는 데이터가 너무 적지만..!😅

그래도 간단한 학습을 위해 작은 데이터셋을 만들어보았습니다!

 

⚙️학습 설정 파일

Axolotl로 학습을 하기 위한 모든 설정은 yaml 파일 안에 작성하면 됩니다. 

제가 이번에 학습을 위해 사용한 설정은 다음과 같습니다. 

# === 모델 설정 ===
base_model: Qwen/Qwen3-0.6B
base_model_config: Qwen/Qwen3-0.6B
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer
trust_remote_code: true
adapter: null

# === 데이터셋 설정 ===
datasets:
	- path: ./honggildong_train.jsonl
	type: alpaca

# 전체 길이 설정 (Qwen3-0.6B는 최대 16384 토큰)
sequence_len: 8192
preprocessing_num_workers: 4
shuffle_merged_datasets: true

# === 출력 디렉토리 설정 ===
output_dir: ./models/Qwen3-0.6B-honggildong/
overwrite_output_dir: true
save_only_model: true
save_safetensors: true

# === 로깅 및 저장 ===
logging_steps: 10
save_strategy: epoch
save_total_limit: 1

# === 훈련 하이퍼파라미터 ===
gradient_accumulation_steps: 1
gradient_checkpointing: true
micro_batch_size: 16
num_epochs: 5
learning_rate: 5.0e-5
lr_scheduler: cosine
warmup_ratio: 0.1
flash_attention: true
bf16: true

# === 최적화 (liger kernel, cut cross entropy) ===
plugins:
  - axolotl.integrations.liger.LigerPlugin
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_layer_norm: true
liger_fused_linear_cross_entropy: false
  
# === 평가 설정 ===
val_set_size: 0.01
eval_strategy: steps
eval_steps: 20
per_device_eval_batch_size: 1

 

설정에 대한 전체 내용들은 아래 링크 내용을 참고해주세요!!

https://docs.axolotl.ai/docs/config-reference.html

Config Reference – Axolotl

 

제가 설정한 내용들에 대해서 하나씩 살펴보겠습니다. 

 

모델 설정

 

• base_model / base_model_config: Hugging Face에서 불러올 사전학습 모델과 설정
• model_type / tokenizer_type: 모델과 토크나이저를 위해 사용할 HuggingFace 클래스
• trust_remote_code: remote code (ex. Qwen의 커스텀 모델) 신뢰 여부
• adapter: null: LoRA 등의 adapter 사용 안 함 = full fine-tuning
  
  

 

데이터 설정

 

• datasets: 학습 데이터 경로 및 형식 설정
  • path: 학습 데이터 경로
  • type: 학습 데이터 형식 (alpaca)
• sequence_len: 입력 시퀀스 길이 (Qwen3는 최대 16384, 여기선 8192).
• preprocessing_num_workers: 전처리 시 병렬 작업 수.
• shuffle_merged_datasets: 학습 데이터 전체를 랜덤하게 섞을지 여부 결정
  
  

출력 디렉토리

 

• output_dir: 학습 모델 저장 경로.
• overwrite_output_dir: 기존 폴더가 존재하는 경우 덮어쓰기 허용
• save_only_model: 옵티마이저 상태 등은 저장 안하고 모델만 저장
• save_safetensors: .safetensors 포맷으로 저장
  
  

로깅 및 체크포인트

• logging_steps: 로그 출력 간격(step 단위)
• save_strategy: 모델 저장 단위, 현재는 매 epoch마다 저장 
• save_total_limit: 저장할 checkpoint 개수 제한
  
  

훈련 하이퍼파라미터

 

• gradient_accumulation_steps: 작은 배치 스텝을 누적해 더 큰 배치처럼 학습
• gradient_checkpointing: 메모리 절약을 위한 중간 gradient 저장 생략
• micro_batch_size: GPU당 배치 크기
• num_epochs: 총 학습 epoch 수
• learning_rate: 학습률
• lr_scheduler:  학습률 스케줄러 (cosine decay)
• warmup_ratio: 학습률 스케줄러의 warmup 단계 비율 (전체 step의 10%)
• flash_attention: 빠른 attention 커널 사용
• bf16: BF16 정밀도로 학습
  
  

 

 

 

최적화 (Liger, Cut Cross Entropy)

 

• plugins: 훈련 시 사용할 Axolotl의 최적화 플러그인.
  • LigerPlugin: LayerNorm, RoPE, MLP를 fused kernel로 최적화
  • CutCrossEntropyPlugin: Softmax 전체 대신 top-k + 정답만 사용하는 loss로 메모리 및 연산량 감소
• liger_~~~: 각 모듈에 대해 Liger-style 최적화를 적용할지 여부 결정
  
  

평가 설정

 

• val_set_size: 전체 학습 데이터 중 1%를 validation에 사용.
• eval_strategy: 평가를 수행할 주기의 단위, 현재는 스텝
• eval_steps: 평가를 수행할 스텝 수, 현재는 20 스텝마다 평가 수행
• per_device_eval_batch_size: 평가시 GPU당 배치 크기

 

 

위 설정들의 내용을 요약하자면 다음과 같습니다!

• Qwen3-0.6B 모델을 8192 context로 full fine-tuning
• alpaca 형식을 가지는 ./honggildong_train.jsonl 데이터로 학습 수행
• Cut Cross Entropy와 Liger를 통한 속도 및 메모리 최적화
• FlashAttention, BF16, Gradient Checkpointing 등의 학습 효율 기법도 추가

이렇게 설정이 완료된 파일을 train_config.yml 이라는 이름으로 저장합니다!

 

📝학습 수행

이제 데이터셋과 학습을 위한 설정이 모두 끝났습니다!! 

Axolotl로 학습을 수행하는 것은 정말 간단합니다!! 

아래와 같이 명령을 입력하면 끝입니다!

axolotl train train_config.yml

 

학습이 잘 되었는지 지표상으로 확인하기 위해 학습 loss와 평가 loss를 살펴보겠습니다. 

wandb를 통해 확인해보았을 때 아래와 같이 학습 loss와 평가 loss 모두 잘 줄어든 것을 확인할 수 있습니다 

 

참고로 axolotl 학습 과정을 wandb로 로깅하기 위해서는 학습 설정에 아래와 같이 추가해주시면 됩니다. 

wandb_entity: entity 이름
wandb_project: project 이름
wandb_name: wandb run 이름

 

🧪추론 수행

이제 학습이 얼마나 잘 되었는지 결과를 한번 확인해보겠습니다!

모델의 성능은 테스트 셋에서 옛 한글을 입력하고 모델이 이를 현대어로 바꿨을 때 실제 정답 현대어와 비교하여 그 정확도를 평가하는 것으로 하겠습니다! 

 

우선 학습 하기 전에 얼마나 못하는지 확인하기 위해 학습한 모델이 아니라 사전학습 된 Qwen3-0.6b 모델을 통해 정확도를 구해보겠습니다. 

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json
from tqdm import tqdm

model_path = "Qwen/Qwen3-0.6B"

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True,
    torch_dtype="auto",
    device_map={"": "cuda:0"},
)

model.eval()

# json 파일을 불러오는 함수
def load_jsonl(filename):
    with open(filename, "r", encoding="utf-8") as f:
        for line in f:
            yield json.loads(line)

# 예측 단어 중 정답과 겹치는 단어의 비율 계산 함수
def get_word_overlap(pred, ref):
    ref_words = set(ref.strip().split())
    pred_words = set(pred.strip().split())
    common = ref_words & pred_words
    return len(common) / len(ref_words)

# 모델의 출력을 얻는 함수
def get_model_output(input_text, model, tokenizer):
    # 프롬프트 템플릿 적용
    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
입력의 옛 한글 내용을 현대의 한글로 변환해줘

### Input:
{input_text}

### Response:
"""
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.01,
            top_p=0.95,
            do_sample=True,
        )
    output_ids = outputs[0][inputs["input_ids"].shape[1]:]
    decoded = tokenizer.decode(output_ids, skip_special_tokens=False).strip()
    decoded = decoded.split("<|im_end|>")[0].strip()
    return decoded

# -- 평가 루프 --
results = []
jsonl_path = "honggildong_test.jsonl"

for item in tqdm(load_jsonl(jsonl_path)):
    input_text = item["input"]
    target_text = item["output"]
    decoded = get_model_output(input_text, model, tokenizer)
    acc = get_word_overlap(decoded, target_text)
    results.append(acc)

accuracy = sum(results) / len(results)
print(f"전체 단어 포함 기준 Accuracy: {accuracy:.3f}")

 

평가 수행 과정은 아래와 같습니다. 

• 평가를 위한 jsonl 데이터 (honggildong_test.jsonl) 불러오기
• 입력과 정답 데이터를 불러온 다음 입력 데이터를 통한 모델의 추론 결과를 얻고 이를 실제 정답과 비교한 정확도 계산
• 전체 정확도의 평균을 계산 (accuracy)

 

제가 이 코드를 실행했을 때 파인튜닝 하지 않은 Qwen 모델의 정확도는 다음과 같았습니다!

• 전체 단어 포함 기준 Accuracy: 0.039

 

다음으로는 학습을 수행한 모델의 성능을 살펴보기 위해 코드에서 model을 다음과 같이 학습된 모델이 저장된 경로로 변경합니다. 

model_path = "./models/Qwen3-0.6B-honggildong"

 

그리고 위 코드를 실행했을 때 결과가 다음과 같았습니다!

• 전체 단어 포함 기준 Accuracy: 0.639

 

무려 정확도가 3.9%에서 63.9%로 대폭 상승했습니다!! 🤩

더 많은 데이터가 있었다면 성능이 더 늘어났겠지만..!

일단 이 정도의 짧은 학습으로도 성능이 크게 증가된다는 것을 확인했습니다!

 

실제 출력 결과도 확인해보기 위해 평가 코드의 get_model_output에 테스트 데이터 중 하나를 아래와 같이 입력해보겠습니다. 

test_input = "이젹의길동의슈단이신츌귀몰ᄒᆞ야팔도의횡ᄒᆡᆼᄒᆞ되능히알ᄌᆡ업ᄂᆞᆫ지라"
print(get_model_output(test_input, model, tokenizer))

 

 

학습 이전의 Qwen3-0.6b와 학습한 Qwen3-0.6b의 결과와 정답을 비교한 결과가 다음과 같습니다!

• 실제 정답: 이때에 길동의 수단이 신출귀몰하여 팔도에 횡행하되 능히 알 자가 없는지라.
• 학습 전 Qwen3-0.6b: 이주시의 길동의슈단이신출귀몰을 찾아팔도의 횡원을 이용하여 알을 전하는지라
• 학습 후 Qwen3-0.6b: 이때에 길동의 수단이 신출귀몰하여 팔도에 횡행하되 능히 알자 없는지라.

 

👋마무리

여기까지의 내용으로 간단한 데이터셋과 학습 설정을 통해 Axolotl로 언어 모델을 SFT (Supervised Finetuning)하는 내용을 살펴보았습니다!

다음 내용은 RLHF 학습을 진행하는 내용을 다뤄볼까합니다!!🤖