mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-27 15:42:05 +00:00
Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 yinghsienwu
|
fb89ec2c07 | ||
|
yinghsienwu
|
569264a0c4 | ||
|
yinghsienwu
|
b6b2e74cd7 | ||
|
yinghsienwu
|
4cb9b7fe38 | ||
|
yinghsienwu
|
4466795a7f | ||
|
yinghsienwu
|
5bcb683e85 | ||
|
yinghsienwu
|
8a41b1d7bf | ||
|
yinghsienwu
|
c5458742ff | ||
|
yinghsienwu
|
b61562c95b | ||
|
yinghsienwu
|
204c214a7e | ||
|
yinghsienwu
|
70b8d8c877 |
@@ -238,7 +238,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
|
||||
|
||||
# Look for the python version specification pattern
|
||||
re_match = re.search(
|
||||
r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
|
||||
"python version = (\d\.\d)", markdown, flags=re.IGNORECASE
|
||||
)
|
||||
if re_match:
|
||||
# get the version number
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
|
||||
notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
|
||||
notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
|
||||
notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
|
||||
.cloud-build/tests/python_version_test.ipynb
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
|
||||
# Ignore model garden dockerfiles:
|
||||
- package-ecosystem: "npm"
|
||||
directory: "/community-content/vertex_model_garden"
|
||||
schedule:
|
||||
interval: "monthly"
|
||||
ignore:
|
||||
- dependency-name: "*"
|
||||
@@ -4,7 +4,7 @@
|
||||
# 2. To lint specific notebooks:
|
||||
# docker run -v ${PWD}:/setup/app gcr.io/python-docs-samples-tests/notebook_linter:latest notebooks/1.ipynb notebooks/2.ipynb
|
||||
|
||||
FROM python:3.13
|
||||
FROM python:3.12
|
||||
|
||||
WORKDIR setup
|
||||
|
||||
|
||||
@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
|
||||
ipython
|
||||
jupyter
|
||||
nbconvert
|
||||
black==25.1.0
|
||||
pyupgrade==3.20.0
|
||||
isort==6.0.1
|
||||
flake8==7.3.0
|
||||
nbqa==1.9.1
|
||||
black==23.12.1
|
||||
pyupgrade==3.15.0
|
||||
isort==5.13.2
|
||||
flake8==7.0.0
|
||||
nbqa==1.7.1
|
||||
|
||||
|
||||
@@ -58,7 +58,7 @@ done
|
||||
# Only check notebooks in test folders modified in this pull request.
|
||||
# Note: Use process substitution to persist the data in the array
|
||||
if [ ${#notebooks[@]} -eq 0 ]; then
|
||||
echo "Checking for changed notebooks using git"
|
||||
echo "Checking for changed notebooked using git"
|
||||
while read -r file || [ -n "$line" ]; do
|
||||
notebooks+=("$file")
|
||||
done < <(git diff --name-only main... | grep '\.ipynb$')
|
||||
|
||||
@@ -1,176 +1,37 @@
|
||||
#  Google Cloud Vertex AI Samples
|
||||
# Google Cloud Vertex AI Samples
|
||||
|
||||
This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
|
||||
[](LICENSE)
|
||||
|
||||
Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/) sample repository.
|
||||
|
||||
## Overview
|
||||
|
||||
[Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
|
||||
|
||||
For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
|
||||
|
||||
## Explore, learn and contribute
|
||||
|
||||
You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI!
|
||||
|
||||
### Explore and learn
|
||||
|
||||
Explore this repository, follow the links in the header section of each of the notebooks to -
|
||||
|
||||
 Open and run the notebook in [Colab](https://colab.google/)\
|
||||
 Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
|
||||
 Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
|
||||
 View the notebook on Github
|
||||
|
||||
### Contribute
|
||||
|
||||
See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
|
||||
|
||||
## Get started
|
||||
|
||||
To get started using Vertex AI, you must have a Google Cloud project.
|
||||
|
||||
- If you don't have a Google Cloud project, you can learn and build on GCP for free using [Free Trail](https://cloud.google.com/free).
|
||||
- Once you have a Google Cloud project, you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).
|
||||
|
||||
The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
|
||||
|
||||
## Repository structure
|
||||
|
||||
```bash
|
||||
├── community-content - Sample code and tutorials contributed by the community
|
||||
├── notebooks
|
||||
│ ├── community - Notebooks contributed by the community
|
||||
│ ├── official - Notebooks demonstrating use of each Vertex AI service
|
||||
│ │ ├── automl
|
||||
│ │ ├── custom
|
||||
│ │ ├── ...
|
||||
│ ├── community - Notebooks contributed by the community
|
||||
│ │ ├── model_garden
|
||||
│ │ ├── ...
|
||||
├── community-content - Sample code and tutorials contributed by the community
|
||||
|
||||
```
|
||||
## Examples
|
||||
|
||||
<!-- markdownlint-disable MD033 -->
|
||||
<table>
|
||||
## Contributing
|
||||
|
||||
<tr>
|
||||
<th style="text-align: center;">Category</th>
|
||||
<th style="text-align: center;">Product</th>
|
||||
<th style="text-align: center;">Description</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Model</td>
|
||||
<td>
|
||||
<a href="notebooks/community/model_garden"><code>Model Garden/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Curated collection of first-party, open-source, and third-party models available on Vertex AI including Gemini, Gemma, Llama 3, Claude 3 and many more.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Data</td>
|
||||
<td>
|
||||
<a href="notebooks/official/feature_store"><code>Feature Store/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Set up and manage online serving using Vertex AI Feature Store.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/datasets"><code>datasets/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Use BigQuery and Data Labeling service with Vertex AI.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Model development</td>
|
||||
<td>
|
||||
<a href="notebooks/official/automl"><code>automl/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Train and make predictions on AutoML models
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/custom"><code>custom/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Create, deploy and serve custom models on Vertex AI
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/ray_on_vertex_ai"><code>ray_on_vertex_ai/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Use Colab Enterprise and Vertex AI SDK for Python to connect to the Ray Cluster.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Deploy and use</td>
|
||||
<td>
|
||||
<a href="notebooks/official/prediction"><code>prediction/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Build, train and deploy models using prebuilt containers for custom training and prediction.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/model_registry"><code>model_registry/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Use Model Registry to create and register a model.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/explainable_ai"><code>Explainable AI/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Use Vertex Explainable AI's feature-based and example-based explanations to explain how or why a model produced a specific prediction.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td></td>
|
||||
<td>
|
||||
<a href="notebooks/official/ml_metadata"><code>ml_metadata/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Record the metadata and artifacts and query that metadata to help analyze, debug, and audit the performance of your ML system.
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Tools</td>
|
||||
<td>
|
||||
<a href="notebooks/official/pipelines"><code>Pipelines/</code></a>
|
||||
</td>
|
||||
<td>
|
||||
Use `Vertex AI Pipelines` and `Google Cloud Pipeline Components` to build, tune, or deploy a custom model.
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
<!-- markdownlint-enable MD033 -->
|
||||
Contributions welcome! See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
|
||||
|
||||
## Getting help
|
||||
|
||||
## Get help
|
||||
|
||||
Please use the [Issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
|
||||
Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
|
||||
|
||||
## Disclaimer
|
||||
|
||||
This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.
|
||||
|
||||
## Feedback
|
||||
|
||||
## References
|
||||
- [Vertex AI Jupyter Notebook tutorials](https://cloud.google.com/vertex-ai/docs/tutorials/jupyter-notebooks)
|
||||
- Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository
|
||||
- [Vertex AI documentaton](https://cloud.google.com/vertex-ai/docs)
|
||||
|
||||
Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.
|
||||
|
||||
@@ -10,7 +10,6 @@
|
||||
/pipeline_components @Ark-kun
|
||||
/pipeline_components/image_ml_model_training @lakeyk
|
||||
/prediction_featurestore_integration @googleapis/vertex-prediction-team
|
||||
/vertex_model_garden/model_oss/notebook_util @minwoo33park
|
||||
/vertex_model_garden/model_oss/util @weigary
|
||||
/vertex_model_garden/model_oss/diffusers @weigary
|
||||
/vertex_model_garden/model_oss/keras @dstnluong-google
|
||||
@@ -20,14 +19,10 @@
|
||||
/vertex_model_garden/model_oss/movinet @KCFindstr
|
||||
/vertex_model_garden/model_oss/data_converter @KCFindstr
|
||||
/vertex_model_garden/model_oss/peft @weigary
|
||||
/vertex_model_garden/model_oss/peft/templates @rayandasoriya
|
||||
/vertex_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
|
||||
/vertex_model_garden/model_oss/tfvision @dstnluong-google
|
||||
/vertex_model_garden/model_oss/fvlm @minwoo33park
|
||||
/vertex_model_garden/model_oss/imagebind @kathyyu-google
|
||||
/vertex_model_garden/model_oss/llava @py4
|
||||
/vertex_model_garden/model_oss/vllm @kathyyu-google
|
||||
/vertex_model_garden/benchmarking_reports @lavraicse
|
||||
/vertex_model_garden/model_oss/autogluon @lavraicse
|
||||
/vertex_distributed_training/a3mega/llama-3-8b-nemo-pretraining @mstyer-google @erwinh85 @mchrestkha
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
absl-py==1.1.0
|
||||
fastapi==0.109.1
|
||||
fastapi==0.75.2
|
||||
uvicorn==0.18.2
|
||||
timm==0.5.4
|
||||
smart_open==6.0.0
|
||||
|
||||
+6
-30
@@ -1,40 +1,16 @@
|
||||
# Stage 1: Build Environment
|
||||
FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime AS builder
|
||||
|
||||
# Install necessary tools and dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y curl gnupg && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
|
||||
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y google-cloud-sdk
|
||||
|
||||
# Copy application code
|
||||
COPY . /trainer
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /trainer
|
||||
|
||||
# Install Python dependencies
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Stage 2: Runtime Environment
|
||||
FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
|
||||
|
||||
# Install Google Cloud SDK
|
||||
RUN apt-get update && \
|
||||
apt-get install -y curl gnupg && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
|
||||
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
|
||||
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y google-cloud-sdk && \
|
||||
apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||
apt-get install google-cloud-sdk -y
|
||||
|
||||
# Copy from the builder stage
|
||||
COPY --from=builder /trainer /trainer
|
||||
COPY . /trainer
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /trainer
|
||||
|
||||
# Set the entry point
|
||||
ENTRYPOINT ["python", "-m", "task"]
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
ENTRYPOINT ["python", "-m", "task"]
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
torch==2.2.0
|
||||
torch==1.13.1
|
||||
torchvision==0.9.1
|
||||
tensorboard==2.5.0
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
torch==2.7.0
|
||||
torch==1.13.1
|
||||
torchvision==0.9.1
|
||||
tensorboard==2.5.0
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
google-cloud-bigquery==2.20.0
|
||||
tensorflow==2.12.1
|
||||
pillow==10.3.0
|
||||
tensorflow==2.7.2
|
||||
pillow==10.0.1
|
||||
tf-agents==0.8.0
|
||||
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
google-cloud-pubsub==2.5.0
|
||||
pillow==10.3.0
|
||||
pillow==10.0.1
|
||||
tf-agents==0.8.0
|
||||
tensorflow==2.12.1
|
||||
tensorflow==2.7.2
|
||||
|
||||
+2
-2
@@ -1,5 +1,5 @@
|
||||
dataclasses==0.6
|
||||
google-cloud-aiplatform==1.8.1
|
||||
tensorflow==2.12.1
|
||||
pillow==10.3.0
|
||||
tensorflow==2.7.2
|
||||
pillow==10.0.1
|
||||
tf-agents==0.8.0
|
||||
+1
-1
@@ -1 +1 @@
|
||||
tensorflow==2.12.1
|
||||
tensorflow==2.7.2
|
||||
+1
-1
@@ -1 +1 @@
|
||||
tensorflow==2.12.1
|
||||
tensorflow==2.7.2
|
||||
-126
@@ -1,126 +0,0 @@
|
||||
# Vertex AI Training: Llama 3.1 8B pre-training using Nvidia A3 Mega VMs (H100)
|
||||
This document provides a step-by-step guide for pre-training a Llama 3.1 8B model on the `en-wiki` dataset using multiple [Vertex AI Custom Training](https://cloud.google.com/vertex-ai/docs/training/overview) `a3-megagpu-8g` nodes.
|
||||
|
||||
We will use a custom container based on NVIDIA's [NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/24.07/overview.html) to demonstrate a scalable, multi-node training workflow. All required artifacts and commands are included.
|
||||
|
||||
## 1. Prerequisites
|
||||
|
||||
### 1.1. Google Cloud Project setup
|
||||
- **Enable APIs:** Ensure the Vertex AI API is [enabled for your project](http://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).
|
||||
- **H100 Mega Quota:** A3 Mega VMs are powered by H100 GPUs. Request quota for `custom_model_training_nvidia_h100_mega_gpus` in one of the [supported regions](https://cloud.google.com/vertex-ai/docs/general/locations#accelerator_support). If using Spot VMs, request `custom_model_training_preemptible_nvidia_h100_mega_gpus` quota instead.
|
||||
- **Reservations (Optional but recommended):** For guaranteed capacity, [create a reservation](https://cloud.google.com/compute/docs/instances/reservations-shared) and ensure the reservation is shared with the Vertex AI service account. This guide requires a minimum of **16 H100 GPUs** (2 full A3 Mega nodes).
|
||||
|
||||
### 1.2. GCS bucket
|
||||
Create a [Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) in the same region where you have quota. If you're using Hierarchical Namespace for your bucket, you may need to update permissions of the Vertex AI Custom Code Service Agent .
|
||||
|
||||
This bucket is used for:
|
||||
- Staging the training application.
|
||||
- Storing model checkpoints and logs.
|
||||
- Storing data if you use your own data.
|
||||
|
||||
|
||||
## 2. Setup & configuration
|
||||
|
||||
### 2.1. Clone the repo
|
||||
First clone the repo into your development environment.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
|
||||
```
|
||||
|
||||
Navigate to the root folder for this sample.
|
||||
|
||||
### 2.2. Environment Setup
|
||||
First, configure your local environment. These variables are used in subsequent commands.
|
||||
|
||||
```bash
|
||||
# Required: Update with your values
|
||||
export PROJECT_ID="<your-project-id>"
|
||||
export REPOSITORY="<your-artifact-registry-repo-name>" # e.g., "my-containers"
|
||||
export BUCKET="<your-gcs-bucket-name>"
|
||||
|
||||
# Optional: Change if needed
|
||||
export REGION="us-central1"
|
||||
|
||||
# --- Do not change the lines below ---
|
||||
export ARTIFACT_REGISTRY="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}"
|
||||
export REPO_ROOT=$(git rev-parse --show-toplevel)
|
||||
```
|
||||
|
||||
## 3. Build and push a docker container image to Artifact Registry
|
||||
Normally, you can use any custom training container on Vertex AI Training. In this example you build a NeMo Docker image that is based on the [Nvidia’s NeMo 24.09](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags) image. Use Cloud Build to build and push the container image.
|
||||
|
||||
This document picked NeMo as the demonstrating container since it’s a widely adopted GPU LLM training framework providing high performance and versatile training functionalities.
|
||||
|
||||
In addition to the base image, some customizations are included to form the final prebuilt image:
|
||||
- Some dependencies are installed to integrate with Vertex AI Training.
|
||||
- An entrypoint script that sets up required environments and calls the training job.
|
||||
- Some patches are applied to the NeMo code to let it load the dataset from a GCS bucket.
|
||||
|
||||
Run this command to build the container and push the container into the Google Artifact Registry.
|
||||
|
||||
```bash
|
||||
cd "${REPO_ROOT}/community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining"
|
||||
export IMAGE_NAME="vertex-nemo-llama"
|
||||
gcloud builds submit . \
|
||||
--project="${PROJECT_ID}" \
|
||||
--region="${REGION}" \
|
||||
--config=docker/cloudbuild.yml \
|
||||
--substitutions="_ARTIFACT_REGISTRY=${ARTIFACT_REGISTRY},_IMAGE_NAME=${IMAGE_NAME}" \
|
||||
--timeout="2h" \
|
||||
--machine-type="e2-highcpu-32"
|
||||
```
|
||||
|
||||
## 4. Launch the Training Job
|
||||
|
||||
|
||||
### 4.1. Job Configuration File
|
||||
Once the container is built, update the job_config.json to set up the training job.
|
||||
File: job_config.json
|
||||
```json
|
||||
{
|
||||
"project_id": "<project-id>",
|
||||
"region": "<region>",
|
||||
"zone": "<zone if using reservation>",
|
||||
"bucket": "<bucket>",
|
||||
"dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
|
||||
"image_uri": "<docker image uri from artifact registry>",
|
||||
"strategy": "spot",
|
||||
"nodes": "2",
|
||||
"machine_type": "a3-megagpu-8g",
|
||||
"gpu_type": "NVIDIA_H100_MEGA_80GB",
|
||||
"gpus_per_node": "8",
|
||||
"recipe_name": "llama3_1_8b_pretrain_a3mega",
|
||||
"job_prefix": "vertex-spot-",
|
||||
"reservation_name": ""
|
||||
}
|
||||
```
|
||||
|
||||
### 4.2 Launch the Training Job
|
||||
|
||||
First, create a Python virtual environment using your tool of choice, then install
|
||||
the requirements specified in `requirements.txt`. Using `pip`, the command would be:
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Now launch the Vertex AI training job using the provided Python script.
|
||||
|
||||
```bash
|
||||
python3 scripts/launch.py --config_file=job_config.json
|
||||
```
|
||||
|
||||
This script reads job_config.json, defines the cluster specification (2 nodes, 8 GPUs each), and submits the custom training job to Vertex AI.
|
||||
|
||||
## 5. Monitor and Clean Up
|
||||
|
||||
### 5.1. Monitoring
|
||||
Vertex AI Console: Track the job's status in the Google Cloud Console under Vertex AI > Training > Custom Jobs.
|
||||
Logs: View detailed logs in Cloud Logging by filtering for your job name.
|
||||
Checkpoints: Model checkpoints are saved to your GCS bucket at the path specified in your training script's configuration.
|
||||
|
||||
### 5.2. Cleaning Up
|
||||
To avoid ongoing charges, delete the resources you created:
|
||||
- The Artifact Registry image.
|
||||
- The contents of the GCS bucket (checkpoints, logs).
|
||||
- The Vertex AI Custom Job will eventually complete or fail, incurring no further cost.
|
||||
-265
@@ -1,265 +0,0 @@
|
||||
# Reference:
|
||||
# https://github.com/NVIDIA/NeMo-Framework-Launcher/blob/24.07/launcher_scripts/conf/training/llama/llama3_1_8b.yaml
|
||||
name: llama3_1_8b_pretrain_a3mega
|
||||
restore_from_path: null # used when starting from a .nemo file
|
||||
|
||||
trainer:
|
||||
devices: 8
|
||||
num_nodes: 1
|
||||
accelerator: gpu
|
||||
precision: bf16
|
||||
logger: false # logger provided by exp_manager
|
||||
enable_checkpointing: false
|
||||
use_distributed_sampler: false
|
||||
max_epochs: -1 # PTL default. In practice, max_steps will be reached first.
|
||||
max_steps: 30 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
|
||||
log_every_n_steps: 1
|
||||
val_check_interval: null
|
||||
limit_val_batches: 1
|
||||
limit_test_batches: 1
|
||||
accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
|
||||
gradient_clip_val: 1.0
|
||||
benchmark: false
|
||||
enable_model_summary: false # default PTL callback for this does not support model parallelism, instead we log manually
|
||||
|
||||
exp_manager:
|
||||
explicit_log_dir: null
|
||||
exp_dir: /data
|
||||
name: ${name}
|
||||
create_dllogger_logger: true
|
||||
dllogger_logger_kwargs:
|
||||
verbose: true
|
||||
stdout: true
|
||||
json_file: "/data/dllogger.json"
|
||||
create_wandb_logger: false
|
||||
wandb_logger_kwargs:
|
||||
project: null
|
||||
name: null
|
||||
resume_if_exists: true
|
||||
resume_ignore_no_checkpoint: true
|
||||
create_checkpoint_callback: false
|
||||
checkpoint_callback_params:
|
||||
monitor: val_loss
|
||||
save_top_k: 3
|
||||
mode: min
|
||||
always_save_nemo: false # saves nemo file during validation, not implemented for model parallel
|
||||
save_nemo_on_train_end: false # not recommended when training large models on clusters with short time limits
|
||||
filename: 'megatron_gpt--{val_loss:.2f}-{step}-{consumed_samples}'
|
||||
model_parallel_size: ${multiply:${model.tensor_model_parallel_size}, ${model.pipeline_model_parallel_size}}
|
||||
seconds_to_sleep: 5 # Allows node_rank!=0 to sleep and let node0 to init, like preparing data
|
||||
|
||||
model:
|
||||
mcore_gpt: true
|
||||
# specify micro_batch_size, global_batch_size, and model parallelism
|
||||
# gradient accumulation will be done automatically based on data_parallel_size
|
||||
micro_batch_size: 1 # limited by GPU memory
|
||||
global_batch_size: 1024 # will use more micro batches to reach global batch size
|
||||
tensor_model_parallel_size: 1 # intra-layer model parallelism
|
||||
pipeline_model_parallel_size: 2 # inter-layer model parallelism
|
||||
context_parallel_size: 1
|
||||
virtual_pipeline_model_parallel_size: null # interleaved pipeline
|
||||
## Sequence Parallelism
|
||||
# Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially
|
||||
# See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
|
||||
sequence_parallel: false
|
||||
|
||||
fsdp: false
|
||||
fsdp_cpu_offload: true
|
||||
fsdp_sharding_strategy: "full" # Method to shard model states. Available options are 'full', 'hybrid', and 'grad'.
|
||||
fsdp_grad_reduce_dtype: "16" # Gradient reduction data type.
|
||||
fsdp_sharded_checkpoint: false # Store and load FSDP shared checkpoint.
|
||||
fsdp_use_orig_params: false # Set to True to use FSDP for specific peft scheme.
|
||||
|
||||
# Distributed checkpoint setup
|
||||
dist_ckpt_format: "torch_dist" # Set to 'torch_dist' to use PyTorch distributed checkpoint format.
|
||||
dist_ckpt_load_on_device: true # whether to load checkpoint weights directly on GPU or to CPU
|
||||
dist_ckpt_parallel_save: true # if true, each worker will write its own part of the dist checkpoint
|
||||
dist_ckpt_parallel_save_within_dp: false # if true, save will be parallelized only within a DP group (whole world otherwise), which might slightly reduce the save overhead
|
||||
dist_ckpt_parallel_load: false # if true, each worker will load part of the dist checkpoint and exchange with NCCL. Might use some extra GPU memory
|
||||
dist_ckpt_torch_dist_multiproc: 2 # number of extra processes per rank used during ckpt save with PyTorch distributed format
|
||||
dist_ckpt_assume_constant_structure: false # set to True only if the state dict structure doesn't change within a single job. Allows caching some computation across checkpoint saves.
|
||||
dist_ckpt_parallel_dist_opt: true # parallel save/load of a DistributedOptimizer. 'True' allows performant save and reshardable checkpoints. Set to 'False' only in order to minimize the number of checkpoint files.
|
||||
dist_ckpt_load_strictness: null # defines checkpoint keys mismatch behavior (only during dist-ckpt load). Choices: assume_ok_unexpected (default - try loading without any check), log_all (log mismatches), raise_all (raise mismatches)
|
||||
|
||||
# model architecture
|
||||
encoder_seq_length: 8192
|
||||
max_position_embeddings: ${.encoder_seq_length}
|
||||
num_layers: 32 # 8b: 32 | 70b: 80 | 405b: 126
|
||||
hidden_size: 4096 # 8b: 4096 | 70b: 8192 | 405b: 16384
|
||||
ffn_hidden_size: 14336 # 8b: 14336 | 70b: 28672 | 405b: 53248
|
||||
num_attention_heads: 32 # 8b: 32 | 70b: 64 | 405b: 128
|
||||
num_query_groups: 8 # Number of query groups for group query attention. If None, normal attention is used. 8b: 8 | 70b: 8 | 405b: 16
|
||||
init_method_std: 0.01 # Standard deviation of the zero mean normal distribution used for weight initialization. 8b: 0.01 | 70b: 0.008944 | 405b: 0.02
|
||||
use_scaled_init_method: true # use scaled residuals initialization
|
||||
hidden_dropout: 0.0 # Dropout probability for hidden state transformer.
|
||||
attention_dropout: 0.0 # Dropout probability for attention
|
||||
ffn_dropout: 0.0 # Dropout probability in the feed-forward layer.
|
||||
kv_channels: null # Projection weights dimension in multi-head attention. Set to hidden_size // num_attention_heads if null
|
||||
apply_query_key_layer_scaling: true # scale Q * K^T by 1 / layer-number.
|
||||
normalization: 'rmsnorm' # Normalization layer to use. Options are 'layernorm', 'rmsnorm'
|
||||
layernorm_epsilon: 1e-5
|
||||
do_layer_norm_weight_decay: false # True means weight decay on all params
|
||||
make_vocab_size_divisible_by: 128 # Pad the vocab size to be divisible by this value for computation efficiency.
|
||||
pre_process: true # add embedding
|
||||
post_process: true # add pooler
|
||||
persist_layer_norm: true # Use of persistent fused layer norm kernel.
|
||||
bias: false # Whether to use bias terms in all weight matrices.
|
||||
activation: 'fast-swiglu' # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
|
||||
headscale: false # Whether to learn extra parameters that scale the output of the each self-attention head.
|
||||
transformer_block_type: 'pre_ln' # Options ['pre_ln', 'post_ln', 'normformer']
|
||||
openai_gelu: false # Use OpenAI's GELU instead of the default GeLU
|
||||
normalize_attention_scores: true # Whether to scale the output Q * K^T by 1 / sqrt(hidden_size_per_head). This arg is provided as a configuration option mostly for compatibility with models that have been weight-converted from HF. You almost always want to se this to True.
|
||||
position_embedding_type: 'rope' # Position embedding type. Options ['learned_absolute', 'rope']
|
||||
rotary_percentage: 1.0 # If using position_embedding_type=rope, then the per head dim is multiplied by this.
|
||||
attention_type: 'multihead' # Attention type. Options ['multihead']
|
||||
share_embeddings_and_output_weights: false # Share embedding and output layer weights.
|
||||
scale_positional_embedding: true # This is false for llama3 models. Only used for >= llama3.1.
|
||||
|
||||
# Use GPT2BPETokenizer for test, because the testing dataset is tokenized by this tokenizer.
|
||||
# https://docs.nvidia.com/nemo-framework/user-guide/24.07/playbooks/singlenodepretrain.html#data-download-and-pre-processing
|
||||
tokenizer:
|
||||
library: megatron
|
||||
type: GPT2BPETokenizer
|
||||
model: null # /path/to/tokenizer.model
|
||||
vocab_file: null
|
||||
merge_file: null
|
||||
delimiter: null # only used for tabular tokenizer
|
||||
sentencepiece_legacy: false # Legacy=True allows you to add special tokens to sentencepiece tokenizers.
|
||||
|
||||
# Mixed precision
|
||||
native_amp_init_scale: 4294967296 # 2 ** 32
|
||||
native_amp_growth_interval: 1000
|
||||
hysteresis: 2 # Gradient scale hysteresis
|
||||
fp32_residual_connection: false # Move residual connections to fp32
|
||||
fp16_lm_cross_entropy: false # Move the cross entropy unreduced loss calculation for lm head to fp16
|
||||
|
||||
# Megatron O2-style half-precision
|
||||
megatron_amp_O2: true # Enable O2-level automatic mixed precision using main parameters
|
||||
grad_allreduce_chunk_size_mb: 125
|
||||
|
||||
# Fusion
|
||||
grad_div_ar_fusion: true # Fuse grad division into torch.distributed.all_reduce. Only used with O2 and no pipeline parallelism..
|
||||
gradient_accumulation_fusion: true # Fuse weight gradient accumulation to GEMMs. Only used with pipeline parallelism and O2.
|
||||
bias_activation_fusion: true # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
|
||||
bias_dropout_add_fusion: true # Use a kernel that fuses the bias addition, dropout and residual connection addition.
|
||||
masked_softmax_fusion: true # Use a kernel that fuses the attention softmax with it's mask.
|
||||
apply_rope_fusion: true # Use a kernel to add rotary positional embeddings. Only used if position_embedding_type=rope
|
||||
cross_entropy_loss_fusion: true
|
||||
|
||||
# Miscellaneous
|
||||
seed: 1234
|
||||
resume_from_checkpoint: null # manually set the checkpoint file to load from
|
||||
use_cpu_initialization: false # Init weights on the CPU (slow for large models)
|
||||
onnx_safe: false # Use work-arounds for known problems with Torch ONNX exporter.
|
||||
apex_transformer_log_level: 30 # Python logging level displays logs with severity greater than or equal to this
|
||||
gradient_as_bucket_view: true # PyTorch DDP argument. Allocate gradients in a contiguous bucket to save memory (less fragmentation and buffer memory)
|
||||
sync_batch_comm: false # Enable stream synchronization after each p2p communication between pipeline stages
|
||||
|
||||
## Activation Checkpointing
|
||||
# NeMo Megatron supports 'selective' activation checkpointing where only the memory intensive part of attention is checkpointed.
|
||||
# These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+).
|
||||
# See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
|
||||
# 'full' will checkpoint the entire transformer layer.
|
||||
activations_checkpoint_granularity: null # 'selective' or 'full'
|
||||
activations_checkpoint_method: null # 'uniform', 'block'
|
||||
# 'uniform' divides the total number of transformer layers and checkpoints the input activation
|
||||
# of each chunk at the specified granularity. When used with 'selective', 'uniform' checkpoints all attention blocks in the model.
|
||||
# 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity
|
||||
activations_checkpoint_num_layers: null
|
||||
# when using 'uniform' this creates groups of transformer layers to checkpoint. Usually set to 1. Increase to save more memory.
|
||||
# when using 'block' this this will checkpoint the first activations_checkpoint_num_layers per pipeline stage.
|
||||
num_micro_batches_with_partial_activation_checkpoints: null
|
||||
# This feature is valid only when used with pipeline-model-parallelism.
|
||||
# When an integer value is provided, it sets the number of micro-batches where only a partial number of Transformer layers get checkpointed
|
||||
# and recomputed within a window of micro-batches. The rest of micro-batches in the window checkpoint all Transformer layers. The size of window is
|
||||
# set by the maximum outstanding micro-batch backpropagations, which varies at different pipeline stages. The number of partial layers to checkpoint
|
||||
# per micro-batch is set by 'activations_checkpoint_num_layers' with 'activations_checkpoint_method' of 'block'.
|
||||
# This feature enables using activation checkpoint at a fraction of micro-batches up to the point of full GPU memory usage.
|
||||
activations_checkpoint_layers_per_pipeline: null
|
||||
# This feature is valid only when used with pipeline-model-parallelism.
|
||||
# When an integer value (rounded down when float is given) is provided, it sets the number of Transformer layers to skip checkpointing at later
|
||||
# pipeline stages. For example, 'activations_checkpoint_layers_per_pipeline' of 3 makes pipeline stage 1 to checkpoint 3 layers less than
|
||||
# stage 0 and stage 2 to checkpoint 6 layers less stage 0, and so on. This is possible because later pipeline stage
|
||||
# uses less GPU memory with fewer outstanding micro-batch backpropagations. Used with 'num_micro_batches_with_partial_activation_checkpoints',
|
||||
# this feature removes most of activation checkpoints at the last pipeline stage, which is the critical execution path.
|
||||
|
||||
## Transformer Engine
|
||||
transformer_engine: true
|
||||
fp8: false # enables fp8 in TransformerLayer forward
|
||||
fp8_e4m3: false # sets fp8_format = recipe.Format.E4M3
|
||||
fp8_hybrid: false # sets fp8_format = recipe.Format.HYBRID
|
||||
fp8_margin: 0 # scaling margin
|
||||
fp8_interval: 1 # scaling update interval
|
||||
fp8_amax_history_len: 1024 # Number of steps for which amax history is recorded per tensor
|
||||
fp8_amax_compute_algo: 'max' # 'most_recent' or 'max'. Algorithm for computing amax from history
|
||||
ub_tp_comm_overlap: false # do not turn on because of b/397797926
|
||||
use_flash_attention: true
|
||||
gc_interval: 100
|
||||
|
||||
## Offloading Activations/Weights to CPU
|
||||
cpu_offloading: false
|
||||
cpu_offloading_num_layers: ${sum:${.num_layers},-1} # This value should be between [1,num_layers-1] as we don't want to offload the final layer's activations and expose any offloading duration for the final layer
|
||||
cpu_offloading_activations: true
|
||||
cpu_offloading_weights: true
|
||||
|
||||
data:
|
||||
# Path to data must be specified by the user.
|
||||
# Supports List, String and Dictionary
|
||||
# List : can override from the CLI: "model.data.data_prefix=[.5,/raid/data/pile/my-gpt3_00_text_document,.5,/raid/data/pile/my-gpt3_01_text_document]",
|
||||
# Or see example below:
|
||||
# data_prefix:
|
||||
# - .5
|
||||
# - /raid/data/pile/my-gpt3_00_text_document
|
||||
# - .5
|
||||
# - /raid/data/pile/my-gpt3_01_text_document
|
||||
# Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
|
||||
# Or see example below:
|
||||
# "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
|
||||
data_prefix: [1.0, /data/hfbpe_gpt_training_data_text_document]
|
||||
index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
|
||||
data_impl: mmap
|
||||
splits_string: 900,50,50
|
||||
seq_length: ${model.encoder_seq_length}
|
||||
skip_warmup: true
|
||||
num_workers: 2
|
||||
dataloader_type: single # cyclic
|
||||
reset_position_ids: false # Reset position ids after end-of-document token
|
||||
reset_attention_mask: false # Reset attention mask after end-of-document token
|
||||
eod_mask_loss: false # Mask loss for the end of document tokens
|
||||
validation_drop_last: true # Set to false if the last partial validation samples is to be consumed
|
||||
no_seqlen_plus_one_input_tokens: false # Set to True to disable fetching (sequence length + 1) input tokens, instead get (sequence length) input tokens and mask the last token
|
||||
pad_samples_to_global_batch_size: false # Set to True if you want to pad the last partial batch with -1's to equal global batch size
|
||||
shuffle_documents: true # Set to False to disable documents shuffling. Sample index will still be shuffled
|
||||
|
||||
# Nsys profiling options
|
||||
nsys_profile:
|
||||
enabled: false
|
||||
start_step: 0 # Global batch to start profiling
|
||||
end_step: 1 # Global batch to end profiling
|
||||
ranks: [0] # Global rank IDs to profile
|
||||
gen_shape: false # Generate model and kernel details including input shapes
|
||||
|
||||
memory_profile:
|
||||
enabled: false
|
||||
start_step: 0
|
||||
end_step: 1
|
||||
ranks: [0]
|
||||
output_path: /data # Must be a dir
|
||||
|
||||
optim:
|
||||
name: distributed_fused_adam # E.g., fused_adam or set _target_: torch.optim.AdamW field
|
||||
lr: 2e-5
|
||||
weight_decay: 0.01
|
||||
betas:
|
||||
- 0.9
|
||||
- 0.98
|
||||
bucket_cap_mb: 125
|
||||
overlap_grad_sync: true
|
||||
overlap_param_sync: true
|
||||
contiguous_grad_buffer: true
|
||||
contiguous_param_buffer: true
|
||||
sched:
|
||||
name: CosineAnnealing
|
||||
warmup_steps: 400
|
||||
constant_steps: 0
|
||||
min_lr: 2e-6
|
||||
-26
@@ -1,26 +0,0 @@
|
||||
# Copyright 2024 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
steps:
|
||||
- name: 'gcr.io/cloud-builders/docker'
|
||||
args:
|
||||
- 'build'
|
||||
- '--tag=${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'
|
||||
- '--file=docker/vertex-dist-recipes.Dockerfile'
|
||||
- '.'
|
||||
automapSubstitutions: true
|
||||
env:
|
||||
- 'DOCKER_BUILDKIT=1'
|
||||
images:
|
||||
- '${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'
|
||||
-41
@@ -1,41 +0,0 @@
|
||||
diff --git a/nemo/collections/nlp/parts/megatron_trainer_builder.py b/nemo/collections/nlp/parts/megatron_trainer_builder.py
|
||||
index b2c85cde4..a3a9670c3 100644
|
||||
--- a/nemo/collections/nlp/parts/megatron_trainer_builder.py
|
||||
+++ b/nemo/collections/nlp/parts/megatron_trainer_builder.py
|
||||
@@ -19,6 +19,7 @@ from lightning_fabric.utilities.exceptions import MisconfigurationException
|
||||
from omegaconf import DictConfig
|
||||
from pytorch_lightning import Trainer
|
||||
from pytorch_lightning.callbacks import ModelSummary
|
||||
+from pytorch_lightning.callbacks import Callback
|
||||
from pytorch_lightning.plugins.environments import TorchElasticEnvironment
|
||||
|
||||
from nemo.collections.common.metrics.perf_metrics import FLOPsMeasurementCallback
|
||||
@@ -38,6 +39,23 @@ from nemo.utils.callbacks.dist_ckpt_io import (
|
||||
AsyncFinalizerCallback,
|
||||
DistributedCheckpointIO,
|
||||
)
|
||||
+from vmg.util.device_stats import gpu_stats_str
|
||||
+
|
||||
+class GpuStatsMon(Callback):
|
||||
+ def on_train_start(self, trainer, pl_module) -> None:
|
||||
+ rank=pl_module.global_rank
|
||||
+ print(f'train_start: {rank=} {gpu_stats_str()}', flush=True)
|
||||
+
|
||||
+ def on_train_batch_start(self, trainer, pl_module, batch, batch_idx) -> None:
|
||||
+ rank=pl_module.global_rank
|
||||
+ print(f'batch_start: {rank=} {gpu_stats_str()}', flush=True)
|
||||
+
|
||||
+ def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None:
|
||||
+ rank=pl_module.global_rank
|
||||
+ print(f'batch_end: {rank=} {gpu_stats_str()}', flush=True)
|
||||
|
||||
|
||||
class MegatronTrainerBuilder:
|
||||
@@ -178,6 +196,7 @@ class MegatronTrainerBuilder:
|
||||
if self.cfg.get('exp_manager', {}).get('log_tflops_per_sec_per_gpu', True):
|
||||
callbacks.append(FLOPsMeasurementCallback(self.cfg))
|
||||
|
||||
+ callbacks.append(GpuStatsMon())
|
||||
return callbacks
|
||||
|
||||
def create_trainer(self, callbacks=None) -> Trainer:
|
||||
-41
@@ -1,41 +0,0 @@
|
||||
diff -ruN old-datasets/blended_megatron_dataset_builder.py datasets/blended_megatron_dataset_builder.py
|
||||
--- old-datasets/blended_megatron_dataset_builder.py 2025-05-02 04:08:45.369199665 +0000
|
||||
+++ datasets/blended_megatron_dataset_builder.py 2025-05-02 04:10:47.369119891 +0000
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
import logging
|
||||
import math
|
||||
+import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Callable, Iterable, List, Optional, Type, Union
|
||||
|
||||
@@ -353,7 +354,7 @@
|
||||
num_dataset_builder_threads = self.config.num_dataset_builder_threads
|
||||
|
||||
if torch.distributed.is_initialized():
|
||||
- rank = torch.distributed.get_rank()
|
||||
+ rank = int(os.getenv("LOCAL_RANK", "0"))
|
||||
# First, build on rank 0
|
||||
if rank == 0:
|
||||
num_workers = num_dataset_builder_threads
|
||||
@@ -475,7 +476,7 @@
|
||||
Optional[Union[DistributedDataset, Iterable]]: The DistributedDataset instantion, the Iterable instantiation, or None
|
||||
"""
|
||||
if torch.distributed.is_initialized():
|
||||
- rank = torch.distributed.get_rank()
|
||||
+ rank = int(os.getenv("LOCAL_RANK", "0"))
|
||||
|
||||
dataset = None
|
||||
|
||||
diff -ruN old-datasets/gpt_dataset.py datasets/gpt_dataset.py
|
||||
--- old-datasets/gpt_dataset.py 2025-05-02 04:08:45.369199665 +0000
|
||||
+++ datasets/gpt_dataset.py 2025-05-02 04:09:30.309170278 +0000
|
||||
@@ -351,7 +351,7 @@
|
||||
|
||||
if not path_to_cache or (
|
||||
not cache_hit
|
||||
- and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
|
||||
+ and (not torch.distributed.is_initialized() or int(os.getenv("LOCAL_RANK", "0")) == 0)
|
||||
):
|
||||
|
||||
log_single_rank(
|
||||
-13
@@ -1,13 +0,0 @@
|
||||
diff --git a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
|
||||
index 8da15148d..005cae6c9 100644
|
||||
--- a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
|
||||
+++ b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
|
||||
@@ -104,6 +104,8 @@ def convert(input_nemo_file, output_hf_file, precision=None, cpu_only=False) ->
|
||||
dummy_trainer = Trainer(devices=1, accelerator='cpu', strategy=NLPDDPStrategy())
|
||||
model_config = MegatronGPTModel.restore_from(input_nemo_file, trainer=dummy_trainer, return_config=True)
|
||||
model_config.tensor_model_parallel_size = 1
|
||||
+ model_config.virtual_pipeline_model_parallel_size = None
|
||||
+ model_config.sequence_parallel = False
|
||||
model_config.pipeline_model_parallel_size = 1
|
||||
if cpu_only:
|
||||
map_location = torch.device('cpu')
|
||||
-24
@@ -1,24 +0,0 @@
|
||||
diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
|
||||
index bfe8ea359..dfeaf93b5 100644
|
||||
--- a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
|
||||
+++ b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
|
||||
@@ -13,6 +13,8 @@
|
||||
# limitations under the License.
|
||||
|
||||
import torch.multiprocessing as mp
|
||||
+import torch.distributed as dist
|
||||
+
|
||||
from omegaconf.omegaconf import OmegaConf
|
||||
|
||||
from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel
|
||||
@@ -76,6 +78,10 @@ def main(cfg) -> None:
|
||||
|
||||
trainer.fit(model)
|
||||
|
||||
+ if dist.is_available() and dist.is_initialized():
|
||||
+ dist.barrier()
|
||||
+ dist.destroy_process_group()
|
||||
+
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
-13
@@ -1,13 +0,0 @@
|
||||
diff --git a/src/utils/training_metrics/process_training_results.py b/src/utils/training_metrics/process_training_results.py
|
||||
index 3e82a66..e61e1d8 100644
|
||||
--- a/src/utils/training_metrics/process_training_results.py
|
||||
+++ b/src/utils/training_metrics/process_training_results.py
|
||||
@@ -134,7 +134,7 @@ def get_average_step_time(file: str, start_step: int, end_step: int) -> float:
|
||||
for line in datajson:
|
||||
if line.get("step") != "PARAMETER":
|
||||
step = line.get("step")
|
||||
- if step >= start_step and step <= end_step:
|
||||
+ if step >= start_step and step <= end_step and "train_step_timing in s" in line["data"]:
|
||||
time_step_accumulator += line["data"].get("train_step_timing in s")
|
||||
num_steps += 1
|
||||
if num_steps == 0:
|
||||
-10
@@ -1,10 +0,0 @@
|
||||
dllogger@git+https://github.com/NVIDIA/dllogger@v1.0.0
|
||||
|
||||
# Fixing these libraries versions to avoid conflicting or broken packages.
|
||||
immutabledict==4.2.1
|
||||
protobuf==4.25.8
|
||||
opencv-python-headless==4.11.0.86
|
||||
docutils==0.16
|
||||
urllib3==2.5.0
|
||||
google-cloud-storage==3.0.0
|
||||
retrying
|
||||
-18
@@ -1,18 +0,0 @@
|
||||
# cuml-cu12==24.8.0 was installed in nemo:24.09
|
||||
# Removing cuml=24.4.0 to avoid conflicting packages.
|
||||
cudf==24.4.0
|
||||
cugraph==24.4.0
|
||||
cugraph-service-server==24.4.0
|
||||
cuml==24.4.0
|
||||
dask-cudf==24.4.0
|
||||
raft-dask==24.4.0
|
||||
cugraph-dgl==24.4.0
|
||||
cugraph-pyg==24.4.0
|
||||
# The following packages are removed temporarily to avoid conflicting packages
|
||||
# and can be brought back if needed.
|
||||
tensorrt-llm==0.12.0
|
||||
img2dataset==1.45.0
|
||||
Sphinx==8.1.3
|
||||
sphinxcontrib-bibtex==2.6.3
|
||||
torchx==0.7.0
|
||||
nemo-run
|
||||
-66
@@ -1,66 +0,0 @@
|
||||
# Dockerfile wrapping NeMo.
|
||||
#
|
||||
# To workaround base nemo docker image using too many layers, we use Multi-stage
|
||||
# build to first collect the additional files we'll need.
|
||||
FROM alpine:latest AS prep_files
|
||||
WORKDIR /workspace
|
||||
RUN mkdir -p configs vdt vdt/util
|
||||
COPY scripts/*.py vdt/
|
||||
COPY scripts/util/*.py vdt/util/
|
||||
COPY configs/* configs/
|
||||
COPY docker/patches/24.09/* vdt/patches/
|
||||
RUN chmod a+rwX -R vdt
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Available tags
|
||||
# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags
|
||||
# It installs NeMo source code in /opt/NeMo folder, with tag=r2.0.0
|
||||
FROM nvcr.io/nvidia/nemo:24.09
|
||||
|
||||
RUN apt-get update && apt-get install -y sudo zsh tmux && \
|
||||
rm -rf /var/lib/apt/lists*
|
||||
|
||||
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | \
|
||||
tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
|
||||
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | \
|
||||
apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
|
||||
apt-get update -y && apt-get install google-cloud-sdk -y && \
|
||||
rm -rf /var/lib/apt/lists*
|
||||
|
||||
# Install libraries with pip
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
|
||||
# We expect this will be run in the root directory of the vertex-dist-recipes repo
|
||||
ARG HOST_SRC_DIR="."
|
||||
|
||||
# The pre-installed NeMo introduces a lot of deps conflicts.
|
||||
# We uninstall the confilicting libs and reinstall some of them as needed.
|
||||
COPY ${HOST_SRC_DIR}/docker/uninstall.txt /tmp/uninstall.txt
|
||||
RUN cat /tmp/uninstall.txt | grep -v '#' | xargs pip uninstall -y
|
||||
COPY ${HOST_SRC_DIR}/docker/requirements.txt /tmp/requirements.txt
|
||||
RUN pip install -r /tmp/requirements.txt
|
||||
|
||||
# Make sure there's no inconsistent pip libraries.
|
||||
RUN pip check
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Copy configs
|
||||
COPY ${HOST_SRC_DIR}/configs/* /opt/NeMo/examples/nlp/language_modeling/conf/
|
||||
|
||||
# Copy all additional files we need from `prep_files` image.
|
||||
COPY --from=prep_files /workspace/ .
|
||||
|
||||
# Install for `src/utils/training_metrics/process_training_results.py` to report
|
||||
# throughput and MFU numbers.
|
||||
RUN git clone https://github.com/AI-Hypercomputer/gpu-recipes.git
|
||||
|
||||
# This hack is needed for multi-node training while not using a sharing file system.
|
||||
RUN patch --verbose -l -d /opt/megatron-lm/megatron/core/datasets -p1 -i /workspace/vdt/patches/local_rank.patch; \
|
||||
git -C /workspace/gpu-recipes apply /workspace/vdt/patches/throughput_calc.patch; \
|
||||
git -C /opt/NeMo apply /workspace/vdt/patches/nemo2hf.patch; \
|
||||
git -C /opt/NeMo apply /workspace/vdt/patches/sigabort.patch;
|
||||
# git -C /opt/NeMo apply /workspace/vdt/patches/gpu_stats.patch;
|
||||
|
||||
# Do not put an entrypoint here. Specify the entrypoint in the docker run script.
|
||||
-16
@@ -1,16 +0,0 @@
|
||||
{
|
||||
"project_id": "<your_project_id>",
|
||||
"region": "us-central1",
|
||||
"zone": "us-central1-c",
|
||||
"bucket": "<your_bucket",
|
||||
"dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
|
||||
"image_uri": "<your_image_uri>",
|
||||
"strategy": "spot",
|
||||
"nodes": "2",
|
||||
"machine_type": "a3-megagpu-8g",
|
||||
"gpu_type": "NVIDIA_H100_MEGA_80GB",
|
||||
"gpus_per_node": "8",
|
||||
"recipe_name": "llama3_1_8b_pretrain_a3mega",
|
||||
"job_prefix": "vertex-ai",
|
||||
"reservation_name": ""
|
||||
}
|
||||
-49
@@ -1,49 +0,0 @@
|
||||
absl-py==2.2.2
|
||||
annotated-types==0.7.0
|
||||
anyio==4.9.0
|
||||
black==25.1.0
|
||||
cachetools==5.5.2
|
||||
certifi==2025.4.26
|
||||
charset-normalizer==3.4.2
|
||||
click==8.1.8
|
||||
docstring_parser==0.16
|
||||
google-api-core==2.24.2
|
||||
google-auth==2.40.1
|
||||
google-cloud-aiplatform==1.92.0
|
||||
google-cloud-bigquery==3.31.0
|
||||
google-cloud-core==2.4.3
|
||||
google-cloud-resource-manager==1.14.2
|
||||
google-cloud-storage==2.19.0
|
||||
google-crc32c==1.7.1
|
||||
google-genai==1.14.0
|
||||
google-resumable-media==2.7.2
|
||||
googleapis-common-protos==1.70.0
|
||||
grpc-google-iam-v1==0.14.2
|
||||
grpcio==1.71.0
|
||||
grpcio-status==1.71.0
|
||||
h11==0.16.0
|
||||
httpcore==1.0.9
|
||||
httpx==0.28.1
|
||||
idna==3.10
|
||||
mypy_extensions==1.1.0
|
||||
numpy==2.2.5
|
||||
packaging==25.0
|
||||
pathspec==0.12.1
|
||||
platformdirs==4.3.8
|
||||
proto-plus==1.26.1
|
||||
protobuf==5.29.4
|
||||
pyasn1==0.6.1
|
||||
pyasn1_modules==0.4.2
|
||||
pydantic==2.11.4
|
||||
pydantic_core==2.33.2
|
||||
python-dateutil==2.9.0.post0
|
||||
pytz==2025.2
|
||||
requests==2.32.4
|
||||
rsa==4.9.1
|
||||
shapely==2.1.0
|
||||
six==1.17.0
|
||||
sniffio==1.3.1
|
||||
typing-inspection==0.4.0
|
||||
typing_extensions==4.13.2
|
||||
urllib3==2.4.0
|
||||
websockets==15.0.1
|
||||
-173
@@ -1,173 +0,0 @@
|
||||
"""Launch script for Vertex distributed training"""
|
||||
|
||||
# Copy the sample_job_config.json file to job_config.json
|
||||
# to define the job parameters.
|
||||
#
|
||||
# Run like this:
|
||||
#
|
||||
# python3 vertex_dist_train/launch.py --config_file=job_config.json
|
||||
#
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import pprint
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, List
|
||||
|
||||
from absl import app, flags
|
||||
from google.cloud import aiplatform
|
||||
from google.cloud.aiplatform_v1.types.custom_job import Scheduling
|
||||
from pytz import timezone
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
flags.DEFINE_string("config_file", None, "Path to JSON config file")
|
||||
flags.DEFINE_boolean(
|
||||
"debug", False, "Debug mode: just print the command, don't run it."
|
||||
)
|
||||
|
||||
|
||||
def launch_job(
|
||||
job_name: str,
|
||||
project: str,
|
||||
region: str,
|
||||
gcs_bucket: str,
|
||||
image_uri: str,
|
||||
entrypoint_cmd: List[str],
|
||||
trainer_args: List[Any],
|
||||
num_nodes: int,
|
||||
machine_type: str,
|
||||
num_gpus_per_node: int,
|
||||
gpu_type: str,
|
||||
strategy: str,
|
||||
reservation_name: str = "",
|
||||
):
|
||||
assert strategy in ("dws", "spot", "reservation")
|
||||
aiplatform.init(
|
||||
project=project, location=region, staging_bucket=gcs_bucket
|
||||
)
|
||||
|
||||
train_job = aiplatform.CustomContainerTrainingJob(
|
||||
display_name=job_name,
|
||||
container_uri=image_uri,
|
||||
command=entrypoint_cmd,
|
||||
)
|
||||
|
||||
job_args = dict(
|
||||
args=trainer_args,
|
||||
enable_web_access=True,
|
||||
replica_count=num_nodes,
|
||||
machine_type=machine_type,
|
||||
accelerator_type=gpu_type,
|
||||
accelerator_count=num_gpus_per_node,
|
||||
boot_disk_size_gb=1000,
|
||||
restart_job_on_worker_restart=True,
|
||||
#restart_job_on_worker_restart=False,
|
||||
)
|
||||
|
||||
if strategy == "spot":
|
||||
job_args.update({"scheduling_strategy": Scheduling.Strategy.SPOT.name})
|
||||
elif strategy == "dws":
|
||||
job_args.update(
|
||||
{"scheduling_strategy": Scheduling.Strategy.FLEX_START.name}
|
||||
)
|
||||
elif strategy == "reservation":
|
||||
assert reservation_name != "", (
|
||||
"If using a reservation, provide the reservation_name in the "
|
||||
"format `projects/{project_id_or_number}/zones/{zone}/"
|
||||
"reservations/{reservation_name}`"
|
||||
)
|
||||
job_args.update(
|
||||
{
|
||||
"reservation_affinity_type": "SPECIFIC_RESERVATION",
|
||||
"reservation_affinity_key": "compute.googleapis.com/reservation-name",
|
||||
"reservation_affinity_values": [reservation_name],
|
||||
}
|
||||
)
|
||||
|
||||
pprint.pprint(job_args)
|
||||
if not FLAGS.debug:
|
||||
train_job.submit(**job_args)
|
||||
|
||||
|
||||
def main(argv: Sequence[str]) -> None:
|
||||
config_file_path = FLAGS.config_file
|
||||
print(f"Reading job config from {config_file_path}")
|
||||
with open(config_file_path, encoding="utf-8") as config_file:
|
||||
config = json.load(config_file)
|
||||
|
||||
project_id = config["project_id"]
|
||||
region = config["region"]
|
||||
zone = config["zone"]
|
||||
bucket = config["bucket"]
|
||||
dataset_bucket = config["dataset_bucket"]
|
||||
n_nodes = int(config["nodes"])
|
||||
machine_type = config["machine_type"]
|
||||
num_gpus_per_node = int(config["gpus_per_node"])
|
||||
gpu_type = config["gpu_type"]
|
||||
reservation_name = config.get("reservation_name")
|
||||
reservation_full_name = (
|
||||
f"projects/{project_id}/zones/{zone}/reservations/{reservation_name}"
|
||||
if "reservation_name" in config
|
||||
else ""
|
||||
)
|
||||
|
||||
strategy = config["strategy"]
|
||||
recipe_name = config["recipe_name"]
|
||||
job_prefix = config["job_prefix"]
|
||||
image_uri = config["image_uri"]
|
||||
|
||||
# Job name
|
||||
timestamp = (
|
||||
datetime.datetime.now()
|
||||
.astimezone(timezone("US/Pacific"))
|
||||
.strftime("%Y%m%d_%H%M%S")
|
||||
)
|
||||
job_name = f"{recipe_name}-{timestamp}"
|
||||
if job_prefix:
|
||||
job_name = f"{job_prefix}-{job_name}"
|
||||
|
||||
base_output_dir = os.path.join("/gcs", bucket, job_name)
|
||||
|
||||
# Training command and args
|
||||
entrypoint_cmd = ["python3", "vdt/run.py"]
|
||||
|
||||
dataset_bucket = f"gs://{config['dataset_bucket']}"
|
||||
|
||||
trainer_args = [
|
||||
f"--train_data_gcs={dataset_bucket}",
|
||||
"/opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py",
|
||||
"--config-path=conf/",
|
||||
f"--config-name={recipe_name}.yaml",
|
||||
f"exp_manager.explicit_log_dir={base_output_dir}",
|
||||
f"exp_manager.dllogger_logger_kwargs.json_file={base_output_dir}/dllogger.json",
|
||||
"+exp_manager.create_tensorboard_logger=true",
|
||||
"exp_manager.create_checkpoint_callback=false",
|
||||
f"trainer.num_nodes={n_nodes}",
|
||||
f"trainer.devices={num_gpus_per_node}",
|
||||
"trainer.max_steps=10",
|
||||
"trainer.log_every_n_steps=1",
|
||||
"model.tokenizer.vocab_file=/data/gpt2-vocab.json",
|
||||
"model.tokenizer.merge_file=/data/gpt2-merges.txt",
|
||||
"model.data.data_prefix=[1.0,/data/hfbpe_gpt_training_data_text_document]",
|
||||
]
|
||||
|
||||
launch_job(
|
||||
job_name=job_name,
|
||||
project=project_id,
|
||||
region=region,
|
||||
gcs_bucket=bucket,
|
||||
image_uri=image_uri,
|
||||
entrypoint_cmd=entrypoint_cmd,
|
||||
trainer_args=trainer_args,
|
||||
num_nodes=n_nodes,
|
||||
machine_type=machine_type,
|
||||
num_gpus_per_node=num_gpus_per_node,
|
||||
gpu_type=gpu_type,
|
||||
strategy=strategy,
|
||||
reservation_name=reservation_full_name,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app.run(main)
|
||||
-85
@@ -1,85 +0,0 @@
|
||||
"""Entrypoint for Vertex Distributed Training container."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from subprocess import STDOUT, check_output, run
|
||||
|
||||
from absl import app, flags, logging
|
||||
from util import cluster_spec
|
||||
|
||||
from retrying import retry
|
||||
|
||||
# PyTorch barrier call which synchronizes all of the nodes before launching the training process.
|
||||
# This makes sure that processes will block until all processes are ready.
|
||||
# Improves the reliability of spot VM usage for multi-node training jobs
|
||||
|
||||
@retry(stop_max_attempt_number=100, wait_exponential_multiplier=1000)
|
||||
def barrier_with_retry() -> None:
|
||||
import torch
|
||||
logging.info("Starting barrier on RANK {}".format(os.environ["RANK"]))
|
||||
torch.distributed.init_process_group()
|
||||
torch.distributed.barrier()
|
||||
torch.distributed.destroy_process_group()
|
||||
logging.info("Finished barrier on RANK {}".format(os.environ["RANK"]))
|
||||
|
||||
def main(unused_argv: Sequence[str]) -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--train_data_gcs",
|
||||
type=str,
|
||||
help="Download training data from gcs path",
|
||||
)
|
||||
args, unknown = parser.parse_known_args()
|
||||
|
||||
for key, val in os.environ.items():
|
||||
logging.info("ENV %s=%s", key, val)
|
||||
|
||||
if args.train_data_gcs:
|
||||
local_dir = "/data"
|
||||
if not os.path.exists(local_dir):
|
||||
os.mkdir(local_dir)
|
||||
logging.info("downloading %s to %s...", args.train_data_gcs, local_dir)
|
||||
check_output(
|
||||
[
|
||||
"gcloud",
|
||||
"storage",
|
||||
"cp",
|
||||
"-r",
|
||||
f"{args.train_data_gcs}/*",
|
||||
local_dir,
|
||||
],
|
||||
stderr=STDOUT,
|
||||
)
|
||||
logging.info("%s downloaded.", args.train_data_gcs)
|
||||
|
||||
primary_node_addr, primary_node_port, node_rank, num_nodes = (
|
||||
cluster_spec.get_cluster_spec()
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"torchrun",
|
||||
"--nproc-per-node=8",
|
||||
f"--nnodes={num_nodes}",
|
||||
f"--node_rank={node_rank}",
|
||||
]
|
||||
if num_nodes > 1:
|
||||
cmd += [
|
||||
"--max-restarts=3",
|
||||
"--rdzv-backend=static",
|
||||
f'--rdzv_id={os.getenv("CLOUD_ML_JOB_ID", primary_node_port)}',
|
||||
f"--rdzv-endpoint={primary_node_addr}:{primary_node_port}",
|
||||
]
|
||||
cmd += unknown
|
||||
|
||||
logging.info("launching with cmd: \n%s", " \\\n".join(cmd))
|
||||
barrier_with_retry()
|
||||
run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.get_absl_handler().python_handler.stream = sys.stdout
|
||||
app.run(
|
||||
main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True)
|
||||
)
|
||||
-81
@@ -1,81 +0,0 @@
|
||||
"""Get cluster info from environment variables."""
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import os
|
||||
|
||||
from absl import logging
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ClusterInfo:
|
||||
"""Contains information about the cluster.
|
||||
|
||||
Attributes:
|
||||
primary_node_addr: The address of the primary node.
|
||||
primary_node_port: The port of the primary node.
|
||||
node_rank: The rank of the node.
|
||||
num_nodes: The number of nodes in the cluster.
|
||||
"""
|
||||
|
||||
primary_node_addr: str | None = None
|
||||
primary_node_port: str | None = None
|
||||
node_rank: int = 0
|
||||
num_nodes: int = 1
|
||||
|
||||
# Allows unpacking operation like
|
||||
# primary_node_addr, primary_node_port, _, _ = ClusterInfo()
|
||||
# See https://stackoverflow.com/a/70753113
|
||||
def __iter__(self):
|
||||
return iter(dataclasses.astuple(self))
|
||||
|
||||
|
||||
def get_cluster_spec() -> ClusterInfo:
|
||||
"""Parses CLUSTER_SPEC environment variable and returns the cluster info.
|
||||
|
||||
Returns:
|
||||
A ClusterInfo object.
|
||||
"""
|
||||
cluster_spec = os.getenv("CLUSTER_SPEC", None)
|
||||
|
||||
# If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
|
||||
if not cluster_spec:
|
||||
cluster_info = ClusterInfo(
|
||||
primary_node_addr=os.getenv("MASTER_ADDR", None),
|
||||
primary_node_port=os.getenv("MASTER_PORT", None),
|
||||
node_rank=int(os.getenv("RANK", "0")),
|
||||
num_nodes=int(os.getenv("NNODES", "1")),
|
||||
)
|
||||
return cluster_info
|
||||
|
||||
cluster_data = json.loads(cluster_spec)
|
||||
# Get primary node info
|
||||
primary_node = cluster_data["cluster"]["workerpool0"][0]
|
||||
logging.info("primary node: %s", primary_node)
|
||||
primary_node_addr, primary_node_port = primary_node.split(":")
|
||||
logging.info("primary node address: %s", primary_node_addr)
|
||||
logging.info("primary node port: %s", primary_node_port)
|
||||
|
||||
# Determine node rank of this machine
|
||||
workerpool = cluster_data["task"]["type"]
|
||||
if workerpool == "workerpool0":
|
||||
node_rank = 0
|
||||
elif workerpool == "workerpool1":
|
||||
# Add 1 for the primary node, since `index` is the index of workerpool1.
|
||||
node_rank = cluster_data["task"]["index"] + 1
|
||||
else:
|
||||
raise ValueError(
|
||||
"Only workerpool0 and workerpool1 are supported. Unknown workerpool:"
|
||||
f" {workerpool}"
|
||||
)
|
||||
logging.info("node rank: %s", node_rank)
|
||||
|
||||
# Calculate total nodes.
|
||||
num_nodes = 1 # For the primary node.
|
||||
if "workerpool1" in cluster_data["cluster"]:
|
||||
num_nodes += len(cluster_data["cluster"]["workerpool1"])
|
||||
logging.info("num nodes: %s", num_nodes)
|
||||
|
||||
return ClusterInfo(
|
||||
primary_node_addr, primary_node_port, node_rank, num_nodes
|
||||
)
|
||||
-59
@@ -1,59 +0,0 @@
|
||||
"""Add tests for cluster_spec.py."""
|
||||
|
||||
import os
|
||||
|
||||
from . import cluster_spec
|
||||
|
||||
|
||||
# TODO(styer): Use pytest instead
|
||||
class ClusterSpecTest(googletest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.curr_env_var = os.environ.copy()
|
||||
|
||||
def tearDown(self):
|
||||
super().tearDown()
|
||||
os.environ = self.curr_env_var
|
||||
|
||||
def test_get_cluster_spec_from_env_vars(self):
|
||||
os.environ["CLUSTER_SPEC"] = ""
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = "8080"
|
||||
os.environ["RANK"] = "0"
|
||||
os.environ["NNODES"] = "2"
|
||||
cluster_info = cluster_spec.get_cluster_spec()
|
||||
self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
|
||||
self.assertEqual(cluster_info.primary_node_port, "8080")
|
||||
self.assertEqual(cluster_info.node_rank, 0)
|
||||
self.assertEqual(cluster_info.num_nodes, 2)
|
||||
|
||||
def test_get_cluster_spec_from_cluster_spec(self):
|
||||
os.environ[
|
||||
"CLUSTER_SPEC"
|
||||
] = """
|
||||
{
|
||||
"cluster": {
|
||||
"workerpool0": [
|
||||
"127.0.0.1:8080"
|
||||
],
|
||||
"workerpool1": [
|
||||
"127.0.0.2:8080",
|
||||
"127.0.0.3:8080"
|
||||
]
|
||||
},
|
||||
"task": {
|
||||
"type": "workerpool1",
|
||||
"index": 0
|
||||
}
|
||||
}
|
||||
"""
|
||||
cluster_info = cluster_spec.get_cluster_spec()
|
||||
self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
|
||||
self.assertEqual(cluster_info.primary_node_port, "8080")
|
||||
self.assertEqual(cluster_info.node_rank, 1)
|
||||
self.assertEqual(cluster_info.num_nodes, 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
googletest.main()
|
||||
@@ -1,15 +0,0 @@
|
||||
# Vertex AI custom prediction routines samples
|
||||
|
||||
## Overview
|
||||
Vertex Custom Prediction Routines(CPR) simplify the process of building custom containers
|
||||
and make local model testing easy. Here are the sameple codes for different libraries.
|
||||
|
||||
|
||||
### Objectives
|
||||
The objective is to provide various samples for Vertex Custom Prediction Routine(CPR).
|
||||
|
||||
|
||||
### Supporting libraries
|
||||
* torch
|
||||
* sklearn
|
||||
* xgboost
|
||||
@@ -1,33 +0,0 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import pickle
|
||||
|
||||
from google.cloud.aiplatform.constants import prediction
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from sklearn.datasets import load_breast_cancer
|
||||
from sklearn.linear_model import RidgeClassifier
|
||||
|
||||
class LinearRegressionPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists(prediction.MODEL_FILENAME_PKL):
|
||||
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
|
||||
else:
|
||||
self._model = RidgeClassifier()
|
||||
X, y = load_breast_cancer(return_X_y=True)
|
||||
self._model.fit(X, y)
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> np.ndarray:
|
||||
instances = prediction_input["instances"]
|
||||
return np.asarray(instances)
|
||||
|
||||
def predict(self, instances: np.ndarray) -> np.ndarray:
|
||||
return self._model.predict(instances)
|
||||
|
||||
def postprocess(self, prediction_results: np.ndarray) -> dict:
|
||||
return {"predictions": prediction_results.tolist()}
|
||||
@@ -1,33 +0,0 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import pickle
|
||||
|
||||
from google.cloud.aiplatform.constants import prediction
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from sklearn.datasets import make_blobs
|
||||
from sklearn.linear_model import LinearRegression
|
||||
|
||||
class LinearRegressionPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists(prediction.MODEL_FILENAME_PKL):
|
||||
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
|
||||
else:
|
||||
self._model = LogisticRegression()
|
||||
X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=1)
|
||||
self._model.fit(X, y)
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> np.ndarray:
|
||||
instances = prediction_input["instances"]
|
||||
return np.asarray(instances)
|
||||
|
||||
def predict(self, instances: np.ndarray) -> np.ndarray:
|
||||
return self._model.predict_proba(instances)
|
||||
|
||||
def postprocess(self, prediction_results: np.ndarray) -> dict:
|
||||
return {"predictions": prediction_results.tolist()}
|
||||
@@ -1,33 +0,0 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import pickle
|
||||
|
||||
from google.cloud.aiplatform.constants import prediction
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from sklearn.linear_model import SGDClassifier
|
||||
|
||||
class SGDClassifierPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists(prediction.MODEL_FILENAME_PKL):
|
||||
self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
|
||||
else:
|
||||
self._model = SGDClassifier(max_iter=5)
|
||||
X = [[0., 0.], [1., 1.]]
|
||||
y = [0, 1]
|
||||
self._model.fit(X, y)
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> np.ndarray:
|
||||
instances = prediction_input["instances"]
|
||||
return np.asarray(instances)
|
||||
|
||||
def predict(self, instances: np.ndarray) -> np.ndarray:
|
||||
return self._model.predict(instances)
|
||||
|
||||
def postprocess(self, prediction_results: np.ndarray) -> dict:
|
||||
return {"predictions": prediction_results.tolist()}
|
||||
@@ -1,34 +0,0 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from torchvision.models import detection, resnet50, ResNet50_Weights
|
||||
from typing import Dict, List
|
||||
|
||||
class ResNetPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists("model.pth.tar"):
|
||||
self.model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
|
||||
stat_dic = torch.load("model.pth.tar")
|
||||
self.model.load_state_dict(stat_dic['state_dict'])
|
||||
else:
|
||||
weights = ResNet50_Weights.DEFAULT
|
||||
self.model = resnet50(weights=weights)
|
||||
self.model.eval()
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> torch.Tensor:
|
||||
instances = prediction_input["instances"]
|
||||
return torch.Tensor(instances)
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, instances: torch.Tensor) -> List[str]:
|
||||
return self._model(instances)
|
||||
|
||||
def postprocess(self, prediction_results: List[str]) -> Dict:
|
||||
return {"predictions": prediction_results}
|
||||
@@ -1,73 +0,0 @@
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
import pickle
|
||||
import torch
|
||||
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from transformers import AutoModelForQuestionAnswering
|
||||
from typing import Dict, List
|
||||
|
||||
class TorchTransformersPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
|
||||
if os.path.isfile("setup_config.json"):
|
||||
with open("setup_config.json") as setup_config_file:
|
||||
self.setup_config = json.load(setup_config_file)
|
||||
|
||||
if os.path.exists("model.pt"):
|
||||
self.model = AutoModelForQuestionAnswering.from_pretrained("model.pt")
|
||||
self.model.eval()
|
||||
else:
|
||||
raise ValueError("One of the following model files must be provided: model.pt.")
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> torch.Tensor:
|
||||
max_length = self.setup_config["max_length"]
|
||||
instances = prediction_input["instances"]
|
||||
question_context = ast.literal_eval(instances)
|
||||
question = question_context["question"]
|
||||
context = question_context["context"]
|
||||
inputs = self.tokenizer.encode_plus(
|
||||
question,
|
||||
context,
|
||||
max_length=int(max_length),
|
||||
pad_to_max_length=True,
|
||||
add_special_tokens=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = inputs["input_ids"]
|
||||
attention_mask = inputs["attention_mask"]
|
||||
return torch.Tensor(input_ids, attention_mask)
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, instances: torch.Tensor) -> List[str]:
|
||||
input_ids, attention_mask = instances
|
||||
outputs = self._model(input_ids, attention_mask)
|
||||
answer_start_scores = outputs.start_logits
|
||||
answer_end_scores = outputs.end_logits
|
||||
|
||||
num_rows, num_cols = answer_start_scores.shape
|
||||
inferences = []
|
||||
for i in range(num_rows):
|
||||
answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)
|
||||
answer_start = torch.argmax(answer_start_scores_one_seq)
|
||||
answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)
|
||||
answer_end = torch.argmax(answer_end_scores_one_seq) + 1
|
||||
prediction = self.tokenizer.convert_tokens_to_string(
|
||||
self.tokenizer.convert_ids_to_tokens(
|
||||
input_ids[i].tolist()[answer_start:answer_end]
|
||||
)
|
||||
)
|
||||
inferences.append(prediction)
|
||||
return inferences
|
||||
|
||||
def postprocess(self, prediction_results: List[str]) -> Dict:
|
||||
return {"predictions": prediction_results}
|
||||
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pickle
|
||||
import xgboost as xgb
|
||||
|
||||
from google.cloud.aiplatform.constants import prediction
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
from sklearn.datasets import make_blobs
|
||||
from xgboost import XGBClassifier
|
||||
|
||||
|
||||
class ClassifierPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists(prediction.MODEL_FILENAME_PKL):
|
||||
booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
|
||||
else:
|
||||
X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=1)
|
||||
model = XGBClassifier()
|
||||
model.fit(X, y)
|
||||
booster = model.get_booster()
|
||||
self._booster = booster
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
|
||||
instances = prediction_input["instances"]
|
||||
return xgb.DMatrix(instances)
|
||||
|
||||
def predict(self, instances: xgb.DMatrix) -> np.ndarray:
|
||||
return self._booster.predict(instances)
|
||||
|
||||
def postprocess(self, prediction_results: np.ndarray) -> dict:
|
||||
return {"predictions": prediction_results.tolist()}
|
||||
@@ -1,41 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pickle
|
||||
import xgboost as xgb
|
||||
|
||||
from google.cloud.aiplatform.constants import prediction
|
||||
from google.cloud.aiplatform.utils import prediction_utils
|
||||
from google.cloud.aiplatform.prediction.predictor import Predictor
|
||||
|
||||
class XGBRankerPredictor(Predictor):
|
||||
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def load(self, artifacts_uri: str) -> None:
|
||||
prediction_utils.download_model_artifacts(artifacts_uri)
|
||||
if os.path.exists(prediction.MODEL_FILENAME_PKL):
|
||||
booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
|
||||
self._booster = booster
|
||||
else:
|
||||
N = 500
|
||||
dates = pd.date_range(start='2023-01-01', end='2023-01-12', periods=N)
|
||||
X = pd.DataFrame(np.random.randn(N, 5), columns=list('ABCDE'), index=dates)
|
||||
y = pd.Series(np.random.randint(0, 10, size=N), index=dates, name='label')
|
||||
group = X.groupby(dates + pd.offsets.MonthEnd(0)).size()
|
||||
sample_weight = pd.Series(np.arange(len(group)), index=group.index)
|
||||
model = xgb.XGBRanker(objective='rank:pairwise', max_depth=3, learning_rate=0.1, booster='gbtree', tree_method='hist', n_jobs=4, n_estimators=50, enable_categorical=False, random_state=42)
|
||||
model.fit(X=X, y=y, group=group, sample_weight=sample_weight, verbose=True)
|
||||
booster = model.get_booster()
|
||||
self._booster = booster
|
||||
|
||||
def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
|
||||
instances = prediction_input["instances"]
|
||||
return xgb.DMatrix(instances)
|
||||
|
||||
def predict(self, instances: xgb.DMatrix) -> np.ndarray:
|
||||
return self._booster.predict(instances, output_margin=False, ntree_limit=0)
|
||||
|
||||
def postprocess(self, prediction_results: np.ndarray) -> dict:
|
||||
return {"predictions": prediction_results.tolist()}
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
# Benchmark report on fine tuning the OpenLLaMA 7B model on Google Cloud Vertex Model Garden
|
||||
|
||||
Gary Wei, Software Engineer, Google Cloud
|
||||
Dustin Luong, Software Engineer, Google Cloud
|
||||
Changyu Zhu, Software Engineer, Google Cloud
|
||||
Genquan Duan, Software Engineer, Google Cloud
|
||||
|
||||
## Introduction
|
||||
|
||||
Fine-tuning of LLMs can be non-trivial to find an optimal configuration of
|
||||
machine types, training parameters, and other hyperparameters that achieves a
|
||||
good balance between cost efficiency and model performance. To facilitate users
|
||||
in conducting tuning experiments, this report benchmarks OpenLLaMA 7B
|
||||
fine-tuning on Google Cloud Vertex Model Garden, demonstrating both efficiency
|
||||
and effectiveness. The observations are general and can be applied to other LLM
|
||||
models.
|
||||
|
||||
We benchmarked fine tuning algorithms [LoRA](https://arxiv.org/abs/2106.09685)
|
||||
and [QLoRA](https://arxiv.org/abs/2305.14314) supported by
|
||||
[huggingface PEFT libraries](https://github.com/huggingface/peft). LoRA, short
|
||||
for Low-Rank Adaptation of Large Language Models, is an improved fine tuning
|
||||
method where instead of fine tuning all the weights that constitute the weight
|
||||
matrix of the pre-trained large language model, two smaller matrices that
|
||||
approximate this larger matrix are fine-tuned. QLoRA is an even more
|
||||
memory-efficient version of LoRA, where the pretrained model is loaded to GPU
|
||||
memory as quantized 4-bit weights, while preserving similar effectiveness to
|
||||
LoRA. We also provide simple scripts and parameter settings to reproduce the
|
||||
results reported in this report.
|
||||
|
||||
In general, there are many factors that affect the performance of fine-tuning
|
||||
experiments, such as hardware settings, parameters, cost, and accuracy. It is
|
||||
impractical to obtain benchmarks for all possible combinations of these factors.
|
||||
Instead, we focus on tuning a subset of related parameters and evaluating their
|
||||
impact on a set of chosen metrics. The evaluation metrics are GPU memory usage,
|
||||
percentage of parameters tuned, tuning speed, cost, and accuracy. The tuning
|
||||
parameters are batch size, lora rank, maximum sequence length, and maximum
|
||||
training steps.
|
||||
|
||||
## Key takeaways
|
||||
|
||||
- **Use QLoRA to minimize the peak GPU requirements**: The QLoRA can
|
||||
significantly reduce the peak GPU memory usage by ~75% compared to LoRA. For
|
||||
OpenLLaMA7b, the peak memory is ~28G for LoRA and ~7G for QLoRA.
|
||||
- **Use LoRA to maximize the tuning speed and minimize the tuning cost**: LoRA
|
||||
is ~66% faster than QLoRA in fine tuning speed. LoRA/QLoRA tuning cost is
|
||||
low generally, while LoRA is even ~40% cheaper than QLoRA with the same
|
||||
parameters. Suggest to use QLoRA for limited GPU memories, and LoRA for
|
||||
limited training budgets. For OpenLLaMA7b, the tuning speed for LoRA/QLoRA
|
||||
~5 samples / 3 samples per second, and the tuning cost for LoRA/QLoRA in 500
|
||||
steps is ~$1/$1.7 on `a2-highgpu-1g` with 1 A100 40G GPU. The tuning cost
|
||||
for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1 V100 GPU, while LoRA
|
||||
could not run because of OOM.
|
||||
- **Use QLoRA to tune models with large sequence lengths**. For OpenLLaMA7b,
|
||||
the max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
|
||||
while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
|
||||
encounter OOM when max sequence length is 1024.
|
||||
- **Both LoRA and QLoRA give similar accuracy improvement after fine tuning.**
|
||||
For OpenLLaMA7b, both LoRA/QLoRA can improve the average accuracy by ~4%
|
||||
evaluating on 3 typical tasks (ARC challenge, HellaSwag and TruthfulQA),
|
||||
after training 1875 steps on dataset
|
||||
[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
|
||||
- **Use a big batch size if GPU memory is not a constraint**. For OpenLLaMA7b
|
||||
with other default parameters, we suggest using a batch size as 24 for
|
||||
QLoRA, but 2 for LoRA when tuning with 1 A100 40G. We also suggest using a
|
||||
batch size as 8 for QLoRA when tuning with 1 V100. Tuning with LoRA and
|
||||
batch size as 1 got OOM and we don't recommend tuning LoRA with 1 V100.
|
||||
|
||||
## Benchmark Details
|
||||
|
||||
### Experiment Setup
|
||||
|
||||
The benchmark dataset is
|
||||
[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
|
||||
The training dataset is directly downloaded from hugging face to the VM, before
|
||||
every experiment.
|
||||
|
||||
The default tuning parameters during benchmark are:
|
||||
|
||||
- Host VM: a2-highgpu-1g
|
||||
- Accelerator type: 1 A100 40G
|
||||
- batch size: 2
|
||||
- lora_rank: 16
|
||||
- max_seq_length: 512
|
||||
- precision_mode: float16
|
||||
- max_train_steps: 500
|
||||
|
||||
For simplicity, we set the precision mode to `float16` when tuning LoRA models,
|
||||
and set the precision to `4bit` for QLoRA.
|
||||
|
||||
Sample script to start fine tuning dockers in a VM on GCP.
|
||||
|
||||
```shell
|
||||
IMAGE_TAG=us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:latest
|
||||
docker run --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0 \
|
||||
--rm --name "test_gpu" -it --pull=always ${IMAGE_TAG} \
|
||||
--task=instruct-lora \
|
||||
--pretrained_model_id=openlm-research/open_llama_7b \
|
||||
--dataset_name="timdettmers/openassistant-guanaco" \
|
||||
--instruct_column_in_dataset="text" \
|
||||
--precision_mode="float16" \
|
||||
--output_dir=<OUTPUT DIR> \
|
||||
--lora_rank=2 \
|
||||
--max_sequence_length=512 \
|
||||
--learning_rate=2e-4 \
|
||||
--max_steps=50
|
||||
```
|
||||
|
||||
### GPU Memory
|
||||
|
||||
In this benchmark, we investigated the impact of batch size, lora rank, and
|
||||
maximum sequence length on GPU memory, and then made recommendations on the
|
||||
maximum batch size for different GPUs.
|
||||
|
||||
#### Peak GPU memory by batch size (GB)
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-batch-size.png" width="600">
|
||||
|
||||
- The QLoRA can significantly reduce the peak GPU memory usage by ~75%
|
||||
compared to LoRA. The peak GPU memory is ~28G for LoRA and ~7G for QLoRA
|
||||
when batch size is 2.
|
||||
- QLoRA can support much larger batch sizes than LoRA
|
||||
- We can use a batch size as 32 for QLoRA, but only 2 for LoRA on 1 A100
|
||||
40G.
|
||||
- We can use a batch size of 8 for QLoRA on 1 V100 GPU. LoRA will fail
|
||||
with OOM even with a batch size of 1.
|
||||
|
||||
#### Peak GPU memory by LoRA rank (GB)
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-lora-rank.png" width="600">
|
||||
|
||||
- Peak GPU memories are quite similar for different LoRA ranks for both
|
||||
LoRA/QLoRA.
|
||||
- The peak GPU memory increasing percentages are very small generally when
|
||||
LoRA rank increases.
|
||||
- The peak GPU memory increases from 28G with LoRA rank 4 to 29.09G with
|
||||
LoRA rank 64, and the increasing percentage is only ~3.9%.
|
||||
|
||||
#### Peak GPU memory by max sequence length for LoRA/QLoRA (GB)
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-max-seq-length.png" width="600">
|
||||
|
||||
- The peak GPU increases quickly when max sequence length increases for both
|
||||
LoRA/QLoRA, and the increasing rate of LoRA is much faster than QLoRA.
|
||||
- For LoRA tuning, the GPU memory increased from 20.5G (max sequence
|
||||
length=256) to 28.2G (max sequence length=512), an increase of ~37%.
|
||||
- For QLoRA tuning, the GPU memory increased from 6.94G (max sequence
|
||||
length=256) to 7.57G (max sequence length=512), an increase of ~9%.
|
||||
- The max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
|
||||
while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
|
||||
encounter OOM when max sequence length is 1024.
|
||||
|
||||
### Fine Tuning Parameters
|
||||
|
||||
This section shows the number/percentage of trainable parameters, and the sizes
|
||||
of the fine tuned models. LoRA and QLoRA differ only in how they represent the
|
||||
precision of their parameters. The total number of parameters and the number of
|
||||
trainable parameters are the same for both methods.
|
||||
|
||||
| LoRA Rank | Finetuned parameters | Total parameters | Trainable Parameter Percentage | Fine tuned model size (MB) |
|
||||
| --------- | -------------------- | ---------------- | ------------------------------ | -------------------------- |
|
||||
| 8 | 2.00E+07 | 6.76E+09 | 0.3% | 76.4 |
|
||||
| 16 | 4.00E+07 | 6.78E+09 | 0.6% | 152.65 |
|
||||
| 32 | 8.00E+07 | 6.82E+09 | 1.2% | 305.15 |
|
||||
| 64 | 1.60E+08 | 6.90E+09 | 2.3% | 610.15 |
|
||||
|
||||
|
||||
LoRA/QLoRA tunes quite a small fraction (only 0.3% with LoRA rank=8) of all
|
||||
parameters, and the tuned models are very small (only 76.4MB with LoRA rank=8).
|
||||
|
||||
### Fine Tuning Speed And Costs
|
||||
|
||||
The fine-tuning speed and cost are affected by various factors, such as the
|
||||
GPUs, LoRA ranks, and max sequence lengths.
|
||||
|
||||
- LoRA is ~66% faster than QLoRA in fine tuning speed. The tuning speed for
|
||||
LoRA/QLoRA ~5 samples / 3 samples per second on 1 A100 40G GPU
|
||||
- Higher LoRA ranks, slower tuning speed for both LoRA/QLoRA.
|
||||
- LoRA tuning speed reduces from ~5 samples per second with LoRA rank as 8
|
||||
to ~4 samples per second with LoRA rank as 64, slowed down by 20%.
|
||||
- QLoRA tuning speed reduces from ~3 samples per second with LoRA rank as
|
||||
8 to ~2.5 samples per second with LoRA rank as 64, slowed down by 17%.
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-vs-lora-rank.png" width="600">
|
||||
|
||||
- Longer sequence lengths, slower tuning speed.
|
||||
- LoRA tuning speed reduces from ~5.56 samples per second with max
|
||||
sequence length as 256 to ~4.84 samples per second with max sequence
|
||||
length as 512 slowed down by 13%.
|
||||
- LoRA tuning speed reduces from ~2.95 samples per second with max
|
||||
sequence length as 256 to ~2.88 samples per second with max sequence
|
||||
length as 512 slowed down by ~2.4%.
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-lora-qlora.png" width="600">
|
||||
|
||||
- LoRA/QLoRA tuning cost is low generally, while LoRA is even ~40% cheaper
|
||||
than QLoRA with the same parameters.
|
||||
- The LoRA/QLoRA fine tuning cost for 500 steps is ~$1/$1.7 on 1 A100 40G.
|
||||
- The tuning cost for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1
|
||||
V100 GPU, while LoRA could not run because of OOM.
|
||||
|
||||
<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-cost-lora-qlora.png" width="600">
|
||||
|
||||
### Accuracy
|
||||
|
||||
We fine tuned Open Llama 7B model with
|
||||
[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco),
|
||||
and report accuracy similar to the
|
||||
[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
|
||||
using
|
||||
[Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
|
||||
[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
|
||||
mainly compares models on ARC, HellaSwag, MMLU, and TruthfulQA. The authors did
|
||||
not publish OpenLLaMA 7B on MMLU
|
||||
([link](https://huggingface.co/openlm-research/open_llama_7b)). Therefore, we
|
||||
only benchmark accuracies on ARC, HellaSwag, and TruthfulQA.
|
||||
|
||||
| | Mean | ARC | HellaSwag | TruthfulQA | Tuning Parameters |
|
||||
| ------------------------------------------------------------ | ---- | ---- | --------- | ---------- | ------------------------------------------------------------ |
|
||||
| OpenLLaMA7B ([Original Report](https://huggingface.co/openlm-research/open_llama_7b)) | 0.49 | 0.41 | 0.73 | 0.34 | n/a |
|
||||
| OpenLLaMA7B ([Re-run with lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)) | 0.51 | 0.47 | 0.72 | 0.35 | n/a |
|
||||
| OpenLLaMA7B+LoRA | 0.56 | 0.48 | 0.74 | 0.45 | LoRA Rank=16; Max Sequence Length=512;Learning Rate=1e-4; Train steps=1875 |
|
||||
| OpenLLaMA7B+QLoRA | 0.53 | 0.45 | 0.73 | 0.42 | LoRA Rank=16; Max Sequence Length=512; Learning Rate=1e-4; Train steps=1875 |
|
||||
|
||||
- The base OpenLLaMA7B model gets better performance (2%) when using the
|
||||
[Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
|
||||
- LoRA/QLoRA can improve the performance by ~2-4% when trained for 1875 steps
|
||||
with learning rate 1e-4.
|
||||
-50
@@ -1,50 +0,0 @@
|
||||
# Dockerfile for serving dockers with AutoGluon.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/autogluon/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
|
||||
|
||||
USER root
|
||||
|
||||
# AutoGluon might require libgomp for some dependencies.
|
||||
RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
|
||||
curl \
|
||||
wget \
|
||||
vim \
|
||||
libgomp1
|
||||
|
||||
# Install AutoGluon and other dependencies.
|
||||
RUN pip install --upgrade pip
|
||||
RUN pip install autogluon==1.0.0
|
||||
RUN pip install flask==3.0.0
|
||||
|
||||
# Dependencies needed to work with GCS.
|
||||
RUN pip install absl-py==2.0.0
|
||||
RUN pip install google-cloud-storage==2.7.0
|
||||
|
||||
# Copy scripts into the container.
|
||||
COPY model_oss/autogluon /autogluon
|
||||
COPY model_oss/util /autogluon/util
|
||||
WORKDIR /autogluon
|
||||
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
RUN wget https://github.com/pallets/flask/blob/main/LICENSE.rst
|
||||
|
||||
# Expose the port the app runs on.
|
||||
EXPOSE 8501
|
||||
|
||||
# Set the working directory to a specific path for consistency.
|
||||
WORKDIR /autogluon
|
||||
|
||||
# Change to a non-root user for security purposes.
|
||||
RUN useradd -m autogluonuser
|
||||
USER autogluonuser
|
||||
|
||||
# Run Flask application.
|
||||
CMD ["python", "serve.py"]
|
||||
-36
@@ -1,36 +0,0 @@
|
||||
# Dockerfile for training dockers with Autogluon.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/autogluon/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
|
||||
|
||||
# Install tools.
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
|
||||
RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
|
||||
apt-utils \
|
||||
curl \
|
||||
wget \
|
||||
git \
|
||||
jq \
|
||||
gnupg \
|
||||
build-essential \
|
||||
tesseract-ocr \
|
||||
vim
|
||||
|
||||
# Install libraries.
|
||||
RUN pip install autogluon==1.0.0
|
||||
|
||||
COPY model_oss/autogluon /autogluon
|
||||
WORKDIR /autogluon
|
||||
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
ENTRYPOINT ["python", "train.py"]
|
||||
@@ -1,87 +0,0 @@
|
||||
r"""AutoGluon serving binary.
|
||||
|
||||
This module sets up a Flask web server for serving predictions from a
|
||||
trained AutoGluon model. The server exposes two endpoints:
|
||||
|
||||
1. `/ping`: A health check endpoint that returns "pong" to
|
||||
indicate that the server is running.
|
||||
2. `/predict`: An endpoint that accepts POST requests with JSON content.
|
||||
Each request should contain one or more instances for which the
|
||||
predictions are desired. The endpoint returns the predictions and
|
||||
associated probabilities in a JSON response.
|
||||
|
||||
The server expects an environment variable `model_path` that points to
|
||||
the directory where the AutoGluon model artifacts are
|
||||
stored. If `model_path` is not provided, it defaults to '/autogluon/models'.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
from autogluon.tabular import TabularPredictor
|
||||
import flask
|
||||
import pandas as pd
|
||||
|
||||
from util import constants
|
||||
from util import fileutils
|
||||
|
||||
_SUCCESS_STATUS = 200
|
||||
_ERROR_STATUS = 500
|
||||
_PORT = 8501
|
||||
|
||||
app = flask.Flask(__name__)
|
||||
# Check the environment variables.
|
||||
model_dir = os.getenv('model_path', '/autogluon/models')
|
||||
logging.info('Model directory passed by the user is: %s', model_dir)
|
||||
# If the model is on GCS then copy it to a local folder first.
|
||||
if model_dir.startswith(constants.GCS_URI_PREFIX):
|
||||
gcs_path = model_dir[len(constants.GCS_URI_PREFIX) :]
|
||||
local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
|
||||
logging.info('Download %s to %s', model_dir, local_model_dir)
|
||||
fileutils.download_gcs_dir_to_local(model_dir, local_model_dir)
|
||||
model_dir = local_model_dir
|
||||
logging.info('Local model directory is: %s', model_dir)
|
||||
|
||||
|
||||
# Load the predictor at startup.
|
||||
predictor = TabularPredictor.load(model_dir)
|
||||
|
||||
|
||||
@app.route('/ping', methods=['GET'])
|
||||
def ping() -> flask.Response:
|
||||
"""Health check route."""
|
||||
return flask.Response('pong', status=_SUCCESS_STATUS)
|
||||
|
||||
|
||||
@app.route('/predict', methods=['POST'])
|
||||
def predict() -> flask.Response:
|
||||
"""Prediction route."""
|
||||
try:
|
||||
# Extract JSON content from the POST request.
|
||||
data = flask.request.get_json(force=True)
|
||||
instances = data.get('instances', [])
|
||||
|
||||
# Convert instances to DataFrame.
|
||||
df_to_predict = pd.DataFrame(instances)
|
||||
|
||||
# Perform prediction.
|
||||
predictions = predictor.predict(df_to_predict).tolist()
|
||||
response = {'predictions': predictions}
|
||||
|
||||
return flask.Response(
|
||||
json.dumps(response),
|
||||
status=_SUCCESS_STATUS,
|
||||
mimetype='application/json',
|
||||
)
|
||||
|
||||
except Exception as e: # pylint: disable=broad-exception-caught
|
||||
return flask.Response(
|
||||
json.dumps({'error': str(e)}),
|
||||
status=_ERROR_STATUS,
|
||||
mimetype='application/json',
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(host='0.0.0.0', port=_PORT)
|
||||
@@ -1,144 +0,0 @@
|
||||
"""AutoGluon training binary. """
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from autogluon.tabular import TabularPredictor
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class BaseConfig:
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
key: value for key, value in self.__dict__.items() if value is not None
|
||||
}
|
||||
|
||||
|
||||
class DataConfig(BaseConfig):
|
||||
|
||||
def __init__(self, train_data_path: Any) -> None:
|
||||
self.train_data_path = train_data_path
|
||||
|
||||
|
||||
class ProblemConfig(BaseConfig):
|
||||
|
||||
def __init__(self, label: Any, problem_type: Any) -> None:
|
||||
self.label = label
|
||||
self.problem_type = problem_type
|
||||
|
||||
|
||||
class EvaluationConfig(BaseConfig):
|
||||
|
||||
def __init__(self, eval_metric: Any) -> None:
|
||||
self.eval_metric = eval_metric
|
||||
|
||||
|
||||
class TrainingConfig(BaseConfig):
|
||||
"""Config for training."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
time_limit: Any,
|
||||
presets: Any,
|
||||
hyperparameters: Any,
|
||||
model_save_path: str,
|
||||
) -> None:
|
||||
self.time_limit = time_limit
|
||||
self.hyperparameters = hyperparameters
|
||||
self.presets = presets
|
||||
self.model_save_path = model_save_path
|
||||
|
||||
|
||||
def parse_args() -> (
|
||||
tuple[DataConfig, ProblemConfig, EvaluationConfig, TrainingConfig]
|
||||
):
|
||||
"""Parse command line arguments."""
|
||||
parser = argparse.ArgumentParser(description="AutoGluon Tabular Predictor")
|
||||
# Add arguments for each config class
|
||||
parser.add_argument(
|
||||
"--train_data_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to the input data CSV file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--label", type=str, required=True, help="Target variable column name."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--problem_type",
|
||||
type=str,
|
||||
choices=["binary", "multiclass", "regression", "quantile"],
|
||||
default=None,
|
||||
help="Problem type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval_metric", type=str, default=None, help="Evaluation metric to use."
|
||||
)
|
||||
# Add arguments for TrainingConfig if needed
|
||||
parser.add_argument(
|
||||
"--time_limit",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Time limit in seconds for training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--presets",
|
||||
type=str,
|
||||
default="medium_quality",
|
||||
help="Presets used for training ",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hyperparameters",
|
||||
type=json.loads,
|
||||
default=None,
|
||||
help="Hyperparameter dictionary in JSON format.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_save_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to save the trained model.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
data_config = DataConfig(train_data_path=args.train_data_path)
|
||||
problem_config = ProblemConfig(
|
||||
label=args.label, problem_type=args.problem_type
|
||||
)
|
||||
eval_config = EvaluationConfig(eval_metric=args.eval_metric)
|
||||
training_config = TrainingConfig(
|
||||
time_limit=args.time_limit,
|
||||
presets=args.presets,
|
||||
hyperparameters=args.hyperparameters,
|
||||
model_save_path=args.model_save_path,
|
||||
)
|
||||
|
||||
return data_config, problem_config, eval_config, training_config
|
||||
|
||||
|
||||
def main() -> None:
|
||||
data_config, problem_config, eval_config, training_config = parse_args()
|
||||
|
||||
# Load the training data.
|
||||
data = pd.read_csv(data_config.train_data_path)
|
||||
|
||||
# Create a TabularPredictor.
|
||||
predictor = TabularPredictor(
|
||||
label=problem_config.label,
|
||||
eval_metric=eval_config.eval_metric,
|
||||
path=training_config.model_save_path,
|
||||
)
|
||||
|
||||
# Fit the model
|
||||
predictor.fit(
|
||||
data,
|
||||
presets=training_config.presets,
|
||||
time_limit=training_config.time_limit,
|
||||
hyperparameters=training_config.hyperparameters,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,25 +0,0 @@
|
||||
# The provided content is a configuration file for the ZipNeRF
|
||||
# PyTorch implementation.
|
||||
|
||||
# Sets the name of the experiment to 'test'.
|
||||
Config.exp_name = 'test'
|
||||
# Specifies the dataset loader, in this case, 'llff' for light field.
|
||||
Config.dataset_loader = 'llff'
|
||||
# Defines the near and far clipping planes for the camera view.
|
||||
Config.near = 0.2
|
||||
Config.far = 1e6
|
||||
# Image downsampling.
|
||||
Config.factor = 4
|
||||
|
||||
# For the model configurations.
|
||||
Model.raydist_fn = 'power_transformation'
|
||||
Model.opaque_background = True
|
||||
|
||||
# Disables the computation of density normals and RGB values, and sets
|
||||
# the grid level dimension to 1 for PropMLP.
|
||||
PropMLP.disable_density_normals = True
|
||||
PropMLP.disable_rgb = True
|
||||
PropMLP.grid_level_dim = 1
|
||||
|
||||
# Disable density normals for NerfMLP
|
||||
NerfMLP.disable_density_normals = True
|
||||
@@ -1,21 +0,0 @@
|
||||
# The provided content is a configuration file for Generative
|
||||
# Latent Optimization (GLO) vectors in the Pytorch implemnetation of ZipNeRF.
|
||||
|
||||
# Specifies the dataset loader, in this case, 'llff' for light field.
|
||||
Config.dataset_loader = 'llff'
|
||||
# Defines the near and far clipping planes for the camera view.
|
||||
Config.near = 0.2
|
||||
Config.far = 1e6
|
||||
# Image downsampling.
|
||||
Config.factor = 4
|
||||
|
||||
# For the model configurations.
|
||||
Model.raydist_fn = 'power_transformation'
|
||||
Model.num_glo_features = 128
|
||||
Model.opaque_background = True
|
||||
|
||||
PropMLP.disable_density_normals = True
|
||||
PropMLP.disable_rgb = True
|
||||
PropMLP.grid_level_dim = 1
|
||||
|
||||
NerfMLP.disable_density_normals = True
|
||||
-18
@@ -1,18 +0,0 @@
|
||||
# The provided content is a configuration file running ZipNeRF
|
||||
# training on 8 gpu machine.
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
distributed_type: MULTI_GPU
|
||||
downcast_bf16: 'no'
|
||||
gpu_ids: all
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
mixed_precision: fp16
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-120
@@ -1,120 +0,0 @@
|
||||
# Dockerfile for ZipNeRF base image.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_base.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
|
||||
|
||||
USER root
|
||||
|
||||
ARG COLMAP_GIT_COMMIT=main
|
||||
ARG CUDA_ARCHITECTURES=60;70;75;80;86
|
||||
|
||||
# Prevent stop building ubuntu at time zone selection.
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
|
||||
curl \
|
||||
g++ \
|
||||
wget \
|
||||
vim \
|
||||
bash \
|
||||
cmake \
|
||||
imagemagick \
|
||||
ninja-build \
|
||||
build-essential \
|
||||
libboost-program-options-dev \
|
||||
libboost-filesystem-dev \
|
||||
libboost-graph-dev \
|
||||
libboost-system-dev \
|
||||
libeigen3-dev \
|
||||
libflann-dev \
|
||||
libfreeimage-dev \
|
||||
libmetis-dev \
|
||||
libgoogle-glog-dev \
|
||||
libgtest-dev \
|
||||
libsqlite3-dev \
|
||||
libglew-dev \
|
||||
qtbase5-dev \
|
||||
libqt5opengl5-dev \
|
||||
libcgal-dev \
|
||||
libceres-dev \
|
||||
git \
|
||||
git-lfs \
|
||||
python3-cffi \
|
||||
python3-cryptography \
|
||||
libffi-dev \
|
||||
python-dev
|
||||
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Install google cloud CLI.
|
||||
RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-430.0.0-linux-x86.tar.gz
|
||||
RUN tar xzf google-cloud-cli-430.0.0-linux-x86.tar.gz
|
||||
RUN ./google-cloud-sdk/install.sh -q
|
||||
# Make sure gsutil will use the default service account.
|
||||
RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
|
||||
|
||||
# Install deps and install gsutil.
|
||||
RUN pip install gsutil==5.27
|
||||
|
||||
# When building colmap in colab, the link error "undefined reference.
|
||||
# to '_glapi_tls_Current'" happens. A solution is to install "libglvnd"
|
||||
# as described in this page https://github.com/colmap/colmap/issues/1271.
|
||||
RUN git clone --depth 1 --branch v1.7.0 https://github.com/NVIDIA/libglvnd && \
|
||||
apt-get install -y libxext-dev libx11-dev x11proto-gl-dev && \
|
||||
cd libglvnd/ && \
|
||||
apt-get install -y autoconf automake libtool && \
|
||||
apt-get install -y libffi-dev && \
|
||||
./autogen.sh && \
|
||||
./configure && \
|
||||
make -j4 && \
|
||||
make install
|
||||
|
||||
RUN apt remove nvidia-cuda-toolkit -y \
|
||||
nvidia-cuda-toolkit \
|
||||
nvidia-cuda-toolkit-gcc
|
||||
|
||||
# Install libraries.
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
RUN python3 -m pip install --upgrade pip
|
||||
|
||||
ENV CUDA_HOME=/usr/local/cuda
|
||||
|
||||
RUN git clone --branch main https://github.com/SuLvXiangXin/zipnerf-pytorch.git
|
||||
# Set current directory to the downloaded 'zipnerf-pytorch' repository.
|
||||
WORKDIR ./zipnerf-pytorch
|
||||
# Using git reset command to pin it down to a specific version.
|
||||
RUN git reset --hard 4de3d21ebb9e15412d36951b56e2d713fddd812b
|
||||
COPY model_oss/cloudnerf/requirements.txt requirements.txt
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
# Install gridencoder extensions and nvdiffrast (for textured mesh).
|
||||
RUN cd .. && \
|
||||
TORCH_CUDA_ARCH_LIST="6.0 7.0 7.5 8.0 8.6+PTX" CXX=g++ pip install ./zipnerf-pytorch/gridencoder
|
||||
|
||||
# Install cuda version of torch_scatter.
|
||||
RUN pip install torch-scatter==2.1.2 -f https://data.pyg.org/whl/torch-2.0.1+cu118.html
|
||||
RUN pip install google-cloud-aiplatform==1.25.0
|
||||
RUN pip install google-cloud-storage==2.9.0
|
||||
|
||||
# Build and install COLMAP.
|
||||
RUN git clone --depth 1 --branch 3.8 https://github.com/colmap/colmap.git
|
||||
RUN cd colmap && \
|
||||
git fetch https://github.com/colmap/colmap.git ${COLMAP_GIT_COMMIT} && \
|
||||
mkdir build && \
|
||||
cd build && \
|
||||
cmake .. -GNinja -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} && \
|
||||
ninja && \
|
||||
ninja install && \
|
||||
cd .. && rm -rf colmap
|
||||
|
||||
RUN git clone --depth 1 --branch v1.0.2 https://github.com/dranjan/python-plyfile.git
|
||||
|
||||
RUN sed -i "20 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py
|
||||
RUN sed -i "21 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py
|
||||
-16
@@ -1,16 +0,0 @@
|
||||
# Dockerfile for ZipNeRF COLMAP image calibration.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/cloudnerf/dockerfile/cloudnerf_pytorch_calibrate.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
|
||||
|
||||
COPY model_oss/cloudnerf/local_colmap_and_resize.sh /workspace/zipnerf-pytorch/scripts/local_colmap_and_resize.sh
|
||||
|
||||
WORKDIR /workspace/zipnerf-pytorch/
|
||||
|
||||
ENTRYPOINT ["bash","scripts/local_colmap_and_resize.sh"]
|
||||
-22
@@ -1,22 +0,0 @@
|
||||
# Dockerfile for ZipNeRF rendering.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_render.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
|
||||
|
||||
COPY model_oss/cloudnerf/render.sh /workspace/zipnerf-pytorch/scripts/render.sh
|
||||
COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
|
||||
COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
|
||||
COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
|
||||
RUN sed -i '324s/.*/ keyframe_names = fp.read().splitlines()/' /workspace/zipnerf-pytorch/internal/camera_utils.py
|
||||
|
||||
ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
|
||||
|
||||
WORKDIR /workspace/zipnerf-pytorch/
|
||||
|
||||
ENTRYPOINT ["bash", "scripts/render.sh"]
|
||||
-21
@@ -1,21 +0,0 @@
|
||||
# Dockerfile for ZipNeRF training.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_train.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
|
||||
|
||||
COPY model_oss/cloudnerf/train.sh /workspace/zipnerf-pytorch/scripts/train.sh
|
||||
COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
|
||||
COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
|
||||
COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
|
||||
|
||||
ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
|
||||
|
||||
WORKDIR /workspace/zipnerf-pytorch/
|
||||
|
||||
ENTRYPOINT ["bash", "scripts/train.sh"]
|
||||
@@ -1,144 +0,0 @@
|
||||
#!/bin/bash
|
||||
# This script runs colmap for scale invariant feature (SIFT) extraction and
|
||||
# matching to map camera extrinsics and intrinsics values for ZipNeRF,
|
||||
# given a folder of images and videos
|
||||
# from a GCS bucket. It uses ffmepg to extract an image from a video at
|
||||
# 1fps. The folder can contain images or videos. If both images and videos
|
||||
# are present, the extracted frames from the videos is added to the images
|
||||
# to create the final combined image dataset.
|
||||
# vv-docker:google3-begin(internal)
|
||||
# TODO(b/314042136): Specify cloudnerf colmap fps.
|
||||
# vv-docker:google3-end
|
||||
|
||||
# Initialize variables.
|
||||
use_gpu=1 # Default to 1 (assuming the docker is run on a machine with GPU)
|
||||
gcs_dataset_path=""
|
||||
gcs_experiment_path=""
|
||||
camera=""
|
||||
|
||||
# This loop processes command-line arguments for configuring the container.
|
||||
# It supports arguments for GPU usage, dataset and experiment paths,
|
||||
# and camera type.
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
-use_gpu)
|
||||
use_gpu="$2"
|
||||
if ! [[ $use_gpu =~ ^[0-9]+$ ]]; then
|
||||
echo "Error: -use_gpu must be an integer."
|
||||
exit 1
|
||||
fi
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-gcs_dataset_path)
|
||||
gcs_dataset_path="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-gcs_experiment_path)
|
||||
gcs_experiment_path="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-camera)
|
||||
camera="$2"
|
||||
if [[ $camera != "OPENCV" && $camera != "OPENCV_FISHEYE" ]]; then
|
||||
echo "Error: -camera must be either 'OPENCV' or 'OPENCV_FISHEYE'."
|
||||
exit 1
|
||||
fi
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
*) # unknown option
|
||||
echo "Unknown option: $1" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
local_folder="dataset_content"
|
||||
images_folder="dataset_images"
|
||||
images_subfolder="images"
|
||||
output_folder="$images_folder/$images_subfolder"
|
||||
|
||||
# Create the local folder if it doesn't exist
|
||||
mkdir -p "$local_folder"
|
||||
mkdir -p "$output_folder"
|
||||
|
||||
# Download the content from the GCS URI
|
||||
gsutil -m cp -r "$gcs_dataset_path"/* "$local_folder/"
|
||||
|
||||
# Process files in the local folder
|
||||
for file in "$local_folder"/*; do
|
||||
if [[ -f "$file" ]]; then
|
||||
# Check if the file is an image (e.g., jpg, png, etc.)
|
||||
if file --mime-type "$file" | grep -q "image"; then
|
||||
# Copy the image to the "images" subfolder within the "dataset_images" folder
|
||||
cp "$file" "$output_folder/$(basename "$file")"
|
||||
elif file --mime-type "$file" | grep -q "video"; then
|
||||
# Use FFmpeg to extract an image every 30 frames from the video
|
||||
ffmpeg -i "$file" -vf "select='not(mod(n,30))'" "$output_folder/$(basename "$file" ."${file##*.}")_%03d.jpg"
|
||||
else
|
||||
echo "Skipping unsupported file: $file"
|
||||
fi
|
||||
fi
|
||||
done
|
||||
|
||||
# Run COLMAP Feature extraction
|
||||
colmap feature_extractor \
|
||||
--database_path "$local_folder"/database.db \
|
||||
--image_path "$output_folder" \
|
||||
--ImageReader.single_camera 1 \
|
||||
--ImageReader.camera_model "$camera" \
|
||||
--SiftExtraction.use_gpu "$use_gpu"
|
||||
|
||||
# Run COLMAP Feature matching
|
||||
colmap exhaustive_matcher \
|
||||
--database_path "$local_folder"/database.db \
|
||||
--SiftMatching.use_gpu "$use_gpu"
|
||||
|
||||
# Bundle adjustment. The default Mapper tolerance is unnecessarily large,
|
||||
# decreasing it speeds up bundle adjustment steps.
|
||||
mkdir -p "$local_folder"/sparse
|
||||
colmap mapper \
|
||||
--database_path "$local_folder"/database.db \
|
||||
--image_path "$output_folder" \
|
||||
--output_path "$local_folder"/sparse \
|
||||
--Mapper.ba_global_function_tolerance=0.000001
|
||||
|
||||
# Downsample images at 1/2, 1/4, 1/8 scales. Save feature matching to
|
||||
# sqlite database.
|
||||
# All input and output images:
|
||||
# $gcs_dataset_path
|
||||
# $gcs_experiment_path/data/images
|
||||
# Downsampled output images:
|
||||
# $gcs_experiment_path/data/images_2/
|
||||
# $gcs_experiment_path/data/images_4/
|
||||
# $gcs_experiment_path/data/images_8/
|
||||
# COLMAP sparse reconstruction files: project.ini, images.bin,
|
||||
# cameras.bin, points3D.bin
|
||||
# $gcs_experiment_path/data/sparse/0/
|
||||
cp -r "$output_folder" "$images_folder"/images_2
|
||||
pushd "$images_folder"/images_2
|
||||
ls | xargs -P 8 -I {} mogrify -resize 50% {}
|
||||
popd
|
||||
gsutil -m cp -r "$images_folder"/images_2/* "$gcs_experiment_path"/data/images_2
|
||||
|
||||
cp -r "$output_folder" "$images_folder"/images_4
|
||||
pushd "$images_folder"/images_4
|
||||
ls | xargs -P 8 -I {} mogrify -resize 25% {}
|
||||
popd
|
||||
gsutil -m cp -r "$images_folder"/images_4/* "$gcs_experiment_path"/data/images_4
|
||||
|
||||
cp -r "$output_folder" "$images_folder"/images_8
|
||||
pushd "$images_folder"/images_8
|
||||
ls | xargs -P 8 -I {} mogrify -resize 12.5% {}
|
||||
popd
|
||||
gsutil -m cp "$images_folder"/images_8/* "$gcs_experiment_path"/data/images_8
|
||||
|
||||
# Copy images and sparse reconstruction files to gcs experiment folder.
|
||||
gsutil -m cp "$images_folder"/images/* "$gcs_experiment_path"/data/images
|
||||
gsutil -m cp -r "$local_folder"/sparse "$gcs_experiment_path"/data
|
||||
gsutil -m cp "$local_folder"/database.db "$gcs_experiment_path"/data
|
||||
|
||||
echo "Processing complete."
|
||||
@@ -1,117 +0,0 @@
|
||||
#!/bin/bash
|
||||
# This script runs rendering for ZipNeRF given an experiment folder
|
||||
# from a GCS bucket with colmap dataset.
|
||||
|
||||
# Initialize associative array for arguments.
|
||||
declare -A args
|
||||
|
||||
# vv-docker:google3-begin(internal)
|
||||
# TODO(b/311468174): Pass gin config file from gcs bucket.
|
||||
# vv-docker:google3-end
|
||||
# Function to parse named arguments.
|
||||
parse_args() {
|
||||
while [[ $# -gt 0 ]]; do
|
||||
key="$1"
|
||||
case $key in
|
||||
-gcs_experiment_path|-gin_config_file|-gcs_keyframes_file)
|
||||
args[$key]="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-training_job_name)
|
||||
training_job_name="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-rendering_job_name)
|
||||
rendering_job_name="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-render_path_frames|-factor|-render_video_fps)
|
||||
args[$key]="$2"
|
||||
if ! [[ ${args[$key]} =~ ^[0-9]+$ ]]; then
|
||||
echo "Error: $key must be an integer."
|
||||
exit 1
|
||||
fi
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
}
|
||||
|
||||
# Function to create a directory if it doesn't exist.
|
||||
create_dir_if_not_exists() {
|
||||
local dir_path=$1
|
||||
if [[ ! -d "$dir_path" ]]; then
|
||||
echo "Creating folder: $dir_path"
|
||||
mkdir "$dir_path"
|
||||
else
|
||||
echo "Folder $dir_path already exists."
|
||||
fi
|
||||
}
|
||||
|
||||
# Function to launch rendering.
|
||||
launch_rendering() {
|
||||
local keyframes_file=$1
|
||||
local render_bindings=(
|
||||
"--gin_configs=${args[-gin_config_file]}"
|
||||
"--gin_bindings=Config.data_dir='${DATASET_PATH}'"
|
||||
"--gin_bindings=Config.exp_name='${EXPERIMENT}'"
|
||||
"--gin_bindings=Config.render_path=True"
|
||||
"--gin_bindings=Config.render_path_frames=${args[-render_path_frames]}"
|
||||
"--gin_bindings=Config.render_video_fps=${args[-render_video_fps]}"
|
||||
"--gin_bindings=Config.factor=${args[-factor]}"
|
||||
)
|
||||
|
||||
if [[ -n $keyframes_file ]]; then
|
||||
render_bindings+=("--gin_bindings=Config.render_spline_keyframes='${keyframes_file}'")
|
||||
fi
|
||||
|
||||
accelerate launch render.py "${render_bindings[@]}"
|
||||
}
|
||||
|
||||
# Parse arguments.
|
||||
parse_args "$@"
|
||||
|
||||
# Extract folder names and paths.
|
||||
scene_folder_name=$(basename "${args[-gcs_experiment_path]}")
|
||||
local_dataset_path="local_dataset"
|
||||
local_experiment_path="exp"
|
||||
exp_folder_name=$(basename "${args[-gcs_experiment_path]}")
|
||||
DATASET_PATH="$local_experiment_path/$exp_folder_name/data"
|
||||
CHECKPOINTS_PATH="$local_experiment_path/$exp_folder_name/checkpoints"
|
||||
OUTPUT_RENDER_PATH="$local_experiment_path/$scene_folder_name/render"
|
||||
EXPERIMENT=$exp_folder_name
|
||||
|
||||
# Create necessary directories.
|
||||
create_dir_if_not_exists "$local_dataset_path"
|
||||
create_dir_if_not_exists "$local_experiment_path"
|
||||
create_dir_if_not_exists "$local_experiment_path/$exp_folder_name"
|
||||
create_dir_if_not_exists "$CHECKPOINTS_PATH"
|
||||
|
||||
# Create the file log_render.txt in the exp folder.
|
||||
touch "$local_experiment_path/$exp_folder_name/log_render.txt"
|
||||
|
||||
# Copy experiment from GCS bucket to local
|
||||
gsutil -m cp -r "${args[-gcs_experiment_path]}/data" "$local_experiment_path/$exp_folder_name" || exit 1
|
||||
gsutil -m cp -r "${args[-gcs_experiment_path]}/checkpoints/${training_job_name}/*" "$CHECKPOINTS_PATH" || exit 1
|
||||
|
||||
# Check and copy keyframes file.
|
||||
if [[ -n ${args[-gcs_keyframes_file]} ]]; then
|
||||
keyframes_file_basename=$(basename "${args[-gcs_keyframes_file]}")
|
||||
local_keyframes_file="$local_dataset_path/$keyframes_file_basename"
|
||||
gsutil cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
|
||||
echo "Local keyframe file: $local_keyframes_file"
|
||||
launch_rendering "$local_keyframes_file"
|
||||
else
|
||||
launch_rendering ""
|
||||
fi
|
||||
|
||||
# Copy rendered data back to GCS.
|
||||
gsutil -m cp -r "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"
|
||||
@@ -1,24 +0,0 @@
|
||||
--find-links https://download.pytorch.org/whl/torch_stable.html
|
||||
|
||||
torch==2.2.0
|
||||
numpy==1.26.1
|
||||
absl_py==2.0.0
|
||||
accelerate==0.24.0
|
||||
gin_config==0.5.0
|
||||
imageio==2.31.6
|
||||
imageio-ffmpeg==0.4.9
|
||||
matplotlib==3.8.0
|
||||
mediapy==1.1.9
|
||||
ninja==1.11.1.1
|
||||
opencv_contrib_python==4.8.1.78
|
||||
opencv_python==4.8.1.78
|
||||
Pillow==10.3.0
|
||||
rawpy==0.18.1
|
||||
scipy==1.11.3
|
||||
scikit-image==0.22.0
|
||||
scikit-learn==1.5.0
|
||||
tensorboard==2.15.0
|
||||
tensorboardX==2.6.2.2
|
||||
tqdm==4.66.3
|
||||
trimesh==4.0.1
|
||||
xatlas==0.0.8
|
||||
@@ -1,94 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Initialize variables.
|
||||
training_job_name=""
|
||||
gcs_experiment_path=""
|
||||
gin_config_file="configs/360.gin"
|
||||
factor=4
|
||||
max_training_steps=25000
|
||||
|
||||
# Parse named arguments.
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
-training_job_name)
|
||||
training_job_name="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-gcs_experiment_path)
|
||||
gcs_experiment_path="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-gin_config_file)
|
||||
gin_config_file="$2"
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-factor)
|
||||
factor="$2"
|
||||
if ! [[ $factor =~ ^[0-9]+$ ]]; then
|
||||
echo "Error: -factor must be an integer."
|
||||
exit 1
|
||||
fi
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
-max_training_steps)
|
||||
max_training_steps="$2"
|
||||
if ! [[ $max_training_steps =~ ^[0-9]+$ ]]; then
|
||||
echo "Error: -max_training_steps must be an integer."
|
||||
exit 1
|
||||
fi
|
||||
shift # past argument
|
||||
shift # past value
|
||||
;;
|
||||
*) # unknown option
|
||||
echo "Unknown option: $1" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Function to create a directory if it doesn't exist.
|
||||
create_dir_if_not_exists() {
|
||||
local dir_path=$1
|
||||
if [[ ! -d "$dir_path" ]]; then
|
||||
echo "Creating folder: $dir_path"
|
||||
mkdir "$dir_path"
|
||||
else
|
||||
echo "Folder $dir_path already exists."
|
||||
fi
|
||||
}
|
||||
|
||||
# Extract folder names and paths.
|
||||
scene_folder_name=$(basename "${gcs_experiment_path}")
|
||||
local_dataset_path="local_dataset"
|
||||
local_experiment_path="exp"
|
||||
DATASET_PATH="$local_experiment_path/$scene_folder_name/data"
|
||||
EXPERIMENT=$scene_folder_name
|
||||
|
||||
# Create necessary directories.
|
||||
create_dir_if_not_exists "$local_dataset_path"
|
||||
create_dir_if_not_exists "$local_experiment_path"
|
||||
create_dir_if_not_exists "$local_experiment_path/$scene_folder_name"
|
||||
|
||||
# Copy experiment from GCS bucket to local.
|
||||
gsutil -m cp -r "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
|
||||
|
||||
echo "GCS Experiment: $gcs_experiment_path"
|
||||
echo "Gin Config File: $gin_config_file"
|
||||
echo "Factor: $factor"
|
||||
echo "Scene: $scene_folder_name"
|
||||
echo "Local Dataset: $DATASET_PATH"
|
||||
echo "Local Experiment: $EXPERIMENT"
|
||||
|
||||
accelerate launch train.py --gin_configs="$gin_config_file" \
|
||||
--gin_bindings="Config.data_dir = '${DATASET_PATH}'" \
|
||||
--gin_bindings="Config.exp_name = '${EXPERIMENT}'" \
|
||||
--gin_bindings="Config.factor = ${factor}" \
|
||||
--gin_bindings="Config.max_steps = ${max_training_steps}"
|
||||
|
||||
gsutil -m rm -r "${gcs_experiment_path}/checkpoints/${training_job_name}"
|
||||
gsutil -m cp -r "$local_experiment_path/$scene_folder_name/config.gin" "${gcs_experiment_path}/${training_job_name}_config.gin"
|
||||
gsutil -m cp -r "$local_experiment_path/$scene_folder_name/checkpoints/*/*" "${gcs_experiment_path}/checkpoints/${training_job_name}"
|
||||
-83
@@ -1,83 +0,0 @@
|
||||
# This Dockerfile converts JAX vision transformer model to
|
||||
# tensorflow saved model format.
|
||||
# Here is an example to build this dockerfile:
|
||||
# PROJECT="your gcp project"
|
||||
# IMAGE_TAG="jax-f-vlm-model-conversion:${USER}-test"
|
||||
# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_model_conversion.Dockerfile . -t "${IMAGE_TAG}"
|
||||
# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
|
||||
|
||||
# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
|
||||
FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install basic libs
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
cmake \
|
||||
curl \
|
||||
wget \
|
||||
sudo \
|
||||
gnupg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
lsb-release \
|
||||
ca-certificates \
|
||||
build-essential \
|
||||
git \
|
||||
libgl1
|
||||
|
||||
# Copy Apache license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Install required libs
|
||||
RUN pip install --upgrade pip
|
||||
|
||||
# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
|
||||
ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
|
||||
RUN git clone -c \
|
||||
remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
|
||||
https://github.com/google-research/google-research --no-checkout --progress \
|
||||
--depth 1
|
||||
WORKDIR ./google-research
|
||||
RUN git sparse-checkout init --cone
|
||||
RUN git sparse-checkout set fvlm
|
||||
RUN git checkout ${COMMIT_ID}
|
||||
|
||||
# The following pip installs are pinned down versions satisfying
|
||||
# fvlm/requirements.txt file.
|
||||
# NOTE: Using `no-deps` flag to avoid overwriting of dependent library
|
||||
# versions. For example, both `chex` and `jax` can overwrite each other's
|
||||
# `jax-lib` version.
|
||||
# Note: The following libraries are pinned down versions of:
|
||||
# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
|
||||
RUN pip install --no-cache-dir tensorflow==2.12.0
|
||||
RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
|
||||
RUN pip install --no-cache-dir numpy==1.23.5
|
||||
RUN pip install --no-cache-dir torch==2.0.1
|
||||
RUN pip install --no-cache-dir torchvision==0.15.2
|
||||
RUN pip install --no-cache-dir opencv-python==4.7.0.72
|
||||
RUN pip install --no-cache-dir tqdm==4.65.0
|
||||
RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
|
||||
RUN pip install --no-cache-dir Pillow==9.5.0
|
||||
RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
|
||||
RUN pip install --no-cache-dir gin-config==0.5.0
|
||||
RUN pip install --no-cache-dir pycocotools==2.0.6
|
||||
RUN pip install --no-cache-dir contextlib2==21.6.0
|
||||
RUN pip install --no-cache-dir ml-collections==0.1.1
|
||||
RUN pip install --no-cache-dir chex==0.1.7
|
||||
RUN pip install --no-cache-dir optax==0.1.5
|
||||
# Dependencies already included. Use no-deps to not update numpy.
|
||||
RUN pip install --no-cache-dir --no-deps flax==0.7.2
|
||||
RUN pip install --no-cache-dir --no-deps clu==0.0.9
|
||||
RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
|
||||
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
|
||||
RUN pip install --no-cache-dir ml-dtypes==0.2.0
|
||||
RUN pip install --no-cache-dir tensorflow_text==2.12.0
|
||||
|
||||
WORKDIR ./fvlm
|
||||
ENV PYTHONPATH ./
|
||||
|
||||
ENTRYPOINT ["python", "export_saved_model.py"]
|
||||
-78
@@ -1,78 +0,0 @@
|
||||
# This Dockerfile trains the F-VLM model on GPU.
|
||||
# Here is an example to build this dockerfile:
|
||||
# PROJECT="your gcp project"
|
||||
# IMAGE_TAG="jax-f-vlm-train:${USER}-test"
|
||||
# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_gpu.Dockerfile . -t "${IMAGE_TAG}"
|
||||
# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
|
||||
# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
|
||||
FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install basic libs
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
cmake \
|
||||
curl \
|
||||
wget \
|
||||
sudo \
|
||||
gnupg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
lsb-release \
|
||||
ca-certificates \
|
||||
build-essential \
|
||||
git
|
||||
|
||||
# Copy Apache license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Install required libs
|
||||
RUN pip install --upgrade pip
|
||||
# The following pip installs are pinned down versions satisfying
|
||||
# fvlm/requirements.txt file.
|
||||
# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
|
||||
# google-research repository.
|
||||
# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
|
||||
ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
|
||||
RUN git clone -c \
|
||||
remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
|
||||
https://github.com/google-research/google-research --no-checkout --progress \
|
||||
--depth 1
|
||||
WORKDIR ./google-research
|
||||
RUN git sparse-checkout init --cone
|
||||
RUN git sparse-checkout set fvlm
|
||||
RUN git checkout ${COMMIT_ID}
|
||||
|
||||
# Note: The following libraries are pinned down versions of:
|
||||
# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
|
||||
RUN pip install --no-cache-dir tensorflow==2.12.0
|
||||
RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
|
||||
RUN pip install --no-cache-dir numpy==1.23.5
|
||||
RUN pip install --no-cache-dir torch==2.0.1
|
||||
RUN pip install --no-cache-dir torchvision==0.15.2
|
||||
RUN pip install --no-cache-dir opencv-python==4.7.0.72
|
||||
RUN pip install --no-cache-dir tqdm==4.65.0
|
||||
RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
|
||||
RUN pip install --no-cache-dir Pillow==9.5.0
|
||||
RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
|
||||
RUN pip install --no-cache-dir gin-config==0.5.0
|
||||
RUN pip install --no-cache-dir pycocotools==2.0.6
|
||||
RUN pip install --no-cache-dir contextlib2==21.6.0
|
||||
RUN pip install --no-cache-dir ml-collections==0.1.1
|
||||
RUN pip install --no-cache-dir chex==0.1.7
|
||||
RUN pip install --no-cache-dir optax==0.1.5
|
||||
# Dependencies already included. Use no-deps to not update numpy.
|
||||
RUN pip install --no-cache-dir --no-deps flax==0.7.2
|
||||
RUN pip install --no-cache-dir --no-deps clu==0.0.9
|
||||
# Installing jax at the very end with GPU support.
|
||||
# NOTE: Not using `no-deps` flag here because we need CUDA support.
|
||||
RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
|
||||
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
|
||||
|
||||
WORKDIR ./fvlm
|
||||
ENV PYTHONPATH ./
|
||||
|
||||
ENTRYPOINT ["python", "train_and_eval.py"]
|
||||
-138
@@ -1,138 +0,0 @@
|
||||
# This Dockerfile trains the F-VLM model on TPU.
|
||||
# Here is an example to build this dockerfile:
|
||||
# PROJECT="your gcp project"
|
||||
# IMAGE_TAG="jax-f-vlm-train-tpu:${USER}-test"
|
||||
# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_tpu.Dockerfile . -t "${IMAGE_TAG}"
|
||||
# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
|
||||
|
||||
FROM python:3.11
|
||||
|
||||
# Get libtpu shared library. See go/what-is-libtpu.
|
||||
RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install basic libs
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
cmake \
|
||||
curl \
|
||||
wget \
|
||||
sudo \
|
||||
gnupg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
lsb-release \
|
||||
ca-certificates \
|
||||
build-essential \
|
||||
git \
|
||||
libgl1
|
||||
|
||||
# Copy Apache license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Install required libs
|
||||
RUN pip install --upgrade pip
|
||||
|
||||
# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
|
||||
# google-research repository.
|
||||
# Using the commit 05ece4b1c97285b48b51fa44321ccb2cb347406a on Dec 11th, 2023.
|
||||
ARG COMMIT_ID=05ece4b1c97285b48b51fa44321ccb2cb347406a
|
||||
RUN git clone -c \
|
||||
remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
|
||||
https://github.com/google-research/google-research --no-checkout --progress \
|
||||
--depth 1
|
||||
WORKDIR ./google-research
|
||||
RUN git sparse-checkout init --cone
|
||||
RUN git sparse-checkout set fvlm
|
||||
RUN git checkout ${COMMIT_ID}
|
||||
|
||||
# Note: The following libraries are pinned down versions of:
|
||||
# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
|
||||
RUN pip install --no-cache-dir ml_dtypes==0.3.1
|
||||
RUN pip install --no-cache-dir tensorstore==0.1.51
|
||||
RUN pip install --no-cache-dir MarkupSafe==2.1.3
|
||||
RUN pip install --no-cache-dir Pillow==9.5.0
|
||||
RUN pip install --no-cache-dir PyYAML==6.0.1
|
||||
RUN pip install --no-cache-dir absl_py==1.4.0
|
||||
RUN pip install --no-cache-dir array_record==0.4.1
|
||||
RUN pip install --no-cache-dir astunparse==1.6.3
|
||||
RUN pip install --no-cache-dir cachetools==5.3.1
|
||||
RUN pip install --no-cache-dir certifi==2023.7.22
|
||||
RUN pip install --no-cache-dir charset_normalizer==3.3.0
|
||||
RUN pip install --no-cache-dir chex==0.1.83
|
||||
RUN pip install --no-cache-dir click==8.1.7
|
||||
RUN pip install --no-cache-dir clip==0.2.0
|
||||
RUN pip install --no-cache-dir clu==0.0.9
|
||||
RUN pip install --no-cache-dir contourpy==1.1.1
|
||||
RUN pip install --no-cache-dir cycler==0.12.1
|
||||
RUN pip install --no-cache-dir dm_tree==0.1.8
|
||||
RUN pip install --no-cache-dir etils==1.5.1
|
||||
RUN pip install --no-cache-dir filelock==3.12.4
|
||||
RUN pip install --no-cache-dir flatbuffers==23.5.26
|
||||
RUN pip install --no-cache-dir flax==0.7.4
|
||||
RUN pip install --no-cache-dir fonttools==4.43.1
|
||||
RUN pip install --no-cache-dir fsspec==2023.9.2
|
||||
RUN pip install --no-cache-dir ftfy==6.1.1
|
||||
RUN pip install --no-cache-dir gast==0.5.4
|
||||
RUN pip install --no-cache-dir gin_config==0.5.0
|
||||
RUN pip install --no-cache-dir google_auth==2.23.3
|
||||
RUN pip install --no-cache-dir google_auth_oauthlib==1.0.0
|
||||
RUN pip install --no-cache-dir google_pasta==0.2.0
|
||||
RUN pip install --no-cache-dir googleapis_common_protos==1.61.0
|
||||
RUN pip install --no-cache-dir grpcio==1.59.0
|
||||
RUN pip install --no-cache-dir h5py==3.10.0
|
||||
RUN pip install --no-cache-dir importlib_resources==6.1.0
|
||||
RUN pip install --no-cache-dir 'jax[tpu]==0.4.18' \
|
||||
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
|
||||
RUN pip install --no-cache-dir jaxlib==0.4.18
|
||||
RUN pip install --no-cache-dir jinja2==3.1.2
|
||||
RUN pip install --no-cache-dir keras==2.14.0
|
||||
RUN pip install --no-cache-dir kiwisolver==1.4.5
|
||||
RUN pip install --no-cache-dir libclang==16.0.6
|
||||
RUN pip install --no-cache-dir markdown==3.5
|
||||
RUN pip install --no-cache-dir matplotlib==3.8.0
|
||||
RUN pip install --no-cache-dir mpmath==1.3.0
|
||||
RUN pip install --no-cache-dir networkx==3.1
|
||||
RUN pip install --no-cache-dir numpy==1.26.0
|
||||
RUN pip install --no-cache-dir nvidia_cublas_cu12==12.1.3.1
|
||||
RUN pip install --no-cache-dir nvidia_cuda_cupti_cu12==12.1.105
|
||||
RUN pip install --no-cache-dir nvidia_cuda_nvrtc_cu12==12.1.105
|
||||
RUN pip install --no-cache-dir nvidia_cuda_runtime_cu12==12.1.105
|
||||
RUN pip install --no-cache-dir nvidia_cudnn_cu12==8.9.2.26
|
||||
RUN pip install --no-cache-dir nvidia_cufft_cu12==11.0.2.54
|
||||
RUN pip install --no-cache-dir nvidia_curand_cu12==10.3.2.106
|
||||
RUN pip install --no-cache-dir nvidia_cusolver_cu12==11.4.5.107
|
||||
RUN pip install --no-cache-dir nvidia_cusparse_cu12==12.1.0.106
|
||||
RUN pip install --no-cache-dir nvidia_nccl_cu12==2.18.1
|
||||
RUN pip install --no-cache-dir nvidia_nvjitlink_cu12==12.2.140
|
||||
RUN pip install --no-cache-dir nvidia_nvtx_cu12==12.1.105
|
||||
RUN pip install --no-cache-dir opencv_python==4.8.1.78
|
||||
RUN pip install --no-cache-dir orbax_checkpoint==0.4.1
|
||||
RUN pip install --no-cache-dir promise==2.3
|
||||
RUN pip install --no-cache-dir protobuf==3.20.3
|
||||
RUN pip install --no-cache-dir psutil==5.9.5
|
||||
RUN pip install --no-cache-dir pyasn1==0.5.0
|
||||
RUN pip install --no-cache-dir pycocotools==2.0.7
|
||||
RUN pip install --no-cache-dir pygments==2.16.1
|
||||
RUN pip install --no-cache-dir regex==2023.10.3
|
||||
RUN pip install --no-cache-dir rich==13.6.0
|
||||
RUN pip install --no-cache-dir scipy==1.11.3
|
||||
RUN pip install --no-cache-dir sympy==1.12
|
||||
RUN pip install --no-cache-dir tensorboard==2.14.1
|
||||
RUN pip install --no-cache-dir tensorboard_data_server==0.7.1
|
||||
RUN pip install --no-cache-dir tensorflow==2.14.0
|
||||
RUN pip install --no-cache-dir tensorflow_datasets==4.9.3
|
||||
RUN pip install --no-cache-dir torch==2.1.0
|
||||
RUN pip install --no-cache-dir torchvision==0.16.0
|
||||
RUN pip install --no-cache-dir urllib3==2.0.6
|
||||
RUN pip install --no-cache-dir wcwidth==0.2.8
|
||||
RUN pip install --no-cache-dir werkzeug==3.0.0
|
||||
RUN pip install --no-cache-dir wheel==0.41.2
|
||||
RUN pip install --no-cache-dir tensorflow_text==2.14.0
|
||||
|
||||
WORKDIR ./fvlm
|
||||
ENV PYTHONPATH ./
|
||||
|
||||
ENTRYPOINT ["python", "train_and_eval.py"]
|
||||
+1
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -1,71 +0,0 @@
|
||||
FROM pytorch/torchserve:0.9.0-gpu
|
||||
|
||||
USER root
|
||||
|
||||
# Install tools.
|
||||
RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
|
||||
curl \
|
||||
wget \
|
||||
vim \
|
||||
git
|
||||
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
ENV INFER_PORT=7080
|
||||
ENV MNG_PORT=7081
|
||||
ENV MODEL_NAME="llava_serving"
|
||||
ENV PATH="/home/model-server/:${PATH}"
|
||||
ENV PATH="/usr/local/cuda-12.1/bin:${PATH}"
|
||||
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
|
||||
ENV NVIDIA_VISIBLE_DEVICES=all
|
||||
|
||||
# Get 'LLaVA' repository from github.
|
||||
RUN git clone https://github.com/haotian-liu/LLaVA /home/model-server/LLaVA
|
||||
WORKDIR /home/model-server/LLaVA
|
||||
# Using git reset command to pin it down to a specific version.
|
||||
RUN git reset --hard 7775b12d6b20cd69089be7a18ea02615a59621cd
|
||||
|
||||
# Install the package.
|
||||
RUN python3 -m pip install --upgrade pip
|
||||
RUN pip install google-cloud-storage==2.13.0
|
||||
RUN pip install absl-py==2.0.0
|
||||
RUN pip install -e .
|
||||
|
||||
# Copy model artifacts.
|
||||
COPY model_oss/llava/handler.py /home/model-server/handler.py
|
||||
COPY model_oss/llava/model_handler_setup.py /home/model-server/model_handler_setup.py
|
||||
COPY model_oss/util/ /home/model-server/util/
|
||||
ENV PYTHONPATH /home/model-server
|
||||
WORKDIR /home/model-server
|
||||
|
||||
# Create torchserve configuration file.
|
||||
RUN echo \
|
||||
"default_response_timeout=1800\n" \
|
||||
"service_envelope=json\n" \
|
||||
"inference_address=http://0.0.0.0:${INFER_PORT}\n" \
|
||||
"management_address=http://0.0.0.0:${MNG_PORT}\n" \
|
||||
"default_workers_per_model=DEFAULT_WORKERS_PER_MODEL" >> /home/model-server/config.properties
|
||||
|
||||
# Expose ports.
|
||||
EXPOSE ${INFER_PORT}
|
||||
EXPOSE ${MNG_PORT}
|
||||
|
||||
# Archive model artifacts and dependencies.
|
||||
# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
|
||||
RUN torch-model-archiver \
|
||||
--model-name=${MODEL_NAME} \
|
||||
--version=1.0 \
|
||||
--handler=/home/model-server/handler.py \
|
||||
--runtime=python3 \
|
||||
--export-path=/home/model-server/model-store \
|
||||
--archive-format=default \
|
||||
--force
|
||||
|
||||
# Run Torchserve HTTP serve to respond to prediction requests.
|
||||
# Use $NUM_GPU workers unless overriden by $TS_NUM_WORKERS
|
||||
CMD ["TOTAL=$(nvidia-smi", "--list-gpus","|","wc","-l)","&&", "TS_NUM_WORKERS=${TS_NUM_WORKERS:-$TOTAL}","&&", "sed","-i","\"s/DEFAULT_WORKERS_PER_MODEL/$TS_NUM_WORKERS/g\"","/home/model-server/config.properties", "&&", \
|
||||
"torchserve", "--start", \
|
||||
"--ts-config", "/home/model-server/config.properties", \
|
||||
"--models", "${MODEL_NAME}=${MODEL_NAME}.mar", \
|
||||
"--model-store", "/home/model-server/model-store"]
|
||||
@@ -1,174 +0,0 @@
|
||||
"""Customer handler for LLava 1.5 OSS model.
|
||||
|
||||
The code is based on here: https://github.com/haotian-liu/LLaVA
|
||||
handler based on:
|
||||
https://github.com/haotian-liu/LLaVA/blob/main/llava/eval/run_llava.py
|
||||
There are two supported variant:
|
||||
1. liuhaotian/llava-v1.5-13b: 13B params
|
||||
2. liuhaotian/llava-v1.5-7b: 7B params
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from llava import constants as llava_constants
|
||||
from llava import conversation
|
||||
from llava import mm_utils
|
||||
from llava.model import builder
|
||||
import model_handler_setup
|
||||
import torch
|
||||
from ts.torch_handler import base_handler
|
||||
|
||||
from util import constants
|
||||
from util import image_format_converter
|
||||
|
||||
|
||||
DEFAULT_MODEL_ID = "liuhaotian/llava-v1.5-7b"
|
||||
|
||||
|
||||
class LlavaHandler(base_handler.BaseHandler):
|
||||
"""Custom handler for LLava model."""
|
||||
|
||||
def initialize(self, context: Any):
|
||||
"""Initializes model, tokenizer, and other components."""
|
||||
self.map_location = model_handler_setup.get_map_location(context=context)
|
||||
self.device = model_handler_setup.get_model_device(
|
||||
map_location=self.map_location, context=context
|
||||
)
|
||||
self.manifest = context.manifest
|
||||
self.model_id = model_handler_setup.get_model_id(
|
||||
default_model_id=DEFAULT_MODEL_ID
|
||||
)
|
||||
|
||||
# Allows 4bit and 8bit quantiziation using BnB nf4.
|
||||
precision = os.environ.get("PRECISION_MODE")
|
||||
load_8bit = precision == constants.PRECISION_MODE_8
|
||||
load_4bit = precision == constants.PRECISION_MODE_4
|
||||
|
||||
self.tokenizer, self.model, self.image_processor, self.context_len = (
|
||||
builder.load_pretrained_model(
|
||||
model_path=self.model_id,
|
||||
model_base=None,
|
||||
model_name=mm_utils.get_model_name_from_path(self.model_id),
|
||||
load_8bit=load_8bit,
|
||||
load_4bit=load_4bit,
|
||||
)
|
||||
)
|
||||
|
||||
def preprocess(self, data: List[Dict[str, Any]]) -> Any:
|
||||
"""Runs the preprocessing to tokenize image and the prompt."""
|
||||
if len(data) > 1:
|
||||
raise ValueError(
|
||||
"LLava original repo currently does not support batch inference."
|
||||
" https://github.com/haotian-liu/LLaVA/issues/754"
|
||||
)
|
||||
data = data[0]
|
||||
prompt, base64_image = data["prompt"], data["base64_image"]
|
||||
|
||||
# Adds proper image token to the prompt.
|
||||
image_token_se = (
|
||||
llava_constants.DEFAULT_IM_START_TOKEN
|
||||
+ llava_constants.DEFAULT_IMAGE_TOKEN
|
||||
+ llava_constants.DEFAULT_IM_END_TOKEN
|
||||
)
|
||||
if llava_constants.IMAGE_PLACEHOLDER in prompt:
|
||||
if self.model.config.mm_use_im_start_end:
|
||||
prompt = re.sub(
|
||||
llava_constants.IMAGE_PLACEHOLDER, image_token_se, prompt
|
||||
)
|
||||
else:
|
||||
prompt = re.sub(
|
||||
llava_constants.IMAGE_PLACEHOLDER,
|
||||
llava_constants.DEFAULT_IMAGE_TOKEN,
|
||||
prompt,
|
||||
)
|
||||
else:
|
||||
if self.model.config.mm_use_im_start_end:
|
||||
prompt = image_token_se + "\n" + prompt
|
||||
else:
|
||||
prompt = llava_constants.DEFAULT_IMAGE_TOKEN + "\n" + prompt
|
||||
|
||||
# Formats the prompt as a conversation to be fed to the model.
|
||||
conv = conversation.conv_llava_v1.copy()
|
||||
conv.append_message(role=conv.roles[0], message=prompt)
|
||||
conv.append_message(role=conv.roles[1], message=None)
|
||||
prompt = conv.get_prompt()
|
||||
|
||||
# Tokenizes the prompt that includes special image token as well.
|
||||
input_ids = (
|
||||
mm_utils.tokenizer_image_token(
|
||||
prompt=prompt,
|
||||
tokenizer=self.tokenizer,
|
||||
image_token_index=llava_constants.IMAGE_TOKEN_INDEX,
|
||||
return_tensors="pt",
|
||||
)
|
||||
.unsqueeze(0)
|
||||
.to(self.device)
|
||||
)
|
||||
|
||||
images = [
|
||||
image_format_converter.base64_to_image(image_str=base64_image).convert(
|
||||
"RGB"
|
||||
)
|
||||
]
|
||||
# Gets the image embedding.
|
||||
images_tensor = mm_utils.process_images(
|
||||
images=images,
|
||||
image_processor=self.image_processor,
|
||||
model_cfg=self.model.config,
|
||||
).to(self.device, dtype=torch.float16)
|
||||
|
||||
self.stop_str = conversation.conv_llava_v1.sep2
|
||||
self.keywords = [self.stop_str]
|
||||
|
||||
return input_ids, images_tensor
|
||||
|
||||
def inference(
|
||||
self, input_ids: List[torch.Tensor], images_tensor: torch.Tensor
|
||||
) -> List[torch.Tensor]:
|
||||
"""Runs the inference."""
|
||||
stopping_criteria = mm_utils.KeywordsStoppingCriteria(
|
||||
keywords=self.keywords, tokenizer=self.tokenizer, input_ids=input_ids
|
||||
)
|
||||
|
||||
with torch.inference_mode():
|
||||
output_ids = self.model.generate(
|
||||
input_ids=input_ids,
|
||||
images=images_tensor,
|
||||
do_sample=False,
|
||||
temperature=0,
|
||||
top_p=None,
|
||||
num_beams=1,
|
||||
max_new_tokens=512,
|
||||
use_cache=True,
|
||||
stopping_criteria=[stopping_criteria],
|
||||
)
|
||||
|
||||
return output_ids
|
||||
|
||||
def postprocess(
|
||||
self, output_ids: List[torch.Tensor], input_token_len: int
|
||||
) -> List[str]:
|
||||
"""Runs the postprocessing to convert token ids to string."""
|
||||
outputs = self.tokenizer.batch_decode(
|
||||
output_ids[:, input_token_len:], skip_special_tokens=True
|
||||
)[0]
|
||||
outputs = outputs.strip()
|
||||
if outputs.endswith(self.stop_str):
|
||||
outputs = outputs[: -len(self.stop_str)]
|
||||
outputs = outputs.strip()
|
||||
|
||||
return [outputs]
|
||||
|
||||
def handle(self, data: List[Dict[str, Any]], context: Any) -> List[str]:
|
||||
"""Handles an incoming request by passing it through `preprocess`, `inference`, and `postprocess`."""
|
||||
input_ids, images_tensor = self.preprocess(data=data)
|
||||
model_output = self.inference(
|
||||
input_ids=input_ids, images_tensor=images_tensor
|
||||
)
|
||||
|
||||
input_token_len = input_ids.shape[1]
|
||||
return self.postprocess(
|
||||
output_ids=model_output, input_token_len=input_token_len
|
||||
)
|
||||
@@ -1,77 +0,0 @@
|
||||
"""Common utility functions for setting up and initializing the model and the handler."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from util import constants
|
||||
from util import fileutils
|
||||
|
||||
|
||||
def get_model_id(default_model_id: str) -> str:
|
||||
"""Gets a model id or a local model path.
|
||||
|
||||
Args:
|
||||
default_model_id: Default model id for the corresponding model set in the
|
||||
handler.
|
||||
|
||||
Returns:
|
||||
str: model id or a local model path.
|
||||
"""
|
||||
# The model id can be either:
|
||||
# 1) a huggingface model card id, like "Salesforce/blip", or
|
||||
# 2) a GCS path to the model files, like "gs://foo/bar".
|
||||
# If it's a model card id, the model will be loaded from huggingface.
|
||||
model_id = (
|
||||
default_model_id
|
||||
if os.environ.get("MODEL_ID") is None
|
||||
else os.environ["MODEL_ID"]
|
||||
)
|
||||
|
||||
# Else it will be downloaded from GCS to local first.
|
||||
# Since the transformers from_pretrained API can't read from GCS.
|
||||
if model_id.startswith(constants.GCS_URI_PREFIX):
|
||||
gcs_path = model_id[len(constants.GCS_URI_PREFIX) :]
|
||||
local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
|
||||
logging.info("Download %s to %s", model_id, local_model_dir)
|
||||
fileutils.download_gcs_dir_to_local(model_id, local_model_dir)
|
||||
model_id = local_model_dir
|
||||
|
||||
return model_id
|
||||
|
||||
|
||||
def get_map_location(context: Any) -> str:
|
||||
"""Gets model map location.
|
||||
|
||||
Args:
|
||||
context: Torchserve worker context.
|
||||
|
||||
Returns:
|
||||
str: Mapping location.
|
||||
"""
|
||||
properties = context.system_properties
|
||||
return (
|
||||
"cuda"
|
||||
if torch.cuda.is_available() and properties.get("gpu_id") is not None
|
||||
else "cpu"
|
||||
)
|
||||
|
||||
|
||||
def get_model_device(map_location: str, context: Any) -> torch.device:
|
||||
"""Gets model accelerator device.
|
||||
|
||||
Args:
|
||||
map_location: Model map location.
|
||||
context: TorchServe worker context.
|
||||
|
||||
Returns:
|
||||
torch.Device: Device to load the model into.
|
||||
"""
|
||||
properties = context.system_properties
|
||||
return torch.device(
|
||||
map_location + ":" + str(properties.get("gpu_id"))
|
||||
if torch.cuda.is_available() and properties.get("gpu_id") is not None
|
||||
else map_location
|
||||
)
|
||||
@@ -1,747 +0,0 @@
|
||||
"""Common util functions for notebook."""
|
||||
|
||||
import base64
|
||||
from collections.abc import Sequence
|
||||
import datetime
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from google import auth
|
||||
from google.cloud import storage
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import requests
|
||||
import tensorflow as tf
|
||||
import yaml
|
||||
|
||||
|
||||
GCS_URI_PREFIX = "gs://"
|
||||
CHECKPOINT_BUCKET = "gs://model_garden_checkpoints"
|
||||
|
||||
|
||||
def convert_numpy_array_to_byte_string_via_tf_tensor(
|
||||
np_array: np.ndarray,
|
||||
) -> str:
|
||||
"""Serializes a numpy array to tensor bytes.
|
||||
|
||||
Args:
|
||||
np_array: A numpy array.
|
||||
|
||||
Returns:
|
||||
A tensor bytes.
|
||||
"""
|
||||
tensor_array = tf.convert_to_tensor(np_array)
|
||||
tensor_byte_string = tf.io.serialize_tensor(tensor_array)
|
||||
return tensor_byte_string.numpy()
|
||||
|
||||
|
||||
def get_jpeg_bytes(local_image_path: str, new_width: int = -1) -> bytes:
|
||||
"""Returns jpeg bytes given an image path and resizes if required.
|
||||
|
||||
Args:
|
||||
local_image_path: A string of local image path.
|
||||
new_width: An integer of new image width.
|
||||
|
||||
Returns:
|
||||
A jpeg bytes.
|
||||
"""
|
||||
image = Image.open(local_image_path)
|
||||
if new_width <= 0:
|
||||
new_image = image
|
||||
else:
|
||||
width, height = image.size
|
||||
print("original input image size: ", width, " , ", height)
|
||||
new_height = int(height * new_width / width)
|
||||
print("new input image size: ", new_width, " , ", new_height)
|
||||
new_image = image.resize((new_width, new_height))
|
||||
buffered = io.BytesIO()
|
||||
new_image.save(buffered, format="JPEG")
|
||||
return buffered.getvalue()
|
||||
|
||||
|
||||
def gcs_fuse_path(path: str) -> str:
|
||||
"""Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.
|
||||
|
||||
Args:
|
||||
path: A string of path.
|
||||
|
||||
Returns:
|
||||
A gcsfuse path.
|
||||
"""
|
||||
path = path.strip()
|
||||
if path.startswith("gs://"):
|
||||
return "/gcs/" + path[5:]
|
||||
return path
|
||||
|
||||
|
||||
def get_job_name_with_datetime(prefix: str) -> str:
|
||||
"""Gets a job name by adding current time to prefix.
|
||||
|
||||
Args:
|
||||
prefix: A string of job name prefix.
|
||||
|
||||
Returns:
|
||||
A job name.
|
||||
"""
|
||||
now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
job_name = f"{prefix}-{now}".replace("_", "-")
|
||||
return job_name
|
||||
|
||||
|
||||
def create_job_name(prefix: str) -> str:
|
||||
"""Creates a job name.
|
||||
|
||||
Args:
|
||||
prefix: A string of job name prefix.
|
||||
|
||||
Returns:
|
||||
A job name.
|
||||
"""
|
||||
user = os.environ.get("USER")
|
||||
now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
job_name = f"{prefix}-{user}-{now}".replace("_", "-")
|
||||
return job_name
|
||||
|
||||
|
||||
def save_subset_annotation(
|
||||
input_annotation_path: str, output_annotation_path: str
|
||||
):
|
||||
"""Saves a subset of COCO annotation json file with CCA 4.0 license.
|
||||
|
||||
Args:
|
||||
input_annotation_path: A string of input annotation path.
|
||||
output_annotation_path: A string of output annotation path.
|
||||
"""
|
||||
|
||||
with open(input_annotation_path) as f:
|
||||
coco_json = json.load(f)
|
||||
|
||||
img_ids = set()
|
||||
images = []
|
||||
annotations = []
|
||||
|
||||
for img in coco_json["images"]:
|
||||
if img["license"] in [4, 5]: # CCA 4.0 license.
|
||||
img_ids.add(img["id"])
|
||||
images.append(img)
|
||||
|
||||
for ann in coco_json["annotations"]:
|
||||
if ann["image_id"] in img_ids:
|
||||
annotations.append(ann)
|
||||
|
||||
new_json = {
|
||||
"info": coco_json["info"],
|
||||
"licenses": coco_json["licenses"],
|
||||
"images": images,
|
||||
"annotations": annotations,
|
||||
"categories": coco_json["categories"],
|
||||
}
|
||||
|
||||
with open(output_annotation_path, "w") as f:
|
||||
json.dump(new_json, f)
|
||||
|
||||
|
||||
def image_to_base64(image: Any, image_format: str = "JPEG") -> str:
|
||||
"""Converts an image to base64.
|
||||
|
||||
Args:
|
||||
image: A PIL.Image instance.
|
||||
image_format: A string of image format.
|
||||
|
||||
Returns:
|
||||
A base64 string.
|
||||
"""
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=image_format)
|
||||
image_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
return image_str
|
||||
|
||||
|
||||
def base64_to_image(image_str: str) -> Any:
|
||||
"""Convert base64 encoded string to an image.
|
||||
|
||||
Args:
|
||||
image_str: A string of base64 encoded image.
|
||||
|
||||
Returns:
|
||||
A PIL.Image instance.
|
||||
"""
|
||||
image = Image.open(io.BytesIO(base64.b64decode(image_str)))
|
||||
return image
|
||||
|
||||
|
||||
def image_grid(imgs: Sequence[Any], rows: int = 2, cols: int = 2) -> Any:
|
||||
"""Creates an image grid.
|
||||
|
||||
Args:
|
||||
imgs: A list of PIL.Image instances.
|
||||
rows: An integer of number of rows.
|
||||
cols: An integer of number of columns.
|
||||
|
||||
Returns:
|
||||
A PIL.Image instance.
|
||||
"""
|
||||
w, h = imgs[0].size
|
||||
grid = Image.new(
|
||||
mode="RGB", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)
|
||||
)
|
||||
for i, img in enumerate(imgs):
|
||||
grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))
|
||||
return grid
|
||||
|
||||
|
||||
def display_image(image: Any):
|
||||
"""Displays an image.
|
||||
|
||||
Args:
|
||||
image: A PIL.Image instance.
|
||||
"""
|
||||
_ = plt.figure(figsize=(20, 15))
|
||||
plt.grid(False)
|
||||
plt.imshow(image)
|
||||
|
||||
|
||||
def download_gcs_file_to_local(gcs_uri: str, local_path: str):
|
||||
"""Download a gcs file to a local path.
|
||||
|
||||
Args:
|
||||
gcs_uri: A string of file path on GCS.
|
||||
local_path: A string of local file path.
|
||||
"""
|
||||
if not gcs_uri.startswith(GCS_URI_PREFIX):
|
||||
raise ValueError(
|
||||
f"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}."
|
||||
)
|
||||
client = storage.Client()
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
with open(local_path, "wb") as f:
|
||||
client.download_blob_to_file(gcs_uri, f)
|
||||
|
||||
|
||||
def download_image(url: str) -> str:
|
||||
"""Downloads an image from the given URL.
|
||||
|
||||
Args:
|
||||
url: The URL of the image to download.
|
||||
|
||||
Returns:
|
||||
base64 encoded image.
|
||||
"""
|
||||
response = requests.get(url)
|
||||
return Image.open(io.BytesIO(response.content)) # pytype: disable=bad-return-type # pillow-102-upgrade
|
||||
|
||||
|
||||
def resize_image(image: Any, new_width: int = 1000) -> Any:
|
||||
"""Resizes an image to a certain width.
|
||||
|
||||
Args:
|
||||
image: The image which has to be resized.
|
||||
new_width: New width of the image.
|
||||
|
||||
Returns:
|
||||
New resized image.
|
||||
"""
|
||||
width, height = image.size
|
||||
new_height = int(height * new_width / width)
|
||||
new_img = image.resize((new_width, new_height))
|
||||
return new_img
|
||||
|
||||
|
||||
def load_img(path: str) -> Any:
|
||||
"""Reads image from path and return PIL.Image instance.
|
||||
|
||||
Args:
|
||||
path: A string of image path.
|
||||
|
||||
Returns:
|
||||
A PIL.Image instance.
|
||||
"""
|
||||
img = tf.io.read_file(path)
|
||||
img = tf.image.decode_jpeg(img, channels=3)
|
||||
return Image.fromarray(np.uint8(img)).convert("RGB")
|
||||
|
||||
|
||||
def decode_image(
|
||||
image_str_tensor: tf.string, new_height: int, new_width: int
|
||||
) -> tf.float32:
|
||||
"""Converts and resizes image bytes to image tensor.
|
||||
|
||||
Args:
|
||||
image_str_tensor: A string of image bytes.
|
||||
new_height: An integer of new image height.
|
||||
new_width: An integer of new image width.
|
||||
|
||||
Returns:
|
||||
An image tensor.
|
||||
"""
|
||||
image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)
|
||||
image = tf.image.resize(image, (new_height, new_width))
|
||||
return image
|
||||
|
||||
|
||||
def get_label_map(label_map_yaml_filepath: str) -> dict[int, str]:
|
||||
"""Returns class id to label mapping given a filepath to the label map.
|
||||
|
||||
Args:
|
||||
label_map_yaml_filepath: A string of label map yaml file path.
|
||||
|
||||
Returns:
|
||||
A dictionary of class id to label mapping.
|
||||
"""
|
||||
with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file:
|
||||
label_map = yaml.safe_load(input_file.read())["label_map"]
|
||||
return label_map
|
||||
|
||||
|
||||
def get_prediction_instances(test_filepath: str, new_width: int = -1) -> Any:
|
||||
"""Generate instance from image path to pass to Vertex AI Endpoint for prediction.
|
||||
|
||||
Args:
|
||||
test_filepath: A string of test image path.
|
||||
new_width: An integer of new image width.
|
||||
|
||||
Returns:
|
||||
A list of instances.
|
||||
"""
|
||||
if new_width <= 0:
|
||||
with tf.io.gfile.GFile(test_filepath, "rb") as input_file:
|
||||
encoded_string = base64.b64encode(input_file.read()).decode("utf-8")
|
||||
else:
|
||||
img = load_img(test_filepath)
|
||||
width, height = img.size
|
||||
print("original input image size: ", width, " , ", height)
|
||||
new_height = int(height * new_width / width)
|
||||
new_img = img.resize((new_width, new_height))
|
||||
print("resized input image size: ", new_width, " , ", new_height)
|
||||
buffered = io.BytesIO()
|
||||
new_img.save(buffered, format="JPEG")
|
||||
encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
|
||||
instances = [{
|
||||
"encoded_image": {"b64": encoded_string},
|
||||
}]
|
||||
return instances
|
||||
|
||||
|
||||
def vqa_predict(
|
||||
endpoint: Any,
|
||||
question_prompts: Sequence[str],
|
||||
image: Any,
|
||||
language_code: str = "en",
|
||||
new_width: int = 1000,
|
||||
use_dedicated_endpoint: bool = False,
|
||||
) -> Sequence[str]:
|
||||
"""Predicts the answer to a question about an image using an Endpoint."""
|
||||
# Resize and convert image to base64 string.
|
||||
resized_image = resize_image(image, new_width)
|
||||
resized_image_base64 = image_to_base64(resized_image)
|
||||
|
||||
instances = []
|
||||
if question_prompts:
|
||||
# Format question prompt
|
||||
question_prompt_format = "answer {} {}\n"
|
||||
for question_prompt in question_prompts:
|
||||
if question_prompt:
|
||||
instances.append({
|
||||
"prompt": question_prompt_format.format(
|
||||
language_code, question_prompt
|
||||
),
|
||||
"image": resized_image_base64,
|
||||
})
|
||||
else:
|
||||
instances.append({
|
||||
"image": resized_image_base64,
|
||||
})
|
||||
|
||||
response = endpoint.predict(
|
||||
instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
|
||||
)
|
||||
return [pred.get("response") for pred in response.predictions]
|
||||
|
||||
|
||||
def caption_predict(
|
||||
endpoint: Any,
|
||||
language_code: str,
|
||||
image: Any,
|
||||
caption_prompt: bool = False,
|
||||
new_width: int = 1000,
|
||||
use_dedicated_endpoint: bool = False,
|
||||
) -> str:
|
||||
"""Predicts a caption for a given image using an Endpoint."""
|
||||
# Resize and convert image to base64 string.
|
||||
resized_image = resize_image(image, new_width)
|
||||
resized_image_base64 = image_to_base64(resized_image)
|
||||
|
||||
instance = {"image": resized_image_base64}
|
||||
|
||||
if caption_prompt:
|
||||
# Format caption prompt
|
||||
caption_prompt_format = "caption {}\n"
|
||||
instance["prompt"] = caption_prompt_format.format(language_code)
|
||||
|
||||
instances = [instance]
|
||||
response = endpoint.predict(
|
||||
instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
|
||||
)
|
||||
return response.predictions[0].get("response")
|
||||
|
||||
|
||||
def ocr_predict(
|
||||
endpoint: Any,
|
||||
ocr_prompt: str,
|
||||
image: Any,
|
||||
new_width: int = 1000,
|
||||
use_dedicated_endpoint: bool = False,
|
||||
) -> str:
|
||||
"""Extracts text from a given image using an Endpoint."""
|
||||
# Resize and convert image to base64 string.
|
||||
resized_image = resize_image(image, new_width)
|
||||
resized_image_base64 = image_to_base64(resized_image)
|
||||
|
||||
instance = {"image": resized_image_base64}
|
||||
if ocr_prompt:
|
||||
instance["prompt"] = ocr_prompt
|
||||
instances = [instance]
|
||||
|
||||
response = endpoint.predict(
|
||||
instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
|
||||
)
|
||||
return response.predictions[0].get("response")
|
||||
|
||||
|
||||
def detect_predict(
|
||||
endpoint: Any,
|
||||
detect_prompt: str,
|
||||
image: Any,
|
||||
new_width: int = 1000,
|
||||
use_dedicated_endpoint: bool = False,
|
||||
) -> str:
|
||||
"""Predicts the answer to a question about an image using an Endpoint."""
|
||||
# Resize and convert image to base64 string.
|
||||
resized_image = resize_image(image, new_width)
|
||||
resized_image_base64 = image_to_base64(resized_image)
|
||||
|
||||
instance = {"image": resized_image_base64}
|
||||
if detect_prompt:
|
||||
instance["prompt"] = detect_prompt
|
||||
instances = [instance]
|
||||
|
||||
response = endpoint.predict(
|
||||
instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
|
||||
)
|
||||
return response.predictions[0].get("response")
|
||||
|
||||
|
||||
def copy_model_artifacts(
|
||||
model_id: str,
|
||||
model_source: str,
|
||||
model_destination: str,
|
||||
) -> None:
|
||||
"""Copies model artifacts from model_source to model_destination.
|
||||
|
||||
model_source and model_destination should be GCS path.
|
||||
|
||||
Args:
|
||||
model_id: The model id.
|
||||
model_source: The source of the model artifact.
|
||||
model_destination: The destination of the model artifact.
|
||||
"""
|
||||
if not model_source.startswith(GCS_URI_PREFIX):
|
||||
raise ValueError(
|
||||
f"{model_source} is not a GCS path starting with {GCS_URI_PREFIX}."
|
||||
)
|
||||
if not model_destination.startswith(GCS_URI_PREFIX):
|
||||
raise ValueError(
|
||||
f"{model_destination} is not a GCS path starting with {GCS_URI_PREFIX}."
|
||||
)
|
||||
model_source = f"{model_source}/{model_id}"
|
||||
model_destination = f"{model_destination}/{model_id}"
|
||||
print("Copying model artifact from ", model_source, " to ", model_destination)
|
||||
subprocess.check_output([
|
||||
"gcloud",
|
||||
"storage",
|
||||
"cp",
|
||||
"-r",
|
||||
model_source,
|
||||
model_destination,
|
||||
])
|
||||
|
||||
|
||||
def get_quota(project_id: str, region: str, resource_id: str) -> int:
|
||||
"""Returns the quota for a resource in a region.
|
||||
|
||||
Args:
|
||||
project_id: The project id.
|
||||
region: The region.
|
||||
resource_id: The resource id.
|
||||
|
||||
Returns:
|
||||
The quota for the resource in the region. Returns -1 if can not figure out
|
||||
the quota.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the command to get quota fails.
|
||||
"""
|
||||
service_endpoint = "aiplatform.googleapis.com"
|
||||
|
||||
command = (
|
||||
"gcloud alpha services quota list"
|
||||
f" --service={service_endpoint} --consumer=projects/{project_id}"
|
||||
f" --filter='{service_endpoint}/{resource_id}' --format=json"
|
||||
)
|
||||
process = subprocess.run(
|
||||
command, shell=True, capture_output=True, text=True, check=True
|
||||
)
|
||||
if process.returncode == 0:
|
||||
quota_data = json.loads(process.stdout)
|
||||
else:
|
||||
raise RuntimeError(f"Error fetching quota data: {process.stderr}")
|
||||
|
||||
if not quota_data or "consumerQuotaLimits" not in quota_data[0]:
|
||||
return -1
|
||||
if (
|
||||
not quota_data[0]["consumerQuotaLimits"]
|
||||
or "quotaBuckets" not in quota_data[0]["consumerQuotaLimits"][0]
|
||||
):
|
||||
return -1
|
||||
all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
|
||||
|
||||
# If the quota data does not have dimensions, it is global quota. However,
|
||||
# global quota may be overridden by regional quota. So we need to check the
|
||||
# global quota first.
|
||||
global_quota = -1
|
||||
if (
|
||||
all_regions_data
|
||||
and "dimensions" not in all_regions_data[0]
|
||||
and "effectiveLimit" in all_regions_data[0]
|
||||
):
|
||||
global_quota = int(all_regions_data[0]["effectiveLimit"])
|
||||
for region_data in all_regions_data:
|
||||
if (
|
||||
region_data.get("dimensions")
|
||||
and region_data["dimensions"]["region"] == region
|
||||
):
|
||||
if "effectiveLimit" in region_data:
|
||||
return int(region_data["effectiveLimit"])
|
||||
else:
|
||||
return 0
|
||||
return global_quota
|
||||
|
||||
|
||||
def get_resource_id(
|
||||
accelerator_type: str,
|
||||
is_for_training: bool,
|
||||
is_spot: bool = False,
|
||||
is_restricted_image: bool = False,
|
||||
is_dynamic_workload_scheduler: bool = False,
|
||||
) -> str:
|
||||
"""Returns the resource id for a given accelerator type and the use case.
|
||||
|
||||
Args:
|
||||
accelerator_type: The accelerator type.
|
||||
is_for_training: Whether the resource is used for training. Set false for
|
||||
serving use case.
|
||||
is_spot: Whether the resource is used with Spot.
|
||||
is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
|
||||
is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
|
||||
Workload Scheduler.
|
||||
|
||||
Returns:
|
||||
The resource id.
|
||||
"""
|
||||
accelerator_suffix_map = {
|
||||
"NVIDIA_TESLA_V100": "nvidia_v100_gpus",
|
||||
"NVIDIA_TESLA_P100": "nvidia_p100_gpus",
|
||||
"NVIDIA_L4": "nvidia_l4_gpus",
|
||||
"NVIDIA_TESLA_A100": "nvidia_a100_gpus",
|
||||
"NVIDIA_A100_80GB": "nvidia_a100_80gb_gpus",
|
||||
"NVIDIA_H100_80GB": "nvidia_h100_gpus",
|
||||
"NVIDIA_H100_MEGA_80GB": "nvidia_h100_mega_gpus",
|
||||
"NVIDIA_H200_141GB": "nvidia_h200_gpus",
|
||||
"NVIDIA_TESLA_T4": "nvidia_t4_gpus",
|
||||
"TPU_V6e": "tpu_v6e",
|
||||
"TPU_V5e": "tpu_v5e",
|
||||
"TPU_V3": "tpu_v3",
|
||||
}
|
||||
default_training_accelerator_map = {
|
||||
key: f"custom_model_training_{accelerator_suffix_map[key]}"
|
||||
for key in accelerator_suffix_map
|
||||
}
|
||||
dws_training_accelerator_map = {
|
||||
key: f"custom_model_training_preemptible_{accelerator_suffix_map[key]}"
|
||||
for key in accelerator_suffix_map
|
||||
}
|
||||
restricted_image_training_accelerator_map = {
|
||||
"NVIDIA_A100_80GB": "restricted_image_training_nvidia_a100_80gb_gpus",
|
||||
}
|
||||
spot_serving_accelerator_map = {
|
||||
key: f"custom_model_serving_preemptible_{accelerator_suffix_map[key]}"
|
||||
for key in accelerator_suffix_map
|
||||
}
|
||||
serving_accelerator_map = {
|
||||
key: f"custom_model_serving_{accelerator_suffix_map[key]}"
|
||||
for key in accelerator_suffix_map
|
||||
}
|
||||
|
||||
if is_for_training:
|
||||
if is_restricted_image and is_dynamic_workload_scheduler:
|
||||
raise ValueError(
|
||||
"Dynamic Workload Scheduler does not work for restricted image"
|
||||
" training."
|
||||
)
|
||||
training_accelerator_map = (
|
||||
restricted_image_training_accelerator_map
|
||||
if is_restricted_image
|
||||
else default_training_accelerator_map
|
||||
)
|
||||
if accelerator_type in training_accelerator_map:
|
||||
if is_dynamic_workload_scheduler:
|
||||
return dws_training_accelerator_map[accelerator_type]
|
||||
else:
|
||||
return training_accelerator_map[accelerator_type]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Could not find accelerator type: {accelerator_type} for training."
|
||||
)
|
||||
else:
|
||||
if is_dynamic_workload_scheduler:
|
||||
raise ValueError("Dynamic Workload Scheduler does not work for serving.")
|
||||
accelerator_map = (
|
||||
spot_serving_accelerator_map if is_spot else serving_accelerator_map
|
||||
)
|
||||
if accelerator_type in accelerator_map:
|
||||
return accelerator_map[accelerator_type]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Could not find accelerator type: {accelerator_type} for serving."
|
||||
)
|
||||
|
||||
|
||||
def check_quota(
|
||||
project_id: str,
|
||||
region: str,
|
||||
accelerator_type: str,
|
||||
accelerator_count: int,
|
||||
is_for_training: bool,
|
||||
is_spot: bool = False,
|
||||
is_restricted_image: bool = False,
|
||||
is_dynamic_workload_scheduler: bool = False,
|
||||
) -> None:
|
||||
"""Checks if the project and the region has the required quota.
|
||||
|
||||
Args:
|
||||
project_id: The project id.
|
||||
region: The region.
|
||||
accelerator_type: The accelerator type.
|
||||
accelerator_count: The number of accelerators to check quota for.
|
||||
is_for_training: Whether the resource is used for training. Set false for
|
||||
serving use case.
|
||||
is_spot: Whether the resource is used with Spot.
|
||||
is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
|
||||
is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
|
||||
Workload Scheduler.
|
||||
"""
|
||||
resource_id = get_resource_id(
|
||||
accelerator_type,
|
||||
is_for_training=is_for_training,
|
||||
is_spot=is_spot,
|
||||
is_restricted_image=is_restricted_image,
|
||||
is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,
|
||||
)
|
||||
quota = get_quota(project_id, region, resource_id)
|
||||
quota_request_instruction = (
|
||||
"Either use "
|
||||
"a different region or request additional quota. Follow "
|
||||
"instructions here "
|
||||
"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota"
|
||||
" to check quota in a region or request additional quota for "
|
||||
"your project."
|
||||
)
|
||||
if quota == -1:
|
||||
raise ValueError(
|
||||
f"Quota not found for: {resource_id} in {region}."
|
||||
f" {quota_request_instruction}"
|
||||
)
|
||||
if quota < accelerator_count:
|
||||
raise ValueError(
|
||||
f"Quota not enough for {resource_id} in {region}: {quota} <"
|
||||
f" {accelerator_count}. {quota_request_instruction}"
|
||||
)
|
||||
|
||||
|
||||
def get_deploy_source() -> str:
|
||||
"""Gets deploy_source string based on running environment."""
|
||||
vertex_product = os.environ.get("VERTEX_PRODUCT", "")
|
||||
match vertex_product:
|
||||
case "COLAB_ENTERPRISE":
|
||||
return "notebook_colab_enterprise"
|
||||
case "WORKBENCH_INSTANCE":
|
||||
return "notebook_workbench"
|
||||
case _:
|
||||
# Legacy workbench, legacy colab, or other custom environments.
|
||||
return "notebook_environment_unspecified"
|
||||
|
||||
|
||||
def _is_operation_done(op_name: str, region: str) -> bool:
|
||||
"""Checks if the operation is done.
|
||||
|
||||
Args:
|
||||
op_name: The name of the operation to poll.
|
||||
region: The region of the operation.
|
||||
|
||||
Returns:
|
||||
True if the operation is done, False otherwise.
|
||||
|
||||
Raises:
|
||||
ValueError: If the operation failed.
|
||||
"""
|
||||
creds, _ = auth.default()
|
||||
auth_req = auth.transport.requests.Request()
|
||||
creds.refresh(auth_req)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {creds.token}",
|
||||
}
|
||||
url = f"https://{region}-aiplatform.googleapis.com/ui/{op_name}"
|
||||
response = requests.get(url, headers=headers)
|
||||
operation_data = response.json()
|
||||
if "error" in operation_data:
|
||||
raise ValueError(f"Operation failed: {operation_data['error']}")
|
||||
return operation_data.get("done", False)
|
||||
|
||||
|
||||
def poll_and_wait(
|
||||
op_name: str, region: str, total_wait: int, interval: int = 60
|
||||
) -> None:
|
||||
"""Polls the operation and waits for it to complete.
|
||||
|
||||
Args:
|
||||
op_name: The name of the operation to poll.
|
||||
region: The region of the operation.
|
||||
total_wait: The total wait time in seconds.
|
||||
interval: The interval between each poll in seconds.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the operation times out.
|
||||
"""
|
||||
start_time = time.time()
|
||||
while True:
|
||||
if _is_operation_done(op_name, region):
|
||||
break
|
||||
time_elapsed = time.time() - start_time
|
||||
if time_elapsed > total_wait:
|
||||
raise TimeoutError(
|
||||
f"Operation timed out after {int(time_elapsed)} seconds."
|
||||
)
|
||||
print(
|
||||
"\rStill waiting for operation... Elapsed time in seconds:"
|
||||
f" {int(time_elapsed):<6}",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
time.sleep(interval)
|
||||
|
||||
-570
@@ -1,570 +0,0 @@
|
||||
"""Functions for dataset validation.
|
||||
|
||||
This tool is used to validate the dataset against the given template.
|
||||
"""
|
||||
|
||||
import json
|
||||
import multiprocessing
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Any, Callable, Dict, Tuple, Union
|
||||
from absl import logging
|
||||
import accelerate
|
||||
import datasets
|
||||
import transformers
|
||||
|
||||
GCS_URI_PREFIX = "gs://"
|
||||
GCSFUSE_URI_PREFIX = "/gcs/"
|
||||
LOCAL_BASE_MODEL_DIR = "/tmp/base_model_dir"
|
||||
LOCAL_TEMPLATE_DIR = "/tmp/template_dir"
|
||||
_TEMPLATE_DIRNAME = "templates"
|
||||
_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME = "vertex-ai-samples"
|
||||
_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR = (
|
||||
"community-content/vertex_model_garden/model_oss/peft/train/vmg/templates"
|
||||
)
|
||||
_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
|
||||
_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
|
||||
_DESCRIPTION_KEY = "description"
|
||||
_SOURCE_KEY = "source"
|
||||
_PROMPT_INPUT_KEY = "prompt_input"
|
||||
_PROMPT_NO_INPUT_KEY = "prompt_no_input"
|
||||
_RESPONSE_SEPARATOR = "response_separator"
|
||||
_INSTRUCTION_SEPARATOR = "instruction_separator"
|
||||
_CHAT_TEMPLATE_KEY = "chat_template"
|
||||
_KNOWN_KEYS = (
|
||||
_DESCRIPTION_KEY,
|
||||
_SOURCE_KEY,
|
||||
_PROMPT_INPUT_KEY,
|
||||
_PROMPT_NO_INPUT_KEY,
|
||||
_RESPONSE_SEPARATOR,
|
||||
_INSTRUCTION_SEPARATOR,
|
||||
_CHAT_TEMPLATE_KEY,
|
||||
)
|
||||
|
||||
|
||||
def is_gcs_path(input_path: str) -> bool:
|
||||
"""Checks if the input path is a Google Cloud Storage (GCS) path.
|
||||
|
||||
Args:
|
||||
input_path: The input path to be checked.
|
||||
|
||||
Returns:
|
||||
True if the input path is a GCS path, False otherwise.
|
||||
"""
|
||||
return input_path is not None and input_path.startswith(GCS_URI_PREFIX)
|
||||
|
||||
|
||||
def force_gcs_fuse_path(gcs_uri: str) -> str:
|
||||
"""Converts gs:// uris to their /gcs/ equivalents. No-op for other uris.
|
||||
|
||||
Args:
|
||||
gcs_uri: The GCS URI to convert.
|
||||
|
||||
Returns:
|
||||
The converted GCS URI.
|
||||
"""
|
||||
if is_gcs_path(gcs_uri):
|
||||
return GCSFUSE_URI_PREFIX + gcs_uri[len(GCS_URI_PREFIX) :]
|
||||
else:
|
||||
return gcs_uri
|
||||
|
||||
|
||||
def download_gcs_uri_to_local(
|
||||
gcs_uri: str,
|
||||
destination_dir: str = LOCAL_BASE_MODEL_DIR,
|
||||
check_path_exists: bool = True,
|
||||
) -> str:
|
||||
"""Downloads GCS URI to local.
|
||||
|
||||
If GCS URI is a directory, gs://some/folder is downloaded to
|
||||
/destination_dir/folder. If GCS URI is a file, gs://some/file is downloaded to
|
||||
/destination_dir/file.
|
||||
|
||||
Args:
|
||||
gcs_uri: GCS URI to download.
|
||||
destination_dir: Local directory directory.
|
||||
check_path_exists: Whether to check if the path exists.
|
||||
|
||||
Returns:
|
||||
Local path to target folder/file.
|
||||
"""
|
||||
target = os.path.join(
|
||||
destination_dir,
|
||||
os.path.basename(os.path.normpath(gcs_uri)),
|
||||
)
|
||||
if check_path_exists and os.path.exists(target):
|
||||
logging.info("File %s already exists.", target)
|
||||
return target
|
||||
if accelerate.PartialState().is_local_main_process:
|
||||
logging.info(
|
||||
"Downloading file(s) from %s to %s...", gcs_uri, destination_dir
|
||||
)
|
||||
if not os.path.exists(destination_dir):
|
||||
os.mkdir(destination_dir)
|
||||
subprocess.check_output([
|
||||
"gsutil",
|
||||
"-m",
|
||||
"cp",
|
||||
"-r",
|
||||
gcs_uri,
|
||||
destination_dir,
|
||||
])
|
||||
logging.info("Downloaded file(s) from %s to %s.", gcs_uri, destination_dir)
|
||||
# Make sure ALL processes process to next step after data downloading is done.
|
||||
# It matters for the main process to wait for other processes as well.
|
||||
accelerate.PartialState().wait_for_everyone()
|
||||
return target
|
||||
|
||||
|
||||
def get_template(template_path: str) -> Dict[str, str]:
|
||||
"""Gets the template dictionary given the file path.
|
||||
|
||||
Args:
|
||||
template_path: Path to the template file.
|
||||
|
||||
Returns:
|
||||
A dictionary of the template.
|
||||
|
||||
Raises:
|
||||
ValueError: If the template file does not exist or contains unknown keys.
|
||||
"""
|
||||
if is_gcs_path(template_path):
|
||||
template_path = force_gcs_fuse_path(template_path)
|
||||
elif not os.path.isfile(template_path):
|
||||
template_path = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
_TEMPLATE_DIRNAME,
|
||||
template_path + ".json",
|
||||
)
|
||||
if not os.path.isfile(template_path):
|
||||
raise ValueError(f"Template file {template_path} does not exist.")
|
||||
with open(template_path, "r") as f:
|
||||
template_json: dict[str, str] = json.load(f)
|
||||
for key in template_json:
|
||||
if key not in _KNOWN_KEYS:
|
||||
raise ValueError(f"Unknown key {key} in template {template_path}.")
|
||||
return template_json
|
||||
|
||||
|
||||
def get_response_separator(template_json: Dict[str, str]) -> Union[str, None]:
|
||||
return template_json.get(_RESPONSE_SEPARATOR, None)
|
||||
|
||||
|
||||
def get_instruction_separator(
|
||||
template_json: Dict[str, str],
|
||||
) -> Union[str, None]:
|
||||
return template_json.get(_INSTRUCTION_SEPARATOR, None)
|
||||
|
||||
|
||||
def _format_template_fn(
|
||||
template: str,
|
||||
input_column: str,
|
||||
tokenizer: transformers.PreTrainedTokenizer | None = None,
|
||||
) -> Callable[[Dict[str, str]], Dict[str, str]]:
|
||||
"""Formats a dataset example according to a template.
|
||||
|
||||
Args:
|
||||
template: Name of the JSON template file under `templates/` or GCS path to
|
||||
the template file.
|
||||
input_column: The input column in the dataset to be used or updated by the
|
||||
template. If it does not exist, the template's `prompt_no_input` will be
|
||||
used, and the input_column will be created.
|
||||
tokenizer: The tokenizer to use for chat_template templates.
|
||||
|
||||
Returns:
|
||||
A function that formats data according to the template.
|
||||
"""
|
||||
template_json = get_template(template)
|
||||
|
||||
if _CHAT_TEMPLATE_KEY not in template_json:
|
||||
|
||||
def format_fn(example: Dict[str, str]) -> Dict[str, str]:
|
||||
format_dict = {key: value for key, value in example.items()}
|
||||
if format_dict.get(input_column):
|
||||
format_str = template_json[_PROMPT_INPUT_KEY]
|
||||
elif _PROMPT_NO_INPUT_KEY in template_json:
|
||||
format_str = template_json[_PROMPT_NO_INPUT_KEY]
|
||||
else:
|
||||
raise KeyError(
|
||||
f"The template {os.path.basename(template)} does not contain"
|
||||
f" {_PROMPT_INPUT_KEY} or {_PROMPT_NO_INPUT_KEY} key."
|
||||
)
|
||||
try:
|
||||
return {input_column: format_str.format(**format_dict)}
|
||||
except KeyError as e:
|
||||
raise KeyError(
|
||||
f"The template {os.path.basename(template)} contains a key {e} in"
|
||||
f" {_PROMPT_INPUT_KEY} or {_PROMPT_NO_INPUT_KEY} that does not"
|
||||
" exist in the dataset example. The dataset example looks like"
|
||||
f" {format_dict}."
|
||||
) from e
|
||||
|
||||
return format_fn
|
||||
elif (
|
||||
_PROMPT_INPUT_KEY in template_json
|
||||
or _PROMPT_NO_INPUT_KEY in template_json
|
||||
):
|
||||
raise ValueError(
|
||||
f"chat_template templates do not support {_PROMPT_INPUT_KEY} or"
|
||||
f" {_PROMPT_NO_INPUT_KEY} templates."
|
||||
)
|
||||
else:
|
||||
if tokenizer is None:
|
||||
raise ValueError("A tokenizer is required for chat_template templates.")
|
||||
# Assign HuggingFace jinja template.
|
||||
tokenizer.chat_template = template_json[_CHAT_TEMPLATE_KEY]
|
||||
|
||||
def format_fn(example: Dict[str, str]) -> Dict[str, str]:
|
||||
try:
|
||||
return {
|
||||
input_column: tokenizer.apply_chat_template(
|
||||
example[input_column],
|
||||
tokenize=False,
|
||||
add_generation_prompt=False,
|
||||
)
|
||||
}
|
||||
except KeyError as e:
|
||||
raise KeyError(
|
||||
f"The template {os.path.basename(template)} contains a key {e} in"
|
||||
f" {_CHAT_TEMPLATE_KEY} that does not exist in the dataset example."
|
||||
) from e
|
||||
|
||||
return format_fn
|
||||
|
||||
|
||||
def _get_split_string(
|
||||
split: str,
|
||||
dataset_percent: int | None = None,
|
||||
dataset_k_rows: int | None = None,
|
||||
) -> str:
|
||||
"""Gets the formatted split string for the dataset.
|
||||
|
||||
This is used to format the split string as per
|
||||
https://huggingface.co/docs/datasets/v2.21.0/loading#slice-splits. Also, this
|
||||
function will only be used to load the partial dataset for validating the
|
||||
dataset against the template.
|
||||
|
||||
Args:
|
||||
split: Split of the dataset.
|
||||
dataset_percent: The percentage of the dataset to load.
|
||||
dataset_k_rows: The top k sequences to load from the dataset.
|
||||
|
||||
Returns:
|
||||
A formatted split string.
|
||||
"""
|
||||
# Validate the dataset_percent and dataset_k_rows values.
|
||||
if dataset_percent and dataset_k_rows:
|
||||
raise ValueError(
|
||||
"You can set either validate_percentage_of_dataset or"
|
||||
" validate_k_rows_of_dataset, but not both."
|
||||
)
|
||||
|
||||
if dataset_percent:
|
||||
logging.info("Loading %d percent of the dataset...", dataset_percent)
|
||||
return f"{split}[:{dataset_percent}%]"
|
||||
|
||||
if dataset_k_rows:
|
||||
logging.info("Loading top %d rows of the dataset...", dataset_k_rows)
|
||||
return f"{split}[:{dataset_k_rows}]"
|
||||
|
||||
return split
|
||||
|
||||
|
||||
def _github_template_path(template: str) -> str:
|
||||
"""Generates the path to the template in the Vertex AI Samples GitHub repo.
|
||||
|
||||
Args:
|
||||
template: Name of the template.
|
||||
|
||||
Returns:
|
||||
The path to the template in the Vertex AI Samples GitHub repo.
|
||||
"""
|
||||
# vertex-ai-samples directory may lie under separate directory depending on
|
||||
# the scratch_dir parameter in the notebook execution environment.
|
||||
vertex_ai_samples_abs_path = os.getcwd().split(
|
||||
_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME
|
||||
)[0]
|
||||
return os.path.join(
|
||||
vertex_ai_samples_abs_path,
|
||||
_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME,
|
||||
_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR,
|
||||
template + ".json",
|
||||
)
|
||||
|
||||
|
||||
def _get_dataset(
|
||||
dataset_name: str,
|
||||
split: str,
|
||||
num_proc: int | None = None,
|
||||
) -> datasets.DatasetDict:
|
||||
"""Gets a dataset.
|
||||
|
||||
Args:
|
||||
dataset_name: Name of the dataset or path to a custom dataset.
|
||||
split: Split of the dataset.
|
||||
num_proc: Number of processors to use.
|
||||
|
||||
Returns:
|
||||
A dataset.
|
||||
"""
|
||||
dataset_name = force_gcs_fuse_path(dataset_name)
|
||||
if os.path.isfile(dataset_name):
|
||||
# Custom dataset.
|
||||
return datasets.load_dataset(
|
||||
"json",
|
||||
data_files=[dataset_name],
|
||||
split=split,
|
||||
num_proc=num_proc,
|
||||
)
|
||||
# HF dataset.
|
||||
return datasets.load_dataset(dataset_name, split=split, num_proc=num_proc)
|
||||
|
||||
|
||||
def should_add_pad_token(model_id: str) -> bool:
|
||||
"""Returns whether the model requires adding a special pad token.
|
||||
|
||||
Args:
|
||||
model_id: The name of the model.
|
||||
|
||||
Returns:
|
||||
True if the model requires adding a special pad token, False otherwise.
|
||||
"""
|
||||
return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
|
||||
|
||||
|
||||
def should_add_eos_token(model_id: str) -> bool:
|
||||
"""Returns whether the model requires adding a special eos token.
|
||||
|
||||
Args:
|
||||
model_id: The name of the model.
|
||||
|
||||
Returns:
|
||||
True if the model requires adding a special eos token, False otherwise.
|
||||
"""
|
||||
return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
|
||||
|
||||
|
||||
def load_tokenizer(
|
||||
pretrained_model_id: str,
|
||||
padding_side: str | None = None,
|
||||
access_token: str | None = None,
|
||||
) -> transformers.AutoTokenizer:
|
||||
"""Loads tokenizer based on `pretrained_model_id`.
|
||||
|
||||
Args:
|
||||
pretrained_model_id: The name of the pretrained model.
|
||||
padding_side: The side to pad the input on.
|
||||
access_token: The access token to use for the tokenizer.
|
||||
|
||||
Returns:
|
||||
The tokenizer.
|
||||
"""
|
||||
tokenizer_kwargs = {}
|
||||
if should_add_eos_token(pretrained_model_id):
|
||||
tokenizer_kwargs["add_eos_token"] = True
|
||||
if padding_side:
|
||||
tokenizer_kwargs["padding_side"] = padding_side
|
||||
|
||||
with accelerate.PartialState().local_main_process_first():
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(
|
||||
pretrained_model_id,
|
||||
trust_remote_code=False,
|
||||
use_fast=True,
|
||||
token=access_token,
|
||||
**tokenizer_kwargs,
|
||||
)
|
||||
|
||||
if should_add_pad_token(pretrained_model_id):
|
||||
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
|
||||
|
||||
return tokenizer
|
||||
|
||||
|
||||
def get_filtered_dataset(
|
||||
dataset: Any,
|
||||
input_column: str,
|
||||
max_seq_length: int,
|
||||
tokenizer: transformers.PreTrainedTokenizer,
|
||||
) -> Any:
|
||||
"""Returns the dataset by removing examples that are longer than max_seq_length.
|
||||
|
||||
Args:
|
||||
dataset: The dataset to filter.
|
||||
input_column: The input column in the dataset to be used.
|
||||
max_seq_length: The maximum sequence length.
|
||||
tokenizer: The tokenizer.
|
||||
"""
|
||||
actual_dataset_length = len(dataset)
|
||||
filtered_dataset = dataset.filter(
|
||||
lambda x: len(tokenizer(x[input_column])["input_ids"]) <= max_seq_length
|
||||
)
|
||||
filtered_dataset_length = len(filtered_dataset)
|
||||
if actual_dataset_length != filtered_dataset_length:
|
||||
examples_removed_percent = (
|
||||
(actual_dataset_length - filtered_dataset_length)
|
||||
* 100
|
||||
/ actual_dataset_length
|
||||
)
|
||||
logging.info(
|
||||
"(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
|
||||
" max-seq-length. Filtering out %d example(s) which are longer than"
|
||||
" max-seq-length.",
|
||||
100 - examples_removed_percent,
|
||||
max_seq_length,
|
||||
examples_removed_percent,
|
||||
actual_dataset_length - filtered_dataset_length,
|
||||
)
|
||||
|
||||
return filtered_dataset
|
||||
|
||||
|
||||
def format_dataset(
|
||||
dataset: datasets.Dataset,
|
||||
input_column: str,
|
||||
template: str = None,
|
||||
tokenizer: transformers.PreTrainedTokenizer | None = None,
|
||||
) -> datasets.Dataset:
|
||||
"""Takes a raw dataset and formats it using a template and tokenizer.
|
||||
|
||||
Args:
|
||||
dataset: The raw (unprocessed) dataset to format.
|
||||
input_column: The input column in the dataset to be used or updaded by the
|
||||
template. If it does not exist, the template's `prompt_no_input` will be
|
||||
used, and the input_column will be created.
|
||||
template: Name of the JSON template file under `templates/` or GCS path to
|
||||
the template file.
|
||||
tokenizer: The tokenizer to use for chat_template templates.
|
||||
|
||||
Returns:
|
||||
A dataset compatible with the template.
|
||||
"""
|
||||
return dataset.map(
|
||||
_format_template_fn(
|
||||
template,
|
||||
input_column=input_column,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def load_dataset_with_template(
|
||||
dataset_name: str,
|
||||
split: str,
|
||||
input_column: str,
|
||||
template: str = None,
|
||||
tokenizer: transformers.PreTrainedTokenizer | None = None,
|
||||
) -> Tuple[Any, Any]:
|
||||
"""Loads dataset with templates.
|
||||
|
||||
Args:
|
||||
dataset_name: Name of the dataset or path to a custom dataset.
|
||||
split: Split of the dataset.
|
||||
input_column: The input column in the dataset to be used or updaded by the
|
||||
template. If it does not exist, the template's `prompt_no_input` will be
|
||||
used, and the input_column will be created.
|
||||
template: Name of the JSON template file under `templates/` or GCS path to
|
||||
the template file.
|
||||
tokenizer: The tokenizer to use for chat_template templates.
|
||||
|
||||
Returns:
|
||||
The raw dataset and the dataset compatible with the template.
|
||||
"""
|
||||
raw = _get_dataset(dataset_name, split=split)
|
||||
if template:
|
||||
templated = format_dataset(raw, input_column, template, tokenizer)
|
||||
else:
|
||||
templated = None
|
||||
|
||||
return raw, templated
|
||||
|
||||
|
||||
def validate_dataset_with_template(
|
||||
dataset_name: str,
|
||||
split: str,
|
||||
input_column: str,
|
||||
template: str,
|
||||
tokenizer: transformers.PreTrainedTokenizer | None = None,
|
||||
max_seq_length: int | None = None,
|
||||
use_multiprocessing: bool = False,
|
||||
validate_percentage_of_dataset: int | None = None,
|
||||
validate_k_rows_of_dataset: int | None = None,
|
||||
) -> Any:
|
||||
"""Validates dataset with templates.
|
||||
|
||||
This function will be used to load the dataset and validate it against the
|
||||
template. In case of validation, we also allow the users to load the dataset
|
||||
partially by allowing them to read x% or top k rows of the dataset. To
|
||||
validate the dataset, the template file must be available in the GCS bucket
|
||||
and the dataset must be available either in the GCS bucket or Hugging Face.
|
||||
|
||||
Args:
|
||||
dataset_name: Name of the dataset or path to a custom dataset.
|
||||
split: Split of the dataset.
|
||||
input_column: The input column in the dataset to be used or updaded by the
|
||||
template. If it does not exist, the template's `prompt_no_input` will be
|
||||
used, and the input_column will be created.
|
||||
template: Name of the JSON template file under `templates/` or GCS path to
|
||||
the template file.
|
||||
tokenizer: The tokenizer to use for chat_template templates.
|
||||
max_seq_length: The maximum sequence length.
|
||||
use_multiprocessing: If True, it will use multiprocessing to load the
|
||||
dataset.
|
||||
validate_percentage_of_dataset: The percentage of the dataset to load.
|
||||
validate_k_rows_of_dataset: The top k sequences to load from the dataset.
|
||||
|
||||
Returns:
|
||||
None if the validation is successful, otherwise returns the error message.
|
||||
"""
|
||||
if not template:
|
||||
raise ValueError("template is required for validate_dataset.")
|
||||
|
||||
if not dataset_name:
|
||||
raise ValueError("dataset_name is empty.")
|
||||
|
||||
if not split:
|
||||
raise ValueError("split is empty.")
|
||||
|
||||
split = _get_split_string(
|
||||
split,
|
||||
validate_percentage_of_dataset,
|
||||
validate_k_rows_of_dataset,
|
||||
)
|
||||
|
||||
num_proc = multiprocessing.cpu_count() if use_multiprocessing else 1
|
||||
|
||||
# gcsfuse cannot be used from the notebook runtime env. Hence, we have
|
||||
# to download dataset and template from gcs to local.
|
||||
if is_gcs_path(dataset_name):
|
||||
dataset_name = download_gcs_uri_to_local(dataset_name, LOCAL_BASE_MODEL_DIR)
|
||||
|
||||
if is_gcs_path(template):
|
||||
template_path = download_gcs_uri_to_local(template, LOCAL_TEMPLATE_DIR)
|
||||
elif os.path.isfile(_github_template_path(template)):
|
||||
template_path = _github_template_path(template)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Template file {template} does not exist. To validate the"
|
||||
" dataset, please provide a valid GCS path for the template or a valid"
|
||||
" template name from"
|
||||
f" https://github.com/GoogleCloudPlatform/{_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME}/tree/main/{_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR}."
|
||||
)
|
||||
|
||||
dataset = format_dataset(
|
||||
_get_dataset(dataset_name, split, num_proc),
|
||||
input_column,
|
||||
template_path,
|
||||
tokenizer,
|
||||
)
|
||||
|
||||
if tokenizer is not None:
|
||||
get_filtered_dataset(
|
||||
dataset=dataset,
|
||||
input_column=input_column,
|
||||
max_seq_length=max_seq_length,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
print(
|
||||
"Dataset {} is compatible with the {} template.".format(
|
||||
os.path.basename(dataset_name), os.path.basename(template)
|
||||
)
|
||||
)
|
||||
@@ -10,15 +10,28 @@ import open_clip
|
||||
import torch
|
||||
from ts.torch_handler.base_handler import BaseHandler
|
||||
|
||||
from util import constants
|
||||
from util import fileutils
|
||||
from util import image_format_converter
|
||||
from google3.cloud.ml.applications.vision.model_garden.model_oss.util import constants
|
||||
from google3.cloud.ml.applications.vision.model_garden.model_oss.util import fileutils
|
||||
from google3.cloud.ml.applications.vision.model_garden.model_oss.util import image_format_converter
|
||||
|
||||
|
||||
@enum.unique
|
||||
class Precision(enum.Enum):
|
||||
AMP = "amp"
|
||||
AMP_BF16 = "amp_bf16"
|
||||
AMP_BFLOAT16 = "amp_bfloat16"
|
||||
BF16 = "bf16"
|
||||
FP16 = "fp16"
|
||||
PURE_BF16 = "pure_bf16"
|
||||
PURE_FP16 = "pure_fp16"
|
||||
FP32 = "fp32"
|
||||
|
||||
|
||||
# Supported checkpoint&model pairs:
|
||||
# https://github.com/mlfoundations/open_clip#pretrained-model-interface
|
||||
_DEFAULT_CHECKPOINT = "openai"
|
||||
_DEFAULT_MODEL = "RN50"
|
||||
_BIOMED_CLIP_MODEL = "microsoft/BiomedCLIP"
|
||||
_DEFAULT_PRECISION = Precision.AMP
|
||||
_ZERO_CLASSIFICATION = "zero-shot-image-classification"
|
||||
_FEATURE_EMBEDDING = "feature-embedding"
|
||||
_VALID_TASKS = frozenset([_ZERO_CLASSIFICATION, _FEATURE_EMBEDDING])
|
||||
@@ -32,21 +45,6 @@ _TEXT_FEATURES_KEY = "text_features"
|
||||
class OpenclipHandler(BaseHandler):
|
||||
"""Custom handler for OpenCLIP."""
|
||||
|
||||
@enum.unique
|
||||
class Precision(enum.Enum):
|
||||
AMP = "amp"
|
||||
AMP_BF16 = "amp_bf16"
|
||||
AMP_BFLOAT16 = "amp_bfloat16"
|
||||
# For the difference between floating points and "pure" floating points, see
|
||||
# https://github.com/mlfoundations/open_clip/blob/0142d279298a4ca0138316286f775fe9d7bdbb94/src/open_clip/factory.py#L232C58-L232C58
|
||||
BF16 = "bf16"
|
||||
FP16 = "fp16"
|
||||
PURE_BF16 = "pure_bf16"
|
||||
PURE_FP16 = "pure_fp16"
|
||||
FP32 = "fp32"
|
||||
|
||||
_DEFAULT_PRECISION = Precision.AMP
|
||||
|
||||
def initialize(self, context: Any):
|
||||
"""Custom initialize."""
|
||||
|
||||
@@ -63,32 +61,29 @@ class OpenclipHandler(BaseHandler):
|
||||
)
|
||||
self.manifest = context.manifest
|
||||
|
||||
self.model_name = os.environ.get("MODEL", None)
|
||||
if not self.model_name:
|
||||
self.model_name = os.environ.get("MODEL_ID", _DEFAULT_MODEL)
|
||||
precision = os.environ.get("PRECISION", self._DEFAULT_PRECISION)
|
||||
checkpoint = os.environ.get("CHECKPOINT")
|
||||
model_name = os.environ.get("MODEL", _DEFAULT_MODEL)
|
||||
precision = os.environ.get("PRECISION", _DEFAULT_PRECISION)
|
||||
checkpoint = os.environ.get("CHECKPOINT", _DEFAULT_CHECKPOINT)
|
||||
self.task = os.environ.get("TASK", _FEATURE_EMBEDDING)
|
||||
if self.task not in _VALID_TASKS:
|
||||
raise ValueError(f"Invalid task: {self.task}.")
|
||||
logging.info(
|
||||
"Handler initializing task:%s, model:%s, precision:%s, checkpoint:%s",
|
||||
self.task,
|
||||
self.model_name,
|
||||
model_name,
|
||||
precision,
|
||||
checkpoint,
|
||||
)
|
||||
|
||||
if fileutils.is_gcs_path(checkpoint):
|
||||
if checkpoint != _DEFAULT_CHECKPOINT:
|
||||
local_fname = os.path.join(constants.LOCAL_MODEL_DIR, "model.pt")
|
||||
fileutils.download_gcs_file_to_local(checkpoint, local_fname)
|
||||
checkpoint = local_fname
|
||||
|
||||
self.model, self.preprocessor = open_clip.create_model_from_pretrained(
|
||||
self.model_name, pretrained=checkpoint, precision=precision
|
||||
self.model, _, self.preprocessor = open_clip.create_model_and_transforms(
|
||||
model_name, pretrained=checkpoint, precision=precision
|
||||
)
|
||||
self.model.to(self.device)
|
||||
self.tokenizer = open_clip.get_tokenizer(self.model_name)
|
||||
self.tokenizer = open_clip.get_tokenizer(model_name)
|
||||
|
||||
self.initialized = True
|
||||
|
||||
@@ -107,32 +102,9 @@ class OpenclipHandler(BaseHandler):
|
||||
processed_list.append(sample)
|
||||
return processed_list
|
||||
|
||||
def _biomedclip_inference(
|
||||
self, data: List[Dict[str, Any]], *args, **kwargs
|
||||
) -> List[List[float]]:
|
||||
"""Inference for BiomedCLIP model."""
|
||||
texts = torch.stack(
|
||||
[item[_TEXT_KEY][0] for item in data if _TEXT_KEY in item]
|
||||
).to(self.map_location)
|
||||
images = torch.stack(
|
||||
[item[_IMAGE_KEY][0] for item in data if _IMAGE_KEY in item]
|
||||
).to(self.map_location)
|
||||
if texts.shape[0] == 0 or images.shape[0] == 0:
|
||||
return []
|
||||
with torch.no_grad():
|
||||
image_features, text_features, logit_scale = self.model(images, texts)
|
||||
logits = (
|
||||
(logit_scale * image_features @ text_features.t())
|
||||
.detach()
|
||||
.softmax(dim=-1)
|
||||
)
|
||||
return logits.cpu().numpy().tolist()
|
||||
|
||||
def inference(
|
||||
self, data: List[Dict[str, Any]], *args, **kwargs
|
||||
) -> List[Dict[str, Any]]:
|
||||
if _BIOMED_CLIP_MODEL in self.model_name:
|
||||
return self._biomedclip_inference(data)
|
||||
feature_list = []
|
||||
with torch.no_grad(), torch.cuda.amp.autocast():
|
||||
for item in data:
|
||||
@@ -148,8 +120,6 @@ class OpenclipHandler(BaseHandler):
|
||||
|
||||
def postprocess(self, features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Postprocess the image/text featreus for downstream task."""
|
||||
if _BIOMED_CLIP_MODEL in self.model_name:
|
||||
return features
|
||||
preds = []
|
||||
if self.task == _FEATURE_EMBEDDING:
|
||||
for item in features:
|
||||
@@ -169,4 +139,4 @@ class OpenclipHandler(BaseHandler):
|
||||
text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
|
||||
preds.append(text_probs.tolist())
|
||||
|
||||
return preds
|
||||
return preds
|
||||
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
"""Causal language modeling with LoRA models."""
|
||||
|
||||
# pylint: disable=g-importing-member
|
||||
|
||||
from datasets import load_dataset
|
||||
from peft import get_peft_model
|
||||
from peft import LoraConfig
|
||||
import torch
|
||||
from torch import nn
|
||||
import transformers
|
||||
from transformers import AutoModelForCausalLM
|
||||
from transformers import AutoTokenizer
|
||||
from transformers import BitsAndBytesConfig
|
||||
from transformers import TrainingArguments
|
||||
from util import constants
|
||||
|
||||
|
||||
def finetune_causal_language_modeling(
|
||||
pretrained_model_id: str,
|
||||
dataset_name: str,
|
||||
output_dir: str,
|
||||
precision_mode: str = None,
|
||||
lora_rank: int = 16,
|
||||
lora_alpha: int = 32,
|
||||
lora_dropout: float = 0.05,
|
||||
warmup_steps: int = 10,
|
||||
max_steps: int = 10,
|
||||
learning_rate: float = 2e-4,
|
||||
local_pretrained_model_id: str = None,
|
||||
) -> None:
|
||||
"""Finetunes causal language modelings."""
|
||||
if precision_mode == constants.PRECISION_MODE_32:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
local_pretrained_model_id
|
||||
if local_pretrained_model_id
|
||||
else pretrained_model_id,
|
||||
torch_dtype=torch.float32,
|
||||
device_map="auto",
|
||||
)
|
||||
elif precision_mode == constants.PRECISION_MODE_16:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
local_pretrained_model_id
|
||||
if local_pretrained_model_id
|
||||
else pretrained_model_id,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
)
|
||||
elif precision_mode == constants.PRECISION_MODE_8:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_8bit=True, int8_threshold=0
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
local_pretrained_model_id
|
||||
if local_pretrained_model_id
|
||||
else pretrained_model_id,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
else:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
local_pretrained_model_id
|
||||
if local_pretrained_model_id
|
||||
else pretrained_model_id,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16,
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
local_pretrained_model_id
|
||||
if local_pretrained_model_id
|
||||
else pretrained_model_id
|
||||
)
|
||||
if "llama" in pretrained_model_id:
|
||||
tokenizer.pad_token = "[PAD]"
|
||||
|
||||
for param in model.parameters():
|
||||
# Freezes the model - train adapters later.
|
||||
param.requires_grad = False
|
||||
if param.ndim == 1:
|
||||
# Casts the small parameters (e.g. layernorm) to fp32 for stability.
|
||||
param.data = param.data.to(torch.float32)
|
||||
|
||||
# Reduces the number of stored activations.
|
||||
model.gradient_checkpointing_enable()
|
||||
model.enable_input_require_grads()
|
||||
|
||||
class CastOutputToFloat(nn.Sequential):
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x).to(torch.float32)
|
||||
|
||||
model.lm_head = CastOutputToFloat(model.lm_head)
|
||||
|
||||
config = LoraConfig(
|
||||
r=lora_rank,
|
||||
lora_alpha=lora_alpha,
|
||||
target_modules=["q_proj", "v_proj"],
|
||||
lora_dropout=lora_dropout,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
|
||||
model = get_peft_model(model, config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
data = load_dataset(dataset_name)
|
||||
data = data.map(
|
||||
lambda samples: tokenizer(samples["quote"]),
|
||||
batched=True,
|
||||
)
|
||||
|
||||
trainer = transformers.Trainer(
|
||||
model=model,
|
||||
train_dataset=data["train"],
|
||||
args=TrainingArguments(
|
||||
per_device_train_batch_size=4,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_steps=warmup_steps,
|
||||
max_steps=max_steps,
|
||||
learning_rate=learning_rate,
|
||||
fp16=True,
|
||||
logging_steps=1,
|
||||
output_dir=output_dir,
|
||||
ddp_find_unused_parameters=False,
|
||||
),
|
||||
data_collator=transformers.DataCollatorForLanguageModeling(
|
||||
tokenizer,
|
||||
mlm=False,
|
||||
),
|
||||
)
|
||||
# Silence the warnings. Please re-enable for inference!
|
||||
model.config.use_cache = False
|
||||
trainer.train()
|
||||
|
||||
model.save_pretrained(output_dir)
|
||||
-28
@@ -1,28 +0,0 @@
|
||||
# Base on pytorch-cuda image.
|
||||
FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
|
||||
|
||||
# Install tools.
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
RUN apt-get update
|
||||
RUN apt-get install -y --no-install-recommends apt-utils
|
||||
RUN apt-get install -y --no-install-recommends curl
|
||||
RUN apt-get install -y --no-install-recommends wget
|
||||
RUN apt-get install -y --no-install-recommends git
|
||||
|
||||
# Install libraries.
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
RUN python3 -m pip install --upgrade pip
|
||||
RUN pip install tokenizers==0.13.3
|
||||
RUN pip install accelerate==0.21.0
|
||||
RUN pip install sentencepiece==0.1.99
|
||||
RUN pip install datasets==2.14.4
|
||||
RUN pip install protobuf==4.24.1
|
||||
|
||||
# Install transformers
|
||||
RUN git clone https://github.com/huggingface/transformers.git
|
||||
WORKDIR transformers
|
||||
# Pin the commit to add-code-llama 08/25/2023
|
||||
RUN git reset --hard 015f8e110d270a0ad42de4ae5b98198d69eb1964
|
||||
RUN pip install -e .
|
||||
|
||||
ENTRYPOINT ["python","src/transformers/models/llama/convert_llama_weights_to_hf.py"]
|
||||
@@ -1,22 +0,0 @@
|
||||
# Dockerfile for Language Model Conversion.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/peft/dockerfile/conversion.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM tensorflow/build:2.14-python3.8
|
||||
|
||||
RUN git clone https://github.com/facebookresearch/llama-recipes.git && \
|
||||
cd llama-recipes && \
|
||||
pip install -r requirements.txt && \
|
||||
pip freeze | grep transformers && \
|
||||
git clone https://github.com/huggingface/transformers.git && \
|
||||
cd transformers && \
|
||||
pip install protobuf
|
||||
|
||||
WORKDIR /llama-recipes/transformers
|
||||
|
||||
ENTRYPOINT ["python","src/transformers/models/llama/convert_llama_weights_to_hf.py"]
|
||||
@@ -7,40 +7,39 @@
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM pytorch/torchserve:0.11.0-gpu
|
||||
FROM pytorch/torchserve:0.7.0-gpu
|
||||
|
||||
USER root
|
||||
|
||||
ENV INFER_PORT=7080
|
||||
ENV MNG_PORT=7081
|
||||
ENV MODEL="peft_serving"
|
||||
ENV infer_port=7080
|
||||
ENV mng_port=7081
|
||||
ENV model_name="peft_serving"
|
||||
ENV PATH="/home/model-server/:${PATH}"
|
||||
|
||||
RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
curl \
|
||||
wget \
|
||||
vim \
|
||||
git \
|
||||
git-lfs
|
||||
RUN git lfs install
|
||||
RUN apt-get autoremove -y
|
||||
|
||||
# Install libraries.
|
||||
ENV PIP_ROOT_USER_ACTION=ignore
|
||||
RUN python3 -m pip install --upgrade pip
|
||||
RUN pip install --upgrade torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
|
||||
RUN pip install --upgrade torch==2.0.1
|
||||
RUN pip install torchvision==0.15.2
|
||||
RUN pip install tokenizers==0.13.3
|
||||
RUN pip install accelerate==0.21.0
|
||||
RUN pip install sentencepiece==0.1.99
|
||||
RUN pip install grpcio-status==1.33.2
|
||||
RUN pip install protobuf==3.19.6
|
||||
RUN pip install peft==0.5.0
|
||||
RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
|
||||
RUN pip install datasets==2.14.4
|
||||
RUN pip install triton==3.0.0
|
||||
RUN pip install triton==2.0.0.dev20221120
|
||||
RUN pip install xformers==0.0.20
|
||||
RUN pip install google-cloud-storage
|
||||
RUN pip install absl-py
|
||||
RUN pip install google-cloud-storage==2.7.0
|
||||
RUN pip install absl-py==1.4.0
|
||||
RUN pip install scipy==1.10.1
|
||||
RUN pip install evaluate==0.4.0
|
||||
RUN pip install scikit-learn==1.2.2
|
||||
@@ -48,43 +47,52 @@ RUN pip install loralib==0.1.1
|
||||
RUN pip install bitsandbytes==0.39.0
|
||||
RUN pip install trl==0.4.4
|
||||
RUN pip install einops==0.6.1
|
||||
RUN pip install optimum==1.13.2
|
||||
RUN pip install auto-gptq==0.4.2
|
||||
RUN pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.7/autoawq-0.1.7+cu118-cp39-cp39-linux_x86_64.whl
|
||||
RUN pip install diffusers==0.27.2
|
||||
RUN pip install tiktoken==0.6.0
|
||||
RUn pip install git+https://github.com/huggingface/transformers.git@76fa17c1663a0efeca7208c20579833365584889
|
||||
RUN pip install pynvml==11.4.0
|
||||
RUN pip install -i https://test.pypi.org/simple/ bitsandbytes
|
||||
|
||||
# Install diffusers from source.
|
||||
RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
|
||||
WORKDIR diffusers
|
||||
RUN pip install -e .
|
||||
WORKDIR /home/model-server
|
||||
|
||||
# Install transformers from source.
|
||||
RUN git clone --depth 1 --branch v4.31.0 https://github.com/huggingface/transformers.git
|
||||
# The patch is used to change the transformers loading model behavior:
|
||||
# 1) For models on Huggingface hub: if the model has multiple shards, each shard
|
||||
# will be downloaded separately and get deleted after loading to GPU.
|
||||
# 2) For models on local disk: if a model bin file is actually a text file
|
||||
# recording a GCS path, the model file will be downloaded and get deleted
|
||||
# after loading to GPU.
|
||||
COPY model_oss/peft/hf_transformers_lazy_download.patch /home/model-server/hf_transformers_lazy_download.patch
|
||||
WORKDIR transformers
|
||||
RUN git apply /home/model-server/hf_transformers_lazy_download.patch
|
||||
RUN pip install -e .
|
||||
WORKDIR /home/model-server
|
||||
|
||||
# Copy license.
|
||||
WORKDIR /home/model-server
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
# Copy model artifacts.
|
||||
COPY model_oss/peft/handler.py /home/model-server/handler.py
|
||||
COPY model_oss/peft/config.properties /home/model-server/config.properties
|
||||
COPY model_oss/util/ /home/model-server/util/
|
||||
COPY model_oss/util/pytorch_startup_prober.sh /model_garden/scripts/pytorch_startup_prober.sh
|
||||
ENV PYTHONPATH /home/model-server/
|
||||
|
||||
# Expose ports.
|
||||
EXPOSE ${INFER_PORT}
|
||||
EXPOSE ${MNG_PORT}
|
||||
EXPOSE ${infer_port}
|
||||
EXPOSE ${mng_port}
|
||||
|
||||
# Set environments.
|
||||
ENV TASK "causal-language-modeling-lora"
|
||||
ENV BASE_MODEL_ID ""
|
||||
ENV MODEL_ID ""
|
||||
ENV BASE_MODEL_ID "openlm-research/open_llama_7b"
|
||||
ENV PRECISION_LOADING_MODE "float16"
|
||||
ENV FINETUNED_LORA_MODEL_PATH ""
|
||||
ENV TRUST_REMOTE_CODE ""
|
||||
|
||||
|
||||
# Archive model artifacts and dependencies.
|
||||
# Do not set --model-file and --serialized-file because model and checkpoint
|
||||
# will be dynamically loaded in handler.py.
|
||||
RUN torch-model-archiver \
|
||||
--model-name=${MODEL} \
|
||||
--model-name=${model_name} \
|
||||
--version=1.0 \
|
||||
--handler=/home/model-server/handler.py \
|
||||
--runtime=python3 \
|
||||
@@ -95,5 +103,5 @@ RUN torch-model-archiver \
|
||||
# Run Torchserve HTTP serve to respond to prediction requests.
|
||||
CMD ["torchserve", "--start", \
|
||||
"--ts-config", "/home/model-server/config.properties", \
|
||||
"--models", "${MODEL}=${MODEL}.mar", \
|
||||
"--models", "${model_name}=${model_name}.mar", \
|
||||
"--model-store", "/home/model-server/model-store"]
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# Dockerfile for PEFT Training.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/peft/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
# Builds GPU docker image of PyTorch
|
||||
# Uses multi-staged approach to reduce size
|
||||
# Stage 1
|
||||
# Use base conda image to reduce time
|
||||
FROM continuumio/miniconda3:latest AS compile-image
|
||||
# Specify py version
|
||||
ENV PYTHON_VERSION=3.8
|
||||
# Install apt libs - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
|
||||
RUN apt-get update && \
|
||||
apt-get install -y curl git wget software-properties-common git-lfs && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists*
|
||||
|
||||
# Install audio-related libraries
|
||||
RUN apt-get update && \
|
||||
apt install -y ffmpeg
|
||||
|
||||
RUN apt install -y libsndfile1-dev
|
||||
RUN git lfs install
|
||||
|
||||
# Create our conda env - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
|
||||
RUN conda create --name peft python=${PYTHON_VERSION} ipython jupyter pip
|
||||
RUN python3 -m pip install --no-cache-dir --upgrade pip
|
||||
|
||||
# Below is copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
|
||||
# We don't install pytorch here yet since CUDA isn't available
|
||||
# instead we use the direct torch wheel
|
||||
ENV PATH /opt/conda/envs/peft/bin:$PATH
|
||||
# Activate our bash shell
|
||||
RUN chsh -s /bin/bash
|
||||
SHELL ["/bin/bash", "-c"]
|
||||
# Activate the conda env and install transformers + accelerate from source
|
||||
RUN source activate peft
|
||||
RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/transformers
|
||||
RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/accelerate
|
||||
RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft#egg=peft[test]
|
||||
RUN python3 -m pip install --no-cache-dir bitsandbytes
|
||||
|
||||
# Stage 2
|
||||
FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
|
||||
COPY --from=compile-image /opt/conda /opt/conda
|
||||
ENV PATH /opt/conda/bin:$PATH
|
||||
|
||||
# Install apt libs
|
||||
RUN apt-get update && \
|
||||
apt-get install -y curl git wget vim && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists*
|
||||
|
||||
# Copy license.
|
||||
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
|
||||
|
||||
RUN echo "source activate peft" >> ~/.profile
|
||||
|
||||
# Install libraries.
|
||||
RUN pip install --upgrade torch==2.0.1
|
||||
RUN pip install torchvision==0.15.2
|
||||
RUN pip install git+https://github.com/huggingface/transformers@de9255de27abfcae4a1f816b904915f0b1e23cd9
|
||||
RUN pip install transformers -U
|
||||
RUN pip install accelerate==0.21.0
|
||||
RUN pip install sentencepiece==0.1.99
|
||||
RUN pip install grpcio-status==1.33.2
|
||||
RUN pip install protobuf==3.19.6
|
||||
RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
|
||||
RUN pip install datasets==2.9.0
|
||||
RUN pip install triton==2.0.0.dev20221120
|
||||
RUN pip install xformers==0.0.20
|
||||
RUN pip install Jinja2==3.1.2
|
||||
RUN pip install ftfy==6.1.1
|
||||
RUN pip install cloudml-hypertune==0.1.0.dev6
|
||||
RUN pip install tensorboard==2.12.0
|
||||
RUN pip install scipy==1.10.1
|
||||
RUN pip install evaluate==0.4.0
|
||||
RUN pip install scikit-learn==1.2.2
|
||||
RUN pip install loralib==0.1.1
|
||||
RUN pip install bitsandbytes==0.39.0
|
||||
RUN pip install trl==0.4.4
|
||||
RUN pip install einops==0.6.1
|
||||
RUN pip install google-cloud-storage==2.7.0
|
||||
|
||||
RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
|
||||
WORKDIR diffusers
|
||||
RUN pip install -e .
|
||||
|
||||
# Switch to diffusers examples folder.
|
||||
WORKDIR examples
|
||||
|
||||
# NOTE: use 'sed' to modify train_text_to_image_lora.py to
|
||||
# fix the bug for accelerator.
|
||||
RUN sed -i \
|
||||
"s#logging_dir=logging_dir#project_dir=logging_dir#g" \
|
||||
text_to_image/train_text_to_image_lora.py
|
||||
|
||||
# Config accelerate.
|
||||
RUN mkdir -p ./vertex_vision_model_garden_peft/
|
||||
COPY model_oss/peft/train.sh ./vertex_vision_model_garden_peft/train.sh
|
||||
COPY model_oss/peft/*.py ./vertex_vision_model_garden_peft/
|
||||
COPY model_oss/util /diffusers/examples/util
|
||||
ENV PYTHONPATH /diffusers/examples/
|
||||
|
||||
# Generate accelerate config at the beginning of docker run.
|
||||
ENTRYPOINT ["python3", "vertex_vision_model_garden_peft/main.py"]
|
||||
@@ -3,17 +3,15 @@
|
||||
# pylint: disable=g-importing-member
|
||||
# pylint: disable=logging-fstring-interpolation
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, List, Tuple
|
||||
from typing import Any, List
|
||||
|
||||
from absl import logging
|
||||
from awq import AutoAWQForCausalLM
|
||||
from diffusers import DPMSolverMultistepScheduler
|
||||
from diffusers import StableDiffusionPipeline
|
||||
from peft import PeftModel
|
||||
from PIL import Image
|
||||
import psutil
|
||||
import torch
|
||||
import transformers
|
||||
from transformers import AutoModelForCausalLM
|
||||
@@ -26,12 +24,6 @@ from util import constants
|
||||
from util import fileutils
|
||||
from util import image_format_converter
|
||||
|
||||
if os.path.exists(constants.SHARED_MEM_DIR):
|
||||
logging.info(
|
||||
"SharedMemorySizeMb: %s",
|
||||
psutil.disk_usage(constants.SHARED_MEM_DIR).free / 1e6,
|
||||
)
|
||||
|
||||
# Tasks
|
||||
TEXT_TO_IMAGE_LORA = "text-to-image-lora"
|
||||
SEQUENCE_CLASSIFICATION_LORA = "sequence-classification-lora"
|
||||
@@ -41,13 +33,8 @@ INSTRUCT_LORA = "instruct-lora"
|
||||
# Inference parameters.
|
||||
_NUM_INFERENCE_STEPS = 25
|
||||
_MAX_LENGTH_DEFAULT = 200
|
||||
_MAX_TOKENS_DEFAULT = None
|
||||
_TEMPERATURE_DEFAULT = 1.0
|
||||
_TOP_P_DEFAULT = 1.0
|
||||
_TOP_K_DEFAULT = 10
|
||||
|
||||
logging.set_verbosity(os.environ.get("LOG_LEVEL", logging.INFO))
|
||||
|
||||
|
||||
class PeftHandler(BaseHandler):
|
||||
"""Custom handler for Peft models."""
|
||||
@@ -72,48 +59,16 @@ class PeftHandler(BaseHandler):
|
||||
"PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
|
||||
)
|
||||
self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
|
||||
trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", None)
|
||||
if trust_remote_code == "false":
|
||||
self.trust_remote_code = False
|
||||
else:
|
||||
self.trust_remote_code = True
|
||||
|
||||
# If present, the path of the model in the container.
|
||||
aip_storage_dir = os.environ.get("AIP_STORAGE_DIR", None)
|
||||
|
||||
# If present, the URI of the model in a google owned GCS bucket.
|
||||
aip_storage_uri = os.environ.get("AIP_STORAGE_URI", None)
|
||||
|
||||
model_id = os.environ.get("MODEL_ID", None)
|
||||
base_model_id = os.environ.get("BASE_MODEL_ID", None)
|
||||
|
||||
self.model_id = None
|
||||
if aip_storage_dir:
|
||||
self.model_id = aip_storage_dir
|
||||
logging.info(f"Loaded base model from AIP_STORAGE_DIR: {self.model_id}.")
|
||||
elif aip_storage_uri:
|
||||
self.model_id = aip_storage_uri
|
||||
logging.info(f"Loaded base model from AIP_STORAGE_URI: {self.model_id}.")
|
||||
elif model_id:
|
||||
self.model_id = model_id
|
||||
logging.info(f"Loaded base model from MODEL_ID: {self.model_id}.")
|
||||
elif base_model_id:
|
||||
# Note: BASE_MODEL_ID has been unified with MODEL_ID.
|
||||
# MODEL_ID should be used whenever possible.
|
||||
self.model_id = base_model_id
|
||||
logging.info(f"Loaded base model from BASE_MODEL_ID: {self.model_id}.")
|
||||
|
||||
self.quantization = os.environ.get("QUANTIZATION", None)
|
||||
|
||||
if not self.model_id:
|
||||
raise ValueError("Base model id is must be set.")
|
||||
if fileutils.is_gcs_path(self.model_id):
|
||||
self.base_model_id = os.environ.get(
|
||||
"BASE_MODEL_ID", "openlm-research/open_llama_7b"
|
||||
)
|
||||
if fileutils.is_gcs_path(self.base_model_id):
|
||||
fileutils.download_gcs_dir_to_local(
|
||||
self.model_id,
|
||||
self.base_model_id,
|
||||
constants.LOCAL_BASE_MODEL_DIR,
|
||||
skip_hf_model_bin=True,
|
||||
)
|
||||
self.model_id = constants.LOCAL_BASE_MODEL_DIR
|
||||
self.base_model_id = constants.LOCAL_BASE_MODEL_DIR
|
||||
self.finetuned_lora_model_path = os.environ.get(
|
||||
"FINETUNED_LORA_MODEL_PATH", ""
|
||||
)
|
||||
@@ -124,39 +79,38 @@ class PeftHandler(BaseHandler):
|
||||
self.finetuned_lora_model_path = constants.LOCAL_MODEL_DIR
|
||||
|
||||
logging.info(
|
||||
f"Using task:{self.task}, base model:{self.model_id}, lora model:"
|
||||
f" {self.finetuned_lora_model_path}, precision {self.precision_mode}."
|
||||
f"Using task:{self.task}, base model:{self.base_model_id}, lora model:"
|
||||
f" {self.finetuned_lora_model_path}, and precision"
|
||||
f" {self.precision_mode}."
|
||||
)
|
||||
|
||||
self.pipeline = None
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
start_time = time.perf_counter()
|
||||
logging.info("Started PEFT handler initialization at: %s", start_time)
|
||||
if self.task == TEXT_TO_IMAGE_LORA:
|
||||
pipeline = StableDiffusionPipeline.from_pretrained(
|
||||
self.model_id, torch_dtype=torch.float16
|
||||
self.base_model_id, torch_dtype=torch.float16
|
||||
)
|
||||
logging.debug("Initialized the base model for text to image.")
|
||||
pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
|
||||
pipeline.scheduler.config
|
||||
)
|
||||
logging.debug("Initialized the scheduler for text to image.")
|
||||
if self.finetuned_lora_model_path:
|
||||
pipeline.unet.load_attn_procs(self.finetuned_lora_model_path)
|
||||
logging.debug("Initialized the LoRA model for text to image.")
|
||||
# This is to reduce GPU memory requirements.
|
||||
pipeline.enable_xformers_memory_efficient_attention()
|
||||
pipeline = pipeline.to(self.map_location)
|
||||
# Reduces memory footprint.
|
||||
pipeline.enable_attention_slicing()
|
||||
if self.finetuned_lora_model_path:
|
||||
pipeline.load_lora_weights(self.finetuned_lora_model_path)
|
||||
logging.debug("Initialized the LoRA model for text to image.")
|
||||
self.pipeline = pipeline
|
||||
logging.info("Initialized the text to image pipelines.")
|
||||
elif self.task == SEQUENCE_CLASSIFICATION_LORA:
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.model_id)
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
|
||||
logging.debug("Initialized the tokenizer for sequence classification.")
|
||||
model = AutoModelForSequenceClassification.from_pretrained(
|
||||
self.model_id, torch_dtype=torch.float16
|
||||
self.base_model_id, torch_dtype=torch.float16
|
||||
)
|
||||
logging.debug("Initialized the base model for sequence classification.")
|
||||
if self.finetuned_lora_model_path:
|
||||
@@ -168,86 +122,55 @@ class PeftHandler(BaseHandler):
|
||||
elif (
|
||||
self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
|
||||
):
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.model_id,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
|
||||
logging.debug("Initialized the tokenizer.")
|
||||
if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
|
||||
if self.quantization == constants.AWQ:
|
||||
model = AutoAWQForCausalLM.from_quantized(
|
||||
self.model_id,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
if self.precision_mode == constants.PRECISION_MODE_32:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.base_model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.float32,
|
||||
device_map="auto",
|
||||
)
|
||||
elif self.precision_mode == constants.PRECISION_MODE_16:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.base_model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
)
|
||||
elif self.precision_mode == constants.PRECISION_MODE_8:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_8bit=True, int8_threshold=0
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.base_model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
elif self.quantization == constants.GPTQ or not self.quantization:
|
||||
if self.precision_mode == constants.PRECISION_MODE_32:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.float32,
|
||||
device_map="auto",
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
elif self.precision_mode == constants.PRECISION_MODE_16B:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="auto",
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
elif self.precision_mode == constants.PRECISION_MODE_16:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
elif self.precision_mode == constants.PRECISION_MODE_8:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_8bit=True, int8_threshold=0
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto",
|
||||
quantization_config=quantization_config,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
return_dict=True,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16,
|
||||
quantization_config=quantization_config,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid QUANTIZATION value: {self.quantization}")
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.base_model_id,
|
||||
return_dict=True,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16,
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
device_map="auto",
|
||||
)
|
||||
except: # pylint: disable=bare-except
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_id,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
device_map="auto",
|
||||
)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
self.base_model_id,
|
||||
return_dict=True,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
)
|
||||
logging.debug("Initialized the base model.")
|
||||
if self.finetuned_lora_model_path:
|
||||
model = PeftModel.from_pretrained(model, self.finetuned_lora_model_path)
|
||||
@@ -263,57 +186,35 @@ class PeftHandler(BaseHandler):
|
||||
raise ValueError(f"Invalid TASK: {self.task}")
|
||||
|
||||
self.initialized = True
|
||||
end_time = time.perf_counter()
|
||||
logging.info("The PEFT handler was initialize at: %s", end_time)
|
||||
logging.info("Handler initiation took %s seconds", end_time - start_time)
|
||||
logging.info("The PEFT handler was initialized.")
|
||||
|
||||
def preprocess(self, data: Any) -> Any:
|
||||
"""Preprocesses input data."""
|
||||
# Assumes that the parameters are same in one request. We parse the
|
||||
# parameters from the first instance for all instances in one request.
|
||||
# For generation length: `max_length` defines the maximum length of the
|
||||
# sequence to be generated, including both input and output tokens.
|
||||
# `max_length` is overridden by `max_new_tokens` if also set.
|
||||
# `max_new_tokens` defines the maximum number of new tokens to generate,
|
||||
# ignoring the current number of tokens.
|
||||
# Reference:
|
||||
# https://github.com/huggingface/transformers/blob/574a5384557b1aaf98ddb13ea9eb0a0ee8ff2cb2/src/transformers/generation/configuration_utils.py#L69-L73
|
||||
max_length = _MAX_LENGTH_DEFAULT
|
||||
max_tokens = _MAX_TOKENS_DEFAULT
|
||||
temperature = _TEMPERATURE_DEFAULT
|
||||
top_p = _TOP_P_DEFAULT
|
||||
top_k = _TOP_K_DEFAULT
|
||||
|
||||
prompts = [item["prompt"] for item in data]
|
||||
if "max_length" in data[0]:
|
||||
max_length = data[0]["max_length"]
|
||||
if "max_tokens" in data[0]:
|
||||
max_tokens = data[0]["max_tokens"]
|
||||
if "temperature" in data[0]:
|
||||
temperature = data[0]["temperature"]
|
||||
if "top_p" in data[0]:
|
||||
top_p = data[0]["top_p"]
|
||||
if "top_k" in data[0]:
|
||||
top_k = data[0]["top_k"]
|
||||
|
||||
return prompts, max_length, max_tokens, temperature, top_p, top_k
|
||||
return prompts, max_length, top_k
|
||||
|
||||
def inference(
|
||||
self, data: Any, *args, **kwargs
|
||||
) -> Tuple[List[str], List[Image.Image]]:
|
||||
def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
|
||||
"""Runs the inference."""
|
||||
prompts, max_length, max_tokens, temperature, top_p, top_k = data
|
||||
prompts, max_length, top_k = data
|
||||
logging.debug(
|
||||
f"Inference prompts={prompts}, max_length={max_length},"
|
||||
f" max_tokens={max_tokens}, temperature={temperature}, top_p={top_p},"
|
||||
f" top_k={top_k}."
|
||||
f"Inference prompts={prompts}, max_length={max_length}, top_k={top_k}."
|
||||
)
|
||||
if self.task == TEXT_TO_IMAGE_LORA:
|
||||
predicted_results = self.pipeline(
|
||||
prompt=prompts, num_inference_steps=_NUM_INFERENCE_STEPS
|
||||
).images
|
||||
elif self.task == SEQUENCE_CLASSIFICATION_LORA:
|
||||
encoded_input = self.tokenizer(prompts, return_tensors="pt", padding=True)
|
||||
encoded_input = self.tokenizer(prompts, return_tensors="pt")
|
||||
encoded_input.to(self.map_location)
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**encoded_input)
|
||||
@@ -325,43 +226,25 @@ class PeftHandler(BaseHandler):
|
||||
predicted_results = self.pipeline(
|
||||
prompts,
|
||||
max_length=max_length,
|
||||
max_new_tokens=max_tokens,
|
||||
do_sample=True,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
num_return_sequences=1,
|
||||
eos_token_id=self.tokenizer.eos_token_id,
|
||||
return_full_text=False,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid TASK: {self.task}")
|
||||
return prompts, predicted_results
|
||||
return predicted_results
|
||||
|
||||
def postprocess(self, data: Any) -> List[str]:
|
||||
"""Postprocesses output data."""
|
||||
prompts, predicted_results = data
|
||||
if self.task == TEXT_TO_IMAGE_LORA:
|
||||
# Converts the images to base64 string.
|
||||
outputs = [
|
||||
image_format_converter.image_to_base64(image)
|
||||
for image in predicted_results
|
||||
image_format_converter.image_to_base64(image) for image in data
|
||||
]
|
||||
elif self.task == SEQUENCE_CLASSIFICATION_LORA:
|
||||
outputs = predicted_results
|
||||
else:
|
||||
outputs = []
|
||||
for prompt, predicted_result in zip(prompts, predicted_results):
|
||||
formatted_output = self._format_text_generation_output(
|
||||
prompt=prompt, output=predicted_result[0]["generated_text"]
|
||||
)
|
||||
outputs.append(formatted_output)
|
||||
outputs = data
|
||||
return outputs
|
||||
|
||||
def _format_text_generation_output(self, prompt: str, output: str) -> str:
|
||||
"""Formats text generation output."""
|
||||
output = output.strip("\n")
|
||||
return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
|
||||
|
||||
|
||||
# pylint: enable=logging-fstring-interpolation
|
||||
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
|
||||
index 45459ed..32527f4 100644
|
||||
--- a/src/transformers/modeling_utils.py
|
||||
+++ b/src/transformers/modeling_utils.py
|
||||
@@ -32,6 +32,8 @@ import torch
|
||||
from packaging import version
|
||||
from torch import Tensor, nn
|
||||
from torch.nn import CrossEntropyLoss
|
||||
+from huggingface_hub import hf_hub_download
|
||||
+from google.cloud import storage
|
||||
|
||||
from .activations import get_activation
|
||||
from .configuration_utils import PretrainedConfig
|
||||
@@ -442,6 +444,29 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
|
||||
"""
|
||||
Reads a PyTorch checkpoint file, returning properly formatted errors if they arise.
|
||||
"""
|
||||
+ delete_download = False
|
||||
+ tmp_dir = "/tmp/model"
|
||||
+ os.makedirs(tmp_dir, exist_ok=True)
|
||||
+ if isinstance(checkpoint_file, dict):
|
||||
+ # Download model file from huggingface
|
||||
+ print(f"==> Download model from HF: {checkpoint_file}")
|
||||
+ checkpoint_file = hf_hub_download(
|
||||
+ local_dir=tmp_dir, local_dir_use_symlinks=False, force_download=True, resume_download=True, **checkpoint_file)
|
||||
+ delete_download = True
|
||||
+ else:
|
||||
+ with open(checkpoint_file, "rb") as f:
|
||||
+ is_gcs_file = (f.read(2) == b"gs")
|
||||
+ if is_gcs_file:
|
||||
+ # Download model file from GCS
|
||||
+ with open(checkpoint_file, "r") as f:
|
||||
+ gcs_file = f.read()
|
||||
+ checkpoint_file = os.path.join(tmp_dir, gcs_file.split("/")[-1])
|
||||
+ print(f"==> Download model from GCS: {gcs_file} to: {checkpoint_file}")
|
||||
+ client = storage.Client()
|
||||
+ with open(checkpoint_file, 'wb') as f:
|
||||
+ client.download_blob_to_file(gcs_file, f)
|
||||
+ delete_download = True
|
||||
+
|
||||
if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
|
||||
# Check format of the archive
|
||||
with safe_open(checkpoint_file, framework="pt") as f:
|
||||
@@ -455,9 +480,9 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
|
||||
raise NotImplementedError(
|
||||
f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
|
||||
)
|
||||
- return safe_load_file(checkpoint_file)
|
||||
+ state_dict = safe_load_file(checkpoint_file)
|
||||
try:
|
||||
- return torch.load(checkpoint_file, map_location="cpu")
|
||||
+ state_dict = torch.load(checkpoint_file, map_location="cpu")
|
||||
except Exception as e:
|
||||
try:
|
||||
with open(checkpoint_file) as f:
|
||||
@@ -478,6 +503,10 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
|
||||
f"at '{checkpoint_file}'. "
|
||||
"If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
|
||||
)
|
||||
+ if delete_download:
|
||||
+ print(f"==> Delete downloaded model: {checkpoint_file}")
|
||||
+ os.remove(checkpoint_file)
|
||||
+ return state_dict
|
||||
|
||||
|
||||
def set_initialized_submodules(model, state_dict_keys):
|
||||
@@ -3179,7 +3208,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
|
||||
return mismatched_keys
|
||||
|
||||
if resolved_archive_file is not None:
|
||||
- folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
|
||||
+ if isinstance(resolved_archive_file, str):
|
||||
+ folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
|
||||
+ else:
|
||||
+ folder = None
|
||||
else:
|
||||
folder = None
|
||||
if device_map is not None and is_safetensors:
|
||||
diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
|
||||
index ffed743..4b15770 100644
|
||||
--- a/src/transformers/utils/hub.py
|
||||
+++ b/src/transformers/utils/hub.py
|
||||
@@ -414,20 +414,34 @@ def cached_file(
|
||||
user_agent = http_user_agent(user_agent)
|
||||
try:
|
||||
# Load from URL or cache if already cached
|
||||
- resolved_file = hf_hub_download(
|
||||
- path_or_repo_id,
|
||||
- filename,
|
||||
- subfolder=None if len(subfolder) == 0 else subfolder,
|
||||
- repo_type=repo_type,
|
||||
- revision=revision,
|
||||
- cache_dir=cache_dir,
|
||||
- user_agent=user_agent,
|
||||
- force_download=force_download,
|
||||
- proxies=proxies,
|
||||
- resume_download=resume_download,
|
||||
- use_auth_token=use_auth_token,
|
||||
- local_files_only=local_files_only,
|
||||
- )
|
||||
+ if filename.endswith(".bin"):
|
||||
+ # NOTE: To save disk we do not download bin file eagerly. Do not support safetensors.
|
||||
+ resolved_file = dict(
|
||||
+ repo_id=path_or_repo_id,
|
||||
+ filename=filename,
|
||||
+ subfolder=None if len(subfolder) == 0 else subfolder,
|
||||
+ repo_type=repo_type,
|
||||
+ revision=revision,
|
||||
+ user_agent=user_agent,
|
||||
+ proxies=proxies,
|
||||
+ use_auth_token=use_auth_token,
|
||||
+ )
|
||||
+ print(f"--> Apply lazy download to bin file: {resolved_file}")
|
||||
+ else:
|
||||
+ resolved_file = hf_hub_download(
|
||||
+ path_or_repo_id,
|
||||
+ filename,
|
||||
+ subfolder=None if len(subfolder) == 0 else subfolder,
|
||||
+ repo_type=repo_type,
|
||||
+ revision=revision,
|
||||
+ cache_dir=cache_dir,
|
||||
+ user_agent=user_agent,
|
||||
+ force_download=force_download,
|
||||
+ proxies=proxies,
|
||||
+ resume_download=resume_download,
|
||||
+ use_auth_token=use_auth_token,
|
||||
+ local_files_only=local_files_only,
|
||||
+ )
|
||||
|
||||
except RepositoryNotFoundError:
|
||||
raise EnvironmentError(
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Instruct/Chat with LoRA models."""
|
||||
|
||||
# pylint: disable=g-importing-member
|
||||
from datasets import load_dataset
|
||||
from peft import LoraConfig
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM
|
||||
from transformers import AutoTokenizer
|
||||
from transformers import BitsAndBytesConfig
|
||||
from transformers import TrainingArguments
|
||||
from trl import SFTTrainer
|
||||
|
||||
|
||||
def finetune_instruct(
|
||||
pretrained_model_id: str,
|
||||
dataset_name: str,
|
||||
output_dir: str,
|
||||
lora_rank: int = 64,
|
||||
lora_alpha: int = 16,
|
||||
lora_dropout: float = 0.1,
|
||||
warmup_ratio: int = 0.03,
|
||||
max_steps: int = 10,
|
||||
max_seq_length: int = 512,
|
||||
learning_rate: float = 2e-4,
|
||||
) -> None:
|
||||
"""Finetunes instruct."""
|
||||
dataset = load_dataset(dataset_name, split="train")
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
)
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
pretrained_model_id,
|
||||
quantization_config=bnb_config,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
model.config.use_cache = False
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
pretrained_model_id, trust_remote_code=True
|
||||
)
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
peft_config = LoraConfig(
|
||||
lora_alpha=lora_alpha,
|
||||
lora_dropout=lora_dropout,
|
||||
r=lora_rank,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
target_modules=[
|
||||
"query_key_value",
|
||||
"dense",
|
||||
"dense_h_to_4h",
|
||||
"dense_4h_to_h",
|
||||
],
|
||||
)
|
||||
|
||||
per_device_train_batch_size = 4
|
||||
gradient_accumulation_steps = 4
|
||||
optim = "paged_adamw_32bit"
|
||||
save_steps = 10
|
||||
logging_steps = 10
|
||||
max_grad_norm = 0.3
|
||||
lr_scheduler_type = "constant"
|
||||
|
||||
training_arguments = TrainingArguments(
|
||||
output_dir=output_dir,
|
||||
per_device_train_batch_size=per_device_train_batch_size,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps,
|
||||
optim=optim,
|
||||
save_steps=save_steps,
|
||||
logging_steps=logging_steps,
|
||||
learning_rate=learning_rate,
|
||||
fp16=True,
|
||||
max_grad_norm=max_grad_norm,
|
||||
max_steps=max_steps,
|
||||
warmup_ratio=warmup_ratio,
|
||||
group_by_length=True,
|
||||
lr_scheduler_type=lr_scheduler_type,
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
train_dataset=dataset,
|
||||
peft_config=peft_config,
|
||||
dataset_text_field="text",
|
||||
max_seq_length=max_seq_length,
|
||||
tokenizer=tokenizer,
|
||||
args=training_arguments,
|
||||
)
|
||||
for name, module in trainer.model.named_modules():
|
||||
if "norm" in name:
|
||||
module = module.to(torch.float32)
|
||||
trainer.train()
|
||||
@@ -0,0 +1,177 @@
|
||||
"""Main function to start PEFT finetuning."""
|
||||
import subprocess
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
from absl import logging
|
||||
|
||||
from peft import causal_language_modeling_lora
|
||||
from peft import instruct_lora
|
||||
from peft import sequence_classification_lora
|
||||
from util import constants
|
||||
from util import fileutils
|
||||
|
||||
_TASK = flags.DEFINE_string(
|
||||
'task',
|
||||
constants.CAUSAL_LANGUAGE_MODELING_LORA,
|
||||
'The supported PEFT tasks.',
|
||||
)
|
||||
|
||||
_PRETRAINED_MODEL_ID = flags.DEFINE_string(
|
||||
'pretrained_model_id',
|
||||
None,
|
||||
'The pretrained model id. Supported models can be causal language modeling'
|
||||
' models from https://github.com/huggingface/peft/tree/main.',
|
||||
required=True,
|
||||
)
|
||||
|
||||
_DATASET_NAME = flags.DEFINE_string(
|
||||
'dataset_name',
|
||||
None,
|
||||
'The dataset name in huggingface.',
|
||||
required=True,
|
||||
)
|
||||
|
||||
_OUTPUT_DIR = flags.DEFINE_string(
|
||||
'output_dir',
|
||||
None,
|
||||
'The output directory.',
|
||||
required=True,
|
||||
)
|
||||
|
||||
_PRECISION_MODE = flags.DEFINE_string(
|
||||
'precision_mode',
|
||||
constants.PRECISION_MODE_16,
|
||||
'Supported finetuning precision_modes are `{}` and `{}`.'.format(
|
||||
constants.PRECISION_MODE_8, constants.PRECISION_MODE_16
|
||||
),
|
||||
)
|
||||
|
||||
_LORA_RANK = flags.DEFINE_integer(
|
||||
'lora_rank',
|
||||
16,
|
||||
'The rank of the update matrices, expressed in int. Lower rank results in'
|
||||
' smaller update matrices with fewer trainable parameters, referring to'
|
||||
' https://huggingface.co/docs/peft/conceptual_guides/lora.',
|
||||
)
|
||||
|
||||
_LORA_ALPHA = flags.DEFINE_integer(
|
||||
'lora_alpha',
|
||||
32,
|
||||
'LoRA scaling factor, referring to'
|
||||
' https://huggingface.co/docs/peft/conceptual_guides/lora.',
|
||||
)
|
||||
|
||||
_LORA_DROPOUT = flags.DEFINE_float(
|
||||
'lora_dropout',
|
||||
0.05,
|
||||
'dropout probability of the LoRA layers, referring to'
|
||||
' https://huggingface.co/docs/peft/task_guides/token-classification-lora.',
|
||||
)
|
||||
|
||||
_WARMUP_STEPS = flags.DEFINE_integer(
|
||||
'warmup_steps',
|
||||
10,
|
||||
'Number of steps for the warmup in the learning rate scheduler.',
|
||||
)
|
||||
|
||||
_WARMUP_RATIO = flags.DEFINE_float(
|
||||
'warmup_ratio',
|
||||
0.03,
|
||||
'The warmup ratio in the learning rate scheduler.',
|
||||
)
|
||||
|
||||
_MAX_STEPS = flags.DEFINE_integer(
|
||||
'max_steps',
|
||||
10,
|
||||
'Total number of training steps.',
|
||||
)
|
||||
|
||||
_MAX_SEQ_LENGTH = flags.DEFINE_integer(
|
||||
'max_seq_length',
|
||||
512,
|
||||
'The maximum sequence length.',
|
||||
)
|
||||
|
||||
_NUM_EPOCHS = flags.DEFINE_integer(
|
||||
'num_epochs',
|
||||
20,
|
||||
'The number of training epochs.',
|
||||
)
|
||||
|
||||
_BATCH_SIZE = flags.DEFINE_integer(
|
||||
'batch_size',
|
||||
32,
|
||||
'The batch size.',
|
||||
)
|
||||
|
||||
_LEARNING_RATE = flags.DEFINE_float(
|
||||
'learning_rate',
|
||||
2e-4,
|
||||
'The learning rate after the potential warmup period.',
|
||||
)
|
||||
|
||||
|
||||
def main(_) -> None:
|
||||
task = _TASK.value
|
||||
pretrained_model_id = _PRETRAINED_MODEL_ID.value
|
||||
local_pretrained_model_id = None
|
||||
if pretrained_model_id.startswith(constants.GCS_URI_PREFIX):
|
||||
logging.info(
|
||||
'Start to copy pretrained models locally: %s.', pretrained_model_id
|
||||
)
|
||||
fileutils.download_gcs_dir_to_local(
|
||||
pretrained_model_id, constants.LOCAL_BASE_MODEL_DIR
|
||||
)
|
||||
local_pretrained_model_id = constants.LOCAL_BASE_MODEL_DIR
|
||||
logging.info(
|
||||
'Finished copying pretrained models locally to: %s.',
|
||||
local_pretrained_model_id,
|
||||
)
|
||||
if task == constants.TEXT_TO_IMAGE_LORA:
|
||||
subprocess.run(['/bin/bash', 'train.sh'], check=True)
|
||||
elif task == constants.SEQUENCE_CLASSIFICATION_LORA:
|
||||
sequence_classification_lora.finetune_sequence_classification(
|
||||
pretrained_model_id=pretrained_model_id,
|
||||
dataset_name=_DATASET_NAME.value,
|
||||
output_dir=_OUTPUT_DIR.value,
|
||||
lora_rank=_LORA_RANK.value,
|
||||
lora_alpha=_LORA_ALPHA.value,
|
||||
lora_dropout=_LORA_DROPOUT.value,
|
||||
num_epochs=_NUM_EPOCHS.value,
|
||||
batch_size=_BATCH_SIZE.value,
|
||||
learning_rate=_LEARNING_RATE.value,
|
||||
)
|
||||
elif task == constants.CAUSAL_LANGUAGE_MODELING_LORA:
|
||||
causal_language_modeling_lora.finetune_causal_language_modeling(
|
||||
pretrained_model_id=pretrained_model_id,
|
||||
dataset_name=_DATASET_NAME.value,
|
||||
output_dir=_OUTPUT_DIR.value,
|
||||
precision_mode=_PRECISION_MODE.value,
|
||||
lora_rank=_LORA_RANK.value,
|
||||
lora_alpha=_LORA_ALPHA.value,
|
||||
lora_dropout=_LORA_DROPOUT.value,
|
||||
warmup_steps=_WARMUP_STEPS.value,
|
||||
max_steps=_MAX_STEPS.value,
|
||||
learning_rate=_LEARNING_RATE.value,
|
||||
local_pretrained_model_id=local_pretrained_model_id,
|
||||
)
|
||||
elif task == constants.INSTRUCT_LORA:
|
||||
instruct_lora.finetune_instruct(
|
||||
pretrained_model_id=pretrained_model_id,
|
||||
dataset_name=_DATASET_NAME.value,
|
||||
output_dir=_OUTPUT_DIR.value,
|
||||
lora_rank=_LORA_RANK.value,
|
||||
lora_alpha=_LORA_ALPHA.value,
|
||||
lora_dropout=_LORA_DROPOUT.value,
|
||||
warmup_ratio=_WARMUP_RATIO.value,
|
||||
max_steps=_MAX_STEPS.value,
|
||||
max_seq_length=_MAX_SEQ_LENGTH.value,
|
||||
learning_rate=_LEARNING_RATE.value,
|
||||
)
|
||||
else:
|
||||
raise ValueError('The task {} is not supported.'.format(task))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(main)
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Sequence classification with LoRA models."""
|
||||
|
||||
# pylint: disable=g-importing-member
|
||||
|
||||
from datasets import load_dataset
|
||||
import evaluate
|
||||
from peft import get_peft_model
|
||||
from peft import LoraConfig
|
||||
import torch
|
||||
from torch.optim import AdamW
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
from transformers import AutoTokenizer
|
||||
from transformers import get_linear_schedule_with_warmup
|
||||
|
||||
|
||||
def finetune_sequence_classification(
|
||||
pretrained_model_id: str,
|
||||
dataset_name: str,
|
||||
output_dir: str,
|
||||
lora_rank: int = 8,
|
||||
lora_alpha: int = 16,
|
||||
lora_dropout: float = 0.1,
|
||||
num_epochs: int = 20,
|
||||
batch_size: int = 32,
|
||||
learning_rate: float = 3e-4,
|
||||
) -> None:
|
||||
"""Finetunes sequence classification."""
|
||||
task = "mrpc"
|
||||
device = "cuda"
|
||||
|
||||
peft_config = LoraConfig(
|
||||
task_type="SEQ_CLS",
|
||||
inference_mode=False,
|
||||
r=lora_rank,
|
||||
lora_alpha=lora_alpha,
|
||||
lora_dropout=lora_dropout,
|
||||
)
|
||||
if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
|
||||
padding_side = "left"
|
||||
else:
|
||||
padding_side = "right"
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
pretrained_model_id, padding_side=padding_side
|
||||
)
|
||||
if getattr(tokenizer, "pad_token_id") is None:
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
|
||||
datasets = load_dataset(dataset_name, task)
|
||||
metric = evaluate.load(dataset_name, task)
|
||||
|
||||
def tokenize_function(examples):
|
||||
# max_length=None => use the model max length (it's actually the default)
|
||||
outputs = tokenizer(
|
||||
examples["sentence1"],
|
||||
examples["sentence2"],
|
||||
truncation=True,
|
||||
max_length=None,
|
||||
)
|
||||
return outputs
|
||||
|
||||
tokenized_datasets = datasets.map(
|
||||
tokenize_function,
|
||||
batched=True,
|
||||
remove_columns=["idx", "sentence1", "sentence2"],
|
||||
)
|
||||
|
||||
# We also rename the 'label' column to 'labels' which is the expected name for
|
||||
# labels by the models of the transformers library.
|
||||
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
|
||||
|
||||
def collate_fn(examples):
|
||||
return tokenizer.pad(examples, padding="longest", return_tensors="pt")
|
||||
|
||||
# Instantiate dataloaders.
|
||||
train_dataloader = DataLoader(
|
||||
tokenized_datasets["train"],
|
||||
shuffle=True,
|
||||
collate_fn=collate_fn,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
eval_dataloader = DataLoader(
|
||||
tokenized_datasets["validation"],
|
||||
shuffle=False,
|
||||
collate_fn=collate_fn,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
|
||||
model = AutoModelForSequenceClassification.from_pretrained(
|
||||
pretrained_model_id, return_dict=True
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
optimizer = AdamW(params=model.parameters(), lr=learning_rate)
|
||||
|
||||
# Instantiate scheduler
|
||||
lr_scheduler = get_linear_schedule_with_warmup(
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
|
||||
num_training_steps=(len(train_dataloader) * num_epochs),
|
||||
)
|
||||
|
||||
model.to(device)
|
||||
for epoch in range(num_epochs):
|
||||
model.train()
|
||||
for _, batch in enumerate(tqdm(train_dataloader)):
|
||||
batch.to(device)
|
||||
outputs = model(**batch)
|
||||
loss = outputs.loss
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
model.eval()
|
||||
for _, batch in enumerate(tqdm(eval_dataloader)):
|
||||
batch.to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**batch)
|
||||
predictions = outputs.logits.argmax(dim=-1)
|
||||
references = batch["labels"]
|
||||
metric.add_batch(
|
||||
predictions=predictions,
|
||||
references=references,
|
||||
)
|
||||
|
||||
eval_metric = metric.compute()
|
||||
print(f"epoch {epoch}:", eval_metric)
|
||||
|
||||
model.save_pretrained(output_dir)
|
||||
@@ -0,0 +1,6 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Setup accelerate config before running trainer.
|
||||
python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
|
||||
|
||||
accelerate launch "$@"
|
||||
-16
@@ -1,16 +0,0 @@
|
||||
# Dockerfile for axolotl training.
|
||||
#
|
||||
# To build:
|
||||
# docker build -f model_oss/peft/train/axolotol/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
|
||||
#
|
||||
# To push to gcr:
|
||||
# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
|
||||
|
||||
FROM winglian/axolotl:main-latest
|
||||
|
||||
RUN mkdir -p ./vertex_vision_model_garden/
|
||||
|
||||
COPY model_oss/peft/train/axolotl/*.py ./vertex_vision_model_garden/
|
||||
|
||||
ENTRYPOINT ["python3", "./vertex_vision_model_garden/train_entrypoint.py"]
|
||||
-20
@@ -1,20 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Run copybara first:
|
||||
# cloud/ml/applications/vision/model_garden/copybara/run_copybara_local.sh
|
||||
# Run docker build:
|
||||
# cloud/ml/applications/vision/model_garden/model_oss/peft/train/axolotl/scripts/build_train_docker.sh
|
||||
|
||||
set -x
|
||||
|
||||
COPYBARA_DIR="/tmp/train_docker/"
|
||||
|
||||
pushd "${COPYBARA_DIR}"
|
||||
|
||||
PROJECT="cloud-nas-260507"
|
||||
IMAGE_TAG="gcr.io/${PROJECT}/axolotl-train:${USER}-test"
|
||||
|
||||
docker build -f model_oss/peft/train/axolotl/dockerfile/train.Dockerfile . -t "${IMAGE_TAG}"
|
||||
docker push "${IMAGE_TAG}"
|
||||
|
||||
popd
|
||||
-88
@@ -1,88 +0,0 @@
|
||||
"""Entrypoint for axolotl train docker."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
|
||||
def _get_multi_node_flags(cluster_spec: str) -> list[str]:
|
||||
"""Returns the multi-node flags."""
|
||||
print(f'CLUSTER_SPEC: {cluster_spec}')
|
||||
|
||||
cluster_data = json.loads(cluster_spec)
|
||||
|
||||
# Get primary node info
|
||||
primary_node = cluster_data['cluster']['workerpool0'][0]
|
||||
print(f'primary node: {primary_node}')
|
||||
primary_node_addr, primary_node_port = primary_node.split(':')
|
||||
print(f'primary node address: {primary_node_addr}')
|
||||
print(f'primary node port: {primary_node_port}')
|
||||
|
||||
# Determine node rank of this machine
|
||||
workerpool = cluster_data['task']['type']
|
||||
if workerpool == 'workerpool0':
|
||||
node_rank = 0
|
||||
else:
|
||||
node_rank = cluster_data['task']['index'] + 1
|
||||
print(f'node rank: {node_rank}')
|
||||
|
||||
# Calculate total nodes
|
||||
num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
|
||||
num_nodes = num_worker_nodes + 1 # Add 1 for the primary node
|
||||
print(f'num nodes: {num_nodes}')
|
||||
|
||||
return [
|
||||
f'--machine_rank={node_rank}',
|
||||
f'--num_machines={num_nodes}',
|
||||
f'--main_process_ip={primary_node_addr}',
|
||||
f'--main_process_port={primary_node_port}',
|
||||
'--max_restarts=0',
|
||||
'--monitor_interval=120',
|
||||
'--dynamo_backend=no',
|
||||
]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--config_file')
|
||||
parser.add_argument('--huggingface_access_token')
|
||||
args, unknown = parser.parse_known_args()
|
||||
|
||||
accelerate_flags = []
|
||||
|
||||
if args.config_file:
|
||||
accelerate_flags.append(f'--config_file={args.config_file}')
|
||||
|
||||
if cluster_spec := os.getenv('CLUSTER_SPEC', default=None):
|
||||
print('========== Launch on cloud multi nodes ==========')
|
||||
accelerate_flags.extend(_get_multi_node_flags(cluster_spec))
|
||||
|
||||
cmd = (
|
||||
[
|
||||
'accelerate',
|
||||
'launch',
|
||||
]
|
||||
+ accelerate_flags
|
||||
+ [
|
||||
'-m',
|
||||
'axolotl.cli.train',
|
||||
]
|
||||
+ unknown
|
||||
)
|
||||
print(f'{cmd=}', flush=True)
|
||||
|
||||
env = os.environ.copy()
|
||||
|
||||
if args.huggingface_access_token:
|
||||
env['HF_TOKEN'] = args.huggingface_access_token
|
||||
|
||||
subprocess.run(
|
||||
cmd,
|
||||
check=True,
|
||||
env=env,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
-95
@@ -1,95 +0,0 @@
|
||||
"""Class that bundles docker related flags."""
|
||||
|
||||
import getpass
|
||||
import os
|
||||
import pwd
|
||||
|
||||
|
||||
class CommandBuilder:
|
||||
"""Base class for building commands."""
|
||||
|
||||
def __init__(self):
|
||||
self._defaults = []
|
||||
self._env_vars = {}
|
||||
|
||||
def add_env_var(self, var: str, val: str) -> None:
|
||||
"""Add environment variable to the command.
|
||||
|
||||
Args:
|
||||
var: environment variable name.
|
||||
val: environment variable value.
|
||||
"""
|
||||
self._env_vars[var] = val
|
||||
|
||||
def add_mount_map(self, host_path, docker_path):
|
||||
pass
|
||||
|
||||
|
||||
class DockerCommandBuilder(CommandBuilder):
|
||||
"""Bundle docker related flags."""
|
||||
|
||||
def __init__(self, docker_uri: str, shm_size: str = '128gb'):
|
||||
super().__init__()
|
||||
self._docker_uri = [docker_uri]
|
||||
self.privilege_mode = []
|
||||
self.entrypoint = []
|
||||
|
||||
self._defaults = [
|
||||
'docker',
|
||||
'run',
|
||||
'--gpus=all',
|
||||
'--net=host',
|
||||
'--rm',
|
||||
f'--shm-size={shm_size}',
|
||||
]
|
||||
|
||||
self._mount_maps = []
|
||||
user = getpass.getuser()
|
||||
# username ends with `_google_com` is managed by ldap and does not have a
|
||||
# corresponding entry in /etc/passwd or /etc/group file. We cannot enable
|
||||
# non-root docker user with below method.
|
||||
if not user.endswith('_google_com'):
|
||||
uid = os.getuid()
|
||||
gid = pwd.getpwuid(uid).pw_gid
|
||||
self._defaults += [
|
||||
f'--user={uid}:{gid}',
|
||||
'--volume=/etc/group:/etc/group:ro',
|
||||
'--volume=/etc/passwd:/etc/passwd:ro',
|
||||
]
|
||||
|
||||
def add_mount_map(self, host_path, docker_path):
|
||||
self._mount_maps.append(f'--volume={host_path}:{docker_path}')
|
||||
|
||||
def add_privilege_mode(self):
|
||||
self.privilege_mode = ['--privileged']
|
||||
|
||||
def add_entrypoint(self, entrypoint: list[str]):
|
||||
self.entrypoint = entrypoint
|
||||
|
||||
def build_cmd(self) -> str:
|
||||
return (
|
||||
self._defaults
|
||||
+ [f'--env={var}={val}' for var, val in self._env_vars.items()]
|
||||
+ self._mount_maps
|
||||
+ self.privilege_mode
|
||||
+ self._docker_uri
|
||||
+ self.entrypoint
|
||||
)
|
||||
|
||||
|
||||
class PythonCommandBuilder(CommandBuilder):
|
||||
"""Bundle Python test command related flags."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._defaults = [
|
||||
'python3',
|
||||
'./vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py',
|
||||
]
|
||||
|
||||
def build_cmd(self) -> str:
|
||||
os.environ.update(self._env_vars)
|
||||
return self._defaults
|
||||
|
||||
def add_entrypoint(self, entrypoint: list[str]):
|
||||
self._defaults = entrypoint
|
||||
@@ -1,471 +0,0 @@
|
||||
"""Test util class."""
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import datetime
|
||||
import inspect
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from absl import flags
|
||||
from absl import logging
|
||||
from absl.testing import parameterized
|
||||
import command_builder
|
||||
import immutabledict
|
||||
import torch
|
||||
|
||||
_DOCKER_URI = flags.DEFINE_string('docker_uri', None, 'docker image uri')
|
||||
|
||||
_DRY_RUN = flags.DEFINE_bool('dry_run', False, 'dry-run the commands')
|
||||
|
||||
_LOCAL_INPUT_DIR = flags.DEFINE_string(
|
||||
'local_input_dir',
|
||||
os.path.expanduser('~/test_input'),
|
||||
'local directory for storing input data.',
|
||||
)
|
||||
|
||||
_LOCAL_OUTPUT_DIR = flags.DEFINE_string(
|
||||
'local_output_dir',
|
||||
'/tmp',
|
||||
'local directory for storing test output.',
|
||||
)
|
||||
|
||||
|
||||
_GCS_INPUT_DIR = flags.DEFINE_string(
|
||||
'gcs_input_dir',
|
||||
'gs://vmg-tuning-docker-test',
|
||||
'GCS directory that stores model checkpoint, dataset and etc.',
|
||||
)
|
||||
|
||||
_GCS_OUTPUT_DIR = flags.DEFINE_string(
|
||||
'gcs_output_dir',
|
||||
'gs://vmg-tuning-docker-test/output',
|
||||
'GCS directory that stores test output.',
|
||||
)
|
||||
|
||||
_GCS_TESTDATA_DIR = 'peft-train-image-test'
|
||||
|
||||
_THROUGHPUT_TEST_EXCEPTIONS = immutabledict.immutabledict({
|
||||
('bm_deepspeed_zero3_8gpu_gemma-2-9b-it_4bit.txt', '12.0'): float('inf'),
|
||||
('bm_fsdp_8gpu_llama3.1-70b-hf_4bit.txt', '20.0'): float('inf'),
|
||||
('bm_deepspeed_zero2_8gpu_gemma-2-2b-it_bfloat16.txt', '12.0'): 20.0,
|
||||
('bm_deepspeed_zero3_8gpu_gemma-2-2b-it_4bit.txt', '4.0'): 20.0,
|
||||
('bm_deepspeed_zero3_8gpu_gemma-2-27b-it_4bit.txt', '4.0'): 20.0,
|
||||
})
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class BenchmarkStats:
|
||||
"""Class to store the benchmark result.
|
||||
|
||||
Attributes:
|
||||
peak_mem: peak memory in GB.
|
||||
throughput: throughput in tokens/sec.
|
||||
"""
|
||||
|
||||
peak_mem: float
|
||||
throughput: float
|
||||
|
||||
|
||||
class TestBase(parameterized.TestCase):
|
||||
"""Test base class that defines how to run commands."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
|
||||
# Create a copy of the environment variables
|
||||
self.old_env_var = copy.deepcopy(os.environ)
|
||||
if _DOCKER_URI.value:
|
||||
self.command_builder = command_builder.DockerCommandBuilder(
|
||||
_DOCKER_URI.value
|
||||
)
|
||||
else:
|
||||
self.command_builder = command_builder.PythonCommandBuilder()
|
||||
|
||||
self.command_builder.add_mount_map(
|
||||
os.path.expanduser('~'), os.path.expanduser('~')
|
||||
)
|
||||
self.command_builder.add_mount_map(
|
||||
self.local_input_dir(), self.local_input_dir()
|
||||
)
|
||||
|
||||
self.task_cmd_builder = None
|
||||
|
||||
def tearDown(self):
|
||||
super().tearDown()
|
||||
# Restore the original environment variables
|
||||
os.environ.clear()
|
||||
os.environ.update(self.old_env_var)
|
||||
|
||||
def cmd(self):
|
||||
return self.command_builder.build_cmd() + self.task_cmd_builder.build_cmd()
|
||||
|
||||
def run_cmd(self) -> int:
|
||||
return run_cmd(self.cmd(), output_file=None)
|
||||
|
||||
def gcs_output_dir(self):
|
||||
return _GCS_OUTPUT_DIR.value
|
||||
|
||||
def local_input_dir(self):
|
||||
"""Returns local input dir in host/docker."""
|
||||
return _LOCAL_INPUT_DIR.value
|
||||
|
||||
def local_output_dir(self):
|
||||
"""Returns local output dir in host/docker."""
|
||||
return _LOCAL_OUTPUT_DIR.value
|
||||
|
||||
def get_testcase_name(self):
|
||||
"""Returns the function name at the calling site."""
|
||||
# https://docs.python.org/3/library/inspect.html#inspect.FrameInfo
|
||||
cur_frame = inspect.currentframe()
|
||||
# https://stackoverflow.com/a/17366561
|
||||
return cur_frame.f_back.f_code.co_name
|
||||
|
||||
|
||||
def get_timestamp():
|
||||
return datetime.datetime.now(datetime.timezone.utc).strftime(
|
||||
'%Y%m%d_%H%M%S%Z'
|
||||
)
|
||||
|
||||
|
||||
def download_from_gcs(gcs_uri: str, local_dir: str):
|
||||
if not os.path.exists(local_dir):
|
||||
os.mkdir(local_dir)
|
||||
subprocess.check_output([
|
||||
'gcloud',
|
||||
'storage',
|
||||
'cp',
|
||||
'-r',
|
||||
gcs_uri,
|
||||
local_dir,
|
||||
])
|
||||
|
||||
|
||||
def get_test_data_path(name: str, download: bool = True) -> str:
|
||||
"""Gets test data path.
|
||||
|
||||
Args:
|
||||
name: name of the test data
|
||||
download: if True, then download data from GCS and returns its local path.
|
||||
|
||||
Returns:
|
||||
test data path.
|
||||
"""
|
||||
if not download:
|
||||
return os.path.join(_GCS_INPUT_DIR.value, name)
|
||||
|
||||
local_data = os.path.join(_LOCAL_INPUT_DIR.value, name)
|
||||
if not os.path.exists(local_data):
|
||||
# If `name` is a file in sub-folders, then create the sub-folders under
|
||||
# `_LOCAL_INPUT_DIR`.
|
||||
local_data_dir = os.path.dirname(local_data)
|
||||
if not os.path.exists(local_data_dir):
|
||||
os.makedirs(local_data_dir)
|
||||
|
||||
download_from_gcs(os.path.join(_GCS_INPUT_DIR.value, name), local_data_dir)
|
||||
|
||||
return local_data
|
||||
|
||||
|
||||
def run_cmd(cmd: list[str], output_file: str = None) -> int:
|
||||
"""Runs the command and returns the return code.
|
||||
|
||||
Args:
|
||||
cmd: The command to run.
|
||||
output_file: The file to write the output to.
|
||||
|
||||
Returns:
|
||||
The return code of the command.
|
||||
"""
|
||||
logging.info('running command: \n%s', ' \\\n'.join(cmd))
|
||||
if _DRY_RUN.value:
|
||||
return 0
|
||||
stdout = sys.stdout if output_file is None else open(output_file, 'w')
|
||||
p = subprocess.Popen(cmd, stdout=stdout, stderr=sys.stderr)
|
||||
try:
|
||||
unused_output, unused_error = p.communicate()
|
||||
return_code = p.returncode
|
||||
except KeyboardInterrupt:
|
||||
p.send_signal(signal.SIGINT)
|
||||
return_code = 0
|
||||
finally:
|
||||
if output_file is not None:
|
||||
stdout.close()
|
||||
return return_code
|
||||
|
||||
|
||||
def get_pretrained_model_name_or_path(model_id: str) -> str:
|
||||
# If `model_id` contains `/`, it is assumed to be HF model or model from GCS.
|
||||
if '/' in model_id:
|
||||
return model_id
|
||||
|
||||
return get_test_data_path(model_id, download=True)
|
||||
|
||||
|
||||
def is_gpu_h100():
|
||||
"""Checks if the GPU is H100."""
|
||||
return 'H100' in torch.cuda.get_device_name()
|
||||
|
||||
|
||||
def is_gpu_a100():
|
||||
"""Checks if the GPU is A100."""
|
||||
return 'A100' in torch.cuda.get_device_name()
|
||||
|
||||
|
||||
def _get_formatted_string(max_seq_length: int) -> str:
|
||||
"""Returns the formatted string for max_seq_length.
|
||||
|
||||
Args:
|
||||
max_seq_length: max sequence length to get the formatted string.
|
||||
|
||||
Returns:
|
||||
formatted string for max_seq_length.
|
||||
"""
|
||||
return f'{max_seq_length/1024.0:.1f}'
|
||||
|
||||
|
||||
def get_benchmark_results(
|
||||
benchmark_file_path: str, max_seq_length: int
|
||||
) -> BenchmarkStats:
|
||||
"""Gets benchmark results from the benchmark file.
|
||||
|
||||
Args:
|
||||
benchmark_file_path: path to the benchmark file.
|
||||
max_seq_length: max sequence length to get the benchmark results.
|
||||
|
||||
Returns:
|
||||
peak_mem: peak memory in GB.
|
||||
throughput: throughput in tokens/sec.
|
||||
"""
|
||||
formatted_max_seq_length = _get_formatted_string(max_seq_length)
|
||||
peak_mem, throughput = None, None
|
||||
with open(benchmark_file_path, 'r') as f:
|
||||
for line in f:
|
||||
if line.startswith(formatted_max_seq_length):
|
||||
metrics = line.split('|')
|
||||
try:
|
||||
peak_mem = float(metrics[1].strip())
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
throughput = float(metrics[2].strip())
|
||||
except ValueError:
|
||||
pass
|
||||
break
|
||||
else:
|
||||
logging.error(
|
||||
'No metrics found for max_seq_length %s in %s',
|
||||
formatted_max_seq_length,
|
||||
benchmark_file_path,
|
||||
)
|
||||
return BenchmarkStats(peak_mem, throughput)
|
||||
|
||||
|
||||
def print_benchmark_file(file_path: str) -> None:
|
||||
"""Prints the contents of the file.
|
||||
|
||||
Args:
|
||||
file_path: path to the file.
|
||||
"""
|
||||
with open(file_path, 'r') as f:
|
||||
for line in f:
|
||||
logging.info(line.strip())
|
||||
|
||||
|
||||
def print_benchmark_results(
|
||||
benchmark_file_path: str, benchmark_type: str
|
||||
) -> None:
|
||||
"""Prints the benchmark results.
|
||||
|
||||
Args:
|
||||
benchmark_file_path: path to the benchmark file.
|
||||
benchmark_type: type of the benchmark.
|
||||
"""
|
||||
benchmark_filename = os.path.basename(benchmark_file_path)
|
||||
logging.info('--------------------------------------------------------------')
|
||||
logging.info('%s benchmark for %s', benchmark_type, benchmark_filename)
|
||||
logging.info('--------------------------------------------------------------')
|
||||
print_benchmark_file(benchmark_file_path)
|
||||
|
||||
|
||||
def _calculate_percent_change(
|
||||
actual_value: float, expected_value: float
|
||||
) -> float:
|
||||
"""Calculates the percent change between the actual and expected values.
|
||||
|
||||
Args:
|
||||
actual_value: actual value to compare.
|
||||
expected_value: expected value to compare.
|
||||
|
||||
Returns:
|
||||
percent change between the actual and expected values.
|
||||
"""
|
||||
return ((actual_value - expected_value) / expected_value) * 100.0
|
||||
|
||||
|
||||
def compare_benchmark_results(
|
||||
expected_benchmark_file_path: str,
|
||||
actual_benchmark_file_path: str,
|
||||
allowed_threshold: float,
|
||||
max_seq_length: int,
|
||||
) -> bool:
|
||||
"""Compares if the benchmark results are the similar.
|
||||
|
||||
Args:
|
||||
expected_benchmark_file_path: path to the expected benchmark file.
|
||||
actual_benchmark_file_path: path to the actual benchmark file.
|
||||
allowed_threshold: allowed percent range of the benchmark results.
|
||||
max_seq_length: max sequence length to get the benchmark results.
|
||||
|
||||
Returns:
|
||||
True if the benchmark results are the similar, False otherwise.
|
||||
"""
|
||||
benchmark_filename = os.path.basename(expected_benchmark_file_path)
|
||||
expected_results = get_benchmark_results(
|
||||
expected_benchmark_file_path, max_seq_length
|
||||
)
|
||||
expected_peak_mem, expected_throughput = (
|
||||
expected_results.peak_mem,
|
||||
expected_results.throughput,
|
||||
)
|
||||
actual_results = get_benchmark_results(
|
||||
actual_benchmark_file_path, max_seq_length
|
||||
)
|
||||
actual_peak_mem, actual_throughput = (
|
||||
actual_results.peak_mem,
|
||||
actual_results.throughput,
|
||||
)
|
||||
formatted_max_seq_length = _get_formatted_string(max_seq_length)
|
||||
|
||||
# Case 1: both peak mem and throughput are None(ideally due to OOM)
|
||||
if expected_peak_mem is None and actual_peak_mem is None:
|
||||
logging.info(
|
||||
'Both peak mem and throughput are None for max_seq_length %d.',
|
||||
max_seq_length,
|
||||
)
|
||||
return True
|
||||
|
||||
check_oom_exception = _THROUGHPUT_TEST_EXCEPTIONS.get(
|
||||
(benchmark_filename, formatted_max_seq_length), 0.0
|
||||
) == float('inf')
|
||||
# Case 2: When something strated to fail recently, or something which failed
|
||||
# before but is working now.
|
||||
if expected_peak_mem is None and actual_peak_mem is not None:
|
||||
if check_oom_exception:
|
||||
return True
|
||||
logging.error(
|
||||
'One of the failing benchmarks in %s is passing now for max_seq_length'
|
||||
' %d. The expected peak mem and throughput are None, but the actual'
|
||||
' peak mem is %f and actual throughput is %f',
|
||||
benchmark_filename,
|
||||
max_seq_length,
|
||||
actual_peak_mem,
|
||||
actual_throughput,
|
||||
)
|
||||
return False
|
||||
if actual_peak_mem is None and expected_peak_mem is not None:
|
||||
if check_oom_exception:
|
||||
return True
|
||||
logging.error(
|
||||
'One of the passing benchmarks in %s is failing now for max_seq_length'
|
||||
' %d. The actual peak mem and throughput are None, but the expected'
|
||||
' peak mem is %f and expected throughput is %f',
|
||||
benchmark_filename,
|
||||
max_seq_length,
|
||||
expected_peak_mem,
|
||||
expected_throughput,
|
||||
)
|
||||
return False
|
||||
# Case 3: When both actual peak mem and throughput lies within the range
|
||||
# of their respective expected values.
|
||||
mem_percent_change = _calculate_percent_change(
|
||||
actual_peak_mem, expected_peak_mem
|
||||
)
|
||||
throughput_percent_change = _calculate_percent_change(
|
||||
actual_throughput, expected_throughput
|
||||
)
|
||||
allowed_threshold = _THROUGHPUT_TEST_EXCEPTIONS.get(
|
||||
(benchmark_filename, formatted_max_seq_length), allowed_threshold
|
||||
)
|
||||
|
||||
if abs(mem_percent_change) > allowed_threshold:
|
||||
logging.error(
|
||||
'The peak memory is changing by more than %f%% for max_seq_length %d.'
|
||||
' Expected: %f, Actual: %f',
|
||||
allowed_threshold,
|
||||
max_seq_length,
|
||||
expected_peak_mem,
|
||||
actual_peak_mem,
|
||||
)
|
||||
return False
|
||||
if abs(throughput_percent_change) > allowed_threshold:
|
||||
logging.error(
|
||||
'The throughput is changing by more than %f%% for max_seq_length %d.'
|
||||
' Expected throughput: %f, Actual throughput: %f',
|
||||
allowed_threshold,
|
||||
max_seq_length,
|
||||
expected_throughput,
|
||||
actual_throughput,
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def check_benchmark_results(
|
||||
actual_benchmark_file_path: str,
|
||||
model_family: str,
|
||||
allowed_threshold: float,
|
||||
max_seq_length: int,
|
||||
) -> bool:
|
||||
"""Checks the benchmark result between the actual and expected benchmark files.
|
||||
|
||||
Args:
|
||||
actual_benchmark_file_path: path to the actual benchmark file.
|
||||
model_family: family of the model.
|
||||
allowed_threshold: allowed range of the benchmark results in percent.
|
||||
max_seq_length: max sequence length to get the benchmark results.
|
||||
|
||||
Returns:
|
||||
True if the benchmark results are the similar, False otherwise.
|
||||
"""
|
||||
benchmark_filename = os.path.basename(actual_benchmark_file_path)
|
||||
get_test_data_path(_GCS_TESTDATA_DIR)
|
||||
expected_benchmark_file_path = os.path.join(
|
||||
_LOCAL_INPUT_DIR.value,
|
||||
_GCS_TESTDATA_DIR,
|
||||
model_family,
|
||||
benchmark_filename,
|
||||
)
|
||||
print_benchmark_results(expected_benchmark_file_path, 'Expected')
|
||||
print_benchmark_results(actual_benchmark_file_path, 'Actual')
|
||||
|
||||
return compare_benchmark_results(
|
||||
expected_benchmark_file_path,
|
||||
actual_benchmark_file_path,
|
||||
allowed_threshold,
|
||||
max_seq_length,
|
||||
)
|
||||
|
||||
|
||||
def list_gcs_directories(bucket: str, directory: str) -> list[str]:
|
||||
"""Lists GCS files."""
|
||||
output = subprocess.check_output([
|
||||
'gcloud',
|
||||
'storage',
|
||||
'ls',
|
||||
f'gs://{bucket}/{directory}',
|
||||
])
|
||||
return output.decode('utf-8').splitlines()
|
||||
|
||||
|
||||
def delete_gcs_object(gcs_directory: str):
|
||||
"""Deletes GCS object."""
|
||||
subprocess.check_output([
|
||||
'gcloud',
|
||||
'storage',
|
||||
'rm',
|
||||
'-r',
|
||||
f'{gcs_directory}',
|
||||
])
|
||||
@@ -1,79 +0,0 @@
|
||||
"""Get cluster info from environment variables."""
|
||||
|
||||
import dataclasses
|
||||
import json
|
||||
import os
|
||||
|
||||
from absl import logging
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ClusterInfo:
|
||||
"""Contains information about the cluster.
|
||||
|
||||
Attributes:
|
||||
primary_node_addr: The address of the primary node.
|
||||
primary_node_port: The port of the primary node.
|
||||
node_rank: The rank of the node.
|
||||
num_nodes: The number of nodes in the cluster.
|
||||
"""
|
||||
|
||||
primary_node_addr: str | None = None
|
||||
primary_node_port: str | None = None
|
||||
node_rank: int = 0
|
||||
num_nodes: int = 1
|
||||
|
||||
# Allows unpacking operation like
|
||||
# primary_node_addr, primary_node_port, _, _ = ClusterInfo()
|
||||
# See https://stackoverflow.com/a/70753113
|
||||
def __iter__(self):
|
||||
return iter(dataclasses.astuple(self))
|
||||
|
||||
|
||||
def get_cluster_spec() -> ClusterInfo:
|
||||
"""Parses CLUSTER_SPEC environment variable and returns the cluster info.
|
||||
|
||||
Returns:
|
||||
A ClusterInfo object.
|
||||
"""
|
||||
cluster_spec = os.getenv('CLUSTER_SPEC', None)
|
||||
|
||||
# If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
|
||||
if not cluster_spec:
|
||||
cluster_info = ClusterInfo(
|
||||
primary_node_addr=os.getenv('MASTER_ADDR', None),
|
||||
primary_node_port=os.getenv('MASTER_PORT', None),
|
||||
node_rank=int(os.getenv('RANK', '0')),
|
||||
num_nodes=int(os.getenv('NNODES', '1')),
|
||||
)
|
||||
return cluster_info
|
||||
|
||||
cluster_data = json.loads(cluster_spec)
|
||||
# Get primary node info
|
||||
primary_node = cluster_data['cluster']['workerpool0'][0]
|
||||
logging.info('primary node: %s', primary_node)
|
||||
primary_node_addr, primary_node_port = primary_node.split(':')
|
||||
logging.info('primary node address: %s', primary_node_addr)
|
||||
logging.info('primary node port: %s', primary_node_port)
|
||||
|
||||
# Determine node rank of this machine
|
||||
workerpool = cluster_data['task']['type']
|
||||
if workerpool == 'workerpool0':
|
||||
node_rank = 0
|
||||
elif workerpool == 'workerpool1':
|
||||
# Add 1 for the primary node, since `index` is the index of workerpool1.
|
||||
node_rank = cluster_data['task']['index'] + 1
|
||||
else:
|
||||
raise ValueError(
|
||||
'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
|
||||
f' {workerpool}'
|
||||
)
|
||||
logging.info('node rank: %s', node_rank)
|
||||
|
||||
# Calculate total nodes.
|
||||
num_nodes = 1 # For the primary node.
|
||||
if 'workerpool1' in cluster_data['cluster']:
|
||||
num_nodes += len(cluster_data['cluster']['workerpool1'])
|
||||
logging.info('num nodes: %s', num_nodes)
|
||||
|
||||
return ClusterInfo(primary_node_addr, primary_node_port, node_rank, num_nodes)
|
||||
@@ -1,24 +0,0 @@
|
||||
"""Utility functions."""
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def run_cmd(cmd: list[str]) -> float:
|
||||
"""Runs the command and logs the output.
|
||||
|
||||
Args:
|
||||
cmd: The command to run.
|
||||
|
||||
Returns:
|
||||
The time it took to run the command.
|
||||
"""
|
||||
cmd_str = ' \\\n'.join(cmd)
|
||||
logging.info('launching cmd: \n%s', cmd_str)
|
||||
start_time = time.time()
|
||||
subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
|
||||
elapsed_time = round(time.time() - start_time, 2)
|
||||
logging.info('Command %s finished in %0.2f seconds.', cmd_str, elapsed_time)
|
||||
return elapsed_time
|
||||
@@ -1,197 +0,0 @@
|
||||
"""Calculate dataset statistics like token, example and character counts."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
import dataclasses
|
||||
import json
|
||||
from typing import Any
|
||||
import datasets
|
||||
import numpy as np
|
||||
import transformers
|
||||
from util import dataset_validation_util
|
||||
|
||||
_MAX_NUM_DATASET_SAMPLES = 6
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SupervisedTuningDatasetBucket:
|
||||
"""Represents a histogram bucket for tuning dataset distribution stats."""
|
||||
|
||||
count: float = 0
|
||||
left: float = 0
|
||||
right: float = 0
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SupervisedTuningDatasetDistribution:
|
||||
"""Represents a histogram with summary statistics for tuning dataset distribution stats."""
|
||||
|
||||
sum: int = 0
|
||||
billable_sum: int = 0
|
||||
min: float = 0
|
||||
max: float = 0
|
||||
mean: float = 0
|
||||
median: float = 0
|
||||
p5: float = 0
|
||||
p95: float = 0
|
||||
buckets: list[SupervisedTuningDatasetBucket] = dataclasses.field(
|
||||
default_factory=list
|
||||
)
|
||||
|
||||
|
||||
# Represents detailed tuning dataset statistics.
|
||||
@dataclasses.dataclass
|
||||
class SupervisedTuningDataStats:
|
||||
"""Represents detailed tuning dataset stats."""
|
||||
|
||||
tuning_dataset_example_count: int = 0
|
||||
total_tuning_character_count: int = 0
|
||||
total_billable_token_count: int = 0
|
||||
tuning_step_count: int = 0
|
||||
# Represents a histogram and some summary statistics of the number of input
|
||||
# tokens across examples.
|
||||
user_input_token_distribution: SupervisedTuningDatasetDistribution | None = (
|
||||
None
|
||||
)
|
||||
# Represents a histogram and some summary statistics for the number of output
|
||||
# tokens across examples.
|
||||
user_output_token_distribution: SupervisedTuningDatasetDistribution | None = (
|
||||
None
|
||||
)
|
||||
# Represents the number of "messages" (a single-turn conversation will have a
|
||||
# single message) across examples.
|
||||
user_message_per_example_distribution: (
|
||||
SupervisedTuningDatasetDistribution | None
|
||||
) = None
|
||||
user_dataset_examples: list[str] = dataclasses.field(default_factory=list)
|
||||
|
||||
|
||||
def get_dataset_stats(
|
||||
*,
|
||||
raw: Any,
|
||||
templated: Any,
|
||||
template: str,
|
||||
tokenizer: transformers.PreTrainedTokenizer,
|
||||
column: str,
|
||||
effective_batch_size: int,
|
||||
) -> Mapping[str, Any]:
|
||||
"""Calculates dataset statistics for managed fine-tuning, e.g., total number of tokens."""
|
||||
tokenized_dataset = templated.map(lambda x: tokenizer(x[column]))
|
||||
inputs = tokenized_dataset["input_ids"]
|
||||
tuning_dataset_example_count = int(len(inputs))
|
||||
total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
|
||||
total_tuning_character_count = int(
|
||||
np.sum([len(ex[column]) for ex in templated])
|
||||
)
|
||||
tuning_step_count = (
|
||||
tuning_dataset_example_count + effective_batch_size - 1
|
||||
) // effective_batch_size
|
||||
|
||||
# Assume that data is represented as ChatCompletions or Vertex Text-Bison
|
||||
# formats to extract per-example input/output tokens.
|
||||
user_inputs = []
|
||||
user_outputs = []
|
||||
user_input_messages_counts = []
|
||||
|
||||
for ex in raw:
|
||||
if "messages" in ex:
|
||||
messages = ex["messages"]
|
||||
if messages:
|
||||
# For ChatCompletions assume the last turn (i.e. the instruction
|
||||
# response) is the expected output.
|
||||
user_inputs.append({**ex, "messages": messages[:-1]})
|
||||
user_outputs.append({**ex, "messages": messages[-1:]})
|
||||
# Exclude everything but the last message for the number of input
|
||||
# messages.
|
||||
user_input_messages_counts.append(len(messages[:-1]))
|
||||
elif "input_text" in ex:
|
||||
# For Vertex Text-Bison, the `output_text` field is the expected output.
|
||||
user_inputs.append({**ex, "output_text": ""})
|
||||
user_outputs.append(
|
||||
{**ex, "input_text": ex["output_text"], "output_text": ""}
|
||||
)
|
||||
# Vertex Text-Bison goes from input -> output; i.e. there is only a single
|
||||
# input "message".
|
||||
user_input_messages_counts.append(1)
|
||||
|
||||
def calc_histogram(
|
||||
counts: Sequence[int],
|
||||
) -> SupervisedTuningDatasetDistribution:
|
||||
mean = np.mean(counts)
|
||||
median = np.median(counts).item()
|
||||
max_count = np.max(counts).item()
|
||||
min_count = np.min(counts).item()
|
||||
count_sum = np.sum(counts).item()
|
||||
p5 = np.percentile(counts, 0.05).item()
|
||||
p95 = np.percentile(counts, 0.95).item()
|
||||
hist, bin_edges = np.histogram(counts, bins=10)
|
||||
|
||||
return SupervisedTuningDatasetDistribution(
|
||||
sum=count_sum,
|
||||
billable_sum=count_sum,
|
||||
min=min_count,
|
||||
max=max_count,
|
||||
mean=mean,
|
||||
median=median,
|
||||
p5=p5,
|
||||
p95=p95,
|
||||
buckets=[
|
||||
SupervisedTuningDatasetBucket(
|
||||
count=hist[i].item(),
|
||||
left=bin_edges[i].item(),
|
||||
right=bin_edges[i + 1].item(),
|
||||
)
|
||||
for i in range(len(hist))
|
||||
],
|
||||
)
|
||||
|
||||
# Tokenize input and output messages separately to generate separate summary
|
||||
# statistics about them.
|
||||
user_input_token_distribution = None
|
||||
if user_inputs:
|
||||
user_input_dataset = dataset_validation_util.format_dataset(
|
||||
datasets.Dataset.from_list(user_inputs), column, template, tokenizer
|
||||
)
|
||||
user_input_tokenized_dataset = user_input_dataset.map(
|
||||
lambda x: tokenizer(x[column])
|
||||
)
|
||||
user_input_tokens = user_input_tokenized_dataset["input_ids"]
|
||||
user_input_token_counts = np.array([len(ex) for ex in user_input_tokens])
|
||||
user_input_token_distribution = calc_histogram(user_input_token_counts)
|
||||
|
||||
user_output_token_distribution = None
|
||||
if user_outputs:
|
||||
user_output_dataset = dataset_validation_util.format_dataset(
|
||||
datasets.Dataset.from_list(user_outputs), column, template, tokenizer
|
||||
)
|
||||
user_output_tokenized_dataset = user_output_dataset.map(
|
||||
lambda x: tokenizer(x[column])
|
||||
)
|
||||
user_output_tokens = user_output_tokenized_dataset["input_ids"]
|
||||
user_output_token_counts = np.array([len(ex) for ex in user_output_tokens])
|
||||
user_output_token_distribution = calc_histogram(user_output_token_counts)
|
||||
|
||||
user_messages_per_example_distribution = None
|
||||
if user_input_messages_counts:
|
||||
user_input_messages_counts = np.array(user_input_messages_counts)
|
||||
user_messages_per_example_distribution = calc_histogram(
|
||||
user_input_messages_counts
|
||||
)
|
||||
|
||||
user_dataset_examples = [
|
||||
json.dumps(ex)
|
||||
for ex in raw.shuffle().select(
|
||||
range(min(len(raw), _MAX_NUM_DATASET_SAMPLES))
|
||||
)
|
||||
]
|
||||
|
||||
dataset_stats = SupervisedTuningDataStats(
|
||||
tuning_dataset_example_count=tuning_dataset_example_count,
|
||||
total_tuning_character_count=total_tuning_character_count,
|
||||
total_billable_token_count=total_billable_token_count,
|
||||
tuning_step_count=tuning_step_count,
|
||||
user_input_token_distribution=user_input_token_distribution,
|
||||
user_output_token_distribution=user_output_token_distribution,
|
||||
user_message_per_example_distribution=user_messages_per_example_distribution,
|
||||
user_dataset_examples=user_dataset_examples,
|
||||
)
|
||||
return dataclasses.asdict(dataset_stats)
|
||||
@@ -1,140 +0,0 @@
|
||||
"""Util functions for reporting device (GPU, CPU) stats."""
|
||||
|
||||
import dataclasses
|
||||
|
||||
import psutil
|
||||
import pynvml
|
||||
import torch
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class GpuStats:
|
||||
"""Holds information about GPU usage stats.
|
||||
|
||||
For memory related, see
|
||||
https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
|
||||
"""
|
||||
|
||||
# device id
|
||||
device_id: int
|
||||
# memory reserved.
|
||||
reserved: float
|
||||
# memory occupied.
|
||||
occupied: float
|
||||
# memory reserved, but not used.
|
||||
unused: float
|
||||
# nvidia-smi usually reports more memory usages than pytorch (for driver,
|
||||
# kernel and etc). `smi_diff` tracks this difference.
|
||||
smi_diff: float
|
||||
# Gpu utilization.
|
||||
util: float
|
||||
|
||||
# Allows unpacking operation like
|
||||
# device_id, reserved, occupied, unused, smi_diff, util = GpuStats(...)
|
||||
# See https://stackoverflow.com/a/70753113
|
||||
def __iter__(self):
|
||||
return iter(dataclasses.astuple(self))
|
||||
|
||||
|
||||
def gpu_stats() -> GpuStats:
|
||||
"""Reports GPU memory usage and utilization."""
|
||||
# See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
|
||||
bytes_per_gb = 1024.0**3
|
||||
device = torch.cuda.current_device()
|
||||
occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
|
||||
reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
|
||||
unused = reserved - occupied
|
||||
|
||||
def smi_mem(device):
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(device)
|
||||
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
return info.used / bytes_per_gb
|
||||
except pynvml.NVMLError:
|
||||
return 0.0
|
||||
|
||||
mem_used_smi = smi_mem(device)
|
||||
smi_diff = mem_used_smi - reserved
|
||||
|
||||
util = torch.cuda.utilization(device)
|
||||
return GpuStats(device, reserved, occupied, unused, smi_diff, util)
|
||||
|
||||
|
||||
def gpu_stats_str(stats: GpuStats | None = None) -> str:
|
||||
if stats is None:
|
||||
stats = gpu_stats()
|
||||
device, reserved, occupied, unused, smi_diff, util = stats
|
||||
return (
|
||||
f"GPU ({device=}) memory: {reserved:.2f}({occupied=:.2f}, {unused=:.2f}),"
|
||||
f" {smi_diff=:.2f} GB. Utilization: {util:.2f}%"
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class CpuStats:
|
||||
"""Holds information about CPU usage stats."""
|
||||
|
||||
# Total CPU virtual memory i.e. virtual memory allocated + unallocated.
|
||||
total_virtual_mem: float
|
||||
# CPU virtual memory available for use.
|
||||
unallocated_virtual_mem: float
|
||||
# CPU virtual memory already used.
|
||||
allocated_virtual_mem: float
|
||||
# Total CPU swap memory i.e. swap memory allocated + unallocated.
|
||||
total_swap_mem: float
|
||||
# CPU swap memory available for use.
|
||||
unallocated_swap_mem: float
|
||||
# CPU swap memory already used.
|
||||
allocated_swap_mem: float
|
||||
# CPU utilization percentage.
|
||||
utilization: float
|
||||
|
||||
|
||||
def cpu_stats() -> CpuStats:
|
||||
"""Reports CPU memory usage and utilization."""
|
||||
|
||||
# https://psutil.readthedocs.io/en/latest/#memory
|
||||
gb = 1024.0**3
|
||||
vmem = psutil.virtual_memory()
|
||||
vmem_total = vmem.total / gb
|
||||
vmem_available = vmem.available / gb
|
||||
vmem_used = vmem_total - vmem_available
|
||||
smem = psutil.swap_memory()
|
||||
swap_total = smem.total / gb
|
||||
swap_free = smem.free / gb
|
||||
swap_used = smem.used / gb
|
||||
# https://psutil.readthedocs.io/en/latest/#psutil.cpu_percent
|
||||
cpu_util = psutil.cpu_percent(interval=1e-6)
|
||||
return CpuStats(
|
||||
total_virtual_mem=vmem_total,
|
||||
unallocated_virtual_mem=vmem_available,
|
||||
allocated_virtual_mem=vmem_used,
|
||||
total_swap_mem=swap_total,
|
||||
unallocated_swap_mem=swap_free,
|
||||
allocated_swap_mem=swap_used,
|
||||
utilization=cpu_util,
|
||||
)
|
||||
|
||||
|
||||
def cpu_stats_str(stats: CpuStats | None = None) -> str:
|
||||
"""Returns a string representation of the CPU stats."""
|
||||
|
||||
if stats is None:
|
||||
stats = cpu_stats()
|
||||
total, occupied, unused = (
|
||||
stats.total_virtual_mem,
|
||||
stats.allocated_virtual_mem,
|
||||
stats.unallocated_virtual_mem,
|
||||
)
|
||||
virtual_mem = (
|
||||
f"CPU virtual memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
|
||||
)
|
||||
total, occupied, unused = (
|
||||
stats.total_swap_mem,
|
||||
stats.allocated_swap_mem,
|
||||
stats.unallocated_swap_mem,
|
||||
)
|
||||
swap_mem = f"CPU swap memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
|
||||
percent = stats.utilization
|
||||
return f"{virtual_mem} {swap_mem} CPU Utilization: {percent:.2f}%"
|
||||
@@ -1,126 +0,0 @@
|
||||
"""Different trainer callbacks for PEFT Trainer."""
|
||||
|
||||
from collections.abc import MutableMapping
|
||||
import math
|
||||
import time
|
||||
|
||||
from absl import logging
|
||||
import accelerate
|
||||
from transformers import TrainingArguments
|
||||
from transformers.trainer_callback import TrainerCallback
|
||||
from transformers.trainer_callback import TrainerControl
|
||||
from transformers.trainer_callback import TrainerState
|
||||
|
||||
from util import device_stats
|
||||
|
||||
|
||||
class TrainerStatsCallback(TrainerCallback):
|
||||
"""Trainer callback to report trainer stats."""
|
||||
|
||||
def __init__(self, max_seq_length, filename=None):
|
||||
self._max_seq_length = max_seq_length
|
||||
self._filename = filename
|
||||
|
||||
self._partial_state = accelerate.PartialState()
|
||||
self._start_time = float('nan')
|
||||
self._prev_time = float('nan')
|
||||
self._peak_mem = 0.0
|
||||
self._avg_throughput = 0.0
|
||||
|
||||
def on_log(
|
||||
self,
|
||||
args: TrainingArguments,
|
||||
state: TrainerState,
|
||||
control: TrainerControl,
|
||||
logs: MutableMapping[str, float] | None = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""Calculates perplexity from train loss.
|
||||
|
||||
Args:
|
||||
args: Arguments passed to the trainer.
|
||||
state: State of the trainer.
|
||||
control: Control of the trainer.
|
||||
logs: A dict of logs from the training loop.
|
||||
**kwargs: Additional keyword arguments, not used in this callback.
|
||||
"""
|
||||
del kwargs # Unused.
|
||||
if self._partial_state.is_main_process:
|
||||
train_loss = logs.get('loss') if logs is not None else None
|
||||
if train_loss is not None:
|
||||
perplexity = round(float(math.exp(train_loss)), 4)
|
||||
logs['perplexity'] = perplexity
|
||||
|
||||
def on_step_end(
|
||||
self,
|
||||
args: TrainingArguments,
|
||||
state: TrainerState,
|
||||
control: TrainerControl,
|
||||
**kwargs,
|
||||
):
|
||||
if self._partial_state.is_main_process:
|
||||
if state.global_step == 1:
|
||||
self._prev_time = time.time()
|
||||
self._prev_num_token = state.num_input_tokens_seen
|
||||
throughput = 0.0
|
||||
else:
|
||||
cur_time = time.time()
|
||||
cur_num_token = state.num_input_tokens_seen
|
||||
throughput = (cur_num_token - self._prev_num_token) / (
|
||||
cur_time - self._prev_time
|
||||
)
|
||||
self._prev_time = cur_time
|
||||
self._prev_num_token = cur_num_token
|
||||
self._avg_throughput += (throughput - self._avg_throughput) / (
|
||||
state.global_step - 1
|
||||
)
|
||||
|
||||
gpu_stats = device_stats.gpu_stats()
|
||||
self._peak_mem = max(
|
||||
gpu_stats.reserved + gpu_stats.smi_diff, self._peak_mem
|
||||
)
|
||||
logging.info(
|
||||
'on_step_end: Throughput: %.2f token/s. %s, %s',
|
||||
throughput,
|
||||
device_stats.gpu_stats_str(gpu_stats),
|
||||
device_stats.cpu_stats_str(),
|
||||
)
|
||||
|
||||
def on_train_begin(
|
||||
self,
|
||||
args: TrainingArguments,
|
||||
state: TrainerState,
|
||||
control: TrainerControl,
|
||||
**kwargs,
|
||||
):
|
||||
if self._partial_state.is_main_process:
|
||||
self._start_time = time.time()
|
||||
logging.info(
|
||||
'on_train_begin: %s, %s',
|
||||
device_stats.gpu_stats_str(),
|
||||
device_stats.cpu_stats_str(),
|
||||
)
|
||||
|
||||
def on_train_end(
|
||||
self,
|
||||
args: TrainingArguments,
|
||||
state: TrainerState,
|
||||
control: TrainerControl,
|
||||
**kwargs,
|
||||
):
|
||||
if self._partial_state.is_main_process:
|
||||
train_time = time.time() - self._start_time
|
||||
throughput = state.num_input_tokens_seen / train_time
|
||||
logging.info(
|
||||
'training time %.2f s, throughput (including overhead, e.g., ckpt'
|
||||
' saving): %.2f token/s, peak_mem: %.2f GB',
|
||||
train_time,
|
||||
throughput,
|
||||
self._peak_mem,
|
||||
)
|
||||
if self._filename:
|
||||
with open(self._filename, 'a') as out_f:
|
||||
out_f.write(
|
||||
f'{self._max_seq_length/1024.0:.1f} | {self._peak_mem:.2f} |'
|
||||
f' {self._avg_throughput:.2f}\n'
|
||||
)
|
||||
-18
@@ -1,18 +0,0 @@
|
||||
group:
|
||||
- vertex
|
||||
task: custom_loglikelihood
|
||||
dataset_path: json
|
||||
dataset_name: null
|
||||
output_type: loglikelihood
|
||||
training_split: null
|
||||
validation_split: null
|
||||
test_split: test
|
||||
doc_to_text: "Request: {{prompt}}\nResponse:"
|
||||
doc_to_target: " {{ground_truth}}"
|
||||
metric_list:
|
||||
- metric: perplexity
|
||||
aggregation: perplexity
|
||||
higher_is_better: false
|
||||
- metric: acc
|
||||
aggregation: mean
|
||||
higher_is_better: true
|
||||
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
distributed_type: MULTI_GPU
|
||||
downcast_bf16: 'no'
|
||||
enable_cpu_affinity: false
|
||||
gpu_ids: all
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
mixed_precision: fp16
|
||||
num_machines: 1
|
||||
num_processes: 4
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
distributed_type: MULTI_GPU
|
||||
downcast_bf16: 'no'
|
||||
enable_cpu_affinity: false
|
||||
gpu_ids: all
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
mixed_precision: fp16
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero2.json
|
||||
zero3_init_flag: true
|
||||
distributed_type: DEEPSPEED
|
||||
downcast_bf16: 'no'
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 4
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero2.json
|
||||
zero3_init_flag: true
|
||||
distributed_type: DEEPSPEED
|
||||
downcast_bf16: 'no'
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero3.json
|
||||
zero3_init_flag: true
|
||||
distributed_type: DEEPSPEED
|
||||
downcast_bf16: 'no'
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 4
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-17
@@ -1,17 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero3.json
|
||||
zero3_init_flag: true
|
||||
distributed_type: DEEPSPEED
|
||||
downcast_bf16: 'no'
|
||||
machine_rank: 0
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-28
@@ -1,28 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
distributed_type: FSDP
|
||||
downcast_bf16: 'no'
|
||||
enable_cpu_affinity: false
|
||||
fsdp_config:
|
||||
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
|
||||
fsdp_transformer_layer_cls_to_wrap: Gemma2DecoderLayer
|
||||
fsdp_backward_prefetch: NO_PREFETCH
|
||||
fsdp_cpu_ram_efficient_loading: true
|
||||
fsdp_forward_prefetch: false
|
||||
fsdp_offload_params: true
|
||||
fsdp_sharding_strategy: FULL_SHARD
|
||||
fsdp_state_dict_type: SHARDED_STATE_DICT
|
||||
fsdp_sync_module_states: true
|
||||
fsdp_use_orig_params: false
|
||||
fsdp_activation_checkpointing: false
|
||||
main_training_function: main
|
||||
mixed_precision: bf16
|
||||
machine_rank: 0
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
-28
@@ -1,28 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
debug: false
|
||||
distributed_type: FSDP
|
||||
downcast_bf16: 'no'
|
||||
enable_cpu_affinity: false
|
||||
fsdp_config:
|
||||
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
|
||||
fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
|
||||
fsdp_backward_prefetch: NO_PREFETCH
|
||||
fsdp_cpu_ram_efficient_loading: true
|
||||
fsdp_forward_prefetch: false
|
||||
fsdp_offload_params: true
|
||||
fsdp_sharding_strategy: FULL_SHARD
|
||||
fsdp_state_dict_type: SHARDED_STATE_DICT
|
||||
fsdp_sync_module_states: true
|
||||
fsdp_use_orig_params: false
|
||||
fsdp_activation_checkpointing: false
|
||||
main_training_function: main
|
||||
mixed_precision: bf16
|
||||
machine_rank: 0
|
||||
num_machines: 16
|
||||
num_processes: 128
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user