00_Setup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8d76552-9f14-4c34-96dd-6be7b2f2df30",
   "metadata": {},
   "source": [
    "# Warm-Up Notebook\n",
    "---\n",
    "\n",
    "You've made it to the first notebook in this workshop! As a good way to check that everything is running correctly, let's preload the model weights to save time later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86cc77c-5be6-46d9-b1ad-45479fa0f504",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b21325cf-fc49-4b69-8115-0b2aa144fc4f",
   "metadata": {
    "tags": []
   },
   "source": [
    "We'll be using [StableLM](https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b) throughout this workshop. To trigger the full download of the weights, set up a pipeline that caches the model with the fast NVMe storage on Anyscale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f937b87-a71f-4630-9955-def529143a58",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "p  = pipeline(model=\"stabilityai/stablelm-tuned-alpha-7b\", task='text-generation', \n",
    "              model_kwargs={'device_map':'auto', 'torch_dtype' : torch.float16, 'cache_dir': '/mnt/local_storage/'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bd41ea7-d6a0-4542-9fd9-9f4f29a487e7",
   "metadata": {},
   "source": [
    "Verify that the model is loaded into GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e1511c-258b-4d46-8438-d0c45deea218",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecc97733-5aa8-41b0-9b2c-276b65a6589a",
   "metadata": {},
   "source": [
    "In production situations, this memory should be freed when the process exits. However, in a notebook (or other long-running dev process environment), it can be useful to purge unneeded data directly.\n",
    "\n",
    "Additionally, we can use 🤗 Accelerate to release any unused GPU memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca91486-64dd-4e9a-8906-5bec952f0185",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "del(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b7d5e1d-88fb-4653-a7cf-177c4332df47",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "\n",
    "accelerator = Accelerator()\n",
    "accelerator.free_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f24a374",
   "metadata": {},
   "source": [
    "Verify that the memory is freed. You can also check out the Ray Dashboard for Node GPU memory time series metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbfc38be-3083-473e-90f2-11e661398e86",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2689d0cd-72b3-44fd-ac63-e62d3dd75975",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ray\n",
    "import socket\n",
    "\n",
    "ip  = socket.gethostbyname(socket.gethostname())\n",
    "my_ip = 'node:' + ip\n",
    "\n",
    "def findOther(node):\n",
    "    try: \n",
    "        node['Resources'][my_ip] \n",
    "    except:\n",
    "        return node['Resources']\n",
    "\n",
    "if not ray.is_initialized():\n",
    "    ray.init()\n",
    "\n",
    "keys = next(filter(findOther, ray.nodes()))['Resources'].keys()\n",
    "worker_resource = next(filter(lambda key:key.startswith('node'), keys))\n",
    "\n",
    "@ray.remote(resources={worker_resource:1})\n",
    "def preload():\n",
    "    p  = pipeline(model=\"stabilityai/stablelm-tuned-alpha-7b\", task='text-generation', \n",
    "              model_kwargs={'device_map':'auto', 'torch_dtype' : torch.float16, 'cache_dir': '/mnt/local_storage/'})\n",
    "    return socket.gethostbyname(socket.gethostname())\n",
    "\n",
    "res = preload.remote()\n",
    "print(ray.get(res))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}