Merge pull request #1160 from rajasekharporeddy:qlora_llama3

PiperOrigin-RevId: 877423188

Author: The tunix Authors

examples/qlora_llama3_gpu.ipynb

@@ -38,6 +38,44 @@
         "!nvidia-smi"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "svAMl1BxH_nC"
+      },
+      "source": [
+        "## Installing Libraries: RESTART RUNTIME AFTER INSTALLATION\n",
+        "\n",
+        "The easiest way to get a working environment is the [NVIDIA NGC JAX container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax), which ships with all dependencies preinstalled. To install the dependencies manually:\n",
+        "\n",
+        "```bash\n",
+        "pip install 'jax[cuda13]' flax optax transformers datasets qwix huggingface_hub wandb\n",
+        "```\n",
+        "\n",
+        "On top of the installation (either container or manual), you will need Tunix:\n",
+        "\n",
+        "```bash\n",
+        "pip install git+https://github.com/google/tunix\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YIEvogq3Lvh2"
+      },
+      "outputs": [],
+      "source": [
+        "# Install necessary libraries\n",
+        "import importlib\n",
+        "\n",
+        "if importlib.util.find_spec(\"tunix\") is None:\n",
+        "  print(\"Required packages not found. Running full installation...\")\n",
+        "  %pip install 'jax[cuda13]' flax optax transformers datasets qwix huggingface_hub wandb\n",
+        "  %pip install git+https://github.com/google/tunix"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -118,7 +156,7 @@
       "source": [
         "# Prefer environment variable if already set\n",
         "\n",
-        "from huggingface_hub.v1.hf_api import whoami\n",
+        "from huggingface_hub import whoami\n",
         "HF_TOKEN = os.environ.get(\"HF_TOKEN\")\n",
         "\n",
         "if HF_TOKEN:\n",
@@ -160,19 +198,7 @@
         "- **Optax**: Gradient processing and optimization library for JAX\n",
         "- **Transformers**: Hugging Face library for tokenizers and model configurations\n",
         "- **Qwix**: Quantization and LoRA utilities for JAX models\n",
-        "- **Tunix**: Training utilities including `PeftTrainer` and `AutoModel` for streamlined fine-tuning\n",
-        "\n",
-        "The easiest way to get a working environment is the [NVIDIA NGC JAX container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax), which ships with all dependencies preinstalled. To install the dependencies manually:\n",
-        "\n",
-        "```bash\n",
-        "pip install 'jax[cuda13]' flax optax transformers datasets qwix\n",
-        "```\n",
-        "\n",
-        "On top of the installation (either container or manual), you will need Tunix:\n",
-        "\n",
-        "```bash\n",
-        "pip install tunix\n",
-        "```"
+        "- **Tunix**: Training utilities including `PeftTrainer` and `AutoModel` for streamlined fine-tuning"
       ]
     },
     {
@@ -197,6 +223,7 @@
         "import qwix\n",
         "from tunix.models.automodel import AutoModel\n",
         "from tunix.sft import peft_trainer, metrics_logger\n",
+        "import wandb\n",
         "\n",
         "print(f\"JAX {jax.__version__} | Devices: {jax.devices()}\")"
       ]
@@ -618,6 +645,9 @@
       },
       "outputs": [],
       "source": [
+        "#Initialize wandb\n",
+        "wandb.init()\n",
+        "\n",
         "# Quick inference test with the fine-tuned LoRA model\n",
         "prompt = \"What is the capital of France?\"\n",
         "messages = [{\"role\": \"user\", \"content\": prompt}]\n",