{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# 지식 증류 실습: 교사 응답으로 학생 모델 만들기\n\n이 노트북은 [사후 훈련 — 지식 증류](../book/post_training.md) 절에서 설명한 **시퀀스 수준(데이터) 증류**를 직접 재현합니다.\n\n거대한 **교사 모델**이 생성한 추론 응답을 데이터셋으로 삼아, 훨씬 작은 **학생 모델**을 일반 SFT로 학습시킵니다. `DeepSeek-R1-Distill` 계열이 만들어진 방식의 축소판으로, **로짓 한 줄 없이 교사의 텍스트만으로 추론 능력이 옮겨 가는지** 확인하는 것이 목표입니다.\n\n| 역할 | 모델 | 도구 |\n| :--- | :--- | :--- |\n| 교사(Teacher) | `qwen3:8b` | [Ollama](../book/ollama.md) |\n| 학생(Student) | `Qwen3-0.6B` | [Unsloth LoRA](../book/unsloth-qwen3-finetuning.md) |\n\n교사 생성과 학생 학습이 분리되어 두 모델을 동시에 적재하지 않으므로, 24GB급 단일 GPU(RTX 3090 등)에서 무리 없이 돌아갑니다.\n\n```mermaid\nflowchart LR\n    Q[\"프롬프트 집합\\n(GSM8K)\"] --> T[\"교사 Qwen3-8B\\n(Ollama)\"]\n    T -->|\"<think> 추론 + 정답\"| D[\"증류 데이터셋\"]\n    D --> S[\"학생 Qwen3-0.6B\\n(Unsloth SFT)\"]\n```"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 0. 환경 준비\n\n학생 학습에는 [미세조정](../book/unsloth-qwen3-finetuning.md) 장과 동일한 스택(`unsloth`, `trl`)을 씁니다. 교사 생성은 공식 **Ollama 도커 컨테이너**(`ollama/ollama`)의 REST API를 `requests`로 호출하므로, 별도의 파이썬 클라이언트 패키지가 필요 없습니다([Ollama API](../book/ollama-api.md) 장과 동일한 방식).\n\n```bash\n# Ollama 도커 컨테이너에 교사 모델 내려받기 (컨테이너 이름이 'ollama'라고 가정)\ndocker exec ollama ollama pull qwen3:8b\n```\n\n엔드포인트는 환경 변수 `OLLAMA_HOST`로 바꿀 수 있고, 기본값은 `http://localhost:11434`입니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# !pip install unsloth trl datasets tqdm"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1. 교사에게 물어볼 프롬프트 준비\n\n증류하려는 능력을 대표하는 질문 집합이 필요합니다. 여기서는 초등 수학 추론 벤치마크인 **GSM8K**의 일부를 씁니다. 정답 레이블은 쓰지 않고 **질문만** 가져옵니다 — 정답은 교사가 만들어 줄 것이기 때문입니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from datasets import load_dataset\n\n# 데모용으로 200문항만 사용합니다. 규모를 키울수록 증류 효과가 커집니다.\nseed = load_dataset('openai/gsm8k', 'main', split='train[:200]')\nprompts = [row['question'] for row in seed]\n\nprint(len(prompts), '개 프롬프트')\nprint(prompts[0])"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2. 교사 응답 생성 (Ollama)\n\n교사는 **추론만** 하므로 학생 학습 환경과 완전히 분리됩니다(별도 프로세스). Qwen3의 사고 과정(`thinking`)까지 받아, 학생이 학습할 Qwen3 대화 템플릿과 같은 `<think>…</think>` 형식으로 합칩니다. 추론 트레이스 자체가 증류의 핵심 신호입니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import os\nimport requests\nfrom tqdm.auto import tqdm\n\n# 공식 Ollama 도커 컨테이너의 REST API 엔드포인트\nOLLAMA_BASE_URL = os.environ.get('OLLAMA_HOST', 'http://localhost:11434')\nTEACHER = 'qwen3:8b'\n\ndef ask_teacher(question):\n    resp = requests.post(\n        f'{OLLAMA_BASE_URL}/api/chat',\n        json={\n            'model': TEACHER,\n            'messages': [{'role': 'user', 'content': question}],\n            'think': True,    # Qwen3의 사고 과정을 message.thinking 으로 분리해 받습니다.\n            'stream': False,\n        },\n        timeout=600,\n    )\n    resp.raise_for_status()\n    msg = resp.json()['message']\n    thinking = msg.get('thinking') or ''\n    answer = msg['content']\n    # Qwen3 대화 템플릿과 동일한 <think> 형식으로 결합합니다.\n    return f'<think>\\n{thinking}\\n</think>\\n\\n{answer}' if thinking else answer\n\nrecords = [{'question': q, 'answer': ask_teacher(q)} for q in tqdm(prompts)]\n\nprint(records[0]['answer'][:500])"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3. 증류 데이터셋 구성\n\n교사 응답을 `[지시 - 응답]` 대화 형식으로 표준화합니다. 이 시점부터는 일반 SFT 데이터셋과 완전히 동일합니다 — 출처가 사람이 아니라 교사 모델일 뿐입니다. 재실행을 아끼기 위해 디스크에 저장해 둡니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from datasets import Dataset\n\ndef to_conversation(rec):\n    return {'conversations': [\n        {'role': 'user', 'content': rec['question']},\n        {'role': 'assistant', 'content': rec['answer']},\n    ]}\n\ndistill_dataset = Dataset.from_list([to_conversation(r) for r in records])\ndistill_dataset.to_json('distill_dataset.jsonl')  # 재사용을 위해 저장\ndistill_dataset"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4. 학생 모델 적재 (Unsloth · Qwen3-0.6B)\n\n교사보다 10배 이상 작은 학생을 4비트로 적재하고 LoRA 어댑터를 붙입니다. 교사 생성이 끝났으므로 더 이상 교사를 메모리에 둘 필요가 없습니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from unsloth import FastLanguageModel\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    'unsloth/Qwen3-0.6B-unsloth-bnb-4bit',\n    max_seq_length=2048, dtype=None, load_in_4bit=True)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=16,\n    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n                    'gate_proj', 'up_proj', 'down_proj'],\n    lora_alpha=16,\n    lora_dropout=0,\n    bias='none',\n    use_gradient_checkpointing='unsloth',\n    random_state=2025,\n)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 5. 대화 템플릿 적용\n\n학생이 사전 학습한 것과 **동일한** Qwen3 템플릿(`qwen-3`)으로 데이터를 렌더링해 단일 `text` 컬럼으로 만듭니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from unsloth.chat_templates import get_chat_template\n\ntokenizer = get_chat_template(tokenizer, chat_template='qwen-3')\n\ndef formatting_prompts_func(examples):\n    texts = [\n        tokenizer.apply_chat_template(\n            conversation, tokenize=False, add_generation_prompt=False)\n        for conversation in examples['conversations']\n    ]\n    return {'text': texts}\n\ntrain_dataset = distill_dataset.map(formatting_prompts_func, batched=True)\nprint(train_dataset[0]['text'][:500])"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 6. 학생 학습 (SFT)\n\n교사 응답을 정답 삼아 학생을 미세조정합니다. 증류셋이 작으므로 여러 epoch를 돌립니다. `train_on_responses_only`로 **assistant 응답 구간에만** 손실을 적용해, 학생이 교사의 *답하는 법*에 집중하게 합니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from trl import SFTConfig, SFTTrainer\nfrom unsloth.chat_templates import train_on_responses_only\n\ntrainer = SFTTrainer(\n    model=model,\n    tokenizer=tokenizer,\n    train_dataset=train_dataset,\n    args=SFTConfig(\n        dataset_text_field='text',\n        max_seq_length=2048,\n        packing=False,\n        per_device_train_batch_size=2,\n        gradient_accumulation_steps=4,\n        warmup_steps=5,\n        num_train_epochs=3,   # 작은 증류셋이므로 여러 epoch\n        learning_rate=2e-4,\n        logging_steps=1,\n        optim='adamw_8bit',\n        weight_decay=0.01,\n        lr_scheduler_type='linear',\n        seed=3407,\n        output_dir='qwen3-0.6b-distilled',\n        report_to='none',\n    ),\n)\n\ntrainer = train_on_responses_only(\n    trainer,\n    instruction_part='<|im_start|>user\\n',\n    response_part='<|im_start|>assistant\\n',\n)\n\ntrainer_stats = trainer.train()"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 7. 증류 결과 확인\n\n학습한 적 없는(held-out) 문제를 학생에게 풀게 합니다. 0.6B 모델이 교사처럼 `<think>` 안에서 단계적으로 추론한 뒤 답을 내놓는지 살핍니다. 증류 전 베이스 학생과 비교하려면 LoRA 없이 같은 모델을 따로 적재해 같은 프롬프트로 생성해 보면 됩니다."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from transformers import TextStreamer\n\nFastLanguageModel.for_inference(model)\n\nquestion = ('Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes '\n            'of babysitting. How much did she earn?')\ninput_ids = tokenizer.apply_chat_template(\n    [{'role': 'user', 'content': question}],\n    add_generation_prompt=True, return_tensors='pt').to('cuda')\n\n_ = model.generate(\n    input_ids, streamer=TextStreamer(tokenizer, skip_prompt=True),\n    max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 8. 저장"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "model.save_pretrained('qwen3-0.6b-distilled-lora')\ntokenizer.save_pretrained('qwen3-0.6b-distilled-lora')"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 마무리\n\n로짓이나 교사 가중치에 접근하지 않고 **교사가 생성한 텍스트만으로** 작은 학생에게 추론 양식을 이식했습니다. 이것이 현대 LLM에서 \"증류\"라 부르는 작업의 실체이며, `R1-Distill` 계열도 같은 원리(규모만 훨씬 큰)로 만들어졌습니다.\n\n### 심화: 화이트박스(로짓) 증류\n\n교사의 출력 **분포 전체**를 KL 발산으로 모방하는 화이트박스 증류는 보통 교사·학생의 **어휘(토크나이저) 정렬**이 걸림돌입니다. 하지만 **Qwen3 패밀리는 0.6B부터 8B까지 토크나이저를 공유**하므로 이 문제가 없습니다. `trl`의 [`GKDTrainer`](https://huggingface.co/docs/trl/gkd_trainer)로 교사 `Qwen3-8B`(4비트)·학생 `Qwen3-0.6B` 조합을 구성하면 KL 손실 기반 증류를 단일 GPU에서 실험할 수 있습니다. 다만 두 모델을 동시에 적재해야 해 메모리·코드가 무거우므로, 본 실습은 시퀀스 수준 증류를 기본으로 두었습니다."
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}