mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
fix(examples): Verify TrainJob Completion (#3331)
* fix(examples): Verify TrainJob Completion Signed-off-by: Andrey Velichkevich <[email protected]> * Rename client to TrainerClient() Signed-off-by: Andrey Velichkevich <[email protected]> --------- Signed-off-by: Andrey Velichkevich <[email protected]>
This commit is contained in:
Andrey Velichkevich
committed by
GitHub
GitHub
parent
72f6a855d7
commit
fbacd9f577
@@ -817,6 +817,26 @@
|
||||
" print(logline)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ccc2af7e-249b-4014-91ed-1a63c9bdc2f2",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ece885fe-5bd7-4233-bb08-66d271122268",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"TrainerClient().wait_for_job_status(name=job_id, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "030ea83c-9b1c-477a-a3f3-72d659066678",
|
||||
|
||||
@@ -561,6 +561,24 @@
|
||||
" print(\"\\n\".join(TrainerClient().get_job_logs(name=job_name, follow=True, step=f\"node-{i}\")))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
@@ -596,7 +614,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.2"
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -103,7 +103,8 @@
|
||||
" # Create the model and load it into the device.\n",
|
||||
" device = torch.device(f\"{device}:{local_rank}\")\n",
|
||||
" model=Net().to(device)\n",
|
||||
" model=torch.compile(model)\n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" model=torch.compile(model)\n",
|
||||
" model = nn.parallel.DistributedDataParallel(model)\n",
|
||||
" optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)\n",
|
||||
"\n",
|
||||
@@ -330,6 +331,26 @@
|
||||
" print(logline, end='')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "366cab2f-636f-4356-bfd5-f44f5cc7ad7f",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "024098e2-2987-4d94-af4a-5055481e3332",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cleanup",
|
||||
@@ -367,7 +388,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.12"
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -51,7 +51,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 19,
|
||||
"execution_count": null,
|
||||
"id": "train-function",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
@@ -90,8 +90,9 @@
|
||||
" else:\n",
|
||||
" device = torch.device(\"cpu\")\n",
|
||||
" model = Net().to(device)\n",
|
||||
" model=torch.compile(model)\n",
|
||||
" \n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" model=torch.compile(model)\n",
|
||||
"\n",
|
||||
" # Load MNIST dataset\n",
|
||||
" dataset = datasets.MNIST(\n",
|
||||
" './data',\n",
|
||||
@@ -277,6 +278,26 @@
|
||||
" print(logline, end='')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "25c0989b-48ea-4aaa-b85a-2cb5b64dc1f0",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "3028b661-567e-45a2-934b-0f1d01314b5c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cleanup",
|
||||
@@ -300,7 +321,7 @@
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
@@ -314,7 +335,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.13.3"
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -755,6 +755,26 @@
|
||||
" print(logline)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ca4ce2f1-7670-40ac-8174-357c0db85851",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "1870b93c-2d9e-4325-9f5c-8ec843f5caf4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"TrainerClient().wait_for_job_status(name=job_id, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "c0695485-1590-49f2-a97a-e003891fc550",
|
||||
|
||||
@@ -580,6 +580,26 @@
|
||||
" print(logline)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "555bb81d-85ac-4133-8c2c-6df99fee36a6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "000f3bdd-703a-432f-9d88-9bb8241385f4",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"TrainerClient().wait_for_job_status(name=job_id, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "bff39226-e207-4daa-87e7-60ac4dd5869d",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -355,6 +355,26 @@
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "6b4a2d2f-3341-4e4e-bcf7-10042b28dbe7",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d7e5b951-7a22-475b-9ee5-38470bfb5737",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "d29925f6-2e6b-4db1-80b7-49fedd5910e9",
|
||||
@@ -401,7 +421,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.16"
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -124,7 +124,8 @@
|
||||
" # Create the model and load it into the device.\n",
|
||||
" device = torch.device(f\"{device}:{local_rank}\")\n",
|
||||
" model =Net().to(device)\n",
|
||||
" model=torch.compile(model)\n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" model=torch.compile(model)\n",
|
||||
" model = nn.parallel.DistributedDataParallel(model)\n",
|
||||
" optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
|
||||
" \n",
|
||||
@@ -643,6 +644,24 @@
|
||||
" print(logline)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
@@ -674,6 +693,18 @@
|
||||
"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.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -38,7 +38,9 @@
|
||||
"id": "e5e86ea307b3eec9",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": "# !pip install -U kubeflow"
|
||||
"source": [
|
||||
"# !pip install -U kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
@@ -114,8 +116,8 @@
|
||||
"Requirement already satisfied: pycparser in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cffi>=1.12->cryptography>=2.1.4->azure-storage-blob>=12->cloudpathlib[all]) (2.22)\n",
|
||||
"Requirement already satisfied: pyasn1<0.7.0,>=0.6.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from pyasn1-modules>=0.2.1->google-auth<3.0dev,>=2.26.1->google-cloud-storage->cloudpathlib[all]) (0.6.1)\n",
|
||||
"\n",
|
||||
"\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m A new release of pip is available: \u001B[0m\u001B[31;49m25.0.1\u001B[0m\u001B[39;49m -> \u001B[0m\u001B[32;49m25.1\u001B[0m\n",
|
||||
"\u001B[1m[\u001B[0m\u001B[34;49mnotice\u001B[0m\u001B[1;39;49m]\u001B[0m\u001B[39;49m To update, run: \u001B[0m\u001B[32;49mpip install --upgrade pip\u001B[0m\n"
|
||||
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1\u001b[0m\n",
|
||||
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -447,21 +449,21 @@
|
||||
" 0%| | 0/40 [00:00<?, ?it/s][rank0]:[W429 01:22:58.895439547 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration, which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())\n",
|
||||
"[node-0]: [rank1]:[W429 01:22:58.895689005 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration, which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())\n",
|
||||
"100%|██████████| 40/40 [02:36<00:00, 4.10s/it]\n",
|
||||
" 0%| | 0/10 [00:00<?, ?it/s]\u001B[A\n",
|
||||
" 20%|██ | 2/10 [00:00<00:03, 2.54it/s]\u001B[A\n",
|
||||
" 30%|███ | 3/10 [00:01<00:03, 1.80it/s]\u001B[A\n",
|
||||
" 40%|████ | 4/10 [00:02<00:03, 1.55it/s]\u001B[A\n",
|
||||
" 50%|█████ | 5/10 [00:03<00:03, 1.43it/s]\u001B[A\n",
|
||||
" 60%|██████ | 6/10 [00:03<00:02, 1.37it/s]\u001B[A\n",
|
||||
" 70%|███████ | 7/10 [00:04<00:02, 1.32it/s]\u001B[A\n",
|
||||
" 80%|████████ | 8/10 [00:05<00:01, 1.30it/s]\u001B[A\n",
|
||||
" 90%|█████████ | 9/10 [00:06<00:00, 1.29it/s]\u001B[A\n",
|
||||
" \u001B[A\n",
|
||||
" 0%| | 0/10 [00:00<?, ?it/s]\u001b[A\n",
|
||||
" 20%|██ | 2/10 [00:00<00:03, 2.54it/s]\u001b[A\n",
|
||||
" 30%|███ | 3/10 [00:01<00:03, 1.80it/s]\u001b[A\n",
|
||||
" 40%|████ | 4/10 [00:02<00:03, 1.55it/s]\u001b[A\n",
|
||||
" 50%|█████ | 5/10 [00:03<00:03, 1.43it/s]\u001b[A\n",
|
||||
" 60%|██████ | 6/10 [00:03<00:02, 1.37it/s]\u001b[A\n",
|
||||
" 70%|███████ | 7/10 [00:04<00:02, 1.32it/s]\u001b[A\n",
|
||||
" 80%|████████ | 8/10 [00:05<00:01, 1.30it/s]\u001b[A\n",
|
||||
" 90%|█████████ | 9/10 [00:06<00:00, 1.29it/s]\u001b[A\n",
|
||||
" \u001b[A\n",
|
||||
"[node-0]: {'eval_loss': 5.543211936950684, 'eval_runtime': 7.9713, 'eval_samples_per_second': 2.509, 'eval_steps_per_second': 1.254, 'epoch': 1.0}\n",
|
||||
"100%|██████████| 40/40 [02:45<00:00, 4.10s/it]\n",
|
||||
"100%|██████████| 10/10 [00:07<00:00, 1.28it/s]\u001B[A\n",
|
||||
"100%|██████████| 10/10 [00:07<00:00, 1.28it/s]\u001b[A\n",
|
||||
"[node-0]: {'train_runtime': 165.12, 'train_samples_per_second': 0.484, 'train_steps_per_second': 0.242, 'train_loss': 5.764264678955078, 'epoch': 1.0}\n",
|
||||
"100%|██████████| 40/40 [02:45<00:00, 4.13s/it]\u001B[A\n"
|
||||
"100%|██████████| 40/40 [02:45<00:00, 4.13s/it]\u001b[A\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -470,6 +472,26 @@
|
||||
" print(logline)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "09b28e88-43ea-4a09-a1aa-795eb4d7b950",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "df59e53f-0512-4331-ac93-4683abc90289",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"TrainerClient().wait_for_job_status(name=job_id, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "66e6f04e-6e0c-4402-86d1-c0ccbe3f8602",
|
||||
|
||||
@@ -1,427 +1,445 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# PyTorch Speech Recognition with Kubeflow Trainer\n",
|
||||
"\n",
|
||||
"This example demonstrates how to train a speech recognition model using the [Google Speech Commands](https://www.tensorflow.org/datasets/catalog/speech_commands) dataset and PyTorch Distributed Data Parallel (DDP).\n",
|
||||
"\n",
|
||||
"It follows a development workflow designed for scale:\n",
|
||||
"1. **Define**: Wrap training logic in a self-contained function.\n",
|
||||
"2. **Test Locally**: Run the function in a local subprocess to verify code correctness.\n",
|
||||
"3. **Scale Distributed**: Submit the same function as a distributed `TrainJob` to a Kubernetes cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. Installation\n",
|
||||
"\n",
|
||||
"You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# !pip install -U kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the model dependencies (PyTorch, Torchaudio, etc)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Install model dependencies\n",
|
||||
"!pip install -q torch torchaudio librosa soundfile tensorboard"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Define Training Function\n",
|
||||
"\n",
|
||||
"We define `train_speech_recognition` to encapsulate the entire training loop.\n",
|
||||
"\n",
|
||||
"**Critical Requirement**: This function must be self-contained. It must import all necessary libraries inside itself because it will be serialized and executed in an isolated environment (container or subprocess)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def train_speech_recognition():\n",
|
||||
" \"\"\"\n",
|
||||
" Trains an Audio Transformer on the Speech Commands dataset using PyTorch DDP.\n",
|
||||
" \"\"\"\n",
|
||||
" import os\n",
|
||||
" import torch\n",
|
||||
" import torch.nn as nn\n",
|
||||
" import torch.distributed as dist\n",
|
||||
" from torch.nn.parallel import DistributedDataParallel as DDP\n",
|
||||
" from torch.utils.data import Dataset, DataLoader, DistributedSampler\n",
|
||||
" import torchaudio\n",
|
||||
" import random\n",
|
||||
" from pathlib import Path\n",
|
||||
" import platform\n",
|
||||
" from datetime import timedelta\n",
|
||||
" import warnings\n",
|
||||
"\n",
|
||||
" # Suppress warnings and progress bars for cleaner output\n",
|
||||
" warnings.filterwarnings('ignore')\n",
|
||||
" os.environ['TORCHAUDIO_DOWNLOAD_PROGRESS'] = '0'\n",
|
||||
"\n",
|
||||
" # --- 1. Environment & Configuration ---\n",
|
||||
" # Detect if we are running on Windows (Local) or Linux (Cluster/Container)\n",
|
||||
" if platform.system() == 'Windows':\n",
|
||||
" data_path = Path(\"C:/data\")\n",
|
||||
" output_dir = Path(\"C:/output\")\n",
|
||||
" else:\n",
|
||||
" data_path = Path.home() / \"data\"\n",
|
||||
" output_dir = Path.home() / \"output\"\n",
|
||||
" \n",
|
||||
" data_path.mkdir(parents=True, exist_ok=True)\n",
|
||||
" output_dir.mkdir(parents=True, exist_ok=True)\n",
|
||||
"\n",
|
||||
" # Hyperparameters\n",
|
||||
" BATCH_SIZE = 32\n",
|
||||
" EPOCHS = 3\n",
|
||||
" LR = 0.001\n",
|
||||
"\n",
|
||||
" # --- 2. DDP Initialization ---\n",
|
||||
" def setup_ddp():\n",
|
||||
" \"\"\"Initializes the distributed process group.\"\"\"\n",
|
||||
" if all(k in os.environ for k in [\"LOCAL_RANK\", \"RANK\", \"WORLD_SIZE\"]):\n",
|
||||
" local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
|
||||
" rank = int(os.environ[\"RANK\"])\n",
|
||||
" world_size = int(os.environ[\"WORLD_SIZE\"])\n",
|
||||
" \n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" backend = \"nccl\"\n",
|
||||
" device = torch.device(\"cuda\", local_rank)\n",
|
||||
" torch.cuda.set_device(device)\n",
|
||||
" else:\n",
|
||||
" backend = \"gloo\"\n",
|
||||
" device = torch.device(\"cpu\")\n",
|
||||
" \n",
|
||||
" dist.init_process_group(backend=backend, timeout=timedelta(minutes=60))\n",
|
||||
" \n",
|
||||
" print(f\"[Rank {rank}/{world_size}] DDP Initialized using {backend}\")\n",
|
||||
" return device, local_rank, rank, world_size\n",
|
||||
" else:\n",
|
||||
" print(\"[Single Process] DDP variables not found. Running in standalone mode.\")\n",
|
||||
" device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
|
||||
" return device, 0, 0, 1\n",
|
||||
"\n",
|
||||
" def cleanup_ddp():\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" dist.destroy_process_group()\n",
|
||||
"\n",
|
||||
" # --- 3. Model Architecture ---\n",
|
||||
" class AudioTransformer(nn.Module):\n",
|
||||
" def __init__(self, num_classes=35, d_model=128, nhead=4, num_layers=2):\n",
|
||||
" super().__init__()\n",
|
||||
" self.encoder = nn.TransformerEncoder(\n",
|
||||
" nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True),\n",
|
||||
" num_layers=num_layers\n",
|
||||
" )\n",
|
||||
" self.classifier = nn.Linear(d_model, num_classes)\n",
|
||||
"\n",
|
||||
" def forward(self, x):\n",
|
||||
" x = x.permute(0, 2, 1)\n",
|
||||
" x = self.encoder(x)\n",
|
||||
" x = x.mean(dim=1)\n",
|
||||
" return self.classifier(x)\n",
|
||||
"\n",
|
||||
" # --- 4. Dataset Wrapper ---\n",
|
||||
" class SpeechDataset(Dataset):\n",
|
||||
" def __init__(self, audio_paths, label_map):\n",
|
||||
" self.audio_paths = audio_paths\n",
|
||||
" self.label_map = label_map\n",
|
||||
" self.transform = torchaudio.transforms.MelSpectrogram(n_mels=128)\n",
|
||||
"\n",
|
||||
" def __len__(self): \n",
|
||||
" return len(self.audio_paths)\n",
|
||||
"\n",
|
||||
" def __getitem__(self, idx):\n",
|
||||
" try:\n",
|
||||
" path, label = self.audio_paths[idx]\n",
|
||||
"\n",
|
||||
" # FIX: Explicitly specify soundfile backend instead of torchcodec\n",
|
||||
" waveform, _ = torchaudio.load(path, backend=\"soundfile\")\n",
|
||||
" \n",
|
||||
" target_len = 16000\n",
|
||||
" if waveform.size(1) < target_len:\n",
|
||||
" waveform = torch.nn.functional.pad(waveform, (0, target_len - waveform.size(1)))\n",
|
||||
" else:\n",
|
||||
" waveform = waveform[:, :target_len]\n",
|
||||
" \n",
|
||||
" spec = self.transform(waveform).squeeze(0)\n",
|
||||
" return spec, self.label_map[label]\n",
|
||||
" except Exception as e:\n",
|
||||
" # Log error but don't crash the worker\n",
|
||||
" # print(f\"Error loading audio file: {e}\") \n",
|
||||
" return torch.zeros(128, 81), 0 \n",
|
||||
"\n",
|
||||
" def collate_fn(batch):\n",
|
||||
" specs, labels = zip(*batch)\n",
|
||||
" return torch.stack(specs), torch.tensor(labels, dtype=torch.long)\n",
|
||||
"\n",
|
||||
" # --- 5. Main Training Logic ---\n",
|
||||
" try:\n",
|
||||
" device, local_rank, rank, world_size = setup_ddp()\n",
|
||||
"\n",
|
||||
" if rank == 0:\n",
|
||||
" print(\"Verifying/Downloading dataset...\")\n",
|
||||
" torchaudio.datasets.SPEECHCOMMANDS(root=str(data_path), download=True)\n",
|
||||
" \n",
|
||||
" if dist.is_initialized():\n",
|
||||
" dist.barrier() \n",
|
||||
"\n",
|
||||
" data_root = data_path / \"SpeechCommands\" / \"speech_commands_v0.02\"\n",
|
||||
" labels = sorted([d.name for d in data_root.iterdir() if d.is_dir() and not d.name.startswith(\"_\")])\n",
|
||||
" label_map = {l: i for i, l in enumerate(labels)}\n",
|
||||
" \n",
|
||||
" all_files = []\n",
|
||||
" for label in labels:\n",
|
||||
" for f in (data_root / label).glob(\"*.wav\"):\n",
|
||||
" all_files.append((str(f), label))\n",
|
||||
" \n",
|
||||
" random.shuffle(all_files)\n",
|
||||
" train_files = all_files[:1000]\n",
|
||||
"\n",
|
||||
" dataset = SpeechDataset(train_files, label_map)\n",
|
||||
" sampler = DistributedSampler(dataset) if dist.is_initialized() else None\n",
|
||||
" loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler, \n",
|
||||
" shuffle=(sampler is None), collate_fn=collate_fn, num_workers=0)\n",
|
||||
"\n",
|
||||
" model = AudioTransformer(num_classes=len(labels)).to(device)\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" model = DDP(model, device_ids=[local_rank] if torch.cuda.is_available() else None)\n",
|
||||
"\n",
|
||||
" optimizer = torch.optim.Adam(model.parameters(), lr=LR)\n",
|
||||
" criterion = nn.CrossEntropyLoss()\n",
|
||||
"\n",
|
||||
" print(f\"[Rank {rank}] Starting training on {len(labels)} classes...\")\n",
|
||||
" for epoch in range(1, EPOCHS + 1):\n",
|
||||
" model.train()\n",
|
||||
" if sampler: \n",
|
||||
" sampler.set_epoch(epoch)\n",
|
||||
" \n",
|
||||
" total_loss = 0.0\n",
|
||||
" for i, (data, target) in enumerate(loader):\n",
|
||||
" data, target = data.to(device), target.to(device)\n",
|
||||
" optimizer.zero_grad()\n",
|
||||
" output = model(data)\n",
|
||||
" loss = criterion(output, target)\n",
|
||||
" loss.backward()\n",
|
||||
" optimizer.step()\n",
|
||||
" total_loss += loss.item()\n",
|
||||
" \n",
|
||||
" avg_loss = total_loss / len(loader)\n",
|
||||
" if rank == 0:\n",
|
||||
" print(f\"Epoch {epoch}/{EPOCHS} | Loss: {avg_loss:.4f}\")\n",
|
||||
"\n",
|
||||
" if rank == 0:\n",
|
||||
" model_path = output_dir / \"speech_model.pt\"\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" torch.save(model.module.state_dict(), model_path)\n",
|
||||
" else:\n",
|
||||
" torch.save(model.state_dict(), model_path)\n",
|
||||
" print(f\"Model saved to {model_path}\")\n",
|
||||
"\n",
|
||||
" cleanup_ddp()\n",
|
||||
"\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"[Error] Training failed: {e}\")\n",
|
||||
" import traceback\n",
|
||||
" traceback.print_exc()\n",
|
||||
" cleanup_ddp()\n",
|
||||
" raise"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 3. Run Training Locally\n",
|
||||
"\n",
|
||||
"Before scaling to a cluster, we test the code locally using the `TrainerClient`. We use the `torch-distributed` runtime which creates a local DDP environment (simulating a cluster node)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer import CustomTrainer, TrainerClient, LocalProcessBackendConfig\n",
|
||||
"\n",
|
||||
"# Initialize the client with LocalProcessBackend to run locally\n",
|
||||
"client = TrainerClient(\n",
|
||||
" backend_config=LocalProcessBackendConfig(cleanup_venv=True)\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"local_runtime = None\n",
|
||||
"for runtime in client.list_runtimes():\n",
|
||||
" if runtime.name == \"torch-distributed\":\n",
|
||||
" local_runtime = runtime\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
"if not local_runtime:\n",
|
||||
" raise ValueError(\"Local runtime 'torch-distributed' not found\")\n",
|
||||
"\n",
|
||||
"print(\"Starting local training...\")\n",
|
||||
"print(\"Note: First run will download the dataset (~2GB), which may take a few minutes.\\n\")\n",
|
||||
"\n",
|
||||
"# Start training (removed torchcodec from packages)\n",
|
||||
"local_job_name = client.train(\n",
|
||||
" trainer=CustomTrainer(\n",
|
||||
" func=train_speech_recognition,\n",
|
||||
" packages_to_install=[\"torch\", \"torchaudio\", \"librosa\", \"soundfile\", \"tensorboard\"],\n",
|
||||
" ),\n",
|
||||
" runtime=local_runtime, \n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(f\"Job {local_job_name} started. Streaming logs...\\n\")\n",
|
||||
"try:\n",
|
||||
" for logline in client.get_job_logs(local_job_name, follow=True):\n",
|
||||
" print(logline, end='')\n",
|
||||
"except KeyboardInterrupt:\n",
|
||||
" print(\"\\n\\nLog streaming interrupted by user.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 4. Run Distributed Training (Kubernetes)\n",
|
||||
"\n",
|
||||
"Now we scale the job to a Kubernetes cluster. This requires:\n",
|
||||
"1. A Kubernetes cluster with the **Kubeflow Training Operator** installed.\n",
|
||||
"2. The **TrainingRuntime** CRD (`torch-distributed`) installed on the cluster.\n",
|
||||
"3. Access via `~/.kube/config`.\n",
|
||||
"\n",
|
||||
"Based on your cluster's resources, we use a minimal configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer import CustomTrainer, TrainerClient\n",
|
||||
"\n",
|
||||
"client = TrainerClient()\n",
|
||||
"\n",
|
||||
"print(\"Submitting distributed TrainJob...\")\n",
|
||||
"\n",
|
||||
"try:\n",
|
||||
" k8s_job_name = client.train(\n",
|
||||
" trainer=CustomTrainer(\n",
|
||||
" func=train_speech_recognition,\n",
|
||||
" num_nodes=1, \n",
|
||||
" resources_per_node={\n",
|
||||
" \"cpu\": \"2\",\n",
|
||||
" \"memory\": \"4Gi\",\n",
|
||||
" },\n",
|
||||
" # Removed torchcodec from packages\n",
|
||||
" packages_to_install=[\"torch\", \"torchaudio\", \"librosa\", \"soundfile\", \"tensorboard\"],\n",
|
||||
" ),\n",
|
||||
" runtime=\"torch-distributed\", \n",
|
||||
" )\n",
|
||||
" print(f\"TrainJob '{k8s_job_name}' submitted successfully!\")\n",
|
||||
" print(f\"\\nYou can check the job status with: kubectl get trainjobs {k8s_job_name}\")\n",
|
||||
"\n",
|
||||
"except Exception as e:\n",
|
||||
" print(f\"Failed to submit TrainJob: {e}\")\n",
|
||||
" import traceback\n",
|
||||
" traceback.print_exc()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 5. Check Status & Logs\n",
|
||||
"\n",
|
||||
"Verify the job status and view real-time logs from the master replica."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!kubectl get trainjobs\n",
|
||||
"!kubectl get pods"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if 'k8s_job_name' in locals():\n",
|
||||
" print(f\"Waiting for job {k8s_job_name} to start running...\")\n",
|
||||
" \n",
|
||||
" try:\n",
|
||||
" # INCREASED TIMEOUT to 1200s (20 minutes)\n",
|
||||
" client.wait_for_job_status(name=k8s_job_name, status={\"Running\"}, timeout=1200)\n",
|
||||
" \n",
|
||||
" print(f\"Job is running! Streaming logs...\\n\")\n",
|
||||
" for logline in client.get_job_logs(k8s_job_name, follow=True):\n",
|
||||
" print(logline, end='')\n",
|
||||
" \n",
|
||||
" except KeyboardInterrupt:\n",
|
||||
" print(\"\\n\\nLog streaming interrupted by user.\")\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"\\nError while waiting for job or streaming logs: {e}\")\n",
|
||||
" # Manual fallback\n",
|
||||
" print(f\"You can manually check logs with: kubectl logs -l training.kubeflow.org/job-name={k8s_job_name}\")\n",
|
||||
"else:\n",
|
||||
" print(\"No Kubernetes job found. Please run the previous cell to create a job first.\")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.13.3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# PyTorch Speech Recognition with Kubeflow Trainer\n",
|
||||
"\n",
|
||||
"This example demonstrates how to train a speech recognition model using the [Google Speech Commands](https://www.tensorflow.org/datasets/catalog/speech_commands) dataset and PyTorch Distributed Data Parallel (DDP).\n",
|
||||
"\n",
|
||||
"It follows a development workflow designed for scale:\n",
|
||||
"1. **Define**: Wrap training logic in a self-contained function.\n",
|
||||
"2. **Test Locally**: Run the function in a local subprocess to verify code correctness.\n",
|
||||
"3. **Scale Distributed**: Submit the same function as a distributed `TrainJob` to a Kubernetes cluster."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 1. Installation\n",
|
||||
"\n",
|
||||
"You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# !pip install -U kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"Install the model dependencies (PyTorch, Torchaudio, etc)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Install model dependencies\n",
|
||||
"!pip install -q torch torchaudio librosa soundfile tensorboard"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 2. Define Training Function\n",
|
||||
"\n",
|
||||
"We define `train_speech_recognition` to encapsulate the entire training loop.\n",
|
||||
"\n",
|
||||
"**Critical Requirement**: This function must be self-contained. It must import all necessary libraries inside itself because it will be serialized and executed in an isolated environment (container or subprocess)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def train_speech_recognition():\n",
|
||||
" \"\"\"\n",
|
||||
" Trains an Audio Transformer on the Speech Commands dataset using PyTorch DDP.\n",
|
||||
" \"\"\"\n",
|
||||
" import os\n",
|
||||
" import torch\n",
|
||||
" import torch.nn as nn\n",
|
||||
" import torch.distributed as dist\n",
|
||||
" from torch.nn.parallel import DistributedDataParallel as DDP\n",
|
||||
" from torch.utils.data import Dataset, DataLoader, DistributedSampler\n",
|
||||
" import torchaudio\n",
|
||||
" import random\n",
|
||||
" from pathlib import Path\n",
|
||||
" import platform\n",
|
||||
" from datetime import timedelta\n",
|
||||
" import warnings\n",
|
||||
"\n",
|
||||
" # Suppress warnings and progress bars for cleaner output\n",
|
||||
" warnings.filterwarnings('ignore')\n",
|
||||
" os.environ['TORCHAUDIO_DOWNLOAD_PROGRESS'] = '0'\n",
|
||||
"\n",
|
||||
" # --- 1. Environment & Configuration ---\n",
|
||||
" # Detect if we are running on Windows (Local) or Linux (Cluster/Container)\n",
|
||||
" if platform.system() == 'Windows':\n",
|
||||
" data_path = Path(\"C:/data\")\n",
|
||||
" output_dir = Path(\"C:/output\")\n",
|
||||
" else:\n",
|
||||
" data_path = Path.home() / \"data\"\n",
|
||||
" output_dir = Path.home() / \"output\"\n",
|
||||
" \n",
|
||||
" data_path.mkdir(parents=True, exist_ok=True)\n",
|
||||
" output_dir.mkdir(parents=True, exist_ok=True)\n",
|
||||
"\n",
|
||||
" # Hyperparameters\n",
|
||||
" BATCH_SIZE = 32\n",
|
||||
" EPOCHS = 3\n",
|
||||
" LR = 0.001\n",
|
||||
"\n",
|
||||
" # --- 2. DDP Initialization ---\n",
|
||||
" def setup_ddp():\n",
|
||||
" \"\"\"Initializes the distributed process group.\"\"\"\n",
|
||||
" if all(k in os.environ for k in [\"LOCAL_RANK\", \"RANK\", \"WORLD_SIZE\"]):\n",
|
||||
" local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
|
||||
" rank = int(os.environ[\"RANK\"])\n",
|
||||
" world_size = int(os.environ[\"WORLD_SIZE\"])\n",
|
||||
" \n",
|
||||
" if torch.cuda.is_available():\n",
|
||||
" backend = \"nccl\"\n",
|
||||
" device = torch.device(\"cuda\", local_rank)\n",
|
||||
" torch.cuda.set_device(device)\n",
|
||||
" else:\n",
|
||||
" backend = \"gloo\"\n",
|
||||
" device = torch.device(\"cpu\")\n",
|
||||
" \n",
|
||||
" dist.init_process_group(backend=backend, timeout=timedelta(minutes=60))\n",
|
||||
" \n",
|
||||
" print(f\"[Rank {rank}/{world_size}] DDP Initialized using {backend}\")\n",
|
||||
" return device, local_rank, rank, world_size\n",
|
||||
" else:\n",
|
||||
" print(\"[Single Process] DDP variables not found. Running in standalone mode.\")\n",
|
||||
" device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
|
||||
" return device, 0, 0, 1\n",
|
||||
"\n",
|
||||
" def cleanup_ddp():\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" dist.destroy_process_group()\n",
|
||||
"\n",
|
||||
" # --- 3. Model Architecture ---\n",
|
||||
" class AudioTransformer(nn.Module):\n",
|
||||
" def __init__(self, num_classes=35, d_model=128, nhead=4, num_layers=2):\n",
|
||||
" super().__init__()\n",
|
||||
" self.encoder = nn.TransformerEncoder(\n",
|
||||
" nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True),\n",
|
||||
" num_layers=num_layers\n",
|
||||
" )\n",
|
||||
" self.classifier = nn.Linear(d_model, num_classes)\n",
|
||||
"\n",
|
||||
" def forward(self, x):\n",
|
||||
" x = x.permute(0, 2, 1)\n",
|
||||
" x = self.encoder(x)\n",
|
||||
" x = x.mean(dim=1)\n",
|
||||
" return self.classifier(x)\n",
|
||||
"\n",
|
||||
" # --- 4. Dataset Wrapper ---\n",
|
||||
" class SpeechDataset(Dataset):\n",
|
||||
" def __init__(self, audio_paths, label_map):\n",
|
||||
" self.audio_paths = audio_paths\n",
|
||||
" self.label_map = label_map\n",
|
||||
" self.transform = torchaudio.transforms.MelSpectrogram(n_mels=128)\n",
|
||||
"\n",
|
||||
" def __len__(self): \n",
|
||||
" return len(self.audio_paths)\n",
|
||||
"\n",
|
||||
" def __getitem__(self, idx):\n",
|
||||
" try:\n",
|
||||
" path, label = self.audio_paths[idx]\n",
|
||||
"\n",
|
||||
" # FIX: Explicitly specify soundfile backend instead of torchcodec\n",
|
||||
" waveform, _ = torchaudio.load(path, backend=\"soundfile\")\n",
|
||||
" \n",
|
||||
" target_len = 16000\n",
|
||||
" if waveform.size(1) < target_len:\n",
|
||||
" waveform = torch.nn.functional.pad(waveform, (0, target_len - waveform.size(1)))\n",
|
||||
" else:\n",
|
||||
" waveform = waveform[:, :target_len]\n",
|
||||
" \n",
|
||||
" spec = self.transform(waveform).squeeze(0)\n",
|
||||
" return spec, self.label_map[label]\n",
|
||||
" except Exception as e:\n",
|
||||
" # Log error but don't crash the worker\n",
|
||||
" # print(f\"Error loading audio file: {e}\") \n",
|
||||
" return torch.zeros(128, 81), 0 \n",
|
||||
"\n",
|
||||
" def collate_fn(batch):\n",
|
||||
" specs, labels = zip(*batch)\n",
|
||||
" return torch.stack(specs), torch.tensor(labels, dtype=torch.long)\n",
|
||||
"\n",
|
||||
" # --- 5. Main Training Logic ---\n",
|
||||
" try:\n",
|
||||
" device, local_rank, rank, world_size = setup_ddp()\n",
|
||||
"\n",
|
||||
" if rank == 0:\n",
|
||||
" print(\"Verifying/Downloading dataset...\")\n",
|
||||
" torchaudio.datasets.SPEECHCOMMANDS(root=str(data_path), download=True)\n",
|
||||
" \n",
|
||||
" if dist.is_initialized():\n",
|
||||
" dist.barrier() \n",
|
||||
"\n",
|
||||
" data_root = data_path / \"SpeechCommands\" / \"speech_commands_v0.02\"\n",
|
||||
" labels = sorted([d.name for d in data_root.iterdir() if d.is_dir() and not d.name.startswith(\"_\")])\n",
|
||||
" label_map = {l: i for i, l in enumerate(labels)}\n",
|
||||
" \n",
|
||||
" all_files = []\n",
|
||||
" for label in labels:\n",
|
||||
" for f in (data_root / label).glob(\"*.wav\"):\n",
|
||||
" all_files.append((str(f), label))\n",
|
||||
" \n",
|
||||
" random.shuffle(all_files)\n",
|
||||
" train_files = all_files[:1000]\n",
|
||||
"\n",
|
||||
" dataset = SpeechDataset(train_files, label_map)\n",
|
||||
" sampler = DistributedSampler(dataset) if dist.is_initialized() else None\n",
|
||||
" loader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler, \n",
|
||||
" shuffle=(sampler is None), collate_fn=collate_fn, num_workers=0)\n",
|
||||
"\n",
|
||||
" model = AudioTransformer(num_classes=len(labels)).to(device)\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" model = DDP(model, device_ids=[local_rank] if torch.cuda.is_available() else None)\n",
|
||||
"\n",
|
||||
" optimizer = torch.optim.Adam(model.parameters(), lr=LR)\n",
|
||||
" criterion = nn.CrossEntropyLoss()\n",
|
||||
"\n",
|
||||
" print(f\"[Rank {rank}] Starting training on {len(labels)} classes...\")\n",
|
||||
" for epoch in range(1, EPOCHS + 1):\n",
|
||||
" model.train()\n",
|
||||
" if sampler: \n",
|
||||
" sampler.set_epoch(epoch)\n",
|
||||
" \n",
|
||||
" total_loss = 0.0\n",
|
||||
" for i, (data, target) in enumerate(loader):\n",
|
||||
" data, target = data.to(device), target.to(device)\n",
|
||||
" optimizer.zero_grad()\n",
|
||||
" output = model(data)\n",
|
||||
" loss = criterion(output, target)\n",
|
||||
" loss.backward()\n",
|
||||
" optimizer.step()\n",
|
||||
" total_loss += loss.item()\n",
|
||||
" \n",
|
||||
" avg_loss = total_loss / len(loader)\n",
|
||||
" if rank == 0:\n",
|
||||
" print(f\"Epoch {epoch}/{EPOCHS} | Loss: {avg_loss:.4f}\")\n",
|
||||
"\n",
|
||||
" if rank == 0:\n",
|
||||
" model_path = output_dir / \"speech_model.pt\"\n",
|
||||
" if dist.is_initialized():\n",
|
||||
" torch.save(model.module.state_dict(), model_path)\n",
|
||||
" else:\n",
|
||||
" torch.save(model.state_dict(), model_path)\n",
|
||||
" print(f\"Model saved to {model_path}\")\n",
|
||||
"\n",
|
||||
" cleanup_ddp()\n",
|
||||
"\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"[Error] Training failed: {e}\")\n",
|
||||
" import traceback\n",
|
||||
" traceback.print_exc()\n",
|
||||
" cleanup_ddp()\n",
|
||||
" raise"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 3. Run Training Locally\n",
|
||||
"\n",
|
||||
"Before scaling to a cluster, we test the code locally using the `TrainerClient`. We use the `torch-distributed` runtime which creates a local DDP environment (simulating a cluster node)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer import CustomTrainer, TrainerClient, LocalProcessBackendConfig\n",
|
||||
"\n",
|
||||
"# Initialize the client with LocalProcessBackend to run locally\n",
|
||||
"client = TrainerClient(\n",
|
||||
" backend_config=LocalProcessBackendConfig(cleanup_venv=True)\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"local_runtime = None\n",
|
||||
"for runtime in client.list_runtimes():\n",
|
||||
" if runtime.name == \"torch-distributed\":\n",
|
||||
" local_runtime = runtime\n",
|
||||
" break\n",
|
||||
"\n",
|
||||
"if not local_runtime:\n",
|
||||
" raise ValueError(\"Local runtime 'torch-distributed' not found\")\n",
|
||||
"\n",
|
||||
"print(\"Starting local training...\")\n",
|
||||
"print(\"Note: First run will download the dataset (~2GB), which may take a few minutes.\\n\")\n",
|
||||
"\n",
|
||||
"# Start training (removed torchcodec from packages)\n",
|
||||
"local_job_name = client.train(\n",
|
||||
" trainer=CustomTrainer(\n",
|
||||
" func=train_speech_recognition,\n",
|
||||
" packages_to_install=[\"torch\", \"torchaudio\", \"librosa\", \"soundfile\", \"tensorboard\"],\n",
|
||||
" ),\n",
|
||||
" runtime=local_runtime, \n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(f\"Job {local_job_name} started. Streaming logs...\\n\")\n",
|
||||
"try:\n",
|
||||
" for logline in client.get_job_logs(local_job_name, follow=True):\n",
|
||||
" print(logline, end='')\n",
|
||||
"except KeyboardInterrupt:\n",
|
||||
" print(\"\\n\\nLog streaming interrupted by user.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 4. Run Distributed Training (Kubernetes)\n",
|
||||
"\n",
|
||||
"Now we scale the job to a Kubernetes cluster. This requires:\n",
|
||||
"1. A Kubernetes cluster with the **Kubeflow Training Operator** installed.\n",
|
||||
"2. The **TrainingRuntime** CRD (`torch-distributed`) installed on the cluster.\n",
|
||||
"3. Access via `~/.kube/config`.\n",
|
||||
"\n",
|
||||
"Based on your cluster's resources, we use a minimal configuration."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer import CustomTrainer, TrainerClient\n",
|
||||
"\n",
|
||||
"client = TrainerClient()\n",
|
||||
"\n",
|
||||
"print(\"Submitting distributed TrainJob...\")\n",
|
||||
"\n",
|
||||
"try:\n",
|
||||
" k8s_job_name = client.train(\n",
|
||||
" trainer=CustomTrainer(\n",
|
||||
" func=train_speech_recognition,\n",
|
||||
" num_nodes=1, \n",
|
||||
" resources_per_node={\n",
|
||||
" \"cpu\": \"2\",\n",
|
||||
" \"memory\": \"4Gi\",\n",
|
||||
" },\n",
|
||||
" # Removed torchcodec from packages\n",
|
||||
" packages_to_install=[\"torch\", \"torchaudio\", \"librosa\", \"soundfile\", \"tensorboard\"],\n",
|
||||
" ),\n",
|
||||
" runtime=\"torch-distributed\", \n",
|
||||
" )\n",
|
||||
" print(f\"TrainJob '{k8s_job_name}' submitted successfully!\")\n",
|
||||
" print(f\"\\nYou can check the job status with: kubectl get trainjobs {k8s_job_name}\")\n",
|
||||
"\n",
|
||||
"except Exception as e:\n",
|
||||
" print(f\"Failed to submit TrainJob: {e}\")\n",
|
||||
" import traceback\n",
|
||||
" traceback.print_exc()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 5. Check Status & Logs\n",
|
||||
"\n",
|
||||
"Verify the job status and view real-time logs from the master replica."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!kubectl get trainjobs\n",
|
||||
"!kubectl get pods"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if 'k8s_job_name' in locals():\n",
|
||||
" print(f\"Waiting for job {k8s_job_name} to start running...\")\n",
|
||||
" \n",
|
||||
" try:\n",
|
||||
" # INCREASED TIMEOUT to 1200s (20 minutes)\n",
|
||||
" client.wait_for_job_status(name=k8s_job_name, status={\"Running\"}, timeout=1200)\n",
|
||||
" \n",
|
||||
" print(f\"Job is running! Streaming logs...\\n\")\n",
|
||||
" for logline in client.get_job_logs(k8s_job_name, follow=True):\n",
|
||||
" print(logline, end='')\n",
|
||||
" \n",
|
||||
" except KeyboardInterrupt:\n",
|
||||
" print(\"\\n\\nLog streaming interrupted by user.\")\n",
|
||||
" except Exception as e:\n",
|
||||
" print(f\"\\nError while waiting for job or streaming logs: {e}\")\n",
|
||||
" # Manual fallback\n",
|
||||
" print(f\"You can manually check logs with: kubectl logs -l training.kubeflow.org/job-name={k8s_job_name}\")\n",
|
||||
"else:\n",
|
||||
" print(\"No Kubernetes job found. Please run the previous cell to create a job first.\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## 6. Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=k8s_job_name, timeout=20)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
|
||||
@@ -332,6 +332,26 @@
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "79ec9db1-af27-453e-84aa-fe23213be81e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "9f708a9e-632b-40d0-91e8-25f065577777",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "042ebbb6",
|
||||
|
||||
@@ -1,248 +1,243 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7c49a6d5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Fine-tune Qwe2.5-1.5B with Alpaca Dataset"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e3b15eff",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This example demonstrates how to fine-tune Qwen2.5-1.5B model with the Alpaca Dataset using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK.\n",
|
||||
"\n",
|
||||
"This notebooks walks you through the prerequisites of using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK, and how to submit TrainJob to bootstrap the fine-tuning workflow.\n",
|
||||
"\n",
|
||||
"Qwen2.5-1.5B: https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct\n",
|
||||
"\n",
|
||||
"Alpaca Dataset: https://huggingface.co/datasets/tatsu-lab/alpaca"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ed4a60eb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Install the Kubeflow SDK\n",
|
||||
"\n",
|
||||
"You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "288ec515",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7211fbf9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Prerequisites\n",
|
||||
"\n",
|
||||
"### Install Official Training Runtimes\n",
|
||||
"\n",
|
||||
"You need to make sure that you've installed the Kubeflow Trainer Controller Manager and Kubeflow Training Runtimes mentioned in the [installation guide](https://www.kubeflow.org/docs/components/trainer/operator-guides/installation/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d35e4fd8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# List all available Kubeflow Training Runtimes.\n",
|
||||
"from kubeflow.trainer import *\n",
|
||||
"from kubeflow_trainer_api import models\n",
|
||||
"\n",
|
||||
"client = TrainerClient()\n",
|
||||
"for runtime in client.list_runtimes():\n",
|
||||
" print(runtime)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "67490f66",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Bootstrap LLM Fine-tuning Workflow\n",
|
||||
"\n",
|
||||
"Kubeflow TrainJob will train the model in the referenced (Cluster)TrainingRuntime."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "641fae4d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"job_name = client.train(\n",
|
||||
" runtime=\"torchtune-qwen2.5-1.5b\",\n",
|
||||
" initializer=Initializer(\n",
|
||||
" dataset=HuggingFaceDatasetInitializer(\n",
|
||||
" storage_uri=\"hf://tatsu-lab/alpaca/data\"\n",
|
||||
" ),\n",
|
||||
" model=HuggingFaceModelInitializer(\n",
|
||||
" storage_uri=\"hf://Qwen/Qwen2.5-1.5B-Instruct\",\n",
|
||||
" )\n",
|
||||
" ),\n",
|
||||
" trainer=BuiltinTrainer(\n",
|
||||
" config=TorchTuneConfig(\n",
|
||||
" dataset_preprocess_config=TorchTuneInstructDataset(\n",
|
||||
" source=DataFormat.PARQUET, split=\"train[:1000]\"\n",
|
||||
" ),\n",
|
||||
" resources_per_node={\n",
|
||||
" \"memory\": \"128G\",\n",
|
||||
" \"gpu\": 2,\n",
|
||||
" },\n",
|
||||
" \n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ee5fbe8e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Wait for running status"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "53eaa65a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"\n",
|
||||
"# Wait for the running status.\n",
|
||||
"client.wait_for_job_status(name=job_name, status={\"Running\"})\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "75a82b76",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Watch the TrainJob Logs\n",
|
||||
"\n",
|
||||
"We can use the `get_job_logs()` API to get the TrainJob logs.\n",
|
||||
"\n",
|
||||
"### Dataset Initializer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "68e9d454",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer.constants import constants\n",
|
||||
"\n",
|
||||
"for line in client.get_job_logs(job_name, follow=True, step=constants.DATASET_INITIALIZER):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9f970c03",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Model Initializer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4ce8c1df",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for line in client.get_job_logs(job_name, follow=True, step=constants.MODEL_INITIALIZER):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b67775ea",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Trainer Node "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ae672d7b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for c in client.get_job(name=job_name).steps:\n",
|
||||
" print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")\n",
|
||||
"\n",
|
||||
"for line in client.get_job_logs(job_name, follow=True):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7768e157",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"job = client.wait_for_job_status(name=job_name, timeout=1200)\n",
|
||||
"\n",
|
||||
"if job.status != \"Complete\":\n",
|
||||
" raise RuntimeError(\"TrainJob failed\")\n",
|
||||
"else:\n",
|
||||
" print(\"TrainJob completed successfully\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "042ebbb6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Get the Fine-tuned Model\n",
|
||||
"\n",
|
||||
"After Trainer node completes the fine-tuning task, the fine-tuned model will be stored into the `/workspace/output` directory, which can be shared across Pods through PVC mounting. You can find it in another Pod's `/<mountDir>/output` directory if you mount the PVC under `/<mountDir>`."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.12.3"
|
||||
}
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7c49a6d5",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Fine-tune Qwe2.5-1.5B with Alpaca Dataset"
|
||||
]
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "e3b15eff",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"This example demonstrates how to fine-tune Qwen2.5-1.5B model with the Alpaca Dataset using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK.\n",
|
||||
"\n",
|
||||
"This notebooks walks you through the prerequisites of using TorchTune `BuiltinTrainer` from Kubeflow Trainer SDK, and how to submit TrainJob to bootstrap the fine-tuning workflow.\n",
|
||||
"\n",
|
||||
"Qwen2.5-1.5B: https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct\n",
|
||||
"\n",
|
||||
"Alpaca Dataset: https://huggingface.co/datasets/tatsu-lab/alpaca"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ed4a60eb",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Install the Kubeflow SDK\n",
|
||||
"\n",
|
||||
"You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "288ec515",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install kubeflow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "7211fbf9",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Prerequisites\n",
|
||||
"\n",
|
||||
"### Install Official Training Runtimes\n",
|
||||
"\n",
|
||||
"You need to make sure that you've installed the Kubeflow Trainer Controller Manager and Kubeflow Training Runtimes mentioned in the [installation guide](https://www.kubeflow.org/docs/components/trainer/operator-guides/installation/)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "d35e4fd8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# List all available Kubeflow Training Runtimes.\n",
|
||||
"from kubeflow.trainer import *\n",
|
||||
"from kubeflow_trainer_api import models\n",
|
||||
"\n",
|
||||
"client = TrainerClient()\n",
|
||||
"for runtime in client.list_runtimes():\n",
|
||||
" print(runtime)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "67490f66",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Bootstrap LLM Fine-tuning Workflow\n",
|
||||
"\n",
|
||||
"Kubeflow TrainJob will train the model in the referenced (Cluster)TrainingRuntime."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "641fae4d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"job_name = client.train(\n",
|
||||
" runtime=\"torchtune-qwen2.5-1.5b\",\n",
|
||||
" initializer=Initializer(\n",
|
||||
" dataset=HuggingFaceDatasetInitializer(\n",
|
||||
" storage_uri=\"hf://tatsu-lab/alpaca/data\"\n",
|
||||
" ),\n",
|
||||
" model=HuggingFaceModelInitializer(\n",
|
||||
" storage_uri=\"hf://Qwen/Qwen2.5-1.5B-Instruct\",\n",
|
||||
" )\n",
|
||||
" ),\n",
|
||||
" trainer=BuiltinTrainer(\n",
|
||||
" config=TorchTuneConfig(\n",
|
||||
" dataset_preprocess_config=TorchTuneInstructDataset(\n",
|
||||
" source=DataFormat.PARQUET, split=\"train[:1000]\"\n",
|
||||
" ),\n",
|
||||
" resources_per_node={\n",
|
||||
" \"memory\": \"128G\",\n",
|
||||
" \"gpu\": 2,\n",
|
||||
" },\n",
|
||||
" \n",
|
||||
" )\n",
|
||||
" )\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "ee5fbe8e",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Wait for running status"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "53eaa65a",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"\n",
|
||||
"# Wait for the running status.\n",
|
||||
"client.wait_for_job_status(name=job_name, status={\"Running\"})\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "75a82b76",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Watch the TrainJob Logs\n",
|
||||
"\n",
|
||||
"We can use the `get_job_logs()` API to get the TrainJob logs.\n",
|
||||
"\n",
|
||||
"### Dataset Initializer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "68e9d454",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from kubeflow.trainer.constants import constants\n",
|
||||
"\n",
|
||||
"for line in client.get_job_logs(job_name, follow=True, step=constants.DATASET_INITIALIZER):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "9f970c03",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Model Initializer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "4ce8c1df",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for line in client.get_job_logs(job_name, follow=True, step=constants.MODEL_INITIALIZER):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "b67775ea",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Trainer Node "
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ae672d7b",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for c in client.get_job(name=job_name).steps:\n",
|
||||
" print(f\"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\\n\")\n",
|
||||
"\n",
|
||||
"for line in client.get_job_logs(job_name, follow=True):\n",
|
||||
" print(line)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7768e157",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"job = client.wait_for_job_status(name=job_name, timeout=1200)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "042ebbb6",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# Get the Fine-tuned Model\n",
|
||||
"\n",
|
||||
"After Trainer node completes the fine-tuning task, the fine-tuned model will be stored into the `/workspace/output` directory, which can be shared across Pods through PVC mounting. You can find it in another Pod's `/<mountDir>/output` directory if you mount the PVC under `/<mountDir>`."
|
||||
]
|
||||
}
|
||||
],
|
||||
"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.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
|
||||
@@ -294,6 +294,24 @@
|
||||
" print(\"\\n\".join(TrainerClient().get_job_logs(name=job_name, follow=True, step=f\"node-{i}\")))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Verify TrainJob Completion\n",
|
||||
"\n",
|
||||
"Check the TrainJob status to ensure it completed successfully."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"client.wait_for_job_status(name=job_name, timeout=20)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
@@ -327,9 +345,9 @@
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbformat_minor": 4,
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.12.2"
|
||||
"version": "3.11.13"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
Reference in New Issue
Block a user