678 lines
22 KiB
Plaintext
678 lines
22 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "cab91cfc",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/home/ubuntu/Power-RAG/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
|
" from .autonotebook import tqdm as notebook_tqdm\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"import copy\n",
|
|
"import dataclasses\n",
|
|
"import os\n",
|
|
"import time\n",
|
|
"import pathlib\n",
|
|
"import itertools\n",
|
|
"import multiprocessing\n",
|
|
"import scipy\n",
|
|
"import numpy as np\n",
|
|
"import pandas as pd\n",
|
|
"import pickle\n",
|
|
"import gzip\n",
|
|
"import threading\n",
|
|
"import queue\n",
|
|
"import pytz\n",
|
|
"import traceback\n",
|
|
"from datetime import datetime\n",
|
|
"from tqdm.auto import tqdm, trange\n",
|
|
"from typing import Any\n",
|
|
"\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"import matplotlib.ticker as mtick\n",
|
|
"%matplotlib inline\n",
|
|
"%config InlineBackend.figure_format='retina'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "8d24fbd7",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Sat Apr 12 00:10:05 2025 \n",
|
|
"+-----------------------------------------------------------------------------------------+\n",
|
|
"| NVIDIA-SMI 550.120 Driver Version: 550.120 CUDA Version: 12.4 |\n",
|
|
"|-----------------------------------------+------------------------+----------------------+\n",
|
|
"| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
|
|
"| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n",
|
|
"| | | MIG M. |\n",
|
|
"|=========================================+========================+======================|\n",
|
|
"| 0 NVIDIA A10G Off | 00000000:00:1E.0 Off | 0 |\n",
|
|
"| 0% 27C P8 15W / 300W | 4MiB / 23028MiB | 0% Default |\n",
|
|
"| | | N/A |\n",
|
|
"+-----------------------------------------+------------------------+----------------------+\n",
|
|
" \n",
|
|
"+-----------------------------------------------------------------------------------------+\n",
|
|
"| Processes: |\n",
|
|
"| GPU GI CI PID Type Process name GPU Memory |\n",
|
|
"| ID ID Usage |\n",
|
|
"|=========================================================================================|\n",
|
|
"| No running processes found |\n",
|
|
"+-----------------------------------------------------------------------------------------+\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"!nvidia-smi"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "538b2c11",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def benchmark(f, *, f_setup=None, min_repeat: int, min_secs: float, tqdm_kwargs: dict | None=None) -> np.ndarray:\n",
|
|
" latency = []\n",
|
|
" \n",
|
|
" # First run, ignore min_secs\n",
|
|
" if f_setup is not None:\n",
|
|
" f_setup()\n",
|
|
" st = time.perf_counter_ns()\n",
|
|
" f()\n",
|
|
" ed = time.perf_counter_ns()\n",
|
|
" latency.append((ed-st)/1e9)\n",
|
|
" \n",
|
|
" # Subsequent runs, until reaching both min_repeat and min_secs\n",
|
|
" min_nanos = int(min_secs * 1e9)\n",
|
|
" start_nanos = time.perf_counter_ns()\n",
|
|
" while True:\n",
|
|
" now_nanos = time.perf_counter_ns()\n",
|
|
" if len(latency) > min_repeat and now_nanos - start_nanos > min_nanos:\n",
|
|
" break\n",
|
|
" if f_setup is not None:\n",
|
|
" f_setup()\n",
|
|
" st = time.perf_counter_ns()\n",
|
|
" f()\n",
|
|
" ed = time.perf_counter_ns()\n",
|
|
" latency.append((ed-st)/1e9)\n",
|
|
" return np.array(latency)\n",
|
|
"\n",
|
|
"def tail_mean(xs, skip=0.2):\n",
|
|
" return xs[int(len(xs) * skip):].mean()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"id": "02c9c9b1",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"<torch.autograd.grad_mode.set_grad_enabled at 0x7c5afc12b850>"
|
|
]
|
|
},
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"import torch\n",
|
|
"torch.set_grad_enabled(False)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "3405fdc7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"nd_list = list(itertools.chain(itertools.product([12, 3], [256])))\n",
|
|
"seqlen_list = [256]\n",
|
|
"bs_list = [2,4,8,16,32,64,128,256,512,1024,2048]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"id": "10dc981a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"[(12, 256), (3, 256)]\n",
|
|
"[256]\n",
|
|
"[2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(nd_list)\n",
|
|
"print(seqlen_list)\n",
|
|
"print(bs_list)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"id": "7e0ee385",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def benchmark_dense(out, nd_list, seqlen_list, bs_list):\n",
|
|
" seqlen_list = [1] + seqlen_list\n",
|
|
" total = len(list(itertools.product(nd_list, seqlen_list, bs_list)))\n",
|
|
" pbar = tqdm(total=total)\n",
|
|
" for (n, d), seqlen in reversed(list(itertools.product(nd_list, seqlen_list))):\n",
|
|
" h = n * d\n",
|
|
" maxbs = max(bs_list)\n",
|
|
" print(maxbs, n, d, seqlen)\n",
|
|
" cache = torch.empty(int(256e6 // 4), dtype=torch.int, device=\"cuda:0\")\n",
|
|
" X = torch.rand((maxbs, seqlen, h), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" W = torch.rand((h, h), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" for bs in reversed(bs_list):\n",
|
|
" pbar.set_postfix(n=n, h=h, d=d, seqlen=seqlen, bs=bs)\n",
|
|
" def run():\n",
|
|
" torch.matmul(X[:bs], W)\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" def clear_cache():\n",
|
|
" cache.zero_()\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" latency = benchmark(run, f_setup=clear_cache, min_repeat=20, min_secs=2)\n",
|
|
" l = tail_mean(latency)\n",
|
|
" out.append({\n",
|
|
" \"n\": n,\n",
|
|
" \"d\": d,\n",
|
|
" \"seqlen\": seqlen,\n",
|
|
" \"bs\": bs,\n",
|
|
" \"latency\": l\n",
|
|
" })\n",
|
|
" pbar.update()\n",
|
|
" del cache, X, W\n",
|
|
" torch.cuda.empty_cache()\n",
|
|
" pbar.close()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"id": "c206a502",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def benchmark_qk_init(out, nd_list, seqlen_list, bs_list):\n",
|
|
" total = len(list(itertools.product(nd_list, seqlen_list, bs_list)))\n",
|
|
" pbar = tqdm(total=total)\n",
|
|
" for (n, d), seqlen in reversed(list(itertools.product(nd_list, seqlen_list))):\n",
|
|
" h = n * d\n",
|
|
" try:\n",
|
|
" maxbs = max(b for b in bs_list if b*n*seqlen*d*2*2+b*n*seqlen**2*2 < 80e9)\n",
|
|
" except ValueError:\n",
|
|
" pbar.update(len(bs_list))\n",
|
|
" continue\n",
|
|
" cache = torch.empty(int(256e6 // 4), dtype=torch.int, device=\"cuda:0\")\n",
|
|
" Qmax = torch.rand((maxbs, n, seqlen, d), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" Kmax = torch.rand((maxbs, n, seqlen, d), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" for bs in reversed(bs_list):\n",
|
|
" pbar.set_postfix(n=n, h=h, d=d, seqlen=seqlen, bs=bs)\n",
|
|
" if bs > maxbs:\n",
|
|
" pbar.update()\n",
|
|
" continue\n",
|
|
" Q = Qmax[:bs]\n",
|
|
" K = Kmax[:bs]\n",
|
|
" def run():\n",
|
|
" torch.bmm(Q.view(bs * n, seqlen, d), K.view(bs * n, seqlen, d).transpose(1, 2))\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" def clear_cache():\n",
|
|
" cache.zero_()\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" latency = benchmark(run, f_setup=clear_cache, min_repeat=20, min_secs=2)\n",
|
|
" l = tail_mean(latency)\n",
|
|
" out.append({\n",
|
|
" \"n\": n,\n",
|
|
" \"d\": d,\n",
|
|
" \"seqlen\": seqlen,\n",
|
|
" \"bs\": bs,\n",
|
|
" \"latency\": l\n",
|
|
" })\n",
|
|
" pbar.update()\n",
|
|
" del cache, Q, K, Qmax, Kmax\n",
|
|
" torch.cuda.empty_cache()\n",
|
|
" pbar.close()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"id": "a3a2103c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def benchmark_qk_ar(out, nd_list, seqlen_list, bs_list):\n",
|
|
" total = len(list(itertools.product(nd_list, seqlen_list, bs_list)))\n",
|
|
" pbar = tqdm(total=total)\n",
|
|
" for (n, d), seqlen in reversed(list(itertools.product(nd_list, seqlen_list))):\n",
|
|
" h = n * d\n",
|
|
" try:\n",
|
|
" maxbs = max(b for b in bs_list if b*n*(1+seqlen)*d*2+b*n*seqlen*2 < 80e9)\n",
|
|
" except ValueError:\n",
|
|
" pbar.update(len(bs_list))\n",
|
|
" continue\n",
|
|
" cache = torch.empty(int(256e6 // 4), dtype=torch.int, device=\"cuda:0\")\n",
|
|
" Qmax = torch.rand((maxbs, n, 1, d), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" Kmax = torch.rand((maxbs, n, seqlen, d), dtype=torch.bfloat16, device=\"cuda:0\")\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" for bs in reversed(bs_list):\n",
|
|
" pbar.set_postfix(n=n, h=h, d=d, seqlen=seqlen, bs=bs)\n",
|
|
" if bs > maxbs:\n",
|
|
" pbar.update()\n",
|
|
" continue\n",
|
|
" Q = Qmax[:bs]\n",
|
|
" K = Kmax[:bs]\n",
|
|
" def run():\n",
|
|
" torch.bmm(Q.view(bs * n, 1, d), K.view(bs * n, seqlen, d).transpose(1, 2))\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" def clear_cache():\n",
|
|
" cache.zero_()\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" latency = benchmark(run, f_setup=clear_cache, min_repeat=20, min_secs=2)\n",
|
|
" l = tail_mean(latency)\n",
|
|
" out.append({\n",
|
|
" \"n\": n,\n",
|
|
" \"d\": d,\n",
|
|
" \"seqlen\": seqlen,\n",
|
|
" \"bs\": bs,\n",
|
|
" \"latency\": l\n",
|
|
" })\n",
|
|
" pbar.update()\n",
|
|
" del cache, Q, K, Qmax, Kmax\n",
|
|
" torch.cuda.empty_cache()\n",
|
|
" pbar.close()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 10,
|
|
"id": "3aaad98a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"data = {}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"id": "18137de3",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
" 0%| | 0/22 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"100%|██████████| 22/22 [00:44<00:00, 2.04s/it, bs=2, d=256, h=3072, n=12, seqlen=256] \n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"db = []\n",
|
|
"benchmark_qk_init(db, nd_list, seqlen_list, bs_list)\n",
|
|
"data[\"qk_init\"] = db"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"id": "26c76e15",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"100%|██████████| 22/22 [00:44<00:00, 2.01s/it, bs=2, d=256, h=3072, n=12, seqlen=256] \n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"db = []\n",
|
|
"benchmark_qk_ar(db, nd_list, seqlen_list, bs_list)\n",
|
|
"data[\"qk_ar\"] = db"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 13,
|
|
"id": "313e36eb",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
" 0%| | 0/44 [00:00<?, ?it/s, bs=2048, d=256, h=768, n=3, seqlen=256]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"2048 3 256 256\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
" 25%|██▌ | 11/44 [00:22<01:06, 2.00s/it, bs=2048, d=256, h=768, n=3, seqlen=1] "
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"2048 3 256 1\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
" 50%|█████ | 22/44 [00:44<00:44, 2.00s/it, bs=2048, d=256, h=3072, n=12, seqlen=256]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"2048 12 256 256\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
" 75%|███████▌ | 33/44 [01:07<00:22, 2.02s/it, bs=2048, d=256, h=3072, n=12, seqlen=1] "
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"2048 12 256 1\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"100%|██████████| 44/44 [01:29<00:00, 2.03s/it, bs=2, d=256, h=3072, n=12, seqlen=1] \n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"db = []\n",
|
|
"benchmark_dense(db, nd_list, seqlen_list, bs_list)\n",
|
|
"data[\"dense\"] = db"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"id": "50c37959",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"with gzip.open(\"data/20230516-transformer-batching1.pkl.gz\", \"wb\") as f:\n",
|
|
" pickle.dump(data, f)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"id": "828ddb54",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"df_dense = (\n",
|
|
" pd.DataFrame.from_dict(data[\"dense\"])\n",
|
|
" .assign(h=lambda x: x[\"n\"] * x[\"d\"])\n",
|
|
" .assign(flop=lambda x: (x[\"bs\"] * x[\"seqlen\"] * x[\"h\"]**2) * 2)\n",
|
|
" .assign(io=lambda x: (x[\"bs\"]*x[\"seqlen\"]*x[\"h\"]*2 + x[\"h\"]**2) * 2/x['latency']/1e9)\n",
|
|
" .assign(intensity=lambda x: x[\"flop\"] / x[\"io\"])\n",
|
|
" .assign(throughput=lambda x: x[\"flop\"] / x[\"latency\"])\n",
|
|
" .assign(series=\"dense\")\n",
|
|
")\n",
|
|
"df_qk_init = (\n",
|
|
" pd.DataFrame.from_dict(data[\"qk_init\"])\n",
|
|
" .assign(h=lambda x: x[\"n\"] * x[\"d\"])\n",
|
|
" .assign(flop=lambda x: (x[\"bs\"]*x[\"n\"]*x[\"d\"]*x[\"seqlen\"]**2) * 2)\n",
|
|
" .assign(io=lambda x: (x[\"bs\"]*x[\"n\"]*(x[\"seqlen\"]*x[\"d\"]*2 + x[\"seqlen\"]**2)) * 2/x['latency']/1e9)\n",
|
|
" .assign(intensity=lambda x: x[\"flop\"] / x[\"io\"])\n",
|
|
" .assign(throughput=lambda x: x[\"flop\"] / x[\"latency\"])\n",
|
|
" .assign(series=\"qk_init\")\n",
|
|
")\n",
|
|
"df_qk_ar = (\n",
|
|
" pd.DataFrame.from_dict(data[\"qk_ar\"])\n",
|
|
" .assign(h=lambda x: x[\"n\"] * x[\"d\"])\n",
|
|
" .assign(flop=lambda x: (x[\"bs\"]*x[\"n\"]*x[\"d\"]*x[\"seqlen\"]) * 2)\n",
|
|
" .assign(io=lambda x: (x[\"bs\"]*x[\"n\"]*(x[\"d\"] + x[\"seqlen\"]*x[\"d\"] + x[\"seqlen\"])) * 2)\n",
|
|
" .assign(intensity=lambda x: x[\"flop\"] / x[\"io\"])\n",
|
|
" .assign(throughput=lambda x: x[\"bs\"] / x[\"latency\"])\n",
|
|
" .assign(series=\"qk_ar\")\n",
|
|
")\n",
|
|
"pd.concat([df_dense, df_qk_init, df_qk_ar]).to_csv(\"data/transformer-batching-microbenchmarks.csv\", index=False)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 39,
|
|
"id": "c296a395",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"<module 'pandas' from '/home/ubuntu/Power-RAG/.venv/lib/python3.10/site-packages/pandas/__init__.py'>"
|
|
]
|
|
},
|
|
"execution_count": 39,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"pd\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "a25cdd5a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "63b8a531",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import transformers"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "af90eff1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def _gen_opt_cfg(n_layers: int, d_model: int, n_heads: int, **kwargs) -> transformers.OPTConfig:\n",
|
|
" return transformers.OPTConfig(\n",
|
|
" num_hidden_layers=n_layers,\n",
|
|
" hidden_size=d_model,\n",
|
|
" ffn_dim=d_model*4,\n",
|
|
" num_attention_heads=n_heads,\n",
|
|
" **kwargs\n",
|
|
" )\n",
|
|
"optcfg = {\n",
|
|
" # https://arxiv.org/pdf/2205.01068.pdf Table 2.1\n",
|
|
" \"125m\": _gen_opt_cfg(12, 768, 12),\n",
|
|
" \"350m\": _gen_opt_cfg(24, 1024, 16),\n",
|
|
" \"760m\": _gen_opt_cfg(24, 1536, 16),\n",
|
|
" \"1.3b\": _gen_opt_cfg(24, 2048, 32),\n",
|
|
" \"2.7b\": _gen_opt_cfg(32, 2560, 32),\n",
|
|
" \"6.7b\": _gen_opt_cfg(32, 4096, 32),\n",
|
|
" \"13b\": _gen_opt_cfg(40, 5120, 40),\n",
|
|
" \"13b_1layer\": _gen_opt_cfg(1, 5120, 40),\n",
|
|
" \"30b\": _gen_opt_cfg(48, 7168, 56),\n",
|
|
" \"66b\": _gen_opt_cfg(64, 9216, 72),\n",
|
|
" \"175b\": _gen_opt_cfg(96, 12288, 96),\n",
|
|
" \"175b_1layer\": _gen_opt_cfg(1, 12288, 96),\n",
|
|
"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5b9ebbec",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def greedy_sample_one(model, input_ids, attention_mask=None, past_key_values=None):\n",
|
|
" bs, tgt_len = input_ids.shape\n",
|
|
" if past_key_values is not None:\n",
|
|
" _bs, _num_heads, src_len, _head_dims = past_key_values[0][0].shape\n",
|
|
" assert bs == _bs\n",
|
|
" else:\n",
|
|
" src_len = 0\n",
|
|
" if attention_mask is None:\n",
|
|
" attention_mask = torch.ones((bs, src_len + tgt_len), device=model.device)\n",
|
|
" ret = model(\n",
|
|
" input_ids=input_ids,\n",
|
|
" attention_mask=attention_mask,\n",
|
|
" past_key_values=past_key_values,\n",
|
|
" use_cache=True, output_hidden_states=False, return_dict=True,\n",
|
|
" )\n",
|
|
" return ret\n",
|
|
"\n",
|
|
"def time_greedy_generate(model, input_ids, new_tokens):\n",
|
|
" ts = []\n",
|
|
" output = input_ids\n",
|
|
" past_key_values = None\n",
|
|
" cache = torch.empty(int(256e6 // 4), dtype=torch.int, device=model.device)\n",
|
|
" attention_mask = torch.ones(input_ids.shape, device=model.device) \n",
|
|
" for _ in range(new_tokens):\n",
|
|
" cache.zero_()\n",
|
|
" torch.cuda.synchronize()\n",
|
|
" st = time.perf_counter_ns()\n",
|
|
" \n",
|
|
" ret = greedy_sample_one(model, input_ids, attention_mask, past_key_values)\n",
|
|
" input_ids = torch.argmax(ret.logits[:, -1, :], axis=-1)[:, None]\n",
|
|
" output = torch.cat([output, input_ids], axis=1)\n",
|
|
" past_key_values = ret.past_key_values\n",
|
|
" attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)\n",
|
|
" \n",
|
|
" torch.cuda.synchronize()\n",
|
|
" ed = time.perf_counter_ns()\n",
|
|
" ts.append((ed-st)/1e9)\n",
|
|
" return np.array(ts)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fc92f940",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"opt_config = optcfg[\"6.7b\"]\n",
|
|
"\n",
|
|
"torch.set_default_dtype(torch.bfloat16)\n",
|
|
"with transformers.modeling_utils.no_init_weights():\n",
|
|
" model = transformers.models.opt.OPTForCausalLM(opt_config).to(\"cuda\")\n",
|
|
"torch.set_default_dtype(torch.float32)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c19fa396",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"db = {}\n",
|
|
"input_tokens = 200\n",
|
|
"new_tokens = 500\n",
|
|
"for bs in tqdm(list(itertools.chain(range(1, 8), range(8, 16, 2), [16]))):\n",
|
|
" x = torch.randint(1000, 10000, (bs, input_tokens), device=model.device)\n",
|
|
" stack = []\n",
|
|
" for _ in range(10):\n",
|
|
" l = time_greedy_generate(model, x, new_tokens=new_tokens)\n",
|
|
" stack.append(l)\n",
|
|
" db[bs] = np.median(np.stack(stack), axis=0)\n",
|
|
" del x\n",
|
|
" torch.cuda.empty_cache()\n",
|
|
"del model\n",
|
|
"torch.cuda.empty_cache()\n",
|
|
"\n",
|
|
"with gzip.open(\"data/20230516-e2e-text-generation-batch.pkl.gz\", \"wb\") as f:\n",
|
|
" pickle.dump(db, f)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"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.10.12"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|