<aside> 프로젝트 경험
</aside>
2026.02.02~2026.02.26 (팀 프로젝트)
개발 주요 사항
LG Aimers 8기 과제로써 Exaone 경량화를 진행
36 / 628 위 기록 / 간발의 차로 본선 진출은 실패
기술 스택
🛠️ 개발 환경 🤖 사용한 모델
Python 3 - Exaone
Pytorch
Distillation
살펴보기
해당 프로젝트는 Exaone의 경량화 성능을 고도화하는 프로젝트였습니다. 평가지표는 KMMLU-REDUX를 포함한 몇가지 비공개 데이터셋으로 진행되었으며, 최종적으로는 628팀중 36위를 기록하였습니다.
Quantization Aware Distillation을 통해 Exaone 1.2B BF16을 Exaone 1.2B FP8로 지식증류하여 성능 극대화를 시도하였으며, Loss에 단순 KL Divergence가 아닌 Entropy를 섞어 Teacher (bf16)모델이 확신있는 토큰만 더 강하게 distillation을 주었습니다.
Teacher 모델의 출력 분포가 0.95, 0.04, 0.01과 같이 특정 토큰에 확률이 집중된 경우에는 높은 확신도를 부여하고, 0.25, 0.23, 0.21과 같이 확률 분포가 평평한 경우에는 낮은 확신도를 부여하였습니다.
확신도(confidence)가 작으면 distillation loss에 곱해지는 가중치가 작아져서 전체 loss에 미치는 영향이 줄어들고, 확신도가 크면 distillation loss에 곱해지는 가중치가 커져서 전체 loss에 미치는 영향이 커집니다.
이를 통해 teacher가 확신하는 토큰에 대해서는 distillation의 영향을 크게 하고, 불확실한 예측에 대해서는 영향을 줄여 student 모델이 보다 의미 있는 정보에 집중하여 학습하도록 유도하였습니다.
LLMQuantizationModifier를 통해 원본 고정밀 BF16 모델을 Teacher로써, 양자화된 FP8 모델을 Student로써 두어 양자화로 인한 성능 저하 복구 방법인 Quantization Aware Distillation을 활용하였습니다. 대회 GPU에 제한이 없었다면 NVFP4 사용을 해보았을텐데, 주최의 GPU가 L4여서 양자화했을때 가장 성능이 괜찮았던 FP8을 두어 경량화 고도화를 시도하였습니다.
QAD와 관련한 정보는 https://arxiv.org/pdf/2601.20088를 참고해주시면 감사하겠습니다.
import os
import torch
import random
import numpy as np
import torch.nn.functional as F
from datasets import load_dataset, Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, set_seed
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
# ===============================
# Settings
# ===============================
MODEL_ID = "LGAI-EXAONE/EXAONE-4.0-1.2B"
DATASET_ID = "LGAI-EXAONE/MANTA-1M"
OUT_DIR = "/workspace/model"
MAX_LEN = 1024
TRAIN_SAMPLES = 512
CALIB_SAMPLES = 64
POOL_SIZE = 20000
MAX_STEPS = 150
LR = 1e-5
WARMUP_STEPS = 10
GRAD_ACCUM = 8
BATCH_SIZE = 1
T = 2.0
SEED = 42
# loss mix
ALPHA_CE = 0.5
ALPHA_KD = 0.5
# ===== CAD knobs =====
CAD_GAMMA = 2.0 # CAD: (1 - H_norm)^gamma, 1.0~2.0 추천
CAD_FLOOR = 0.05 # CAD: weight가 0으로 죽지 않게 최소값 (0~0.1)
CAD_EPS = 1e-9 # CAD: log 안정화
def make_reproducible(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
set_seed(seed)
# ===============================
# 샘플 필터
# ===============================
def is_valid_sample(convs):
if len(convs) < 2:
return False
if convs[-1]["role"] != "assistant":
return False
user = convs[-2]["content"]
asst = convs[-1]["content"]
if not isinstance(user, str) or not isinstance(asst, str):
return False
if len(user) < 20 or len(asst) < 20:
return False
return True
# ===============================
# Answer-only 전처리 (CE용)
# ===============================
def build_answer_only_features(tokenizer, conversations, max_len=512):
full_text = tokenizer.apply_chat_template(
conversations,
tokenize=False,
add_generation_prompt=False
)
prompt_text = tokenizer.apply_chat_template(
conversations[:-1],
tokenize=False,
add_generation_prompt=True
)
full_enc = tokenizer(
full_text,
truncation=True,
max_length=max_len,
padding="max_length"
)
prompt_enc = tokenizer(
prompt_text,
truncation=True,
max_length=max_len,
padding=False
)
input_ids = full_enc["input_ids"]
attention_mask = full_enc["attention_mask"]
prompt_len = min(len(prompt_enc["input_ids"]), len(input_ids))
labels = input_ids.copy()
# prompt 영역 마스킹 (CE는 answer-only)
for i in range(prompt_len):
labels[i] = -100
# pad 마스킹
for i, m in enumerate(attention_mask):
if m == 0:
labels[i] = -100
answer_tokens = sum(1 for x in labels if x != -100)
if answer_tokens < 8:
return None
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
}
# ===============================
# KD Trainer (CAD: entropy-weighted KL)
# ===============================
class KDTrainer(Trainer):
def __init__(self, teacher_model, *args, **kwargs):
super().__init__(*args, **kwargs)
self.teacher_model = teacher_model.eval()
for p in self.teacher_model.parameters():
p.requires_grad_(False)
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
out_s = model(**inputs)
logits_s = out_s.logits
loss_ce = out_s.loss
with torch.no_grad():
out_t = self.teacher_model(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"]
)
logits_t = out_t.logits
# ===== CAD KD =====
zs = logits_s[:, :-1, :] / T
zt = logits_t[:, :-1, :] / T
log_ps = F.log_softmax(zs, dim=-1)
p_t = F.softmax(zt, dim=-1)
kl_tok = F.kl_div(log_ps, p_t, reduction="none").sum(dim=-1)
mask = (inputs["attention_mask"][:, 1:] == 1)
# ---- entropy 계산 ----
entropy = -(p_t * torch.log(p_t + 1e-9)).sum(dim=-1)
V = p_t.size(-1)
max_entropy = torch.log(torch.tensor(float(V), device=entropy.device))
entropy_norm = entropy / (max_entropy + 1e-9)
conf = 1.0 - entropy_norm
weight = conf.pow(CAD_GAMMA).clamp_min(CAD_FLOOR)
weighted_kl = kl_tok * weight
denom = mask.sum().clamp_min(1)
loss_kd = (weighted_kl.masked_select(mask).sum() / denom) * (T ** 2)
total_loss = ALPHA_CE * loss_ce + ALPHA_KD * loss_kd
return (total_loss, out_s) if return_outputs else total_loss
# ===============================
# Main
# ===============================
def main():
make_reproducible(SEED)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
teacher = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True
)
student = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True
)
print("Loading pool...")
pool = load_dataset(DATASET_ID, split=f"train[:{POOL_SIZE}]")
filtered = [x for x in pool if is_valid_sample(x["conversations"])]
random.shuffle(filtered)
filtered = filtered[:TRAIN_SAMPLES]
train_ds = Dataset.from_list(filtered)
# 전처리 (CE는 answer-only labels)
features = []
for ex in train_ds:
feat = build_answer_only_features(tokenizer, ex["conversations"], MAX_LEN)
if feat is not None:
features.append(feat)
train_tok = Dataset.from_list(features)
# Calib (학습 포맷과 일관)
calib_ds = train_ds.shuffle(seed=SEED).select(range(min(len(train_ds), CALIB_SAMPLES)))
calib_ds = calib_ds.map(lambda x: {
"text": tokenizer.apply_chat_template(
x["conversations"],
tokenize=False,
add_generation_prompt=True
)
})
# FP8 oneshot
recipe = [QuantizationModifier(targets="Linear", scheme="FP8", ignore=["lm_head"])]
oneshot(
model=student,
dataset=calib_ds,
recipe=recipe,
max_seq_length=MAX_LEN,
num_calibration_samples=CALIB_SAMPLES,
)
training_args = TrainingArguments(
output_dir=OUT_DIR,
per_device_train_batch_size=BATCH_SIZE,
gradient_accumulation_steps=GRAD_ACCUM,
max_steps=MAX_STEPS,
learning_rate=LR,
warmup_steps=WARMUP_STEPS,
lr_scheduler_type="cosine",
bf16=True,
gradient_checkpointing=True,
logging_steps=10,
logging_strategy="steps",
save_strategy="no",
optim="adamw_torch",
weight_decay=0.0,
adam_beta2=0.98,
adam_epsilon=1e-6,
max_grad_norm=0.3,
)
student.config.use_cache = False
student.gradient_checkpointing_enable()
trainer = KDTrainer(
model=student,
teacher_model=teacher,
args=training_args,
train_dataset=train_tok,
)
trainer.train()
student.gradient_checkpointing_disable()
student.config.use_cache = True
student.eval()
os.makedirs(OUT_DIR, exist_ok=True)
student.save_pretrained(OUT_DIR, save_compressed=True)
tokenizer.save_pretrained(OUT_DIR)
if __name__ == "__main__":
main()