gptj_prompt_generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "747464e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import datetime\n",
    "import json\n",
    "from gptj_utils import GPTJ_PrefixTune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "19c34b64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load soft prompt tuning model...\n",
      "Start freezing params...\n",
      "========================\n",
      "transformer.wte.learned_embedding\n",
      "exclude transformer.wte.learned_embedding from freezing!\n",
      "transformer.wte.wte.weight\n",
      "transformer.h.0.ln_1.weight\n",
      "transformer.h.0.ln_1.bias\n",
      "transformer.h.0.attn.attention.k_proj.weight\n",
      "transformer.h.0.attn.attention.v_proj.weight\n",
      "transformer.h.0.attn.attention.q_proj.weight\n",
      "transformer.h.0.attn.attention.out_proj.weight\n",
      "transformer.h.0.mlp.c_fc.weight\n",
      "transformer.h.0.mlp.c_fc.bias\n",
      "transformer.h.0.mlp.c_proj.weight\n",
      "transformer.h.0.mlp.c_proj.bias\n",
      "transformer.h.1.ln_1.weight\n",
      "transformer.h.1.ln_1.bias\n",
      "transformer.h.1.attn.attention.k_proj.weight\n",
      "transformer.h.1.attn.attention.v_proj.weight\n",
      "transformer.h.1.attn.attention.q_proj.weight\n",
      "transformer.h.1.attn.attention.out_proj.weight\n",
      "transformer.h.1.mlp.c_fc.weight\n",
      "transformer.h.1.mlp.c_fc.bias\n",
      "transformer.h.1.mlp.c_proj.weight\n",
      "transformer.h.1.mlp.c_proj.bias\n",
      "transformer.h.2.ln_1.weight\n",
      "transformer.h.2.ln_1.bias\n",
      "transformer.h.2.attn.attention.k_proj.weight\n",
      "transformer.h.2.attn.attention.v_proj.weight\n",
      "transformer.h.2.attn.attention.q_proj.weight\n",
      "transformer.h.2.attn.attention.out_proj.weight\n",
      "transformer.h.2.mlp.c_fc.weight\n",
      "transformer.h.2.mlp.c_fc.bias\n",
      "transformer.h.2.mlp.c_proj.weight\n",
      "transformer.h.2.mlp.c_proj.bias\n",
      "transformer.h.3.ln_1.weight\n",
      "transformer.h.3.ln_1.bias\n",
      "transformer.h.3.attn.attention.k_proj.weight\n",
      "transformer.h.3.attn.attention.v_proj.weight\n",
      "transformer.h.3.attn.attention.q_proj.weight\n",
      "transformer.h.3.attn.attention.out_proj.weight\n",
      "transformer.h.3.mlp.c_fc.weight\n",
      "transformer.h.3.mlp.c_fc.bias\n",
      "transformer.h.3.mlp.c_proj.weight\n",
      "transformer.h.3.mlp.c_proj.bias\n",
      "transformer.h.4.ln_1.weight\n",
      "transformer.h.4.ln_1.bias\n",
      "transformer.h.4.attn.attention.k_proj.weight\n",
      "transformer.h.4.attn.attention.v_proj.weight\n",
      "transformer.h.4.attn.attention.q_proj.weight\n",
      "transformer.h.4.attn.attention.out_proj.weight\n",
      "transformer.h.4.mlp.c_fc.weight\n",
      "transformer.h.4.mlp.c_fc.bias\n",
      "transformer.h.4.mlp.c_proj.weight\n",
      "transformer.h.4.mlp.c_proj.bias\n",
      "transformer.h.5.ln_1.weight\n",
      "transformer.h.5.ln_1.bias\n",
      "transformer.h.5.attn.attention.k_proj.weight\n",
      "transformer.h.5.attn.attention.v_proj.weight\n",
      "transformer.h.5.attn.attention.q_proj.weight\n",
      "transformer.h.5.attn.attention.out_proj.weight\n",
      "transformer.h.5.mlp.c_fc.weight\n",
      "transformer.h.5.mlp.c_fc.bias\n",
      "transformer.h.5.mlp.c_proj.weight\n",
      "transformer.h.5.mlp.c_proj.bias\n",
      "transformer.h.6.ln_1.weight\n",
      "transformer.h.6.ln_1.bias\n",
      "transformer.h.6.attn.attention.k_proj.weight\n",
      "transformer.h.6.attn.attention.v_proj.weight\n",
      "transformer.h.6.attn.attention.q_proj.weight\n",
      "transformer.h.6.attn.attention.out_proj.weight\n",
      "transformer.h.6.mlp.c_fc.weight\n",
      "transformer.h.6.mlp.c_fc.bias\n",
      "transformer.h.6.mlp.c_proj.weight\n",
      "transformer.h.6.mlp.c_proj.bias\n",
      "transformer.h.7.ln_1.weight\n",
      "transformer.h.7.ln_1.bias\n",
      "transformer.h.7.attn.attention.k_proj.weight\n",
      "transformer.h.7.attn.attention.v_proj.weight\n",
      "transformer.h.7.attn.attention.q_proj.weight\n",
      "transformer.h.7.attn.attention.out_proj.weight\n",
      "transformer.h.7.mlp.c_fc.weight\n",
      "transformer.h.7.mlp.c_fc.bias\n",
      "transformer.h.7.mlp.c_proj.weight\n",
      "transformer.h.7.mlp.c_proj.bias\n",
      "transformer.h.8.ln_1.weight\n",
      "transformer.h.8.ln_1.bias\n",
      "transformer.h.8.attn.attention.k_proj.weight\n",
      "transformer.h.8.attn.attention.v_proj.weight\n",
      "transformer.h.8.attn.attention.q_proj.weight\n",
      "transformer.h.8.attn.attention.out_proj.weight\n",
      "transformer.h.8.mlp.c_fc.weight\n",
      "transformer.h.8.mlp.c_fc.bias\n",
      "transformer.h.8.mlp.c_proj.weight\n",
      "transformer.h.8.mlp.c_proj.bias\n",
      "transformer.h.9.ln_1.weight\n",
      "transformer.h.9.ln_1.bias\n",
      "transformer.h.9.attn.attention.k_proj.weight\n",
      "transformer.h.9.attn.attention.v_proj.weight\n",
      "transformer.h.9.attn.attention.q_proj.weight\n",
      "transformer.h.9.attn.attention.out_proj.weight\n",
      "transformer.h.9.mlp.c_fc.weight\n",
      "transformer.h.9.mlp.c_fc.bias\n",
      "transformer.h.9.mlp.c_proj.weight\n",
      "transformer.h.9.mlp.c_proj.bias\n",
      "transformer.h.10.ln_1.weight\n",
      "transformer.h.10.ln_1.bias\n",
      "transformer.h.10.attn.attention.k_proj.weight\n",
      "transformer.h.10.attn.attention.v_proj.weight\n",
      "transformer.h.10.attn.attention.q_proj.weight\n",
      "transformer.h.10.attn.attention.out_proj.weight\n",
      "transformer.h.10.mlp.c_fc.weight\n",
      "transformer.h.10.mlp.c_fc.bias\n",
      "transformer.h.10.mlp.c_proj.weight\n",
      "transformer.h.10.mlp.c_proj.bias\n",
      "transformer.h.11.ln_1.weight\n",
      "transformer.h.11.ln_1.bias\n",
      "transformer.h.11.attn.attention.k_proj.weight\n",
      "transformer.h.11.attn.attention.v_proj.weight\n",
      "transformer.h.11.attn.attention.q_proj.weight\n",
      "transformer.h.11.attn.attention.out_proj.weight\n",
      "transformer.h.11.mlp.c_fc.weight\n",
      "transformer.h.11.mlp.c_fc.bias\n",
      "transformer.h.11.mlp.c_proj.weight\n",
      "transformer.h.11.mlp.c_proj.bias\n",
      "transformer.h.12.ln_1.weight\n",
      "transformer.h.12.ln_1.bias\n",
      "transformer.h.12.attn.attention.k_proj.weight\n",
      "transformer.h.12.attn.attention.v_proj.weight\n",
      "transformer.h.12.attn.attention.q_proj.weight\n",
      "transformer.h.12.attn.attention.out_proj.weight\n",
      "transformer.h.12.mlp.c_fc.weight\n",
      "transformer.h.12.mlp.c_fc.bias\n",
      "transformer.h.12.mlp.c_proj.weight\n",
      "transformer.h.12.mlp.c_proj.bias\n",
      "transformer.h.13.ln_1.weight\n",
      "transformer.h.13.ln_1.bias\n",
      "transformer.h.13.attn.attention.k_proj.weight\n",
      "transformer.h.13.attn.attention.v_proj.weight\n",
      "transformer.h.13.attn.attention.q_proj.weight\n",
      "transformer.h.13.attn.attention.out_proj.weight\n",
      "transformer.h.13.mlp.c_fc.weight\n",
      "transformer.h.13.mlp.c_fc.bias\n",
      "transformer.h.13.mlp.c_proj.weight\n",
      "transformer.h.13.mlp.c_proj.bias\n",
      "transformer.h.14.ln_1.weight\n",
      "transformer.h.14.ln_1.bias\n",
      "transformer.h.14.attn.attention.k_proj.weight\n",
      "transformer.h.14.attn.attention.v_proj.weight\n",
      "transformer.h.14.attn.attention.q_proj.weight\n",
      "transformer.h.14.attn.attention.out_proj.weight\n",
      "transformer.h.14.mlp.c_fc.weight\n",
      "transformer.h.14.mlp.c_fc.bias\n",
      "transformer.h.14.mlp.c_proj.weight\n",
      "transformer.h.14.mlp.c_proj.bias\n",
      "transformer.h.15.ln_1.weight\n",
      "transformer.h.15.ln_1.bias\n",
      "transformer.h.15.attn.attention.k_proj.weight\n",
      "transformer.h.15.attn.attention.v_proj.weight\n",
      "transformer.h.15.attn.attention.q_proj.weight\n",
      "transformer.h.15.attn.attention.out_proj.weight\n",
      "transformer.h.15.mlp.c_fc.weight\n",
      "transformer.h.15.mlp.c_fc.bias\n",
      "transformer.h.15.mlp.c_proj.weight\n",
      "transformer.h.15.mlp.c_proj.bias\n",
      "transformer.h.16.ln_1.weight\n",
      "transformer.h.16.ln_1.bias\n",
      "transformer.h.16.attn.attention.k_proj.weight\n",
      "transformer.h.16.attn.attention.v_proj.weight\n",
      "transformer.h.16.attn.attention.q_proj.weight\n",
      "transformer.h.16.attn.attention.out_proj.weight\n",
      "transformer.h.16.mlp.c_fc.weight\n",
      "transformer.h.16.mlp.c_fc.bias\n",
      "transformer.h.16.mlp.c_proj.weight\n",
      "transformer.h.16.mlp.c_proj.bias\n",
      "transformer.h.17.ln_1.weight\n",
      "transformer.h.17.ln_1.bias\n",
      "transformer.h.17.attn.attention.k_proj.weight\n",
      "transformer.h.17.attn.attention.v_proj.weight\n",
      "transformer.h.17.attn.attention.q_proj.weight\n",
      "transformer.h.17.attn.attention.out_proj.weight\n",
      "transformer.h.17.mlp.c_fc.weight\n",
      "transformer.h.17.mlp.c_fc.bias\n",
      "transformer.h.17.mlp.c_proj.weight\n",
      "transformer.h.17.mlp.c_proj.bias\n",
      "transformer.h.18.ln_1.weight\n",
      "transformer.h.18.ln_1.bias\n",
      "transformer.h.18.attn.attention.k_proj.weight\n",
      "transformer.h.18.attn.attention.v_proj.weight\n",
      "transformer.h.18.attn.attention.q_proj.weight\n",
      "transformer.h.18.attn.attention.out_proj.weight\n",
      "transformer.h.18.mlp.c_fc.weight\n",
      "transformer.h.18.mlp.c_fc.bias\n",
      "transformer.h.18.mlp.c_proj.weight\n",
      "transformer.h.18.mlp.c_proj.bias\n",
      "transformer.h.19.ln_1.weight\n",
      "transformer.h.19.ln_1.bias\n",
      "transformer.h.19.attn.attention.k_proj.weight\n",
      "transformer.h.19.attn.attention.v_proj.weight\n",
      "transformer.h.19.attn.attention.q_proj.weight\n",
      "transformer.h.19.attn.attention.out_proj.weight\n",
      "transformer.h.19.mlp.c_fc.weight\n",
      "transformer.h.19.mlp.c_fc.bias\n",
      "transformer.h.19.mlp.c_proj.weight\n",
      "transformer.h.19.mlp.c_proj.bias\n",
      "transformer.h.20.ln_1.weight\n",
      "transformer.h.20.ln_1.bias\n",
      "transformer.h.20.attn.attention.k_proj.weight\n",
      "transformer.h.20.attn.attention.v_proj.weight\n",
      "transformer.h.20.attn.attention.q_proj.weight\n",
      "transformer.h.20.attn.attention.out_proj.weight\n",
      "transformer.h.20.mlp.c_fc.weight\n",
      "transformer.h.20.mlp.c_fc.bias\n",
      "transformer.h.20.mlp.c_proj.weight\n",
      "transformer.h.20.mlp.c_proj.bias\n",
      "transformer.h.21.ln_1.weight\n",
      "transformer.h.21.ln_1.bias\n",
      "transformer.h.21.attn.attention.k_proj.weight\n",
      "transformer.h.21.attn.attention.v_proj.weight\n",
      "transformer.h.21.attn.attention.q_proj.weight\n",
      "transformer.h.21.attn.attention.out_proj.weight\n",
      "transformer.h.21.mlp.c_fc.weight\n",
      "transformer.h.21.mlp.c_fc.bias\n",
      "transformer.h.21.mlp.c_proj.weight\n",
      "transformer.h.21.mlp.c_proj.bias\n",
      "transformer.h.22.ln_1.weight\n",
      "transformer.h.22.ln_1.bias\n",
      "transformer.h.22.attn.attention.k_proj.weight\n",
      "transformer.h.22.attn.attention.v_proj.weight\n",
      "transformer.h.22.attn.attention.q_proj.weight\n",
      "transformer.h.22.attn.attention.out_proj.weight\n",
      "transformer.h.22.mlp.c_fc.weight\n",
      "transformer.h.22.mlp.c_fc.bias\n",
      "transformer.h.22.mlp.c_proj.weight\n",
      "transformer.h.22.mlp.c_proj.bias\n",
      "transformer.h.23.ln_1.weight\n",
      "transformer.h.23.ln_1.bias\n",
      "transformer.h.23.attn.attention.k_proj.weight\n",
      "transformer.h.23.attn.attention.v_proj.weight\n",
      "transformer.h.23.attn.attention.q_proj.weight\n",
      "transformer.h.23.attn.attention.out_proj.weight\n",
      "transformer.h.23.mlp.c_fc.weight\n",
      "transformer.h.23.mlp.c_fc.bias\n",
      "transformer.h.23.mlp.c_proj.weight\n",
      "transformer.h.23.mlp.c_proj.bias\n",
      "transformer.h.24.ln_1.weight\n",
      "transformer.h.24.ln_1.bias\n",
      "transformer.h.24.attn.attention.k_proj.weight\n",
      "transformer.h.24.attn.attention.v_proj.weight\n",
      "transformer.h.24.attn.attention.q_proj.weight\n",
      "transformer.h.24.attn.attention.out_proj.weight\n",
      "transformer.h.24.mlp.c_fc.weight\n",
      "transformer.h.24.mlp.c_fc.bias\n",
      "transformer.h.24.mlp.c_proj.weight\n",
      "transformer.h.24.mlp.c_proj.bias\n",
      "transformer.h.25.ln_1.weight\n",
      "transformer.h.25.ln_1.bias\n",
      "transformer.h.25.attn.attention.k_proj.weight\n",
      "transformer.h.25.attn.attention.v_proj.weight\n",
      "transformer.h.25.attn.attention.q_proj.weight\n",
      "transformer.h.25.attn.attention.out_proj.weight\n",
      "transformer.h.25.mlp.c_fc.weight\n",
      "transformer.h.25.mlp.c_fc.bias\n",
      "transformer.h.25.mlp.c_proj.weight\n",
      "transformer.h.25.mlp.c_proj.bias\n",
      "transformer.h.26.ln_1.weight\n",
      "transformer.h.26.ln_1.bias\n",
      "transformer.h.26.attn.attention.k_proj.weight\n",
      "transformer.h.26.attn.attention.v_proj.weight\n",
      "transformer.h.26.attn.attention.q_proj.weight\n",
      "transformer.h.26.attn.attention.out_proj.weight\n",
      "transformer.h.26.mlp.c_fc.weight\n",
      "transformer.h.26.mlp.c_fc.bias\n",
      "transformer.h.26.mlp.c_proj.weight\n",
      "transformer.h.26.mlp.c_proj.bias\n",
      "transformer.h.27.ln_1.weight\n",
      "transformer.h.27.ln_1.bias\n",
      "transformer.h.27.attn.attention.k_proj.weight\n",
      "transformer.h.27.attn.attention.v_proj.weight\n",
      "transformer.h.27.attn.attention.q_proj.weight\n",
      "transformer.h.27.attn.attention.out_proj.weight\n",
      "transformer.h.27.mlp.c_fc.weight\n",
      "transformer.h.27.mlp.c_fc.bias\n",
      "transformer.h.27.mlp.c_proj.weight\n",
      "transformer.h.27.mlp.c_proj.bias\n",
      "transformer.ln_f.weight\n",
      "transformer.ln_f.bias\n",
      "lm_head.weight\n",
      "lm_head.bias\n",
      "========================\n",
      "[2021-09-28 21:43:49,674] [INFO] [distributed.py:46:init_distributed] Initializing torch distributed with backend: nccl\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-09-28 21:43:49,738] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed info: version=0.4.1, git-hash=unknown, git-branch=unknown\n",
      "[2021-09-28 21:44:02,848] [INFO] [utils.py:11:_initialize_parameter_parallel_groups] data_parallel_size: 1, parameter_parallel_size: 1\n",
      "[2021-09-28 21:44:02,925] [INFO] [engine.py:172:__init__] DeepSpeed Flops Profiler Enabled: False\n",
      "Using /home/fellow/.cache/torch_extensions as PyTorch extensions root...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file /home/fellow/.cache/torch_extensions/cpu_adam/build.ninja...\n",
      "Building extension module cpu_adam...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "Loading extension module cpu_adam...\n",
      "Time to load cpu_adam op: 1.6930701732635498 seconds\n",
      "[2021-09-28 21:44:06,588] [INFO] [engine.py:701:_configure_optimizer] Using DeepSpeed Optimizer param name adam as basic optimizer\n",
      "[2021-09-28 21:44:06,589] [INFO] [engine.py:706:_configure_optimizer] DeepSpeed Basic Optimizer = DeepSpeedCPUAdam\n",
      "Checking ZeRO support for optimizer=DeepSpeedCPUAdam type=<class 'deepspeed.ops.adam.cpu_adam.DeepSpeedCPUAdam'>\n",
      "[2021-09-28 21:44:06,590] [INFO] [logging.py:60:log_dist] [Rank 0] Creating fp16 ZeRO stage 2 optimizer\n",
      "[2021-09-28 21:44:06,591] [INFO] [stage2.py:105:__init__] Reduce bucket size 200000000\n",
      "[2021-09-28 21:44:06,591] [INFO] [stage2.py:106:__init__] Allgather bucket size 200000000\n",
      "[2021-09-28 21:44:06,592] [INFO] [stage2.py:107:__init__] CPU Offload: True\n",
      "Using /home/fellow/.cache/torch_extensions as PyTorch extensions root...\n",
      "Emitting ninja build file /home/fellow/.cache/torch_extensions/utils/build.ninja...\n",
      "Building extension module utils...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "Loading extension module utils...\n",
      "Time to load utils op: 1.2017416954040527 seconds\n",
      "[2021-09-28 21:44:56,694] [INFO] [stage2.py:409:__init__] optimizer state initialized\n",
      "[2021-09-28 21:44:56,697] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed Final Optimizer = adam\n",
      "[2021-09-28 21:44:56,698] [INFO] [engine.py:504:_configure_lr_scheduler] DeepSpeed using configured LR scheduler = WarmupLR\n",
      "[2021-09-28 21:44:56,699] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupLR object at 0x7f9e3462d3d0>\n",
      "[2021-09-28 21:44:56,699] [INFO] [logging.py:60:log_dist] [Rank 0] step=0, skipped=0, lr=[1e-05], mom=[[0.8, 0.999]]\n",
      "[2021-09-28 21:44:56,700] [INFO] [config.py:900:print] DeepSpeedEngine configuration:\n",
      "[2021-09-28 21:44:56,702] [INFO] [config.py:904:print]   activation_checkpointing_config  {\n",
      "    \"partition_activations\": false, \n",
      "    \"contiguous_memory_optimization\": false, \n",
      "    \"cpu_checkpointing\": false, \n",
      "    \"number_checkpoints\": null, \n",
      "    \"synchronize_checkpoint_boundary\": false, \n",
      "    \"profile\": false\n",
      "}\n",
      "[2021-09-28 21:44:56,702] [INFO] [config.py:904:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}\n",
      "[2021-09-28 21:44:56,703] [INFO] [config.py:904:print]   allreduce_always_fp32 ........ False\n",
      "[2021-09-28 21:44:56,703] [INFO] [config.py:904:print]   amp_enabled .................. False\n",
      "[2021-09-28 21:44:56,704] [INFO] [config.py:904:print]   amp_params ................... False\n",
      "[2021-09-28 21:44:56,705] [INFO] [config.py:904:print]   checkpoint_tag_validation_enabled  True\n",
      "[2021-09-28 21:44:56,705] [INFO] [config.py:904:print]   checkpoint_tag_validation_fail  False\n",
      "[2021-09-28 21:44:56,706] [INFO] [config.py:904:print]   disable_allgather ............ False\n",
      "[2021-09-28 21:44:56,706] [INFO] [config.py:904:print]   dump_state ................... False\n",
      "[2021-09-28 21:44:56,706] [INFO] [config.py:904:print]   dynamic_loss_scale_args ...... {'init_scale': 4096, 'scale_window': 1000, 'delayed_shift': 2, 'min_scale': 1}\n",
      "[2021-09-28 21:44:56,707] [INFO] [config.py:904:print]   eigenvalue_enabled ........... False\n",
      "[2021-09-28 21:44:56,707] [INFO] [config.py:904:print]   eigenvalue_gas_boundary_resolution  1\n",
      "[2021-09-28 21:44:56,707] [INFO] [config.py:904:print]   eigenvalue_layer_name ........ bert.encoder.layer\n",
      "[2021-09-28 21:44:56,708] [INFO] [config.py:904:print]   eigenvalue_layer_num ......... 0\n",
      "[2021-09-28 21:44:56,708] [INFO] [config.py:904:print]   eigenvalue_max_iter .......... 100\n",
      "[2021-09-28 21:44:56,709] [INFO] [config.py:904:print]   eigenvalue_stability ......... 1e-06\n",
      "[2021-09-28 21:44:56,709] [INFO] [config.py:904:print]   eigenvalue_tol ............... 0.01\n",
      "[2021-09-28 21:44:56,709] [INFO] [config.py:904:print]   eigenvalue_verbose ........... False\n",
      "[2021-09-28 21:44:56,710] [INFO] [config.py:904:print]   elasticity_enabled ........... False\n",
      "[2021-09-28 21:44:56,710] [INFO] [config.py:904:print]   flops_profiler_config ........ {\n",
      "    \"enabled\": false, \n",
      "    \"profile_step\": 1, \n",
      "    \"module_depth\": -1, \n",
      "    \"top_modules\": 1, \n",
      "    \"detailed\": true, \n",
      "    \"output_file\": null\n",
      "}\n",
      "[2021-09-28 21:44:56,710] [INFO] [config.py:904:print]   fp16_enabled ................. True\n",
      "[2021-09-28 21:44:56,711] [INFO] [config.py:904:print]   fp16_mixed_quantize .......... False\n",
      "[2021-09-28 21:44:56,711] [INFO] [config.py:904:print]   global_rank .................. 0\n",
      "[2021-09-28 21:44:56,712] [INFO] [config.py:904:print]   gradient_accumulation_steps .. 1\n",
      "[2021-09-28 21:44:56,712] [INFO] [config.py:904:print]   gradient_clipping ............ 0.0\n",
      "[2021-09-28 21:44:56,712] [INFO] [config.py:904:print]   gradient_predivide_factor .... 1.0\n",
      "[2021-09-28 21:44:56,713] [INFO] [config.py:904:print]   initial_dynamic_scale ........ 4096\n",
      "[2021-09-28 21:44:56,713] [INFO] [config.py:904:print]   loss_scale ................... 0\n",
      "[2021-09-28 21:44:56,716] [INFO] [config.py:904:print]   memory_breakdown ............. False\n",
      "[2021-09-28 21:44:56,716] [INFO] [config.py:904:print]   optimizer_legacy_fusion ...... False\n",
      "[2021-09-28 21:44:56,717] [INFO] [config.py:904:print]   optimizer_name ............... adam\n",
      "[2021-09-28 21:44:56,717] [INFO] [config.py:904:print]   optimizer_params ............. {'lr': 1e-05, 'betas': [0.8, 0.999], 'eps': 1e-08, 'weight_decay': 3e-07}\n",
      "[2021-09-28 21:44:56,717] [INFO] [config.py:904:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}\n",
      "[2021-09-28 21:44:56,718] [INFO] [config.py:904:print]   pld_enabled .................. False\n",
      "[2021-09-28 21:44:56,718] [INFO] [config.py:904:print]   pld_params ................... False\n",
      "[2021-09-28 21:44:56,718] [INFO] [config.py:904:print]   prescale_gradients ........... False\n",
      "[2021-09-28 21:44:56,719] [INFO] [config.py:904:print]   quantize_change_rate ......... 0.001\n",
      "[2021-09-28 21:44:56,719] [INFO] [config.py:904:print]   quantize_groups .............. 1\n",
      "[2021-09-28 21:44:56,719] [INFO] [config.py:904:print]   quantize_offset .............. 1000\n",
      "[2021-09-28 21:44:56,720] [INFO] [config.py:904:print]   quantize_period .............. 1000\n",
      "[2021-09-28 21:44:56,720] [INFO] [config.py:904:print]   quantize_rounding ............ 0\n",
      "[2021-09-28 21:44:56,721] [INFO] [config.py:904:print]   quantize_start_bits .......... 16\n",
      "[2021-09-28 21:44:56,721] [INFO] [config.py:904:print]   quantize_target_bits ......... 8\n",
      "[2021-09-28 21:44:56,721] [INFO] [config.py:904:print]   quantize_training_enabled .... False\n",
      "[2021-09-28 21:44:56,722] [INFO] [config.py:904:print]   quantize_type ................ 0\n",
      "[2021-09-28 21:44:56,722] [INFO] [config.py:904:print]   quantize_verbose ............. False\n",
      "[2021-09-28 21:44:56,722] [INFO] [config.py:904:print]   scheduler_name ............... WarmupLR\n",
      "[2021-09-28 21:44:56,723] [INFO] [config.py:904:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 3e-05, 'warmup_num_steps': 500}\n",
      "[2021-09-28 21:44:56,723] [INFO] [config.py:904:print]   sparse_attention ............. None\n",
      "[2021-09-28 21:44:56,724] [INFO] [config.py:904:print]   sparse_gradients_enabled ..... False\n",
      "[2021-09-28 21:44:56,724] [INFO] [config.py:904:print]   steps_per_print .............. 2000\n",
      "[2021-09-28 21:44:56,724] [INFO] [config.py:904:print]   tensorboard_enabled .......... False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-09-28 21:44:56,725] [INFO] [config.py:904:print]   tensorboard_job_name ......... DeepSpeedJobName\n",
      "[2021-09-28 21:44:56,725] [INFO] [config.py:904:print]   tensorboard_output_path ...... \n",
      "[2021-09-28 21:44:56,725] [INFO] [config.py:904:print]   train_batch_size ............. 16\n",
      "[2021-09-28 21:44:56,726] [INFO] [config.py:904:print]   train_micro_batch_size_per_gpu  16\n",
      "[2021-09-28 21:44:56,726] [INFO] [config.py:904:print]   use_quantizer_kernel ......... False\n",
      "[2021-09-28 21:44:56,726] [INFO] [config.py:904:print]   wall_clock_breakdown ......... False\n",
      "[2021-09-28 21:44:56,727] [INFO] [config.py:904:print]   world_size ................... 1\n",
      "[2021-09-28 21:44:56,727] [INFO] [config.py:904:print]   zero_allow_untested_optimizer  False\n",
      "[2021-09-28 21:44:56,728] [INFO] [config.py:904:print]   zero_config .................. {\n",
      "    \"stage\": 2, \n",
      "    \"contiguous_gradients\": true, \n",
      "    \"reduce_scatter\": true, \n",
      "    \"reduce_bucket_size\": 2.000000e+08, \n",
      "    \"allgather_partitions\": true, \n",
      "    \"allgather_bucket_size\": 2.000000e+08, \n",
      "    \"overlap_comm\": true, \n",
      "    \"load_from_fp32_weights\": true, \n",
      "    \"elastic_checkpoint\": true, \n",
      "    \"offload_param\": null, \n",
      "    \"offload_optimizer\": {\n",
      "        \"device\": \"cpu\", \n",
      "        \"nvme_path\": null, \n",
      "        \"buffer_count\": 4, \n",
      "        \"pin_memory\": false, \n",
      "        \"pipeline_read\": false, \n",
      "        \"pipeline_write\": false, \n",
      "        \"fast_init\": false\n",
      "    }, \n",
      "    \"sub_group_size\": 1.000000e+09, \n",
      "    \"prefetch_bucket_size\": 5.000000e+07, \n",
      "    \"param_persistence_threshold\": 1.000000e+05, \n",
      "    \"max_live_parameters\": 1.000000e+09, \n",
      "    \"max_reuse_distance\": 1.000000e+09, \n",
      "    \"gather_fp16_weights_on_model_save\": false, \n",
      "    \"ignore_unused_parameters\": true, \n",
      "    \"legacy_stage1\": false\n",
      "}\n",
      "[2021-09-28 21:44:56,728] [INFO] [config.py:904:print]   zero_enabled ................. True\n",
      "[2021-09-28 21:44:56,729] [INFO] [config.py:904:print]   zero_optimization_stage ...... 2\n",
      "[2021-09-28 21:44:56,729] [INFO] [config.py:906:print]   json = {\n",
      "    \"micro_batch_per_gpu\": 1, \n",
      "    \"train_batch_size\": 16, \n",
      "    \"steps_per_print\": 2.000000e+03, \n",
      "    \"fp16\": {\n",
      "        \"enabled\": true, \n",
      "        \"loss_scale\": 0, \n",
      "        \"loss_scale_window\": 1000, \n",
      "        \"initial_scale_power\": 12, \n",
      "        \"hysteresis\": 2, \n",
      "        \"min_loss_scale\": 1\n",
      "    }, \n",
      "    \"zero_optimization\": {\n",
      "        \"stage\": 2, \n",
      "        \"allgather_partitions\": true, \n",
      "        \"allgather_bucket_size\": 2.000000e+08, \n",
      "        \"overlap_comm\": true, \n",
      "        \"reduce_scatter\": true, \n",
      "        \"reduce_bucket_size\": 2.000000e+08, \n",
      "        \"contiguous_gradients\": true, \n",
      "        \"cpu_offload\": true\n",
      "    }, \n",
      "    \"optimizer\": {\n",
      "        \"type\": \"Adam\", \n",
      "        \"params\": {\n",
      "            \"lr\": 1e-05, \n",
      "            \"betas\": [0.8, 0.999], \n",
      "            \"eps\": 1e-08, \n",
      "            \"weight_decay\": 3e-07\n",
      "        }\n",
      "    }, \n",
      "    \"scheduler\": {\n",
      "        \"type\": \"WarmupLR\", \n",
      "        \"params\": {\n",
      "            \"warmup_min_lr\": 0, \n",
      "            \"warmup_max_lr\": 3e-05, \n",
      "            \"warmup_num_steps\": 500\n",
      "        }\n",
      "    }, \n",
      "    \"wall_clock_breakdown\": false\n",
      "}\n",
      "Using /home/fellow/.cache/torch_extensions as PyTorch extensions root...\n",
      "No modifications detected for re-loaded extension module utils, skipping build step...\n",
      "Loading extension module utils...\n",
      "Time to load utils op: 0.004884004592895508 seconds\n",
      "self.eos_token_id=50256\n",
      "self.bos_token_id=50256\n",
      "self.pad_token_id=None\n"
     ]
    }
   ],
   "source": [
    "gptj = GPTJ_PrefixTune.from_pretrained(\n",
    "    './model_webnlg/checkpoint-3250/',\n",
    "    main_checkpoint_override=\"/export/data/gptj/j6b_ckpt/\",\n",
    "    deepspeed_config='ds_config_stage2_gptj_gen.json',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5f4258f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_dataset_file = \"/export/WebNLG/webnlg-dataset/webnlg_challenge_2017/val.json\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f096834e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_inputs(dataset_file, model):\n",
    "\n",
    "    with open(dataset_file) as f:\n",
    "        lines_dict = json.load(f)\n",
    "\n",
    "    full_rela_lst = []\n",
    "    full_src_lst = []\n",
    "    full_tgt_lst = []\n",
    "\n",
    "    for i, example in enumerate(lines_dict['entries']):\n",
    "        sents = example[str(i + 1)]['lexicalisations']\n",
    "        triples = example[str(i + 1)]['modifiedtripleset']\n",
    "\n",
    "        rela_lst = []\n",
    "        temp_triples = ''\n",
    "        for j, tripleset in enumerate(triples):\n",
    "            subj, rela, obj = tripleset['subject'], tripleset['property'], tripleset['object']\n",
    "            rela_lst.append(rela)\n",
    "            temp_triples += ' | '\n",
    "            temp_triples += '{} : {} : {}'.format(subj, rela, obj)\n",
    "\n",
    "        for sent in sents:\n",
    "            if sent[\"comment\"] == 'good':\n",
    "                full_tgt_lst.append(sent[\"lex\"])\n",
    "                full_src_lst.append(temp_triples)\n",
    "                full_rela_lst.append(rela_lst)\n",
    "\n",
    "    edited_sents = []\n",
    "    for src, tgt in zip(full_src_lst, full_tgt_lst):\n",
    "        sent = ' {} {} '.format(src, model.tokenizer.bos_token)\n",
    "        edited_sents.append(sent)\n",
    "        \n",
    "    return edited_sents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "99dd84ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "sents = get_inputs(val_dataset_file, gptj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b1ebc70e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'  | Andrews_County_Airport : owner : Andrews_County,_Texas <|endoftext|> '"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sents[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9f442bf6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50256"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eos_token_id = gptj.tokenizer.encode(\"<|endoftext|>\")[0]\n",
    "eos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "65028156",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7ab31ce3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# =================== #\n",
      "\n",
      "Andrews County Airport is owned by Andrews County, Texas.  \n",
      "\n",
      "<|endoftext|>\n",
      "\n",
      "\n",
      "\n",
      "# =================== #\n",
      "\n",
      "Andrews County Airport is owned by Andrews County, Texas.  \n",
      "\n",
      "<|endoftext|>\n",
      "\n",
      "\n",
      "\n",
      "# =================== #\n",
      "\n",
      "Andrews County Airport is owned by Andrews County, Texas.  \n",
      "\n",
      "<|endoftext|>\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Duration = 1.606154\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    text = sents[idx]\n",
    "    start = datetime.datetime.now()\n",
    "    out = gptj.generate(\n",
    "        1,\n",
    "        text=text,\n",
    "        max_length=256,\n",
    "        #num_beams=2,\n",
    "        do_sample=True,\n",
    "        temperature=0.1,\n",
    "        top_k=5,\n",
    "        top_p=0.95,\n",
    "        no_repeat_ngram_size=2, \n",
    "        early_stopping=True,\n",
    "        num_return_sequences=3,\n",
    "        use_cache=False,\n",
    "        eos_token_id=eos_token_id,\n",
    "        return_only_generated=True,\n",
    "    )\n",
    "    duration = datetime.datetime.now() - start\n",
    "    for o in out:\n",
    "        print(f\"# =================== #\\n{o[:]}\\n\\n\\n\")\n",
    "    print(f\"\\n\\nDuration = {duration.total_seconds()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2e7d7ba1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'  | Andrews_County_Airport : owner : Andrews_County,_Texas <|endoftext|> '"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sents[idx]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5de7098f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gptj.s_wte.n_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ead1d7d9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}