gpt neo with deepspeed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "import transformers as tr\n",
    "import os\n",
    "\n",
    "model_name='EleutherAI/gpt-neo-125M'\n",
    "\n",
    "os.environ['MASTER_ADDR'] = 'localhost'\n",
    "os.environ['MASTER_PORT'] = '9994'\n",
    "os.environ['RANK'] = \"0\"\n",
    "os.environ['LOCAL_RANK'] = \"0\"\n",
    "os.environ['WORLD_SIZE'] = \"1\"\n",
    "\n",
    "bos_token='<|endoftext|>'\n",
    "eos_token='<|endoftext|>'\n",
    "pad_token='<|pad|>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda  11.1 \n",
      "device  Tesla T4\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print('cuda ',torch.version.cuda,\n",
    "      '\\ndevice ', torch.cuda.get_device_name(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Embedding(50258, 768)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = tr.GPT2Tokenizer.from_pretrained(model_name,    \n",
    "                            bos_token=bos_token,\n",
    "                            eos_token=eos_token,\n",
    "                            pad_token=pad_token,)\n",
    "model = tr.GPTNeoForCausalLM.from_pretrained(model_name).cuda()\n",
    "model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=['bla blablabla bla bla']*30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f47219c253c14d278d78a0de44db63fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max_length = max([len(tokenizer.encode(_)) for _ in tqdm(data)])\n",
    "max_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "331f18fa4f68462dbb1a16840f87389e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from torch.utils.data import Dataset\n",
    "\n",
    "class PythonDataset(Dataset):\n",
    "    def __init__(self, txt_list, tokenizer, max_length):\n",
    "        self.input_ids = []\n",
    "        self.attn_masks = []\n",
    "        self.labels = []\n",
    "        for txt in tqdm(txt_list):\n",
    "            # Encode the descriptions using the GPT-Neo tokenizer\n",
    "            encodings_dict = tokenizer( bos_token+txt +    \n",
    "                                        eos_token,\n",
    "                                        truncation=True,\n",
    "                                        max_length=max_length, \n",
    "                                        padding='max_length',\n",
    "                                      )\n",
    "            input_ids = torch.tensor(encodings_dict['input_ids'])    \n",
    "            self.input_ids.append(input_ids)\n",
    "            mask = torch.tensor(encodings_dict['attention_mask'])\n",
    "            self.attn_masks.append(mask)\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx],self.attn_masks[idx]\n",
    "\n",
    "\n",
    "data_collator = tr.DataCollatorForLanguageModeling(\n",
    "        tokenizer=tokenizer, mlm=False,\n",
    "    )\n",
    "\n",
    "dataset = PythonDataset(data, tokenizer, max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30\n",
      "27 3\n"
     ]
    }
   ],
   "source": [
    "len_dataset=len(dataset)\n",
    "print(len_dataset)\n",
    "train_size = int(0.9 * len_dataset)\n",
    "\n",
    "train_dataset, val_dataset = torch.utils.data.dataset.random_split(dataset, [train_size, len_dataset - train_size])\n",
    "print(len(train_dataset),len(val_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-29 21:53:00,608] [INFO] [distributed.py:47:init_distributed] Initializing torch distributed with backend: nccl\n"
     ]
    }
   ],
   "source": [
    "save_dir='./results'\n",
    "training_args = tr.TrainingArguments(output_dir=save_dir, num_train_epochs=5, logging_steps=300, save_steps=300,\n",
    "                                  per_device_train_batch_size=1, per_device_eval_batch_size=1,warmup_steps=50,\n",
    "                                     learning_rate=0.001,adam_epsilon=1e-06,fp16=True,\n",
    "                                  weight_decay=0.01, logging_dir=f'{save_dir}/logs', deepspeed='./ds_config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using amp fp16 backend\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-06-29 21:53:00,650] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed info: version=0.4.1, git-hash=unknown, git-branch=unknown\n",
      "[2021-06-29 21:53:00,659] [INFO] [utils.py:13:_initialize_parameter_parallel_groups] data_parallel_size: 1, parameter_parallel_size: 1\n"
     ]
    }
   ],
   "source": [
    "trainer = tr.Trainer(model=model, \n",
    "                     args=training_args,  \n",
    "                  train_dataset=train_dataset,\n",
    "                  eval_dataset=val_dataset, \n",
    "                  data_collator=lambda data: \n",
    "              {'input_ids': torch.stack([f[0] for f in data]),       \n",
    "               'attention_mask': torch.stack([f[1] for f in data]),\n",
    "               'labels': torch.stack([f[0] for f in data])}\n",
    "                    )\n",
    "# Start training process!\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# at this point I cannot interrupt kenrel"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "esemala",
   "language": "python",
   "name": "esemala"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}