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/templates/01_batch_inference/README.md
+14
View File
@@ -0,0 +1,14 @@
# Scaling Batch Inference with Ray Data
| Template Specification | Description |
| ---------------------- | ----------- |
| Summary | This template walks through GPU batch inference on an image dataset. |
| Time to Run | Less than 5 minutes to compute predictions on the dataset. |
| Minimum Compute Requirements | No hard requirements. The default is 4 nodes, each with 1 NVIDIA T4 GPU. |
| Cluster Environment | This template uses the latest Anyscale-provided Ray ML image using Python 3.9: [`anyscale/ray-ml:latest-py39-gpu`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=01_batch_inference). If you want to change to a different cluster environment, make sure that it's based on this image.|
## Getting Started
**When the workspace is up and running, start coding by clicking on the Jupyter or VS Code icon above. Open the `start.ipynb` file and follow the instructions there.**
By the end, we will have classified around 10k images with a PyTorch model.
doc/source/templates/01_batch_inference/start.ipynb
+744
View File
File diff suppressed because one or more lines are too long
doc/source/templates/02_many_model_training/README.md
+16
View File
@@ -0,0 +1,16 @@
# Scaling Many Model Training with Ray Tune
| Template Specification | Description |
| ---------------------- | ----------- |
| Summary | This template demonstrates how to parallelize the training of hundreds of time-series forecasting models with [Ray Tune](https://docs.ray.io/en/latest/tune/index.html). The template uses the `statsforecast` library to fit models to partitions of the M4 forecasting competition dataset. |
| Time to Run | Around 5 minutes to train all models. |
| Minimum Compute Requirements | No hard requirements. The default is 8 nodes with 8 CPUs each. |
| Cluster Environment | This template uses the latest Anyscale-provided Ray ML image using Python 3.9, [`anyscale/ray-ml:latest-py39-gpu`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=many_model_training_readme), with some extra requirements from `requirements.txt` installed on top. If you want to change to a different cluster environment, make sure that it's based on this image and includes all packages listed in the `requirements.txt` file. |
## Getting Started
**When the workspace is up and running, start coding by clicking on the Jupyter or VS Code icon above. Open the `start.ipynb` file and follow the instructions there.**
The end result of the template is fitting multiple models on each dataset partition, then determining the best model based on cross-validation metrics. Then, using the best model, we can generate forecasts like the ones shown below:
![Forecasts](https://github-production-user-asset-6210df.s3.amazonaws.com/3887863/239091118-2413f399-4636-40cf-8b12-8d3ce15f5ce1.png)
doc/source/templates/02_many_model_training/requirements.txt
+1
View File
@@ -0,0 +1 @@
statsforecast==1.5.0
doc/source/templates/02_many_model_training/start.ipynb
+396
View File
@@ -0,0 +1,396 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "98e0d4f3",
"metadata": {},
"source": [
"# Scaling Many Model Training with Ray Tune\n",
"\n",
"| Template Specification | Description |\n",
"| ---------------------- | ----------- |\n",
"| Summary | This template demonstrates how to parallelize the training of hundreds of time-series forecasting models with [Ray Tune](https://docs.ray.io/en/latest/tune/index.html). The template uses the `statsforecast` library to fit models to partitions of the M4 forecasting competition dataset. |\n",
"| Time to Run | Around 5 minutes to train all models. |\n",
"| Minimum Compute Requirements | No hard requirements. The default is 8 nodes with 8 CPUs each. |\n",
"| Cluster Environment | This template uses the latest Anyscale-provided Ray ML image using Python 3.9: [`anyscale/ray-ml:latest-py39-gpu`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=many_model_training_start_ipynb), with some extra requirements from `requirements.txt` installed on top. If you want to change to a different cluster environment, make sure that it's based on this image and includes all packages listed in the `requirements.txt` file. |\n",
"\n",
"The end result of the template is fitting multiple models on each dataset partition, then determining the best model based on cross-validation metrics. Then, using the best model, you can generate forecasts like the ones shown below:\n",
"\n",
"![Forecasts](https://github-production-user-asset-6210df.s3.amazonaws.com/3887863/239091118-2413f399-4636-40cf-8b12-8d3ce15f5ce1.png)\n",
"\n",
"\n",
"In many model training, the focus is on training models on multiple subsets of\n",
"a dataset, rather than training a single model on the entire dataset. Each model is trained on an independent\n",
"dataset partition, allowing Ray to parallelize the workload by running multiple\n",
"training jobs concurrently, instead of sequentially training each model.\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "08e65f8d",
"metadata": {},
"source": [
"> Slot in your code below wherever you see the ✂️ icon to build off of this template!\n",
">\n",
"> The framework and data format used in this template can be easily replaced to suit your own application!"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "52aa4f70",
"metadata": {},
"source": [
"## Set up the dependencies\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "488cd257",
"metadata": {},
"source": [
"When running in a distributed Ray Cluster, all nodes need to have access to dependencies.\n",
"For this, we'll use `pip install --user` to install the necessary requirements. On an Anyscale Workspace, this is configured to install packages to a shared filesystem that will be available to all nodes in the cluster.\n",
"\n",
"```\n",
"pip install --user -r requirements.txt\n",
"```\n",
"\n",
"After installing all the requirements, we'll start with some imports."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b5ac5876",
"metadata": {},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"import pandas as pd\n",
"from statsforecast import StatsForecast\n",
"from statsforecast.models import AutoARIMA, AutoETS, MSTL\n",
"\n",
"from ray import train, tune\n",
"from ray.train import RunConfig\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "44a9dc54",
"metadata": {},
"source": [
"## Define the custom training function\n",
"\n",
"Next, we define the custom training function that fits the forecasting models and\n",
"computes evaluation metrics.\n",
"Ray Tune will distribute this code across the cluster and schedule for as many training\n",
"jobs as possible to execute in parallel, considering the available cluster resources."
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "060ee3ce",
"metadata": {},
"source": [
"> ✂️ Replace this with your own training logic to run per dataset partition.\n",
">\n",
"> The only additional Ray Tune code that is added is the `train.report`\n",
"> at the end of the training function. This reports metrics for Ray Tune to log,\n",
"> which can be analyzed after the run finishes."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "faaa0dad",
"metadata": {},
"outputs": [],
"source": [
"n_cv_windows = 1\n",
"\n",
"# Try two different types of forecasting models per dataset partition.\n",
"# The dataset contains hourly records, so the `season_length` is 24 hours.\n",
"models = [\n",
" AutoETS(season_length=24),\n",
" MSTL(season_length=24, trend_forecaster=AutoARIMA()),\n",
"]\n",
"\n",
"# See the appendix for info on setting resource requirements for each trial.\n",
"cpus_per_trial = len(models) * n_cv_windows\n",
"\n",
"\n",
"def train_fn(config: dict):\n",
" # First, define some helper functions for fetching data and computing eval metrics.\n",
"\n",
" def get_m5_partition(unique_id: str) -> pd.DataFrame:\n",
" df = pd.read_parquet(\n",
" \"https://datasets-nixtla.s3.amazonaws.com/m4-hourly.parquet\"\n",
" )\n",
" df = df[df[\"unique_id\"] == unique_id]\n",
" return df.dropna()\n",
"\n",
" def evaluate_cross_validation(df: pd.DataFrame) -> pd.DataFrame:\n",
" from sklearn.metrics import mean_squared_error\n",
"\n",
" models = df.drop(columns=[\"ds\", \"cutoff\", \"y\"]).columns.tolist()\n",
" evals = []\n",
" for model in models:\n",
" eval_ = (\n",
" df.groupby([\"unique_id\", \"cutoff\"])\n",
" # Calculate the Root Mean Squared Error (RMSE)\n",
" .apply(\n",
" lambda x: mean_squared_error(\n",
" x[\"y\"].values, x[model].values, squared=False\n",
" )\n",
" ).to_frame()\n",
" )\n",
" eval_.columns = [model]\n",
" evals.append(eval_)\n",
" evals = pd.concat(evals, axis=1)\n",
" evals = evals.groupby([\"unique_id\"]).mean(numeric_only=True)\n",
" evals[\"best_model\"] = evals.idxmin(axis=1)\n",
" return evals\n",
"\n",
" # Later, we will set up Ray Tune to populate `config['data_partition_id']`.\n",
" # Use this value to determine which partition of the dataset to use.\n",
" data_partition_id = config[\"data_partition_id\"]\n",
" train_df = get_m5_partition(data_partition_id)\n",
"\n",
" forecast_horizon = 24 # Forecast the next 24 hours\n",
"\n",
" sf = StatsForecast(\n",
" df=train_df,\n",
" models=models,\n",
" freq=\"H\",\n",
" # Set the number of cores used by statsforecast to the\n",
" # number of CPUs assigned to the trial!\n",
" n_jobs=cpus_per_trial,\n",
" )\n",
" cv_df = sf.cross_validation(\n",
" h=forecast_horizon,\n",
" step_size=forecast_horizon,\n",
" n_windows=n_cv_windows,\n",
" )\n",
"\n",
" eval_df = evaluate_cross_validation(df=cv_df)\n",
" best_model = eval_df[\"best_model\"][data_partition_id]\n",
" forecast_mse = eval_df[best_model][data_partition_id]\n",
"\n",
" if data_partition_id == \"H1\":\n",
" # For the first data partition, plot forecasts of the best model.\n",
" forecast_df = sf.forecast(h=forecast_horizon)\n",
" fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
" plot_df = pd.concat([train_df, forecast_df]).set_index(\"ds\")\n",
" plot_df[[\"y\", best_model]].plot(ax=ax)\n",
" ax.set_title(f\"Forecast for data partition: {data_partition_id}\")\n",
" ax.set_xlabel(f\"Timestamp [ds]\")\n",
" ax.set_ylabel(f\"Target [y]\")\n",
" ax.get_figure().savefig(\"prediction.png\")\n",
"\n",
" # Report the best-performing model and its corresponding eval metric.\n",
" train.report({\"forecast_mse\": forecast_mse, \"best_model\": best_model})\n",
"\n",
"\n",
"trainable = tune.with_resources(train_fn, resources={\"CPU\": cpus_per_trial})\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "421eb6f6",
"metadata": {},
"source": [
"## Define the data partitions to train on"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "89741e7a",
"metadata": {},
"source": [
"In this template, we consider the dataset partition ID as a hyperparameter, and we leverage Ray Tune to parallelize the execution of our training function across each dataset partition.\n",
"\n",
"> ✂️ Modify the hyperparameter search space `param_space` to enable your training function to configure the dataset! This is how `config['data_partition_id']` from earlier gets populated."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1e9f2825",
"metadata": {},
"outputs": [],
"source": [
"# First, pull the list of unique IDs used to partition the dataset.\n",
"data_partition_ids = list(\n",
" pd.read_parquet(\n",
" \"https://datasets-nixtla.s3.amazonaws.com/m4-hourly.parquet\",\n",
" columns=[\"unique_id\"],\n",
" )[\"unique_id\"].unique()\n",
")\n",
"print(f\"Training on a total of {len(data_partition_ids)} dataset partitions.\")\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "21bccbcc",
"metadata": {},
"outputs": [],
"source": [
"param_space = {\n",
" \"data_partition_id\": tune.grid_search(data_partition_ids),\n",
"}\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "13b4dd3e",
"metadata": {},
"source": [
"Run many model training using Ray Tune!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1ef8245",
"metadata": {},
"outputs": [],
"source": [
"tuner = tune.Tuner(\n",
" trainable,\n",
" param_space=param_space,\n",
" # Experiment results are saved to a shared filesystem available to all nodes.\n",
" run_config=RunConfig(storage_path=\"/mnt/cluster_storage\"),\n",
")\n",
"result_grid = tuner.fit()\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "ba1a07d0",
"metadata": {},
"source": [
"View the reported results of all trials as a dataframe."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d7baa29a",
"metadata": {},
"outputs": [],
"source": [
"results_df = result_grid.get_dataframe()\n",
"results_df\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "5ed8df5c",
"metadata": {},
"source": [
"## View one of the model forecasts\n",
"\n",
"We saved an image of the forecast generated by the best model trained on the first dataset partition `'H1'`.\n",
"Let's find that file and display it!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3909636a",
"metadata": {},
"outputs": [],
"source": [
"from IPython.display import Image, display\n",
"import os\n",
"\n",
"for result in result_grid:\n",
" # Find the result associated with the run that saved a forecast plot.\n",
" if result.config[\"data_partition_id\"] == \"H1\":\n",
" display(Image(os.path.join(result.path, \"prediction.png\")))\n",
" break\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "0c67dfdb",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"This template is a quickstart to using [Ray Tune](https://docs.ray.io/en/latest/tune/index.html) for many model training. See [this blog post](https://www.anyscale.com/blog/training-one-million-machine-learning-models-in-record-time-with-ray) for more information on the benefits of performing many model training with Ray!\n",
"\n",
"At a high level, this template showed how to do the following:\n",
"\n",
"1. [Define the training function for a single partition of data.](https://docs.ray.io/en/latest/tune/tutorials/tune-run.html)\n",
"2. [Define a Tune search space to run training over many partitions of data.](https://docs.ray.io/en/latest/tune/tutorials/tune-search-spaces.html)\n",
"3. [Extract the best model per dataset partition from the Tune experiment output.](https://docs.ray.io/en/latest/tune/examples/tune_analyze_results.html)\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "df5fb149",
"metadata": {},
"source": [
"### Appendix\n",
"\n",
"#### Specifying required resources\n",
"\n",
"`tune.with_resources` was used to specify the resources needed to launch one of our training jobs.\n",
"Feel free to change this to the resources required by your application! You can also comment out the `tune.with_resources` block to assign `1 CPU` (the default) to each trial.\n",
"\n",
"Note that the number of CPUs to assign a trial is dependent on the workload.\n",
"In this template, `statsforecast` has a `n_jobs` configuration that determines the number of CPU cores to use for performing the model fitting and cross-validation *within a trial*. So, we should set `n_jobs = cpus_per_trial`. We chose to set the parallelism equal to the total number of models that are fitted during cross-validation: `M model types * N temporal cross-validation windows = 2 * 1 = 2`.\n",
"\n",
"See [Ray Tune's guide on assigning resources](https://docs.ray.io/en/latest/tune/tutorials/tune-resources.html) for more information."
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "dd48618e",
"metadata": {},
"source": []
}
],
"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.8.13"
},
"vscode": {
"interpreter": {
"hash": "265d195fda5292fe8f69c6e37c435a5634a1ed3b6799724e66a975f68fa21517"
}
}
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/templates/03_serving_stable_diffusion/README.md
+80
View File
@@ -0,0 +1,80 @@
# Serving a Stable Diffusion Model with Ray Serve
| Template Specification | Description |
| ---------------------- | ----------- |
| Summary | This app provides users a one click production option for serving a pre-trained Stable Diffusion model from Hugging Face. It leverages [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) to deploy locally and the built-in IDE integration on an Anyscale Workspace so you can iterate and add additional logic to the app. You can then use a simple CLI to deploy to production with [Anyscale Services](https://docs.anyscale.com/productionize/services/get-started?utm_source=ray_docs&utm_medium=docs&utm_campaign=stable_diffusion). |
| Time to Run | Around 2 minutes to setup the models and generate your first image(s). Less than 10 seconds for every subsequent round of image generation (depending on the image size). |
| Minimum Compute Requirements | At least 1 GPU node with 1 NVIDIA A10 GPU. |
| Cluster Environment | This template uses a docker image built on top of the latest Anyscale-provided Ray 2.9 image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=stable_diffusion). See the appendix below for more details. |
## Get Started
**When the workspace is up and running, start coding by clicking on the Jupyter or VS Code icon above. Open the `start.ipynb` file and follow the instructions there.**
By the end, we'll have an application that generates images using stable diffusion for a given prompt!
The application will look something like this:
```text
Enter a prompt (or 'q' to quit): twin peaks sf in basquiat painting style
Generating image(s)...
Generated 4 image(s) in 8.75 seconds to the directory: 58b298d9
```
![Example output](https://github-production-user-asset-6210df.s3.amazonaws.com/3887863/239090189-dc1f1b7b-2fa0-4886-ae12-ca5d35b8ebc9.png)
## Deploying on Anyscale Service
This template also includes an example for deploying stable diffusion in production with a FastAPI server. In order to run it locally on your workspace run:
```bash
serve run app:entrypoint
```
Query the serve application:
```bash
python query.py
```
To deploy to a production endpoint on Anyscale run:
```bash
anyscale service rollout -f service.yaml --name {ENTER_NAME_FOR_SERVICE}
```
You can find the link to the service in the logs of the `anyscale service rollout` command. Something like:
```
(anyscale +2.9s) View the service in the UI at https://console.anyscale.com/services/service_gxr3cfmqn2gethuuiusv2zif.
```
You can call the service programmatically (see the instruction from top right corner's Query button) or using the web interface.
![api-doc-image](https://user-images.githubusercontent.com/21118851/204909023-9e3fac37-40c0-44e3-bfe0-4db502e30c2e.png)
1. Wait for the service to be in a "Running" state.
2. In the "Deployments" section, find the "APIIngress" row, click the "View" under "API Docs".
3. You should now see a OpenAPI rendered documentation page.
4. Click the `/imagine` endpoint, then "Try it out" to enable calling it via the interactive API browser.
5. Fill in your prompt and click execute.
## Appendix
### Advanced: Build off of this template's cluster environment
#### Option 1: Build a new cluster environment on Anyscale
Find a `cluster_env.yaml` file in the working directory of the template. Feel free to modify this YAML to include more requirements, then follow [this guide](https://docs.anyscale.com/configure/dependency-management/cluster-environments#creating-a-cluster-environment?utm_source=ray_docs&utm_medium=docs&utm_campaign=stable_diffusion) to create a new cluster environment with the `anyscale` CLI .
Finally, update your workspace's cluster environment to this new one after it's done building.
#### Option 2: Build a new docker image with your own infrastructure
Use the following `docker pull` command if you want to manually build a new Docker image based off of this one.
```bash
docker pull us-docker.pkg.dev/anyscale-workspace-templates/workspace-templates/serve-stable-diffusion-model-ray-serve:latest
```
doc/source/templates/03_serving_stable_diffusion/app.py
+62
View File
@@ -0,0 +1,62 @@
from io import BytesIO
from ray import serve
from fastapi import FastAPI
from fastapi.responses import Response
import torch
from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline
import logging
app = FastAPI()
logger = logging.getLogger("ray.serve")
@serve.deployment(num_replicas=1)
@serve.ingress(app)
class APIIngress:
def __init__(self, diffusion_model_handle) -> None:
self.handle = diffusion_model_handle
@app.get(
"/imagine",
responses={200: {"content": {"image/png": {}}}},
response_class=Response,
)
async def generate(self, prompt: str, img_size: int = 512):
assert len(prompt), "prompt parameter cannot be empty"
image = await self.handle.generate.remote(prompt, img_size=img_size)
file_stream = BytesIO()
image.save(file_stream, "PNG")
return Response(content=file_stream.getvalue(), media_type="image/png")
@serve.deployment(
ray_actor_options={"num_gpus": 1, "num_cpus": 1},
max_ongoing_requests=2,
autoscaling_config={
"min_replicas": 1,
"max_replicas": 3,
"target_ongoing_requests": 1,
},
)
class StableDiffusionV2:
def __init__(self):
model_id = "stabilityai/stable-diffusion-2"
scheduler = EulerDiscreteScheduler.from_pretrained(
model_id, subfolder="scheduler"
)
self.pipe = StableDiffusionPipeline.from_pretrained(
model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
)
self.pipe = self.pipe.to("cuda")
def generate(self, prompt: str, img_size: int = 512):
assert len(prompt), "prompt parameter cannot be empty"
logger.info("Prompt: [%s]", prompt)
image = self.pipe(prompt, height=img_size, width=img_size).images[0]
return image
entrypoint = APIIngress.bind(StableDiffusionV2.bind())
doc/source/templates/03_serving_stable_diffusion/cluster_env.yaml
+22
View File
@@ -0,0 +1,22 @@
# See https://hub.docker.com/r/anyscale/ray for full list of
# available Ray, Python, and CUDA versions.
base_image: anyscale/ray:2.9.0-py39-cu118
env_vars: {}
debian_packages: []
python:
pip_packages:
- accelerate==0.20.3
- diffusers==0.17.1
- fastapi==0.97.0
- ipywidgets
- matplotlib==3.7.1
- numpy==1.24.3
- torch==2.0.1
- transformers==4.30.1
conda_packages: []
post_build_cmds: []
doc/source/templates/03_serving_stable_diffusion/query.py
+15
View File
@@ -0,0 +1,15 @@
import requests
endpoint = "http://localhost:8000/imagine"
def generate_image(prompt, image_size):
req = {"prompt": prompt, "img_size": image_size}
resp = requests.get(endpoint, params=req)
return resp.content
image = generate_image("twin peaks sf in basquiat painting style", 640)
filename = "image.png"
with open(filename, "wb") as f:
f.write(image)
doc/source/templates/03_serving_stable_diffusion/service.yaml
+7
View File
@@ -0,0 +1,7 @@
name: "stable-diffusion-service"
ray_serve_config:
applications:
- name: stable-diffusion
import_path: app:entrypoint
runtime_env:
working_dir: "."
doc/source/templates/03_serving_stable_diffusion/start.ipynb
+456
View File
@@ -0,0 +1,456 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"id": "597c13c0",
"metadata": {},
"source": [
"# Serving a Stable Diffusion Model with Ray Serve\n",
"\n",
"| Template Specification | Description |\n",
"| ---------------------- | ----------- |\n",
"| Summary | This template loads a pretrained stable diffusion model from HuggingFace and serves it to a local endpoint as a [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) deployment. |\n",
"| Time to Run | Around 2 minutes to setup the models and generate your first image(s). Less than 10 seconds for every subsequent round of image generation (depending on the image size). |\n",
"| Minimum Compute Requirements | At least 1 GPU node. The default is 4 nodes, each with 1 NVIDIA T4 GPU. |\n",
"| Cluster Environment | This template uses a custom docker image built on top of the Anyscale-provided Ray image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview). See the appendix in the `README` for more details. |\n",
"\n",
"By the end, we'll have an application that generates images using stable diffusion for a given prompt!\n",
"\n",
"The application will look something like this:\n",
"\n",
"```text\n",
"Enter a prompt (or 'q' to quit): twin peaks sf in basquiat painting style\n",
"\n",
"Generating image(s)...\n",
"\n",
"Generated 4 image(s) in 8.75 seconds to the directory: 58b298d9\n",
"```\n",
"\n",
"![Example output](https://github-production-user-asset-6210df.s3.amazonaws.com/3887863/239090189-dc1f1b7b-2fa0-4886-ae12-ca5d35b8ebc9.png)\n",
"> Slot in your code below wherever you see the ✂️ icon to build off of this template!\n",
">\n",
"> The framework and data format used in this template can be easily replaced to suit your own application!\n",
"\n",
"We'll start with some imports and initialize Ray:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "59842da3",
"metadata": {},
"outputs": [],
"source": [
"from fastapi import FastAPI\n",
"from fastapi.responses import Response\n",
"from io import BytesIO\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import os\n",
"import requests\n",
"import time\n",
"import uuid\n",
"\n",
"import ray\n",
"from ray import serve\n",
"\n",
"ray.init()\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "520ef4d7",
"metadata": {},
"source": [
"## Deploy the Ray Serve application locally\n",
"\n",
"First, we define the Ray Serve application with the model loading and inference logic. This includes setting up:\n",
"- The `/imagine` API endpoint that we query to generate the image.\n",
"- The stable diffusion model loaded inside a Ray Serve Deployment.\n",
" We'll specify the *number of model replicas* to keep active in our Ray cluster. These model replicas can process incoming requests concurrently.\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "de6318ac",
"metadata": {},
"source": [
"> ✂️ Replace these values to change the number of model replicas to serve, as well as the GPU resources required by each replica.\n",
">\n",
"> With more model replicas, more images can be generated in parallel!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90eca147",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"NUM_REPLICAS: int = 4\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "be41ca9e",
"metadata": {},
"outputs": [],
"source": [
"if NUM_REPLICAS > ray.available_resources()[\"GPU\"]:\n",
" print(\n",
" \"Your cluster does not currently have enough resources to run with these settings. \"\n",
" \"Consider decreasing the number of workers, or decreasing the resources needed \"\n",
" \"per worker. Ignore this if your cluster auto-scales.\"\n",
" )\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "89eb3e2c",
"metadata": {},
"source": [
"First, we define the Ray Serve Deployment, which will load a stable diffusion model and perform inference with it.\n",
"\n",
"> ✂️ Modify this block to load your own model, and change the `generate` method to perform your own online inference logic!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f203efd4",
"metadata": {},
"outputs": [],
"source": [
"@serve.deployment(\n",
" ray_actor_options={\"num_gpus\": 1},\n",
" num_replicas=NUM_REPLICAS,\n",
")\n",
"class StableDiffusionV2:\n",
" def __init__(self):\n",
" # <Replace with your own model loading logic>\n",
" import torch\n",
" from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline\n",
"\n",
" model_id = \"stabilityai/stable-diffusion-2\"\n",
" scheduler = EulerDiscreteScheduler.from_pretrained(\n",
" model_id, subfolder=\"scheduler\"\n",
" )\n",
" self.pipe = StableDiffusionPipeline.from_pretrained(\n",
" model_id, scheduler=scheduler, revision=\"fp16\", torch_dtype=torch.float16\n",
" )\n",
" self.pipe = self.pipe.to(\"cuda\")\n",
"\n",
" def generate(self, prompt: str, img_size: int = 776):\n",
" # <Replace with your own model inference logic>\n",
" assert len(prompt), \"prompt parameter cannot be empty\"\n",
" image = self.pipe(prompt, height=img_size, width=img_size).images[0]\n",
" return image\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "0134aa54",
"metadata": {},
"source": [
"Next, we'll define the actual API endpoint to live at `/imagine`.\n",
"\n",
"> ✂️ Modify this block to change the endpoint URL, response schema, and add any post-processing logic needed from your model output!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6f80fee2",
"metadata": {},
"outputs": [],
"source": [
"app = FastAPI()\n",
"\n",
"\n",
"@serve.deployment(num_replicas=1)\n",
"@serve.ingress(app)\n",
"class APIIngress:\n",
" def __init__(self, diffusion_model_handle) -> None:\n",
" self.handle = diffusion_model_handle\n",
"\n",
" @app.get(\n",
" \"/imagine\",\n",
" responses={200: {\"content\": {\"image/png\": {}}}},\n",
" response_class=Response,\n",
" )\n",
" async def generate(self, prompt: str, img_size: int = 776):\n",
" assert len(prompt), \"prompt parameter cannot be empty\"\n",
"\n",
" image = await self.handle.generate.remote(prompt, img_size=img_size)\n",
"\n",
" file_stream = BytesIO()\n",
" image.save(file_stream, \"PNG\")\n",
" return Response(content=file_stream.getvalue(), media_type=\"image/png\")\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "61b8916d",
"metadata": {},
"source": [
"Now, we deploy the Ray Serve application locally at `http://localhost:8000`!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dfc2e244",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"entrypoint = APIIngress.bind(StableDiffusionV2.bind())\n",
"\n",
"# Shutdown any existing Serve replicas, if they're still around.\n",
"serve.shutdown()\n",
"serve.run(entrypoint, name=\"serving_stable_diffusion_template\")\n",
"print(\"Done setting up replicas! Now accepting requests...\")\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "757678cc",
"metadata": {},
"source": [
"## Make requests to the endpoint\n",
"\n",
"Next, we'll build a simple client to submit prompts as HTTP requests to the local endpoint at `http://localhost:8000/imagine`."
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "008976b5",
"metadata": {},
"source": [
"Start the client script in the next few cells, and generate your first image!\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "67ad095b",
"metadata": {},
"outputs": [],
"source": [
"endpoint = \"http://localhost:8000/imagine\"\n",
"\n",
"\n",
"@ray.remote(num_cpus=0)\n",
"def generate_image(prompt, image_size):\n",
" req = {\"prompt\": prompt, \"img_size\": image_size}\n",
" resp = requests.get(endpoint, params=req)\n",
" return resp.content\n",
"\n",
"\n",
"def show_images(filenames):\n",
" fig, axs = plt.subplots(1, len(filenames), figsize=(4 * len(filenames), 4))\n",
" for i, filename in enumerate(filenames):\n",
" ax = axs if len(filenames) == 1 else axs[i]\n",
" ax.imshow(plt.imread(filename))\n",
" ax.axis(\"off\")\n",
" plt.show()\n",
"\n",
"\n",
"def main(\n",
" interactive: bool = False,\n",
" prompt: str = \"twin peaks sf in basquiat painting style\",\n",
" num_images: int = 4,\n",
" image_size: int = 640,\n",
"):\n",
" try:\n",
" requests.get(endpoint, timeout=0.1)\n",
" except Exception as e:\n",
" raise RuntimeWarning(\n",
" \"Did you setup the Ray Serve model replicas with `serve.run` \"\n",
" \"in a previous cell?\"\n",
" ) from e\n",
"\n",
" generation_times = []\n",
" while True:\n",
" prompt = (\n",
" prompt\n",
" if not interactive\n",
" else input(f\"\\nEnter a prompt (or 'q' to quit): \")\n",
" )\n",
" if prompt.lower() == \"q\":\n",
" break\n",
"\n",
" print(\"\\nGenerating image(s)...\\n\")\n",
" start = time.time()\n",
"\n",
" # Make `num_images` requests to the endpoint at once!\n",
" images = ray.get(\n",
" [generate_image.remote(prompt, image_size) for _ in range(num_images)]\n",
" )\n",
"\n",
" dirname = f\"{uuid.uuid4().hex[:8]}\"\n",
" os.makedirs(dirname)\n",
" filenames = []\n",
" for i, image in enumerate(images):\n",
" filename = os.path.join(dirname, f\"{i}.png\")\n",
" with open(filename, \"wb\") as f:\n",
" f.write(image)\n",
" filenames.append(filename)\n",
"\n",
" elapsed = time.time() - start\n",
" generation_times.append(elapsed)\n",
" print(\n",
" f\"\\nGenerated {len(images)} image(s) in {elapsed:.2f} seconds to \"\n",
" f\"the directory: {dirname}\\n\"\n",
" )\n",
" show_images(filenames)\n",
" if not interactive:\n",
" break\n",
" return np.mean(generation_times) if generation_times else -1\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "c8949cc7",
"metadata": {},
"source": [
"Once the stable diffusion model finishes generating your image(s), it will be included in the HTTP response body.\n",
"The client saves all the images in a local directory for you to view, and they'll also show up in the notebook cell!"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "3e29193b",
"metadata": {},
"source": [
"> ✂️ Replace this value to change the number of images to generate per prompt.\n",
">\n",
"> Each image will be generated starting from a different set of random noise,\n",
"> so you'll be able to see multiple options per prompt!\n",
">\n",
"> Try starting with `NUM_IMAGES_PER_PROMPT` equal to `NUM_REPLICAS` from earlier.\n",
">\n",
"> You can choose to run this interactively, or submit a single `PROMPT`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dd20a52d",
"metadata": {},
"outputs": [],
"source": [
"NUM_IMAGES_PER_PROMPT: int = NUM_REPLICAS\n",
"\n",
"# Control the output size: (IMAGE_SIZE, IMAGE_SIZE)\n",
"# The stable diffusion model requires `IMAGE_SIZE` to be a multiple of 8.\n",
"# NOTE: Generated image quality degrades rapidly if you reduce the size too much.\n",
"IMAGE_SIZE: int = 640\n",
"\n",
"INTERACTIVE: bool = False\n",
"PROMPT = \"twin peaks sf in basquiat painting style\"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"mean_generation_time = main(\n",
" interactive=INTERACTIVE,\n",
" prompt=PROMPT,\n",
" num_images=NUM_IMAGES_PER_PROMPT,\n",
" image_size=IMAGE_SIZE,\n",
")\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "fb124968",
"metadata": {},
"source": [
"You've successfully served a stable diffusion model!\n",
"You can modify this template and iterate your model deployment directly on your cluster within your Anyscale Workspace,\n",
"testing with the local endpoint."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9e360cf9",
"metadata": {},
"outputs": [],
"source": [
"# Shut down the model replicas once you're done!\n",
"serve.shutdown()\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "880c2d6f",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"This template used [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) to serve many replicas of a stable diffusion model. \n",
"\n",
"At a high level, this template showed how to:\n",
"1. Define a Ray Serve deployment to load a HuggingFace model and perform inference.\n",
"2. Set up a local endpoint to accept and route requests to the different model replicas.\n",
"3. Make multiple requests in parallel to generate many images at a time.\n",
"\n",
"See this [getting started guide](https://docs.ray.io/en/latest/serve/getting_started.html) for a more detailed walkthrough of Ray Serve."
]
},
{
"attachments": {},
"cell_type": "markdown",
"id": "bcc69b2d",
"metadata": {},
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "ray_dev_py38",
"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.8.13"
},
"vscode": {
"interpreter": {
"hash": "265d195fda5292fe8f69c6e37c435a5634a1ed3b6799724e66a975f68fa21517"
}
}
},
"nbformat": 4,
"nbformat_minor": 5
}
doc/source/templates/04_finetuning_llms_with_deepspeed/README.md
+237
View File
@@ -0,0 +1,237 @@
# Fine-tuning Llama-2 series models with Deepspeed, Accelerate, and Ray Train TorchTrainer
| Template Specification | Description |
| ---------------------- | ----------- |
| Summary | This template, demonstrates how to perform fine-tuning (full parameter or LoRA) for Llama-2 series models (7B, 13B, and 70B) using TorchTrainer with the DeepSpeed ZeRO-3 strategy. |
| Time to Run | 1 epoch (3.5M tokens) training wall-clock time: ~14 min. for 7B, ~26 min. for 13B, and ~190 min. for 70B (see the setup details below) |
| Minimum Compute Requirements | 16xg5.4xlarge for worker nodes for 7B model, 4xg5.12xlarge nodes for 13B model, and 4xg5.48xlarge (or 2xp4de.24xlarge) nodes for 70B|
| Cluster Environment | This template uses a Docker image built on top of the latest Anyscale-provided Ray image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=finetuning_llms). |
## Getting Started
For a full-parameter fine-tuning of 7B models, set up a cluster on AWS with the following settings:
| | num | instance type | GPU per node | GPU Memory | CPU Memory |
|------------|-----|---------------|--------------|------------|------------|
| Head node | 1 | m5.xlarge | - | - | - |
| Worker node| 16 | g5.4xlarge | 1 x A10G | 24 GB | 64 GB |
And launch the following script to fine-tune LLaMA 2 7B:
```
./run_llama_ft.sh --size=7b --as-test
```
The flag `--as-test` is for demo / testing purposes as it runs through only one forward and backward pass of the model. The model loading, and remote checkpointing would still run.
Similarly for 13B you need a different compute config.
| | num | instance type | GPU per node | GPU Memory | CPU Memory |
|------------|-----|---------------|--------------|------------|------------|
| Head node | 1 | m5.xlarge | - | - | - |
| Worker node| 4 | g5.12xlarge | 4 x A10G | 24 GB | 64 GB |
```
./run_llama_ft.sh --size=13b [--as-test]
```
## What is happening under the hood?
### Downloading the pre-trained checkpoint on to all GPU nodes.
The pre-trained models for these models is quite large (12.8G for 7B model and 128G for 70B model). In order to make loading these models faster, we have mirrored the weights on to an AWS S3 bucket which can result in up 10GB/s download speed if the aws configs are setup correctly.
### Cloud storage
Similarly the checkpoints during training can be quite large and we would like to be able to save those checkpoints to the familiar huggingface format so that we can serve it conveniently. The fine-tuning script in this template uses Ray Train Checkpointing to sync the checkpoints created by each node back to a centralized cloud storage on AWS S3. The final file structure for each checkpoint will have a look similar to the following structure:
```
aws s3 ls s3://<bucket_path>/checkpoint_00000
├── .is_checkpoint
├── .metadata.pkl
├── .tune_metadata
├── _metadata.meta.pkl
├── _preprocessor
├── _preprocessor.meta.pkl
├── added_tokens.json
├── config.json
├── generation_config.json
├── model-00001-of-00002.safetensors
├── model-00002-of-00002.safetensors
├── model.safetensors
├── model.safetensors.index.json
├── special_tokens_map.json
├── tokenizer.json
├── tokenizer.model
└── tokenizer_config.json
```
After training we can use [RayLLM](https://github.com/ray-project/ray-llm) to deploy our fine-tuned LLM by providing the checkpoint path stored on cloud directly.
### Creating the dataset
The main fine-tuning script is written in a general format that would require you to provide a `jsonl` file for train and test datasets in addition to a `json` file listing the special tokens used in your dataset.
For example each row in your dataset might be formated like the following:
```
{"input": "<ASSISTANT>How can I help you?</ASSISTANT><USER>how is the weather?</USER>}
```
And the special tokens can be:
```
{"tokens": ["<ASSISTANT>", "</ASSISTANT>", "<USER>", "</USER>"]}
```
Depending on the dataset you want to fine-tune on, the tokenization and dataset pre-processing will likely need to be adjusted. The current code is configured to train on the Grade School Math 8k (GSM8K) dataset. By running the code below we create three files that are needed to launch the training script with.
```
python create_dataset.py
>>> data/train.jsonl # 7.4k training data
>>> data/test.jsonl # 1.3k test data
>>> tokens.json # a list of special tokens
```
This dataset is trained with a context length of 512 which includes excessive padding to keep all samples limited to 512 tokens. This means that the training dataset has 3.5 M tokens.
### Launching fine-tuning
The script is written using Ray Train + Deepspeed integration via accelerate API. The script is general enough that it can be used to fine-tune all released sizes of Llama-2 models.
The command for seeing all the options is:
```
python finetune_hf_llm.py --help
```
This script was tested across three model sizes on the following cluster configurations on Anyscale platform.
| Model Size | Base HF Model ID | Batch size per device | GPUs | Time per epoch (min.) |
|------------|------------------------------|-----------------------|----------------|-----------------------|
| 7B | `meta-llama/Llama-2-7b-hf` | 16 | 16x A10G (24G) | ~14 min. |
| 13B | `meta-llama/Llama-2-13b-hf` | 16 | 16x A10G (24G) | ~26 min. |
| 70B | `meta-llama/Llama-2-70b-hf` | 8 | 32x A10G (24G) | ~190 min. |
To launch a full fine-tuning you can use the following command:
```
./run_llama_ft.sh --size=7b
```
### Launching LoRA fine-tuning
You can utilize [LoRA](https://arxiv.org/abs/2106.09685) to achieve more resource efficient fine-tuning results than full-parameter fine-tuning, but unlocking smaller instance types and more efficient model serving. To launch a LoRA fine-tuning, you can use the following command or similar commands for other model sizes:
```
./run_llama_ft.sh --size=7b --lora
```
Fine-tuning a model with LoRA results in a checkpoint containing only the fine-tuned weights. As an example, the default Llama 2 LoRA configuration should yield a 42/64/202MB checkpoint for 7B/13B/70B models. If we want to evaluate the model after training, we can merge the model weights with the original (non-fine-tuned) model. We provide a script to merge the fine-tuned weights with the original weights to produce a full-parameter checkpoint. The script has high CPU memory requirements because it requires us to load all parameters into memory at the same time, 13GB/24GB/152GB for 7B/13B/70B models. Downloading and loading the original weights should take ~1min/~2min/~10min each on a p4de.24xlarge instance. You can run the script as follows:
```
python merge_lora_weights.py --model-name=7b --checkpoint=<path to your checkpoint> --output-path=<desired output path>
```
This leaves a self-contained LoRA fine-tuned model, config and tokenizer at the desired output path.
### Guideline on how to pick node instances when A100s are not available.
Here is the suggested cluster config for each workload:
7B:
```
head_node_type:
name: head_node_type
instance_type: m5.xlarge
worker_node_types:
- name: gpu_worker
instance_type: g5.4xlarge
min_workers: 0
max_workers: 16
use_spot: false
```
13B:
```
head_node_type:
name: head_node_type
instance_type: m5.xlarge
worker_node_types:
- name: gpu_worker
instance_type: g5.12xlarge
min_workers: 0
max_workers: 4
use_spot: false
```
70B:
```
head_node_type:
name: head_node_type
instance_type: m5.xlarge
worker_node_types:
- name: gpu_worker
instance_type: g5.48xlarge
min_workers: 0
max_workers: 4
use_spot: false
```
There are two things that you should consider when choosint the cluster configurations:
1. CPU RAM requirement for optimizer state and parameter offloading
Deepspeed offers [Zero-offload](https://www.deepspeed.ai/tutorials/zero-offload/) which allows offloading the optimizer or parameter states to the CPU memory for more memory efficient training. We have enabled this by default in our deepspeed configs used for this workspace template.
This method creates extra CPU RAM requirements on the machines. A rule of thumb for this implementation is that it needs O(18M/N*K) CPU RAM where M is the model size, N is the number shards, and K is the number of GPUs on a single machine.
For example, for 70B model, on an 8xA100 machine with 8-way sharding you would need `18 * (70 / 8) * 8 = 1.26 TB` of CPU RAM. This is not available on a single machine of 8xA100s. But if we use two 8xA100 machines instead, with 16-way sharding we would need `18 * (70 / 16) * 8 = 630 GB` of CPU RAM which is accessible.
Another example: For 70B model, on an 4xA10G machine with 32-way sharding you would need `18 * (70 / 32) * 4 = 158 GB` of CPU RAM on each machine. If you use 8xA10G machines instead you would need `18 * (70 / 32) * 8 = 316 GB` of CPU RAM on each machine.
So availability of enough CPU RAM is very important when using optimizer state offloading.
2. CPU RAM requirement during checkpointing
During checkpointing in the middle of training, we have to aggregate the weights from all the shards back to rank 0 so that it can save the model. We can also save the weights of each shard independently and aggregate the weights later offline. The extra CPU memory requirement would not get solved tho.
Emprically the implementation that `accelerate` provides needs `O(4M)` CPU RAM on rank 0 machine where M is the model size. This would mean that for 70B we need 280GB of CPU on top of what we needed before (e.g. due to CPU offloading). This requirement is only for rank 0 though and not any other machine. So it's important to schedule this process on a machine with this much of RAM while the other processes can get scheduled on machines with lower RAM requirements.
For example, for 70B model, with 32-way sharding on a machine with 8xA10Gs (g5.48xlarge), you need 280G (because of checkpointing) and 315 GB (because of optimizer state offloading) making the total memory requirement ~595 GB.
Ray provides an easy way to control which process gets launched on what machine type. To do this, in your cluster config add a custom label for those machines that satisfies the CPU RAM requirement of rank 0 and call them `large_cpu_mem` instances. Then in our script we specify the custom tag as a resource requirement for the `trainer` actor which is in the same machine that rank zero process will get executed on.
```
scaling_config=air.ScalingConfig(
# "large_cpu_mem" is the tag used to identify this machine type in the
# cluster config.
trainer_resources={"large_cpu_mem": 0.01},
num_workers=args.num_devices,
use_gpu=True,
resources_per_worker={"GPU": 1},
)
```
### Submiting a production job
You can easily submit a production job using the following command:
```
python create_job_yaml.py --size=7b --output-path=./job.yaml
```
This will create a job yaml file that you can use to submit a production job on Anyscale platform.
```
anyscale job submit job.yaml
```
doc/source/templates/04_finetuning_llms_with_deepspeed/compute_configs/aws_70b_v1.yaml
+22
View File
@@ -0,0 +1,22 @@
region: us-west1
allowed_azs: [any]
head_node_type:
name: head_node_type
instance_type: g5.48xlarge
resources:
custom_resources:
large_cpu_mem: 1
worker_node_types:
- name: gpu_worker
instance_type: g5.48xlarge
min_workers: 3
max_workers: 3
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/04_finetuning_llms_with_deepspeed/compute_configs/aws_70b_v2.yaml
+28
View File
@@ -0,0 +1,28 @@
region: us-west1
allowed_azs: [any]
head_node_type:
name: head_node_type
instance_type: g5.48xlarge
resources:
custom_resources:
large_cpu_mem: 1
worker_node_types:
- name: large_gpu_worker
instance_type: g5.48xlarge
min_workers: 2
max_workers: 2
use_spot: false
- name: medium_gpu_worker
instance_type: g5.24xlarge
min_workers: 2
max_workers: 2
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/04_finetuning_llms_with_deepspeed/compute_configs/aws_7b_or_13b.yaml
+20
View File
@@ -0,0 +1,20 @@
# Autoscale to 16 g5.4xlarge --> 16 A10Gs
region: us-west1
allowed_azs: [any]
head_node_type:
name: head_node
instance_type: m5.xlarge
worker_node_types:
- name: worker_node
instance_type: g5.4xlarge
min_workers: 0
max_workers: 16
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/04_finetuning_llms_with_deepspeed/compute_configs/gce_7b_or_13b.yaml
+16
View File
@@ -0,0 +1,16 @@
head_node_type:
name: head_node_type
instance_type: g2-standard-32-nvidia-l4-1
resources:
custom_resources:
large_cpu_mem: 1
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-16-nvidia-l4-1
min_workers: 15
max_workers: 15
use_spot: false
resources:
custom_resources:
medium_cpu_mem: 1
doc/source/templates/04_finetuning_llms_with_deepspeed/create_dataset.py
+31
View File
@@ -0,0 +1,31 @@
from datasets import load_dataset
import json
import os
dataset = load_dataset("gsm8k", "main")
dataset_splits = {"train": dataset["train"], "test": dataset["test"]}
def main():
if not os.path.exists("data"):
os.mkdir("data")
with open("data/tokens.json", "w") as f:
tokens = {}
tokens["tokens"] = ["<START_Q>", "<END_Q>", "<START_A>", "<END_A>"]
f.write(json.dumps(tokens))
for key, ds in dataset_splits.items():
with open(f"data/{key}.jsonl", "w") as f:
for item in ds:
newitem = {}
newitem["input"] = (
f"<START_Q>{item['question']}<END_Q>"
f"<START_A>{item['answer']}<END_A>"
)
f.write(json.dumps(newitem) + "\n")
if __name__ == "__main__":
main()
doc/source/templates/04_finetuning_llms_with_deepspeed/create_job_yaml.py
+79
View File
@@ -0,0 +1,79 @@
from argparse import ArgumentParser
import yaml
import os
import pathlib
def _parse_args():
parser = ArgumentParser()
parser.add_argument(
"--size",
type=str,
default="7b",
choices=["7b", "13b", "70b"],
help="Size of the model to train",
)
parser.add_argument(
"--as-test", action="store_true", help="Whether to run in test mode"
)
parser.add_argument(
"--max-retries",
type=int,
default=0,
help="Number of times to retry the job if it fails",
)
parser.add_argument(
"--output-path",
type=str,
default="./job.yaml",
help="The path that job yaml should be stored.",
)
parser.add_argument("--compute-config", type=str, help="Path to the compute config")
parser.add_argument(
"--cluster-env-build-id",
type=str,
help="The build-id of the cluster env to use",
)
return parser.parse_args()
def main():
pargs = _parse_args()
# Resolve compute config
compute_config_kwargs = {}
if pargs.compute_config:
with open(pargs.compute_config, "r") as f:
compute_config = yaml.safe_load(f)
compute_config.update(
{
"cloud_id": os.environ["ANYSCALE_CLOUD_ID"],
}
)
compute_config_kwargs.update(compute_config=compute_config)
# Resolve cluster env config
cluster_env_config_kwargs = {}
if pargs.cluster_env_build_id:
cluster_env_config_kwargs.update(build_id=pargs.cluster_env_build_id)
base_cmd = f"chmod +x ./run_llama_ft.sh && ./run_llama_ft.sh --size={pargs.size}"
job_config = {
"name": f"llama-2-{pargs.size}",
"entrypoint": base_cmd + (" --as-test" if pargs.as_test else ""),
"max_retries": pargs.max_retries,
**compute_config_kwargs,
**cluster_env_config_kwargs,
}
pathlib.Path(os.path.dirname(pargs.output_path)).mkdir(parents=True, exist_ok=True)
with open(pargs.output_path, "w") as f:
yaml.safe_dump(job_config, f)
print("Job config written to ", pargs.output_path)
print("To submit the job, run:")
print(f"anyscale job submit {pargs.output_path}")
if __name__ == "__main__":
main()
doc/source/templates/04_finetuning_llms_with_deepspeed/deepspeed_configs/zero_3_llama_2_13b.json
+35
View File
@@ -0,0 +1,35 @@
{
"fp16": {
"enabled": "auto"
},
"bf16": {
"enabled": "auto"
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e6,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true,
"round_robin_gradients": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 10,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
doc/source/templates/04_finetuning_llms_with_deepspeed/deepspeed_configs/zero_3_llama_2_70b.json
+28
View File
@@ -0,0 +1,28 @@
{
"fp16": {
"enabled": false
},
"bf16": {
"enabled": true
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": false
},
"overlap_comm": true,
"contiguous_gradients": true,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"gather_16bit_weights_on_model_save": true,
"round_robin_gradients": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 10,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
doc/source/templates/04_finetuning_llms_with_deepspeed/deepspeed_configs/zero_3_llama_2_70b_nvme_offload.json
+31
View File
@@ -0,0 +1,31 @@
{
"fp16": {
"enabled": "auto"
},
"bf16": {
"enabled": "auto"
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "nvme",
"nvme_path": "/mnt/local_storage/zero",
"pin_memory": false,
"buffer_count": 4,
"fast_init": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"gather_16bit_weights_on_model_save": true,
"round_robin_gradients": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 10,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
doc/source/templates/04_finetuning_llms_with_deepspeed/deepspeed_configs/zero_3_llama_2_7b.json
+35
View File
@@ -0,0 +1,35 @@
{
"fp16": {
"enabled": "auto"
},
"bf16": {
"enabled": "auto"
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": 5e8,
"stage3_prefetch_bucket_size": 5e8,
"stage3_param_persistence_threshold": 1e6,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true,
"round_robin_gradients": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 10,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
doc/source/templates/04_finetuning_llms_with_deepspeed/finetune_hf_llm.py
+773
View File
@@ -0,0 +1,773 @@
import argparse
from filelock import FileLock
import functools
import json
import math
import os
from pathlib import Path
import re
import tempfile
import time
import tree
from typing import Tuple
try:
import deepspeed # noqa: F401
except ImportError as e:
raise RuntimeError(
"Please install deepspeed with `pip install --user deepspeed`."
) from e
from accelerate import Accelerator, DeepSpeedPlugin
from accelerate.utils import DummyOptim, DummyScheduler, set_seed
import torch
import torch.nn as nn
import tqdm
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
get_linear_schedule_with_warmup,
)
from peft import LoraConfig, get_peft_model
import ray
from ray import train
import ray.util.scheduling_strategies
from ray.train.torch import TorchTrainer
from ray.train import Checkpoint
from utils import (
get_checkpoint_and_refs_dir,
get_mirror_link,
download_model,
get_download_path,
)
OPTIM_BETAS = (0.9, 0.999)
OPTIM_EPS = 1e-8
NUM_WARMUP_STEPS = 10
OPTIM_WEIGHT_DECAY = 0.0
ATTENTION_LAYER_NAME = "self_attn"
def get_expected_lora_num_parameters(
model, lora_config: LoraConfig, attn_layer_name: str = ATTENTION_LAYER_NAME
):
"""Calculate the expected number of parameters for lora finetuning."""
sum_params = 0
num_attention_layers = 0
modules = model.named_modules()
loraified_modules = 0
# We calculate the number of parameters we need for lora finetuning by calculating
# the sizes of the deecomposed weight matrices according to the paper.
for full_name, target in modules:
layer_name = full_name.split(".")[-1]
if layer_name == attn_layer_name:
# Detected another attention layer (for example, llama 2 70b should have 80
# of these)
num_attention_layers += 1
elif layer_name in lora_config.modules_to_save:
# Detect another non-lora module to save, which will also contribute to the
# number of checkpointed parameters. This will result in one set of
# trainable parameters "<layer>.original_module.weight" and another one with
# "<layer>.modules_to_save.default.weight"
# Therefore, each layer contributes 2 x the number of actual elements in
# that layer.
sum_params += 2 * target.weight.numel()
print(
"Found non-lora-layer to checkpoint: ",
layer_name,
" with num params ",
target.weight.numel(),
)
else:
for module_name in lora_config.target_modules:
if layer_name == module_name:
loraified_modules += 1
if isinstance(target, nn.Linear):
# Target is attention weight
sum_params += (
target.in_features + target.out_features
) * lora_config.r
elif isinstance(target, nn.Embedding):
# Target is linear weight
sum_params += (
target.embedding_dim + target.num_embeddings
) * lora_config.r
print(
f"Detected {num_attention_layers} attention layers, containing"
f" {loraified_modules} modules to modify according to LoRA's `target_modules`."
f" This should yield {sum_params} trainable parameters."
)
return sum_params
def get_number_of_params(model: nn.Module):
sum = 0
for name, param in model.named_parameters():
if param.requires_grad:
sum += param.numel()
return sum
def collate_fn(batch, tokenizer, block_size, device):
out_batch = tokenizer(
list(batch["input"]),
padding="max_length",
max_length=block_size,
truncation=True,
return_tensors="pt",
)
out_batch["labels"] = out_batch["input_ids"].clone()
out_batch = tree.map_structure(lambda x: x.to(device), out_batch)
return out_batch
def get_pretrained_path(model_id: str):
mirror_uri = get_mirror_link(model_id)
ckpt_path, _ = get_checkpoint_and_refs_dir(
model_id=model_id, bucket_uri=mirror_uri, s3_sync_args=["--no-sign-request"]
)
return ckpt_path
def get_tokenizer(model_name, special_tokens):
pretrained_path = get_pretrained_path(model_name)
# Context for legacy=True: https://github.com/huggingface/transformers/issues/25176
tokenizer = AutoTokenizer.from_pretrained(pretrained_path, legacy=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.add_tokens(special_tokens, special_tokens=True)
return tokenizer
def evaluate(
*, model, eval_ds, accelerator, bsize, ds_kwargs, as_test: bool = False
) -> Tuple[float, float]:
model.eval()
losses = []
eval_dataloader = eval_ds.iter_torch_batches(batch_size=bsize, **ds_kwargs)
eval_ds_len = len(list(eval_ds.iter_batches(batch_size=1)))
for step, batch in tqdm.tqdm(
enumerate(eval_dataloader), total=eval_ds_len // (bsize + 1)
):
with torch.no_grad():
outputs = model(**batch)
loss = outputs.loss
# The tensors are gathered by concatenating them on the first dimension, so we
# add a new dimension to the scalar loss to get a tensor of shape (K,) for K
# workers.
losses.append(accelerator.gather(loss[None]))
if as_test:
break
# We stack losses so that we have a tensor of shape (T, K) where T is the number of
# steps and K is the number of workers.
losses = torch.stack(losses)
try:
eval_loss = torch.mean(losses).item()
perplexity = math.exp(eval_loss)
except OverflowError:
perplexity = float("inf")
return perplexity, eval_loss
def _test_tokenizer(model_name):
# This function tests that adding special tokens does not
# result in un-expected tokenization
# Context: https://github.com/huggingface/transformers/issues/25176
tokenizer = get_tokenizer(model_name=model_name, special_tokens=["<REPR_END>"])
testoutput = tokenizer("<REPR_END>inform")["input_ids"]
expected = tokenizer("inform")["input_ids"]
assert testoutput[-1] == expected[-1], (
"The tokenizer is not working as expected with special tokens, "
f"testoutput={testoutput}, expected={expected}"
)
def checkpoint_model(
checkpoint_folder, ckpt_id, model, epoch, last_global_step, **kwargs
):
"""Utility function for checkpointing model + optimizer dictionaries
The main purpose for this is to be able to resume training from that instant again.
"""
checkpoint_state_dict = {
"epoch": epoch,
"last_global_step": last_global_step,
}
# Add extra kwargs too
checkpoint_state_dict.update(kwargs)
# In here model will be a DeepspeedEngine object
model.save_checkpoint(checkpoint_folder, ckpt_id, checkpoint_state_dict)
status_msg = (
f"checkpointing: checkpoint_folder={checkpoint_folder}, ckpt_id={ckpt_id}"
)
print(status_msg)
def training_function(kwargs: dict):
print("training_function called")
# Train has a bug somewhere that causes ACCELERATE_TORCH_DEVICE to not be set
# properly on multi-gpu nodes
cuda_visible_device = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
local_rank = int(os.environ["LOCAL_RANK"])
device_id = cuda_visible_device[local_rank]
os.environ["ACCELERATE_TORCH_DEVICE"] = f"cuda:{device_id}"
config = kwargs["config"]
args = argparse.Namespace(**kwargs["args"])
special_tokens = kwargs.get("special_tokens", [])
model_id = config["model_name"]
# We need to download the model weights on this machine if they don't exit.
# We need to acquire a lock to ensure that only one process downloads the model
bucket_uri = get_mirror_link(model_id)
download_path = get_download_path(model_id)
base_path = Path(download_path).parent
base_path.mkdir(parents=True, exist_ok=True)
lock_file = str(base_path / f'{model_id.replace("/", "--")}.lock')
with FileLock(lock_file):
download_model(
model_id=model_id, bucket_uri=bucket_uri, s3_sync_args=["--no-sign-request"]
)
# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
lr = config["lr"]
num_epochs = int(config["num_epochs"])
seed = int(config["seed"])
batch_size = int(config["batch_size"])
gradient_accumulation_steps = int(config["gradient_accumulation_steps"])
# Get deepspeed config to setup the batch size per device
ds_plugin = config["ds_plugin"]
ds_plugin.hf_ds_config.config["train_micro_batch_size_per_gpu"] = batch_size
# Initialize accelerator
accelerator = Accelerator(
deepspeed_plugin=ds_plugin,
gradient_accumulation_steps=gradient_accumulation_steps,
mixed_precision=args.mx,
)
set_seed(seed)
# train_ds is the local shard for this model
train_ds = train.get_dataset_shard("train")
valid_ds = train.get_dataset_shard("valid")
train_ds_len = len(list(train_ds.iter_batches(batch_size=1)))
_test_tokenizer(args.model_name)
tokenizer = get_tokenizer(model_name=args.model_name, special_tokens=special_tokens)
collate_partial = functools.partial(
collate_fn,
tokenizer=tokenizer,
block_size=config["block_size"],
device=accelerator.device,
)
pretrained_path = get_pretrained_path(model_id)
print(f"Loading model from {pretrained_path} ...")
s = time.time()
model = AutoModelForCausalLM.from_pretrained(
pretrained_path,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
# `use_cache=True` is incompatible with gradient checkpointing.
use_cache=False,
use_flash_attention_2=True,
)
print(f"Done loading model in {time.time() - s} seconds.")
model.resize_token_embeddings(len(tokenizer))
if config["lora"]:
# Apply LoRA
s = time.time()
lora_config = LoraConfig(**config["lora_config"])
expected_num_parameters = get_expected_lora_num_parameters(
lora_config=lora_config, model=model
)
print(f"Attempting to apply LoRA config: {lora_config}")
model.enable_input_require_grads()
model = get_peft_model(model, lora_config)
num_parameters = get_number_of_params(model)
if num_parameters != expected_num_parameters:
raise ValueError(
f"Expected {expected_num_parameters} parameters, got {num_parameters} "
f"parameters. LoRA-ification failed."
)
print(
f"LoRA-ification done in {time.time() - s} seconds. Estimated checkpoint "
f"size (fp16): {num_parameters * 2 / 1e6} MB"
)
print(f"Number of checkpointed parameters: {get_number_of_params(model)}")
print("Model initialized with pretrained weights. Training starting...")
if not args.no_grad_ckpt:
model.gradient_checkpointing_enable()
optimizer_cls = (
torch.optim.AdamW
if accelerator.state.deepspeed_plugin is None
or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
else DummyOptim
)
optimizer = optimizer_cls(
model.parameters(),
lr=lr,
betas=OPTIM_BETAS,
weight_decay=OPTIM_WEIGHT_DECAY,
eps=OPTIM_EPS,
)
# Instantiate scheduler
# Creates Dummy Scheduler if `scheduler` was specified in the config file or
# else, creates `args.lr_scheduler_type` Scheduler
# get train and valid dataset lengths
num_steps_per_epoch = math.ceil(train_ds_len / args.batch_size_per_device)
total_training_steps = (
num_steps_per_epoch * num_epochs // gradient_accumulation_steps
)
if (
accelerator.state.deepspeed_plugin is None
or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
):
lr_scheduler = get_linear_schedule_with_warmup(
optimizer=optimizer,
num_warmup_steps=NUM_WARMUP_STEPS * args.num_devices,
num_training_steps=total_training_steps * args.num_devices,
)
else:
lr_scheduler = DummyScheduler(
optimizer,
warmup_num_steps=NUM_WARMUP_STEPS * args.num_devices,
total_num_steps=total_training_steps * args.num_devices,
)
# Prepare everything
# There is no specific order to remember, we just need to unpack the objects in the
# same order we gave them to the prepare method.
s = time.time()
model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
print(f"Prepare done in {time.time() - s} seconds.")
# Now we train the model
if accelerator.is_main_process:
print("Starting training ...")
print("Number of batches on main process", train_ds_len // batch_size)
for epoch in range(num_epochs):
fwd_time_sum, bwd_time_sum, optim_step_time_sum = 0, 0, 0
s_epoch = time.time()
model.train()
loss_sum = torch.tensor(0.0).to(accelerator.device)
train_dataloader = train_ds.iter_torch_batches(
batch_size=batch_size,
collate_fn=collate_partial,
)
for step, batch in tqdm.tqdm(
enumerate(train_dataloader), total=train_ds_len // batch_size + 1
):
# We could avoid this line since we set the accelerator with
# `device_placement=True`.
with accelerator.accumulate(model):
s_fwd = time.time()
outputs = model(**batch)
loss = outputs.loss
loss_sum += loss.item()
e_fwd = time.time()
fwd_time = e_fwd - s_fwd
fwd_time_sum += fwd_time
s_bwd = time.time()
accelerator.backward(loss)
e_bwd = time.time()
bwd_time = e_bwd - s_bwd
bwd_time_sum += bwd_time
s_opt_step = time.time()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
e_opt_step = time.time()
optim_step_time_sum += e_opt_step - s_opt_step
if accelerator.is_main_process:
accelerator.print(
f"[epoch {epoch} step {step}] "
f"loss: {loss.item()} step-time: {e_opt_step - s_fwd}"
)
aggregated_loss = torch.mean(accelerator.gather(loss[None])).item()
if config["as_test"]:
break
# as long as this is not the last step report here
if step != (train_ds_len // batch_size - 1):
train.report(
{
"epoch": epoch,
"iteration": step,
"train_loss_batch": aggregated_loss,
"avg_train_loss_epoch": None,
"eval_loss": None,
"perplexity": None,
"num_iterations": step + 1,
"train_time_per_epoch": None,
"eval_time_per_epoch": None,
"fwd_time": fwd_time,
"bwd_time": bwd_time,
"avg_fwd_time_per_epoch": None,
"avg_bwd_time_per_epoch": None,
"learning_rate": lr_scheduler.get_lr()[0],
}
)
e_epoch = time.time()
accelerator.print("Train time per epoch: ", e_epoch - s_epoch)
eval_s_epoch = time.time()
print("Running evaluation ...")
perplex, eloss = evaluate(
model=model,
eval_ds=valid_ds,
accelerator=accelerator,
bsize=config["eval_batch_size"],
ds_kwargs={"collate_fn": collate_partial},
as_test=config["as_test"],
)
accelerator.print("Eval result loss", eloss)
accelerator.print("Eval perplex", perplex)
eval_e_epoch = time.time()
accelerator.print("Eval time per epoch: ", eval_e_epoch - eval_s_epoch)
accelerator.print("avg fwd time: ", fwd_time_sum / (step + 1))
accelerator.print("avg bwd time: ", bwd_time_sum / (step + 1))
accelerator.print("avg opt step time: ", optim_step_time_sum / (step + 1))
metrics = {
"epoch": epoch,
"iteration": step,
"train_loss_batch": aggregated_loss,
"avg_train_loss_epoch": loss_sum.item() / (step + 1),
"eval_loss": eloss,
"perplexity": perplex,
"num_iterations": step + 1,
"train_time_per_epoch": e_epoch - s_epoch,
"eval_time_per_epoch": eval_e_epoch - eval_s_epoch,
"fwd_time": fwd_time,
"bwd_time": bwd_time,
"avg_fwd_time_per_epoch": fwd_time_sum / (step + 1),
"avg_bwd_time_per_epoch": bwd_time_sum / (step + 1),
"learning_rate": lr_scheduler.get_lr()[0],
}
with tempfile.TemporaryDirectory(dir=args.output_dir) as temp_checkpoint_dir:
accelerator.print(f"Saving the model locally at {temp_checkpoint_dir}")
accelerator.wait_for_everyone()
checkpoint_save_start = time.perf_counter()
if accelerator.is_main_process:
print("Saving tokenizer and config.")
tokenizer.save_pretrained(temp_checkpoint_dir)
accelerator.wait_for_everyone()
# Checkpointing strategy 1: Distributed checkpointing
# This checkpointing method makes deepspeed checkpoints on each node
# and then Ray Train will aggregate them to a central s3 bucket.
# It should be done on all processes (not just the Rank 0)
# aggregate_on_rank_0 = False
# checkpoint_model(
# checkpoint_folder=tempdir,
# ckpt_id=epoch,
# model=model,
# epoch=epoch,
# last_global_step=step
# )
# Checkpointing strategy 2: Aggregate model on the rank 0 worker then upload
aggregate_on_rank_0 = True
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(
temp_checkpoint_dir,
is_main_process=accelerator.is_main_process,
save_function=accelerator.save,
safe_serialization=True,
state_dict=accelerator.get_state_dict(model),
)
accelerator.wait_for_everyone()
print("Checkpoint save time: ", time.perf_counter() - checkpoint_save_start)
checkpoint_upload_start = time.perf_counter()
# Create the checkpoint object to report to Ray Train and upload to storage.
# If we aggregated the model on rank 0, we only need to report
# the checkpoint from the rank 0 worker, since all other checkpoint
# directories are empty (`save_pretrained` was a noop for other workers).
if aggregate_on_rank_0:
checkpoint = (
Checkpoint.from_directory(temp_checkpoint_dir)
if accelerator.is_main_process
else None
)
else:
# Distributed checkpointing should upload shards from each worker.
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
# Note: After `train.report`, in the case of remote storage,
# the checkpoint directory will be uploaded to the remote storage.
train.report(metrics, checkpoint=checkpoint)
print(
"Checkpoint upload time: ",
time.perf_counter() - checkpoint_upload_start,
)
print(
"Total checkpointing time: ",
time.perf_counter() - checkpoint_save_start,
)
if perplex < args.stop_perplexity:
print(f"Perplexity reached {perplex} < {args.stop_perplexity}. Stopping.")
break
if config["as_test"]:
break
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of training script.")
parser.add_argument(
"--mx",
type=str,
default="bf16",
choices=["no", "fp16", "bf16", "fp8"],
help="Whether to use mixed precision. Choose"
"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
"and an NVIDIA Ampere GPU.",
)
parser.add_argument(
"--batch-size-per-device",
"-bs",
type=int,
default=16,
help="Batch size to use per device.",
)
parser.add_argument(
"--stop-perplexity",
default=0,
type=float,
help="Target perplexity to reach after which to stop training. Default is 0. "
"If 0, training will not stop on perplexity.",
)
parser.add_argument(
"--eval-batch-size-per-device",
type=int,
default=64,
help="Batch size to use per device (For evaluation).",
)
parser.add_argument(
"--num-devices", "-nd", type=int, default=4, help="Number of devices to use."
)
parser.add_argument(
"--grad_accum", type=int, default=1, help="Gradient accumulation steps."
)
parser.add_argument("--train_path", type=str, help="Path to training jsonl file")
parser.add_argument("--test_path", type=str, help="Path to testing jsonl file")
parser.add_argument(
"--special_token_path", type=str, help="Path to token json file"
)
parser.add_argument(
"--no-grad-ckpt",
action="store_true",
help="If passed, will not use gradient checkpointing.",
)
parser.add_argument("--output_dir", type=str, help="Path to output directory.")
parser.add_argument(
"--model_name", default="meta-llama/Llama-2-7b-chat-hf", type=str
)
parser.add_argument(
"--num-epochs", type=int, default=1, help="Number of epochs to train for."
)
parser.add_argument(
"--num-checkpoints-to-keep",
type=int,
help=(
"Number of checkpoints to keep, if None, all checkpoints will be kept, "
"if set to n>=1, the top n checkpoint with min. evaluation perplexity "
"will be kept."
),
default=None,
)
parser.add_argument("--lr", type=float, default=5e-6, help="Learning rate to use.")
parser.add_argument(
"--ctx-len",
type=int,
default=512,
help="Maximum context length for the model input sequences.",
)
parser.add_argument(
"--as-test",
action="store_true",
help="If passed, will run the script in test mode.",
)
parser.add_argument(
"--ds-config",
type=str,
default="./deepspeed_configs/zero_3_llama_2_7b.json",
help="Deepspeed config json to use.",
)
parser.add_argument(
"--lora",
action="store_true",
default=False,
help="If passed, will enable parameter efficient fine-tuning with LoRA ("
"https://arxiv.org/pdf/2106.09685.pdf).",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
if not args.output_dir:
raise ValueError("--output_dir must be specified")
# update the config with args so that we have access to them.
config = vars(args)
config.update(
**{
"lr": args.lr,
"num_epochs": args.num_epochs,
"seed": 42,
"batch_size": args.batch_size_per_device,
"gradient_accumulation_steps": args.grad_accum,
"model_name": args.model_name,
"block_size": args.ctx_len,
"eval_batch_size": args.eval_batch_size_per_device,
}
)
# Add LoRA config if needed
if args.lora:
with open("./lora_configs/lora.json", "r") as json_file:
lora_config = json.load(json_file)
config["lora_config"] = lora_config
# Add deepspeed plugin to the config
ds_plugin = DeepSpeedPlugin(hf_ds_config=config.get("ds_config"))
config.update(ds_plugin=ds_plugin)
ray.init(
runtime_env={
"env_vars": {"HF_HOME": "/mnt/local_storage/.cache/huggingface"},
"working_dir": ".",
}
)
# Read data
train_ds = ray.data.read_json(args.train_path)
if args.test_path is not None:
valid_ds = ray.data.read_json(args.test_path)
else:
valid_ds = None
# json file
with open(args.special_token_path, "r") as json_file:
special_tokens = json.load(json_file)["tokens"]
assert (
"ANYSCALE_ARTIFACT_STORAGE" in os.environ
), "ANYSCALE_ARTIFACT_STORAGE env var must be set!"
artifact_storage = os.environ["ANYSCALE_ARTIFACT_STORAGE"]
user_name = re.sub(r"\s+", "__", os.environ.get("ANYSCALE_USERNAME", "user"))
storage_path = (
f"{artifact_storage}/{user_name}/ft_llms_with_deepspeed/{args.model_name}"
)
trial_name = f"{args.model_name}".split("/")[-1]
if args.lora:
trial_name += "-lora"
trainer = TorchTrainer(
training_function,
train_loop_config={
"config": config,
"args": vars(args),
"special_tokens": special_tokens,
},
run_config=train.RunConfig(
storage_path=storage_path,
checkpoint_config=train.CheckpointConfig(
num_to_keep=args.num_checkpoints_to_keep,
checkpoint_score_attribute="perplexity",
checkpoint_score_order="min",
),
),
scaling_config=train.ScalingConfig(
num_workers=args.num_devices,
use_gpu=True,
resources_per_worker={"GPU": 1},
),
datasets={"train": train_ds, "valid": valid_ds},
dataset_config=ray.train.DataConfig(datasets_to_split=["train", "valid"]),
)
result: train.Result = trainer.fit()
# `best_checkpoints` are sorted in increasing score order.
# (Ex: in this case, negative perplexity, since we set `checkpoint_score_order=min`)
best_checkpoint, best_checkpoint_metrics = result.best_checkpoints[-1]
print("Results are stored at:")
print(result.path)
print("Best checkpoint is stored at:")
print(best_checkpoint)
print(f"With perplexity: {best_checkpoint_metrics['perplexity']}")
if __name__ == "__main__":
main()
doc/source/templates/04_finetuning_llms_with_deepspeed/lora_configs/lora.json
+11
View File
@@ -0,0 +1,11 @@
{
"r": 8,
"lora_alpha": 16,
"lora_dropout": 0.05,
"target_modules": ["gate_proj", "up_proj", "down_proj"],
"task_type": "CAUSAL_LM",
"modules_to_save": [],
"bias": "none",
"fan_in_fan_out": false,
"init_lora_weights": true
}
doc/source/templates/04_finetuning_llms_with_deepspeed/merge_lora_weights.py
+157
View File
@@ -0,0 +1,157 @@
"""
This script merges the weights of a LoRA checkpoint with the base model weights
to create a single model that can be used for model evaluation.
"""
import torch
import argparse
import time
import peft
from pathlib import Path
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
StoppingCriteriaList,
)
from utils import download_model, get_mirror_link, get_checkpoint_and_refs_dir
# In addition to merging the lora weights, you can also formulate a prompt for the
# model here to quickly test it after merging
TEST_EVAL = False
TEST_PROMPT = (
"<START_Q>Natalia sold clips to 48 of her friends in April, and then "
"she sold half as many clips in May. How many clips did Natalia sell "
"altogether in April and May?<END_Q><START_A>"
)
STOP_TOKEN = "<END_A>"
def parse_args():
parser = argparse.ArgumentParser(description="Simple example of training script.")
parser.add_argument(
"--output-path",
type=str,
help="Path to output directory. Defaults to the original checkpoint directory.",
required=True,
)
parser.add_argument("--model-name", required=True, type=str, help="7b, 13b or 70b.")
parser.add_argument(
"--checkpoint",
type=str,
required=True,
help="Path to checkpoint containing the LoRA weights.",
)
args = parser.parse_args()
return args
def test_eval(model, tokenizer):
"""Query the model with a single prompt to sanity check it."""
print("Starting model evaluation...")
model.eval()
model.to("cuda")
print("Prompting model with prompt : ", TEST_PROMPT)
input_ids = tokenizer(TEST_PROMPT, return_tensors="pt")["input_ids"].to("cuda")
stop_token_embedding = tokenizer(
STOP_TOKEN, return_tensors="pt", add_special_tokens=False
)["input_ids"].to("cuda")
def custom_stopping_criteria(embeddings, *args, **kwargs) -> bool:
return stop_token_embedding in embeddings
stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])
with torch.no_grad():
generation_output = model.generate(
input_ids=input_ids,
output_scores=True,
max_new_tokens=500,
stopping_criteria=stopping_criteria,
)
decoded = tokenizer.batch_decode(generation_output)
print("Outputs: ", decoded)
def main():
args = parse_args()
# Sanity checks
if not Path(args.checkpoint).exists():
raise ValueError(f"Checkpoint {args.checkpoint} does not exist.")
if not args.output_path:
args.output_path = Path(args.checkpoint) / "merged_model"
print(f"Output path not specified. Using {args.output_path}")
Path(args.output_path).mkdir(parents=True, exist_ok=True)
# Load orignal model
s = time.time()
model_id = f"meta-llama/Llama-2-{args.model_name}-hf"
s3_bucket = get_mirror_link(model_id)
ckpt_path, _ = get_checkpoint_and_refs_dir(model_id=model_id, bucket_uri=s3_bucket)
print(f"Downloading original model {model_id} from {s3_bucket} to {ckpt_path} ...")
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, legacy=True)
tokenizer.save_pretrained(Path(args.output_path))
print(f"Saved tokenizer to {args.output_path}")
download_model(
model_id=model_id,
bucket_uri=s3_bucket,
s3_sync_args=["--no-sign-request"],
)
print(f"Downloading to {ckpt_path} finished after {time.time() - s} seconds.")
print(f"Loading original model from {ckpt_path} ...")
s2 = time.time()
model = AutoModelForCausalLM.from_pretrained(
ckpt_path,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
use_cache=False,
)
model.resize_token_embeddings(len(tokenizer))
print(f"Done downloading and loading model after {time.time() - s2} seconds.")
print("Loading and merging peft weights...")
s3 = time.time()
# Load LoRA weights
model: peft.PeftModel = peft.PeftModel.from_pretrained(
model=model,
model_id=args.checkpoint,
)
# Merge weights and save
model = model.merge_and_unload()
output_path = Path(args.output_path)
model.save_pretrained(output_path, safe_serialization=True)
model.config.save_pretrained(output_path)
print(f"Saved merged model to {args.output_path} after {time.time() - s3} seconds.")
print(f"This script took {time.time() - s} seconds to execute.")
if TEST_EVAL:
test_eval(model, tokenizer)
if __name__ == "__main__":
main()
doc/source/templates/04_finetuning_llms_with_deepspeed/run_llama_ft.sh
Executable
+98
View File
@@ -0,0 +1,98 @@
#!/bin/bash
# Function to check if data directory exists, if not, run create_dataset.py
check_and_create_dataset() {
local data_dir=$1
if [ ! -d "${data_dir}" ]; then
echo "Data directory not found. Creating dataset..."
if ! python create_dataset.py; then
echo "Failed to create dataset. Exiting..."
exit 1
fi
fi
}
# Function to fine-tune the model
fine_tune() {
local bs=$1
local nd=$2
local model_name=$3
local output_dir=$4
local ds_config=$5
local train_path=$6
local test_path=$7
local token_path=$8
local params=("${@:9}")
echo "Fine-tuning model..."
if ! python finetune_hf_llm.py \
-bs "${bs}" \
-nd "${nd}" \
--model_name "${model_name}" \
--output_dir "${output_dir}" \
--ds-config "${ds_config}" \
--train_path "${train_path}" \
--test_path "${test_path}" \
--special_token_path "${token_path}" \
--num-checkpoints-to-keep 1 \
--num-epochs 3 \
"${params[@]}"; then
echo "Failed to fine-tune the model. Exiting..."
exit 1
fi
}
# Variables for cleaner handling
BASE_DIR="/mnt/local_storage"
DATA_DIR="./data"
TRAIN_PATH="${DATA_DIR}/train.jsonl"
TEST_PATH="${DATA_DIR}/test.jsonl"
TOKEN_PATH="${DATA_DIR}/tokens.json"
# Parse arguments
SIZE=""
for arg in "$@"
do
key=${arg%%=*}
value=${arg#*=}
if [[ "$key" == "--size" ]]; then
SIZE=${value};
elif [[ "$arg" == "--as-test" ]]; then
params+=("--as-test");
elif [[ "$arg" == "--lora" ]]; then
params+=("--lora");
# Lora usually requires a lower learning rate
params+=("--lr");
params+=("1e-4");
fi
done
# Batch size and node count
case $SIZE in
"7b")
BS=16
ND=16
;;
"13b")
BS=16
ND=16
;;
"70b")
BS=8
ND=32
;;
*)
echo "Invalid size: ${SIZE}"
exit 1
;;
esac
# Model related variables
MODEL_ID="meta-llama/Llama-2-${SIZE}-hf"
CONFIG_DIR="./deepspeed_configs/zero_3_llama_2_${SIZE}.json"
check_and_create_dataset "${DATA_DIR}"
fine_tune "$BS" "$ND" "$MODEL_ID" "$BASE_DIR" "$CONFIG_DIR" "$TRAIN_PATH" "$TEST_PATH" "$TOKEN_PATH" "${params[@]}"
echo "Process completed."
doc/source/templates/04_finetuning_llms_with_deepspeed/utils.py
+85
View File
@@ -0,0 +1,85 @@
from typing import List, Optional
import os
import subprocess
import logging
logger = logging.getLogger(__name__)
def get_hash_from_bucket(
bucket_uri: str, s3_sync_args: Optional[List[str]] = None
) -> str:
s3_sync_args = s3_sync_args or []
subprocess.run(
["aws", "s3", "cp", "--quiet"]
+ s3_sync_args
+ [os.path.join(bucket_uri, "refs", "main"), "."],
check=True,
)
with open(os.path.join(".", "main"), "r") as f:
f_hash = f.read().strip()
return f_hash
def get_checkpoint_and_refs_dir(
model_id: str,
bucket_uri: str,
s3_sync_args: Optional[List[str]] = None,
mkdir: bool = False,
) -> str:
from transformers.utils.hub import TRANSFORMERS_CACHE
f_hash = get_hash_from_bucket(bucket_uri, s3_sync_args)
path = os.path.join(TRANSFORMERS_CACHE, f"models--{model_id.replace('/', '--')}")
refs_dir = os.path.join(path, "refs")
checkpoint_dir = os.path.join(path, "snapshots", f_hash)
if mkdir:
os.makedirs(refs_dir, exist_ok=True)
os.makedirs(checkpoint_dir, exist_ok=True)
return checkpoint_dir, refs_dir
def get_download_path(model_id: str):
from transformers.utils.hub import TRANSFORMERS_CACHE
path = os.path.join(TRANSFORMERS_CACHE, f"models--{model_id.replace('/', '--')}")
return path
def download_model(
model_id: str,
bucket_uri: str,
s3_sync_args: Optional[List[str]] = None,
tokenizer_only: bool = False,
) -> None:
"""
Download a model from an S3 bucket and save it in TRANSFORMERS_CACHE for
seamless interoperability with Hugging Face's Transformers library.
The downloaded model may have a 'hash' file containing the commit hash corresponding
to the commit on Hugging Face Hub.
"""
s3_sync_args = s3_sync_args or []
path = get_download_path(model_id)
cmd = (
["aws", "s3", "sync"]
+ s3_sync_args
+ (["--exclude", "*", "--include", "*token*"] if tokenizer_only else [])
+ [bucket_uri, path]
)
print(f"RUN({cmd})")
subprocess.run(cmd)
print("done")
def get_mirror_link(model_id: str) -> str:
return f"s3://llama-2-weights/models--{model_id.replace('/', '--')}"
doc/source/templates/05_dreambooth_finetuning/README.md
+99
View File
@@ -0,0 +1,99 @@
# DreamBooth fine-tuning of Stable Diffusion with Ray Train
| Template Specification | Description |
| ---------------------- | ----------- |
| Summary | This example shows how to do [DreamBooth fine-tuning](https://dreambooth.github.io/) of a Stable Diffusion model using Ray Train for data-parallel training with many workers and Ray Data for data ingestion. Use one of the provided datasets, or supply your own photos. By the end of this example, you'll be able to generate images of your subject in a variety of situations, just by feeding in a text prompt! |
| Time to Run | ~10-15 minutes to generate a regularization dataset and fine-tune the model on photos of your subject. |
| Minimum Compute Requirements | At least 1 GPUs, where each GPU has >= 24GB GRAM. The default is 1 node with 4 GPUS: A10G GPU (AWS) or L4 GPU (GCE). |
| Cluster Environment | This template uses a Docker image built on top of the latest Anyscale-provided Ray image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview?utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning). See the appendix below for more details. |
![Dreambooth fine-tuning sample results](https://raw.githubusercontent.com/ray-project/ray/workspace_templates_2.6.1/doc/source/templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_example.png)
## Run the example
This README will only contain minimal instructions on running this example on Anyscale. See [the guide on the Ray documentation](https://docs.ray.io/en/latest/train/examples/pytorch/dreambooth_finetuning.html) for a step-by-step walkthrough of the training code.
You can get started fine-tuning on a sample dog dataset with default settings with the following commands:
```bash
chmod +x ./dreambooth_run.sh
./dreambooth_run.sh
```
## Customizing the example
Here are a few modifications to the `dreambooth_run.sh` script that you may want to make:
1. The image dataset of your subject. This example provides two sample datasets, but you can also supply your own directory of 4-5 images, as well as the general class your subject falls under. For example, the dog dataset contains images of one particular puppy, and the general class this subject falls under is `dog`.
- Modify the `$CLASS_NAME` and `$INSTANCE_DIR` environment variables.
2. The `$DATA_PREFIX` that the pre-trained model is downloaded to. This directory is also where the training dataset and the fine-tuned model checkpoint are written at the end of training.
- If you add more worker nodes to the cluster, you should `$DATA_PREFIX` to a shared NFS filesystem such as `/mnt/cluster_storage`. See [this doc](https://docs.anyscale.com/develop/workspaces/storage#storage-shared-across-nodes?utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning) for all the options.
- Note that each run of the script will overwrite the fine-tuned model checkpoint from the previous run, so consider changing the `$DATA_PREFIX` environment variable on each run if you don't want to lose the models/data of previous runs.
3. The `$NUM_WORKERS` variable sets the number of data-parallel workers used during fine-tuning. The default is 2 workers (2 workers, each using 1 GPU), and you should increase this number if you add more GPU worker nodes to the cluster.
4. Setting `--num_epochs` and `--max_train_steps` determines the number of fine-tuning steps to take.
- Depending on the batch size and number of data-parallel workers, one epoch will run for a certain number of steps. The run will terminate when one of these values (epoch vs. total number of steps) is reached.
5. `generate.py` is used to generate stable diffusion images after loading the model from a checkpoint. You should modify the prompt at the end to be something more interesting, rather than just a photo of your subject.
6. If you want to launch another fine-tuning run, you may want to run *only* the `python train.py ...` command. Running the bash script will start from the beginning (generating another regularization dataset).
7. Use the following command for LoRA fine-tuning.
```bash
python train.py \
--model_dir=$ORIG_MODEL_PATH \
--output_dir=$TUNED_MODEL_DIR \
--instance_images_dir=$IMAGES_OWN_DIR \
--instance_prompt="photo of $UNIQUE_TOKEN $CLASS_NAME" \
--class_images_dir=$IMAGES_REG_DIR \
--class_prompt="photo of a $CLASS_NAME" \
--train_batch_size=2 \
--lr=1e-4 \ # Note a much higher learning rate here!
--num_epochs=10 \
--max_train_steps=400 \
--num_workers $NUM_WORKERS
--use_lora
```
## Interact with the fine-tuned model
### Generate images with a script
Use the `generate.py` script to generate images with a prompt. Replace the variables with the values that you used in the fine-tuning script. See `run_model_flags` in `flags.py` for a full list of available command line arguments to pass to the script.
```bash
python generate.py \
--model_dir=$TUNED_MODEL_DIR \
--output_dir=$IMAGES_NEW_DIR \
--prompts="photo of a $UNIQUE_TOKEN $CLASS_NAME" \
--num_samples_per_prompt=5
```
To generate images using LoRA fine-tuned model:
```bash
python generate.py \
--model_dir=$ORIG_MODEL_PATH \
--lora_weights_dir=$TUNED_MODEL_DIR \
--output_dir=$IMAGES_NEW_DIR \
--prompts="photo of a $UNIQUE_TOKEN $CLASS_NAME" \
--num_samples_per_prompt=5
```
### Generate images interactively in a notebook
See the `playground.ipynb` notebook for a more interactive way to generate images with the fine-tuned model. Click on the Jupyter icon on the workspace page and open the notebook. *Note: The widgets in this notebook don't work in VS Code, so please use Jupyter!*
## Appendix
### Advanced: Build off of this template's cluster environment
#### Option 1: Build a new cluster environment on Anyscale
The `dreambooth/requirements.txt` file lists the requirements. Feel free to modify this file to include more requirements, then follow [this guide](https://docs.anyscale.com/configure/dependency-management/cluster-environments#creating-a-cluster-environment?utm_source=ray_docs&utm_medium=docs&utm_campaign=dreambooth_finetuning) to create a new cluster environment with the `anyscale` CLI . Paste the requirements into the cluster environment YAML.
Finally, update the workspace's cluster environment to this environment after it's done building.
#### Option 2: Build a new docker image with your own infrastructure
Use the following `docker pull` command if you want to manually build a new Docker image based off of this one.
```bash
docker pull us-docker.pkg.dev/anyscale-workspace-templates/workspace-templates/dreambooth-finetuning:latest
```
doc/source/templates/05_dreambooth_finetuning/configs/Dockerfile
+9
View File
@@ -0,0 +1,9 @@
# Run `docker build` with this from the 05_dreambooth_finetuning directory
FROM anyscale/ray:latest-py39-cu118
COPY dreambooth/requirements.txt ./
RUN pip install --no-cache-dir -U -r requirements.txt
RUN echo "Testing Ray Import..." && python -c "import ray"
RUN ray --version
doc/source/templates/05_dreambooth_finetuning/configs/compute/aws.yaml
+7
View File
@@ -0,0 +1,7 @@
head_node_type:
name: head_node_type
instance_type: g5.12xlarge
worker_node_types: []
max_workers: 0
doc/source/templates/05_dreambooth_finetuning/configs/compute/gce.yaml
+7
View File
@@ -0,0 +1,7 @@
head_node_type:
name: head_node_type
instance_type: g2-standard-48-nvidia-l4-4
worker_node_types: []
max_workers: 0
doc/source/templates/05_dreambooth_finetuning/dreambooth/cache_model.py
+20
View File
@@ -0,0 +1,20 @@
# Cache model files to a local directory
import os
from huggingface_hub import snapshot_download
from flags import cache_model_flags
def cache(args):
os.makedirs(args.model_dir, exist_ok=True)
snapshot_download(
repo_id=args.model_name, revision=args.revision, cache_dir=args.model_dir
)
if __name__ == "__main__":
args = cache_model_flags().parse_args()
cache(args)
doc/source/templates/05_dreambooth_finetuning/dreambooth/dataset.py
+162
View File
@@ -0,0 +1,162 @@
from typing import Dict
import numpy as np
import pandas as pd
import torch
from ray.data import read_images
from torchvision import transforms
from transformers import AutoTokenizer
def get_train_dataset(args, image_resolution=512):
"""Build a Dataset for fine-tuning DreamBooth model."""
# Load a directory of images as a Ray Dataset
instance_dataset = read_images(args.instance_images_dir)
class_dataset = read_images(args.class_images_dir)
# We now duplicate the instance images multiple times to make the
# two sets contain exactly the same number of images.
# This is so we can zip them up during training to compute the
# prior preserving loss in one pass.
#
# Example: If we have 200 class images (for regularization) and 4 instance
# images of our subject, then we'll duplicate the instance images 50 times
# so that our dataset looks like:
#
# instance_image_0, class_image_0
# instance_image_1, class_image_1
# instance_image_2, class_image_2
# instance_image_3, class_image_3
# instance_image_0, class_image_4
# instance_image_1, class_image_5
# ...
dup_times = class_dataset.count() // instance_dataset.count()
instance_dataset = instance_dataset.map_batches(
lambda df: pd.concat([df] * dup_times), batch_format="pandas"
)
# Load tokenizer for tokenizing the image prompts.
tokenizer = AutoTokenizer.from_pretrained(
pretrained_model_name_or_path=args.model_dir,
subfolder="tokenizer",
)
def _tokenize(prompt):
return tokenizer(
prompt,
truncation=True,
padding="max_length",
max_length=tokenizer.model_max_length,
return_tensors="pt",
).input_ids.numpy()
# Get the token ids for both prompts.
class_prompt_ids = _tokenize(args.class_prompt)[0]
instance_prompt_ids = _tokenize(args.instance_prompt)[0]
# START: image preprocessing
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Resize(
image_resolution,
interpolation=transforms.InterpolationMode.BILINEAR,
antialias=True,
),
transforms.RandomCrop(image_resolution),
# use the appropriate mean and std for your dataset
transforms.Normalize([0.5], [0.5]),
]
)
def transform_image(
batch: Dict[str, np.ndarray], output_column_name: str
) -> Dict[str, np.ndarray]:
transformed_tensors = [transform(image).numpy() for image in batch["image"]]
batch[output_column_name] = transformed_tensors
return batch
# END: image preprocessing
# START: Apply preprocessing steps as Ray Dataset operations
# For each dataset:
# - perform image preprocessing
# - drop the original image column
# - add a new column with the tokenized prompts
instance_dataset = (
instance_dataset.map_batches(
transform_image, fn_kwargs={"output_column_name": "instance_image"}
)
.drop_columns(["image"])
.add_column(
"instance_prompt_ids", lambda df: pd.Series([instance_prompt_ids] * len(df))
)
)
# END: Apply preprocessing steps as Ray Dataset operations
class_dataset = (
class_dataset.map_batches(
transform_image, fn_kwargs={"output_column_name": "class_image"}
)
.drop_columns(["image"])
.add_column(
"class_prompt_ids", lambda df: pd.Series([class_prompt_ids] * len(df))
)
)
# --- Ray Data
# We may have too many duplicates of the instance images, so limit the
# dataset size so that len(instance_dataset) == len(class_dataset)
final_size = min(instance_dataset.count(), class_dataset.count())
# Now, zip the images up.
train_dataset = (
instance_dataset.limit(final_size)
.repartition(final_size)
.zip(class_dataset.limit(final_size).repartition(final_size))
)
print("Training dataset schema after pre-processing:")
print(train_dataset.schema())
return train_dataset.random_shuffle()
def collate(batch, dtype):
"""Build Torch training batch.
B = batch size
(C, W, H) = (channels, width, height)
L = max length in tokens of the text guidance input
Input batch schema (see `get_train_dataset` on how this was setup):
instance_images: (B, C, W, H)
class_images: (B, C, W, H)
instance_prompt_ids: (B, L)
class_prompt_ids: (B, L)
Output batch schema:
images: (2 * B, C, W, H)
All instance images in the batch come before the class images:
[instance_images[0], ..., instance_images[B-1], class_images[0], ...]
prompt_ids: (2 * B, L)
Prompt IDs are ordered the same way as the images.
During training, a batch will be chunked into 2 sub-batches for
prior preserving loss calculation.
"""
images = torch.cat([batch["instance_image"], batch["class_image"]], dim=0)
images = images.to(memory_format=torch.contiguous_format).to(dtype)
batch_size = len(batch["instance_prompt_ids"])
prompt_ids = torch.cat(
[batch["instance_prompt_ids"], batch["class_prompt_ids"]], dim=0
).reshape(batch_size * 2, -1)
return {
"images": images,
"prompt_ids": prompt_ids, # token ids should stay int.
}
doc/source/templates/05_dreambooth_finetuning/dreambooth/download_example_dataset.py
+14
View File
@@ -0,0 +1,14 @@
from huggingface_hub import snapshot_download
import os
import sys
local_dir = sys.argv[1]
os.makedirs(local_dir, exist_ok=True)
snapshot_download(
"diffusers/dog-example",
local_dir=local_dir,
repo_type="dataset",
ignore_patterns=".gitattributes",
)
doc/source/templates/05_dreambooth_finetuning/dreambooth/flags.py
+158
View File
@@ -0,0 +1,158 @@
import argparse
def train_arguments():
"""Commandline arguments for running DreamBooth training script."""
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_dir",
type=str,
default=None,
required=True,
help="Path to a pretrained huggingface Stable Diffusion model.",
)
parser.add_argument(
"--output_dir",
type=str,
default=None,
required=True,
help="Directory where trained models or LoRA weights are saved.",
)
parser.add_argument(
"--use_lora", default=False, action="store_true", help="Use LoRA."
)
parser.add_argument(
"--instance_images_dir",
type=str,
default=None,
required=True,
help=(
"Directory where a few images of the instance to be fine tuned "
"into the model are saved."
),
)
parser.add_argument(
"--instance_prompt",
type=str,
default=None,
required=True,
help=("Prompt for creating the instance images."),
)
parser.add_argument(
"--class_images_dir",
type=str,
default=None,
required=True,
help=(
"Directory where images of similar objects for preserving "
"model priors are saved."
),
)
parser.add_argument(
"--class_prompt",
type=str,
default=None,
required=True,
help=("Prompt for creating the class images."),
)
parser.add_argument(
"--train_batch_size", type=int, default=1, help="Train batch size."
)
parser.add_argument("--lr", type=float, default=5e-6, help="Train learning rate.")
parser.add_argument(
"--num_epochs", type=int, default=4, help="Number of epochs to train."
)
parser.add_argument(
"--max_train_steps",
type=int,
default=800,
help="Maximum number of fine-tuning update steps to take.",
)
parser.add_argument(
"--prior_loss_weight",
type=float,
default=1.0,
help="The weight for prior preservation loss.",
)
parser.add_argument(
"--max_grad_norm", type=float, default=1.0, help="Maximum gradient norm."
)
parser.add_argument("--num_workers", type=int, default=2, help="Number of workers.")
return parser
def cache_model_flags():
"""Commandline arguments for running local model caching script."""
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_dir",
type=str,
default=None,
required=True,
help="Directory to write the cached model files.",
)
parser.add_argument(
"--model_name",
type=str,
default="CompVis/stable-diffusion-v1-4",
help="Name of the huggingface model.",
)
parser.add_argument(
"--revision",
type=str,
default="3857c45b7d4e78b3ba0f39d4d7f50a2a05aa23d4",
help="Revision of the huggingface model repo to cache.",
)
return parser
def run_model_flags():
"""Commandline arguments for running a tuned DreamBooth model."""
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_dir",
type=str,
default=None,
required=True,
help="Directory of the tuned model files.",
)
parser.add_argument(
"--output_dir",
type=str,
default=None,
required=True,
help="Directory to save the generated images.",
)
parser.add_argument(
"--prompts",
type=str,
default=None,
required=True,
help="Comma separated prompt strings for generating the images.",
)
parser.add_argument(
"--num_samples_per_prompt",
type=int,
default=1,
help="Number of images to generate for each prompt.",
)
parser.add_argument(
"--use_ray_data",
default=False,
action="store_true",
help=(
"Enable using Ray Data to use multiple GPU workers to perform inference."
),
)
parser.add_argument(
"--lora_weights_dir",
default=None,
help=("The directory where `pytorch_lora_weights.bin` is stored."),
)
return parser
doc/source/templates/05_dreambooth_finetuning/dreambooth/generate.py
+80
View File
@@ -0,0 +1,80 @@
import hashlib
from os import path
import time
import torch
import ray
from flags import run_model_flags
from generate_utils import get_pipeline
def run(args):
class StableDiffusionCallable:
def __init__(self, model_dir, output_dir, lora_weights_dir=None):
print(f"Loading model from {model_dir}")
self.pipeline = get_pipeline(model_dir, lora_weights_dir)
self.pipeline.set_progress_bar_config(disable=True)
if torch.cuda.is_available():
self.pipeline.to("cuda")
self.output_dir = output_dir
def __call__(self, batch):
filenames = []
for i, prompt in zip(batch["idx"], batch["prompt"]):
# Generate 1 image at a time to reduce memory consumption.
for image in self.pipeline(prompt).images:
hash_image = hashlib.sha256(image.tobytes()).hexdigest()
image_filename = path.join(self.output_dir, f"{i}-{hash_image}.jpg")
image.save(image_filename)
print(f"Saved {image_filename}")
filenames.append(image_filename)
return {"filename": filenames}
prompts = args.prompts.split(",")
start_time = time.time()
num_samples = len(prompts) * args.num_samples_per_prompt
if args.use_ray_data:
# Use Ray Data to perform batch inference to generate many images in parallel
prompts_with_idxs = []
for prompt in prompts:
prompts_with_idxs.extend(
[
{"idx": i, "prompt": prompt}
for i in range(args.num_samples_per_prompt)
]
)
prompt_ds = ray.data.from_items(prompts_with_idxs)
num_workers = 4
# Run the batch inference by consuming output with `take_all`.
prompt_ds.map_batches(
StableDiffusionCallable,
compute=ray.data.ActorPoolStrategy(size=num_workers),
fn_constructor_args=(args.model_dir, args.output_dir),
num_gpus=1,
batch_size=num_samples // num_workers,
).take_all()
else:
# Generate images one by one
stable_diffusion_predictor = StableDiffusionCallable(
args.model_dir, args.output_dir, args.lora_weights_dir
)
for prompt in prompts:
for i in range(args.num_samples_per_prompt):
stable_diffusion_predictor({"idx": [i], "prompt": [prompt]})
elapsed = time.time() - start_time
print(
f"Generated and saved {num_samples} images to {args.output_dir} in "
f"{elapsed} seconds."
)
if __name__ == "__main__":
args = run_model_flags().parse_args()
run(args)
doc/source/templates/05_dreambooth_finetuning/dreambooth/generate_utils.py
+26
View File
@@ -0,0 +1,26 @@
from diffusers import DiffusionPipeline
from diffusers.loaders import LoraLoaderMixin
import torch
def load_lora_weights(unet, text_encoder, input_dir):
lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
LoraLoaderMixin.load_lora_into_unet(
lora_state_dict, network_alphas=network_alphas, unet=unet
)
LoraLoaderMixin.load_lora_into_text_encoder(
lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder
)
return unet, text_encoder
def get_pipeline(model_dir, lora_weights_dir=None):
pipeline = DiffusionPipeline.from_pretrained(model_dir, torch_dtype=torch.float16)
if lora_weights_dir:
unet = pipeline.unet
text_encoder = pipeline.text_encoder
print(f"Loading LoRA weights from {lora_weights_dir}")
unet, text_encoder = load_lora_weights(unet, text_encoder, lora_weights_dir)
pipeline.unet = unet
pipeline.text_encoder = text_encoder
return pipeline
doc/source/templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 504 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/dreambooth_training.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 14 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/lego-car/1.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 109 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/lego-car/2.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 108 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/lego-car/3.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 112 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/lego-car/4.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 97 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/images/lego-car/5.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 120 KiB

doc/source/templates/05_dreambooth_finetuning/dreambooth/requirements.txt
+12
View File
@@ -0,0 +1,12 @@
accelerate==0.20.3
bitsandbytes==0.39.1
diffusers==0.19.3
flax==0.6.11
ipywidgets
huggingface_hub==0.19.4
jax==0.4.17
jaxlib==0.4.17
numpy==1.24.4
torch==2.0.1
torchvision==0.15.2
transformers==4.30.2
doc/source/templates/05_dreambooth_finetuning/dreambooth/train.py
+352
View File
@@ -0,0 +1,352 @@
from typing import Dict
import itertools
from diffusers import (
AutoencoderKL,
DDPMScheduler,
DiffusionPipeline,
UNet2DConditionModel,
)
# LoRA related imports begin ##
from diffusers.loaders import (
LoraLoaderMixin,
text_encoder_lora_state_dict,
)
from diffusers.models.attention_processor import (
AttnAddedKVProcessor,
AttnAddedKVProcessor2_0,
LoRAAttnAddedKVProcessor,
LoRAAttnProcessor,
LoRAAttnProcessor2_0,
SlicedAttnAddedKVProcessor,
)
# LoRA related imports end ##
from diffusers.utils.import_utils import is_xformers_available
from ray.train import ScalingConfig
from ray import train
from ray.train.torch import TorchTrainer
import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from transformers import CLIPTextModel
from dataset import collate, get_train_dataset
from flags import train_arguments
LORA_RANK = 4
def prior_preserving_loss(model_pred, target, weight):
# Chunk the noise and model_pred into two parts and compute
# the loss on each part separately.
model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
target, target_prior = torch.chunk(target, 2, dim=0)
# Compute instance loss
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
# Compute prior loss
prior_loss = F.mse_loss(
model_pred_prior.float(), target_prior.float(), reduction="mean"
)
# Add the prior loss to the instance loss.
return loss + weight * prior_loss
def get_target(scheduler, noise, latents, timesteps):
"""Get the target for loss depending on the prediction type."""
pred_type = scheduler.config.prediction_type
if pred_type == "epsilon":
return noise
if pred_type == "v_prediction":
return scheduler.get_velocity(latents, noise, timesteps)
raise ValueError(f"Unknown prediction type {pred_type}")
def add_lora_layers(unet, text_encoder):
"""Add LoRA layers for unet and text encoder.
`unet` and `text_encoder` will be modified in place.
Returns:
The LoRA parameters for unet and text encoder correspondingly.
"""
unet_lora_attn_procs = {}
unet_lora_parameters = []
for name, attn_processor in unet.attn_processors.items():
cross_attention_dim = (
None
if name.endswith("attn1.processor")
else unet.config.cross_attention_dim
)
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
if isinstance(
attn_processor,
(AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0),
):
lora_attn_processor_class = LoRAAttnAddedKVProcessor
else:
lora_attn_processor_class = (
LoRAAttnProcessor2_0
if hasattr(F, "scaled_dot_product_attention")
else LoRAAttnProcessor
)
module = lora_attn_processor_class(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
rank=LORA_RANK,
)
unet_lora_attn_procs[name] = module
unet_lora_parameters.extend(module.parameters())
unet.set_attn_processor(unet_lora_attn_procs)
text_lora_parameters = LoraLoaderMixin._modify_text_encoder(
text_encoder, dtype=torch.float32, rank=LORA_RANK
)
return unet_lora_parameters, text_lora_parameters
def load_models(config):
"""Load pre-trained Stable Diffusion models."""
# Load all models in bfloat16 to save GRAM.
# For models that are only used for inferencing,
# full precision is also not required.
dtype = torch.bfloat16
text_encoder = CLIPTextModel.from_pretrained(
args.model_dir,
subfolder="text_encoder",
torch_dtype=dtype,
)
noise_scheduler = DDPMScheduler.from_pretrained(
config["model_dir"],
subfolder="scheduler",
torch_dtype=dtype,
)
# VAE is only used for inference, keeping weights in full precision is not required.
vae = AutoencoderKL.from_pretrained(
config["model_dir"],
subfolder="vae",
torch_dtype=dtype,
)
# We are not training VAE part of the model.
vae.requires_grad_(False)
# Convert unet to bf16 to save GRAM.
unet = UNet2DConditionModel.from_pretrained(
config["model_dir"],
subfolder="unet",
torch_dtype=dtype,
)
if is_xformers_available():
unet.enable_xformers_memory_efficient_attention()
if not config["use_lora"]:
unet_trainable_parameters = unet.parameters()
text_trainable_parameters = text_encoder.parameters()
else:
text_encoder.requires_grad_(False)
unet.requires_grad_(False)
unet_trainable_parameters, text_trainable_parameters = add_lora_layers(
unet, text_encoder
)
text_encoder.train()
unet.train()
torch.cuda.empty_cache()
return (
text_encoder,
noise_scheduler,
vae,
unet,
unet_trainable_parameters,
text_trainable_parameters,
)
def train_fn(config):
# Load pre-trained models.
(
text_encoder,
noise_scheduler,
vae,
unet,
unet_trainable_parameters,
text_trainable_parameters,
) = load_models(config)
text_encoder = train.torch.prepare_model(text_encoder)
unet = train.torch.prepare_model(unet)
# manually move to device as `prepare_model` can't be used on
# non-training models.
vae = vae.to(train.torch.get_device())
# Use the regular AdamW optimizer to work with bfloat16 weights.
optimizer = torch.optim.AdamW(
itertools.chain(unet_trainable_parameters, text_trainable_parameters),
lr=config["lr"],
)
train_dataset = train.get_dataset_shard("train")
# Train!
num_train_epochs = config["num_epochs"]
print(f"Running {num_train_epochs} epochs.")
global_step = 0
for _ in range(num_train_epochs):
if global_step >= config["max_train_steps"]:
print(f"Stopping training after reaching {global_step} steps...")
break
for _, batch in enumerate(
train_dataset.iter_torch_batches(
batch_size=config["train_batch_size"],
device=train.torch.get_device(),
)
):
batch = collate(batch, torch.bfloat16)
optimizer.zero_grad()
# Convert images to latent space
latents = vae.encode(batch["images"]).latent_dist.sample() * 0.18215
# Sample noise that we'll add to the latents
noise = torch.randn_like(latents)
bsz = latents.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(
0,
noise_scheduler.config.num_train_timesteps,
(bsz,),
device=latents.device,
)
timesteps = timesteps.long()
# Add noise to the latents according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
# Get the text embedding for conditioning
encoder_hidden_states = text_encoder(batch["prompt_ids"])[0]
# Predict the noise residual.
model_pred = unet(
noisy_latents.to(train.torch.get_device()),
timesteps.to(train.torch.get_device()),
encoder_hidden_states.to(train.torch.get_device()),
).sample
target = get_target(noise_scheduler, noise, latents, timesteps)
loss = prior_preserving_loss(
model_pred, target, config["prior_loss_weight"]
)
loss.backward()
# Gradient clipping before optimizer stepping.
clip_grad_norm_(
itertools.chain(unet_trainable_parameters, text_trainable_parameters),
config["max_grad_norm"],
)
optimizer.step() # Step all optimizers.
global_step += 1
results = {
"step": global_step,
"loss": loss.detach().item(),
}
train.report(results)
if global_step >= config["max_train_steps"]:
break
# END: Training loop
# Create pipeline using the trained modules and save it.
if train.get_context().get_world_rank() == 0:
if not config["use_lora"]:
pipeline = DiffusionPipeline.from_pretrained(
config["model_dir"],
text_encoder=text_encoder.module,
unet=unet.module,
)
pipeline.save_pretrained(config["output_dir"])
else:
save_lora_weights(unet.module, text_encoder.module, config["output_dir"])
def unet_attn_processors_state_dict(unet) -> Dict[str, torch.tensor]:
"""
Returns:
a state dict containing just the attention processor parameters.
"""
attn_processors = unet.attn_processors
attn_processors_state_dict = {}
for attn_processor_key, attn_processor in attn_processors.items():
for parameter_key, parameter in attn_processor.state_dict().items():
param_name = f"{attn_processor_key}.{parameter_key}"
attn_processors_state_dict[param_name] = parameter
return attn_processors_state_dict
def save_lora_weights(unet, text_encoder, output_dir):
unet_lora_layers_to_save = None
text_encoder_lora_layers_to_save = None
unet_lora_layers_to_save = unet_attn_processors_state_dict(unet)
text_encoder_lora_layers_to_save = text_encoder_lora_state_dict(text_encoder)
LoraLoaderMixin.save_lora_weights(
output_dir,
unet_lora_layers=unet_lora_layers_to_save,
text_encoder_lora_layers=text_encoder_lora_layers_to_save,
)
if __name__ == "__main__":
args = train_arguments().parse_args()
# Build training dataset.
train_dataset = get_train_dataset(args)
print(f"Loaded training dataset (size: {train_dataset.count()})")
# Train with Ray Train TorchTrainer.
trainer = TorchTrainer(
train_fn,
train_loop_config=vars(args),
scaling_config=ScalingConfig(
use_gpu=True,
num_workers=args.num_workers,
),
datasets={
"train": train_dataset,
},
)
result = trainer.fit()
print(result)
doc/source/templates/05_dreambooth_finetuning/dreambooth_run.sh
Executable
+158
View File
@@ -0,0 +1,158 @@
#!/bin/bash
# shellcheck disable=SC2086
set -xe
# Step 0
pushd dreambooth || true
# Step 0 cont
# __preparation_start__
# TODO: If running on multiple nodes, change this path to a shared directory (ex: NFS)
export DATA_PREFIX="/tmp"
export ORIG_MODEL_NAME="CompVis/stable-diffusion-v1-4"
export ORIG_MODEL_HASH="b95be7d6f134c3a9e62ee616f310733567f069ce"
export ORIG_MODEL_DIR="$DATA_PREFIX/model-orig"
export ORIG_MODEL_PATH="$ORIG_MODEL_DIR/models--${ORIG_MODEL_NAME/\//--}/snapshots/$ORIG_MODEL_HASH"
export TUNED_MODEL_DIR="$DATA_PREFIX/model-tuned"
export IMAGES_REG_DIR="$DATA_PREFIX/images-reg"
export IMAGES_OWN_DIR="$DATA_PREFIX/images-own"
export IMAGES_NEW_DIR="$DATA_PREFIX/images-new"
# TODO: Add more worker nodes and increase NUM_WORKERS for more data-parallelism
export NUM_WORKERS=2
mkdir -p $ORIG_MODEL_DIR $TUNED_MODEL_DIR $IMAGES_REG_DIR $IMAGES_OWN_DIR $IMAGES_NEW_DIR
# __preparation_end__
# Unique token to identify our subject (e.g., a random dog vs. our unqtkn dog)
export UNIQUE_TOKEN="unqtkn"
skip_image_setup=false
use_lora=false
# parse args
for arg in "$@"; do
case $arg in
--skip_image_setup)
echo "Option --skip_image_setup is set"
skip_image_setup=true
;;
--lora)
echo "Option --lora is set"
use_lora=true
;;
*)
echo "Invalid option: $arg"
;;
esac
done
# Step 1
# __cache_model_start__
python cache_model.py --model_dir=$ORIG_MODEL_DIR --model_name=$ORIG_MODEL_NAME --revision=$ORIG_MODEL_HASH
# __cache_model_end__
download_image() {
# Step 2
# __supply_own_images_start__
# Only uncomment one of the following:
# Option 1: Use the dog dataset ---------
export CLASS_NAME="dog"
python download_example_dataset.py ./images/dog
export INSTANCE_DIR=./images/dog
# ---------------------------------------
# Option 2: Use the lego car dataset ----
# export CLASS_NAME="car"
# export INSTANCE_DIR=./images/lego-car
# ---------------------------------------
# Option 3: Use your own images ---------
# export CLASS_NAME="<class-of-your-subject>"
# export INSTANCE_DIR="/path/to/images/of/subject"
# ---------------------------------------
# Copy own images into IMAGES_OWN_DIR
cp -rf $INSTANCE_DIR/* "$IMAGES_OWN_DIR/"
# __supply_own_images_end__
# Clear reg dir
rm -rf "$IMAGES_REG_DIR"/*.jpg
# Step 3: START
python generate.py \
--model_dir=$ORIG_MODEL_PATH \
--output_dir=$IMAGES_REG_DIR \
--prompts="photo of a $CLASS_NAME" \
--num_samples_per_prompt=200 \
--use_ray_data
# Step 3: END
}
# Skip step 2 and 3 if skip_image_setup=true
if $skip_image_setup; then
echo "Skipping image downloading..."
else
download_image
fi
if [ "$use_lora" = false ]; then
echo "Start full-finetuning..."
# Step 4: START
python train.py \
--model_dir=$ORIG_MODEL_PATH \
--output_dir=$TUNED_MODEL_DIR \
--instance_images_dir=$IMAGES_OWN_DIR \
--instance_prompt="photo of $UNIQUE_TOKEN $CLASS_NAME" \
--class_images_dir=$IMAGES_REG_DIR \
--class_prompt="photo of a $CLASS_NAME" \
--train_batch_size=2 \
--lr=5e-6 \
--num_epochs=4 \
--max_train_steps=200 \
--num_workers $NUM_WORKERS
# Step 4: END
else
echo "Start LoRA finetuning..."
python train.py \
--use_lora \
--model_dir=$ORIG_MODEL_PATH \
--output_dir=$TUNED_MODEL_DIR \
--instance_images_dir=$IMAGES_OWN_DIR \
--instance_prompt="photo of $UNIQUE_TOKEN $CLASS_NAME" \
--class_images_dir=$IMAGES_REG_DIR \
--class_prompt="photo of a $CLASS_NAME" \
--train_batch_size=2 \
--lr=1e-4 \
--num_epochs=4 \
--max_train_steps=200 \
--num_workers $NUM_WORKERS
fi
# Clear new dir
rm -rf "$IMAGES_NEW_DIR"/*.jpg
if [ "$use_lora" = false ]; then
# Step 5: START
python generate.py \
--model_dir=$TUNED_MODEL_DIR \
--output_dir=$IMAGES_NEW_DIR \
--prompts="photo of a $UNIQUE_TOKEN $CLASS_NAME in a bucket" \
--num_samples_per_prompt=5
# Step 5: END
else
python generate.py \
--model_dir=$ORIG_MODEL_PATH \
--lora_weights_dir=$TUNED_MODEL_DIR \
--output_dir=$IMAGES_NEW_DIR \
--prompts="photo of a $UNIQUE_TOKEN $CLASS_NAME in a bucket" \
--num_samples_per_prompt=5
fi
# Save artifact
mkdir -p /tmp/artifacts
cp -f "$IMAGES_NEW_DIR"/0-*.jpg /tmp/artifacts/example_out.jpg
# Exit
popd || true
doc/source/templates/05_dreambooth_finetuning/playground.ipynb
+196
View File
@@ -0,0 +1,196 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Generate images with your fine-tuned Stable Diffusion model\n",
"\n",
"You should use this notebook to interactively generate images, after you've already fine-tuned a stable diffusion model and have a model checkpoint available to load. See the README for instructions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# TODO: Change this to the path of your fine-tuned model checkpoint!\n",
"# This is the $TUNED_MODEL_DIR variable defined in the run script.\n",
"TUNED_MODEL_PATH = \"/tmp/model-tuned\"\n",
"# TODO: Set the following variables if you fine-tuned with LoRA.\n",
"ORIG_MODEL_PATH = \"/tmp/model-orig/models--CompVis--stable-diffusion-v1-4/snapshots/b95be7d6f134c3a9e62ee616f310733567f069ce/\"\n",
"LORA_WEIGHTS_DIR = \"/tmp/model-tuned\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"First, load the model checkpoint as a HuggingFace 🤗 pipeline.\n",
"Load the model onto a GPU and define a function to generate images from a text prompt."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from os import environ\n",
"\n",
"import torch\n",
"from diffusers import DiffusionPipeline\n",
"\n",
"from dreambooth.generate_utils import load_lora_weights, get_pipeline\n",
"\n",
"pipeline = None\n",
"\n",
"def on_full_ft():\n",
" global pipeline\n",
" pipeline = get_pipeline(TUNED_MODEL_PATH)\n",
" pipeline.to(\"cuda\")\n",
" \n",
"def on_lora_ft():\n",
" assert ORIG_MODEL_PATH\n",
" assert LORA_WEIGHTS_DIR\n",
" global pipeline\n",
" pipeline = get_pipeline(ORIG_MODEL_PATH, LORA_WEIGHTS_DIR)\n",
" pipeline.to(\"cuda\")\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"def generate(\n",
" pipeline: DiffusionPipeline,\n",
" prompt: str,\n",
" img_size: int = 512,\n",
" num_samples: int = 1,\n",
") -> list:\n",
" return pipeline([prompt] * num_samples, height=img_size, width=img_size).images\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Try out your model!\n",
"\n",
"Now, play with your fine-tuned diffusion model through this simple GUI."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import time\n",
"import ipywidgets as widgets\n",
"from IPython.display import display, clear_output\n",
"\n",
"# TODO: When giving prompts, make sure to include your subject's unique identifier,\n",
"# as well as its class name.\n",
"# For example, if your subject's unique identifier is \"unqtkn\" and is a dog,\n",
"# you can give the prompt \"photo of a unqtkn dog on the beach\".\n",
"\n",
"# IPython GUI Layouts\n",
"\n",
"output = widgets.Output()\n",
"toggle_buttons = widgets.ToggleButtons(\n",
" options=[\"Full fine-tuning\",\"LoRA fine-tuning\"],\n",
" disabled=False,\n",
" button_style='', # 'success', 'info', 'warning', 'danger' or ''\n",
" value=None,\n",
" # layout=widgets.Layout(width='100px')\n",
")\n",
"\n",
"def toggle_callback(change):\n",
" with output:\n",
" clear_output()\n",
" if change[\"new\"] == \"Full fine-tuning\":\n",
" on_full_ft()\n",
" else:\n",
" on_lora_ft()\n",
" \n",
"toggle_buttons.observe(toggle_callback, names=\"value\")\n",
" \n",
"input_text = widgets.Text(\n",
" value=\"photo of a unqtkn dog on the beach\",\n",
" placeholder=\"\",\n",
" description=\"Prompt:\",\n",
" disabled=False,\n",
" layout=widgets.Layout(width=\"500px\"),\n",
")\n",
"\n",
"button = widgets.Button(description=\"Generate!\")\n",
"\n",
"# Define button click event\n",
"def on_button_clicked(b):\n",
" with output:\n",
" clear_output()\n",
" print(\"Generating images...\")\n",
" print(\n",
" \"(The output image may be completely black if it's filtered by \"\n",
" \"HuggingFace diffusers safety checkers.)\"\n",
" )\n",
" start_time = time.time()\n",
" images = generate(pipeline=pipeline, prompt=input_text.value, num_samples=2)\n",
" display(*images)\n",
" finish_time = time.time()\n",
" print(f\"Completed in {finish_time - start_time} seconds.\")\n",
"\n",
"button.on_click(on_button_clicked)\n",
"\n",
"# Display the widgets\n",
"display(toggle_buttons, widgets.HBox([input_text, button]), output)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# release memory properly\n",
"del pipeline \n",
"torch.cuda.empty_cache()"
]
}
],
"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.8.13"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
doc/source/templates/README.md
+84
View File
@@ -0,0 +1,84 @@
# Ray Starter Templates
These templates are a set of minimal examples that are quick and easy to run and customize.
Although the templates may include some machine learning framework-specific code, the individual code blocks are meant to be swapped in with your own application logic. The templates just serve as skeletons that showcase popular applications of Ray.
## Running on a Ray Cluster
<!-- TODO(justinvyu): Add in OSS cluster support. -->
Coming soon...
## Contributing Guide
To add a template:
1. Add your template as a directory somewhere in `doc/source/templates`.
For example:
```text
ray/
doc/source/templates/
<name-of-your-template>/
README.md
<name-of-your-template>.ipynb
requirements.txt (Optional)
templates.yaml
```
Your template does not need to be a Jupyter notebook. It can also be presented as a Python script with `README` instructions of how to run.
2. Add a release test for the template in `release/release_tests.yaml` (for both AWS and GCE). For Data tests, use `release/release_data_tests.yaml` instead.
See the section on workspace templates for an example. Note that the cluster env and compute config are a little different for release tests. Use the files in the `doc/source/templates/testing/release` folder.
The release test compute configs contain placeholders for regions and cloud ids that our CI infra will fill in. The cluster env builds a nightly docker image with all the required dependencies.
3. Add an entry to `doc/source/templates/templates.yaml` that links to your template.
See the top of the `templates.yaml` file for something to copy-paste and fill in your own values.
When you specify the template's compute config, see `doc/source/templates/configs` for shared configs. You can also create custom compute configs (of the same format as these shared ones).
For handling dependencies:
- If your template requires any special dependencies that are not included in a base image that you chose, be sure to list and provide instructions to install the necessary dependencies within the notebook. See `02_many_model_training` for an example.
- If your template requires a custom docker image, be sure to mention this in the `README` and link the docker image URL somewhere. See `03_serving_stable_diffusion` for an example.
4. Run a validation script on `templates.yaml` to make sure that the paths you specified are all valid and all yamls are properly formatted.
**Note:** This will also run in CI, but you can check quickly by running the validation script.
```bash
$ python doc/source/templates/testing/validate.py
Success!
```
5. Success! Your template is ready for review.
<!-- 2. Add another copy of the template that includes test-specific code and a smoke-test version if applicable.
**Note:** The need for a second test copy is temporary. Only one notebook will be needed from 2.5 onward, since the test-specific code will be filtered out.
**Label all test-specific code with the `remove-cell` Jupyter notebook tag.**
**Put this test copy in `doc/source/templates/tests/<name-of-your-template>.ipynb`.**
3. List the smoke-test version of the template in `doc/BUILD` under the templates section. This will configure the smoke-test version to run in pre-merge CI.
Set the `SMOKE_TEST` environment variable, which should be used in your template to **to make the template work for a single CI instance.** This environment variable can also be used to conditionally set certain smoke test parameters (like limiting dataset size).
**Make sure that you tag the test with `"gpu"` if required, and any other tags needed for special dependencies.**
```python
py_test_run_all_notebooks(
size = "large",
include = ["source/templates/tests/batch_inference.ipynb"],
exclude = [],
data = ["//doc:workspace_templates"],
tags = ["exclusive", "team:ml", "ray_air", "gpu"],
env = {"SMOKE_TEST": "1"},
)
``` -->
doc/source/templates/configs/compute/cpu/aws.yaml
+11
View File
@@ -0,0 +1,11 @@
# 8 m5.2xlarge nodes --> 64 CPUs
head_node_type:
name: head_node_type
instance_type: m5.2xlarge
worker_node_types:
- name: cpu_worker
instance_type: m5.2xlarge
min_workers: 7
max_workers: 7
use_spot: false
doc/source/templates/configs/compute/cpu/gce.yaml
+11
View File
@@ -0,0 +1,11 @@
# 8 n2-standard-8 nodes --> 64 CPUs
head_node_type:
name: head_node_type
instance_type: n2-standard-8
worker_node_types:
- name: cpu_worker
instance_type: n2-standard-8
min_workers: 7
max_workers: 7
use_spot: false
doc/source/templates/configs/compute/gpu/aws.yaml
+11
View File
@@ -0,0 +1,11 @@
# 3 g5.4xlarge nodes --> 48 CPUs, 3 GPUs
head_node_type:
name: head_node_type
instance_type: g5.4xlarge
worker_node_types:
- name: gpu_worker
instance_type: g5.4xlarge
min_workers: 0
max_workers: 3
use_spot: false
doc/source/templates/configs/compute/gpu/gce.yaml
+11
View File
@@ -0,0 +1,11 @@
# 4 n1-standard-8-nvidia-tesla-t4-1 nodes --> 32 CPUs, 4 GPUs
head_node_type:
name: head_node_type
instance_type: n1-standard-8-nvidia-tesla-t4-1
worker_node_types:
- name: gpu_worker
instance_type: n1-standard-8-nvidia-tesla-t4-1
min_workers: 3
max_workers: 3
use_spot: false
doc/source/templates/testing/cluster_envs/02_many_model_training.yaml
+16
View File
@@ -0,0 +1,16 @@
base_image: anyscale/ray-ml:nightly-py39-gpu
env_vars: {}
post_build_cmds:
# Install Ray
- pip3 uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
- {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}
python:
pip_packages:
- statsforecast==1.5.0
# numba doesn't support numpy > 1.24
# See: https://github.com/numba/numba/issues/8698
# NOTE: This is only an issue when `statsforecast` is installed with
# `pip install -U`, which is what's happening for this cluster env.
- numpy<1.25.0
doc/source/templates/testing/cluster_envs/03_serving_stable_diffusion.yaml
+20
View File
@@ -0,0 +1,20 @@
base_image: anyscale/ray:nightly-py39-cu118
debian_packages:
- curl
post_build_cmds:
# Install Ray
- pip3 uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
- {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}
python:
pip_packages:
- accelerate==0.20.3
- diffusers==0.17.1
- fastapi==0.97.0
- ipywidgets
- matplotlib==3.7.1
- numpy==1.24.3
- torch==2.0.1
- transformers==4.30.1
doc/source/templates/testing/cluster_envs/04_finetuning_llms_with_deepspeed.yaml
+33
View File
@@ -0,0 +1,33 @@
base_image: anyscale/ray:nightly-py39-cu118
env_vars: {}
debian_packages:
- libaio1
python:
pip_packages: [
peft==0.7.0,
deepspeed,
fairscale,
transformers>=4.31.0,
dataset,
accelerate,
evaluate,
bitsandbytes,
wandb,
pytorch-lightning,
protobuf,
torchmetrics,
lm_eval==0.3.0,
tiktoken==0.1.2,
sentencepiece,
]
conda_packages: []
post_build_cmds:
- pip uninstall bitsandbytes -y || true
- pip install torch==2.1.1 --index-url https://download.pytorch.org/whl/cu118
# Install Ray
- pip3 uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
- {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}
doc/source/templates/testing/cluster_envs/default_cluster_env_nightly_ml_py39.yaml
+9
View File
@@ -0,0 +1,9 @@
base_image: anyscale/ray-ml:nightly-py39-gpu
env_vars: {}
debian_packages:
- curl
post_build_cmds:
# Install Ray
- pip3 uninstall -y ray || true && pip3 install -U {{ env["RAY_WHEELS"] | default("ray") }}
- {{ env["RAY_WHEELS_SANITY_CHECK"] | default("echo No Ray wheels sanity check") }}
doc/source/templates/testing/compute_configs/04_finetuning_llms_with_deepspeed/aws_70b_v1.yaml
+23
View File
@@ -0,0 +1,23 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west-2
head_node_type:
name: head_node_type
instance_type: g5.48xlarge
resources:
custom_resources:
large_cpu_mem: 1
worker_node_types:
- name: gpu_worker
instance_type: g5.48xlarge
min_workers: 3
max_workers: 3
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/testing/compute_configs/04_finetuning_llms_with_deepspeed/aws_70b_v2.yaml
+29
View File
@@ -0,0 +1,29 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west-2
head_node_type:
name: head_node_type
instance_type: g5.48xlarge
resources:
custom_resources:
large_cpu_mem: 1
worker_node_types:
- name: large_gpu_worker
instance_type: g5.48xlarge
min_workers: 2
max_workers: 2
use_spot: false
- name: medium_gpu_worker
instance_type: g5.24xlarge
min_workers: 2
max_workers: 2
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/testing/compute_configs/04_finetuning_llms_with_deepspeed/aws_7b.yaml
+21
View File
@@ -0,0 +1,21 @@
# Autoscale to 16 g5.4xlarge --> 16 A10Gs
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west-2
head_node_type:
name: head_node
instance_type: m5.xlarge
worker_node_types:
- name: worker_node
instance_type: g5.4xlarge
min_workers: 0
max_workers: 16
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/testing/compute_configs/04_finetuning_llms_with_deepspeed/gce_7b.yaml
+16
View File
@@ -0,0 +1,16 @@
cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
region: us-west1
allowed_azs:
- us-west1-b
head_node_type:
name: head_node_type
instance_type: n2-standard-16
worker_node_types:
- name: gpu_worker
instance_type: g2-standard-16-nvidia-l4-1
min_workers: 0
max_workers: 16
use_spot: false
doc/source/templates/testing/compute_configs/cpu/aws.yaml
+21
View File
@@ -0,0 +1,21 @@
cloud_id: {{ env["ANYSCALE_CLOUD_ID"] }}
region: us-west-2
# 8 m5.2xlarge nodes --> 64 CPUs
head_node_type:
name: head_node_type
instance_type: m5.2xlarge
worker_node_types:
- name: cpu_worker
instance_type: m5.2xlarge
min_workers: 7
max_workers: 7
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/testing/compute_configs/cpu/gce.yaml
+17
View File
@@ -0,0 +1,17 @@
cloud_id: {{ env["ANYSCALE_CLOUD_ID"] }}
region: us-west1
allowed_azs:
- us-west1-b
# 8 n2-standard-8 nodes --> 64 CPUs
head_node_type:
name: head_node_type
instance_type: n2-standard-8
worker_node_types:
- name: cpu_worker
instance_type: n2-standard-8
min_workers: 7
max_workers: 7
use_spot: false
doc/source/templates/testing/compute_configs/gpu/aws.yaml
+21
View File
@@ -0,0 +1,21 @@
cloud_id: {{ env["ANYSCALE_CLOUD_ID"] }}
region: us-west-2
# 4 g4dn.2xlarge nodes --> 32 CPUs, 4 GPUs
head_node_type:
name: head_node_type
instance_type: g4dn.2xlarge
worker_node_types:
- name: gpu_worker
instance_type: g4dn.2xlarge
min_workers: 3
max_workers: 3
use_spot: false
advanced_configurations_json:
TagSpecifications:
- ResourceType: "instance"
Tags:
- Key: ttl-hours
Value: '24'
doc/source/templates/testing/compute_configs/gpu/gce.yaml
+17
View File
@@ -0,0 +1,17 @@
cloud_id: {{ env["ANYSCALE_CLOUD_ID"] }}
region: us-west1
allowed_azs:
- us-west1-b
# 4 n1-standard-8-nvidia-tesla-t4-1 nodes --> 32 CPUs, 4 GPUs
head_node_type:
name: head_node_type
instance_type: n1-standard-8-nvidia-tesla-t4-1
worker_node_types:
- name: gpu_worker
instance_type: n1-standard-8-nvidia-tesla-t4-1
min_workers: 3
max_workers: 3
use_spot: false
doc/source/templates/testing/docker/03_serving_stable_diffusion/Dockerfile
+9
View File
@@ -0,0 +1,9 @@
# Dockerfile used to create the docker image for `03_serving_stable_diffusion`.
FROM anyscale/ray:latest-py39-cu118
COPY requirements.txt ./
RUN pip install --no-cache-dir -U -r requirements.txt
RUN echo "Testing Ray Import..." && python -c "import ray"
RUN ray --version
doc/source/templates/testing/docker/03_serving_stable_diffusion/requirements.txt
+8
View File
@@ -0,0 +1,8 @@
accelerate==0.20.3
diffusers==0.17.1
fastapi==0.97.0
ipywidgets
matplotlib==3.7.1
numpy==1.24.3
torch==2.0.1
transformers==4.30.1
doc/source/templates/testing/docker/04_finetuning_llms_with_deepspeed/Dockerfile
+17
View File
@@ -0,0 +1,17 @@
# Dockerfile used to create the docker image for `04_finetuning_llms_with_deepspeed`.
FROM anyscale/ray:2.9.0-py310-cu121
COPY requirements.txt ./
RUN sudo apt-get update
RUN sudo apt-get install -y libaio1
RUN pip install --upgrade pip
# We need pydantic at this version to install deepspeed 0.10.2 (as part of the requirements.txt)
RUN pip install pydantic==1.10.7
RUN pip install -U -r requirements.txt
RUN pip install torch==2.1.1 --index-url https://download.pytorch.org/whl/cu121
RUN pip install flash-attn==2.4.2 --no-build-isolation
RUN echo "Testing Ray Import..." && python -c "import ray"
RUN ray --version
doc/source/templates/testing/docker/04_finetuning_llms_with_deepspeed/requirements.txt
+12
View File
@@ -0,0 +1,12 @@
deepspeed==0.10.2
fairscale
transformers>=4.36.2
dataset
accelerate
evaluate
wandb
pytorch-lightning
protobuf
torchmetrics
sentencepiece
peft==0.7.0