diff --git a/notebook/0_tokenizer.ipynb b/notebook/0_tokenizer.ipynb
index 8bde653..da14045 100644
--- a/notebook/0_tokenizer.ipynb
+++ b/notebook/0_tokenizer.ipynb
@@ -132,7 +132,7 @@
     "            data = json.loads(line)\n",
     "            yield data['text']\n",
     "\n",
-    "data_path = './demo/tokenizer_data.jsonl'"
+    "data_path = './toydata/tokenizer_data.jsonl'"
    ]
   },
   {
@@ -220,7 +220,7 @@
     {
      "data": {
       "text/plain": [
-       "['./demo/tokenizer_demo\\\\vocab.json', './demo/tokenizer_demo\\\\merges.txt']"
+       "['./model/toy_tokenizer\\\\vocab.json', './model/toy_tokenizer\\\\merges.txt']"
       ]
      },
      "execution_count": 8,
@@ -231,7 +231,7 @@
    "source": [
     "import os\n",
     "\n",
-    "tokenizer_dir = \"./demo/tokenizer_demo\"\n",
+    "tokenizer_dir = \"./model/toy_tokenizer\"\n",
     "os.makedirs(tokenizer_dir, exist_ok=True)\n",
     "tokenizer.save(os.path.join(tokenizer_dir, \"tokenizer.json\")) # At this point, you will see a file named tokenizer.json under tokenizer_dir\n",
     "tokenizer.model.save(tokenizer_dir) # generate vocab.json & merges.txt"
@@ -332,7 +332,7 @@
    "source": [
     "from transformers import AutoTokenizer\n",
     "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(\"./demo/tokenizer_demo\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"./model/toy_tokenizer\")\n",
     "msg = [{\"text\": \"失去的东西就要学着去接受，学着放下。\"}]\n",
     "new_msg = tokenizer.apply_chat_template(\n",
     "    msg,\n",
diff --git a/notebook/1_model.ipynb b/notebook/1_model.ipynb
index 5ca2d42..c262827 100644
--- a/notebook/1_model.ipynb
+++ b/notebook/1_model.ipynb
@@ -27,7 +27,7 @@
     "\n",
     "作者提供了对其 MiniMind Dense Model 模型结构的可视化：\n",
     "\n",
-    "![](../images/LLM-structure.png)"
+    "![image](../images/LLM-structure.png)"
    ]
   },
   {
@@ -42,7 +42,7 @@
     "import inspect\n",
     "import time\n",
     "\n",
-    "from demo.LMConfig import LMConfig\n",
+    "from model.LMConfig import LMConfig\n",
     "from typing import Any, Optional, Tuple, List\n",
     "import numpy as np\n",
     "import torch\n",
@@ -263,7 +263,7 @@
     "\n",
     "具体来说, GQA 中将 h 个查询头分为 G 组, 每组 包含 h / G 个查询头，并共享一个公共的键和值.\n",
     "\n",
-    "![img](./demo/gqa.png)\n",
+    "![img](./images/gqa.png)\n",
     "\n",
     "GQA 相比传统的 MHA， 减少了键和值的数量，降低了计算量和内存开销，提高了推理速度.\n",
     "\n",
@@ -311,7 +311,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from demo.LMConfig import LMConfig\n",
     "from typing import Any,  Optional,  Tuple,  List\n",
     "import torch.nn as nn\n",
     "import math\n",
@@ -800,9 +799,9 @@
       "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
       "<==== forward propagation completed, num_kv_cache = 2\n",
       "\n",
-      "new tokens list :[tensor([[2747]]), tensor([[2065]]), tensor([[5211]])]\n",
+      "new tokens list :[tensor([[5081]]), tensor([[4500]]), tensor([[850]])]\n",
       "\n",
-      "生成结果：tensor([[   1,    3,    5,    7, 2747, 2065, 5211]])\n"
+      "生成结果：tensor([[   1,    3,    5,    7, 5081, 4500,  850]])\n"
      ]
     }
    ],
@@ -988,9 +987,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "shape: topk_idx = torch.Size([64, 2]), topk_weight = torch.Size([64, 2]), aux_loss = 0.10118068754673004\n",
-      "token 0 选择的专家：idx = tensor([3, 2]), weight = tensor([3.7838e-21, 2.9849e-21], grad_fn=<SelectBackward0>)\n",
-      "辅助损失：aux_loss = 0.10118068754673004\n"
+      "shape: topk_idx = torch.Size([64, 2]), topk_weight = torch.Size([64, 2]), aux_loss = 0.10163114219903946\n",
+      "token 0 选择的专家：idx = tensor([2, 0]), weight = tensor([5.8955e-21, 1.9281e-21], grad_fn=<SelectBackward0>)\n",
+      "辅助损失：aux_loss = 0.10163114219903946\n"
      ]
     }
    ],
@@ -1364,9 +1363,9 @@
       "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
       "<==== forward propagation completed, num_kv_cache = 2\n",
       "\n",
-      "new tokens list :[tensor([[1233]]), tensor([[2740]]), tensor([[2772]])]\n",
+      "new tokens list :[tensor([[4065]]), tensor([[431]]), tensor([[2541]])]\n",
       "\n",
-      "生成结果：tensor([[   1,    3,    5,    7, 1233, 2740, 2772]])\n"
+      "生成结果：tensor([[   1,    3,    5,    7, 4065,  431, 2541]])\n"
      ]
     }
    ],
diff --git a/notebook/2_dataset.ipynb b/notebook/2_dataset.ipynb
index 2fbd650..392eabb 100644
--- a/notebook/2_dataset.ipynb
+++ b/notebook/2_dataset.ipynb
@@ -99,7 +99,7 @@
    ],
    "source": [
     "# 我们可以查看一下 demo 中提供的数据\n",
-    "path_pretrain = './demo/pretrain_data.jsonl'\n",
+    "path_pretrain = './toydata/pretrain_data.jsonl'\n",
     "\n",
     "with open(path_pretrain, 'r', encoding='utf-8') as f:\n",
     "    for line_num, line in enumerate(f, 1):\n",
@@ -228,7 +228,7 @@
    ],
    "source": [
     "# 我们可以查看一下 demo 中提供的数据\n",
-    "path_sft = './demo/sft_data.jsonl'\n",
+    "path_sft = './toydata/sft_data.jsonl'\n",
     "\n",
     "with open(path_sft, 'r', encoding='utf-8') as f:\n",
     "    for line_num, line in enumerate(f, 1):\n",
@@ -391,7 +391,7 @@
    ],
    "source": [
     "# 我们可以查看一下 demo 中提供的数据\n",
-    "path_dpo = './demo/dpo_data.jsonl'\n",
+    "path_dpo = './toydata/dpo_data.jsonl'\n",
     "\n",
     "with open(path_sft, 'r', encoding='utf-8') as f:\n",
     "    for line_num, line in enumerate(f, 1):\n",
diff --git a/notebook/3_pretrain.ipynb b/notebook/3_pretrain.ipynb
index d6fb2d9..173db88 100644
--- a/notebook/3_pretrain.ipynb
+++ b/notebook/3_pretrain.ipynb
@@ -37,9 +37,9 @@
     "\n",
     "from transformers import AutoTokenizer\n",
     "\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import PretrainDataset"
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import PretrainDataset"
    ]
   },
   {
@@ -90,7 +90,7 @@
     "    n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/pretrain_data.jsonl' # 数据集路径"
+    "    data_path: str = './toydata/pretrain_data.jsonl' # 数据集路径"
    ]
   },
   {
@@ -147,7 +147,7 @@
      "output_type": "stream",
      "text": [
       "LLM总参数量：8.915 百万\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001E4890FE1D0>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001C6BF33F280>\n"
      ]
     }
    ],
@@ -171,7 +171,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 7,
    "id": "60c8c970-f448-4ea0-b21b-8f929e03f6fa",
    "metadata": {},
    "outputs": [
@@ -331,7 +331,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch:[1/1](0/1) loss:8.853 lr:0.000550000000 epoch_Time:0.0min:\n"
+      "Epoch:[1/1](0/1) loss:8.959 lr:0.000550000000 epoch_Time:0.0min:\n"
      ]
     }
    ],
@@ -343,7 +343,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "id": "6f257b9e-3a1e-4686-9516-3f9eeb7d0654",
    "metadata": {},
    "outputs": [],
diff --git a/notebook/4_sft.ipynb b/notebook/4_sft.ipynb
index 235f895..ff9d79a 100644
--- a/notebook/4_sft.ipynb
+++ b/notebook/4_sft.ipynb
@@ -39,9 +39,9 @@
     "from torch.nn.parallel import DistributedDataParallel\n",
     "from torch.utils.data import DataLoader, DistributedSampler\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import SFTDataset"
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import SFTDataset"
    ]
   },
   {
@@ -92,7 +92,7 @@
     "    n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/sft_data.jsonl' # 数据集路径"
+    "    data_path: str = './toydata/sft_data.jsonl' # 数据集路径"
    ]
   },
   {
@@ -155,7 +155,7 @@
      "output_type": "stream",
      "text": [
       "LLM总参数量：8.915 百万\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000002D0DD01D600>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000024234BBE950>\n"
      ]
     }
    ],
@@ -341,7 +341,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch:[1/1](0/1) loss:8.869 lr:0.000550000000 epoch_Time:0.0min:\n"
+      "Epoch:[1/1](0/1) loss:8.902 lr:0.000550000000 epoch_Time:0.0min:\n"
      ]
     }
    ],
diff --git a/notebook/5_dpo.ipynb b/notebook/5_dpo.ipynb
index bab8cea..06dc24c 100644
--- a/notebook/5_dpo.ipynb
+++ b/notebook/5_dpo.ipynb
@@ -39,9 +39,9 @@
     "from torch.nn.parallel import DistributedDataParallel\n",
     "from torch.utils.data import DataLoader, DistributedSampler\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import DPODataset"
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import DPODataset"
    ]
   },
   {
@@ -93,7 +93,7 @@
     "    n_layers: int = 1 # MiniMind Block 数量 模型超参数 | 由于 dpo 要加载两个模型 我们出于演示目的设定 n_layers = 1\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/dpo_data.jsonl' # 数据集路径"
+    "    data_path: str = './toydata/dpo_data.jsonl' # 数据集路径"
    ]
   },
   {
@@ -165,7 +165,7 @@
      "output_type": "stream",
      "text": [
       "LLM总参数量：6.096 百万\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x00000264370DC7F0>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001AF0BBB71C0>\n"
      ]
     }
    ],
diff --git a/notebook/6_lora.ipynb b/notebook/6_lora.ipynb
index bb98b13..0ec199d 100644
--- a/notebook/6_lora.ipynb
+++ b/notebook/6_lora.ipynb
@@ -13,7 +13,7 @@
     "\n",
     "具体来说，LoRA 它冻结了预训练模型 W 的权重，并注入可训练的秩分解矩阵 A 与 B，在微调时，只训练降维矩阵 A 和 升维矩阵 B，微调结束后，将 AB 与 W 进行叠加.\n",
     "\n",
-    "![image](./demo/lora.png)\n",
+    "![images](./images/lora.png)\n",
     "\n",
     "其中，用随机高斯分布进行初始化 A，用 0 矩阵初始化 B，从而保证训练开始时旁路矩阵为 0 矩阵.\n",
     "\n",
@@ -142,7 +142,9 @@
    "cell_type": "code",
    "execution_count": 5,
    "id": "eb24e24f-da09-4cd7-b4ce-7fe482194f75",
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -237,9 +239,9 @@
     "from contextlib import nullcontext\n",
     "from torch.utils.data import DataLoader, DistributedSampler\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import SFTDataset\n",
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import SFTDataset\n",
     "\n",
     "warnings.filterwarnings('ignore')"
    ]
@@ -282,7 +284,7 @@
     "    n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/lora_data.jsonl' # 数据集路径\n",
+    "    data_path: str = './toydata/lora_data.jsonl' # 数据集路径\n",
     "    lora_name: str = 'lora_identity' # 根据任务保存成lora_(英文/医学/心理...)"
    ]
   },
@@ -344,7 +346,7 @@
       "apply lora on module: layers.0.attention.wo\n",
       "apply lora on module: layers.1.attention.wq\n",
       "apply lora on module: layers.1.attention.wo\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000018FB27CA050>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000022F5786D4E0>\n"
      ]
     }
    ],
@@ -540,7 +542,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch:[1/1](0/1) loss:8.894 lr:0.000550000000 epoch_Time:0.0min:\n"
+      "Epoch:[1/1](0/1) loss:8.883 lr:0.000550000000 epoch_Time:0.0min:\n"
      ]
     }
    ],
diff --git a/notebook/7_distill.ipynb b/notebook/7_distill.ipynb
index 2d2744b..6b7b589 100644
--- a/notebook/7_distill.ipynb
+++ b/notebook/7_distill.ipynb
@@ -39,9 +39,9 @@
     "from torch.nn.parallel import DistributedDataParallel\n",
     "from torch.utils.data import DataLoader, DistributedSampler\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import SFTDataset\n",
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import SFTDataset\n",
     "\n",
     "warnings.filterwarnings('ignore')"
    ]
@@ -84,7 +84,7 @@
     "    n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/sft_data.jsonl' # 数据集路径"
+    "    data_path: str = './toydata/sft_data.jsonl' # 数据集路径"
    ]
   },
   {
@@ -157,7 +157,7 @@
      "text": [
       "学生模型(LLM)总参数量：6.096 百万\n",
       "教师模型(LLM)总参数量：17.305 百万\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000002AE220D7880>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000021E5B6E7010>\n"
      ]
     }
    ],
@@ -373,7 +373,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch:[0/0](0/1) loss:0.3301 lr:0.000550000000 epoch_Time:0.0min:\n"
+      "Epoch:[0/0](0/1) loss:0.3257 lr:0.000550000000 epoch_Time:0.0min:\n"
      ]
     }
    ],
diff --git a/notebook/8_reason.ipynb b/notebook/8_reason.ipynb
index aaff2f8..e7329c9 100644
--- a/notebook/8_reason.ipynb
+++ b/notebook/8_reason.ipynb
@@ -49,9 +49,9 @@
     "from torch.nn.parallel import DistributedDataParallel\n",
     "from torch.utils.data import DataLoader, DistributedSampler\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "from demo.model import MiniMindLM\n",
-    "from demo.LMConfig import LMConfig\n",
-    "from demo.dataset import SFTDataset\n",
+    "from model.model import MiniMindLM\n",
+    "from model.LMConfig import LMConfig\n",
+    "from model.dataset import SFTDataset\n",
     "\n",
     "warnings.filterwarnings('ignore')"
    ]
@@ -94,7 +94,7 @@
     "    n_layers: int = 2 # MiniMind Block 数量 模型超参数\n",
     "    max_seq_len: int = 512 # 序列长度阈值\n",
     "    use_moe: bool = False # 是否启用混合专家\n",
-    "    data_path: str = './demo/r1_data.jsonl' # 数据集路径"
+    "    data_path: str = './toydata/r1_data.jsonl' # 数据集路径"
    ]
   },
   {
@@ -154,7 +154,7 @@
      "output_type": "stream",
      "text": [
       "LLM总参数量：8.915 百万\n",
-      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000021EAB98B940>\n"
+      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001DAD0FFD630>\n"
      ]
     }
    ],
@@ -314,7 +314,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch:[1/1](0/1) loss:11.907 lr:0.000550000000 epoch_Time:0.0min:\n"
+      "Epoch:[1/1](0/1) loss:11.656 lr:0.000550000000 epoch_Time:0.0min:\n"
      ]
     }
    ],
diff --git a/notebook/demo/gqa.png b/notebook/images/gqa.png
similarity index 100%
rename from notebook/demo/gqa.png
rename to notebook/images/gqa.png
diff --git a/notebook/demo/lora.png b/notebook/images/lora.png
similarity index 100%
rename from notebook/demo/lora.png
rename to notebook/images/lora.png
diff --git a/notebook/demo/LMConfig.py b/notebook/model/LMConfig.py
similarity index 100%
rename from notebook/demo/LMConfig.py
rename to notebook/model/LMConfig.py
diff --git a/notebook/demo/dataset.py b/notebook/model/dataset.py
similarity index 100%
rename from notebook/demo/dataset.py
rename to notebook/model/dataset.py
diff --git a/notebook/demo/model.py b/notebook/model/model.py
similarity index 100%
rename from notebook/demo/model.py
rename to notebook/model/model.py
diff --git a/notebook/demo/dpo_data.jsonl b/notebook/toydata/dpo_data.jsonl
similarity index 100%
rename from notebook/demo/dpo_data.jsonl
rename to notebook/toydata/dpo_data.jsonl
diff --git a/notebook/demo/lora_data.jsonl b/notebook/toydata/lora_data.jsonl
similarity index 100%
rename from notebook/demo/lora_data.jsonl
rename to notebook/toydata/lora_data.jsonl
diff --git a/notebook/demo/pretrain_data.jsonl b/notebook/toydata/pretrain_data.jsonl
similarity index 100%
rename from notebook/demo/pretrain_data.jsonl
rename to notebook/toydata/pretrain_data.jsonl
diff --git a/notebook/demo/r1_data.jsonl b/notebook/toydata/r1_data.jsonl
similarity index 100%
rename from notebook/demo/r1_data.jsonl
rename to notebook/toydata/r1_data.jsonl
diff --git a/notebook/demo/sft_data.jsonl b/notebook/toydata/sft_data.jsonl
similarity index 100%
rename from notebook/demo/sft_data.jsonl
rename to notebook/toydata/sft_data.jsonl
diff --git a/notebook/demo/tokenizer_data.jsonl b/notebook/toydata/tokenizer_data.jsonl
similarity index 100%
rename from notebook/demo/tokenizer_data.jsonl
rename to notebook/toydata/tokenizer_data.jsonl