{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "8d1d3984-edef-4142-89ad-7ebddf1696ca",
"metadata": {
"execution": {
"iopub.execute_input": "2024-07-29T16:52:29.081889Z",
"iopub.status.busy": "2024-07-29T16:52:29.081436Z",
"iopub.status.idle": "2024-07-29T16:52:29.087301Z",
"shell.execute_reply": "2024-07-29T16:52:29.086777Z",
"shell.execute_reply.started": "2024-07-29T16:52:29.081867Z"
}
},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\""
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "d684d892",
"metadata": {},
"outputs": [],
"source": [
"import sys\n",
"\n",
"sys.path.insert(0, \"/home/tony/Work/neon/sunoGPT\")\n",
"from utils.dpo_data_utils import get_batch"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "4fbb3ae8",
"metadata": {},
"outputs": [],
"source": [
"from suno_utils.utils.text import (\n",
" normalize_whitespace,\n",
" read_json,\n",
" read_jsonl,\n",
" write_json,\n",
" write_jsonl,\n",
")\n",
"import numpy as np\n",
"import pandas as pd"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "98c5907d-abc6-4c93-929d-344a96031d50",
"metadata": {},
"outputs": [],
"source": [
"# model_path = \"/app/suno/checkpoints/2024-07-04_04-14-58/2b_75k_infer.pt\"\n",
"model_path = \"/app/suno/data/dpo/models/model_30b_fix_ft2_20k.pt\"\n",
"tokenizer_path = \"/app/suno/models/chirp_v2/tokenizer_60k.json\""
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "796cf784",
"metadata": {},
"outputs": [],
"source": [
"DATA_DIR = \"/app/suno/data/dpo/30b_t1_v7\"\n",
"\n",
"test_metas = read_jsonl(os.path.join(DATA_DIR, \"meta_val.jsonl\"))\n",
"test_info = read_json(os.path.join(DATA_DIR, \"info_val.json\"))\n",
"test_data = np.memmap(\n",
" os.path.join(DATA_DIR, \"data_val.bin\"), dtype=np.uint16, mode=\"r\"\n",
").reshape(-1, 6016, 13)\n",
"\n",
"assert len(test_data) == len(test_metas)\n",
"assert test_data[:100, :, 0].min() >= 0\n",
"assert test_data[:100, :, 0].max() <= 4000\n",
"assert test_data[:100, :, 1:].min() >= 0\n",
"assert test_data[:100, :, 1:].max() <= 2048"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "f2acee85",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[2024-07-31_04:05:46]: Failed to import xformers.\n",
"[2024-07-31_04:05:46]: Failed to import flash_attn RMSNorm. Falling back to torch RMSNorm.\n"
]
}
],
"source": [
"from modules.gpt import GPT, GPTConfig, GPTTrainConfig\n",
"from transformers import PreTrainedTokenizerFast\n",
"import torch\n",
"from tokenizers import AddedToken"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "2dfd8c6c",
"metadata": {},
"outputs": [],
"source": [
"\n",
"gptconf = GPTConfig()\n",
"gpttrainconf = GPTTrainConfig()"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "e72ff15f",
"metadata": {},
"outputs": [],
"source": [
"gpttrainconf.layer_init = False"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "8fbc647a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[2024-07-31_04:07:30]: number of parameters: 1265M\n"
]
}
],
"source": [
"model = GPT(gptconf, gpttrainconf)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "b7740c1b",
"metadata": {},
"outputs": [],
"source": [
"def _load_model(ckpt_path, tokenizer_path, device, use_tp=False):\n",
" assert ckpt_path is not None\n",
"\n",
" # load tokenizer\n",
" if tokenizer_path is not None:\n",
" assert os.path.exists(tokenizer_path)\n",
" tokenizer = PreTrainedTokenizerFast(\n",
" tokenizer_file=tokenizer_path,\n",
" unk_token=\"[UNK]\",\n",
" pad_token=\"[PAD]\",\n",
" )\n",
" tokenizer.add_special_tokens({\"additional_special_tokens\": [AddedToken(\"\\n\")]})\n",
" else:\n",
" tokenizer = None\n",
"\n",
" # load model\n",
" checkpoint = torch.load(str(ckpt_path), mmap=True)\n",
" model_args = checkpoint[\"model_args\"]\n",
" state_dict = checkpoint[\"model\"]\n",
" # fixup checkpoint\n",
" if \"dropout\" in model_args:\n",
" del model_args[\"dropout\"]\n",
" unwanted_prefix = \"_orig_mod.\"\n",
" for k, v in list(state_dict.items()):\n",
" if k.startswith(unwanted_prefix):\n",
" state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)\n",
" # fix checkpoint based on gqa change\n",
" if \"n_embd\" in model_args:\n",
" n_emb = model_args[\"n_embd\"]\n",
" new_state_dict = {}\n",
" for k, v in state_dict.items():\n",
" if \"c_attn\" in k:\n",
" new_state_dict[k.replace(\"c_attn\", \"c_attn_q\")] = v[: n_emb * 1]\n",
" new_state_dict[k.replace(\"c_attn\", \"c_attn_k\")] = v[\n",
" n_emb * 1 : n_emb * 2\n",
" ]\n",
" new_state_dict[k.replace(\"c_attn\", \"c_attn_v\")] = v[n_emb * 2 :]\n",
" else:\n",
" new_state_dict[k] = v\n",
" model_args[\"n_kv_head\"] = None\n",
" model_args[\"d_head\"] = n_emb // checkpoint[\"model_args\"][\"n_head\"]\n",
" del model_args[\"n_embd\"]\n",
" state_dict = new_state_dict\n",
" del new_state_dict\n",
" # load model\n",
" # with torch.device(\"meta\"):\n",
" gptconf = GPTConfig(**model_args)\n",
" gpttrainconf = GPTTrainConfig()\n",
" model = GPT(gptconf, gpttrainconf)\n",
" # # verify checkpoint\n",
" # extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())\n",
" # extra_keys = set([k for k in extra_keys if not k.endswith(\".attn.bias\")])\n",
" # missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())\n",
" # missing_keys = set([k for k in missing_keys if not k.endswith(\".attn.bias\")])\n",
"\n",
" # # ignore attn keys since we use a pre_hook to convert them\n",
" # extra_keys = [k for k in extra_keys if \"attn\" not in k]\n",
" # missing_keys = [k for k in missing_keys if \"attn\" not in k]\n",
"\n",
" # if len(extra_keys) != 0:\n",
" # raise ValueError(f\"extra keys found: {extra_keys}\")\n",
" # if len(missing_keys) != 0:\n",
" # raise ValueError(f\"missing keys: {missing_keys}\")\n",
"\n",
" # if \"int8\" in str(ckpt_path):\n",
" # print(\"Using int8 weight-only quantization!\")\n",
" # from suno_utils.gpt.quantize import WeightOnlyInt8QuantHandler\n",
"\n",
" # simple_quantizer = WeightOnlyInt8QuantHandler(model)\n",
" # model = simple_quantizer.convert_for_runtime()\n",
" # set up model\n",
" model.load_state_dict(state_dict, strict=False, assign=True)\n",
" print(f\"model loaded: {ckpt_path}\")\n",
" del checkpoint, state_dict\n",
" n_params = model.get_num_params()\n",
" print(f\"model loaded: {round(n_params/1e6,1)}M params\")\n",
" model.eval()\n",
" model.model_args = model_args\n",
"\n",
" if device == \"cuda\" and torch.cuda.is_bf16_supported():\n",
" model.to(\n",
" device, dtype=torch.bfloat16\n",
" ) # this should be fine since it was trained with AMP\n",
" else:\n",
" model.to(device)\n",
" return model, tokenizer"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "ff0afa26",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[2024-07-30_17:47:47]: number of parameters: 32020M\n",
"model loaded: /app/suno/data/dpo/models/model_30b_fix_ft2_20k.pt\n",
"model loaded: 32020.4M params\n"
]
}
],
"source": [
"model, tokenizer = _load_model(\n",
" ckpt_path=model_path, tokenizer_path=tokenizer_path, device=\"cuda\", use_tp=False\n",
")\n",
"cfg = model.config"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "494cbaf4",
"metadata": {},
"outputs": [],
"source": [
"# Import necessary modules and constants\n",
"import torch\n",
"\n",
"# Define missing constants\n",
"train_cfg = cfg # Assuming train_cfg is the same as cfg for this example\n",
"eval_loss_batch_size = 2 # Adjust as needed\n",
"eval_loss_batch_size_tokens = eval_loss_batch_size * cfg.block_size\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"device_type = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
"\n",
"# Build eval_loss_data_sampling_info similar to train_dpo.py\n",
"eval_loss_data_sampling_info = {\n",
" \"cfg\": cfg,\n",
" \"train_cfg\": train_cfg,\n",
" \"batch_size\": eval_loss_batch_size,\n",
" \"batch_size_tokens\": eval_loss_batch_size_tokens,\n",
" \"tokenizer_fp\": tokenizer_path,\n",
" \"device\": device,\n",
" \"device_type\": device_type,\n",
" \"train\": {\n",
" \"data\": test_data,\n",
" \"metas\": test_metas,\n",
" \"infos\": test_info,\n",
" \"artist_to_songs\": None, # Assuming this is not needed for evaluation\n",
" \"names\": [\"train\"], # Assuming a single validation dataset\n",
" \"weights\": [1.0], # Assuming equal weight\n",
" \"idx_lists\": [list(range(len(test_data)))],\n",
" \"all_idx_lists\": list(range(len(test_data))),\n",
" },\n",
" \"val\": {\n",
" \"data\": test_data,\n",
" \"metas\": test_metas,\n",
" \"infos\": test_info,\n",
" \"artist_to_songs\": None, # Assuming this is not needed for evaluation\n",
" \"names\": [\"val\"], # Assuming a single validation dataset\n",
" \"weights\": [1.0], # Assuming equal weight\n",
" \"idx_lists\": [list(range(len(test_data)))],\n",
" \"all_idx_lists\": list(range(len(test_data))),\n",
" },\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "e8f10309",
"metadata": {},
"outputs": [],
"source": [
"from torch.nn import functional as F"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "88f44a3a",
"metadata": {},
"outputs": [],
"source": [
"def get_loss_stored(\n",
" ref_pos_logits_sem, ref_pos_logits_coarse, curr_y, idxs: list, split: str\n",
"):\n",
" \"\"\"This will update the index to loss lookup dictionary under the hood.\n",
"\n",
" curr_y: batch, ncode, time\n",
" \"\"\"\n",
" local_ref_loss_lookups = {}\n",
" local_ref_loss_lookups[\"train\"] = {}\n",
" local_ref_loss_lookups[\"val\"] = {}\n",
" for curr_index, curr_idx in enumerate(idxs):\n",
" loss_dict = {}\n",
" loss_dict[\"z_loss\"] = 0\n",
" for n_semantic in range(cfg.semantic_n_codebooks):\n",
" logits = ref_pos_logits_sem[curr_index][n_semantic]\n",
" loss_dict[f\"semantic_{n_semantic}\"] = (\n",
" F.cross_entropy(\n",
" logits.reshape(-1, logits.size(-1)),\n",
" curr_y[curr_index, n_semantic, :].reshape(-1),\n",
" ignore_index=-1,\n",
" )\n",
" .detach()\n",
" .float()\n",
" .cpu()\n",
" .numpy()\n",
" .item()\n",
" )\n",
" loss_dict[\"z_loss\"] += (\n",
" (torch.logsumexp(logits, dim=-1) ** 2).mean().detach().float().cpu()\n",
" )\n",
" for n_coarse in range(cfg.coarse_n_codebooks):\n",
" logits = ref_pos_logits_coarse[curr_index][n_coarse]\n",
" n2 = cfg.semantic_n_codebooks + n_coarse\n",
" loss_dict[f\"coarse_{n_coarse}\"] = (\n",
" F.cross_entropy(\n",
" logits.reshape(-1, logits.size(-1)),\n",
" curr_y[curr_index, n2, :].reshape(-1),\n",
" ignore_index=-1,\n",
" )\n",
" .detach()\n",
" .float()\n",
" .cpu()\n",
" .numpy()\n",
" .item()\n",
" )\n",
" loss_dict[\"z_loss\"] += (\n",
" (torch.logsumexp(logits, dim=-1) ** 2).mean().detach().float().cpu()\n",
" )\n",
" loss_dict[\"orig_idx\"] = curr_idx\n",
" loss_dict[\"z_loss\"] = loss_dict[\"z_loss\"].numpy().item()\n",
" local_ref_loss_lookups[split][curr_idx] = loss_dict\n",
" return local_ref_loss_lookups"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "7dcc483f",
"metadata": {},
"outputs": [],
"source": [
"def get_loss_displayed(test_start_rank, eval_loss_batch_size):\n",
" eval_loss_batch_size_tokens = eval_loss_batch_size * cfg.block_size\n",
" eval_loss_data_sampling_info[\"batch_size\"] = eval_loss_batch_size\n",
" eval_loss_data_sampling_info[\"batch_size_tokens\"] = eval_loss_batch_size_tokens\n",
" batch_row_idx_list = list(range(len(test_data)))[\n",
" test_start_rank * eval_loss_batch_size : (test_start_rank + 1)\n",
" * eval_loss_batch_size\n",
" ]\n",
" print(batch_row_idx_list)\n",
" split = \"train\"\n",
"\n",
" output_idx_list, _, X, Y = get_batch(\n",
" eval_loss_data_sampling_info,\n",
" split,\n",
" min_text_offs=0,\n",
" suppress_text=False,\n",
" use_private=False,\n",
" dummy_data=False,\n",
" inference=True,\n",
" return_idx=True,\n",
" abs_row_idx=batch_row_idx_list, # use the abs idx of the whole dataset\n",
" )\n",
"\n",
" X = X.to(device)\n",
" Y = Y.to(device)\n",
" with torch.no_grad():\n",
" model.eval()\n",
" ref_logits_smenatic, ref_logits_coarse = model(\n",
" X, y=Y, return_logits=True, last_only=False\n",
" )\n",
"\n",
" local_ref_loss_lookups_4 = get_loss_stored(\n",
" ref_logits_smenatic,\n",
" ref_logits_coarse,\n",
" curr_y=Y,\n",
" idxs=batch_row_idx_list,\n",
" split=split,\n",
" )\n",
"\n",
" # Convert the dictionary to a DataFrame\n",
" df = pd.DataFrame.from_dict(local_ref_loss_lookups_4[\"train\"], orient=\"index\")\n",
"\n",
" # Reset the index to make the original index a column\n",
" df = df.reset_index().rename(columns={\"index\": \"sample_id\"})\n",
"\n",
" # Display the DataFrame\n",
" display(df)\n",
" return X"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "66fa97bc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[0]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/tony/anaconda3/envs/suno_env/lib/python3.10/site-packages/torch/backends/cuda/__init__.py:393: FutureWarning: torch.backends.cuda.sdp_kernel() is deprecated. In the future, this context manager will be removed. Please see, torch.nn.attention.sdpa_kernel() for the new context manager, with updated signature.\n",
" warnings.warn(\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" sample_id | \n",
" z_loss | \n",
" semantic_0 | \n",
" coarse_0 | \n",
" orig_idx | \n",
" coarse_1 | \n",
" coarse_2 | \n",
" coarse_3 | \n",
" coarse_4 | \n",
" coarse_5 | \n",
" coarse_6 | \n",
" coarse_7 | \n",
" coarse_8 | \n",
" coarse_9 | \n",
" coarse_10 | \n",
" coarse_11 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 0 | \n",
" 324.25 | \n",
" 0.992188 | \n",
" 1.757812 | \n",
" 0 | \n",
" 2.734375 | \n",
" 3.28125 | \n",
" 3.609375 | \n",
" 3.8125 | \n",
" 3.796875 | \n",
" 3.71875 | \n",
" 3.765625 | \n",
" 3.734375 | \n",
" 3.78125 | \n",
" 3.828125 | \n",
" 3.921875 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" sample_id | \n",
" z_loss | \n",
" semantic_0 | \n",
" coarse_0 | \n",
" orig_idx | \n",
" coarse_1 | \n",
" coarse_2 | \n",
" coarse_3 | \n",
" coarse_4 | \n",
" coarse_5 | \n",
" coarse_6 | \n",
" coarse_7 | \n",
" coarse_8 | \n",
" coarse_9 | \n",
" coarse_10 | \n",
" coarse_11 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 1 | \n",
" 429.3125 | \n",
" 1.421875 | \n",
" 2.515625 | \n",
" 1 | \n",
" 3.25 | \n",
" 4.15625 | \n",
" 4.3125 | \n",
" 4.3125 | \n",
" 4.5625 | \n",
" 4.40625 | \n",
" 4.4375 | \n",
" 4.5 | \n",
" 4.5 | \n",
" 4.65625 | \n",
" 4.71875 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" sample_id | \n",
" z_loss | \n",
" semantic_0 | \n",
" coarse_0 | \n",
" orig_idx | \n",
" coarse_1 | \n",
" coarse_2 | \n",
" coarse_3 | \n",
" coarse_4 | \n",
" coarse_5 | \n",
" coarse_6 | \n",
" coarse_7 | \n",
" coarse_8 | \n",
" coarse_9 | \n",
" coarse_10 | \n",
" coarse_11 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 0 | \n",
" 324.2500 | \n",
" 0.992188 | \n",
" 1.757812 | \n",
" 0 | \n",
" 2.734375 | \n",
" 3.28125 | \n",
" 3.609375 | \n",
" 3.8125 | \n",
" 3.796875 | \n",
" 3.71875 | \n",
" 3.765625 | \n",
" 3.734375 | \n",
" 3.78125 | \n",
" 3.828125 | \n",
" 3.921875 | \n",
"
\n",
" \n",
" | 1 | \n",
" 1 | \n",
" 429.3125 | \n",
" 1.421875 | \n",
" 2.515625 | \n",
" 1 | \n",
" 3.250000 | \n",
" 4.15625 | \n",
" 4.312500 | \n",
" 4.3125 | \n",
" 4.562500 | \n",
" 4.40625 | \n",
" 4.437500 | \n",
" 4.500000 | \n",
" 4.50000 | \n",
" 4.656250 | \n",
" 4.718750 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" sample_id | \n",
" z_loss | \n",
" semantic_0 | \n",
" coarse_0 | \n",
" orig_idx | \n",
" coarse_1 | \n",
" coarse_2 | \n",
" coarse_3 | \n",
" coarse_4 | \n",
" coarse_5 | \n",
" coarse_6 | \n",
" coarse_7 | \n",
" coarse_8 | \n",
" coarse_9 | \n",
" coarse_10 | \n",
" coarse_11 | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 0 | \n",
" 324.25000 | \n",
" 0.992188 | \n",
" 1.757812 | \n",
" 0 | \n",
" 2.734375 | \n",
" 3.281250 | \n",
" 3.609375 | \n",
" 3.812500 | \n",
" 3.796875 | \n",
" 3.718750 | \n",
" 3.765625 | \n",
" 3.734375 | \n",
" 3.781250 | \n",
" 3.828125 | \n",
" 3.921875 | \n",
"
\n",
" \n",
" | 1 | \n",
" 1 | \n",
" 429.31250 | \n",
" 1.421875 | \n",
" 2.515625 | \n",
" 1 | \n",
" 3.250000 | \n",
" 4.156250 | \n",
" 4.312500 | \n",
" 4.312500 | \n",
" 4.562500 | \n",
" 4.406250 | \n",
" 4.437500 | \n",
" 4.500000 | \n",
" 4.500000 | \n",
" 4.656250 | \n",
" 4.718750 | \n",
"
\n",
" \n",
" | 2 | \n",
" 2 | \n",
" 335.40625 | \n",
" 1.132812 | \n",
" 1.492188 | \n",
" 2 | \n",
" 2.609375 | \n",
" 3.296875 | \n",
" 3.546875 | \n",
" 3.515625 | \n",
" 3.593750 | \n",
" 3.578125 | \n",
" 3.593750 | \n",
" 3.578125 | \n",
" 3.609375 | \n",
" 3.625000 | \n",
" 3.687500 | \n",
"
\n",
" \n",
" | 3 | \n",
" 3 | \n",
" 197.03125 | \n",
" 1.289062 | \n",
" 1.937500 | \n",
" 3 | \n",
" 3.343750 | \n",
" 4.250000 | \n",
" 4.375000 | \n",
" 4.531250 | \n",
" 4.500000 | \n",
" 4.500000 | \n",
" 4.375000 | \n",
" 4.375000 | \n",
" 4.406250 | \n",
" 4.406250 | \n",
" 4.437500 | \n",
"
\n",
" \n",
"
\n",
"