wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/train/examples/intel_gaudi/bert.ipynb
+416
View File
@@ -0,0 +1,416 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_bert_training)=\n",
"# BERT Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-bert\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-bert\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this notebook, we will train a BERT model for sequence classification using the Yelp review full dataset. We will use the `transformers` and `datasets` libraries from Hugging Face, along with `ray.train` for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install the following dependencies to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook transformers datasets evaluate\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
" warnings.warn(\n"
]
}
],
"source": [
"# Import necessary libraries\n",
"\n",
"import os\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"from torch import nn\n",
"from torch.utils.data import DataLoader\n",
"from tqdm import tqdm\n",
"\n",
"import numpy as np\n",
"import evaluate\n",
"from datasets import load_dataset\n",
"import transformers\n",
"from transformers import (\n",
" Trainer,\n",
" TrainingArguments,\n",
" AutoTokenizer,\n",
" AutoModelForSequenceClassification,\n",
")\n",
"\n",
"import ray.train\n",
"from ray.train import ScalingConfig\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Metrics Setup\n",
"\n",
"We will use accuracy as our evaluation metric. The `compute_metrics` function will calculate the accuracy of our model's predictions."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Metrics\n",
"metric = evaluate.load(\"accuracy\")\n",
"\n",
"def compute_metrics(eval_pred):\n",
" logits, labels = eval_pred\n",
" predictions = np.argmax(logits, axis=-1)\n",
" return metric.compute(predictions=predictions, references=labels)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training. It handles data loading, tokenization, model initialization, and the training loop. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def train_func_per_worker(config: Dict):\n",
" \n",
" # Datasets\n",
" dataset = load_dataset(\"yelp_review_full\")\n",
" tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
" \n",
" def tokenize_function(examples):\n",
" return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
"\n",
" lr = config[\"lr\"]\n",
" epochs = config[\"epochs\"]\n",
" batch_size = config[\"batch_size_per_worker\"]\n",
"\n",
" train_dataset = dataset[\"train\"].select(range(1000)).map(tokenize_function, batched=True)\n",
" eval_dataset = dataset[\"test\"].select(range(1000)).map(tokenize_function, batched=True)\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = {}\n",
" dataloaders[\"train\"] = torch.utils.data.DataLoader(\n",
" train_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
" dataloaders[\"test\"] = torch.utils.data.DataLoader(\n",
" eval_dataset, \n",
" shuffle=True, \n",
" collate_fn=transformers.default_data_collator, \n",
" batch_size=batch_size\n",
" )\n",
"\n",
" # Obtain HPU device automatically\n",
" device = ray.train.torch.get_device()\n",
"\n",
" # Prepare model and optimizer\n",
" model = AutoModelForSequenceClassification.from_pretrained(\n",
" \"bert-base-cased\", num_labels=5\n",
" )\n",
" model = model.to(device)\n",
" \n",
" optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)\n",
"\n",
" # Start training loops\n",
" for epoch in range(epochs):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"test\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" # breakpoint()\n",
" for batch in dataloaders[phase]:\n",
" batch = {k: v.to(device) for k, v in batch.items()}\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" \n",
" outputs = model(**batch)\n",
" loss = outputs.loss\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
" print(f\"train epoch:[{epoch}]\\tloss:{loss:.6f}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_bert` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def train_bert(num_workers=2):\n",
" global_batch_size = 8\n",
"\n",
" train_config = {\n",
" \"lr\": 1e-3,\n",
" \"epochs\": 10,\n",
" \"batch_size_per_worker\": global_batch_size // num_workers,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" # start your ray cluster\n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_bert` function to start the training process. You can adjust the number of workers to use.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_bert(num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"Downloading builder script: 100%|██████████| 4.20k/4.20k [00:00<00:00, 27.0MB/s]\n",
"2025-03-03 03:37:08,776 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=75123) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=74734) Started distributed worker processes: \n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75123) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=74734) - (node_id=eef984cd0cd96cce50bad1b1dab12e19c809047f10be3c829524a3d1, ip=100.83.111.228, pid=75122) world_rank=1, local_rank=1, node_rank=0\n",
"Generating train split: 0%| | 0/650000 [00:00<?, ? examples/s]\n",
"Generating train split: 7%|▋ | 45000/650000 [00:00<00:01, 435976.18 examples/s]\n",
"Generating train split: 15%|█▍ | 95000/650000 [00:00<00:01, 469481.51 examples/s]\n",
"Generating train split: 23%|██▎ | 150000/650000 [00:00<00:01, 477676.99 examples/s]\n",
"Generating train split: 31%|███ | 203000/650000 [00:00<00:00, 493746.70 examples/s]\n",
"Generating train split: 43%|████▎ | 279000/650000 [00:00<00:00, 499340.09 examples/s]\n",
"Generating train split: 55%|█████▍ | 355000/650000 [00:00<00:00, 498613.65 examples/s]\n",
"Generating train split: 66%|██████▋ | 431000/650000 [00:00<00:00, 497799.19 examples/s]\n",
"Generating train split: 78%|███████▊ | 506000/650000 [00:01<00:00, 495696.93 examples/s]\n",
"Generating train split: 86%|████████▌ | 556000/650000 [00:01<00:00, 494508.05 examples/s]\n",
"Generating train split: 94%|█████████▎| 609000/650000 [00:01<00:00, 490725.53 examples/s]\n",
"Generating train split: 100%|██████████| 650000/650000 [00:01<00:00, 494916.42 examples/s]\n",
"Generating test split: 0%| | 0/50000 [00:00<?, ? examples/s]\n",
"Generating test split: 100%|██████████| 50000/50000 [00:00<00:00, 509619.87 examples/s]\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3998.33 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 4051.80 examples/s]\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3869.20 examples/s]\n",
"(RayTrainWorker pid=75123) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75123) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"Map: 0%| | 0/1000 [00:00<?, ? examples/s] [repeated 3x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"Map: 100%|██████████| 1000/1000 [00:00<00:00, 3782.66 examples/s] [repeated 2x across cluster]\n",
"(RayTrainWorker pid=75123) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=75123) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=75123) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=75123) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=75123) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=75123) Num CPU Cores : 160\n",
"(RayTrainWorker pid=75123) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=75123) ------------------------------------------------------------------------------\n",
"2025-03-03 03:41:04,658 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-37-11' in 0.0020s.\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-37-11\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-37-06_983992_65223/artifacts/2025-03-03_03-37-11/TorchTrainer_2025-03-03_03-37-11/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭─────────────────────────────────────────────────╮\n",
"│ Training config │\n",
"├─────────────────────────────────────────────────┤\n",
"│ train_loop_config/batch_size_per_worker 4 │\n",
"│ train_loop_config/epochs 10 │\n",
"│ train_loop_config/lr 0.001 │\n",
"╰─────────────────────────────────────────────────╯\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.979938\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.756611 [repeated 36x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.643875 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[0] loss:1.416416 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.272513 [repeated 107x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:2.086884 [repeated 155x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[1] loss:1.426217 [repeated 178x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[1] loss:0.991381 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.294097 [repeated 28x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.386306 [repeated 169x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[2] loss:1.190416 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.171733 [repeated 130x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.287821 [repeated 152x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[3] loss:1.055692 [repeated 179x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[3] loss:1.677789 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.942071 [repeated 19x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:1.592500 [repeated 167x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[4] loss:0.936934 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:2.465384 [repeated 141x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.659170 [repeated 156x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[5] loss:1.850438 [repeated 180x across cluster]\n",
"(RayTrainWorker pid=75122) train epoch:[5] loss:1.101623 [repeated 160x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:2.125591 [repeated 18x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.612838 [repeated 170x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[6] loss:1.759160 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.338552 [repeated 139x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.467959 [repeated 157x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[7] loss:1.682137 [repeated 181x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.395805 [repeated 162x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.527835 [repeated 153x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[8] loss:1.672311 [repeated 177x across cluster]\n",
"(RayTrainWorker pid=75123) \n",
"(RayTrainWorker pid=75122) train epoch:[8] loss:1.093186 [repeated 166x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.457587 [repeated 13x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.727377 [repeated 171x across cluster]\n",
"(RayTrainWorker pid=75123) train epoch:[9] loss:1.694001 [repeated 182x across cluster]\n",
"\n",
"Training completed after 0 iterations at 2025-03-03 03:41:04. Total running time: 3min 53s\n",
"\n",
"Training result: Result(\n",
" metrics={},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-37-11/TorchTrainer_ca6cf_00000_0_2025-03-03_03-37-11',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=75122) Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"(RayTrainWorker pid=75122) You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"(RayTrainWorker pid=75122) train epoch:[9] loss:0.417845 [repeated 136x across cluster]\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"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"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/llama.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Fine-tuning Llama-2 Model with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will:\n",
"- fine-tuning a [Llama-2-7b](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf) model by using Intel Gaudi accelerators with DDP method\n",
"- fine-tuning a [Llama-2-70b](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) model by using Intel Gaudi accelerators with DeepSpeed method\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training. We will use dataset [tatsu-lab/alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca).\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this fine-tuning example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- LoRA training.\n",
"- DDP or DeepSpeed based method.\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- Llama-2-7b/Llama-2-70b model.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import copy\n",
"from typing import Dict\n",
"\n",
"import torch\n",
"\n",
"import datasets\n",
"import transformers\n",
"from transformers import DataCollatorForLanguageModeling\n",
"\n",
"from tqdm import tqdm\n",
"\n",
"import peft\n",
"\n",
"from optimum.habana import GaudiTrainer, GaudiConfig, GaudiTrainingArguments\n",
"from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare Dataset Function\n",
"\n",
"Preprocessing the raw dataset's each line with specified format."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset(raw_datasets):\n",
"\n",
" PROMPT_DICT = {\n",
" \"prompt_with_input\": (\n",
" \"Below is an instruction that describes a task, paired with an input that provides further context. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input}\\n\\n### Response:\"\n",
" ),\n",
" \"prompt_without_input\": (\n",
" \"Below is an instruction that describes a task. \"\n",
" \"Write a response that appropriately completes the request.\\n\\n\"\n",
" \"### Instruction:\\n{instruction}\\n\\n### Response:\"\n",
" ),\n",
" }\n",
"\n",
" def create_prompts(examples):\n",
" prompts = {}\n",
" prompts[\"source\"] = []\n",
" prompts[\"target\"] = []\n",
" for example in examples:\n",
" prompt_template = (\n",
" PROMPT_DICT[\"prompt_with_input\"] if example[\"input\"] != \"\" else PROMPT_DICT[\"prompt_without_input\"]\n",
" )\n",
" source = prompt_template.format_map(example)\n",
" prompts[\"source\"].append(source)\n",
" prompts[\"target\"].append(example[\"output\"])\n",
" return prompts\n",
"\n",
" # Preprocessing the datasets.\n",
" for key in raw_datasets:\n",
" prompts = create_prompts(raw_datasets[key])\n",
" columns_to_be_removed = list(raw_datasets[key].features.keys())\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_sources\", prompts[\"source\"])\n",
" raw_datasets[key] = raw_datasets[key].add_column(\"prompt_targets\", prompts[\"target\"])\n",
" raw_datasets[key] = raw_datasets[key].remove_columns(columns_to_be_removed)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset to Tokenizer Function\n",
"\n",
"Tokenize each line in dataset by model tokenizer.\n",
"\n",
"In example codes, we concatenate the dataset's line content to accelerate training speed.\n",
"\n",
"All datasets are processed as \"train\" datasets, no evaluation datasets are sampled from raw_datasets."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def preprocess_dataset_to_tokenizer(raw_datasets, tokenizer):\n",
" max_seq_length = 512\n",
" tokenizer.pad_token_id = 0\n",
" tokenizer.eos_token_id = 1\n",
" tokenizer.bos_token_id = 2\n",
"\n",
" def tokenize(prompt, add_eos_token=True):\n",
" results = tokenizer(\n",
" prompt,\n",
" truncation=True,\n",
" max_length=max_seq_length,\n",
" padding=False,\n",
" return_tensors=None,\n",
" )\n",
" for i in range(len(results[\"input_ids\"])):\n",
" if (\n",
" results[\"input_ids\"][i][-1] != tokenizer.eos_token_id\n",
" and len(results[\"input_ids\"][i]) < max_seq_length\n",
" and add_eos_token\n",
" ):\n",
" results[\"input_ids\"][i].append(tokenizer.eos_token_id)\n",
" results[\"attention_mask\"][i].append(1)\n",
"\n",
" results[\"labels\"] = copy.deepcopy(results[\"input_ids\"])\n",
" results[\"input_id_len\"] = [len(result) for result in results[\"input_ids\"]]\n",
" return results\n",
"\n",
" def preprocess_function(examples):\n",
" keys = list(examples.data.keys())\n",
" if len(keys) != 2:\n",
" raise ValueError(\"Unsupported dataset format\")\n",
"\n",
" st = [s + t for s, t in zip(examples[keys[0]], examples[keys[1]])]\n",
"\n",
" examples_tokenized = tokenize(st)\n",
" input_ids = examples_tokenized[\"input_ids\"]\n",
" labels = examples_tokenized[\"labels\"]\n",
" return {\n",
" \"input_ids\": input_ids,\n",
" \"labels\": labels,\n",
" \"attention_mask\": examples_tokenized[\"attention_mask\"],\n",
" }\n",
"\n",
" tokenized_datasets = raw_datasets.map(\n",
" preprocess_function,\n",
" batched=True,\n",
" load_from_cache_file=True,\n",
" )\n",
"\n",
" def concatenate_data(dataset, max_seq_length):\n",
" concatenated_dataset = {}\n",
" for column in dataset.features:\n",
" concatenated_data = [item for sample in dataset[column] for item in sample]\n",
" reshaped_data = [\n",
" concatenated_data[i * max_seq_length : (i + 1) * max_seq_length]\n",
" for i in range(len(concatenated_data) // max_seq_length)\n",
" ]\n",
" concatenated_dataset[column] = reshaped_data\n",
" return datasets.Dataset.from_dict(concatenated_dataset)\n",
"\n",
" tokenized_datasets_ = tokenized_datasets[\"train\"].remove_columns([\"prompt_sources\", \"prompt_targets\"])\n",
" tokenized_datasets[\"train\"] = concatenate_data(tokenized_datasets_, max_seq_length)\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare training arguments\n",
"\n",
"here some arguments are hard coded, you can pass arguments from `config`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_training_args(config: Dict):\n",
" # prepare execution mode config\n",
" execution_mode = config[\"execution_mode\"]\n",
" use_lazy_mode = True if execution_mode == \"lazy\" else False\n",
" torch_compile_backend = \"hpu_backend\" if execution_mode == \"eager.compile\" else None\n",
"\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
"\n",
" return GaudiTrainingArguments(deepspeed=deepspeed,\n",
" output_dir=config[\"output\"],\n",
" do_train=True,\n",
" do_eval=False,\n",
" per_device_train_batch_size=config[\"batch_size_per_worker\"],\n",
" bf16=True,\n",
" learning_rate=config[\"lr\"],\n",
" save_strategy=\"no\",\n",
" torch_compile_backend=torch_compile_backend,\n",
" evaluation_strategy=\"no\",\n",
" lr_scheduler_type=\"cosine\",\n",
" num_train_epochs=config[\"epochs\"],\n",
" use_lazy_mode=use_lazy_mode,\n",
" use_habana=True,\n",
" pipelining_fwd_bwd=True,\n",
" save_only_model=True,\n",
" gradient_checkpointing=True,\n",
" warmup_ratio=0.03,\n",
" throughput_warmup_steps=3,\n",
" logging_steps=5)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare model\n",
"\n",
"1. download model from huggingface or read model from local directory.\n",
"2. convert model to lora model.\n",
"3. move model to HPU device.\n",
"\n",
"If you doesn't want to fine-tune with LoRA, just remove LoRA conversion step."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def prepare_model(config: Dict, device):\n",
" # prepare from pretrained model\n",
" deepspeed = config[\"deepspeed\"] if \"deepspeed\" in config else None\n",
" if deepspeed is not None:\n",
" auto_config = transformers.AutoConfig.from_pretrained(config[\"model\"], use_cache=False, revision=\"main\", use_auth_token=None, trust_remote_code=None)\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], config=auto_config, **config[\"model_config\"])\n",
" model.generation_config.attn_softmax_bf16 = True\n",
" model.generation_config.use_flash_attention = True\n",
" else:\n",
" model = transformers.AutoModelForCausalLM.from_pretrained(config[\"model\"], **config[\"model_config\"])\n",
" model.enable_input_require_grads()\n",
"\n",
" # convert to peft model for lora training\n",
" peft_config = peft.LoraConfig(**config[\"lora_config\"])\n",
" model = peft.get_peft_model(model, peft_config)\n",
"\n",
" model.to(dtype=config[\"model_config\"][\"torch_dtype\"], device=device)\n",
"\n",
" return model\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"\n",
"- preparing training args, an instance of `GaudiTrainingArguments`.\n",
"- loading datasets and preprocess datasets, just load the first 4096 item as training datasets.\n",
"- loading pretrained model as tokenizer, and process datasets to tokenizer.\n",
"- loading pretrained model.\n",
"- preparing data collator and gaidu_config.\n",
"- preparing instance of `GaudiTrainer`.\n",
"- calling `train()` to train model.\n",
"- saving model results.\n",
"\n",
"Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"- Detect HPU and set the device.\n",
"- Initialize the habana PyTorch backend.\n",
"- Initialize the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_func_per_worker(config: Dict):\n",
" # adapt transformers to gaudi\n",
" adapt_transformers_to_gaudi()\n",
"\n",
" # prepare training arguments\n",
" training_args = prepare_training_args(config)\n",
"\n",
" # prepare datasets\n",
" # here we use dataset \"tatsu-lab/alpaca\" from huggingface\n",
" raw_datasets = datasets.DatasetDict({\"train\": datasets.load_dataset(\"tatsu-lab/alpaca\", split='train[0:4096]')})\n",
" preprocess_dataset(raw_datasets)\n",
"\n",
" # prepare tokenizer\n",
" tokenizer = transformers.AutoTokenizer.from_pretrained(config[\"model\"])\n",
" tokenized_datasets = preprocess_dataset_to_tokenizer(raw_datasets, tokenizer)\n",
"\n",
" # prepare model\n",
" model = prepare_model(config, training_args.device)\n",
"\n",
" # prepare data collator\n",
" data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors=\"pt\", mlm=False)\n",
"\n",
" # prepare gaudi config\n",
" gaudi_config = GaudiConfig()\n",
" gaudi_config.use_fused_adam = True\n",
" gaudi_config.use_fused_clip_norm = True\n",
"\n",
" # instance GaudiTrainer\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=tokenized_datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=data_collator,\n",
" compute_metrics=None,\n",
" preprocess_logits_for_metrics=None,\n",
" )\n",
"\n",
" train_result = trainer.train()\n",
" print(f\"train_result = {train_result}\")\n",
" trainer.save_model()\n",
"\n",
" return train_result"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"The `train_llama` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Set training method, supported method are:\n",
" - \"ddp\"\n",
" - \"deepspeed\"\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def train_llama(num_workers, execution_mode, training_method):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" # deepspeed config, can also place it to config file\n",
" deepspeed_config = {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" }\n",
"\n",
" # Preparing train configurations\n",
" train_config = {\n",
" \"execution_mode\": execution_mode,\n",
" \"model\": \"/root/models/models--meta-llama--Llama-2-70b-chat-hf/snapshots/e9149a12809580e8602995856f8098ce973d1080/\",\n",
" \"model_config\": {\"torch_dtype\": torch.bfloat16, \"trust_remote_code\": False, \"use_auth_token\": None},\n",
" \"lora_config\": {\"task_type\": \"CAUSAL_LM\", \"r\": 8, \"lora_alpha\": 32, \"lora_dropout\": 0.1, \"target_modules\": [\"q_proj\", \"v_proj\"]},\n",
" \"lr\": 1e-4,\n",
" \"epochs\": 2,\n",
" \"batch_size_per_worker\": 8,\n",
" \"output\": \"/tmp/ray/\",\n",
" \"deepspeed\": deepspeed_config if training_method == \"deepspeed\" else None,\n",
" }\n",
"\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
"\n",
" # start your ray cluster\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_func_per_worker,\n",
" train_loop_config=train_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_llama` function to start the training process. You can adjust the number of workers to use, and the execution mode for HPU."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_DEVICES, such as\n",
"# os.environ[\"HABANA_VISIBLE_DEVICES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"# training_method are [\"ddp\", \"deepspeed\"]\n",
"training_method = \"deepspeed\"\n",
"if training_method == \"deepspeed\":\n",
" os.environ[\"PT_HPU_MAX_COMPOUND_OP_SIZE\"] = \"10\"\n",
" os.environ[\"DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED\"] = \"1\"\n",
"\n",
"# here use 4 HPUs\n",
"train_llama(num_workers=4, execution_mode=execution_mode, training_method=training_method)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Final output\n",
"\n",
"### For DDP on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=123181) {'loss': 1.8051, 'grad_norm': 0.6015625, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.6754, 'grad_norm': 0.408203125, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.568, 'grad_norm': 0.4453125, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.4934, 'grad_norm': 0.4609375, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3965, 'grad_norm': 0.3515625, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.3461, 'grad_norm': 0.34765625, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2924, 'grad_norm': 0.32421875, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2643, 'grad_norm': 0.33203125, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.263, 'grad_norm': 0.318359375, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2502, 'grad_norm': 0.275390625, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2161, 'grad_norm': 0.2734375, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=123181) {'loss': 1.2517, 'grad_norm': 0.294921875, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 13.64, 'max_memory_allocated (GB)': 48.92, 'total_memory_available (GB)': 94.62}\n",
"```\n",
"\n",
"### For DeepSpeed on HPUs\n",
"- Llama-2-70b-chat-hf\n",
"- 4 HPU\n",
"- LoRA\n",
"\n",
"``` bash\n",
"(RayTrainWorker pid=50067) {'loss': 1.662, 'grad_norm': 0.36514782905578613, 'learning_rate': 9.938441702975689e-05, 'epoch': 0.16, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.46, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.6047, 'grad_norm': 0.396455317735672, 'learning_rate': 9.567727288213005e-05, 'epoch': 0.32, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4974, 'grad_norm': 0.49250370264053345, 'learning_rate': 8.885729807284856e-05, 'epoch': 0.48, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.4078, 'grad_norm': 0.49840453267097473, 'learning_rate': 7.938926261462366e-05, 'epoch': 0.65, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.57, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.315, 'grad_norm': 0.3432576656341553, 'learning_rate': 6.7918397477265e-05, 'epoch': 0.81, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.2651, 'grad_norm': 0.32175061106681824, 'learning_rate': 5.522642316338268e-05, 'epoch': 0.97, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1947, 'grad_norm': 0.3646097481250763, 'learning_rate': 4.2178276747988446e-05, 'epoch': 1.13, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1534, 'grad_norm': 0.4598522186279297, 'learning_rate': 2.9663167846209998e-05, 'epoch': 1.29, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1404, 'grad_norm': 0.2677183449268341, 'learning_rate': 1.8533980447508137e-05, 'epoch': 1.45, 'memory_allocated (GB)': 32.75, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1283, 'grad_norm': 0.32087600231170654, 'learning_rate': 9.549150281252633e-06, 'epoch': 1.61, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.0877, 'grad_norm': 0.28305548429489136, 'learning_rate': 3.3209786751399187e-06, 'epoch': 1.77, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=50067) {'loss': 1.1238, 'grad_norm': 0.25713953375816345, 'learning_rate': 2.7390523158633554e-07, 'epoch': 1.94, 'memory_allocated (GB)': 32.86, 'max_memory_allocated (GB)': 94.59, 'total_memory_available (GB)': 94.62}\n",
"\n",
"```"
]
}
],
"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.10.12"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/llama_pretrain.ipynb
+598
View File
@@ -0,0 +1,598 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Llama model pre-training on Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-llama_pretrain\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-llama_pretrain\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will pre-train a [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b) model by using Intel Gaudi accelerators.\n",
"\n",
"We will use PyTorch for model training and Ray for distributed training.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"Basic features for this pre-training example are:\n",
"- Running on HPUs, support three execution mode: [\"lazy\", \"eager\", \"eager.compile\"](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html).\n",
"- Pre-training llama model use configuration [huggyllama/llama-7b](https://huggingface.co/huggyllama/llama-7b)\n",
"- [`GaudiTrainer`](https://github.com/huggingface/optimum-habana/blob/main/optimum/habana/transformers/trainer.py) based training.\n",
"- DeepSpeed based pre-training.\n",
"- Ray based resource scheduling and management."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare environment\n",
"This example run on single node with 4 HPUs.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Habana drivers and container runtime.\n",
"\n",
"### Get docker image\n",
"``` bash\n",
"# more available docker image can be found here: https://vault.habana.ai/ui/native/gaudi-docker\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"### Run docker image\n",
"``` bash\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"# maybe should mapping your workspace volumns\n",
"```\n",
"### Install dependency\n",
"``` bash\n",
"# \"optimum-habana>1.11.1\" if execution mode \"eager\" or \"eager.compile\" \n",
"# \"ray>=2.20.0\"\n",
"pip install ray[train] notebook transformers datasets evaluate peft accelerate scikit-learn optimum-habana\n",
"\n",
"# install deepspeed\n",
"pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.20.0\n",
"\n",
"# this notebook verfied with packages' version:\n",
"# transformers==4.45.2\n",
"# datasets==3.3.2\n",
"# evaluate==0.4.3\n",
"# peft==0.14.0\n",
"# accelerate==0.33.0\n",
"# scikit-learn==1.6.1\n",
"# optimum-habana==1.15.0\n",
"\n",
"# deepspeed==0.16.1+hpu.synapse.v1.20.0\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Import necessary libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#!/usr/bin/env python\n",
"\n",
"import os\n",
"from typing import Any, Dict\n",
"from torch.utils.data import DataLoader\n",
"\n",
"import transformers\n",
"from itertools import chain\n",
"from datasets import load_dataset\n",
"from transformers import default_data_collator\n",
"from transformers.testing_utils import CaptureLogger\n",
"from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments\n",
"from optimum.habana.utils import set_seed"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Build datasets\n",
"\n",
"Download and load dataset from huggingface.co"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_datasets(config):\n",
" dataset_name = config[\"name\"] \n",
" dataset_config_name = config[\"config_name\"]\n",
"\n",
" # Downloading and loading a dataset from the hub.\n",
" raw_datasets = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" if \"validation\" not in raw_datasets.keys():\n",
" raw_datasets[\"validation\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[:{data_args.validation_split_percentage}%]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
" raw_datasets[\"train\"] = load_dataset(\n",
" dataset_name,\n",
" dataset_config_name,\n",
" split=f\"train[{data_args.validation_split_percentage}%:]\",\n",
" cache_dir=None,\n",
" token=None,\n",
" streaming=False,\n",
" )\n",
"\n",
" return raw_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load tokenizer\n",
"\n",
"Download vocabulary from huggingface.co."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_tokenizer(config):\n",
" name = config[\"name\"]\n",
" tokenizer_kwargs = {\n",
" \"cache_dir\": None,\n",
" \"use_fast\": True,\n",
" \"revision\": \"main\",\n",
" \"token\": None,\n",
" \"trust_remote_code\": False,\n",
" }\n",
" return transformers.AutoTokenizer.from_pretrained(name, **tokenizer_kwargs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Tokenize dataset\n",
"\n",
"tokenize word to token ids."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def tokenize_dataset(datasets, tokenizer):\n",
" column_names = list(datasets[\"train\"].features)\n",
" text_column_name = \"text\" if \"text\" in column_names else column_names[0]\n",
"\n",
" tok_logger = transformers.utils.logging.get_logger(\"transformers.tokenization_utils_base\")\n",
"\n",
" def tokenize_function(examples):\n",
" with CaptureLogger(tok_logger) as cl:\n",
" output = tokenizer(examples[text_column_name])\n",
" # clm input could be much much longer than block_size\n",
" if \"Token indices sequence length is longer than the\" in cl.out:\n",
" tok_logger.warning(\n",
" \"^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits\"\n",
" \" before being passed to the model.\"\n",
" )\n",
" return output\n",
"\n",
" tokenized_datasets = datasets.map(\n",
" tokenize_function,\n",
" batched=True,\n",
" num_proc=None,\n",
" remove_columns=column_names,\n",
" load_from_cache_file=True,\n",
" desc=\"Running tokenizer on dataset\",\n",
" )\n",
"\n",
" return tokenized_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Group dataset\n",
"\n",
"This preprocssing will concatenate all texts from our dataset and generate chunks of block_size, and will pre-train model much faster."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def group_dataset(config, datasets, tokenizer):\n",
" config_name = config[\"name\"]\n",
" auto_config = transformers.AutoConfig.from_pretrained(config_name)\n",
" max_pos_embeddings = auto_config.max_position_embeddings\n",
" block_size = tokenizer.model_max_length\n",
" if block_size > max_pos_embeddings:\n",
" print(\n",
" f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n",
" f\"Using block_size={min(1024, max_pos_embeddings)} instead. You can change that default value by passing --block_size xxx.\"\n",
" )\n",
" if max_pos_embeddings > 0:\n",
" block_size = min(1024, max_pos_embeddings)\n",
" else:\n",
" block_size = 1024\n",
"\n",
" # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.\n",
" def group_texts(examples):\n",
" # Concatenate all texts.\n",
" concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n",
" total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
" # We drop the small remainder, and if the total_length < block_size we exclude this batch and return an empty dict.\n",
" # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.\n",
" total_length = (total_length // block_size) * block_size\n",
" # Split by chunks of max_len.\n",
" result = {\n",
" k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
" for k, t in concatenated_examples.items()\n",
" }\n",
" result[\"labels\"] = result[\"input_ids\"].copy()\n",
" return result\n",
"\n",
" lm_datasets = datasets.map(\n",
" group_texts,\n",
" batched=True,\n",
" num_proc=None,\n",
" load_from_cache_file=True,\n",
" desc=f\"Grouping texts in chunks of {block_size}\",\n",
" )\n",
" return lm_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load model\n",
"\n",
"Download and load pre-configed model from huggingface.co, the detail model configurations in config.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def load_model(config):\n",
" name = config[\"name\"]\n",
" model_config = config.get(\"config\", {})\n",
" auto_config = transformers.AutoConfig.from_pretrained(\n",
" pretrained_model_name_or_path=name, **model_config\n",
" )\n",
" model = transformers.AutoModelForCausalLM.from_config(auto_config, trust_remote_code=False)\n",
"\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare trainer\n",
"\n",
"Instance Trainer with `model`, `gaudi_config`, `training_args`, `tokenizer`\n",
"\n",
"No evaluation dataset passed, just training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def get_trainer(training_args, datasets, tokenizer, model):\n",
" gaudi_config = GaudiConfig.from_pretrained(\n",
" training_args.gaudi_config_name, revision=\"main\",\n",
" )\n",
"\n",
" trainer = GaudiTrainer(\n",
" model=model,\n",
" gaudi_config=gaudi_config,\n",
" args=training_args,\n",
" train_dataset=datasets[\"train\"],\n",
" eval_dataset=None,\n",
" tokenizer=tokenizer,\n",
" data_collator=default_data_collator,\n",
" )\n",
" return trainer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Function\n",
"\n",
"This function will be executed by each worker during training, with following steps:\n",
"- prepare GaudiTrainingArguments object.\n",
"- load datasets from huggingface.co.\n",
"- load pre-configed tokenizer from huggingface.co.\n",
"- tokenize dataset with loaded model tokenizer.\n",
"- concatenate all texts from our dataset and generate chunks of block_size.\n",
"- instance object of `GaudiTrainer` with training_args, datasets, tokenizer, and model.\n",
"- call `train` of trainer.\n",
"- save model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def pretrain_llama(config: Dict[str, Any]):\n",
"\n",
" training_args = GaudiTrainingArguments(**config[\"training_args\"])\n",
" set_seed(training_args.seed)\n",
"\n",
" raw_datasets = load_datasets(config[\"datasets\"])\n",
"\n",
" tokenizer = load_tokenizer(config[\"tokenizer\"])\n",
"\n",
" tokenized_datasets = tokenize_dataset(raw_datasets, tokenizer)\n",
"\n",
" tokenized_datasets = group_dataset(config[\"model\"], tokenized_datasets, tokenizer)\n",
"\n",
" model = load_model(config[\"model\"])\n",
"\n",
" trainer = get_trainer(training_args, tokenized_datasets, tokenizer, model)\n",
"\n",
" result = trainer.train()\n",
" trainer.save_model()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `main` function sets up the distributed training environment using Ray and starts the training process. To enable training using HPU, we only need to make the following changes:\n",
"- Set the exectuion mode for training, supported execution mode are:\n",
"\n",
" - \"lazy\": Deferred execution of graphs, comprising of ops delivered from script op by op similar to Eager mode. It gives the Eager mode experience with performance on Gaudi. Unlike Eager Mode with torch.compile, graph is analyzed in each iteration leading to a higher CPU usage.\n",
" - \"eager\": Op-by-op execution as defined in standard PyTorch Eager mode scripts.\n",
" - \"eager.compile\": Eager mode extended with `torch.compile` - Similar to Eager mode but extended with wrapping complete or part of model (such as a function) into a graph. Parts that are not wrapped are executed eagerly.\n",
"\n",
" More detail theory can be found [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PyTorch_Gaudi_Theory_of_Operations.html), and detail performance results can be found [here](https://developer.habana.ai/get-started/habana-models-performance/)\n",
"- Require an HPU for each worker in ScalingConfig\n",
"- Set backend to `hccl` in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def main(num_workers, execution_mode):\n",
" import ray\n",
" from ray.train import ScalingConfig\n",
" from ray.train.torch import TorchTrainer, TorchConfig\n",
"\n",
" pretrain_config = {\n",
" \"datasets\": {\n",
" \"name\": \"wikitext\",\n",
" \"config_name\": \"wikitext-2-raw-v1\",\n",
" },\n",
" \"tokenizer\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {}\n",
" },\n",
" \"model\": {\n",
" \"name\": \"huggyllama/llama-7b\",\n",
" \"config\": {\n",
" \"torch_dtype\": \"bfloat16\",\n",
" },\n",
" },\n",
" \"training_args\": {\n",
" \"per_device_train_batch_size\": 1,\n",
" \"do_train\": True,\n",
" \"save_strategy\": \"no\",\n",
" \"output_dir\": \"/tmp/ray/pretrain-llama-2\",\n",
" \"logging_steps\": 1,\n",
" \"gaudi_config_name\": \"Habana/llama\",\n",
" \"use_habana\": True,\n",
" \"throughput_warmup_steps\": 3,\n",
" \"use_lazy_mode\": True,\n",
" \"overwrite_output_dir\": True,\n",
" \"seed\": 42,\n",
" \"bf16\": True,\n",
" \"report_to\":'tensorboard',\n",
" \"deepspeed\": {\n",
" \"steps_per_print\": 64,\n",
" \"train_batch_size\": \"auto\",\n",
" \"train_micro_batch_size_per_gpu\": \"auto\",\n",
" \"gradient_accumulation_steps\": \"auto\",\n",
" \"bf16\": {\n",
" \"enabled\": True\n",
" },\n",
" \"gradient_clipping\": 1.0,\n",
" \"zero_optimization\": {\n",
" \"stage\": 3,\n",
" \"overlap_comm\": False,\n",
" \"reduce_scatter\": False,\n",
" \"contiguous_gradients\": False,\n",
" \"stage3_gather_16bit_weights_on_model_save\": True\n",
" }\n",
" },\n",
" },\n",
" }\n",
"\n",
" # if execution mode is eager with compile, must spcified with a compile backend\n",
" if execution_mode == \"eager.compile\":\n",
" pretrain_config[\"training_args\"].update({\"torch_compile_backend\": \"hpu_backend\"})\n",
"\n",
" scaling_config = ScalingConfig(num_workers=num_workers,\n",
" use_gpu=False,\n",
" resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
"\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend=\"hccl\")\n",
"\n",
" ray.init()\n",
"\n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=pretrain_llama,\n",
" train_loop_config=pretrain_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `main` function to start the pre-training process.\n",
"\n",
"Before calling `main` function, you must set some environment variables.\n",
"\n",
"1. The visiable devices. Environment variable `HABANA_VISIBLE_DEVICES` and `HABANA_VISIBLE_MODULES` are used to control the HPU device visiable to application, you must set this two environment variable properly. For more detail usage of `HABANA_VISIBLE_DEVICES`, `HABANA_VISIBLE_MODULES`, please visit [here](https://docs.habana.ai/en/latest/PyTorch/Reference/PT_Multiple_Tenants_on_HPU/Multiple_Dockers_each_with_Single_Workload.html)\n",
"\n",
"2. The execution mode. Different execution mode has different runtime performance. The default execution mode is lazy mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set some environment variables\n",
"os.environ[\"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\"] = \"0\"\n",
"# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var\n",
"# you must set HABANA_VISIBLE_MODULES, such as\n",
"# os.environ[\"HABANA_VISIBLE_MODULES\"] = \"0,1,2,3\"\n",
"\n",
"# execution_mode are [\"lazy\", \"eager\", \"eager.compile\"]\n",
"execution_mode = \"lazy\"\n",
"os.environ[\"PT_HPU_LAZY_MODE\"] = \"1\" if execution_mode == \"lazy\" else \"0\"\n",
"\n",
"main(num_workers=4, execution_mode=execution_mode)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"\n",
"...\n",
"\n",
"RayTrainWorker pid=36561) Setting up process group for: env:// [rank=0, world_size=4]\n",
"(TorchTrainer pid=36054) Started distributed worker processes: \n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36561) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36562) world_rank=1, local_rank=1, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36563) world_rank=2, local_rank=2, node_rank=0\n",
"(TorchTrainer pid=36054) - (node_id=409da2dba1dc3e5b8e58a2b766a4a19d90e7879c28c2fb13644148b8, ip=100.83.111.228, pid=36564) world_rank=3, local_rank=3, node_rank=0\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=36561) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=36561) PT_HPU_LAZY_ACC_PAR_MODE = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=36561) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=36561) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=36561) Num CPU Cores : 160\n",
"(RayTrainWorker pid=36561) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=36561) ------------------------------------------------------------------------------\n",
"\n",
"...\n",
"\n",
"(RayTrainWorker pid=36561) {'loss': 4.1052, 'grad_norm': 2.225008249282837, 'learning_rate': 8.26086956521739e-06, 'epoch': 2.5, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0472, 'grad_norm': 2.0701019763946533, 'learning_rate': 8.212560386473431e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.097, 'grad_norm': 2.119075059890747, 'learning_rate': 8.164251207729469e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7035, 'grad_norm': 2.1802899837493896, 'learning_rate': 8.115942028985508e-06, 'epoch': 2.51, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.242, 'grad_norm': 1.9516953229904175, 'learning_rate': 8.067632850241547e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9594, 'grad_norm': 2.0580222606658936, 'learning_rate': 8.019323671497584e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.3415, 'grad_norm': 2.192605495452881, 'learning_rate': 7.971014492753623e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9739, 'grad_norm': 2.0198025703430176, 'learning_rate': 7.922705314009662e-06, 'epoch': 2.52, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1624, 'grad_norm': 2.0957565307617188, 'learning_rate': 7.874396135265701e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9744, 'grad_norm': 2.1159448623657227, 'learning_rate': 7.82608695652174e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1127, 'grad_norm': 2.159834623336792, 'learning_rate': 7.777777777777777e-06, 'epoch': 2.53, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0588, 'grad_norm': 2.106534004211426, 'learning_rate': 7.729468599033817e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8734, 'grad_norm': 2.445814371109009, 'learning_rate': 7.681159420289856e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0278, 'grad_norm': 2.0376927852630615, 'learning_rate': 7.632850241545895e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9643, 'grad_norm': 2.1097891330718994, 'learning_rate': 7.584541062801932e-06, 'epoch': 2.54, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1384, 'grad_norm': 2.157325267791748, 'learning_rate': 7.536231884057972e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9982, 'grad_norm': 2.230065107345581, 'learning_rate': 7.48792270531401e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0988, 'grad_norm': 2.355875015258789, 'learning_rate': 7.439613526570048e-06, 'epoch': 2.55, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0514, 'grad_norm': 2.1178295612335205, 'learning_rate': 7.391304347826088e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.9858, 'grad_norm': 2.089723825454712, 'learning_rate': 7.342995169082126e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1548, 'grad_norm': 2.308490753173828, 'learning_rate': 7.294685990338164e-06, 'epoch': 2.56, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0356, 'grad_norm': 1.9994627237319946, 'learning_rate': 7.246376811594203e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7696, 'grad_norm': 1.9719663858413696, 'learning_rate': 7.1980676328502416e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0157, 'grad_norm': 2.1598856449127197, 'learning_rate': 7.1497584541062814e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0113, 'grad_norm': 1.997869849205017, 'learning_rate': 7.10144927536232e-06, 'epoch': 2.57, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1048, 'grad_norm': 2.099222183227539, 'learning_rate': 7.053140096618358e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0048, 'grad_norm': 2.100231885910034, 'learning_rate': 7.004830917874397e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.0302, 'grad_norm': 2.18204402923584, 'learning_rate': 6.956521739130435e-06, 'epoch': 2.58, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.7227, 'grad_norm': 2.190962553024292, 'learning_rate': 6.908212560386473e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.1111, 'grad_norm': 2.349518060684204, 'learning_rate': 6.859903381642513e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.024, 'grad_norm': 2.5497331619262695, 'learning_rate': 6.811594202898551e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 3.8844, 'grad_norm': 2.3125178813934326, 'learning_rate': 6.7632850241545894e-06, 'epoch': 2.59, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"(RayTrainWorker pid=36561) {'loss': 4.2208, 'grad_norm': 2.1103923320770264, 'learning_rate': 6.7149758454106285e-06, 'epoch': 2.6, 'memory_allocated (GB)': 28.87, 'max_memory_allocated (GB)': 94.26, 'total_memory_available (GB)': 94.62}\n",
"\n",
"...\n",
"```"
]
}
],
"metadata": {
"language_info": {
"name": "python"
},
"orphan": true
},
"nbformat": 4,
"nbformat_minor": 2
}
doc/source/train/examples/intel_gaudi/resnet.ipynb
+610
View File
@@ -0,0 +1,610 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"(hpu_resnet_training)=\n",
"# ResNet Model Training with Intel Gaudi\n",
"\n",
"<a id=\"try-anyscale-quickstart-intel_gaudi-resnet\" href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=intel_gaudi-resnet\">\n",
" <img src=\"../../../_static/img/run-on-anyscale.svg\" alt=\"try-anyscale-quickstart\">\n",
"</a>\n",
"<br></br>\n",
"\n",
"In this Jupyter notebook, we will train a ResNet-50 model to classify images of ants and bees using HPU. We will use PyTorch for model training and Ray for distributed training. The dataset will be downloaded and processed using torchvision's datasets and transforms.\n",
"\n",
"[Intel Gaudi AI Processors (HPUs)](https://habana.ai) are AI hardware accelerators designed by Intel Habana Labs. For more information, see [Gaudi Architecture](https://docs.habana.ai/en/latest/Gaudi_Overview/index.html) and [Gaudi Developer Docs](https://developer.habana.ai/).\n",
"\n",
"## Configuration\n",
"\n",
"A node with Gaudi/Gaudi2 installed is required to run this example. Both Gaudi and Gaudi2 have 8 HPUs. We will use 2 workers to train the model, each using 1 HPU.\n",
"\n",
"We recommend using a prebuilt container to run these examples. To run a container, you need Docker. See [Install Docker Engine](https://docs.docker.com/engine/install/) for installation instructions.\n",
"\n",
"Next, follow [Run Using Containers](https://docs.habana.ai/en/latest/Installation_Guide/Bare_Metal_Fresh_OS.html?highlight=installer#run-using-containers) to install the Gaudi drivers and container runtime.\n",
"\n",
"Next, start the Gaudi container:\n",
"```bash\n",
"docker pull vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.20.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest\n",
"```\n",
"\n",
"Inside the container, install Ray and Jupyter to run this notebook.\n",
"```bash\n",
"pip install ray[train] notebook\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from typing import Dict\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"import torch\n",
"from filelock import FileLock\n",
"from torch import nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets, transforms, models\n",
"from tqdm import tqdm\n",
"\n",
"import ray\n",
"import ray.train as train\n",
"from ray.train import ScalingConfig, Checkpoint\n",
"from ray.train.torch import TorchTrainer\n",
"from ray.train.torch import TorchConfig\n",
"from ray.runtime_env import RuntimeEnv\n",
"\n",
"import habana_frameworks.torch.core as htcore"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Define Data Transforms\n",
"\n",
"We will set up the data transforms for preprocessing images for training and validation. This includes random cropping, flipping, and normalization for the training set, and resizing and normalization for the validation set."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# Data augmentation and normalization for training\n",
"# Just normalization for validation\n",
"data_transforms = {\n",
" \"train\": transforms.Compose([\n",
" transforms.RandomResizedCrop(224),\n",
" transforms.RandomHorizontalFlip(),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
" \"val\": transforms.Compose([\n",
" transforms.Resize(256),\n",
" transforms.CenterCrop(224),\n",
" transforms.ToTensor(),\n",
" transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
" ]),\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Download Function\n",
"\n",
"We will define a function to download the Hymenoptera dataset. This dataset contains images of ants and bees for a binary classification problem."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def download_datasets():\n",
" os.system(\"wget https://download.pytorch.org/tutorial/hymenoptera_data.zip >/dev/null 2>&1\")\n",
" os.system(\"unzip hymenoptera_data.zip >/dev/null 2>&1\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dataset Preparation Function\n",
"\n",
"After downloading the dataset, we need to build PyTorch datasets for training and validation. The `build_datasets` function will apply the previously defined transforms and create the datasets."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def build_datasets():\n",
" torch_datasets = {}\n",
" for split in [\"train\", \"val\"]:\n",
" torch_datasets[split] = datasets.ImageFolder(\n",
" os.path.join(\"./hymenoptera_data\", split), data_transforms[split]\n",
" )\n",
" return torch_datasets"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model Initialization Functions\n",
"\n",
"We will define two functions to initialize our model. The `initialize_model` function will load a pre-trained ResNet-50 model and replace the final classification layer for our binary classification task. The `initialize_model_from_checkpoint` function will load a model from a saved checkpoint if available."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def initialize_model():\n",
" # Load pretrained model params\n",
" model = models.resnet50(pretrained=True)\n",
"\n",
" # Replace the original classifier with a new Linear layer\n",
" num_features = model.fc.in_features\n",
" model.fc = nn.Linear(num_features, 2)\n",
"\n",
" # Ensure all params get updated during finetuning\n",
" for param in model.parameters():\n",
" param.requires_grad = True\n",
" return model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluation Function\n",
"\n",
"To assess the performance of our model during training, we define an `evaluate` function. This function computes the number of correct predictions by comparing the predicted labels with the true labels."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def evaluate(logits, labels):\n",
" _, preds = torch.max(logits, 1)\n",
" corrects = torch.sum(preds == labels).item()\n",
" return corrects"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training Loop Function\n",
"\n",
"This function defines the training loop that will be executed by each worker. It includes downloading the dataset, preparing data loaders, initializing the model, and running the training and validation phases. Compared to a training function for GPU, no changes are needed to port to HPU. Internally, Ray Train does these things:\n",
"\n",
"* Detect HPU and set the device.\n",
"\n",
"* Initializes the habana PyTorch backend.\n",
"\n",
"* Initializes the habana distributed backend."
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"def train_loop_per_worker(configs):\n",
" import warnings\n",
"\n",
" warnings.filterwarnings(\"ignore\")\n",
"\n",
" # Calculate the batch size for a single worker\n",
" worker_batch_size = configs[\"batch_size\"] // train.get_context().get_world_size()\n",
"\n",
" # Download dataset once on local rank 0 worker\n",
" if train.get_context().get_local_rank() == 0:\n",
" download_datasets()\n",
" torch.distributed.barrier()\n",
"\n",
" # Build datasets on each worker\n",
" torch_datasets = build_datasets()\n",
"\n",
" # Prepare dataloader for each worker\n",
" dataloaders = dict()\n",
" dataloaders[\"train\"] = DataLoader(\n",
" torch_datasets[\"train\"], batch_size=worker_batch_size, shuffle=True\n",
" )\n",
" dataloaders[\"val\"] = DataLoader(\n",
" torch_datasets[\"val\"], batch_size=worker_batch_size, shuffle=False\n",
" )\n",
"\n",
" # Distribute\n",
" dataloaders[\"train\"] = train.torch.prepare_data_loader(dataloaders[\"train\"])\n",
" dataloaders[\"val\"] = train.torch.prepare_data_loader(dataloaders[\"val\"])\n",
"\n",
" # Obtain HPU device automatically\n",
" device = train.torch.get_device()\n",
"\n",
" # Prepare DDP Model, optimizer, and loss function\n",
" model = initialize_model()\n",
" model = model.to(device)\n",
"\n",
" optimizer = optim.SGD(\n",
" model.parameters(), lr=configs[\"lr\"], momentum=configs[\"momentum\"]\n",
" )\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" # Start training loops\n",
" for epoch in range(configs[\"num_epochs\"]):\n",
" # Each epoch has a training and validation phase\n",
" for phase in [\"train\", \"val\"]:\n",
" if phase == \"train\":\n",
" model.train() # Set model to training mode\n",
" else:\n",
" model.eval() # Set model to evaluate mode\n",
"\n",
" running_loss = 0.0\n",
" running_corrects = 0\n",
"\n",
" for inputs, labels in dataloaders[phase]:\n",
" inputs = inputs.to(device)\n",
" labels = labels.to(device)\n",
"\n",
" # zero the parameter gradients\n",
" optimizer.zero_grad()\n",
"\n",
" # forward\n",
" with torch.set_grad_enabled(phase == \"train\"):\n",
" # Get model outputs and calculate loss\n",
" outputs = model(inputs)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # backward + optimize only if in training phase\n",
" if phase == \"train\":\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # calculate statistics\n",
" running_loss += loss.item() * inputs.size(0)\n",
" running_corrects += evaluate(outputs, labels)\n",
"\n",
" size = len(torch_datasets[phase]) // train.get_context().get_world_size()\n",
" epoch_loss = running_loss / size\n",
" epoch_acc = running_corrects / size\n",
"\n",
" if train.get_context().get_world_rank() == 0:\n",
" print(\n",
" \"Epoch {}-{} Loss: {:.4f} Acc: {:.4f}\".format(\n",
" epoch, phase, epoch_loss, epoch_acc\n",
" )\n",
" )\n",
"\n",
" # Report metrics and checkpoint every epoch\n",
" if phase == \"val\":\n",
" train.report(\n",
" metrics={\"loss\": epoch_loss, \"acc\": epoch_acc},\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Main Training Function\n",
"\n",
"The `train_resnet` function sets up the distributed training environment using Ray and starts the training process. It specifies the batch size, number of epochs, learning rate, and momentum for the SGD optimizer. To enable training using HPU, we only need to make the following changes:\n",
"* Require an HPU for each worker in ScalingConfig\n",
"* Set backend to \"hccl\" in TorchConfig"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def train_resnet(num_workers=2):\n",
" global_batch_size = 16\n",
"\n",
" train_loop_config = {\n",
" \"input_size\": 224, # Input image size (224 x 224)\n",
" \"batch_size\": 32, # Batch size for training\n",
" \"num_epochs\": 10, # Number of epochs to train for\n",
" \"lr\": 0.001, # Learning Rate\n",
" \"momentum\": 0.9, # SGD optimizer momentum\n",
" }\n",
" # Configure computation resources\n",
" # In ScalingConfig, require an HPU for each worker\n",
" scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={\"CPU\": 1, \"HPU\": 1})\n",
" # Set backend to hccl in TorchConfig\n",
" torch_config = TorchConfig(backend = \"hccl\")\n",
" \n",
" ray.init()\n",
" \n",
" # Initialize a Ray TorchTrainer\n",
" trainer = TorchTrainer(\n",
" train_loop_per_worker=train_loop_per_worker,\n",
" train_loop_config=train_loop_config,\n",
" torch_config=torch_config,\n",
" scaling_config=scaling_config,\n",
" )\n",
"\n",
" result = trainer.fit()\n",
" print(f\"Training result: {result}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Start Training\n",
"\n",
"Finally, we call the `train_resnet` function to start the training process. You can adjust the number of workers to use. Before running this cell, ensure that Ray is properly set up in your environment to handle distributed training.\n",
"\n",
"Note: the following warning is fine, and is resolved in SynapseAI version 1.14.0+:\n",
"```text\n",
"/usr/local/lib/python3.10/dist-packages/torch/distributed/distributed_c10d.py:252: UserWarning: Device capability of hccl unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"train_resnet(num_workers=2) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Possible outputs\n",
"\n",
"``` text\n",
"2025-03-03 03:32:12,620 INFO worker.py:1841 -- Started a local Ray instance.\n",
"/usr/local/lib/python3.10/dist-packages/ray/tune/impl/tuner_internal.py:125: RayDeprecationWarning: The `RunConfig` class should be imported from `ray.tune` when passing it to the Tuner. Please update your imports. See this issue for more context and migration options: https://github.com/ray-project/ray/issues/49454. Disable these warnings by setting the environment variable: RAY_TRAIN_ENABLE_V2_MIGRATION_WARNINGS=0\n",
" _log_deprecation_warning(\n",
"(RayTrainWorker pid=63669) Setting up process group for: env:// [rank=0, world_size=2]\n",
"(TorchTrainer pid=63280) Started distributed worker processes: \n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63669) world_rank=0, local_rank=0, node_rank=0\n",
"(TorchTrainer pid=63280) - (node_id=9f2c34ea47fe405f3227e9168aa857f81655a83e95fd6be359fd76db, ip=100.83.111.228, pid=63668) world_rank=1, local_rank=1, node_rank=0\n",
"(RayTrainWorker pid=63669) ============================= HABANA PT BRIDGE CONFIGURATION =========================== \n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_RECIPE_CACHE_CONFIG = ,false,1024\n",
"(RayTrainWorker pid=63669) PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807\n",
"(RayTrainWorker pid=63669) PT_HPU_LAZY_ACC_PAR_MODE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_EAGER_COLLECTIVE_PIPELINE_ENABLE = 1\n",
"(RayTrainWorker pid=63669) PT_HPU_ENABLE_LAZY_COLLECTIVES = 0\n",
"(RayTrainWorker pid=63669) ---------------------------: System Configuration :---------------------------\n",
"(RayTrainWorker pid=63669) Num CPU Cores : 160\n",
"(RayTrainWorker pid=63669) CPU RAM : 1056374420 KB\n",
"(RayTrainWorker pid=63669) ------------------------------------------------------------------------------\n",
"(RayTrainWorker pid=63668) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 9%|▊ | 8.38M/97.8M [00:00<00:01, 87.7MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 193MB/s]\n",
"100%|██████████| 97.8M/97.8M [00:00<00:00, 203MB/s]\n",
"\n",
"View detailed results here: /root/ray_results/TorchTrainer_2025-03-03_03-32-15\n",
"To visualize your results with TensorBoard, run: `tensorboard --logdir /tmp/ray/session_2025-03-03_03-32-10_695011_53838/artifacts/2025-03-03_03-32-15/TorchTrainer_2025-03-03_03-32-15/driver_artifacts`\n",
"\n",
"Training started with configuration:\n",
"╭──────────────────────────────────────╮\n",
"│ Training config │\n",
"├──────────────────────────────────────┤\n",
"│ train_loop_config/batch_size 32 │\n",
"│ train_loop_config/input_size 224 │\n",
"│ train_loop_config/lr 0.001 │\n",
"│ train_loop_config/momentum 0.9 │\n",
"│ train_loop_config/num_epochs 10 │\n",
"╰──────────────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-train Loss: 0.6574 Acc: 0.6066\n",
"\n",
"Training finished iteration 1 at 2025-03-03 03:32:45. Total running time: 29s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 24.684 │\n",
"│ time_total_s 24.684 │\n",
"│ training_iteration 1 │\n",
"│ acc 0.71053 │\n",
"│ loss 0.51455 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 0-val Loss: 0.5146 Acc: 0.7105\n",
"(RayTrainWorker pid=63669) Epoch 1-train Loss: 0.5016 Acc: 0.7541\n",
"\n",
"Training finished iteration 2 at 2025-03-03 03:32:46. Total running time: 31s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39649 │\n",
"│ time_total_s 26.0805 │\n",
"│ training_iteration 2 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.30218 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 1-val Loss: 0.3022 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 2-train Loss: 0.3130 Acc: 0.9180\n",
"\n",
"Training finished iteration 3 at 2025-03-03 03:32:47. Total running time: 32s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37042 │\n",
"│ time_total_s 27.4509 │\n",
"│ training_iteration 3 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.22201 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 2-val Loss: 0.2220 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 3-train Loss: 0.2416 Acc: 0.9262\n",
"\n",
"Training finished iteration 4 at 2025-03-03 03:32:49. Total running time: 34s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.38353 │\n",
"│ time_total_s 28.8345 │\n",
"│ training_iteration 4 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.17815 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 3-val Loss: 0.1782 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 4-train Loss: 0.1900 Acc: 0.9508\n",
"\n",
"Training finished iteration 5 at 2025-03-03 03:32:50. Total running time: 35s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37318 │\n",
"│ time_total_s 30.2077 │\n",
"│ training_iteration 5 │\n",
"│ acc 0.93421 │\n",
"│ loss 0.17063 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 4-val Loss: 0.1706 Acc: 0.9342\n",
"(RayTrainWorker pid=63669) Epoch 5-train Loss: 0.1346 Acc: 0.9672\n",
"\n",
"Training finished iteration 6 at 2025-03-03 03:32:52. Total running time: 36s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37999 │\n",
"│ time_total_s 31.5876 │\n",
"│ training_iteration 6 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.1552 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 5-val Loss: 0.1552 Acc: 0.9605\n",
"(RayTrainWorker pid=63669) Epoch 6-train Loss: 0.1184 Acc: 0.9672\n",
"\n",
"Training finished iteration 7 at 2025-03-03 03:32:53. Total running time: 38s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.39198 │\n",
"│ time_total_s 32.9796 │\n",
"│ training_iteration 7 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14702 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 6-val Loss: 0.1470 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 7-train Loss: 0.0864 Acc: 0.9836\n",
"\n",
"Training finished iteration 8 at 2025-03-03 03:32:54. Total running time: 39s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.3736 │\n",
"│ time_total_s 34.3532 │\n",
"│ training_iteration 8 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14443 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 7-val Loss: 0.1444 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 8-train Loss: 0.1085 Acc: 0.9590\n",
"\n",
"Training finished iteration 9 at 2025-03-03 03:32:56. Total running time: 40s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.37868 │\n",
"│ time_total_s 35.7319 │\n",
"│ training_iteration 9 │\n",
"│ acc 0.94737 │\n",
"│ loss 0.14194 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 8-val Loss: 0.1419 Acc: 0.9474\n",
"(RayTrainWorker pid=63669) Epoch 9-train Loss: 0.0829 Acc: 0.9754\n",
"\n",
"2025-03-03 03:32:58,628 INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/root/ray_results/TorchTrainer_2025-03-03_03-32-15' in 0.0028s.\n",
"Training finished iteration 10 at 2025-03-03 03:32:57. Total running time: 42s\n",
"╭───────────────────────────────╮\n",
"│ Training result │\n",
"├───────────────────────────────┤\n",
"│ checkpoint_dir_name │\n",
"│ time_this_iter_s 1.36497 │\n",
"│ time_total_s 37.0969 │\n",
"│ training_iteration 10 │\n",
"│ acc 0.96053 │\n",
"│ loss 0.14297 │\n",
"╰───────────────────────────────╯\n",
"(RayTrainWorker pid=63669) Epoch 9-val Loss: 0.1430 Acc: 0.9605\n",
"\n",
"Training completed after 10 iterations at 2025-03-03 03:32:58. Total running time: 43s\n",
"\n",
"Training result: Result(\n",
" metrics={'loss': 0.1429688463869848, 'acc': 0.9605263157894737},\n",
" path='/root/ray_results/TorchTrainer_2025-03-03_03-32-15/TorchTrainer_19fd8_00000_0_2025-03-03_03-32-15',\n",
" filesystem='local',\n",
" checkpoint=None\n",
")\n",
"(RayTrainWorker pid=63669) Downloading: \"https://download.pytorch.org/models/resnet50-0676ba61.pth\" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth\n",
" 0%| | 0.00/97.8M [00:00<?, ?B/s]\n",
" 68%|██████▊ | 66.1M/97.8M [00:00<00:00, 160MB/s] [repeated 6x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)\n",
"```"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.10.12 64-bit",
"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"
},
"orphan": true,
"vscode": {
"interpreter": {
"hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/train/examples/intel_gaudi/sd.ipynb
+494
View File
File diff suppressed because one or more lines are too long