sLLM (Korean)

한국어 소형 언어모델. 24층 · d_model 768 · GQA(12 q / 4 kv) · sliding window(512) + 6층마다 global 혼합 어텐션 · MTP(multi-token prediction) 헤드를 얹은 디코더입니다. SentencePiece 32,768 vocab.

파라미터 178.6M (임베딩 25.2M 제외 시 153.4M). MTP 헤드는 학습에서 보조 손실로 쓰이고, 추론에서는 self-speculative decoding의 draft 생성기로 재활용됩니다.

층 / d_model / FFN	24 / 768 / 2048
어텐션	GQA 12 q · 4 kv head, head_dim 64
컨텍스트	2,048 (sliding window 512, 6층마다 global)
RoPE θ	local 10,000 / global 1,000,000
MTP	16 offset, 손실 가중치 0.2
vocab	32,768 (pad 0, bos 1, eos 2)

학습

코퍼스	한국어 76.86M 토큰, 필터 후 ≈74.86M
배치	8 × grad_accum 4 × 2,048 = 65,536 토큰/step
1 epoch	≈1,142 step
목표	4 epoch ≈ 4,570 step
best 체크포인트	step 4,500 (≈3.9 epoch, 누적 ≈295M 토큰)
처리량	22.9k tok/s, 피크 메모리 70.5 GiB

step 4,500 기준 손실입니다. 총 손실은 main + 0.2 × mtp입니다.

	main	mtp	총
train	1.9419	5.8022	3.1023
val	3.0717	6.4412	—

val perplexity 21.6 (= exp(3.0717)). safetensors 헤더의 best_val 메타데이터가 이 val main 값입니다.

파일

파일	내용
model.safetensors	학습된 가중치 (fp32, 714 MB)
config.json	ModelConfig.to_dict() 결과
model.py	아키텍처 정의 (로드에 필요)
spm.model	SentencePiece 토크나이저
onnx/model.onnx	프리필/디코드 공용 ONNX 트렁크 (708 MB)
onnx/mtp_head.onnx	MTP 헤드 (speculative decoding용, 110 MB)
onnx/config.json	위와 동일한 config

로드

KoreanSLLM은 transformers의 PreTrainedModel을 상속하지 않는 순수 nn.Module이므로 AutoModel.from_pretrained로는 로드되지 않습니다. 파일을 직접 받아 씁니다.

import json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

repo = "PostAlign/sLLM"

# model.py 를 받아서 import
import importlib.util, sys
spec = importlib.util.spec_from_file_location("kslm", hf_hub_download(repo, "model.py"))
kslm = importlib.util.module_from_spec(spec); sys.modules["kslm"] = kslm
spec.loader.exec_module(kslm)

cfg = kslm.ModelConfig(**json.load(open(hf_hub_download(repo, "config.json"))))
model = kslm.KoreanSLLM(cfg)
model.load_state_dict(load_file(hf_hub_download(repo, "model.safetensors")))
model.eval()

import sentencepiece as spm
tok = spm.SentencePieceProcessor(model_file=hf_hub_download(repo, "spm.model"))

ONNX

브라우저(WebGPU)·모바일(NNAPI/CoreML) 가속 추론용 그래프입니다. 프리필과 디코드가 같은 그래프이며, 프리필은 past_len=0, 디코드는 seq=1로 호출합니다.

입력  input_ids (batch, seq) int64
      past_k    (24, batch, 4, past_len, 64)
      past_v    (24, batch, 4, past_len, 64)
출력  logits    (batch, seq, 32768)
      hidden    (batch, seq, 768)        # mtp_head.onnx 의 입력
      present_k (24, batch, 4, past_len + seq, 64)
      present_v (24, batch, 4, past_len + seq, 64)

mtp_head.onnx는 (hidden, offset_idx)를 받아 (logits, mtp_hidden)을 냅니다. speculative decoding의 수용/거부 판정 루프는 그래프가 아니라 호스트 코드에 둡니다.

export 직후 onnxruntime으로 PyTorch 원본과 대조했습니다.

프리필(T=24, P=0) 로짓 최대오차: 1.621e-05
greedy 12스텝 토큰열 일치: True
MTP offset 0..15 로짓 최대오차: 5.722e-06

한계

• max_seq_len=2048이 생성 시 하드 컷입니다.
• ONNX 그래프는 KV 캐시를 매 스텝 전체 복사합니다(in-place 상태가 없음). 긴 컨텍스트에서 병목입니다.
• ONNX logits가 전 위치에 대해 나오므로 긴 프리필에서 출력이 큽니다(2,048 컨텍스트 fp32 기준 268 MB).
• train/val 간극이 큽니다. step 4,500에서 train main 1.9419 (ppl 6.97) 대 val main 3.0717 (ppl 21.6)로, 같은 코퍼스를 약 4회 반복한 데 따른 과적합이 뚜렷합니다. 데이터를 늘리지 않고 step만 더 밀어도 val이 개선되기 어렵습니다.
• 데이터가 모델 크기에 비해 적습니다. 고유 토큰 74.86M은 178M 파라미터 기준 흔히 쓰는 경험칙(파라미터당 ~20토큰, 약 3.6B)에 한참 못 미칩니다. 성능 상한이 데이터에 걸려 있습니다.
• 소규모·짧은 학습 모델입니다. 사실성·안전성 보장이 없으며, 벤치마크 평가를 수행하지 않았습니다.