05-accelerate-deepspeed-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 完成数据预处理将数据整理成prompt的格式\n",
    "from datasets import load_dataset\n",
    "import json\n",
    "import os\n",
    "\n",
    "\n",
    "data_files = {\n",
    "    \"train\": \"alpaca_data.json\",\n",
    "}\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=data_files)\n",
    "\n",
    "def main():\n",
    "    if not os.path.exists(\"data\"):\n",
    "        os.mkdir(\"data\")\n",
    "\n",
    "\n",
    "    with open(f\"data/alpaca_data.jsonl\", \"w\") as f:\n",
    "        for item in dataset[\"train\"]:\n",
    "            newitem = {}\n",
    "            if item[\"input\"]:\n",
    "                newitem[\"source\"] = (\n",
    "                    \"Below is an instruction that describes a task, paired with an input that provides further context. \" + \"\\n\"+ \\\n",
    "                    \"Write a response that appropriately completes the request.\\n\\n\" + \\\n",
    "                    \"### Instruction:\\n\" + item[\"instruction\"] + \"\\n\\n\" + \\\n",
    "                    \"### Input:\\n\" + item[\"input\"] + \"\\n\\n\"  \\\n",
    "                    \"### Response:\\n\" + item[\"output\"]\n",
    "                )\n",
    "            else:\n",
    "                newitem[\"source\"] = (\n",
    "                    \"Below is an instruction that describes a task. \" + \"\\n\"+ \\\n",
    "                    \"Write a response that appropriately completes the request.\\n\\n\" + \\\n",
    "                    \"### Instruction:\\n\" + item[\"instruction\"] + \"\\n\\n\" + \\\n",
    "                    \"### Response:\\n\" + item[\"output\"]\n",
    "                )\n",
    "            f.write(json.dumps(newitem) + \"\\n\")\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "os.environ[\"HF_ENDPOINT\"] = \"https://hf.neolink-ai.com\"\n",
    "os.environ[\"HTTPS_PROXY\"] = \"http://10.161.0.82:7899/\"\n",
    "\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from accelerate import Accelerator\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, get_scheduler, DataCollatorWithPadding\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "llm_data_files = {\n",
    "    \"train\": \"data/alpaca_data.jsonl\",\n",
    "}\n",
    "\n",
    "llm_dataset = load_dataset(\"json\", data_files=llm_data_files)\n",
    "max_length = 128\n",
    "checkpoints = \"meta-llama/Llama-3.2-1B\"\n",
    "special_tokens = json.load(open(\"data/tokens.json\"))[\"tokens\"]\n",
    "accelerator = Accelerator()\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(checkpoints, legacy=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.add_tokens(special_tokens, special_tokens=True)\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    # 在使用tokenizer function的模式下，需要将labels设置为input_ids,默认情况下 是不存在labels 这样会导致loss为nan\n",
    "    out_batch = tokenizer(examples[\"source\"], truncation=True, padding=\"max_length\", max_length=max_length, return_tensors=\"pt\")\n",
    "    out_batch[\"labels\"] = out_batch[\"input_ids\"]\n",
    "    return out_batch\n",
    "\n",
    "tokenized_datasets = llm_dataset.map(tokenize_function, batched=True)\n",
    "print(tokenized_datasets)\n",
    "tokenized_datasets = tokenized_datasets.remove_columns([\"source\"])\n",
    "print(tokenized_datasets)\n",
    "tokenized_datasets.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"labels\"])\n",
    "\n",
    "collate_fn = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "train_dataloader = DataLoader(tokenized_datasets[\"train\"], shuffle=True, batch_size=8, collate_fn=collate_fn)\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(checkpoints, torch_dtype=torch.bfloat16)\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "model.gradient_checkpointing_enable()\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)\n",
    "\n",
    "\n",
    "num_epochs = 3\n",
    "num_training_steps = num_epochs * len(train_dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    \"linear\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0,\n",
    "    num_training_steps=num_training_steps,\n",
    ")\n",
    "\n",
    "model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(model, optimizer, train_dataloader, lr_scheduler)\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        outputs = model(**batch)\n",
    "        loss = outputs.loss\n",
    "        accelerator.backward(loss)\n",
    "\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        progress_bar.update(1)\n",
    "        if step % 10 == 0:\n",
    "            print(f\"epoch: {epoch}, step: {step}, loss: {loss}\")\n"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}