{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a8b9aefc04d48d490f872996c22f513",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0938, Accuracy: 0.0156\n",
      "Step 32, Loss: 1.5938, Accuracy: 0.2988\n",
      "Step 64, Loss: 1.7031, Accuracy: 0.3203\n",
      "Step 96, Loss: 1.7188, Accuracy: 0.2539\n",
      "Step 128, Loss: 1.5938, Accuracy: 0.3320\n",
      "Step 160, Loss: 1.5781, Accuracy: 0.3340\n",
      "Step 192, Loss: 1.5781, Accuracy: 0.3535\n",
      "Step 224, Loss: 1.6641, Accuracy: 0.3652\n",
      "Step 256, Loss: 1.5625, Accuracy: 0.3984\n",
      "Step 288, Loss: 1.6094, Accuracy: 0.3105\n",
      "Step 320, Loss: 1.5703, Accuracy: 0.3926\n",
      "Step 352, Loss: 1.5703, Accuracy: 0.3691\n",
      "Step 384, Loss: 1.5078, Accuracy: 0.3906\n",
      "Step 416, Loss: 1.5234, Accuracy: 0.4238\n",
      "Step 448, Loss: 1.5703, Accuracy: 0.3887\n",
      "Step 480, Loss: 1.5234, Accuracy: 0.3828\n",
      "Step 512, Loss: 1.5312, Accuracy: 0.4004\n",
      "Step 544, Loss: 1.5781, Accuracy: 0.3594\n",
      "Step 576, Loss: 1.4844, Accuracy: 0.4395\n",
      "Step 608, Loss: 1.5234, Accuracy: 0.4258\n",
      "Step 640, Loss: 1.5781, Accuracy: 0.3535\n",
      "Step 672, Loss: 1.5312, Accuracy: 0.3750\n",
      "Step 704, Loss: 1.5156, Accuracy: 0.4102\n",
      "Step 736, Loss: 1.5391, Accuracy: 0.4043\n",
      "Step 768, Loss: 1.6016, Accuracy: 0.3496\n",
      "Step 800, Loss: 1.4922, Accuracy: 0.3945\n",
      "Step 832, Loss: 1.5312, Accuracy: 0.4238\n",
      "Step 864, Loss: 1.6250, Accuracy: 0.3672\n",
      "Step 896, Loss: 1.5703, Accuracy: 0.3613\n",
      "Step 928, Loss: 1.5625, Accuracy: 0.3711\n",
      "Step 960, Loss: 1.5781, Accuracy: 0.3398\n",
      "Step 992, Loss: 1.5312, Accuracy: 0.3672\n",
      "Step 1024, Loss: 1.5703, Accuracy: 0.3711\n",
      "Step 1056, Loss: 1.5469, Accuracy: 0.3691\n",
      "Step 1088, Loss: 1.5703, Accuracy: 0.3984\n",
      "Step 1120, Loss: 1.5547, Accuracy: 0.3965\n",
      "Step 1152, Loss: 1.5469, Accuracy: 0.3691\n",
      "Step 1184, Loss: 1.5391, Accuracy: 0.3887\n",
      "Step 1216, Loss: 1.5859, Accuracy: 0.3770\n",
      "Step 1248, Loss: 1.5703, Accuracy: 0.3789\n",
      "Step 1280, Loss: 1.5312, Accuracy: 0.4258\n",
      "Step 1312, Loss: 1.5156, Accuracy: 0.3789\n",
      "Step 1344, Loss: 1.5781, Accuracy: 0.3633\n",
      "Step 1376, Loss: 1.5625, Accuracy: 0.3613\n",
      "Step 1408, Loss: 1.5703, Accuracy: 0.3438\n",
      "Step 1440, Loss: 1.5312, Accuracy: 0.3945\n",
      "Step 1472, Loss: 1.4844, Accuracy: 0.4902\n",
      "Step 1504, Loss: 1.5312, Accuracy: 0.3926\n",
      "Step 1536, Loss: 1.5156, Accuracy: 0.4121\n",
      "Step 1568, Loss: 1.5859, Accuracy: 0.3379\n",
      "Step 1600, Loss: 1.5625, Accuracy: 0.3691\n",
      "Step 1632, Loss: 1.5000, Accuracy: 0.3906\n",
      "Step 1664, Loss: 1.5312, Accuracy: 0.3809\n",
      "Step 1696, Loss: 1.5547, Accuracy: 0.3516\n",
      "Step 1728, Loss: 1.5625, Accuracy: 0.3398\n",
      "Step 1760, Loss: 1.5625, Accuracy: 0.3965\n",
      "Step 1792, Loss: 1.5078, Accuracy: 0.4160\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[38]\u001b[39m\u001b[32m, line 167\u001b[39m\n\u001b[32m    165\u001b[39m \u001b[38;5;66;03m# 反向传播\u001b[39;00m\n\u001b[32m    166\u001b[39m loss.backward()\n\u001b[32m--> \u001b[39m\u001b[32m167\u001b[39m \u001b[43moptimizer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    168\u001b[39m lr_scheduler.step()\n\u001b[32m    169\u001b[39m optimizer.zero_grad()\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/optim/lr_scheduler.py:137\u001b[39m, in \u001b[36mLRScheduler.__init__.<locals>.patch_track_step_called.<locals>.wrap_step.<locals>.wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m    135\u001b[39m opt = opt_ref()\n\u001b[32m    136\u001b[39m opt._opt_called = \u001b[38;5;28;01mTrue\u001b[39;00m  \u001b[38;5;66;03m# type: ignore[union-attr]\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m137\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m.\u001b[49m\u001b[34;43m__get__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mopt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mopt\u001b[49m\u001b[43m.\u001b[49m\u001b[34;43m__class__\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/optim/optimizer.py:487\u001b[39m, in \u001b[36mOptimizer.profile_hook_step.<locals>.wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m    482\u001b[39m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    483\u001b[39m             \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[32m    484\u001b[39m                 \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m must return None or a tuple of (new_args, new_kwargs), but got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresult\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    485\u001b[39m             )\n\u001b[32m--> \u001b[39m\u001b[32m487\u001b[39m out = \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    488\u001b[39m \u001b[38;5;28mself\u001b[39m._optimizer_step_code()\n\u001b[32m    490\u001b[39m \u001b[38;5;66;03m# call optimizer step post hooks\u001b[39;00m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py:116\u001b[39m, in \u001b[36mcontext_decorator.<locals>.decorate_context\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m    113\u001b[39m \u001b[38;5;129m@functools\u001b[39m.wraps(func)\n\u001b[32m    114\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[34mdecorate_context\u001b[39m(*args, **kwargs):\n\u001b[32m    115\u001b[39m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[32m--> \u001b[39m\u001b[32m116\u001b[39m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.local/lib/python3.11/site-packages/transformers/optimization.py:647\u001b[39m, in \u001b[36mAdamW.step\u001b[39m\u001b[34m(self, closure)\u001b[39m\n\u001b[32m    644\u001b[39m \u001b[38;5;66;03m# Decay the first and second moment running average coefficient\u001b[39;00m\n\u001b[32m    645\u001b[39m \u001b[38;5;66;03m# In-place operations to update the averages at the same time\u001b[39;00m\n\u001b[32m    646\u001b[39m exp_avg.mul_(beta1).add_(grad, alpha=(\u001b[32m1.0\u001b[39m - beta1))\n\u001b[32m--> \u001b[39m\u001b[32m647\u001b[39m \u001b[43mexp_avg_sq\u001b[49m\u001b[43m.\u001b[49m\u001b[43mmul_\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbeta2\u001b[49m\u001b[43m)\u001b[49m.addcmul_(grad, grad, value=\u001b[32m1.0\u001b[39m - beta2)\n\u001b[32m    648\u001b[39m denom = exp_avg_sq.sqrt().add_(group[\u001b[33m\"\u001b[39m\u001b[33meps\u001b[39m\u001b[33m\"\u001b[39m])\n\u001b[32m    650\u001b[39m step_size = group[\u001b[33m\"\u001b[39m\u001b[33mlr\u001b[39m\u001b[33m\"\u001b[39m]\n",
      "\u001b[31mKeyboardInterrupt\u001b[39m: "
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "\n",
    "from itertools import permutations\n",
    "\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    \"\"\"\n",
    "    Generates data samples with step-by-step first element tracking.\n",
    "    \n",
    "    :param k: Length of each swap sequence.\n",
    "    :param num_samples: Number of samples to generate.\n",
    "    :return: A list of tuples. Each tuple contains a swap sequence and a list of first elements after each swap.\n",
    "    \"\"\"\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        # 随机生成交换序列\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        \n",
    "        # 初始排列：(0,1,2,3,4)\n",
    "        current_perm = tuple(range(n_elements ))\n",
    "        \n",
    "        # 记录每次交换后的第一个元素\n",
    "        first_elements = []\n",
    "        \n",
    "        # 应用交换序列并记录每次交换后的第一个元素\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "        \n",
    "        # 将 swap_sequence 和 label 转换为张量\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "\n",
    "        # 添加到数据中\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "# 自定义 Dataset 类\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples  # 每次只返回 num_samples 条数据\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # 每次访问时重新生成数据，保证每个 epoch 数据不同\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 参数设置\n",
    "config = fla.models.TransformerConfig(\n",
    "    vocab_size=128,\n",
    "    hidden_size=1024,\n",
    "    num_hidden_layers=8,\n",
    "    num_heads=8,\n",
    "    num_kv_heads=1,\n",
    "    max_position_embeddings=2048,\n",
    "    layer_norm_eps=1e-6,\n",
    "    use_rotary_pos_emb=None,\n",
    "    fuse_cross_entropy = False,\n",
    "    rope_theta = 1e4,\n",
    ")\n",
    "\n",
    "model =AutoModelForCausalLM.from_config(config)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model.to(device).to(torch.bfloat16)\n",
    "\n",
    "# 超参数\n",
    "num_epochs = 2000\n",
    "batch_size = 32\n",
    "k = 16\n",
    "num_samples_per_batch = batch_size  # 每次生成的样本数 = batch size\n",
    "\n",
    "# 数据集和数据加载器\n",
    "dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# 优化器和学习率调度器\n",
    "optimizer = AdamW(model.parameters(), lr=1e-3,weight_decay=0.01)\n",
    "num_training_steps = num_epochs * len(dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    ")\n",
    "\n",
    "# 训练循环\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "model.train()\n",
    "step = 0\n",
    "for epoch in range(num_epochs):\n",
    "    for batch in dataloader:\n",
    "        # 将输入移至设备\n",
    "        # print(batch)\n",
    "\n",
    "        # assert 1==2\n",
    "        input_ids = batch[\"input_ids\"].to(device)\n",
    "        labels = batch[\"labels\"].to(device)\n",
    "\n",
    "        # 添加 dummy attention_mask，防止警告\n",
    "        # attention_mask = (input_ids != 0).long().to(device)\n",
    "\n",
    "        # 前向传播\n",
    "        outputs = model(\n",
    "            input_ids=input_ids,\n",
    "            labels=labels,\n",
    "            return_dict = True\n",
    "            # attention_mask=attention_mask\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "\n",
    "        # 准确率计算\n",
    "        logits = outputs.logits\n",
    "        # print(outputs)\n",
    "        predictions = torch.argmax(logits, dim=-1)\n",
    "        correct = (predictions == labels).sum().item()\n",
    "        total = labels.numel()\n",
    "        accuracy = correct / total\n",
    "\n",
    "        # 反向传播\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 打印信息\n",
    "        if step%32==0:\n",
    "            print(f\"Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "        step += 1\n",
    "\n",
    "        progress_bar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class S5DataModule:\n",
    "    \"\"\"\n",
    "    S5 DataModule: 负责生成 S5 置换跟踪数据，\n",
    "    并提供 Dataset 与 DataLoader 接口。\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        n_elements: int = 5,\n",
    "        k: int = 4,\n",
    "        num_samples: int = 1000,\n",
    "        batch_size: int = 32,\n",
    "        shuffle: bool = True,\n",
    "        num_workers: int = 0,\n",
    "        pin_memory: bool = False,\n",
    "    ):\n",
    "        self.n_elements = n_elements\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "        # 预先生成所有可能的 swap 对\n",
    "        self.swap_pairs = [\n",
    "            (i, j)\n",
    "            for i in range(self.n_elements)\n",
    "            for j in range(i + 1, self.n_elements)\n",
    "        ]\n",
    "\n",
    "        # DataLoader 配置\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.num_workers = num_workers\n",
    "        self.pin_memory = pin_memory\n",
    "\n",
    "        # 立即生成一次完整数据集\n",
    "        self.data = self._generate_all_data()\n",
    "\n",
    "    def _apply_swap(self, perm: tuple[int, ...], swap_idx: int) -> tuple[int, ...]:\n",
    "        \"\"\"对当前置换 perm 执行一次 swap_pairs[swap_idx] 操作。\"\"\"\n",
    "        i, j = self.swap_pairs[swap_idx]\n",
    "        lst = list(perm)\n",
    "        lst[i], lst[j] = lst[j], lst[i]\n",
    "        return tuple(lst)\n",
    "\n",
    "    def _generate_one_sample(self) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        生成单条样本：\n",
    "        - input_ids: 大小 (k,) 的 swap 索引序列\n",
    "        - labels:   大小 (k,) 的每一步置换后第一个元素\n",
    "        \"\"\"\n",
    "        swap_sequence = [\n",
    "            random.randint(0, len(self.swap_pairs) - 1)\n",
    "            for _ in range(self.k)\n",
    "        ]\n",
    "        perm = tuple(range(self.n_elements))\n",
    "        first_elems = []\n",
    "\n",
    "        for idx in swap_sequence:\n",
    "            perm = self._apply_swap(perm, idx)\n",
    "            first_elems.append(perm[0])\n",
    "\n",
    "        return (\n",
    "            torch.tensor(swap_sequence, dtype=torch.long),\n",
    "            torch.tensor(first_elems, dtype=torch.long),\n",
    "        )\n",
    "\n",
    "    def _generate_all_data(self) -> list[tuple[torch.Tensor, torch.Tensor]]:\n",
    "        \"\"\"一次性生成 num_samples 条样本，存入内存。\"\"\"\n",
    "        return [self._generate_one_sample() for _ in range(self.num_samples)]\n",
    "\n",
    "    class _InternalDataset(Dataset):\n",
    "        \"\"\"内部用 Dataset 直接从预生成的数据列表读取。\"\"\"\n",
    "        def __init__(self, data):\n",
    "            self.data = data\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.data)\n",
    "\n",
    "        def __getitem__(self, idx):\n",
    "            input_ids, labels = self.data[idx]\n",
    "            return {\"input_ids\": input_ids, \"labels\": labels}\n",
    "\n",
    "    def get_dataloader(self) -> DataLoader:\n",
    "        \"\"\"返回一个标准的 PyTorch DataLoader。\"\"\"\n",
    "        ds = self._InternalDataset(self.data)\n",
    "        return DataLoader(\n",
    "            ds,\n",
    "            batch_size=self.batch_size,\n",
    "            shuffle=self.shuffle,\n",
    "            num_workers=self.num_workers,\n",
    "            pin_memory=self.pin_memory,\n",
    "        )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-05-14 11:03:29.171567: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-05-14 11:03:29.303341: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1e559abb31e4f189556e93b5700bb9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 5.2188, Accuracy: 0.0000\n",
      "1_layers - Step 32, Loss: 1.6016, Accuracy: 0.3359\n",
      "1_layers - Step 64, Loss: 1.6016, Accuracy: 0.3340\n",
      "1_layers - Step 96, Loss: 1.5391, Accuracy: 0.4062\n",
      "1_layers - Step 128, Loss: 1.5859, Accuracy: 0.3125\n",
      "1_layers - Step 160, Loss: 1.5625, Accuracy: 0.3691\n",
      "1_layers - Step 192, Loss: 1.5469, Accuracy: 0.4023\n",
      "1_layers - Step 224, Loss: 1.5078, Accuracy: 0.3945\n",
      "1_layers - Step 256, Loss: 1.5547, Accuracy: 0.3867\n",
      "1_layers - Step 288, Loss: 1.5625, Accuracy: 0.3770\n",
      "1_layers - Step 320, Loss: 1.6016, Accuracy: 0.3438\n",
      "1_layers - Step 352, Loss: 1.5859, Accuracy: 0.3770\n",
      "1_layers - Step 384, Loss: 1.5469, Accuracy: 0.3848\n",
      "1_layers - Step 416, Loss: 1.5312, Accuracy: 0.3809\n",
      "1_layers - Step 448, Loss: 1.5469, Accuracy: 0.3867\n",
      "1_layers - Step 480, Loss: 1.5625, Accuracy: 0.3828\n",
      "1_layers - Step 512, Loss: 1.5703, Accuracy: 0.3750\n",
      "1_layers - Step 544, Loss: 1.5781, Accuracy: 0.3672\n",
      "1_layers - Step 576, Loss: 1.5391, Accuracy: 0.4023\n",
      "1_layers - Step 608, Loss: 1.5469, Accuracy: 0.3809\n",
      "1_layers - Step 640, Loss: 1.5156, Accuracy: 0.4434\n",
      "1_layers - Step 672, Loss: 1.5469, Accuracy: 0.4082\n",
      "1_layers - Step 704, Loss: 1.5469, Accuracy: 0.3906\n",
      "1_layers - Step 736, Loss: 1.5703, Accuracy: 0.3242\n",
      "1_layers - Step 768, Loss: 1.5938, Accuracy: 0.3438\n",
      "1_layers - Step 800, Loss: 1.5859, Accuracy: 0.4023\n",
      "1_layers - Step 832, Loss: 1.4766, Accuracy: 0.4043\n",
      "1_layers - Step 864, Loss: 1.5391, Accuracy: 0.3926\n",
      "1_layers - Step 896, Loss: 1.6016, Accuracy: 0.3223\n",
      "1_layers - Step 928, Loss: 1.5703, Accuracy: 0.3828\n",
      "1_layers - Step 960, Loss: 1.5391, Accuracy: 0.3984\n",
      "1_layers - Step 992, Loss: 1.5703, Accuracy: 0.3125\n",
      "1_layers - Step 1024, Loss: 1.5625, Accuracy: 0.4121\n",
      "1_layers - Step 1056, Loss: 1.5234, Accuracy: 0.4082\n",
      "1_layers - Step 1088, Loss: 1.5547, Accuracy: 0.4043\n",
      "1_layers - Step 1120, Loss: 1.5547, Accuracy: 0.3535\n",
      "1_layers - Step 1152, Loss: 1.5156, Accuracy: 0.4219\n",
      "1_layers - Step 1184, Loss: 1.5703, Accuracy: 0.4004\n",
      "1_layers - Step 1216, Loss: 1.5859, Accuracy: 0.3418\n",
      "1_layers - Step 1248, Loss: 1.5625, Accuracy: 0.4043\n",
      "1_layers - Step 1280, Loss: 1.5469, Accuracy: 0.3750\n",
      "1_layers - Step 1312, Loss: 1.5547, Accuracy: 0.3984\n",
      "1_layers - Step 1344, Loss: 1.5312, Accuracy: 0.3848\n",
      "1_layers - Step 1376, Loss: 1.5391, Accuracy: 0.4355\n",
      "1_layers - Step 1408, Loss: 1.5469, Accuracy: 0.3906\n",
      "1_layers - Step 1440, Loss: 1.5703, Accuracy: 0.3828\n",
      "1_layers - Step 1472, Loss: 1.5859, Accuracy: 0.3496\n",
      "1_layers - Step 1504, Loss: 1.5547, Accuracy: 0.3438\n",
      "1_layers - Step 1536, Loss: 1.5625, Accuracy: 0.3574\n",
      "1_layers - Step 1568, Loss: 1.5469, Accuracy: 0.3926\n",
      "1_layers - Step 1600, Loss: 1.5078, Accuracy: 0.4375\n",
      "1_layers - Step 1632, Loss: 1.5312, Accuracy: 0.3809\n",
      "1_layers - Step 1664, Loss: 1.5312, Accuracy: 0.4102\n",
      "1_layers - Step 1696, Loss: 1.5156, Accuracy: 0.4414\n",
      "1_layers - Step 1728, Loss: 1.4922, Accuracy: 0.4102\n",
      "1_layers - Step 1760, Loss: 1.5234, Accuracy: 0.4453\n",
      "1_layers - Step 1792, Loss: 1.5078, Accuracy: 0.4238\n",
      "1_layers - Step 1824, Loss: 1.5469, Accuracy: 0.4297\n",
      "1_layers - Step 1856, Loss: 1.5469, Accuracy: 0.3672\n",
      "1_layers - Step 1888, Loss: 1.5156, Accuracy: 0.4629\n",
      "1_layers - Step 1920, Loss: 1.5156, Accuracy: 0.4531\n",
      "1_layers - Step 1952, Loss: 1.5234, Accuracy: 0.3887\n",
      "1_layers - Step 1984, Loss: 1.4922, Accuracy: 0.4238\n",
      "\n",
      "=== Training with 2 layers ===\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6e29515aa474419859c262e048d686e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 5.0312, Accuracy: 0.0000\n",
      "2_layers - Step 32, Loss: 1.6406, Accuracy: 0.2832\n",
      "2_layers - Step 64, Loss: 1.6016, Accuracy: 0.3379\n",
      "2_layers - Step 96, Loss: 1.5938, Accuracy: 0.3613\n",
      "2_layers - Step 128, Loss: 1.6797, Accuracy: 0.3066\n",
      "2_layers - Step 160, Loss: 1.5469, Accuracy: 0.3945\n",
      "2_layers - Step 192, Loss: 1.5703, Accuracy: 0.3438\n",
      "2_layers - Step 224, Loss: 1.5781, Accuracy: 0.3750\n",
      "2_layers - Step 256, Loss: 1.5312, Accuracy: 0.3984\n",
      "2_layers - Step 288, Loss: 1.5547, Accuracy: 0.3672\n",
      "2_layers - Step 320, Loss: 1.5078, Accuracy: 0.3789\n",
      "2_layers - Step 352, Loss: 1.6094, Accuracy: 0.3359\n",
      "2_layers - Step 384, Loss: 1.5781, Accuracy: 0.3789\n",
      "2_layers - Step 416, Loss: 1.5547, Accuracy: 0.3984\n",
      "2_layers - Step 448, Loss: 1.6016, Accuracy: 0.3613\n",
      "2_layers - Step 480, Loss: 1.5000, Accuracy: 0.4121\n",
      "2_layers - Step 512, Loss: 1.5781, Accuracy: 0.4355\n",
      "2_layers - Step 544, Loss: 1.5234, Accuracy: 0.3867\n",
      "2_layers - Step 576, Loss: 1.4922, Accuracy: 0.4375\n",
      "2_layers - Step 608, Loss: 1.5547, Accuracy: 0.3418\n",
      "2_layers - Step 640, Loss: 1.5625, Accuracy: 0.4062\n",
      "2_layers - Step 672, Loss: 1.5625, Accuracy: 0.3926\n",
      "2_layers - Step 704, Loss: 1.5469, Accuracy: 0.4023\n",
      "2_layers - Step 736, Loss: 1.5547, Accuracy: 0.3652\n",
      "2_layers - Step 768, Loss: 1.5156, Accuracy: 0.4688\n",
      "2_layers - Step 800, Loss: 1.4609, Accuracy: 0.4688\n",
      "2_layers - Step 832, Loss: 1.5781, Accuracy: 0.3887\n",
      "2_layers - Step 864, Loss: 1.5938, Accuracy: 0.3984\n",
      "2_layers - Step 896, Loss: 1.5391, Accuracy: 0.4258\n",
      "2_layers - Step 928, Loss: 1.5078, Accuracy: 0.3906\n",
      "2_layers - Step 960, Loss: 1.5469, Accuracy: 0.4355\n",
      "2_layers - Step 992, Loss: 1.5547, Accuracy: 0.4492\n",
      "2_layers - Step 1024, Loss: 1.5234, Accuracy: 0.4199\n",
      "2_layers - Step 1056, Loss: 1.5078, Accuracy: 0.4355\n",
      "2_layers - Step 1088, Loss: 1.5391, Accuracy: 0.4180\n",
      "2_layers - Step 1120, Loss: 1.5625, Accuracy: 0.3789\n",
      "2_layers - Step 1152, Loss: 1.5391, Accuracy: 0.4316\n",
      "2_layers - Step 1184, Loss: 1.5312, Accuracy: 0.3926\n",
      "2_layers - Step 1216, Loss: 1.5781, Accuracy: 0.3770\n",
      "2_layers - Step 1248, Loss: 1.5234, Accuracy: 0.4277\n",
      "2_layers - Step 1280, Loss: 1.5625, Accuracy: 0.4258\n",
      "2_layers - Step 1312, Loss: 1.5547, Accuracy: 0.4238\n",
      "2_layers - Step 1344, Loss: 1.5000, Accuracy: 0.4434\n",
      "2_layers - Step 1376, Loss: 1.5312, Accuracy: 0.4082\n",
      "2_layers - Step 1408, Loss: 1.5469, Accuracy: 0.4375\n",
      "2_layers - Step 1440, Loss: 1.4922, Accuracy: 0.4199\n",
      "2_layers - Step 1472, Loss: 1.5391, Accuracy: 0.4219\n",
      "2_layers - Step 1504, Loss: 1.5312, Accuracy: 0.4219\n",
      "2_layers - Step 1536, Loss: 1.5156, Accuracy: 0.4863\n",
      "2_layers - Step 1568, Loss: 1.5547, Accuracy: 0.4004\n",
      "2_layers - Step 1600, Loss: 1.5234, Accuracy: 0.4551\n",
      "2_layers - Step 1632, Loss: 1.5234, Accuracy: 0.4473\n",
      "2_layers - Step 1664, Loss: 1.5469, Accuracy: 0.4277\n",
      "2_layers - Step 1696, Loss: 1.5156, Accuracy: 0.4473\n",
      "2_layers - Step 1728, Loss: 1.5547, Accuracy: 0.4121\n",
      "2_layers - Step 1760, Loss: 1.5156, Accuracy: 0.4238\n",
      "2_layers - Step 1792, Loss: 1.5000, Accuracy: 0.4805\n",
      "2_layers - Step 1824, Loss: 1.5391, Accuracy: 0.4102\n",
      "2_layers - Step 1856, Loss: 1.5078, Accuracy: 0.4160\n",
      "2_layers - Step 1888, Loss: 1.4844, Accuracy: 0.4082\n",
      "2_layers - Step 1920, Loss: 1.5000, Accuracy: 0.4199\n",
      "2_layers - Step 1952, Loss: 1.5625, Accuracy: 0.4082\n",
      "2_layers - Step 1984, Loss: 1.5000, Accuracy: 0.4434\n",
      "\n",
      "=== Training with 4 layers ===\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f57e2af1c93f41adba33c94cdc96c1f3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.0625, Accuracy: 0.0000\n",
      "4_layers - Step 32, Loss: 1.6719, Accuracy: 0.2363\n",
      "4_layers - Step 64, Loss: 1.6016, Accuracy: 0.3281\n",
      "4_layers - Step 96, Loss: 1.6016, Accuracy: 0.3047\n",
      "4_layers - Step 128, Loss: 1.5000, Accuracy: 0.3359\n",
      "4_layers - Step 160, Loss: 1.5703, Accuracy: 0.3789\n",
      "4_layers - Step 192, Loss: 1.6016, Accuracy: 0.3184\n",
      "4_layers - Step 224, Loss: 1.5234, Accuracy: 0.4023\n",
      "4_layers - Step 256, Loss: 1.5703, Accuracy: 0.3672\n",
      "4_layers - Step 288, Loss: 1.5781, Accuracy: 0.3633\n",
      "4_layers - Step 320, Loss: 1.5312, Accuracy: 0.4180\n",
      "4_layers - Step 352, Loss: 1.5781, Accuracy: 0.3477\n",
      "4_layers - Step 384, Loss: 1.5703, Accuracy: 0.3477\n",
      "4_layers - Step 416, Loss: 1.5391, Accuracy: 0.4121\n",
      "4_layers - Step 448, Loss: 1.5625, Accuracy: 0.3398\n",
      "4_layers - Step 480, Loss: 1.5547, Accuracy: 0.4043\n",
      "4_layers - Step 512, Loss: 1.5391, Accuracy: 0.3984\n",
      "4_layers - Step 544, Loss: 1.5625, Accuracy: 0.3926\n",
      "4_layers - Step 576, Loss: 1.4844, Accuracy: 0.4453\n",
      "4_layers - Step 608, Loss: 1.5312, Accuracy: 0.4023\n",
      "4_layers - Step 640, Loss: 1.5469, Accuracy: 0.3691\n",
      "4_layers - Step 672, Loss: 1.4609, Accuracy: 0.5000\n",
      "4_layers - Step 704, Loss: 1.4922, Accuracy: 0.4785\n",
      "4_layers - Step 736, Loss: 1.4766, Accuracy: 0.4238\n",
      "4_layers - Step 768, Loss: 1.5547, Accuracy: 0.3340\n",
      "4_layers - Step 800, Loss: 1.5078, Accuracy: 0.4082\n",
      "4_layers - Step 832, Loss: 1.5234, Accuracy: 0.4004\n",
      "4_layers - Step 864, Loss: 1.5312, Accuracy: 0.4043\n",
      "4_layers - Step 896, Loss: 1.5312, Accuracy: 0.3809\n",
      "4_layers - Step 928, Loss: 1.4844, Accuracy: 0.4238\n",
      "4_layers - Step 960, Loss: 1.5547, Accuracy: 0.3730\n",
      "4_layers - Step 992, Loss: 1.5391, Accuracy: 0.4316\n",
      "4_layers - Step 1024, Loss: 1.5234, Accuracy: 0.4160\n",
      "4_layers - Step 1056, Loss: 1.5625, Accuracy: 0.4121\n",
      "4_layers - Step 1088, Loss: 1.5469, Accuracy: 0.3809\n",
      "4_layers - Step 1120, Loss: 1.5312, Accuracy: 0.4023\n",
      "4_layers - Step 1152, Loss: 1.5078, Accuracy: 0.4023\n",
      "4_layers - Step 1184, Loss: 1.4766, Accuracy: 0.4570\n",
      "4_layers - Step 1216, Loss: 1.5547, Accuracy: 0.3828\n",
      "4_layers - Step 1248, Loss: 1.5000, Accuracy: 0.4414\n",
      "4_layers - Step 1280, Loss: 1.5234, Accuracy: 0.4141\n",
      "4_layers - Step 1312, Loss: 1.5234, Accuracy: 0.4883\n",
      "4_layers - Step 1344, Loss: 1.5703, Accuracy: 0.4219\n",
      "4_layers - Step 1376, Loss: 1.5234, Accuracy: 0.3887\n",
      "4_layers - Step 1408, Loss: 1.5391, Accuracy: 0.4414\n",
      "4_layers - Step 1440, Loss: 1.5469, Accuracy: 0.3594\n",
      "4_layers - Step 1472, Loss: 1.5391, Accuracy: 0.3867\n",
      "4_layers - Step 1504, Loss: 1.5078, Accuracy: 0.4258\n",
      "4_layers - Step 1536, Loss: 1.5078, Accuracy: 0.3926\n",
      "4_layers - Step 1568, Loss: 1.5156, Accuracy: 0.4375\n",
      "4_layers - Step 1600, Loss: 1.5000, Accuracy: 0.4219\n",
      "4_layers - Step 1632, Loss: 1.5312, Accuracy: 0.4238\n",
      "4_layers - Step 1664, Loss: 1.5156, Accuracy: 0.4648\n",
      "4_layers - Step 1696, Loss: 1.5156, Accuracy: 0.4219\n",
      "4_layers - Step 1728, Loss: 1.5234, Accuracy: 0.4512\n",
      "4_layers - Step 1760, Loss: 1.4375, Accuracy: 0.4902\n",
      "4_layers - Step 1792, Loss: 1.4844, Accuracy: 0.4648\n",
      "4_layers - Step 1824, Loss: 1.5078, Accuracy: 0.4492\n",
      "4_layers - Step 1856, Loss: 1.5938, Accuracy: 0.3594\n",
      "4_layers - Step 1888, Loss: 1.5156, Accuracy: 0.4551\n",
      "4_layers - Step 1920, Loss: 1.5000, Accuracy: 0.4375\n",
      "4_layers - Step 1952, Loss: 1.5078, Accuracy: 0.4570\n",
      "4_layers - Step 1984, Loss: 1.5156, Accuracy: 0.4258\n",
      "\n",
      "=== Training with 8 layers ===\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbf9fc0cd14741f68e3ba9713793533d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 5.0000, Accuracy: 0.0039\n",
      "8_layers - Step 32, Loss: 1.6406, Accuracy: 0.2285\n",
      "8_layers - Step 64, Loss: 1.6250, Accuracy: 0.2871\n",
      "8_layers - Step 96, Loss: 1.6406, Accuracy: 0.2207\n",
      "8_layers - Step 128, Loss: 1.6250, Accuracy: 0.3008\n",
      "8_layers - Step 160, Loss: 1.5859, Accuracy: 0.3125\n",
      "8_layers - Step 192, Loss: 1.5703, Accuracy: 0.3477\n",
      "8_layers - Step 224, Loss: 1.6406, Accuracy: 0.3242\n",
      "8_layers - Step 256, Loss: 1.5703, Accuracy: 0.3770\n",
      "8_layers - Step 288, Loss: 1.5625, Accuracy: 0.4199\n",
      "8_layers - Step 320, Loss: 1.6016, Accuracy: 0.3633\n",
      "8_layers - Step 352, Loss: 1.5312, Accuracy: 0.4531\n",
      "8_layers - Step 384, Loss: 1.5781, Accuracy: 0.4023\n",
      "8_layers - Step 416, Loss: 1.5703, Accuracy: 0.4023\n",
      "8_layers - Step 448, Loss: 1.5312, Accuracy: 0.4688\n",
      "8_layers - Step 480, Loss: 1.5078, Accuracy: 0.4082\n",
      "8_layers - Step 512, Loss: 1.5625, Accuracy: 0.3770\n",
      "8_layers - Step 544, Loss: 1.5625, Accuracy: 0.3984\n",
      "8_layers - Step 576, Loss: 1.5234, Accuracy: 0.3945\n",
      "8_layers - Step 608, Loss: 1.5547, Accuracy: 0.3906\n",
      "8_layers - Step 640, Loss: 1.5312, Accuracy: 0.4297\n",
      "8_layers - Step 672, Loss: 1.5156, Accuracy: 0.3789\n",
      "8_layers - Step 704, Loss: 1.5469, Accuracy: 0.4082\n",
      "8_layers - Step 736, Loss: 1.5938, Accuracy: 0.3809\n",
      "8_layers - Step 768, Loss: 1.5625, Accuracy: 0.3613\n",
      "8_layers - Step 800, Loss: 1.5547, Accuracy: 0.3691\n",
      "8_layers - Step 832, Loss: 1.5781, Accuracy: 0.3945\n",
      "8_layers - Step 864, Loss: 1.4531, Accuracy: 0.5059\n",
      "8_layers - Step 896, Loss: 1.5469, Accuracy: 0.3691\n",
      "8_layers - Step 928, Loss: 1.5469, Accuracy: 0.3809\n",
      "8_layers - Step 960, Loss: 1.5391, Accuracy: 0.4062\n",
      "8_layers - Step 992, Loss: 1.5703, Accuracy: 0.3887\n",
      "8_layers - Step 1024, Loss: 1.5703, Accuracy: 0.3848\n",
      "8_layers - Step 1056, Loss: 1.5000, Accuracy: 0.4512\n",
      "8_layers - Step 1088, Loss: 1.5312, Accuracy: 0.4160\n",
      "8_layers - Step 1120, Loss: 1.5391, Accuracy: 0.4316\n",
      "8_layers - Step 1152, Loss: 1.5000, Accuracy: 0.4453\n",
      "8_layers - Step 1184, Loss: 1.5781, Accuracy: 0.3594\n",
      "8_layers - Step 1216, Loss: 1.5078, Accuracy: 0.4258\n",
      "8_layers - Step 1248, Loss: 1.5234, Accuracy: 0.4219\n",
      "8_layers - Step 1280, Loss: 1.5000, Accuracy: 0.4590\n",
      "8_layers - Step 1312, Loss: 1.4688, Accuracy: 0.4395\n",
      "8_layers - Step 1344, Loss: 1.4922, Accuracy: 0.4551\n",
      "8_layers - Step 1376, Loss: 1.5234, Accuracy: 0.4141\n",
      "8_layers - Step 1408, Loss: 1.4844, Accuracy: 0.4375\n",
      "8_layers - Step 1440, Loss: 1.5156, Accuracy: 0.4668\n",
      "8_layers - Step 1472, Loss: 1.5078, Accuracy: 0.4805\n",
      "8_layers - Step 1504, Loss: 1.5312, Accuracy: 0.4727\n",
      "8_layers - Step 1536, Loss: 1.5078, Accuracy: 0.4395\n",
      "8_layers - Step 1568, Loss: 1.5234, Accuracy: 0.4004\n",
      "8_layers - Step 1600, Loss: 1.5000, Accuracy: 0.4570\n",
      "8_layers - Step 1632, Loss: 1.4766, Accuracy: 0.4785\n",
      "8_layers - Step 1664, Loss: 1.4688, Accuracy: 0.4668\n",
      "8_layers - Step 1696, Loss: 1.4688, Accuracy: 0.5273\n",
      "8_layers - Step 1728, Loss: 1.5078, Accuracy: 0.4160\n",
      "8_layers - Step 1760, Loss: 1.4844, Accuracy: 0.4902\n",
      "8_layers - Step 1792, Loss: 1.4688, Accuracy: 0.4648\n",
      "8_layers - Step 1824, Loss: 1.5391, Accuracy: 0.4160\n",
      "8_layers - Step 1856, Loss: 1.5000, Accuracy: 0.4180\n",
      "8_layers - Step 1888, Loss: 1.5000, Accuracy: 0.4727\n",
      "8_layers - Step 1920, Loss: 1.5469, Accuracy: 0.4707\n",
      "8_layers - Step 1952, Loss: 1.4766, Accuracy: 0.4844\n",
      "8_layers - Step 1984, Loss: 1.4844, Accuracy: 0.4590\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 128872 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.TransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy=False,\n",
    "        rope_theta=1e4,\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    # plt.title(\"Accuracy vs Training Steps for Different Layer Counts\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.pdf\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers\n",
      "2_layers\n",
      "4_layers\n",
      "8_layers\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 128872 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 绘图\n",
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    print(key)\n",
    "    accs = log_dict[key]['acc']\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key)\n",
    "\n",
    "# plt.title(\"Accuracy vs Training Steps for Different Layer Counts\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(\"Transformer\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "946d348ae136421da2b6d3f8cfb29902",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 5.2500, Accuracy: 0.0000\n",
      "1_layers - Step 32, Loss: 1.6328, Accuracy: 0.2969\n",
      "1_layers - Step 64, Loss: 1.7188, Accuracy: 0.2793\n",
      "1_layers - Step 96, Loss: 1.5000, Accuracy: 0.4902\n",
      "1_layers - Step 128, Loss: 1.6484, Accuracy: 0.3828\n",
      "1_layers - Step 160, Loss: 1.4844, Accuracy: 0.5000\n",
      "1_layers - Step 192, Loss: 1.5234, Accuracy: 0.4785\n",
      "1_layers - Step 224, Loss: 1.3359, Accuracy: 0.9160\n",
      "1_layers - Step 256, Loss: 1.2812, Accuracy: 0.9844\n",
      "1_layers - Step 288, Loss: 1.2422, Accuracy: 0.9941\n",
      "1_layers - Step 320, Loss: 1.2891, Accuracy: 0.9902\n",
      "1_layers - Step 352, Loss: 1.2891, Accuracy: 1.0000\n",
      "1_layers - Step 384, Loss: 1.2500, Accuracy: 1.0000\n",
      "1_layers - Step 416, Loss: 1.2656, Accuracy: 0.9980\n",
      "1_layers - Step 448, Loss: 1.2422, Accuracy: 0.9961\n",
      "1_layers - Step 480, Loss: 1.2656, Accuracy: 0.9980\n",
      "1_layers - Step 512, Loss: 1.2891, Accuracy: 1.0000\n",
      "1_layers - Step 544, Loss: 1.2812, Accuracy: 1.0000\n",
      "1_layers - Step 576, Loss: 1.2969, Accuracy: 1.0000\n",
      "1_layers - Step 608, Loss: 1.1797, Accuracy: 1.0000\n",
      "1_layers - Step 640, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 672, Loss: 1.2266, Accuracy: 1.0000\n",
      "1_layers - Step 704, Loss: 1.2891, Accuracy: 1.0000\n",
      "1_layers - Step 736, Loss: 1.2266, Accuracy: 0.9961\n",
      "1_layers - Step 768, Loss: 1.2812, Accuracy: 1.0000\n",
      "1_layers - Step 800, Loss: 1.2656, Accuracy: 1.0000\n",
      "1_layers - Step 832, Loss: 1.2109, Accuracy: 1.0000\n",
      "1_layers - Step 864, Loss: 1.2578, Accuracy: 0.9961\n",
      "1_layers - Step 896, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 928, Loss: 1.2422, Accuracy: 0.9980\n",
      "1_layers - Step 960, Loss: 1.1562, Accuracy: 1.0000\n",
      "1_layers - Step 992, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1024, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 1056, Loss: 1.1953, Accuracy: 1.0000\n",
      "1_layers - Step 1088, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1120, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1152, Loss: 1.2578, Accuracy: 0.9961\n",
      "1_layers - Step 1184, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1216, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1248, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 1280, Loss: 1.2266, Accuracy: 1.0000\n",
      "1_layers - Step 1312, Loss: 1.2109, Accuracy: 0.9980\n",
      "1_layers - Step 1344, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1376, Loss: 1.1953, Accuracy: 1.0000\n",
      "1_layers - Step 1408, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1440, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 1472, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1504, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1536, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1568, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1600, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1632, Loss: 1.1875, Accuracy: 1.0000\n",
      "1_layers - Step 1664, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1696, Loss: 1.1875, Accuracy: 1.0000\n",
      "1_layers - Step 1728, Loss: 1.2969, Accuracy: 1.0000\n",
      "1_layers - Step 1760, Loss: 1.1328, Accuracy: 1.0000\n",
      "1_layers - Step 1792, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1824, Loss: 1.1953, Accuracy: 1.0000\n",
      "1_layers - Step 1856, Loss: 1.1719, Accuracy: 1.0000\n",
      "1_layers - Step 1888, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 1920, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1952, Loss: 1.2109, Accuracy: 1.0000\n",
      "1_layers - Step 1984, Loss: 1.2578, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 2 layers ===\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "672493ba67e14d1a8948c830c27e2b49",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 5.1250, Accuracy: 0.0000\n",
      "2_layers - Step 32, Loss: 1.5547, Accuracy: 0.3848\n",
      "2_layers - Step 64, Loss: 1.5781, Accuracy: 0.3496\n",
      "2_layers - Step 96, Loss: 1.5547, Accuracy: 0.4668\n",
      "2_layers - Step 128, Loss: 1.5781, Accuracy: 0.3828\n",
      "2_layers - Step 160, Loss: 1.3906, Accuracy: 0.6992\n",
      "2_layers - Step 192, Loss: 1.2891, Accuracy: 0.9668\n",
      "2_layers - Step 224, Loss: 1.2500, Accuracy: 0.9863\n",
      "2_layers - Step 256, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 288, Loss: 1.2109, Accuracy: 0.9980\n",
      "2_layers - Step 320, Loss: 1.2734, Accuracy: 0.9922\n",
      "2_layers - Step 352, Loss: 1.2344, Accuracy: 0.9980\n",
      "2_layers - Step 384, Loss: 1.1641, Accuracy: 0.9980\n",
      "2_layers - Step 416, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 448, Loss: 1.2188, Accuracy: 1.0000\n",
      "2_layers - Step 480, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 512, Loss: 1.2266, Accuracy: 0.9941\n",
      "2_layers - Step 544, Loss: 1.2109, Accuracy: 1.0000\n",
      "2_layers - Step 576, Loss: 1.2891, Accuracy: 0.9980\n",
      "2_layers - Step 608, Loss: 1.2891, Accuracy: 1.0000\n",
      "2_layers - Step 640, Loss: 1.1875, Accuracy: 1.0000\n",
      "2_layers - Step 672, Loss: 1.1953, Accuracy: 0.9980\n",
      "2_layers - Step 704, Loss: 1.2109, Accuracy: 1.0000\n",
      "2_layers - Step 736, Loss: 1.2188, Accuracy: 1.0000\n",
      "2_layers - Step 768, Loss: 1.2500, Accuracy: 0.9902\n",
      "2_layers - Step 800, Loss: 1.1719, Accuracy: 1.0000\n",
      "2_layers - Step 832, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 864, Loss: 1.2812, Accuracy: 0.9980\n",
      "2_layers - Step 896, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 928, Loss: 1.1797, Accuracy: 0.9980\n",
      "2_layers - Step 960, Loss: 1.1797, Accuracy: 1.0000\n",
      "2_layers - Step 992, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1024, Loss: 1.3125, Accuracy: 1.0000\n",
      "2_layers - Step 1056, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 1088, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1120, Loss: 1.2812, Accuracy: 0.9961\n",
      "2_layers - Step 1152, Loss: 1.1953, Accuracy: 0.9980\n",
      "2_layers - Step 1184, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 1216, Loss: 1.2812, Accuracy: 1.0000\n",
      "2_layers - Step 1248, Loss: 1.1875, Accuracy: 1.0000\n",
      "2_layers - Step 1280, Loss: 1.2422, Accuracy: 0.9980\n",
      "2_layers - Step 1312, Loss: 1.1797, Accuracy: 1.0000\n",
      "2_layers - Step 1344, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 1376, Loss: 1.3047, Accuracy: 1.0000\n",
      "2_layers - Step 1408, Loss: 1.2578, Accuracy: 1.0000\n",
      "2_layers - Step 1440, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1472, Loss: 1.1875, Accuracy: 1.0000\n",
      "2_layers - Step 1504, Loss: 1.3281, Accuracy: 1.0000\n",
      "2_layers - Step 1536, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1568, Loss: 1.1641, Accuracy: 1.0000\n",
      "2_layers - Step 1600, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1632, Loss: 1.1953, Accuracy: 1.0000\n",
      "2_layers - Step 1664, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1696, Loss: 1.1953, Accuracy: 1.0000\n",
      "2_layers - Step 1728, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1792, Loss: 1.1797, Accuracy: 1.0000\n",
      "2_layers - Step 1824, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1856, Loss: 1.3516, Accuracy: 1.0000\n",
      "2_layers - Step 1888, Loss: 1.1562, Accuracy: 1.0000\n",
      "2_layers - Step 1920, Loss: 1.2578, Accuracy: 1.0000\n",
      "2_layers - Step 1952, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1984, Loss: 1.1875, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 4 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "670d7175808e462b9c33d4b9c9a65225",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.3750, Accuracy: 0.0039\n",
      "4_layers - Step 32, Loss: 1.5859, Accuracy: 0.3398\n",
      "4_layers - Step 64, Loss: 1.5703, Accuracy: 0.3379\n",
      "4_layers - Step 96, Loss: 1.5859, Accuracy: 0.3887\n",
      "4_layers - Step 128, Loss: 1.4766, Accuracy: 0.5156\n",
      "4_layers - Step 160, Loss: 1.3516, Accuracy: 0.8652\n",
      "4_layers - Step 192, Loss: 1.3281, Accuracy: 0.9590\n",
      "4_layers - Step 224, Loss: 1.2969, Accuracy: 0.9746\n",
      "4_layers - Step 256, Loss: 1.2969, Accuracy: 0.9961\n",
      "4_layers - Step 288, Loss: 1.2969, Accuracy: 0.9941\n",
      "4_layers - Step 320, Loss: 1.3672, Accuracy: 0.9980\n",
      "4_layers - Step 352, Loss: 1.2578, Accuracy: 0.9941\n",
      "4_layers - Step 384, Loss: 1.3125, Accuracy: 1.0000\n",
      "4_layers - Step 416, Loss: 1.2500, Accuracy: 0.9980\n",
      "4_layers - Step 448, Loss: 1.2266, Accuracy: 0.9980\n",
      "4_layers - Step 480, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 512, Loss: 1.2344, Accuracy: 0.9980\n",
      "4_layers - Step 544, Loss: 1.2656, Accuracy: 0.9980\n",
      "4_layers - Step 576, Loss: 1.3047, Accuracy: 1.0000\n",
      "4_layers - Step 608, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 640, Loss: 1.2109, Accuracy: 1.0000\n",
      "4_layers - Step 672, Loss: 1.1953, Accuracy: 0.9980\n",
      "4_layers - Step 704, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 736, Loss: 1.1953, Accuracy: 1.0000\n",
      "4_layers - Step 768, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 800, Loss: 1.2109, Accuracy: 0.9980\n",
      "4_layers - Step 832, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 864, Loss: 1.2109, Accuracy: 1.0000\n",
      "4_layers - Step 896, Loss: 1.3047, Accuracy: 1.0000\n",
      "4_layers - Step 928, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 960, Loss: 1.1953, Accuracy: 1.0000\n",
      "4_layers - Step 992, Loss: 1.1484, Accuracy: 1.0000\n",
      "4_layers - Step 1024, Loss: 1.1562, Accuracy: 1.0000\n",
      "4_layers - Step 1056, Loss: 1.2891, Accuracy: 1.0000\n",
      "4_layers - Step 1088, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1120, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1152, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 1184, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 1216, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1248, Loss: 1.1875, Accuracy: 1.0000\n",
      "4_layers - Step 1280, Loss: 1.1797, Accuracy: 1.0000\n",
      "4_layers - Step 1312, Loss: 1.1875, Accuracy: 1.0000\n",
      "4_layers - Step 1344, Loss: 1.2109, Accuracy: 1.0000\n",
      "4_layers - Step 1376, Loss: 1.2812, Accuracy: 1.0000\n",
      "4_layers - Step 1408, Loss: 1.3047, Accuracy: 1.0000\n",
      "4_layers - Step 1440, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1472, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1504, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1536, Loss: 1.2891, Accuracy: 1.0000\n",
      "4_layers - Step 1568, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1600, Loss: 1.3047, Accuracy: 1.0000\n",
      "4_layers - Step 1632, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1664, Loss: 1.3125, Accuracy: 1.0000\n",
      "4_layers - Step 1696, Loss: 1.1719, Accuracy: 1.0000\n",
      "4_layers - Step 1728, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1760, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 1792, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1824, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1856, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1888, Loss: 1.2891, Accuracy: 1.0000\n",
      "4_layers - Step 1920, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1952, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1984, Loss: 1.2266, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 8 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e9d2e7efaf804795ac91eade7b8ede2d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 5.0312, Accuracy: 0.0020\n",
      "8_layers - Step 32, Loss: 1.6562, Accuracy: 0.2617\n",
      "8_layers - Step 64, Loss: 1.6094, Accuracy: 0.3320\n",
      "8_layers - Step 96, Loss: 1.6016, Accuracy: 0.3711\n",
      "8_layers - Step 128, Loss: 1.5078, Accuracy: 0.5215\n",
      "8_layers - Step 160, Loss: 1.3203, Accuracy: 0.8770\n",
      "8_layers - Step 192, Loss: 1.3047, Accuracy: 0.9922\n",
      "8_layers - Step 224, Loss: 1.2109, Accuracy: 0.9922\n",
      "8_layers - Step 256, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 288, Loss: 1.2656, Accuracy: 0.9980\n",
      "8_layers - Step 320, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 352, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 384, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 416, Loss: 1.2031, Accuracy: 0.9980\n",
      "8_layers - Step 448, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 480, Loss: 1.2500, Accuracy: 0.9883\n",
      "8_layers - Step 512, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 544, Loss: 1.1406, Accuracy: 0.9980\n",
      "8_layers - Step 576, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 608, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 640, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 672, Loss: 1.2109, Accuracy: 0.9980\n",
      "8_layers - Step 704, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 736, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 768, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 800, Loss: 1.2109, Accuracy: 1.0000\n",
      "8_layers - Step 832, Loss: 1.1797, Accuracy: 1.0000\n",
      "8_layers - Step 864, Loss: 1.1641, Accuracy: 1.0000\n",
      "8_layers - Step 896, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 928, Loss: 1.1719, Accuracy: 1.0000\n",
      "8_layers - Step 960, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 992, Loss: 1.1484, Accuracy: 1.0000\n",
      "8_layers - Step 1024, Loss: 1.2578, Accuracy: 1.0000\n",
      "8_layers - Step 1056, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 1088, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 1120, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1152, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 1184, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1216, Loss: 1.1953, Accuracy: 1.0000\n",
      "8_layers - Step 1248, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1280, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1312, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1344, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1376, Loss: 1.1797, Accuracy: 1.0000\n",
      "8_layers - Step 1408, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1440, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1472, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1504, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 1536, Loss: 1.2891, Accuracy: 1.0000\n",
      "8_layers - Step 1568, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 1600, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1632, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 1664, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1696, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 1728, Loss: 1.1562, Accuracy: 1.0000\n",
      "8_layers - Step 1760, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 1792, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1824, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1856, Loss: 1.1719, Accuracy: 1.0000\n",
      "8_layers - Step 1888, Loss: 1.2891, Accuracy: 1.0000\n",
      "8_layers - Step 1920, Loss: 1.1562, Accuracy: 1.0000\n",
      "8_layers - Step 1952, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1984, Loss: 1.1953, Accuracy: 1.0000\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 58528 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'linear',\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "        # if epoch==400:\n",
    "        #     print('length come to 128')\n",
    "        #     dataset = S5Dataset(k=128, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "        # if epoch==500:\n",
    "        #     print('length come to 256')\n",
    "        #     dataset = S5Dataset(k=256, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "        # if epoch==600:\n",
    "        #     print('length come to 512')\n",
    "        #     dataset = S5Dataset(k=512, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "        # if epoch==700:\n",
    "        #     print('length come to 1024')\n",
    "        #     dataset = S5Dataset(k=1024, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54924 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 绘图\n",
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    accs = log_dict[key]['acc']\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key)\n",
    "\n",
    "# plt.title(\"Delta Transformer\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb00d97575d845c383e6ab659ac5d7f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 5.0938, Accuracy: 0.0000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 32, Loss: 1.6094, Accuracy: 0.3047\n",
      "1_layers - Step 64, Loss: 1.5938, Accuracy: 0.3633\n",
      "1_layers - Step 96, Loss: 1.5625, Accuracy: 0.4180\n",
      "1_layers - Step 128, Loss: 1.5938, Accuracy: 0.3633\n",
      "1_layers - Step 160, Loss: 1.5391, Accuracy: 0.3945\n",
      "1_layers - Step 192, Loss: 1.6250, Accuracy: 0.3965\n",
      "1_layers - Step 224, Loss: 1.4062, Accuracy: 0.8066\n",
      "1_layers - Step 256, Loss: 1.3203, Accuracy: 0.9609\n",
      "1_layers - Step 288, Loss: 1.3047, Accuracy: 0.9766\n",
      "1_layers - Step 320, Loss: 1.2344, Accuracy: 0.9902\n",
      "1_layers - Step 352, Loss: 1.2500, Accuracy: 0.9824\n",
      "1_layers - Step 384, Loss: 1.2500, Accuracy: 0.9980\n",
      "1_layers - Step 416, Loss: 1.2109, Accuracy: 1.0000\n",
      "1_layers - Step 448, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 480, Loss: 1.3359, Accuracy: 0.9883\n",
      "1_layers - Step 512, Loss: 1.1953, Accuracy: 1.0000\n",
      "1_layers - Step 544, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 576, Loss: 1.2969, Accuracy: 0.9883\n",
      "1_layers - Step 608, Loss: 1.3281, Accuracy: 1.0000\n",
      "1_layers - Step 640, Loss: 1.1953, Accuracy: 1.0000\n",
      "1_layers - Step 672, Loss: 1.2109, Accuracy: 0.9961\n",
      "1_layers - Step 704, Loss: 1.3125, Accuracy: 1.0000\n",
      "1_layers - Step 736, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 768, Loss: 1.2188, Accuracy: 0.9980\n",
      "1_layers - Step 800, Loss: 1.1875, Accuracy: 1.0000\n",
      "1_layers - Step 832, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 864, Loss: 1.2500, Accuracy: 1.0000\n",
      "1_layers - Step 896, Loss: 1.2266, Accuracy: 1.0000\n",
      "1_layers - Step 928, Loss: 1.2891, Accuracy: 1.0000\n",
      "1_layers - Step 960, Loss: 1.2109, Accuracy: 1.0000\n",
      "1_layers - Step 992, Loss: 1.1641, Accuracy: 1.0000\n",
      "1_layers - Step 1024, Loss: 1.2969, Accuracy: 0.9980\n",
      "1_layers - Step 1056, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1088, Loss: 1.2812, Accuracy: 1.0000\n",
      "1_layers - Step 1120, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 1152, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1184, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 1216, Loss: 1.1562, Accuracy: 1.0000\n",
      "1_layers - Step 1248, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1280, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 1312, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1344, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1376, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1408, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1440, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1472, Loss: 1.2500, Accuracy: 1.0000\n",
      "1_layers - Step 1504, Loss: 1.2500, Accuracy: 1.0000\n",
      "1_layers - Step 1536, Loss: 1.2031, Accuracy: 1.0000\n",
      "1_layers - Step 1568, Loss: 1.2969, Accuracy: 1.0000\n",
      "1_layers - Step 1600, Loss: 1.2812, Accuracy: 1.0000\n",
      "1_layers - Step 1632, Loss: 1.3203, Accuracy: 1.0000\n",
      "1_layers - Step 1664, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 1696, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1728, Loss: 1.2109, Accuracy: 1.0000\n",
      "1_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "1_layers - Step 1792, Loss: 1.2188, Accuracy: 1.0000\n",
      "1_layers - Step 1824, Loss: 1.1328, Accuracy: 1.0000\n",
      "1_layers - Step 1856, Loss: 1.1719, Accuracy: 1.0000\n",
      "1_layers - Step 1888, Loss: 1.2734, Accuracy: 1.0000\n",
      "1_layers - Step 1920, Loss: 1.2578, Accuracy: 1.0000\n",
      "1_layers - Step 1952, Loss: 1.2422, Accuracy: 1.0000\n",
      "1_layers - Step 1984, Loss: 1.2500, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 2 layers ===\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40583f497926409e89f55f476ae68800",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 4.8438, Accuracy: 0.0059\n",
      "2_layers - Step 32, Loss: 1.6094, Accuracy: 0.3008\n",
      "2_layers - Step 64, Loss: 1.6172, Accuracy: 0.3340\n",
      "2_layers - Step 96, Loss: 1.5469, Accuracy: 0.3984\n",
      "2_layers - Step 128, Loss: 1.5547, Accuracy: 0.4570\n",
      "2_layers - Step 160, Loss: 1.5000, Accuracy: 0.5527\n",
      "2_layers - Step 192, Loss: 1.2422, Accuracy: 0.9727\n",
      "2_layers - Step 224, Loss: 1.2734, Accuracy: 0.9980\n",
      "2_layers - Step 256, Loss: 1.2734, Accuracy: 0.9785\n",
      "2_layers - Step 288, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 320, Loss: 1.1953, Accuracy: 1.0000\n",
      "2_layers - Step 352, Loss: 1.3047, Accuracy: 1.0000\n",
      "2_layers - Step 384, Loss: 1.2578, Accuracy: 1.0000\n",
      "2_layers - Step 416, Loss: 1.2422, Accuracy: 0.9844\n",
      "2_layers - Step 448, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 480, Loss: 1.2578, Accuracy: 1.0000\n",
      "2_layers - Step 512, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 544, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 576, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 608, Loss: 1.2266, Accuracy: 0.9941\n",
      "2_layers - Step 640, Loss: 1.2969, Accuracy: 1.0000\n",
      "2_layers - Step 672, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 704, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 736, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 768, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 800, Loss: 1.2578, Accuracy: 0.9980\n",
      "2_layers - Step 832, Loss: 1.2812, Accuracy: 0.9980\n",
      "2_layers - Step 864, Loss: 1.2578, Accuracy: 0.9980\n",
      "2_layers - Step 896, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 928, Loss: 1.1641, Accuracy: 0.9961\n",
      "2_layers - Step 960, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 992, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1024, Loss: 1.2578, Accuracy: 1.0000\n",
      "2_layers - Step 1056, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 1088, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1120, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1152, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1184, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1216, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1248, Loss: 1.1484, Accuracy: 1.0000\n",
      "2_layers - Step 1280, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 1312, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1344, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 1376, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 1408, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1440, Loss: 1.1562, Accuracy: 1.0000\n",
      "2_layers - Step 1472, Loss: 1.1719, Accuracy: 1.0000\n",
      "2_layers - Step 1504, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1536, Loss: 1.1562, Accuracy: 1.0000\n",
      "2_layers - Step 1568, Loss: 1.3047, Accuracy: 1.0000\n",
      "2_layers - Step 1600, Loss: 1.2188, Accuracy: 1.0000\n",
      "2_layers - Step 1632, Loss: 1.1641, Accuracy: 1.0000\n",
      "2_layers - Step 1664, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 1696, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1728, Loss: 1.1562, Accuracy: 1.0000\n",
      "2_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1792, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1824, Loss: 1.3203, Accuracy: 1.0000\n",
      "2_layers - Step 1856, Loss: 1.2812, Accuracy: 1.0000\n",
      "2_layers - Step 1888, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1920, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1952, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1984, Loss: 1.2031, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 4 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1420700ac5bb49e9918bc0a6a0605b46",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.0000, Accuracy: 0.0039\n",
      "4_layers - Step 32, Loss: 1.6719, Accuracy: 0.2871\n",
      "4_layers - Step 64, Loss: 1.6719, Accuracy: 0.2871\n",
      "4_layers - Step 96, Loss: 1.4844, Accuracy: 0.4707\n",
      "4_layers - Step 128, Loss: 1.5625, Accuracy: 0.4785\n",
      "4_layers - Step 160, Loss: 1.3750, Accuracy: 0.6738\n",
      "4_layers - Step 192, Loss: 1.2812, Accuracy: 0.9062\n",
      "4_layers - Step 224, Loss: 1.2344, Accuracy: 0.9785\n",
      "4_layers - Step 256, Loss: 1.2812, Accuracy: 0.9824\n",
      "4_layers - Step 288, Loss: 1.2812, Accuracy: 0.9922\n",
      "4_layers - Step 320, Loss: 1.3203, Accuracy: 0.9980\n",
      "4_layers - Step 352, Loss: 1.2812, Accuracy: 1.0000\n",
      "4_layers - Step 384, Loss: 1.1875, Accuracy: 0.9980\n",
      "4_layers - Step 416, Loss: 1.2422, Accuracy: 0.9980\n",
      "4_layers - Step 448, Loss: 1.2656, Accuracy: 0.9922\n",
      "4_layers - Step 480, Loss: 1.2578, Accuracy: 0.9980\n",
      "4_layers - Step 512, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 544, Loss: 1.2344, Accuracy: 0.9961\n",
      "4_layers - Step 576, Loss: 1.3281, Accuracy: 0.9961\n",
      "4_layers - Step 608, Loss: 1.3047, Accuracy: 1.0000\n",
      "4_layers - Step 640, Loss: 1.1875, Accuracy: 1.0000\n",
      "4_layers - Step 672, Loss: 1.2734, Accuracy: 1.0000\n",
      "4_layers - Step 704, Loss: 1.2344, Accuracy: 0.9961\n",
      "4_layers - Step 736, Loss: 1.2891, Accuracy: 1.0000\n",
      "4_layers - Step 768, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 800, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 832, Loss: 1.2969, Accuracy: 1.0000\n",
      "4_layers - Step 864, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 896, Loss: 1.1797, Accuracy: 1.0000\n",
      "4_layers - Step 928, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 960, Loss: 1.2422, Accuracy: 0.9980\n",
      "4_layers - Step 992, Loss: 1.1797, Accuracy: 1.0000\n",
      "4_layers - Step 1024, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 1056, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1088, Loss: 1.1953, Accuracy: 1.0000\n",
      "4_layers - Step 1120, Loss: 1.2734, Accuracy: 1.0000\n",
      "4_layers - Step 1152, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1184, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1216, Loss: 1.1875, Accuracy: 1.0000\n",
      "4_layers - Step 1248, Loss: 1.2891, Accuracy: 1.0000\n",
      "4_layers - Step 1280, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1312, Loss: 1.1797, Accuracy: 1.0000\n",
      "4_layers - Step 1344, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1376, Loss: 1.2109, Accuracy: 1.0000\n",
      "4_layers - Step 1408, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1440, Loss: 1.1953, Accuracy: 1.0000\n",
      "4_layers - Step 1472, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1504, Loss: 1.1719, Accuracy: 1.0000\n",
      "4_layers - Step 1536, Loss: 1.3359, Accuracy: 1.0000\n",
      "4_layers - Step 1568, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 1600, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1632, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1664, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 1696, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1728, Loss: 1.1875, Accuracy: 1.0000\n",
      "4_layers - Step 1760, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1792, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1824, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 1856, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1888, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1920, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1952, Loss: 1.1250, Accuracy: 1.0000\n",
      "4_layers - Step 1984, Loss: 1.2266, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 8 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84198407778342169a0ebdd2e8de6541",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 4.9375, Accuracy: 0.0000\n",
      "8_layers - Step 32, Loss: 1.5938, Accuracy: 0.2773\n",
      "8_layers - Step 64, Loss: 1.6172, Accuracy: 0.3145\n",
      "8_layers - Step 96, Loss: 1.6250, Accuracy: 0.3848\n",
      "8_layers - Step 128, Loss: 1.5469, Accuracy: 0.3984\n",
      "8_layers - Step 160, Loss: 1.4453, Accuracy: 0.6465\n",
      "8_layers - Step 192, Loss: 1.2422, Accuracy: 0.9551\n",
      "8_layers - Step 224, Loss: 1.2422, Accuracy: 0.9961\n",
      "8_layers - Step 256, Loss: 1.2422, Accuracy: 0.9941\n",
      "8_layers - Step 288, Loss: 1.2969, Accuracy: 0.9980\n",
      "8_layers - Step 320, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 352, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 384, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 416, Loss: 1.3281, Accuracy: 1.0000\n",
      "8_layers - Step 448, Loss: 1.2891, Accuracy: 1.0000\n",
      "8_layers - Step 480, Loss: 1.2578, Accuracy: 0.9980\n",
      "8_layers - Step 512, Loss: 1.2969, Accuracy: 1.0000\n",
      "8_layers - Step 544, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 576, Loss: 1.2734, Accuracy: 0.9980\n",
      "8_layers - Step 608, Loss: 1.2578, Accuracy: 0.9961\n",
      "8_layers - Step 640, Loss: 1.2969, Accuracy: 0.9980\n",
      "8_layers - Step 672, Loss: 1.2734, Accuracy: 0.9980\n",
      "8_layers - Step 704, Loss: 1.1875, Accuracy: 0.9980\n",
      "8_layers - Step 736, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 768, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 800, Loss: 1.2891, Accuracy: 1.0000\n",
      "8_layers - Step 832, Loss: 1.1953, Accuracy: 1.0000\n",
      "8_layers - Step 864, Loss: 1.1797, Accuracy: 1.0000\n",
      "8_layers - Step 896, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 928, Loss: 1.1641, Accuracy: 1.0000\n",
      "8_layers - Step 960, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 992, Loss: 1.2578, Accuracy: 0.9980\n",
      "8_layers - Step 1024, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1056, Loss: 1.1797, Accuracy: 1.0000\n",
      "8_layers - Step 1088, Loss: 1.1562, Accuracy: 1.0000\n",
      "8_layers - Step 1120, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1152, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1184, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1216, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 1248, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 1280, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1312, Loss: 1.1875, Accuracy: 1.0000\n",
      "8_layers - Step 1344, Loss: 1.2578, Accuracy: 1.0000\n",
      "8_layers - Step 1376, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1408, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 1440, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1472, Loss: 1.1875, Accuracy: 1.0000\n",
      "8_layers - Step 1504, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1536, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1568, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 1600, Loss: 1.1484, Accuracy: 1.0000\n",
      "8_layers - Step 1632, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1664, Loss: 1.1641, Accuracy: 1.0000\n",
      "8_layers - Step 1696, Loss: 1.1797, Accuracy: 1.0000\n",
      "8_layers - Step 1728, Loss: 1.1250, Accuracy: 1.0000\n",
      "8_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 1792, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 1824, Loss: 1.2578, Accuracy: 1.0000\n",
      "8_layers - Step 1856, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1888, Loss: 1.1875, Accuracy: 1.0000\n",
      "8_layers - Step 1920, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 1952, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1984, Loss: 1.2500, Accuracy: 1.0000\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57596 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'round',  #######设置kernel,选择有'linear','round','softmax_v1','softmax_v2','relu'\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison_round.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54048 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# 绘图\n",
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    accs = log_dict[key]['acc']\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key)\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison_round.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac05a8a07e9b4c2daf479d8d7a501e1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 4.9688, Accuracy: 0.0098\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 32, Loss: 1.6016, Accuracy: 0.3223\n",
      "1_layers - Step 64, Loss: 1.6016, Accuracy: 0.3262\n",
      "1_layers - Step 96, Loss: 1.5859, Accuracy: 0.3613\n",
      "1_layers - Step 128, Loss: 1.6250, Accuracy: 0.3535\n",
      "1_layers - Step 160, Loss: 1.5547, Accuracy: 0.3750\n",
      "1_layers - Step 192, Loss: 1.6328, Accuracy: 0.3262\n",
      "1_layers - Step 224, Loss: 1.5391, Accuracy: 0.4395\n",
      "1_layers - Step 256, Loss: 1.5547, Accuracy: 0.4531\n",
      "1_layers - Step 288, Loss: 1.5391, Accuracy: 0.4844\n",
      "1_layers - Step 320, Loss: 1.4688, Accuracy: 0.4902\n",
      "1_layers - Step 352, Loss: 1.4375, Accuracy: 0.4609\n",
      "1_layers - Step 384, Loss: 1.5469, Accuracy: 0.4160\n",
      "1_layers - Step 416, Loss: 1.5781, Accuracy: 0.4023\n",
      "1_layers - Step 448, Loss: 1.5078, Accuracy: 0.4531\n",
      "1_layers - Step 480, Loss: 1.5859, Accuracy: 0.4004\n",
      "1_layers - Step 512, Loss: 1.5547, Accuracy: 0.4473\n",
      "1_layers - Step 544, Loss: 1.6016, Accuracy: 0.4121\n",
      "1_layers - Step 576, Loss: 1.5234, Accuracy: 0.4277\n",
      "1_layers - Step 608, Loss: 1.5234, Accuracy: 0.4785\n",
      "1_layers - Step 640, Loss: 1.4688, Accuracy: 0.4941\n",
      "1_layers - Step 672, Loss: 1.5000, Accuracy: 0.4863\n",
      "1_layers - Step 704, Loss: 1.4609, Accuracy: 0.5176\n",
      "1_layers - Step 736, Loss: 1.5469, Accuracy: 0.4668\n",
      "1_layers - Step 768, Loss: 1.4609, Accuracy: 0.5156\n",
      "1_layers - Step 800, Loss: 1.4922, Accuracy: 0.5117\n",
      "1_layers - Step 832, Loss: 1.4453, Accuracy: 0.5410\n",
      "1_layers - Step 864, Loss: 1.5078, Accuracy: 0.4707\n",
      "1_layers - Step 896, Loss: 1.4297, Accuracy: 0.5605\n",
      "1_layers - Step 928, Loss: 1.4531, Accuracy: 0.5586\n",
      "1_layers - Step 960, Loss: 1.4844, Accuracy: 0.5469\n",
      "1_layers - Step 992, Loss: 1.4375, Accuracy: 0.5684\n",
      "1_layers - Step 1024, Loss: 1.4297, Accuracy: 0.6035\n",
      "1_layers - Step 1056, Loss: 1.5469, Accuracy: 0.5430\n",
      "1_layers - Step 1088, Loss: 1.4375, Accuracy: 0.6367\n",
      "1_layers - Step 1120, Loss: 1.4453, Accuracy: 0.6309\n",
      "1_layers - Step 1152, Loss: 1.4688, Accuracy: 0.5957\n",
      "1_layers - Step 1184, Loss: 1.4609, Accuracy: 0.5801\n",
      "1_layers - Step 1216, Loss: 1.5000, Accuracy: 0.5742\n",
      "1_layers - Step 1248, Loss: 1.4219, Accuracy: 0.6172\n",
      "1_layers - Step 1280, Loss: 1.4531, Accuracy: 0.5840\n",
      "1_layers - Step 1312, Loss: 1.4531, Accuracy: 0.6211\n",
      "1_layers - Step 1344, Loss: 1.3984, Accuracy: 0.6758\n",
      "1_layers - Step 1376, Loss: 1.4688, Accuracy: 0.6621\n",
      "1_layers - Step 1408, Loss: 1.4219, Accuracy: 0.6445\n",
      "1_layers - Step 1440, Loss: 1.3984, Accuracy: 0.6504\n",
      "1_layers - Step 1472, Loss: 1.4219, Accuracy: 0.5957\n",
      "1_layers - Step 1504, Loss: 1.3906, Accuracy: 0.6562\n",
      "1_layers - Step 1536, Loss: 1.3984, Accuracy: 0.6309\n",
      "1_layers - Step 1568, Loss: 1.4219, Accuracy: 0.6797\n",
      "1_layers - Step 1600, Loss: 1.4688, Accuracy: 0.6621\n",
      "1_layers - Step 1632, Loss: 1.4141, Accuracy: 0.6875\n",
      "1_layers - Step 1664, Loss: 1.4141, Accuracy: 0.6777\n",
      "1_layers - Step 1696, Loss: 1.4375, Accuracy: 0.6816\n",
      "1_layers - Step 1728, Loss: 1.3906, Accuracy: 0.6660\n",
      "1_layers - Step 1760, Loss: 1.4219, Accuracy: 0.6484\n",
      "1_layers - Step 1792, Loss: 1.4219, Accuracy: 0.6465\n",
      "1_layers - Step 1824, Loss: 1.4062, Accuracy: 0.6699\n",
      "1_layers - Step 1856, Loss: 1.4453, Accuracy: 0.6699\n",
      "1_layers - Step 1888, Loss: 1.4219, Accuracy: 0.7246\n",
      "1_layers - Step 1920, Loss: 1.4141, Accuracy: 0.7266\n",
      "1_layers - Step 1952, Loss: 1.3672, Accuracy: 0.6660\n",
      "1_layers - Step 1984, Loss: 1.3984, Accuracy: 0.6543\n",
      "\n",
      "=== Training with 2 layers ===\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e8702eb8d31463a88529bae4b776d9a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 4.9375, Accuracy: 0.0020\n",
      "2_layers - Step 32, Loss: 1.6016, Accuracy: 0.3496\n",
      "2_layers - Step 64, Loss: 1.6875, Accuracy: 0.2871\n",
      "2_layers - Step 96, Loss: 1.5391, Accuracy: 0.4121\n",
      "2_layers - Step 128, Loss: 1.6719, Accuracy: 0.3887\n",
      "2_layers - Step 160, Loss: 1.6016, Accuracy: 0.3965\n",
      "2_layers - Step 192, Loss: 1.5234, Accuracy: 0.4531\n",
      "2_layers - Step 224, Loss: 1.5469, Accuracy: 0.4297\n",
      "2_layers - Step 256, Loss: 1.4531, Accuracy: 0.5137\n",
      "2_layers - Step 288, Loss: 1.5078, Accuracy: 0.4609\n",
      "2_layers - Step 320, Loss: 1.5234, Accuracy: 0.4395\n",
      "2_layers - Step 352, Loss: 1.5469, Accuracy: 0.4199\n",
      "2_layers - Step 384, Loss: 1.5703, Accuracy: 0.4492\n",
      "2_layers - Step 416, Loss: 1.4922, Accuracy: 0.5117\n",
      "2_layers - Step 448, Loss: 1.5234, Accuracy: 0.4531\n",
      "2_layers - Step 480, Loss: 1.4766, Accuracy: 0.5137\n",
      "2_layers - Step 512, Loss: 1.5156, Accuracy: 0.5234\n",
      "2_layers - Step 544, Loss: 1.4922, Accuracy: 0.5723\n",
      "2_layers - Step 576, Loss: 1.4609, Accuracy: 0.5527\n",
      "2_layers - Step 608, Loss: 1.5000, Accuracy: 0.5156\n",
      "2_layers - Step 640, Loss: 1.4922, Accuracy: 0.5645\n",
      "2_layers - Step 672, Loss: 1.4531, Accuracy: 0.5977\n",
      "2_layers - Step 704, Loss: 1.4062, Accuracy: 0.6230\n",
      "2_layers - Step 736, Loss: 1.3828, Accuracy: 0.7520\n",
      "2_layers - Step 768, Loss: 1.2656, Accuracy: 0.8711\n",
      "2_layers - Step 800, Loss: 1.3047, Accuracy: 0.9277\n",
      "2_layers - Step 832, Loss: 1.3125, Accuracy: 0.9648\n",
      "2_layers - Step 864, Loss: 1.3047, Accuracy: 0.9688\n",
      "2_layers - Step 896, Loss: 1.2734, Accuracy: 0.9844\n",
      "2_layers - Step 928, Loss: 1.2656, Accuracy: 0.9902\n",
      "2_layers - Step 960, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 992, Loss: 1.2891, Accuracy: 1.0000\n",
      "2_layers - Step 1024, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1056, Loss: 1.1641, Accuracy: 1.0000\n",
      "2_layers - Step 1088, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1120, Loss: 1.2109, Accuracy: 0.9941\n",
      "2_layers - Step 1152, Loss: 1.2891, Accuracy: 1.0000\n",
      "2_layers - Step 1184, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 1216, Loss: 1.1797, Accuracy: 1.0000\n",
      "2_layers - Step 1248, Loss: 1.1875, Accuracy: 0.9980\n",
      "2_layers - Step 1280, Loss: 1.1875, Accuracy: 1.0000\n",
      "2_layers - Step 1312, Loss: 1.3281, Accuracy: 1.0000\n",
      "2_layers - Step 1344, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1376, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 1408, Loss: 1.2734, Accuracy: 1.0000\n",
      "2_layers - Step 1440, Loss: 1.2891, Accuracy: 1.0000\n",
      "2_layers - Step 1472, Loss: 1.2266, Accuracy: 1.0000\n",
      "2_layers - Step 1504, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1536, Loss: 1.1953, Accuracy: 1.0000\n",
      "2_layers - Step 1568, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 1600, Loss: 1.1875, Accuracy: 1.0000\n",
      "2_layers - Step 1632, Loss: 1.2031, Accuracy: 1.0000\n",
      "2_layers - Step 1664, Loss: 1.2500, Accuracy: 1.0000\n",
      "2_layers - Step 1696, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1728, Loss: 1.2109, Accuracy: 1.0000\n",
      "2_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "2_layers - Step 1792, Loss: 1.1797, Accuracy: 1.0000\n",
      "2_layers - Step 1824, Loss: 1.2656, Accuracy: 1.0000\n",
      "2_layers - Step 1856, Loss: 1.2188, Accuracy: 1.0000\n",
      "2_layers - Step 1888, Loss: 1.3125, Accuracy: 1.0000\n",
      "2_layers - Step 1920, Loss: 1.2891, Accuracy: 1.0000\n",
      "2_layers - Step 1952, Loss: 1.2422, Accuracy: 1.0000\n",
      "2_layers - Step 1984, Loss: 1.2109, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 4 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e248fe0d02334b1690dee4e3ca715baf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.0938, Accuracy: 0.0078\n",
      "4_layers - Step 32, Loss: 1.5625, Accuracy: 0.3203\n",
      "4_layers - Step 64, Loss: 1.5859, Accuracy: 0.3730\n",
      "4_layers - Step 96, Loss: 1.5703, Accuracy: 0.4004\n",
      "4_layers - Step 128, Loss: 1.5391, Accuracy: 0.4434\n",
      "4_layers - Step 160, Loss: 1.5859, Accuracy: 0.3633\n",
      "4_layers - Step 192, Loss: 1.5391, Accuracy: 0.4199\n",
      "4_layers - Step 224, Loss: 1.6328, Accuracy: 0.3848\n",
      "4_layers - Step 256, Loss: 1.5547, Accuracy: 0.4258\n",
      "4_layers - Step 288, Loss: 1.5469, Accuracy: 0.4238\n",
      "4_layers - Step 320, Loss: 1.5234, Accuracy: 0.4551\n",
      "4_layers - Step 352, Loss: 1.5000, Accuracy: 0.4531\n",
      "4_layers - Step 384, Loss: 1.5234, Accuracy: 0.4727\n",
      "4_layers - Step 416, Loss: 1.5469, Accuracy: 0.4922\n",
      "4_layers - Step 448, Loss: 1.5000, Accuracy: 0.4863\n",
      "4_layers - Step 480, Loss: 1.4453, Accuracy: 0.4922\n",
      "4_layers - Step 512, Loss: 1.5156, Accuracy: 0.5586\n",
      "4_layers - Step 544, Loss: 1.4531, Accuracy: 0.5938\n",
      "4_layers - Step 576, Loss: 1.3828, Accuracy: 0.7109\n",
      "4_layers - Step 608, Loss: 1.4922, Accuracy: 0.6602\n",
      "4_layers - Step 640, Loss: 1.3047, Accuracy: 0.8945\n",
      "4_layers - Step 672, Loss: 1.3203, Accuracy: 0.9316\n",
      "4_layers - Step 704, Loss: 1.2500, Accuracy: 0.9766\n",
      "4_layers - Step 736, Loss: 1.2656, Accuracy: 0.9453\n",
      "4_layers - Step 768, Loss: 1.2109, Accuracy: 0.9766\n",
      "4_layers - Step 800, Loss: 1.1484, Accuracy: 0.9922\n",
      "4_layers - Step 832, Loss: 1.2812, Accuracy: 0.9902\n",
      "4_layers - Step 864, Loss: 1.2188, Accuracy: 0.9902\n",
      "4_layers - Step 896, Loss: 1.2656, Accuracy: 0.9727\n",
      "4_layers - Step 928, Loss: 1.2578, Accuracy: 0.9980\n",
      "4_layers - Step 960, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 992, Loss: 1.2109, Accuracy: 0.9980\n",
      "4_layers - Step 1024, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1056, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1088, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1120, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1152, Loss: 1.2031, Accuracy: 1.0000\n",
      "4_layers - Step 1184, Loss: 1.2969, Accuracy: 1.0000\n",
      "4_layers - Step 1216, Loss: 1.2734, Accuracy: 1.0000\n",
      "4_layers - Step 1248, Loss: 1.1797, Accuracy: 1.0000\n",
      "4_layers - Step 1280, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1312, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1344, Loss: 1.2734, Accuracy: 1.0000\n",
      "4_layers - Step 1376, Loss: 1.2578, Accuracy: 1.0000\n",
      "4_layers - Step 1408, Loss: 1.2109, Accuracy: 1.0000\n",
      "4_layers - Step 1440, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1472, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1504, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1536, Loss: 1.1719, Accuracy: 1.0000\n",
      "4_layers - Step 1568, Loss: 1.2812, Accuracy: 1.0000\n",
      "4_layers - Step 1600, Loss: 1.3125, Accuracy: 1.0000\n",
      "4_layers - Step 1632, Loss: 1.2500, Accuracy: 1.0000\n",
      "4_layers - Step 1664, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1696, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1728, Loss: 1.2266, Accuracy: 1.0000\n",
      "4_layers - Step 1760, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1792, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1824, Loss: 1.2188, Accuracy: 1.0000\n",
      "4_layers - Step 1856, Loss: 1.2344, Accuracy: 1.0000\n",
      "4_layers - Step 1888, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1920, Loss: 1.2656, Accuracy: 1.0000\n",
      "4_layers - Step 1952, Loss: 1.2422, Accuracy: 1.0000\n",
      "4_layers - Step 1984, Loss: 1.2344, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 8 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11b89b222aff4faa9e4138ed471dc3bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 4.9375, Accuracy: 0.0098\n",
      "8_layers - Step 32, Loss: 1.6406, Accuracy: 0.2676\n",
      "8_layers - Step 64, Loss: 1.6094, Accuracy: 0.3496\n",
      "8_layers - Step 96, Loss: 1.7031, Accuracy: 0.3203\n",
      "8_layers - Step 128, Loss: 1.5703, Accuracy: 0.4219\n",
      "8_layers - Step 160, Loss: 1.5625, Accuracy: 0.3867\n",
      "8_layers - Step 192, Loss: 1.5781, Accuracy: 0.4023\n",
      "8_layers - Step 224, Loss: 1.5391, Accuracy: 0.4648\n",
      "8_layers - Step 256, Loss: 1.4609, Accuracy: 0.4746\n",
      "8_layers - Step 288, Loss: 1.5938, Accuracy: 0.3730\n",
      "8_layers - Step 320, Loss: 1.5234, Accuracy: 0.5078\n",
      "8_layers - Step 352, Loss: 1.5781, Accuracy: 0.4531\n",
      "8_layers - Step 384, Loss: 1.5078, Accuracy: 0.4707\n",
      "8_layers - Step 416, Loss: 1.5078, Accuracy: 0.4727\n",
      "8_layers - Step 448, Loss: 1.5078, Accuracy: 0.4805\n",
      "8_layers - Step 480, Loss: 1.5000, Accuracy: 0.5293\n",
      "8_layers - Step 512, Loss: 1.4688, Accuracy: 0.5859\n",
      "8_layers - Step 544, Loss: 1.3750, Accuracy: 0.6582\n",
      "8_layers - Step 576, Loss: 1.5000, Accuracy: 0.6660\n",
      "8_layers - Step 608, Loss: 1.4219, Accuracy: 0.8691\n",
      "8_layers - Step 640, Loss: 1.2734, Accuracy: 0.9121\n",
      "8_layers - Step 672, Loss: 1.3281, Accuracy: 0.9336\n",
      "8_layers - Step 704, Loss: 1.2734, Accuracy: 0.9707\n",
      "8_layers - Step 736, Loss: 1.3047, Accuracy: 0.9648\n",
      "8_layers - Step 768, Loss: 1.2188, Accuracy: 0.9805\n",
      "8_layers - Step 800, Loss: 1.2344, Accuracy: 0.9883\n",
      "8_layers - Step 832, Loss: 1.2109, Accuracy: 0.9922\n",
      "8_layers - Step 864, Loss: 1.2734, Accuracy: 0.9961\n",
      "8_layers - Step 896, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 928, Loss: 1.1719, Accuracy: 1.0000\n",
      "8_layers - Step 960, Loss: 1.2031, Accuracy: 1.0000\n",
      "8_layers - Step 992, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1024, Loss: 1.2109, Accuracy: 1.0000\n",
      "8_layers - Step 1056, Loss: 1.1875, Accuracy: 1.0000\n",
      "8_layers - Step 1088, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1120, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1152, Loss: 1.2891, Accuracy: 1.0000\n",
      "8_layers - Step 1184, Loss: 1.2578, Accuracy: 1.0000\n",
      "8_layers - Step 1216, Loss: 1.3047, Accuracy: 1.0000\n",
      "8_layers - Step 1248, Loss: 1.2422, Accuracy: 1.0000\n",
      "8_layers - Step 1280, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1312, Loss: 1.1875, Accuracy: 1.0000\n",
      "8_layers - Step 1344, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1376, Loss: 1.2188, Accuracy: 1.0000\n",
      "8_layers - Step 1408, Loss: 1.2344, Accuracy: 1.0000\n",
      "8_layers - Step 1440, Loss: 1.1953, Accuracy: 1.0000\n",
      "8_layers - Step 1472, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1504, Loss: 1.2812, Accuracy: 1.0000\n",
      "8_layers - Step 1536, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1568, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1600, Loss: 1.1641, Accuracy: 1.0000\n",
      "8_layers - Step 1632, Loss: 1.2656, Accuracy: 1.0000\n",
      "8_layers - Step 1664, Loss: 1.2734, Accuracy: 1.0000\n",
      "8_layers - Step 1696, Loss: 1.1953, Accuracy: 1.0000\n",
      "8_layers - Step 1728, Loss: 1.1641, Accuracy: 1.0000\n",
      "8_layers - Step 1760, Loss: 1.3281, Accuracy: 1.0000\n",
      "8_layers - Step 1792, Loss: 1.2109, Accuracy: 1.0000\n",
      "8_layers - Step 1824, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1856, Loss: 1.2109, Accuracy: 1.0000\n",
      "8_layers - Step 1888, Loss: 1.2266, Accuracy: 1.0000\n",
      "8_layers - Step 1920, Loss: 1.2500, Accuracy: 1.0000\n",
      "8_layers - Step 1952, Loss: 1.1953, Accuracy: 1.0000\n",
      "8_layers - Step 1984, Loss: 1.1953, Accuracy: 1.0000\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 110556 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'relu',  #######设置kernel,选择有'linear','round','softmax_v1','softmax_v2','relu'\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison_round.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-05-14 11:37:44.149393: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-05-14 11:37:44.184105: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8922479a2c1a483fa12e2e2735e5bf23",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 5.0000, Accuracy: 0.0117\n",
      "1_layers - Step 32, Loss: 1.5625, Accuracy: 0.3594\n",
      "1_layers - Step 64, Loss: 1.5781, Accuracy: 0.4121\n",
      "1_layers - Step 96, Loss: 1.6641, Accuracy: 0.3379\n",
      "1_layers - Step 128, Loss: 1.5938, Accuracy: 0.3574\n",
      "1_layers - Step 160, Loss: 1.5703, Accuracy: 0.3711\n",
      "1_layers - Step 192, Loss: 1.5547, Accuracy: 0.4316\n",
      "1_layers - Step 224, Loss: 1.5781, Accuracy: 0.3770\n",
      "1_layers - Step 256, Loss: 1.5312, Accuracy: 0.3965\n",
      "1_layers - Step 288, Loss: 1.6016, Accuracy: 0.3340\n",
      "1_layers - Step 320, Loss: 1.5469, Accuracy: 0.3867\n",
      "1_layers - Step 352, Loss: 1.5469, Accuracy: 0.4336\n",
      "1_layers - Step 384, Loss: 1.5234, Accuracy: 0.4531\n",
      "1_layers - Step 416, Loss: 1.5156, Accuracy: 0.4277\n",
      "1_layers - Step 448, Loss: 1.4609, Accuracy: 0.4668\n",
      "1_layers - Step 480, Loss: 1.5703, Accuracy: 0.4082\n",
      "1_layers - Step 512, Loss: 1.5391, Accuracy: 0.4629\n",
      "1_layers - Step 544, Loss: 1.4453, Accuracy: 0.5078\n",
      "1_layers - Step 576, Loss: 1.5312, Accuracy: 0.5059\n",
      "1_layers - Step 608, Loss: 1.4609, Accuracy: 0.6484\n",
      "1_layers - Step 640, Loss: 1.3984, Accuracy: 0.6895\n",
      "1_layers - Step 672, Loss: 1.4922, Accuracy: 0.6133\n",
      "1_layers - Step 704, Loss: 1.4453, Accuracy: 0.6875\n",
      "1_layers - Step 736, Loss: 1.3828, Accuracy: 0.7031\n",
      "1_layers - Step 768, Loss: 1.3906, Accuracy: 0.7812\n",
      "1_layers - Step 800, Loss: 1.4531, Accuracy: 0.7070\n",
      "1_layers - Step 832, Loss: 1.3594, Accuracy: 0.7656\n",
      "1_layers - Step 864, Loss: 1.4375, Accuracy: 0.7227\n",
      "1_layers - Step 896, Loss: 1.3828, Accuracy: 0.7129\n",
      "1_layers - Step 928, Loss: 1.4531, Accuracy: 0.7207\n",
      "1_layers - Step 960, Loss: 1.3750, Accuracy: 0.7324\n",
      "1_layers - Step 992, Loss: 1.4219, Accuracy: 0.7617\n",
      "1_layers - Step 1024, Loss: 1.3594, Accuracy: 0.7930\n",
      "1_layers - Step 1056, Loss: 1.3281, Accuracy: 0.7988\n",
      "1_layers - Step 1088, Loss: 1.3047, Accuracy: 0.7480\n",
      "1_layers - Step 1120, Loss: 1.3125, Accuracy: 0.7910\n",
      "1_layers - Step 1152, Loss: 1.2812, Accuracy: 0.8164\n",
      "1_layers - Step 1184, Loss: 1.3672, Accuracy: 0.7676\n",
      "1_layers - Step 1216, Loss: 1.3438, Accuracy: 0.8008\n",
      "1_layers - Step 1248, Loss: 1.3828, Accuracy: 0.8145\n",
      "1_layers - Step 1280, Loss: 1.3516, Accuracy: 0.7988\n",
      "1_layers - Step 1312, Loss: 1.3125, Accuracy: 0.7930\n",
      "1_layers - Step 1344, Loss: 1.3828, Accuracy: 0.7656\n",
      "1_layers - Step 1376, Loss: 1.3438, Accuracy: 0.7832\n",
      "1_layers - Step 1408, Loss: 1.2969, Accuracy: 0.8164\n",
      "1_layers - Step 1440, Loss: 1.3594, Accuracy: 0.8164\n",
      "1_layers - Step 1472, Loss: 1.3281, Accuracy: 0.7949\n",
      "1_layers - Step 1504, Loss: 1.2734, Accuracy: 0.8184\n",
      "1_layers - Step 1536, Loss: 1.3750, Accuracy: 0.7910\n",
      "1_layers - Step 1568, Loss: 1.3359, Accuracy: 0.8145\n",
      "1_layers - Step 1600, Loss: 1.3594, Accuracy: 0.7988\n",
      "1_layers - Step 1632, Loss: 1.2734, Accuracy: 0.8203\n",
      "1_layers - Step 1664, Loss: 1.3516, Accuracy: 0.8184\n",
      "1_layers - Step 1696, Loss: 1.3047, Accuracy: 0.7832\n",
      "1_layers - Step 1728, Loss: 1.3359, Accuracy: 0.8164\n",
      "1_layers - Step 1760, Loss: 1.3281, Accuracy: 0.7832\n",
      "1_layers - Step 1792, Loss: 1.3516, Accuracy: 0.8438\n",
      "1_layers - Step 1824, Loss: 1.3594, Accuracy: 0.7969\n",
      "1_layers - Step 1856, Loss: 1.3594, Accuracy: 0.7832\n",
      "1_layers - Step 1888, Loss: 1.2812, Accuracy: 0.7871\n",
      "1_layers - Step 1920, Loss: 1.3594, Accuracy: 0.7949\n",
      "1_layers - Step 1952, Loss: 1.3281, Accuracy: 0.8184\n",
      "1_layers - Step 1984, Loss: 1.3750, Accuracy: 0.8438\n",
      "\n",
      "=== Training with 2 layers ===\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2eee41a8b25e447498bac972a8d346b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 5.1562, Accuracy: 0.0137\n",
      "2_layers - Step 32, Loss: 1.6172, Accuracy: 0.3105\n",
      "2_layers - Step 64, Loss: 1.6641, Accuracy: 0.2793\n",
      "2_layers - Step 96, Loss: 1.6328, Accuracy: 0.3145\n",
      "2_layers - Step 128, Loss: 1.6016, Accuracy: 0.3789\n",
      "2_layers - Step 160, Loss: 1.5938, Accuracy: 0.3320\n",
      "2_layers - Step 192, Loss: 1.5938, Accuracy: 0.3359\n",
      "2_layers - Step 224, Loss: 1.6016, Accuracy: 0.3477\n",
      "2_layers - Step 256, Loss: 1.5391, Accuracy: 0.3867\n",
      "2_layers - Step 288, Loss: 1.5781, Accuracy: 0.4082\n",
      "2_layers - Step 320, Loss: 1.5625, Accuracy: 0.4238\n",
      "2_layers - Step 352, Loss: 1.5547, Accuracy: 0.4473\n",
      "2_layers - Step 384, Loss: 1.5000, Accuracy: 0.4902\n",
      "2_layers - Step 416, Loss: 1.5469, Accuracy: 0.4883\n",
      "2_layers - Step 448, Loss: 1.4844, Accuracy: 0.4727\n",
      "2_layers - Step 480, Loss: 1.4375, Accuracy: 0.5938\n",
      "2_layers - Step 512, Loss: 1.4688, Accuracy: 0.5957\n",
      "2_layers - Step 544, Loss: 1.4688, Accuracy: 0.5938\n",
      "2_layers - Step 576, Loss: 1.4688, Accuracy: 0.6973\n",
      "2_layers - Step 608, Loss: 1.4453, Accuracy: 0.7109\n",
      "2_layers - Step 640, Loss: 1.3984, Accuracy: 0.7695\n",
      "2_layers - Step 672, Loss: 1.3750, Accuracy: 0.7852\n",
      "2_layers - Step 704, Loss: 1.3672, Accuracy: 0.7461\n",
      "2_layers - Step 736, Loss: 1.4375, Accuracy: 0.7617\n",
      "2_layers - Step 768, Loss: 1.4062, Accuracy: 0.7520\n",
      "2_layers - Step 800, Loss: 1.2891, Accuracy: 0.8008\n",
      "2_layers - Step 832, Loss: 1.4062, Accuracy: 0.7383\n",
      "2_layers - Step 864, Loss: 1.3281, Accuracy: 0.7754\n",
      "2_layers - Step 896, Loss: 1.4062, Accuracy: 0.7168\n",
      "2_layers - Step 928, Loss: 1.3906, Accuracy: 0.7363\n",
      "2_layers - Step 960, Loss: 1.4219, Accuracy: 0.7402\n",
      "2_layers - Step 992, Loss: 1.4219, Accuracy: 0.7188\n",
      "2_layers - Step 1024, Loss: 1.3359, Accuracy: 0.7637\n",
      "2_layers - Step 1056, Loss: 1.3828, Accuracy: 0.7461\n",
      "2_layers - Step 1088, Loss: 1.3594, Accuracy: 0.7754\n",
      "2_layers - Step 1120, Loss: 1.3828, Accuracy: 0.7754\n",
      "2_layers - Step 1152, Loss: 1.3750, Accuracy: 0.7969\n",
      "2_layers - Step 1184, Loss: 1.3672, Accuracy: 0.7715\n",
      "2_layers - Step 1216, Loss: 1.3125, Accuracy: 0.8535\n",
      "2_layers - Step 1248, Loss: 1.3203, Accuracy: 0.8086\n",
      "2_layers - Step 1280, Loss: 1.3516, Accuracy: 0.7812\n",
      "2_layers - Step 1312, Loss: 1.3438, Accuracy: 0.8066\n",
      "2_layers - Step 1344, Loss: 1.3516, Accuracy: 0.8301\n",
      "2_layers - Step 1376, Loss: 1.3594, Accuracy: 0.8223\n",
      "2_layers - Step 1408, Loss: 1.3203, Accuracy: 0.8672\n",
      "2_layers - Step 1440, Loss: 1.3438, Accuracy: 0.8066\n",
      "2_layers - Step 1472, Loss: 1.3516, Accuracy: 0.8125\n",
      "2_layers - Step 1504, Loss: 1.2734, Accuracy: 0.8379\n",
      "2_layers - Step 1536, Loss: 1.2500, Accuracy: 0.8398\n",
      "2_layers - Step 1568, Loss: 1.2969, Accuracy: 0.8594\n",
      "2_layers - Step 1600, Loss: 1.2812, Accuracy: 0.8574\n",
      "2_layers - Step 1632, Loss: 1.2578, Accuracy: 0.8184\n",
      "2_layers - Step 1664, Loss: 1.2969, Accuracy: 0.8301\n",
      "2_layers - Step 1696, Loss: 1.2422, Accuracy: 0.8105\n",
      "2_layers - Step 1728, Loss: 1.3359, Accuracy: 0.8496\n",
      "2_layers - Step 1760, Loss: 1.3438, Accuracy: 0.8477\n",
      "2_layers - Step 1792, Loss: 1.3359, Accuracy: 0.8633\n",
      "2_layers - Step 1824, Loss: 1.3359, Accuracy: 0.8164\n",
      "2_layers - Step 1856, Loss: 1.2969, Accuracy: 0.8730\n",
      "2_layers - Step 1888, Loss: 1.4062, Accuracy: 0.7988\n",
      "2_layers - Step 1920, Loss: 1.3125, Accuracy: 0.8047\n",
      "2_layers - Step 1952, Loss: 1.3828, Accuracy: 0.8008\n",
      "2_layers - Step 1984, Loss: 1.2812, Accuracy: 0.8516\n",
      "\n",
      "=== Training with 4 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04cd67b7484b475ba02fb3f056f2ce06",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.1250, Accuracy: 0.0039\n",
      "4_layers - Step 32, Loss: 1.6094, Accuracy: 0.3477\n",
      "4_layers - Step 64, Loss: 1.5859, Accuracy: 0.3496\n",
      "4_layers - Step 96, Loss: 1.6094, Accuracy: 0.3301\n",
      "4_layers - Step 128, Loss: 1.5859, Accuracy: 0.4121\n",
      "4_layers - Step 160, Loss: 1.5859, Accuracy: 0.2734\n",
      "4_layers - Step 192, Loss: 1.5703, Accuracy: 0.3242\n",
      "4_layers - Step 224, Loss: 1.5938, Accuracy: 0.3516\n",
      "4_layers - Step 256, Loss: 1.5625, Accuracy: 0.3789\n",
      "4_layers - Step 288, Loss: 1.5312, Accuracy: 0.4316\n",
      "4_layers - Step 320, Loss: 1.5859, Accuracy: 0.3691\n",
      "4_layers - Step 352, Loss: 1.5391, Accuracy: 0.4336\n",
      "4_layers - Step 384, Loss: 1.5000, Accuracy: 0.4727\n",
      "4_layers - Step 416, Loss: 1.4531, Accuracy: 0.5156\n",
      "4_layers - Step 448, Loss: 1.5938, Accuracy: 0.3906\n",
      "4_layers - Step 480, Loss: 1.5469, Accuracy: 0.4805\n",
      "4_layers - Step 512, Loss: 1.4531, Accuracy: 0.5879\n",
      "4_layers - Step 544, Loss: 1.5000, Accuracy: 0.4570\n",
      "4_layers - Step 576, Loss: 1.4141, Accuracy: 0.6387\n",
      "4_layers - Step 608, Loss: 1.4062, Accuracy: 0.6484\n",
      "4_layers - Step 640, Loss: 1.4844, Accuracy: 0.5918\n",
      "4_layers - Step 672, Loss: 1.4453, Accuracy: 0.6543\n",
      "4_layers - Step 704, Loss: 1.3672, Accuracy: 0.7402\n",
      "4_layers - Step 736, Loss: 1.3828, Accuracy: 0.6914\n",
      "4_layers - Step 768, Loss: 1.4219, Accuracy: 0.7363\n",
      "4_layers - Step 800, Loss: 1.3984, Accuracy: 0.7324\n",
      "4_layers - Step 832, Loss: 1.4062, Accuracy: 0.7324\n",
      "4_layers - Step 864, Loss: 1.3047, Accuracy: 0.8145\n",
      "4_layers - Step 896, Loss: 1.3984, Accuracy: 0.7363\n",
      "4_layers - Step 928, Loss: 1.3750, Accuracy: 0.7539\n",
      "4_layers - Step 960, Loss: 1.4609, Accuracy: 0.6953\n",
      "4_layers - Step 992, Loss: 1.3203, Accuracy: 0.7715\n",
      "4_layers - Step 1024, Loss: 1.3203, Accuracy: 0.7871\n",
      "4_layers - Step 1056, Loss: 1.4062, Accuracy: 0.7754\n",
      "4_layers - Step 1088, Loss: 1.3906, Accuracy: 0.7695\n",
      "4_layers - Step 1120, Loss: 1.3438, Accuracy: 0.8203\n",
      "4_layers - Step 1152, Loss: 1.3672, Accuracy: 0.7695\n",
      "4_layers - Step 1184, Loss: 1.2734, Accuracy: 0.8398\n",
      "4_layers - Step 1216, Loss: 1.2812, Accuracy: 0.8359\n",
      "4_layers - Step 1248, Loss: 1.3125, Accuracy: 0.8184\n",
      "4_layers - Step 1280, Loss: 1.2500, Accuracy: 0.8301\n",
      "4_layers - Step 1312, Loss: 1.3125, Accuracy: 0.8242\n",
      "4_layers - Step 1344, Loss: 1.3438, Accuracy: 0.8418\n",
      "4_layers - Step 1376, Loss: 1.3984, Accuracy: 0.7305\n",
      "4_layers - Step 1408, Loss: 1.3359, Accuracy: 0.8359\n",
      "4_layers - Step 1440, Loss: 1.3672, Accuracy: 0.7910\n",
      "4_layers - Step 1472, Loss: 1.3594, Accuracy: 0.7891\n",
      "4_layers - Step 1504, Loss: 1.3281, Accuracy: 0.8223\n",
      "4_layers - Step 1536, Loss: 1.3672, Accuracy: 0.8164\n",
      "4_layers - Step 1568, Loss: 1.2734, Accuracy: 0.8281\n",
      "4_layers - Step 1600, Loss: 1.2812, Accuracy: 0.8848\n",
      "4_layers - Step 1632, Loss: 1.3203, Accuracy: 0.8242\n",
      "4_layers - Step 1664, Loss: 1.3906, Accuracy: 0.8125\n",
      "4_layers - Step 1696, Loss: 1.3281, Accuracy: 0.8086\n",
      "4_layers - Step 1728, Loss: 1.2891, Accuracy: 0.8477\n",
      "4_layers - Step 1760, Loss: 1.3516, Accuracy: 0.8242\n",
      "4_layers - Step 1792, Loss: 1.3438, Accuracy: 0.8398\n",
      "4_layers - Step 1824, Loss: 1.3906, Accuracy: 0.7969\n",
      "4_layers - Step 1856, Loss: 1.3359, Accuracy: 0.8398\n",
      "4_layers - Step 1888, Loss: 1.3750, Accuracy: 0.8418\n",
      "4_layers - Step 1920, Loss: 1.3125, Accuracy: 0.8555\n",
      "4_layers - Step 1952, Loss: 1.2969, Accuracy: 0.8418\n",
      "4_layers - Step 1984, Loss: 1.3047, Accuracy: 0.8594\n",
      "\n",
      "=== Training with 8 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d950112e61a43468f459ec7ed8d6d5a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 5.0312, Accuracy: 0.0059\n",
      "8_layers - Step 32, Loss: 1.5938, Accuracy: 0.3555\n",
      "8_layers - Step 64, Loss: 1.5625, Accuracy: 0.3809\n",
      "8_layers - Step 96, Loss: 1.6250, Accuracy: 0.3594\n",
      "8_layers - Step 128, Loss: 1.7188, Accuracy: 0.2422\n",
      "8_layers - Step 160, Loss: 1.6406, Accuracy: 0.1992\n",
      "8_layers - Step 192, Loss: 1.6328, Accuracy: 0.2891\n",
      "8_layers - Step 224, Loss: 1.5781, Accuracy: 0.3438\n",
      "8_layers - Step 256, Loss: 1.6016, Accuracy: 0.3926\n",
      "8_layers - Step 288, Loss: 1.5859, Accuracy: 0.3398\n",
      "8_layers - Step 320, Loss: 1.5234, Accuracy: 0.4180\n",
      "8_layers - Step 352, Loss: 1.5156, Accuracy: 0.4277\n",
      "8_layers - Step 384, Loss: 1.5625, Accuracy: 0.3848\n",
      "8_layers - Step 416, Loss: 1.5547, Accuracy: 0.3984\n",
      "8_layers - Step 448, Loss: 1.5781, Accuracy: 0.3887\n",
      "8_layers - Step 480, Loss: 1.5469, Accuracy: 0.4258\n",
      "8_layers - Step 512, Loss: 1.5156, Accuracy: 0.4863\n",
      "8_layers - Step 544, Loss: 1.5703, Accuracy: 0.3828\n",
      "8_layers - Step 576, Loss: 1.5625, Accuracy: 0.4453\n",
      "8_layers - Step 608, Loss: 1.5547, Accuracy: 0.4473\n",
      "8_layers - Step 640, Loss: 1.4844, Accuracy: 0.4551\n",
      "8_layers - Step 672, Loss: 1.5625, Accuracy: 0.3809\n",
      "8_layers - Step 704, Loss: 1.5703, Accuracy: 0.3672\n",
      "8_layers - Step 736, Loss: 1.5078, Accuracy: 0.4551\n",
      "8_layers - Step 768, Loss: 1.5234, Accuracy: 0.4629\n",
      "8_layers - Step 800, Loss: 1.5625, Accuracy: 0.3887\n",
      "8_layers - Step 832, Loss: 1.4844, Accuracy: 0.4512\n",
      "8_layers - Step 864, Loss: 1.5781, Accuracy: 0.3516\n",
      "8_layers - Step 896, Loss: 1.5547, Accuracy: 0.4199\n",
      "8_layers - Step 928, Loss: 1.5547, Accuracy: 0.3750\n",
      "8_layers - Step 960, Loss: 1.5156, Accuracy: 0.4414\n",
      "8_layers - Step 992, Loss: 1.5234, Accuracy: 0.3867\n",
      "8_layers - Step 1024, Loss: 1.4766, Accuracy: 0.5332\n",
      "8_layers - Step 1056, Loss: 1.4844, Accuracy: 0.4902\n",
      "8_layers - Step 1088, Loss: 1.5156, Accuracy: 0.4609\n",
      "8_layers - Step 1120, Loss: 1.5234, Accuracy: 0.4707\n",
      "8_layers - Step 1152, Loss: 1.5703, Accuracy: 0.4629\n",
      "8_layers - Step 1184, Loss: 1.4766, Accuracy: 0.5078\n",
      "8_layers - Step 1216, Loss: 1.5234, Accuracy: 0.4668\n",
      "8_layers - Step 1248, Loss: 1.5078, Accuracy: 0.5371\n",
      "8_layers - Step 1280, Loss: 1.4688, Accuracy: 0.5684\n",
      "8_layers - Step 1312, Loss: 1.5156, Accuracy: 0.4785\n",
      "8_layers - Step 1344, Loss: 1.5078, Accuracy: 0.5410\n",
      "8_layers - Step 1376, Loss: 1.4375, Accuracy: 0.5391\n",
      "8_layers - Step 1408, Loss: 1.5234, Accuracy: 0.4609\n",
      "8_layers - Step 1440, Loss: 1.5078, Accuracy: 0.4219\n",
      "8_layers - Step 1472, Loss: 1.4922, Accuracy: 0.5645\n",
      "8_layers - Step 1504, Loss: 1.4531, Accuracy: 0.5098\n",
      "8_layers - Step 1536, Loss: 1.5234, Accuracy: 0.4746\n",
      "8_layers - Step 1568, Loss: 1.4922, Accuracy: 0.5352\n",
      "8_layers - Step 1600, Loss: 1.4609, Accuracy: 0.5840\n",
      "8_layers - Step 1632, Loss: 1.4766, Accuracy: 0.5469\n",
      "8_layers - Step 1664, Loss: 1.4844, Accuracy: 0.5879\n",
      "8_layers - Step 1696, Loss: 1.4609, Accuracy: 0.5156\n",
      "8_layers - Step 1728, Loss: 1.4375, Accuracy: 0.5488\n",
      "8_layers - Step 1760, Loss: 1.4453, Accuracy: 0.5703\n",
      "8_layers - Step 1792, Loss: 1.5078, Accuracy: 0.5312\n",
      "8_layers - Step 1824, Loss: 1.4531, Accuracy: 0.6094\n",
      "8_layers - Step 1856, Loss: 1.4844, Accuracy: 0.6035\n",
      "8_layers - Step 1888, Loss: 1.4453, Accuracy: 0.6094\n",
      "8_layers - Step 1920, Loss: 1.4531, Accuracy: 0.5918\n",
      "8_layers - Step 1952, Loss: 1.4609, Accuracy: 0.5664\n",
      "8_layers - Step 1984, Loss: 1.4609, Accuracy: 0.5605\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 149204 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'exp',  #######设置kernel,选择有'linear','round','softmax_v1','softmax_v2','relu'\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison_exp.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 145460 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# 绘图\n",
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    accs = log_dict[key]['acc']\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key)\n",
    "    \n",
    "# plt.title(\"Delta Transformer\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison_exp.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43cfddb3a76145ffbb70a052782b93e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 1_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1_layers - Step 0, Loss: 5.0938, Accuracy: 0.0000\n",
      "1_layers - Step 32, Loss: 1.5469, Accuracy: 0.3457\n",
      "1_layers - Step 64, Loss: 1.6875, Accuracy: 0.3066\n",
      "1_layers - Step 96, Loss: 1.6250, Accuracy: 0.3359\n",
      "1_layers - Step 128, Loss: 1.6172, Accuracy: 0.3574\n",
      "1_layers - Step 160, Loss: 1.5469, Accuracy: 0.4141\n",
      "1_layers - Step 192, Loss: 1.5547, Accuracy: 0.3711\n",
      "1_layers - Step 224, Loss: 1.5625, Accuracy: 0.3633\n",
      "1_layers - Step 256, Loss: 1.5859, Accuracy: 0.3340\n",
      "1_layers - Step 288, Loss: 1.5547, Accuracy: 0.3730\n",
      "1_layers - Step 320, Loss: 1.5312, Accuracy: 0.3848\n",
      "1_layers - Step 352, Loss: 1.5938, Accuracy: 0.3809\n",
      "1_layers - Step 384, Loss: 1.5469, Accuracy: 0.3730\n",
      "1_layers - Step 416, Loss: 1.6094, Accuracy: 0.4043\n",
      "1_layers - Step 448, Loss: 1.5156, Accuracy: 0.4082\n",
      "1_layers - Step 480, Loss: 1.5703, Accuracy: 0.3789\n",
      "1_layers - Step 512, Loss: 1.5234, Accuracy: 0.3965\n",
      "1_layers - Step 544, Loss: 1.5312, Accuracy: 0.3828\n",
      "1_layers - Step 576, Loss: 1.5547, Accuracy: 0.3242\n",
      "1_layers - Step 608, Loss: 1.5312, Accuracy: 0.3984\n",
      "1_layers - Step 640, Loss: 1.5391, Accuracy: 0.3945\n",
      "1_layers - Step 672, Loss: 1.4688, Accuracy: 0.3965\n",
      "1_layers - Step 704, Loss: 1.5078, Accuracy: 0.4375\n",
      "1_layers - Step 736, Loss: 1.5625, Accuracy: 0.3320\n",
      "1_layers - Step 768, Loss: 1.5469, Accuracy: 0.4004\n",
      "1_layers - Step 800, Loss: 1.5312, Accuracy: 0.4043\n",
      "1_layers - Step 832, Loss: 1.5547, Accuracy: 0.3770\n",
      "1_layers - Step 864, Loss: 1.5234, Accuracy: 0.4043\n",
      "1_layers - Step 896, Loss: 1.5078, Accuracy: 0.4121\n",
      "1_layers - Step 928, Loss: 1.5469, Accuracy: 0.4082\n",
      "1_layers - Step 960, Loss: 1.5312, Accuracy: 0.3867\n",
      "1_layers - Step 992, Loss: 1.5000, Accuracy: 0.4395\n",
      "1_layers - Step 1024, Loss: 1.5312, Accuracy: 0.3770\n",
      "1_layers - Step 1056, Loss: 1.5625, Accuracy: 0.4141\n",
      "1_layers - Step 1088, Loss: 1.5391, Accuracy: 0.3652\n",
      "1_layers - Step 1120, Loss: 1.5469, Accuracy: 0.3945\n",
      "1_layers - Step 1152, Loss: 1.5078, Accuracy: 0.4023\n",
      "1_layers - Step 1184, Loss: 1.5000, Accuracy: 0.4180\n",
      "1_layers - Step 1216, Loss: 1.5234, Accuracy: 0.3887\n",
      "1_layers - Step 1248, Loss: 1.5156, Accuracy: 0.4004\n",
      "1_layers - Step 1280, Loss: 1.4844, Accuracy: 0.4062\n",
      "1_layers - Step 1312, Loss: 1.5625, Accuracy: 0.3770\n",
      "1_layers - Step 1344, Loss: 1.5234, Accuracy: 0.3906\n",
      "1_layers - Step 1376, Loss: 1.5000, Accuracy: 0.3887\n",
      "1_layers - Step 1408, Loss: 1.5234, Accuracy: 0.3848\n",
      "1_layers - Step 1440, Loss: 1.5234, Accuracy: 0.3672\n",
      "1_layers - Step 1472, Loss: 1.5156, Accuracy: 0.4316\n",
      "1_layers - Step 1504, Loss: 1.5156, Accuracy: 0.4082\n",
      "1_layers - Step 1536, Loss: 1.5234, Accuracy: 0.4219\n",
      "1_layers - Step 1568, Loss: 1.4844, Accuracy: 0.4160\n",
      "1_layers - Step 1600, Loss: 1.4531, Accuracy: 0.4531\n",
      "1_layers - Step 1632, Loss: 1.5781, Accuracy: 0.3242\n",
      "1_layers - Step 1664, Loss: 1.5000, Accuracy: 0.4043\n",
      "1_layers - Step 1696, Loss: 1.4375, Accuracy: 0.4707\n",
      "1_layers - Step 1728, Loss: 1.5312, Accuracy: 0.4082\n",
      "1_layers - Step 1760, Loss: 1.5391, Accuracy: 0.4160\n",
      "1_layers - Step 1792, Loss: 1.4766, Accuracy: 0.4492\n",
      "1_layers - Step 1824, Loss: 1.5547, Accuracy: 0.4102\n",
      "1_layers - Step 1856, Loss: 1.5234, Accuracy: 0.3945\n",
      "1_layers - Step 1888, Loss: 1.4844, Accuracy: 0.4766\n",
      "1_layers - Step 1920, Loss: 1.5156, Accuracy: 0.3887\n",
      "1_layers - Step 1952, Loss: 1.5312, Accuracy: 0.4023\n",
      "1_layers - Step 1984, Loss: 1.4688, Accuracy: 0.4648\n",
      "\n",
      "=== Training with 2 layers ===\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7927bd52b0584cdcb3191cfb1e417add",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 2_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2_layers - Step 0, Loss: 4.6562, Accuracy: 0.0391\n",
      "2_layers - Step 32, Loss: 1.6172, Accuracy: 0.3086\n",
      "2_layers - Step 64, Loss: 1.6094, Accuracy: 0.3008\n",
      "2_layers - Step 96, Loss: 1.6328, Accuracy: 0.3633\n",
      "2_layers - Step 128, Loss: 1.5625, Accuracy: 0.3438\n",
      "2_layers - Step 160, Loss: 1.5781, Accuracy: 0.3477\n",
      "2_layers - Step 192, Loss: 1.5938, Accuracy: 0.3535\n",
      "2_layers - Step 224, Loss: 1.6016, Accuracy: 0.3418\n",
      "2_layers - Step 256, Loss: 1.5547, Accuracy: 0.3574\n",
      "2_layers - Step 288, Loss: 1.5703, Accuracy: 0.3477\n",
      "2_layers - Step 320, Loss: 1.5156, Accuracy: 0.3809\n",
      "2_layers - Step 352, Loss: 1.5469, Accuracy: 0.3594\n",
      "2_layers - Step 384, Loss: 1.5859, Accuracy: 0.3613\n",
      "2_layers - Step 416, Loss: 1.5234, Accuracy: 0.4043\n",
      "2_layers - Step 448, Loss: 1.5859, Accuracy: 0.3477\n",
      "2_layers - Step 480, Loss: 1.5078, Accuracy: 0.4277\n",
      "2_layers - Step 512, Loss: 1.6094, Accuracy: 0.3535\n",
      "2_layers - Step 544, Loss: 1.5859, Accuracy: 0.4062\n",
      "2_layers - Step 576, Loss: 1.5625, Accuracy: 0.4336\n",
      "2_layers - Step 608, Loss: 1.5000, Accuracy: 0.4297\n",
      "2_layers - Step 640, Loss: 1.5234, Accuracy: 0.4121\n",
      "2_layers - Step 672, Loss: 1.5469, Accuracy: 0.3867\n",
      "2_layers - Step 704, Loss: 1.4922, Accuracy: 0.4570\n",
      "2_layers - Step 736, Loss: 1.5781, Accuracy: 0.3555\n",
      "2_layers - Step 768, Loss: 1.4922, Accuracy: 0.4199\n",
      "2_layers - Step 800, Loss: 1.5703, Accuracy: 0.3926\n",
      "2_layers - Step 832, Loss: 1.5156, Accuracy: 0.3906\n",
      "2_layers - Step 864, Loss: 1.4844, Accuracy: 0.4688\n",
      "2_layers - Step 896, Loss: 1.5703, Accuracy: 0.3418\n",
      "2_layers - Step 928, Loss: 1.5312, Accuracy: 0.4219\n",
      "2_layers - Step 960, Loss: 1.5234, Accuracy: 0.4199\n",
      "2_layers - Step 992, Loss: 1.5312, Accuracy: 0.4277\n",
      "2_layers - Step 1024, Loss: 1.5234, Accuracy: 0.4883\n",
      "2_layers - Step 1056, Loss: 1.5547, Accuracy: 0.4297\n",
      "2_layers - Step 1088, Loss: 1.4922, Accuracy: 0.4492\n",
      "2_layers - Step 1120, Loss: 1.5234, Accuracy: 0.4043\n",
      "2_layers - Step 1152, Loss: 1.5078, Accuracy: 0.5254\n",
      "2_layers - Step 1184, Loss: 1.5469, Accuracy: 0.4277\n",
      "2_layers - Step 1216, Loss: 1.4922, Accuracy: 0.4512\n",
      "2_layers - Step 1248, Loss: 1.4688, Accuracy: 0.4746\n",
      "2_layers - Step 1280, Loss: 1.4844, Accuracy: 0.5117\n",
      "2_layers - Step 1312, Loss: 1.4922, Accuracy: 0.4727\n",
      "2_layers - Step 1344, Loss: 1.5469, Accuracy: 0.4004\n",
      "2_layers - Step 1376, Loss: 1.4375, Accuracy: 0.5293\n",
      "2_layers - Step 1408, Loss: 1.5078, Accuracy: 0.4668\n",
      "2_layers - Step 1440, Loss: 1.5312, Accuracy: 0.4746\n",
      "2_layers - Step 1472, Loss: 1.4922, Accuracy: 0.4746\n",
      "2_layers - Step 1504, Loss: 1.4219, Accuracy: 0.5293\n",
      "2_layers - Step 1536, Loss: 1.5234, Accuracy: 0.4648\n",
      "2_layers - Step 1568, Loss: 1.4688, Accuracy: 0.5195\n",
      "2_layers - Step 1600, Loss: 1.5078, Accuracy: 0.4570\n",
      "2_layers - Step 1632, Loss: 1.4922, Accuracy: 0.4980\n",
      "2_layers - Step 1664, Loss: 1.4453, Accuracy: 0.5137\n",
      "2_layers - Step 1696, Loss: 1.5078, Accuracy: 0.4609\n",
      "2_layers - Step 1728, Loss: 1.4922, Accuracy: 0.5430\n",
      "2_layers - Step 1760, Loss: 1.4688, Accuracy: 0.5293\n",
      "2_layers - Step 1792, Loss: 1.4531, Accuracy: 0.5488\n",
      "2_layers - Step 1824, Loss: 1.4766, Accuracy: 0.5254\n",
      "2_layers - Step 1856, Loss: 1.4453, Accuracy: 0.5625\n",
      "2_layers - Step 1888, Loss: 1.5000, Accuracy: 0.5215\n",
      "2_layers - Step 1920, Loss: 1.4688, Accuracy: 0.4766\n",
      "2_layers - Step 1952, Loss: 1.4453, Accuracy: 0.5273\n",
      "2_layers - Step 1984, Loss: 1.4688, Accuracy: 0.5078\n",
      "\n",
      "=== Training with 4 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7dacd9bca0be4eeca5445bd7682848b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 4_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4_layers - Step 0, Loss: 5.0625, Accuracy: 0.0039\n",
      "4_layers - Step 32, Loss: 1.6250, Accuracy: 0.2949\n",
      "4_layers - Step 64, Loss: 1.5469, Accuracy: 0.3438\n",
      "4_layers - Step 96, Loss: 1.6016, Accuracy: 0.3320\n",
      "4_layers - Step 128, Loss: 1.6562, Accuracy: 0.3828\n",
      "4_layers - Step 160, Loss: 1.6484, Accuracy: 0.3398\n",
      "4_layers - Step 192, Loss: 1.4766, Accuracy: 0.4746\n",
      "4_layers - Step 224, Loss: 1.4844, Accuracy: 0.4629\n",
      "4_layers - Step 256, Loss: 1.5312, Accuracy: 0.5078\n",
      "4_layers - Step 288, Loss: 1.5781, Accuracy: 0.4180\n",
      "4_layers - Step 320, Loss: 1.5078, Accuracy: 0.5352\n",
      "4_layers - Step 352, Loss: 1.5312, Accuracy: 0.4434\n",
      "4_layers - Step 384, Loss: 1.4844, Accuracy: 0.5137\n",
      "4_layers - Step 416, Loss: 1.4141, Accuracy: 0.6094\n",
      "4_layers - Step 448, Loss: 1.4688, Accuracy: 0.5547\n",
      "4_layers - Step 480, Loss: 1.5078, Accuracy: 0.5156\n",
      "4_layers - Step 512, Loss: 1.4375, Accuracy: 0.5840\n",
      "4_layers - Step 544, Loss: 1.5078, Accuracy: 0.5254\n",
      "4_layers - Step 576, Loss: 1.5156, Accuracy: 0.5332\n",
      "4_layers - Step 608, Loss: 1.4531, Accuracy: 0.5801\n",
      "4_layers - Step 640, Loss: 1.4297, Accuracy: 0.6133\n",
      "4_layers - Step 672, Loss: 1.4453, Accuracy: 0.5703\n",
      "4_layers - Step 704, Loss: 1.4219, Accuracy: 0.5840\n",
      "4_layers - Step 736, Loss: 1.4219, Accuracy: 0.5820\n",
      "4_layers - Step 768, Loss: 1.4609, Accuracy: 0.6016\n",
      "4_layers - Step 800, Loss: 1.4453, Accuracy: 0.5586\n",
      "4_layers - Step 832, Loss: 1.4141, Accuracy: 0.6113\n",
      "4_layers - Step 864, Loss: 1.4062, Accuracy: 0.6660\n",
      "4_layers - Step 896, Loss: 1.4141, Accuracy: 0.6211\n",
      "4_layers - Step 928, Loss: 1.4688, Accuracy: 0.5918\n",
      "4_layers - Step 960, Loss: 1.4453, Accuracy: 0.6211\n",
      "4_layers - Step 992, Loss: 1.4219, Accuracy: 0.6211\n",
      "4_layers - Step 1024, Loss: 1.3672, Accuracy: 0.6719\n",
      "4_layers - Step 1056, Loss: 1.4219, Accuracy: 0.6660\n",
      "4_layers - Step 1088, Loss: 1.3984, Accuracy: 0.6406\n",
      "4_layers - Step 1120, Loss: 1.4219, Accuracy: 0.6406\n",
      "4_layers - Step 1152, Loss: 1.3828, Accuracy: 0.6445\n",
      "4_layers - Step 1184, Loss: 1.4844, Accuracy: 0.5996\n",
      "4_layers - Step 1216, Loss: 1.4531, Accuracy: 0.6055\n",
      "4_layers - Step 1248, Loss: 1.4141, Accuracy: 0.5566\n",
      "4_layers - Step 1280, Loss: 1.4766, Accuracy: 0.6152\n",
      "4_layers - Step 1312, Loss: 1.4219, Accuracy: 0.6777\n",
      "4_layers - Step 1344, Loss: 1.4219, Accuracy: 0.6719\n",
      "4_layers - Step 1376, Loss: 1.4375, Accuracy: 0.6504\n",
      "4_layers - Step 1408, Loss: 1.4062, Accuracy: 0.6816\n",
      "4_layers - Step 1440, Loss: 1.4219, Accuracy: 0.6699\n",
      "4_layers - Step 1472, Loss: 1.3984, Accuracy: 0.6621\n",
      "4_layers - Step 1504, Loss: 1.4141, Accuracy: 0.6641\n",
      "4_layers - Step 1536, Loss: 1.3438, Accuracy: 0.7051\n",
      "4_layers - Step 1568, Loss: 1.4453, Accuracy: 0.7129\n",
      "4_layers - Step 1600, Loss: 1.4219, Accuracy: 0.6582\n",
      "4_layers - Step 1632, Loss: 1.3828, Accuracy: 0.6699\n",
      "4_layers - Step 1664, Loss: 1.3047, Accuracy: 0.6973\n",
      "4_layers - Step 1696, Loss: 1.3828, Accuracy: 0.6758\n",
      "4_layers - Step 1728, Loss: 1.3750, Accuracy: 0.6465\n",
      "4_layers - Step 1760, Loss: 1.3984, Accuracy: 0.6406\n",
      "4_layers - Step 1792, Loss: 1.3047, Accuracy: 0.6953\n",
      "4_layers - Step 1824, Loss: 1.3047, Accuracy: 0.6934\n",
      "4_layers - Step 1856, Loss: 1.4141, Accuracy: 0.6992\n",
      "4_layers - Step 1888, Loss: 1.3984, Accuracy: 0.6914\n",
      "4_layers - Step 1920, Loss: 1.4531, Accuracy: 0.6465\n",
      "4_layers - Step 1952, Loss: 1.3672, Accuracy: 0.6523\n",
      "4_layers - Step 1984, Loss: 1.3984, Accuracy: 0.6641\n",
      "\n",
      "=== Training with 8 layers ===\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "112ed158ab6141ba8a3f95f54ff48e7f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 8_layers:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8_layers - Step 0, Loss: 5.0312, Accuracy: 0.0059\n",
      "8_layers - Step 32, Loss: 1.5547, Accuracy: 0.4062\n",
      "8_layers - Step 64, Loss: 1.6484, Accuracy: 0.3105\n",
      "8_layers - Step 96, Loss: 1.6016, Accuracy: 0.3457\n",
      "8_layers - Step 128, Loss: 1.6016, Accuracy: 0.3340\n",
      "8_layers - Step 160, Loss: 1.5547, Accuracy: 0.3750\n",
      "8_layers - Step 192, Loss: 1.5312, Accuracy: 0.3652\n",
      "8_layers - Step 224, Loss: 1.6016, Accuracy: 0.3848\n",
      "8_layers - Step 256, Loss: 1.5547, Accuracy: 0.3730\n",
      "8_layers - Step 288, Loss: 1.5078, Accuracy: 0.4707\n",
      "8_layers - Step 320, Loss: 1.5781, Accuracy: 0.3750\n",
      "8_layers - Step 352, Loss: 1.5469, Accuracy: 0.3613\n",
      "8_layers - Step 384, Loss: 1.5000, Accuracy: 0.4199\n",
      "8_layers - Step 416, Loss: 1.5547, Accuracy: 0.3633\n",
      "8_layers - Step 448, Loss: 1.4922, Accuracy: 0.4453\n",
      "8_layers - Step 480, Loss: 1.5469, Accuracy: 0.4023\n",
      "8_layers - Step 512, Loss: 1.5547, Accuracy: 0.3945\n",
      "8_layers - Step 544, Loss: 1.5469, Accuracy: 0.4258\n",
      "8_layers - Step 576, Loss: 1.5469, Accuracy: 0.4355\n",
      "8_layers - Step 608, Loss: 1.5547, Accuracy: 0.4082\n",
      "8_layers - Step 640, Loss: 1.5156, Accuracy: 0.4668\n",
      "8_layers - Step 672, Loss: 1.5156, Accuracy: 0.4375\n",
      "8_layers - Step 704, Loss: 1.5547, Accuracy: 0.4004\n",
      "8_layers - Step 736, Loss: 1.5156, Accuracy: 0.4316\n",
      "8_layers - Step 768, Loss: 1.4922, Accuracy: 0.5332\n",
      "8_layers - Step 800, Loss: 1.4922, Accuracy: 0.4590\n",
      "8_layers - Step 832, Loss: 1.5391, Accuracy: 0.4453\n",
      "8_layers - Step 864, Loss: 1.4922, Accuracy: 0.4531\n",
      "8_layers - Step 896, Loss: 1.4141, Accuracy: 0.5723\n",
      "8_layers - Step 928, Loss: 1.4844, Accuracy: 0.5234\n",
      "8_layers - Step 960, Loss: 1.5000, Accuracy: 0.4668\n",
      "8_layers - Step 992, Loss: 1.4844, Accuracy: 0.4668\n",
      "8_layers - Step 1024, Loss: 1.4531, Accuracy: 0.4883\n",
      "8_layers - Step 1056, Loss: 1.5000, Accuracy: 0.4805\n",
      "8_layers - Step 1088, Loss: 1.4844, Accuracy: 0.4902\n",
      "8_layers - Step 1120, Loss: 1.5234, Accuracy: 0.5234\n",
      "8_layers - Step 1152, Loss: 1.4609, Accuracy: 0.4941\n",
      "8_layers - Step 1184, Loss: 1.4766, Accuracy: 0.4941\n",
      "8_layers - Step 1216, Loss: 1.5078, Accuracy: 0.4883\n",
      "8_layers - Step 1248, Loss: 1.4609, Accuracy: 0.5078\n",
      "8_layers - Step 1280, Loss: 1.5156, Accuracy: 0.4922\n",
      "8_layers - Step 1312, Loss: 1.4453, Accuracy: 0.5762\n",
      "8_layers - Step 1344, Loss: 1.4844, Accuracy: 0.4863\n",
      "8_layers - Step 1376, Loss: 1.4531, Accuracy: 0.5684\n",
      "8_layers - Step 1408, Loss: 1.4453, Accuracy: 0.5781\n",
      "8_layers - Step 1440, Loss: 1.4922, Accuracy: 0.5527\n",
      "8_layers - Step 1472, Loss: 1.4766, Accuracy: 0.5273\n",
      "8_layers - Step 1504, Loss: 1.4531, Accuracy: 0.5898\n",
      "8_layers - Step 1536, Loss: 1.4453, Accuracy: 0.5723\n",
      "8_layers - Step 1568, Loss: 1.4062, Accuracy: 0.6113\n",
      "8_layers - Step 1600, Loss: 1.4297, Accuracy: 0.5645\n",
      "8_layers - Step 1632, Loss: 1.4141, Accuracy: 0.5781\n",
      "8_layers - Step 1664, Loss: 1.4531, Accuracy: 0.5645\n",
      "8_layers - Step 1696, Loss: 1.4844, Accuracy: 0.5195\n",
      "8_layers - Step 1728, Loss: 1.4219, Accuracy: 0.6211\n",
      "8_layers - Step 1760, Loss: 1.4609, Accuracy: 0.6094\n",
      "8_layers - Step 1792, Loss: 1.4375, Accuracy: 0.6055\n",
      "8_layers - Step 1824, Loss: 1.3906, Accuracy: 0.5879\n",
      "8_layers - Step 1856, Loss: 1.4531, Accuracy: 0.5820\n",
      "8_layers - Step 1888, Loss: 1.4219, Accuracy: 0.5996\n",
      "8_layers - Step 1920, Loss: 1.4062, Accuracy: 0.5723\n",
      "8_layers - Step 1952, Loss: 1.3984, Accuracy: 0.6152\n",
      "8_layers - Step 1984, Loss: 1.4141, Accuracy: 0.6016\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 166496 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'softmax_v2',  #######设置kernel,选择有'linear','round','softmax_v1','softmax_v2','relu'\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    # plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison_exp.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = False,\n",
    "        kernel = 'softmax_v2',  #######设置kernel,选择有'linear','round','softmax_v1','softmax_v2','relu'\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = 16\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1, 2, 4, 8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison_exp.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e3c9dffc8b44326be3df51ffae28ba5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 32 to 256:   0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32 to 256 - Step 0, Loss: 5.2812, Accuracy: 0.0020\n",
      "32 to 256 - Step 32, Loss: 1.6250, Accuracy: 0.2627\n",
      "32 to 256 - Step 64, Loss: 1.6328, Accuracy: 0.2598\n",
      "32 to 256 - Step 96, Loss: 1.5859, Accuracy: 0.3057\n",
      "32 to 256 - Step 128, Loss: 1.5781, Accuracy: 0.3154\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32 to 256 - Step 160, Loss: 1.5703, Accuracy: 0.3662\n",
      "32 to 256 - Step 192, Loss: 1.5781, Accuracy: 0.3574\n",
      "32 to 256 - Step 224, Loss: 1.5234, Accuracy: 0.4570\n",
      "32 to 256 - Step 256, Loss: 1.3203, Accuracy: 0.9014\n",
      "setp 281, length come to 64\n",
      "32 to 256 - Step 288, Loss: 1.2734, Accuracy: 0.9639\n",
      "setp 291, length come to 128\n",
      "32 to 256 - Step 320, Loss: 1.2812, Accuracy: 0.9519\n",
      "setp 334, length come to 256\n",
      "32 to 256 - Step 352, Loss: 1.3359, Accuracy: 0.8527\n",
      "32 to 256 - Step 384, Loss: 1.2969, Accuracy: 0.8988\n",
      "32 to 256 - Step 416, Loss: 1.2969, Accuracy: 0.9001\n",
      "32 to 256 - Step 448, Loss: 1.2891, Accuracy: 0.9143\n",
      "32 to 256 - Step 480, Loss: 1.2891, Accuracy: 0.9296\n",
      "32 to 256 - Step 512, Loss: 1.2969, Accuracy: 0.9069\n",
      "32 to 256 - Step 544, Loss: 1.2812, Accuracy: 0.9196\n",
      "32 to 256 - Step 576, Loss: 1.2812, Accuracy: 0.9270\n",
      "32 to 256 - Step 608, Loss: 1.2734, Accuracy: 0.9180\n",
      "32 to 256 - Step 640, Loss: 1.2891, Accuracy: 0.9359\n",
      "32 to 256 - Step 672, Loss: 1.2891, Accuracy: 0.9296\n",
      "32 to 256 - Step 704, Loss: 1.2812, Accuracy: 0.9407\n",
      "32 to 256 - Step 736, Loss: 1.2734, Accuracy: 0.9221\n",
      "32 to 256 - Step 768, Loss: 1.2891, Accuracy: 0.9347\n",
      "32 to 256 - Step 800, Loss: 1.2578, Accuracy: 0.9415\n",
      "32 to 256 - Step 832, Loss: 1.2578, Accuracy: 0.9485\n",
      "32 to 256 - Step 864, Loss: 1.2812, Accuracy: 0.9427\n",
      "32 to 256 - Step 896, Loss: 1.2656, Accuracy: 0.9478\n",
      "32 to 256 - Step 928, Loss: 1.2734, Accuracy: 0.9508\n",
      "32 to 256 - Step 960, Loss: 1.2578, Accuracy: 0.9537\n",
      "32 to 256 - Step 992, Loss: 1.2578, Accuracy: 0.9513\n",
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9701c1bd8bc3486f95995a255421c080",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 256:   0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "256 - Step 0, Loss: 4.5312, Accuracy: 0.0494\n",
      "256 - Step 32, Loss: 1.6172, Accuracy: 0.2069\n",
      "256 - Step 64, Loss: 1.6250, Accuracy: 0.2061\n",
      "256 - Step 96, Loss: 1.6328, Accuracy: 0.2021\n",
      "256 - Step 128, Loss: 1.6172, Accuracy: 0.2089\n",
      "256 - Step 160, Loss: 1.6094, Accuracy: 0.2200\n",
      "256 - Step 192, Loss: 1.6172, Accuracy: 0.1946\n",
      "256 - Step 224, Loss: 1.6172, Accuracy: 0.2107\n",
      "256 - Step 256, Loss: 1.6094, Accuracy: 0.2018\n",
      "256 - Step 288, Loss: 1.6094, Accuracy: 0.2194\n",
      "256 - Step 320, Loss: 1.6016, Accuracy: 0.2228\n",
      "256 - Step 352, Loss: 1.6094, Accuracy: 0.2173\n",
      "256 - Step 384, Loss: 1.6094, Accuracy: 0.2220\n",
      "256 - Step 416, Loss: 1.6094, Accuracy: 0.2207\n",
      "256 - Step 448, Loss: 1.6094, Accuracy: 0.2267\n",
      "256 - Step 480, Loss: 1.6016, Accuracy: 0.2150\n",
      "256 - Step 512, Loss: 1.6016, Accuracy: 0.2180\n",
      "256 - Step 544, Loss: 1.6094, Accuracy: 0.2173\n",
      "256 - Step 576, Loss: 1.6094, Accuracy: 0.2185\n",
      "256 - Step 608, Loss: 1.6016, Accuracy: 0.2356\n",
      "256 - Step 640, Loss: 1.6016, Accuracy: 0.2297\n",
      "256 - Step 672, Loss: 1.6016, Accuracy: 0.2317\n",
      "256 - Step 704, Loss: 1.6016, Accuracy: 0.2286\n",
      "256 - Step 736, Loss: 1.5938, Accuracy: 0.2426\n",
      "256 - Step 768, Loss: 1.6016, Accuracy: 0.2405\n",
      "256 - Step 800, Loss: 1.5703, Accuracy: 0.3037\n",
      "256 - Step 832, Loss: 1.5078, Accuracy: 0.4917\n",
      "256 - Step 864, Loss: 1.4141, Accuracy: 0.7549\n",
      "256 - Step 896, Loss: 1.3359, Accuracy: 0.8779\n",
      "256 - Step 928, Loss: 1.3047, Accuracy: 0.8926\n",
      "256 - Step 960, Loss: 1.3047, Accuracy: 0.9043\n",
      "256 - Step 992, Loss: 1.3047, Accuracy: 0.9100\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 89472 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None,length = 16):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "        kl2 = True,\n",
    "        kernel = 'round',\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 1000\n",
    "    batch_size = 32\n",
    "    try:\n",
    "        k = int(length)\n",
    "        end = k\n",
    "    except:\n",
    "        k = int(length.split(' ')[0])\n",
    "        end = int(length.split(' ')[-1])\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "        if accuracy>0.99 and end>k:\n",
    "            k*=2\n",
    "            print(f'setp {step}, length come to {k}')\n",
    "            dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for length in ['32 to 256',256]:\n",
    "        print(f\"\\n=== Training with 1 layers ===\")\n",
    "        train_model_with_layers(1, run_name=f\"{length}\", log_dict=log_dict,length = length)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0de36c3737204a3d889b35e60f71f59b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 32 to 256:   0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32 to 256 - Step 0, Loss: 4.7500, Accuracy: 0.0234\n",
      "32 to 256 - Step 32, Loss: 1.6250, Accuracy: 0.2637\n",
      "32 to 256 - Step 64, Loss: 1.6328, Accuracy: 0.2480\n",
      "32 to 256 - Step 96, Loss: 1.6797, Accuracy: 0.2432\n",
      "32 to 256 - Step 128, Loss: 1.6484, Accuracy: 0.2959\n",
      "32 to 256 - Step 160, Loss: 1.5391, Accuracy: 0.3604\n",
      "32 to 256 - Step 192, Loss: 1.5703, Accuracy: 0.3096\n",
      "32 to 256 - Step 224, Loss: 1.5391, Accuracy: 0.4385\n",
      "32 to 256 - Step 256, Loss: 1.5078, Accuracy: 0.5537\n",
      "32 to 256 - Step 288, Loss: 1.3047, Accuracy: 0.9561\n",
      "32 to 256 - Step 320, Loss: 1.2578, Accuracy: 0.9844\n",
      "setp 325, length come to 64\n",
      "setp 341, length come to 128\n",
      "setp 352, length come to 256\n",
      "32 to 256 - Step 352, Loss: 1.3047, Accuracy: 0.9191\n",
      "32 to 256 - Step 384, Loss: 1.2500, Accuracy: 0.9779\n",
      "32 to 256 - Step 416, Loss: 1.2422, Accuracy: 0.9836\n",
      "32 to 256 - Step 448, Loss: 1.2344, Accuracy: 0.9896\n",
      "32 to 256 - Step 480, Loss: 1.2500, Accuracy: 0.9851\n",
      "32 to 256 - Step 512, Loss: 1.2422, Accuracy: 0.9983\n",
      "32 to 256 - Step 544, Loss: 1.2344, Accuracy: 0.9929\n",
      "32 to 256 - Step 576, Loss: 1.2344, Accuracy: 0.9886\n",
      "32 to 256 - Step 608, Loss: 1.2422, Accuracy: 0.9971\n",
      "32 to 256 - Step 640, Loss: 1.2422, Accuracy: 0.9950\n",
      "32 to 256 - Step 672, Loss: 1.2344, Accuracy: 0.9924\n",
      "32 to 256 - Step 704, Loss: 1.2344, Accuracy: 0.9993\n",
      "32 to 256 - Step 736, Loss: 1.2500, Accuracy: 0.9957\n",
      "32 to 256 - Step 768, Loss: 1.2344, Accuracy: 0.9969\n",
      "32 to 256 - Step 800, Loss: 1.2500, Accuracy: 0.9977\n",
      "32 to 256 - Step 832, Loss: 1.2109, Accuracy: 1.0000\n",
      "32 to 256 - Step 864, Loss: 1.2422, Accuracy: 0.9916\n",
      "32 to 256 - Step 896, Loss: 1.2188, Accuracy: 1.0000\n",
      "32 to 256 - Step 928, Loss: 1.2344, Accuracy: 0.9934\n",
      "32 to 256 - Step 960, Loss: 1.2266, Accuracy: 1.0000\n",
      "32 to 256 - Step 992, Loss: 1.2266, Accuracy: 1.0000\n",
      "\n",
      "=== Training with 1 layers ===\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3018a044b3f47f082377c0282993679",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training 256:   0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "256 - Step 0, Loss: 4.7812, Accuracy: 0.0370\n",
      "256 - Step 32, Loss: 1.6250, Accuracy: 0.2104\n",
      "256 - Step 64, Loss: 1.6250, Accuracy: 0.2106\n",
      "256 - Step 96, Loss: 1.6328, Accuracy: 0.2025\n",
      "256 - Step 128, Loss: 1.6172, Accuracy: 0.1976\n",
      "256 - Step 160, Loss: 1.6172, Accuracy: 0.2006\n",
      "256 - Step 192, Loss: 1.6094, Accuracy: 0.2135\n",
      "256 - Step 224, Loss: 1.6172, Accuracy: 0.2252\n",
      "256 - Step 256, Loss: 1.6094, Accuracy: 0.2246\n",
      "256 - Step 288, Loss: 1.6094, Accuracy: 0.2098\n",
      "256 - Step 320, Loss: 1.6094, Accuracy: 0.2006\n",
      "256 - Step 352, Loss: 1.6094, Accuracy: 0.2045\n",
      "256 - Step 384, Loss: 1.6094, Accuracy: 0.2218\n",
      "256 - Step 416, Loss: 1.6016, Accuracy: 0.2107\n",
      "256 - Step 448, Loss: 1.6250, Accuracy: 0.2067\n",
      "256 - Step 480, Loss: 1.6094, Accuracy: 0.2131\n",
      "256 - Step 512, Loss: 1.6016, Accuracy: 0.2233\n",
      "256 - Step 544, Loss: 1.6016, Accuracy: 0.2373\n",
      "256 - Step 576, Loss: 1.6016, Accuracy: 0.2129\n",
      "256 - Step 608, Loss: 1.6094, Accuracy: 0.2072\n",
      "256 - Step 640, Loss: 1.6016, Accuracy: 0.2198\n",
      "256 - Step 672, Loss: 1.6016, Accuracy: 0.2239\n",
      "256 - Step 704, Loss: 1.6016, Accuracy: 0.2179\n",
      "256 - Step 736, Loss: 1.6016, Accuracy: 0.2173\n",
      "256 - Step 768, Loss: 1.6016, Accuracy: 0.2268\n",
      "256 - Step 800, Loss: 1.6016, Accuracy: 0.2288\n",
      "256 - Step 832, Loss: 1.6016, Accuracy: 0.2328\n",
      "256 - Step 864, Loss: 1.6016, Accuracy: 0.2130\n",
      "256 - Step 896, Loss: 1.6016, Accuracy: 0.2268\n",
      "256 - Step 928, Loss: 1.6016, Accuracy: 0.2102\n",
      "256 - Step 960, Loss: 1.6016, Accuracy: 0.2245\n",
      "256 - Step 992, Loss: 1.6016, Accuracy: 0.2247\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79668 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from itertools import permutations\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 数据生成相关逻辑保持不变\n",
    "n_elements = 5\n",
    "swap_pairs = [(i, j) for i in range(n_elements) for j in range(i+1, n_elements)]\n",
    "\n",
    "def apply_swap(perm, swap_idx):\n",
    "    i, j = swap_pairs[swap_idx]\n",
    "    perm = list(perm)\n",
    "    perm[i], perm[j] = perm[j], perm[i]\n",
    "    return tuple(perm)\n",
    "\n",
    "def generate_S5_data(k, num_samples):\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        swap_sequence = [random.randint(0, len(swap_pairs)-1) for _ in range(k)]\n",
    "        current_perm = tuple(range(n_elements))\n",
    "        first_elements = []\n",
    "\n",
    "        for swap_idx in swap_sequence:\n",
    "            current_perm = apply_swap(current_perm, swap_idx)\n",
    "            first_elements.append(current_perm[0])\n",
    "\n",
    "        input_ids = torch.tensor(swap_sequence, dtype=torch.long)\n",
    "        labels = torch.tensor(first_elements, dtype=torch.long)\n",
    "        data.append((input_ids, labels))\n",
    "\n",
    "    return data\n",
    "\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = generate_S5_data(self.k, self.num_samples)\n",
    "        input_ids, labels = data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None,length = 16):\n",
    "    # 配置模型\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=1,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = -1,\n",
    "        kl2 = True,\n",
    "        kernel = 'round',\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    model.to(device).to(torch.bfloat16)\n",
    "\n",
    "    # 参数设置\n",
    "    num_epochs = 1000\n",
    "    batch_size = 32\n",
    "    try:\n",
    "        k = int(length)\n",
    "        end = k\n",
    "    except:\n",
    "        k = int(length.split(' ')[0])\n",
    "        end = int(length.split(' ')[-1])\n",
    "    num_samples_per_batch = batch_size\n",
    "\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    progress_bar = tqdm(range(num_training_steps), desc=f\"Training {run_name}\")\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    for epoch in range(num_epochs):\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "\n",
    "            outputs = model(input_ids=input_ids, labels=labels, return_dict=True)\n",
    "            loss = outputs.loss\n",
    "\n",
    "            logits = outputs.logits\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "\n",
    "            if step % 32 == 0:\n",
    "                print(f\"{run_name} - Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "            step += 1\n",
    "            progress_bar.update(1)\n",
    "\n",
    "        if accuracy>0.99 and end>k:\n",
    "            k*=2\n",
    "            print(f'setp {step}, length come to {k}')\n",
    "            dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for length in ['32 to 256',256]:\n",
    "        print(f\"\\n=== Training with 1 layers ===\")\n",
    "        train_model_with_layers(1, run_name=f\"{length}\", log_dict=log_dict,length = length)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    plt.title(\"Delta Transformer\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81232 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    if key == str(128) or key== str(32):\n",
    "        continue\n",
    "    accs = log_dict[key]['acc'][:1000]\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key) \n",
    "#%% x = 325 342 ,352分别画一个竖着的线，把训练区间划分为4段，分别是32，64，128，256\n",
    "ranges = [(0, 325), (325, 342), (342, 352), (352, 999)]\n",
    "colors = ['lightblue', 'lightgreen', 'lightcoral', 'wheat']\n",
    "for (start, end), color in zip(ranges, colors):\n",
    "    plt.axvspan(start, end, color=color, alpha=0.2)   \n",
    "\n",
    "# plt.title(\"Delta Transformer\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 90572 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 6))\n",
    "for key in log_dict:\n",
    "    if key == str(128) or key== str(32):\n",
    "        continue\n",
    "    accs = log_dict[key]['acc'][:1000]\n",
    "    steps = np.arange(len(accs))\n",
    "    plt.plot(steps, accs, label=key) \n",
    "#%% x = 285,301,325分别画一个竖着的线，把训练区间划分为4段，分别是32，64，128，256\n",
    "    \n",
    "\n",
    "# plt.title(\"Delta Transformer\")\n",
    "plt.xlabel(\"Training Step\")\n",
    "plt.ylabel(\"Accuracy\")\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.savefig(\"accuracy_comparison.png\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "fileId": "cee38b77-89c0-4f5f-bc5e-248340d9fc75",
  "filePath": "/mnt/bn/seed-aws-va/xumingyu/vis/swap.ipynb",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
