doc/source/train/examples/deepspeed/dolly_v2_deepspeed_instruction_fine_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(dolly_v2_deepspeed_instruction_finetune)=\n",
    "\n",
    "# Dolly-V2-3B Instruction Fine-Tuning with Ray Train and DeepSpeed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this demonstration, we'll show how to use the Ray Train for Dolly V2 3B model instruction fine-tuning using the deep-speed framework. Please uncomment the next two cells and install the following libraries dependencies.\n",
    "\n",
    "This work builds upon [existing efforts](https://github.com/ray-project/ray/blob/master/doc/source/train/examples/deepspeed/gptj_deepspeed_fine_tuning.ipynb) by incorporating an instruction fine-tuning component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# conda install -c conda-forge mpi4py\n",
    "# conda install gcc gxx_linux-64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install \"ray==2.5.1\" \"accelerate==0.16.0\" \"datasets==2.12.0\" \"transformers==4.26.0\"  \"torch==1.13.0\" \"deepspeed==0.9.2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray.data\n",
    "import ray\n",
    "from datasets import load_dataset\n",
    "from transformers import Trainer, TrainingArguments\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForCausalLM,\n",
    ")\n",
    "from transformers.utils.logging import enable_progress_bar\n",
    "import torch\n",
    "\n",
    "import transformers\n",
    "from ray.train.huggingface.transformers import (\n",
    "    prepare_trainer,\n",
    "    RayTrainReportCallback,\n",
    ")\n",
    "from ray.train.torch import TorchTrainer\n",
    "from ray.train import RunConfig, ScalingConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Set up Ray <a name=\"setup\"></a>\n",
    "\n",
    "First, we will use 2 workers, each being assigned 1 GPU and 28 CPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"databricks/dolly-v2-3b\"\n",
    "use_gpu = True\n",
    "num_workers = 2\n",
    "cpus_per_worker = 28"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fq_ray_ip = #<replace-this-with-your-ray-server-ip-address>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip_env = {\n",
    "    \"pip\": [\n",
    "        \"datasets==2.12.0\",\n",
    "        \"evaluate==0.4.0\",\n",
    "        \"accelerate==0.16.0\",\n",
    "        \"transformers==4.26.0\",\n",
    "        \"torch==1.13.0\",\n",
    "        \"deepspeed==0.9.2\",\n",
    "        \"ipython==8.14.0\",\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conda_env = {\n",
    "    \"conda\": {\n",
    "        \"dependencies\": [\"mpi4py\", \"pip\", pip_env]\n",
    "    }  # pip install mpi4py won't work, use conda install instead\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init(\n",
    "    f\"ray://{fq_ray_ip}:10001\",  # Note: the port and ip-address depends on your ray server setup.\n",
    "    runtime_env=conda_env,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the dataset <a name=\"load\"></a>\n",
    "\n",
    "We will be fine-tuning the model on the [`alpaca-cleaned` dataset](https://datasets-server.huggingface.co/splits?dataset=yahma%2Falpaca-cleaned), comprised of 51,000 lines of Q&A. The aim will be to make the databricks model better at generating answer by following the instruction.\n",
    "\n",
    "We will use `generate_prompt` function to prepare our dataset for instruction fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_dataset = load_dataset(\"yahma/alpaca-cleaned\")\n",
    "current_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_prompt(data_point):\n",
    "    # ref: https://github.com/tloen/alpaca-lora\n",
    "    if data_point[\"instruction\"]:\n",
    "        return f\"\"\"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
    "\n",
    "### Instruction:\n",
    "{data_point[\"instruction\"]}\n",
    "\n",
    "### Input:\n",
    "{data_point[\"input\"]}\n",
    "\n",
    "### Response:\n",
    "{data_point[\"output\"]}\"\"\"\n",
    "    else:\n",
    "        return f\"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n",
    "\n",
    "### Instruction:\n",
    "{data_point[\"instruction\"]}\n",
    "\n",
    "### Response:\n",
    "{data_point[\"output\"]}\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=\"left\")\n",
    "tokenizer.pad_token_id = 0\n",
    "CUTOFF_LEN = 128\n",
    "\n",
    "current_dataset = current_dataset.shuffle().map(\n",
    "    lambda data_point: tokenizer(\n",
    "        generate_prompt(data_point),\n",
    "        truncation=True,\n",
    "        max_length=CUTOFF_LEN,\n",
    "        padding=\"max_length\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray_datasets = ray.data.from_huggingface(current_dataset[\"train\"])\n",
    "ray_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instruction fine-tuning the model with Ray Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def trainer_init_per_worker(config):\n",
    "    batch_size = config.get(\"batch_size\", 1)\n",
    "    epochs = config.get(\"epochs\", 1)\n",
    "    warmup_steps = config.get(\"warmup_steps\", 0)\n",
    "    learning_rate = config.get(\"learning_rate\", 0.00002)\n",
    "    weight_decay = config.get(\"weight_decay\", 0.01)\n",
    "\n",
    "    deepspeed = {\n",
    "        \"fp16\": {\n",
    "            \"enabled\": \"auto\",\n",
    "            \"initial_scale_power\": 8,\n",
    "        },\n",
    "        \"optimizer\": {\n",
    "            \"type\": \"AdamW\",\n",
    "            \"params\": {\n",
    "                \"lr\": \"auto\",\n",
    "                \"betas\": \"auto\",\n",
    "                \"eps\": \"auto\",\n",
    "            },\n",
    "        },\n",
    "        \"zero_optimization\": {\n",
    "            \"stage\": 3,\n",
    "            \"offload_optimizer\": {\n",
    "                \"device\": \"cpu\",\n",
    "                \"pin_memory\": True,\n",
    "            },\n",
    "            \"offload_param\": {\n",
    "                \"device\": \"cpu\",\n",
    "                \"pin_memory\": True,\n",
    "            },\n",
    "            \"overlap_comm\": True,\n",
    "            \"contiguous_gradients\": True,\n",
    "            \"reduce_bucket_size\": \"auto\",\n",
    "            \"stage3_prefetch_bucket_size\": \"auto\",\n",
    "            \"stage3_param_persistence_threshold\": \"auto\",\n",
    "            \"gather_16bit_weights_on_model_save\": True,\n",
    "            \"round_robin_gradients\": True,\n",
    "        },\n",
    "        \"gradient_accumulation_steps\": \"auto\",\n",
    "        \"gradient_clipping\": \"auto\",\n",
    "        \"steps_per_print\": 10,\n",
    "        \"train_batch_size\": \"auto\",\n",
    "        \"train_micro_batch_size_per_gpu\": \"auto\",\n",
    "        \"wall_clock_breakdown\": False,\n",
    "    }\n",
    "    \n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=\"left\")\n",
    "    tokenizer.pad_token_id = 0  # unk. we want this to be different from the eos token\n",
    "    \n",
    "    print(\"Loading model\")\n",
    "    \n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name, \n",
    "        # device_map=\"auto\",\n",
    "        torch_dtype=torch.float16,\n",
    "    )\n",
    "\n",
    "    print(\"Model loaded\")\n",
    "    \n",
    "    train_ds = ray.train.get_dataset_shard(\"train\")\n",
    "    train_ds_iterable = train_ds.iter_torch_batches(\n",
    "        batch_size=batch_size,\n",
    "    )\n",
    "    \n",
    "    print(f\"batch_size: {batch_size}\")\n",
    "    print(\"Preparing training arguments\")\n",
    "    \n",
    "    training_args = TrainingArguments(\n",
    "        output_dir=\"deepspeed-dolly\",\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        logging_steps=1,\n",
    "        max_steps=steps_per_epoch * epochs,\n",
    "        learning_rate=learning_rate,\n",
    "        weight_decay=weight_decay,\n",
    "        warmup_steps=warmup_steps,\n",
    "        num_train_epochs=epochs,\n",
    "        push_to_hub=False,\n",
    "        disable_tqdm=False,\n",
    "        fp16=True,\n",
    "        gradient_accumulation_steps=16,\n",
    "        deepspeed=deepspeed,\n",
    "    )\n",
    "\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=\"left\")\n",
    "    tokenizer.pad_token_id = 0\n",
    "\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        torch_dtype=torch.float16,\n",
    "    )\n",
    "    \n",
    "    print(\"Model loaded\")\n",
    "    \n",
    "    enable_progress_bar()\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=train_ds_iterable,\n",
    "        tokenizer=tokenizer,\n",
    "        data_collator=transformers.DataCollatorForLanguageModeling(\n",
    "            tokenizer, mlm=False\n",
    "        ),\n",
    "    )\n",
    "    \n",
    "    #Add callback to report checkpoints to Ray Train\n",
    "    trainer.add_callback(RayTrainReportCallback())\n",
    "    trainer = prepare_trainer(trainer)\n",
    "    \n",
    "    print(\"Start training\")\n",
    "    trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "train_ds_size = ray_datasets.count()\n",
    "steps_per_epoch = train_ds_size // (batch_size * num_workers)\n",
    "trainer = TorchTrainer(\n",
    "    train_loop_per_worker=trainer_init_per_worker,\n",
    "    trainer_init_config={\n",
    "        \"batch_size\": batch_size,  # batch_size per device\n",
    "        \"epochs\": 2,\n",
    "        \"steps_per_epoch\": steps_per_epoch,\n",
    "    },\n",
    "    scaling_config=ScalingConfig(\n",
    "        num_workers=num_workers,\n",
    "        use_gpu=use_gpu,\n",
    "        resources_per_worker={\n",
    "            \"GPU\": 1,\n",
    "            \"CPU\": cpus_per_worker,\n",
    "        },  # NOTE: huggingface transformers only support 1 GPU per worker.\n",
    "    ),\n",
    "    run_config=RunConfig(\n",
    "        storage_path = \"<your s3 bucket>\"\n",
    "    ),\n",
    "    datasets={\n",
    "        \"train\": ray_datasets,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Finally, we call the `~ray.train.torch.TorchTrainer.fit` method to start training with Ray Train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = trainer.fit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [.conda-ray-310]",
   "language": "python",
   "name": "conda-env-.conda-ray-310-py"
  },
  "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.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}