From a111e3c0aeccd8d76573fd2596a8e7a0760e9e19 Mon Sep 17 00:00:00 2001
From: mymusise <mymusise1@gmail.com>
Date: Mon, 24 Apr 2023 11:05:52 +0800
Subject: [PATCH] add ppo

Signed-off-by: mymusise <mymusise1@gmail.com>
---
 ppo_training_glm.ipynb | 445 +++++++++++++++++++++++++++++++++++++++++
 utils/__init__.py      |   0
 utils/trainer.py       |  90 +++++++++
 3 files changed, 535 insertions(+)
 create mode 100644 ppo_training_glm.ipynb
 create mode 100644 utils/__init__.py
 create mode 100644 utils/trainer.py

diff --git a/ppo_training_glm.ipynb b/ppo_training_glm.ipynb
new file mode 100644
index 0000000..33ee273
--- /dev/null
+++ b/ppo_training_glm.ipynb
@@ -0,0 +1,445 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModel\n",
+    "from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## init reward model "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from reward_model import RewardModel\n",
+    "from transformers import AutoTokenizer\n",
+    "from peft import PeftModel\n",
+    "from torch.nn.utils import skip_init\n",
+    "import torch\n",
+    "\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"THUDM/chatglm-6b\", trust_remote_code=True)\n",
+    "\n",
+    "\n",
+    "reward_model = RewardModel.from_pretrained(\"THUDM/chatglm-6b\", load_in_8bit=True, device_map='auto')\n",
+    "\n",
+    "## load score weight\n",
+    "\n",
+    "reward_model = PeftModel.from_pretrained(reward_model, './output/reward_model/', load_in_8bit=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "weight torch.Size([1, 4096]) tensor([0.0771, 0.0723, 0.1037, 0.1068, 0.0667], device='cuda:0',\n",
+      "       dtype=torch.float16) tensor([0.0782, 0.0766, 0.0637, 0.0989, 0.1059], device='cuda:0',\n",
+      "       dtype=torch.float16) tensor(0.0939, device='cuda:0', dtype=torch.float16) tensor(0.0825, device='cuda:0', dtype=torch.float16)\n"
+     ]
+    }
+   ],
+   "source": [
+    "class CastOutputToHalf(torch.nn.Sequential):\n",
+    "    def forward(self, x):\n",
+    "        return super().forward(x).half()\n",
+    "\n",
+    "\n",
+    "reward_model.gradient_checkpointing_disable()\n",
+    "\n",
+    "reward_model.base_model.model.score.load_state_dict(torch.load(\"output/reward_model/score.bin\"))\n",
+    "\n",
+    "for k, p in reward_model.base_model.model.score.named_parameters():\n",
+    "        print(k, p.shape, p[0, :5], p[0, -5:], p[0][:20].mean(), p[0][-20:].mean())\n",
+    "\n",
+    "# reward_model.score = CastOutputToHalf(reward_model.score)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "weight torch.Size([1, 4096]) tensor([0.0771, 0.0723, 0.1037, 0.1068, 0.0667], device='cuda:0',\n",
+      "       dtype=torch.float16) tensor([0.0782, 0.0766, 0.0637, 0.0989, 0.1059], device='cuda:0',\n",
+      "       dtype=torch.float16) tensor(0.0939, device='cuda:0', dtype=torch.float16) tensor(0.0825, device='cuda:0', dtype=torch.float16)\n"
+     ]
+    }
+   ],
+   "source": [
+    "for k, p in reward_model.named_parameters():\n",
+    "    p.requires_grad = False\n",
+    "\n",
+    "for k, p in reward_model.score.named_parameters():\n",
+    "        print(k, p.shape, p[0, :5], p[0, -5:], p[0][:20].mean(), p[0][-20:].mean())\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## init actor"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Explicitly passing a `revision` is encouraged when loading a configuration with custom code to ensure no malicious code has been contributed in a newer revision.\n",
+      "Explicitly passing a `revision` is encouraged when loading a model with custom code to ensure no malicious code has been contributed in a newer revision.\n",
+      "Overriding torch_dtype=None with `torch_dtype=torch.float16` due to requirements of `bitsandbytes` to enable model loading in mixed int8. Either pass torch_dtype=torch.float16 or don't pass this argument at all to remove this warning.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fb2b1ec062f54a38b5dc0768b2c4d98c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_int8_training\n",
+    "\n",
+    "model_name = \"THUDM/chatglm-6b\"\n",
+    "\n",
+    "pretrained_model = AutoModel.from_pretrained(model_name, load_in_8bit=True, trust_remote_code=True, device_map='auto')\n",
+    "\n",
+    "\n",
+    "## SFT pretrained_model with LoRA\n",
+    "\n",
+    "pretrained_model.gradient_checkpointing_enable()\n",
+    "pretrained_model.enable_input_require_grads()\n",
+    "pretrained_model.is_parallelizable = True\n",
+    "pretrained_model.model_parallel = True\n",
+    "# model.lm_head = CastOutputToFloat(model.lm_head)\n",
+    "pretrained_model.config.use_cache = (\n",
+    "    False  # silence the warnings. Please re-enable for inference!\n",
+    ")\n",
+    "pretrained_model = prepare_model_for_int8_training(pretrained_model)\n",
+    "\n",
+    "# setup peft\n",
+    "peft_config = LoraConfig(\n",
+    "    task_type=TaskType.CAUSAL_LM,\n",
+    "    inference_mode=False,\n",
+    "    r=8,\n",
+    "    lora_alpha=32,\n",
+    "    lora_dropout=0.1,\n",
+    ")\n",
+    "pretrained_model = get_peft_model(pretrained_model, peft_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def print_trainable_parameters(model):\n",
+    "    \"\"\"\n",
+    "    Prints the number of trainable parameters in the model.\n",
+    "    \"\"\"\n",
+    "    trainable_params = 0\n",
+    "    all_param = 0\n",
+    "    for _, param in model.named_parameters():\n",
+    "        all_param += param.numel()\n",
+    "        if param.requires_grad:\n",
+    "            trainable_params += param.numel()\n",
+    "    print(\n",
+    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 3674113 || all params: 6176960513 || trainable%: 0.059480920952424424\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = AutoModelForCausalLMWithValueHead.from_pretrained(pretrained_model)\n",
+    "\n",
+    "model.gradient_checkpointing_disable = model.pretrained_model.gradient_checkpointing_disable\n",
+    "model.gradient_checkpointing_enable = model.pretrained_model.gradient_checkpointing_enable\n",
+    "\n",
+    "print_trainable_parameters(model)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## init dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Explicitly passing a `revision` is encouraged when loading a model with custom code to ensure no malicious code has been contributed in a newer revision.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from transformers import AutoTokenizer\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"THUDM/chatglm-6b\", trust_remote_code=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset, Dataset\n",
+    "\n",
+    "\n",
+    "dataset = load_dataset(\"BelleGroup/train_0.5M_CN\", split='train')\n",
+    "dataset = Dataset.from_dict(dataset[:1000])\n",
+    "dataset = dataset.rename_columns({'instruction': 'query'})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9b6f5c7654dc493e9f4f5f38651a792f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "def encode_data(sample):\n",
+    "    sample['input_ids'] = tokenizer.encode(sample[\"query\"], max_length=512, truncation=True)\n",
+    "    return sample\n",
+    "\n",
+    "dataset = dataset.map(encode_data)\n",
+    "\n",
+    "dataset.set_format(type=\"torch\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## init ppo trainer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = PPOConfig(\n",
+    "    model_name=model_name,\n",
+    "    learning_rate=1e-5,\n",
+    "    log_with=\"all\",  # wandb and tensorboard\n",
+    "    accelerator_kwargs={\"logging_dir\":\"output/ppo/\"},\n",
+    "    mini_batch_size=2,\n",
+    "    batch_size=2,\n",
+    "    gradient_accumulation_steps=1,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def collator(data):\n",
+    "    return dict((key, [d[key] for d in data]) for key in data[0])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from utils.trainer import ChatGLMPPOTrainer\n",
+    "\n",
+    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)\n",
+    "\n",
+    "ppo_trainer = ChatGLMPPOTrainer(\n",
+    "    config, model, ref_model=None, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "0it [00:00, ?it/s]The dtype of attention mask (torch.int64) is not bool\n",
+      "/home/mymusise/pro/stable-diffusion-webui/venv/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:717: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ../torch/csrc/utils/tensor_new.cpp:230.)\n",
+      "  tensor = as_tensor(value)\n",
+      "1it [00:31, 31.72s/it]"
+     ]
+    }
+   ],
+   "source": [
+    "from tqdm import tqdm \n",
+    "from trl.core import LengthSampler\n",
+    "\n",
+    "\n",
+    "generation_kwargs = {\n",
+    "    \"min_length\": -1,\n",
+    "    \"top_k\": 0.0,\n",
+    "    \"top_p\": 1.0,\n",
+    "    \"do_sample\": True,\n",
+    "    \"pad_token_id\": tokenizer.eos_token_id,\n",
+    "    \"eos_token_id\": -1,\n",
+    "}\n",
+    "output_min_length = 4\n",
+    "output_max_length = 16\n",
+    "output_length_sampler = LengthSampler(output_min_length, output_max_length)\n",
+    "\n",
+    "\n",
+    "for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n",
+    "    query_tensors = batch[\"input_ids\"]\n",
+    "\n",
+    "    model.gradient_checkpointing_disable()\n",
+    "    model.pretrained_model.config.use_cache = True\n",
+    "    # Get response from Causal LM\n",
+    "    response_tensors = []\n",
+    "    for query in query_tensors:\n",
+    "        gen_len = output_length_sampler()\n",
+    "        generation_kwargs[\"max_new_tokens\"] = gen_len\n",
+    "        response = ppo_trainer.generate(\n",
+    "            query_tensors[0], **generation_kwargs\n",
+    "        )\n",
+    "        response_tensors.append(response.squeeze()[-gen_len:])\n",
+    "\n",
+    "    # Compute sentiment score\n",
+    "    batch[\"response\"] = [tokenizer.decode(ids) for ids in response_tensors]\n",
+    "    texts = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]\n",
+    "    # pipe_outputs = sentiment_pipe(texts, **sent_kwargs)\n",
+    "    # rewards = [torch.tensor(output[1][\"score\"]) for output in pipe_outputs]\n",
+    "\n",
+    "    contents = [torch.cat([qids, rids]) for qids, rids in zip(batch['input_ids'], response_tensors)]\n",
+    "\n",
+    "    rewards = []\n",
+    "    for c in contents:\n",
+    "        rewards.append(reward_model(c.reshape(1, -1))[0].sum())\n",
+    "\n",
+    "    # Run PPO step\n",
+    "    model.gradient_checkpointing_enable()\n",
+    "    model.pretrained_model.config.use_cache = False\n",
+    "\n",
+    "    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)\n",
+    "    ppo_trainer.log_stats(stats, batch, rewards)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## save model after ppo"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save_pretrained(\"output/ppo/\")"
+   ]
+  }
+ ],
+ "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.8.12"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "25273a2a68c96ebac13d7fb9e0db516f9be0772777a0507fe06d682a441a3ba7"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/utils/__init__.py b/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/utils/trainer.py b/utils/trainer.py
new file mode 100644
index 0000000..76562f1
--- /dev/null
+++ b/utils/trainer.py
@@ -0,0 +1,90 @@
+import torch
+
+from trl.trainer.ppo_trainer import (
+    logprobs_from_logits,
+    PreTrainedModelWrapper,
+    PPODecorators,
+    PPOTrainer,
+)
+
+
+class ChatGLMPPOTrainer(PPOTrainer):
+    def pad_input_ids(self, ids, max_length):
+        # _ids = torch.ones(max_length, dtype=torch.int8) * self.model.config.pad_token_id
+        _ids = torch.ones(max_length, dtype=torch.int64) * 3
+        _ids[: len(ids)] = ids
+        return _ids
+
+    def prepare_model_inputs(self, queries: torch.Tensor, responses: torch.Tensor):
+        input_ids = [torch.cat([q, r]) for q, r in zip(queries, responses)]
+        max_length = max([len(ids) for ids in input_ids])
+        input_data = self.data_collator(
+            [{"input_ids": self.pad_input_ids(ids, max_length)} for ids in input_ids]
+        ).to(self.current_device)
+
+        input_data.pop("labels", None)  # we don't want to compute LM losses
+
+        return input_data
+
+    @PPODecorators.empty_cuda_cache()
+    def batched_forward_pass(
+        self,
+        model: PreTrainedModelWrapper,
+        queries: torch.Tensor,
+        responses: torch.Tensor,
+        model_inputs: dict,
+    ):
+        """
+        Calculate model outputs in multiple batches.
+
+        Args:
+            queries (`torch.LongTensor`):
+                List of tensors containing the encoded queries, shape (`batch_size`, `query_length`)
+            responses (`torch.LongTensor`):
+                List of tensors containing the encoded responses, shape (`batch_size`, `response_length`)
+        Returns:
+            (tuple):
+                - all_logprobs (`torch.FloatTensor`): Log probabilities of the responses,
+                    shape (`batch_size`, `response_length`)
+                - all_ref_logprobs (`torch.FloatTensor`): Log probabilities of the responses,
+                    shape (`batch_size`, `response_length`)
+                - all_values (`torch.FloatTensor`): Values of the responses, shape (`batch_size`, `response_length`)
+        """
+        bs = len(queries)
+        fbs = self.config.mini_batch_size
+        all_logprobs = []
+        all_logits = []
+        all_masks = []
+        all_values = []
+
+        for i in range(int(bs / fbs)):
+            input_kwargs = {
+                key: value[i * fbs : (i + 1) * fbs]
+                for key, value in model_inputs.items()
+            }
+            logits, _, values = model(**input_kwargs)
+            input_ids = input_kwargs["input_ids"]
+
+            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
+            masks = torch.zeros_like(input_ids)
+            for i, ids in enumerate(input_ids):
+                start_i = (ids == self.model.config.bos_token_id).nonzero()
+                end_i = (ids == self.model.config.eos_token_id).nonzero()
+                if len(end_i):
+                    end_i = end_i[0][0] + 1
+                else:
+                    end_i = None
+                masks[i][start_i:end_i] = 1
+
+            all_logits.append(logits)
+            all_values.append(values.rot90())
+            all_logprobs.append(logprobs)
+            all_masks.append(masks)
+
+        return (
+            torch.cat(all_logprobs),
+            torch.cat(all_logits)[:, :-1],
+            torch.cat(all_values)[:, :-1],
+            torch.cat(all_masks)[:, :-1],
+        )
+