xiaohei/model_garden
1
0
Fork 0
mirror of https://github.com/GoogleCloudPlatform/vertex-ai-samples.git synced 2026-09-26 14:42:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update PEFT train docker code (#3613)

Browse Source 
Co-authored-by: minwoopark <minwoopark@google.com>
This commit is contained in:
Minwoo Park
2024-10-02 19:13:30 +00:00
committed by GitHub
co-authored by minwoopark
parent e500e70580
commit 24c001eaf6
69 changed files with 5645 additions and 829 deletions
Download Patch File Download Diff File Expand all files Collapse all files
community-content/vertex_model_garden/model_oss/peft/causal_language_modeling_lora.py
-144
View File
@@ -1,144 +0,0 @@
"""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 typing import List
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,
target_modules: List[str] = constants.CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES,
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=target_modules,
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)
community-content/vertex_model_garden/model_oss/peft/dockerfile/code_llama_convert.Dockerfile
+28
View File
@@ -0,0 +1,28 @@
# 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"]
community-content/vertex_model_garden/model_oss/peft/dockerfile/conversion.Dockerfile
+22
View File
@@ -0,0 +1,22 @@
# 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"]
community-content/vertex_model_garden/model_oss/peft/dockerfile/serve.Dockerfile
+28 -36
View File
@@ -7,39 +7,40 @@
# 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.7.0-gpu
FROM pytorch/torchserve:0.11.0-gpu
USER root
ENV infer_port=7080
ENV mng_port=7081
ENV model_name="peft_serving"
ENV INFER_PORT=7080
ENV MNG_PORT=7081
ENV MODEL="peft_serving"
ENV PATH="/home/model-server/:${PATH}"
RUN apt-get update && apt-get install -y --no-install-recommends \
RUN apt-get update && apt-get -y upgrade && 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
RUN pip install --upgrade torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
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 python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
RUN pip install peft==0.5.0
RUN pip install datasets==2.14.4
RUN pip install triton==2.0.0.dev20221120
RUN pip install triton==3.0.0
RUN pip install xformers==0.0.20
RUN pip install google-cloud-storage==2.7.0
RUN pip install absl-py==1.4.0
RUN pip install google-cloud-storage
RUN pip install absl-py
RUN pip install scipy==1.10.1
RUN pip install evaluate==0.4.0
RUN pip install scikit-learn==1.2.2
@@ -47,52 +48,43 @@ 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
# 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
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
# 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 MODEL_ID "openlm-research/open_llama_7b"
ENV BASE_MODEL_ID ""
ENV MODEL_ID ""
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_name} \
--model-name=${MODEL} \
--version=1.0 \
--handler=/home/model-server/handler.py \
--runtime=python3 \
@@ -103,5 +95,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_name}=${model_name}.mar", \
"--models", "${MODEL}=${MODEL}.mar", \
"--model-store", "/home/model-server/model-store"]
community-content/vertex_model_garden/model_oss/peft/dockerfile/train.Dockerfile
-111
View File
@@ -1,111 +0,0 @@
# 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"]
community-content/vertex_model_garden/model_oss/peft/handler.py
+50 -15
View File
@@ -72,15 +72,39 @@ class PeftHandler(BaseHandler):
"PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
)
self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
self.base_model_id = os.environ.get("BASE_MODEL_ID", None)
self.model_id = self.base_model_id
if not self.base_model_id:
self.model_id = os.environ.get("MODEL_ID", "")
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)
logging.info(f"Load base model id from MODEL_ID:{self.model_id}.")
if not self.model_id:
self.model_id = os.environ.get("AIP_STORAGE_URI", "")
logging.info(f"Load base model id from AIP_STORAGE_URI: {self.model_id}.")
if not self.model_id:
raise ValueError("Base model id is must be set.")
if fileutils.is_gcs_path(self.model_id):
@@ -101,8 +125,7 @@ class PeftHandler(BaseHandler):
logging.info(
f"Using task:{self.task}, base model:{self.model_id}, lora model:"
f" {self.finetuned_lora_model_path}, and precision"
f" {self.precision_mode}."
f" {self.finetuned_lora_model_path}, precision {self.precision_mode}."
)
self.pipeline = None
@@ -145,11 +168,18 @@ class PeftHandler(BaseHandler):
elif (
self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
):
tokenizer = AutoTokenizer.from_pretrained(self.model_id)
tokenizer = AutoTokenizer.from_pretrained(
self.model_id,
trust_remote_code=self.trust_remote_code,
)
logging.debug("Initialized the tokenizer.")
if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
if self.quantization == constants.AWQ:
model = AutoAWQForCausalLM.from_quantized(self.model_id)
model = AutoAWQForCausalLM.from_quantized(
self.model_id,
trust_remote_code=self.trust_remote_code,
)
elif self.quantization == constants.GPTQ or not self.quantization:
if self.precision_mode == constants.PRECISION_MODE_32:
model = AutoModelForCausalLM.from_pretrained(
@@ -157,6 +187,7 @@ class PeftHandler(BaseHandler):
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(
@@ -164,6 +195,7 @@ class PeftHandler(BaseHandler):
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(
@@ -171,6 +203,7 @@ class PeftHandler(BaseHandler):
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(
@@ -182,6 +215,7 @@ class PeftHandler(BaseHandler):
torch_dtype=torch.float16,
device_map="auto",
quantization_config=quantization_config,
trust_remote_code=self.trust_remote_code,
)
else:
quantization_config = BitsAndBytesConfig(
@@ -195,6 +229,7 @@ class PeftHandler(BaseHandler):
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}")
@@ -203,14 +238,14 @@ class PeftHandler(BaseHandler):
model = AutoModelForCausalLM.from_pretrained(
self.model_id,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
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=True,
trust_remote_code=self.trust_remote_code,
device_map="auto",
)
logging.debug("Initialized the base model.")
@@ -329,4 +364,4 @@ class PeftHandler(BaseHandler):
return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
# pylint: enable=logging-fstring-interpolation
# pylint: enable=logging-fstring-interpolation
community-content/vertex_model_garden/model_oss/peft/hf_transformers_lazy_download.patch
-131
View File
@@ -1,131 +0,0 @@
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(
community-content/vertex_model_garden/model_oss/peft/instruct_lora.py
-95
View File
@@ -1,95 +0,0 @@
"""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
from typing import List
from util import constants
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,
target_modules: List[str] = constants.INSTRUCT_LORA_TARGET_MODULES,
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=target_modules,
)
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()
community-content/vertex_model_garden/model_oss/peft/main.py
-185
View File
@@ -1,185 +0,0 @@
"""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.',
)
_TARGET_MODULES = flags.DEFINE_list(
'target_modules',
constants.CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES,
'The comma separated list of target modules for LoRa training.',
)
_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,
target_modules=_TARGET_MODULES.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,
target_modules=_TARGET_MODULES.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)
community-content/vertex_model_garden/model_oss/peft/templates/README.md
-98
View File
@@ -1,98 +0,0 @@
# Vertex Model Garden Training Dataset Template
## Overview
Vertex Model Garden training provides templates for streamlined preprocessing of
datasets. Although datasets often have intricate structures, the supported LLM
models accept only flat strings. A template facilitates parsing a dataset and
preprocessing it to be compatible with the model.
When fine-tuning a pretrained model, it is advisable to maintain the same format
as the original training data. A template helps replicate the format, ensuring
consistency and potentially enhancing the fine-tuning process.
Both multi-turn messages and single instruction-response pairs are supported.
Multi-turn messages are accommodated using a more general `chat_template` field,
whereas simple instruction-response pair datasets are supported through the
`prompt_input` field.
A template is a JSON file consisting of string key-value pairs. Refer to the
following for the definitions of the supported fields.
## Template field documentation
**description**: An explanation of the template.
**source**: Information about the origin of the template.
**chat_template**: A
[jinja template](https://jinja.palletsprojects.com/en/3.1.x/templates/) that can
be used to parse a chat dataset. This is the same format as
[HF chat templates](https://huggingface.co/docs/transformers/main/en/chat_templating).
To create a chat_template, use the `messages` variable to be filled with the
sample. The flag `--instruct_column_in_dataset` identifies which column will be
passed to the `messages` variable in the chat_template. This field is mutually
exclusive with `prompt_input` and `prompt_no_input`.
**prompt_input**: A string template that is used when value for the input column
exists in the sample. It should be able to be formatted with the
[str.format](https://docs.python.org/3/library/stdtypes.html#str.format) method.
The input column is specified with the flag `--instruct_column_in_dataset`. Used
for instruction dataset. This field is mutually exclusive with `chat_template`.
**prompt_no_input**: A string template that is used when value for the input
column does not exist in the sample. It should be able to be formatted with the
[str.format](https://docs.python.org/3/library/stdtypes.html#str.format) method.
The input column is specified with the flag `--instruct_column_in_dataset`. Used
for instruction dataset. This field is mutually exclusive with `chat_template`.
**instruction_separator**: A unique string used to indicate the start of the
instructions. If not specified, every token after response_separator will be
treated as a response, and every token before the first response_separator will
be treated as instruction.
**response_separator**: A unique string used to indicate the start of the
response. This field is required if `--completion_only` flag is set to `True`.
## Example templates
- See the list of all supported templates [here](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content/vertex_model_garden/model_oss/peft/train/vmg/templates).
- For an example with `chat_template` see the JSON template below.
```
{
"description": "Chat template used by Llama 3.",
"source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
"chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
"instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
"response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
}
```
- For an example with `prompt_input` see the JSON template below. In this case
the flag `--instruct_column_in_dataset=text` should be set, and there must
be a column named `text` in the dataset.
```
{
"description": "Template for openassistant-guanaco dataset.",
"source": "https://huggingface.co/datasets/timdettmers/openassistant-guanaco",
"prompt_input": "{text}",
"instruction_separator": "### Human:",
"response_separator": "### Assistant:"
}
```
- For an example with `prompt_no_input` see the JSON template below. In this
case the flag `--instruct_column_in_dataset=input` should be set, and there
must be columns named `input` and `instruction` in the dataset.
```
{
"description": "Template used by Alpaca-LoRA.",
"source": "https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json",
"prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
"prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
"response_separator": "### Response:"
}
```
community-content/vertex_model_garden/model_oss/peft/templates/openai-chat.json
-7
View File
@@ -1,7 +0,0 @@
{
"description": "Template used by openai chat.",
"source": "https://platform.openai.com/docs/api-reference/fine-tuning/chat-input",
"chat_template": "{% set loop_messages = messages %}{% set content = '' %}{% for message in loop_messages %}{% set content = content ~ '\n\n<|start_header_id|>' ~ message.role ~ '<|end_header_id|>\n\n' %}{% if message.content is string %}{% set content = content ~ message.content|trim ~ '<|eot_id|>' %}{% else %}{% set content = content ~ message.content|join(' ', attribute='text')|trim ~ '<|eot_id|>' %}{% endif %}{% if loop.index0 == 0 %}{% set content = bos_token ~ content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
"instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
"response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
}
community-content/vertex_model_garden/model_oss/peft/train.sh
-6
View File
@@ -1,6 +0,0 @@
#!/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 "$@"
community-content/vertex_model_garden/model_oss/peft/train/axolotl/dockerfile/train.Dockerfile
+16
View File
@@ -0,0 +1,16 @@
# 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"]
community-content/vertex_model_garden/model_oss/peft/train/axolotl/scripts/build_train_docker.sh
Executable
+20
View File
@@ -0,0 +1,20 @@
#!/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
community-content/vertex_model_garden/model_oss/peft/train/axolotl/train_entrypoint.py
+88
View File
@@ -0,0 +1,88 @@
"""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()
community-content/vertex_model_garden/model_oss/peft/train/tests/docker_command_builder.py
+45
View File
@@ -0,0 +1,45 @@
"""Class that bundles docker related flags."""
import getpass
import os
import pwd
class DockerCommandBuilder:
"""Bundle docker related flags."""
def __init__(self, docker_uri, shm_size='128gb'):
self._docker_uri = [docker_uri]
self._defaults = [
'docker',
'run',
'--gpus=all',
'--net=host',
'--rm',
f'--shm-size={shm_size}',
]
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',
]
self._env_vars = []
self._mount_maps = []
def add_env_var(self, var, val):
self._env_vars.append(f'--env={var}={val}')
def add_mount_map(self, host_path, docker_path):
self._mount_maps.append(f'--volume={host_path}:{docker_path}')
def build_cmd(self) -> str:
return self._defaults + self._env_vars + self._mount_maps + self._docker_uri
community-content/vertex_model_garden/model_oss/peft/train/tests/instruct_lora_command_builder.py
+429
View File
@@ -0,0 +1,429 @@
# pylint: disable=W,C,R
# DO NOT MODIFY: this file is auto-generated
# See go/vmg-oss-peft-tests#command-builder-genpy
class InstructLoraCommandBuilder:
def __init__(self):
self._config_file = None
self._task = None
self._pretrained_model_id = None
self._dataset_name = None
self._train_split_name = None
self._template = None
self._instruct_column_in_dataset = None
self._output_dir = None
self._merge_base_and_lora_output_dir = None
self._logging_output_dir = None
self._per_device_train_batch_size = None
self._gradient_accumulation_steps = None
self._lora_rank = None
self._lora_alpha = None
self._lora_dropout = None
self._max_steps = None
self._num_epochs = None
self._max_seq_length = None
self._learning_rate = None
self._lr_scheduler_type = None
self._precision_mode = None
self._train_precision = None
self._enable_gradient_checkpointing = None
self._use_example_packing = None
self._attn_implementation = None
self._optimizer = None
self._warmup_ratio = None
self._report_to = None
self._save_steps = None
self._logging_steps = None
self._huggingface_access_token = None
self._eval_dataset_path = None
self._eval_column = None
self._eval_template = None
self._eval_split = None
self._eval_steps = None
self._eval_tasks = None
self._eval_metric_name = None
self._completion_only = None
self._max_grad_norm = None
self._logger_level = None
self._benchmark_out_file = None
self._tuning_data_stats_file = None
self._enable_peft = None
self._merge_model_precision_mode = None
self._target_modules = None
@property
def config_file(self):
return self._config_file
@config_file.setter
def config_file(self, val: str):
self._config_file = val
@property
def task(self):
return self._task
@task.setter
def task(self, val: str):
self._task = val
@property
def pretrained_model_id(self):
return self._pretrained_model_id
@pretrained_model_id.setter
def pretrained_model_id(self, val: str):
self._pretrained_model_id = val
@property
def train_dataset(self):
return self._dataset_name
@train_dataset.setter
def train_dataset(self, val: str):
self._dataset_name = val
@property
def train_split_name(self):
return self._train_split_name
@train_split_name.setter
def train_split_name(self, val: str):
self._train_split_name = val
@property
def template(self):
return self._template
@template.setter
def template(self, val: str):
self._template = val
@property
def instruct_column(self):
return self._instruct_column_in_dataset
@instruct_column.setter
def instruct_column(self, val: str):
self._instruct_column_in_dataset = val
@property
def ckpt_dir(self):
return self._output_dir
@ckpt_dir.setter
def ckpt_dir(self, val: str):
self._output_dir = val
@property
def merged_model_dir(self):
return self._merge_base_and_lora_output_dir
@merged_model_dir.setter
def merged_model_dir(self, val: str):
self._merge_base_and_lora_output_dir = val
@property
def logging_dir(self):
return self._logging_output_dir
@logging_dir.setter
def logging_dir(self, val: str):
self._logging_output_dir = val
@property
def per_device_batch_size(self):
return self._per_device_train_batch_size
@per_device_batch_size.setter
def per_device_batch_size(self, val: int):
self._per_device_train_batch_size = val
@property
def gradient_accumulation_steps(self):
return self._gradient_accumulation_steps
@gradient_accumulation_steps.setter
def gradient_accumulation_steps(self, val: int):
self._gradient_accumulation_steps = val
@property
def lora_rank(self):
return self._lora_rank
@lora_rank.setter
def lora_rank(self, val: int):
self._lora_rank = val
@property
def lora_alpha(self):
return self._lora_alpha
@lora_alpha.setter
def lora_alpha(self, val: int):
self._lora_alpha = val
@property
def lora_dropout(self):
return self._lora_dropout
@lora_dropout.setter
def lora_dropout(self, val: float):
self._lora_dropout = val
@property
def max_steps(self):
return self._max_steps
@max_steps.setter
def max_steps(self, val: int):
self._max_steps = val
@property
def num_epochs(self):
return self._num_epochs
@num_epochs.setter
def num_epochs(self, val: float):
self._num_epochs = val
@property
def max_seq_length(self):
return self._max_seq_length
@max_seq_length.setter
def max_seq_length(self, val: int):
self._max_seq_length = val
@property
def learning_rate(self):
return self._learning_rate
@learning_rate.setter
def learning_rate(self, val: float):
self._learning_rate = val
@property
def lr_scheduler_type(self):
return self._lr_scheduler_type
@lr_scheduler_type.setter
def lr_scheduler_type(self, val: str):
self._lr_scheduler_type = val
@property
def load_precision(self):
return self._precision_mode
@load_precision.setter
def load_precision(self, val: str):
self._precision_mode = val
@property
def train_precision(self):
return self._train_precision
@train_precision.setter
def train_precision(self, val: str):
self._train_precision = val
@property
def gradient_checkpointing(self):
return self._enable_gradient_checkpointing
@gradient_checkpointing.setter
def gradient_checkpointing(self, val: bool):
self._enable_gradient_checkpointing = val
@property
def example_packing(self):
return self._use_example_packing
@example_packing.setter
def example_packing(self, val: bool):
self._use_example_packing = val
@property
def attn_implementation(self):
return self._attn_implementation
@attn_implementation.setter
def attn_implementation(self, val: str):
self._attn_implementation = val
@property
def optimizer(self):
return self._optimizer
@optimizer.setter
def optimizer(self, val: str):
self._optimizer = val
@property
def warmup_ratio(self):
return self._warmup_ratio
@warmup_ratio.setter
def warmup_ratio(self, val: float):
self._warmup_ratio = val
@property
def report_to(self):
return self._report_to
@report_to.setter
def report_to(self, val: str):
self._report_to = val
@property
def save_steps(self):
return self._save_steps
@save_steps.setter
def save_steps(self, val: int):
self._save_steps = val
@property
def logging_steps(self):
return self._logging_steps
@logging_steps.setter
def logging_steps(self, val: int):
self._logging_steps = val
@property
def huggingface_access_token(self):
return self._huggingface_access_token
@huggingface_access_token.setter
def huggingface_access_token(self, val: str):
self._huggingface_access_token = val
@property
def eval_dataset(self):
return self._eval_dataset_path
@eval_dataset.setter
def eval_dataset(self, val: str):
self._eval_dataset_path = val
@property
def eval_instruct_column(self):
return self._eval_column
@eval_instruct_column.setter
def eval_instruct_column(self, val: str):
self._eval_column = val
@property
def eval_template(self):
return self._eval_template
@eval_template.setter
def eval_template(self, val: str):
self._eval_template = val
@property
def eval_split_name(self):
return self._eval_split
@eval_split_name.setter
def eval_split_name(self, val: str):
self._eval_split = val
@property
def eval_steps(self):
return self._eval_steps
@eval_steps.setter
def eval_steps(self, val: int):
self._eval_steps = val
@property
def eval_tasks(self):
return self._eval_tasks
@eval_tasks.setter
def eval_tasks(self, val: str):
self._eval_tasks = val
@property
def eval_metric_name(self):
return self._eval_metric_name
@eval_metric_name.setter
def eval_metric_name(self, val: str):
self._eval_metric_name = val
@property
def completion_only(self):
return self._completion_only
@completion_only.setter
def completion_only(self, val: bool):
self._completion_only = val
@property
def max_grad_norm(self):
return self._max_grad_norm
@max_grad_norm.setter
def max_grad_norm(self, val: float):
self._max_grad_norm = val
@property
def logger_level(self):
return self._logger_level
@logger_level.setter
def logger_level(self, val: str):
self._logger_level = val
@property
def benchmark_out_file(self):
return self._benchmark_out_file
@benchmark_out_file.setter
def benchmark_out_file(self, val: str):
self._benchmark_out_file = val
@property
def tuning_data_stats_file(self):
return self._tuning_data_stats_file
@tuning_data_stats_file.setter
def tuning_data_stats_file(self, val: str):
self._tuning_data_stats_file = val
@property
def enable_peft(self):
return self._enable_peft
@enable_peft.setter
def enable_peft(self, val: bool):
self._enable_peft = val
@property
def merge_model_precision_mode(self):
return self._merge_model_precision_mode
@merge_model_precision_mode.setter
def merge_model_precision_mode(self, val: str):
self._merge_model_precision_mode = val
@property
def target_modules(self):
return self._target_modules
@target_modules.setter
def target_modules(self, val: str):
self._target_modules = val
def build_cmd(self) -> str:
cmd = []
for k, v in self.__dict__.items():
if v is not None:
cmd.append(f'--{k[1:]}={v}')
return cmd
community-content/vertex_model_garden/model_oss/peft/train/tests/quantize_model_command_builder.py
+142
View File
@@ -0,0 +1,142 @@
# pylint: disable=W,C,R
# DO NOT MODIFY: this file is auto-generated
# See go/vmg-oss-peft-tests#command-builder-genpy
class QuantizeModelCommandBuilder:
def __init__(self):
self._task = None
self._pretrained_model_id = None
self._quantization_method = None
self._quantization_precision_mode = None
self._quantization_dataset_name = None
self._text_column_in_quantization_dataset = None
self._quantization_output_dir = None
self._device_map = None
self._max_memory = None
self._group_size = None
self._desc_act = None
self._damp_percent = None
self._cache_examples_on_gpu = None
self._awq_version = None
@property
def task(self):
return self._task
@task.setter
def task(self, val: str):
self._task = val
@property
def pretrained_model_id(self):
return self._pretrained_model_id
@pretrained_model_id.setter
def pretrained_model_id(self, val: str):
self._pretrained_model_id = val
@property
def quantization_method(self):
return self._quantization_method
@quantization_method.setter
def quantization_method(self, val: str):
self._quantization_method = val
@property
def quantization_precision_mode(self):
return self._quantization_precision_mode
@quantization_precision_mode.setter
def quantization_precision_mode(self, val: str):
self._quantization_precision_mode = val
@property
def quantization_dataset_name(self):
return self._quantization_dataset_name
@quantization_dataset_name.setter
def quantization_dataset_name(self, val: str):
self._quantization_dataset_name = val
@property
def text_column_in_quantization_dataset(self):
return self._text_column_in_quantization_dataset
@text_column_in_quantization_dataset.setter
def text_column_in_quantization_dataset(self, val: str):
self._text_column_in_quantization_dataset = val
@property
def quantization_output_dir(self):
return self._quantization_output_dir
@quantization_output_dir.setter
def quantization_output_dir(self, val: str):
self._quantization_output_dir = val
@property
def device_map(self):
return self._device_map
@device_map.setter
def device_map(self, val: str):
self._device_map = val
@property
def max_memory(self):
return self._max_memory
@max_memory.setter
def max_memory(self, val: str):
self._max_memory = val
@property
def group_size(self):
return self._group_size
@group_size.setter
def group_size(self, val: int):
self._group_size = val
@property
def desc_act(self):
return self._desc_act
@desc_act.setter
def desc_act(self, val: bool):
self._desc_act = val
@property
def damp_percent(self):
return self._damp_percent
@damp_percent.setter
def damp_percent(self, val: float):
self._damp_percent = val
@property
def cache_examples_on_gpu(self):
return self._cache_examples_on_gpu
@cache_examples_on_gpu.setter
def cache_examples_on_gpu(self, val: bool):
self._cache_examples_on_gpu = val
@property
def awq_version(self):
return self._awq_version
@awq_version.setter
def awq_version(self, val: str):
self._awq_version = val
def build_cmd(self) -> str:
cmd = []
for k, v in self.__dict__.items():
if v is not None:
cmd.append(f'--{k[1:]}={v}')
return cmd
community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_adapters.py
+106
View File
@@ -0,0 +1,106 @@
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
"""Tests adapters of PEFT train docker."""
import inspect
import os
import time
from absl.testing import absltest
from absl.testing import parameterized
import instruct_lora_command_builder as task_cmd_builder
from safetensors import safe_open
import test_util
class AdapterTest(test_util.TestBase):
# Needs to be accessible outside docker to check artifacts.
_TEST_OUTPUT_DIR = os.path.expanduser('~/output')
_MODULES_NEED_TO_BE_EXCLUDED_IN_ADAPTER = ['lm_head', 'embed_tokens']
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.test_suite_output_dir = os.path.join(
cls._TEST_OUTPUT_DIR,
os.path.splitext(os.path.basename(__file__))[0],
cls.__class__.__name__,
)
def setUp(self):
super().setUp()
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'peft_train_sample.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'input_text'
self.task_cmd_builder.template = 'llama3-text-bison'
self.task_cmd_builder.gradient_accumulation_steps = 1
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.max_steps = 1
self.task_cmd_builder.max_seq_length = 256
self.task_cmd_builder.load_precision = '4bit'
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.save_steps = 10
self.task_cmd_builder.max_steps = 3
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
def setup_output_dir(self, testcase_name: str):
testcase_output_dir = os.path.join(
self.test_suite_output_dir, testcase_name, test_util.get_timestamp()
)
self.task_cmd_builder.ckpt_dir = os.path.join(
testcase_output_dir, 'adapter'
)
self.task_cmd_builder.logging_dir = os.path.join(
testcase_output_dir, 'logs'
)
def check_adapter_for_bad_modules(self, adapter_path):
unwanted_modules = set()
with safe_open(adapter_path, framework='pt', device='cpu') as f:
for key in f.keys():
for module in self._MODULES_NEED_TO_BE_EXCLUDED_IN_ADAPTER:
if module in key:
unwanted_modules.add(key)
assert (
not unwanted_modules
), f'Adapter includes unwanted modules: {unwanted_modules}'
@parameterized.named_parameters(
('llama3.1-8b', 'llama3.1-8b-hf'),
('llama3.1-70b', 'llama3.1-70b-hf'),
('llama2-7b', 'llama2-7b-hf'),
)
def test_llama_adapters(self, model_name):
test_function_name = inspect.stack()[0][3]
self.setup_output_dir(f'{test_function_name}-{model_name}')
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 1000)
adapter = os.path.join(
self.task_cmd_builder.ckpt_dir,
'checkpoint-final/adapter_model.safetensors',
)
self.check_adapter_for_bad_modules(adapter)
if __name__ == '__main__':
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_features.py
+290
View File
@@ -0,0 +1,290 @@
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
"""Tests various features of PEFT train docker."""
import os
import time
from absl.testing import absltest
from absl.testing import parameterized
import instruct_lora_command_builder as task_cmd_builder
import test_util
class GcsUploadDownloadTest(test_util.TestBase):
def setUp(self):
super().setUp()
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'peft_train_sample.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'input_text'
self.task_cmd_builder.template = 'llama3-text-bison'
self.task_cmd_builder.gradient_accumulation_steps = 1
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.max_steps = 1
self.task_cmd_builder.max_seq_length = 256
self.task_cmd_builder.load_precision = '4bit'
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.ckpt_dir = '/tmp'
@parameterized.named_parameters(
(
'llama3_8b_gcs',
'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf',
),
('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
)
def test_model_download_single_process(self, pretrained_model_id):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(pretrained_model_id)
)
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 5 * 60.0)
@parameterized.named_parameters(
(
'llama3_8b_gcs',
'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf',
),
('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
)
def test_model_download_multi_process(self, pretrained_model_id):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(pretrained_model_id)
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 5 * 60.0)
def test_70b_model_download(self):
self.task_cmd_builder.pretrained_model_id = (
'gs://vertex-model-garden-public-us/llama3/llama3-70b-hf'
)
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 10 * 60.0)
@parameterized.named_parameters(
('merged-without-upload', '/tmp/merged'),
('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
)
def test_model_merge(self, output_dir):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
ckpt_dir = os.path.join(
output_dir,
f'output-{test_util.get_timestamp()}',
)
self.task_cmd_builder.ckpt_dir = ckpt_dir
self.task_cmd_builder.merged_model_dir = os.path.join(ckpt_dir, 'merged')
self.task_cmd_builder.logging_dir = '/tmp/logging'
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 5 * 60.0)
def test_model_fp8_conversion(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
ckpt_dir = f'/tmp/output/output-{test_util.get_timestamp()}'
self.task_cmd_builder.ckpt_dir = ckpt_dir
self.task_cmd_builder.merged_model_dir = os.path.join(ckpt_dir, 'merged')
self.task_cmd_builder.logging_dir = os.path.join(ckpt_dir, 'logging')
self.task_cmd_builder.merge_model_precision_mode = 'float8'
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 5 * 60.0)
@parameterized.named_parameters(
('merged-without-upload', '/tmp/merged'),
('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
)
def test_model_merge_and_upload_deepspeed(self, merged_model_dir):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/deepspeed_zero3_8gpu.yaml'
)
self.task_cmd_builder.merged_model_dir = os.path.join(
merged_model_dir, f'merged-{test_util.get_timestamp()}'
)
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 5 * 60.0)
@parameterized.named_parameters(
('save-only-last', 10),
('save-multiple-times', 1),
)
def test_llama3_8b_save_and_merge_8_gpus_fsdp(self, save_steps):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.save_steps = save_steps
self.task_cmd_builder.max_steps = 3
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.task_cmd_builder.merged_model_dir = '/tmp/merged'
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 9 * 60.0)
class TemplateAndDataStatsTest(test_util.TestBase):
def setUp(self):
super().setUp()
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.gradient_accumulation_steps = 1
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.max_steps = 1
self.task_cmd_builder.max_seq_length = 256
self.task_cmd_builder.load_precision = '4bit'
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.ckpt_dir = '/tmp'
def test_openai_chat_template(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'openai-multi-chat-example-data.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'messages'
self.task_cmd_builder.template = 'llama3'
self.assertEqual(self.run_cmd(), 0)
def test_openai_completion_template(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'openai-completion-example-data.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'prompt'
self.task_cmd_builder.template = 'openai-completion'
self.assertEqual(self.run_cmd(), 0)
def test_data_stats_chat_template(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'openai-multi-chat-example-data.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'messages'
self.task_cmd_builder.template = 'llama3'
self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
self.assertEqual(self.run_cmd(), 0)
def test_data_stats_completion_template(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'openai-completion-example-data.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'prompt'
self.task_cmd_builder.template = 'openai-completion'
self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
self.assertEqual(self.run_cmd(), 0)
class TargetModulesTest(test_util.TestBase):
def setUp(self):
super().setUp()
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'peft_train_sample.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'input_text'
self.task_cmd_builder.template = 'llama3-text-bison'
self.task_cmd_builder.gradient_accumulation_steps = 1
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.max_steps = 1
self.task_cmd_builder.max_seq_length = 256
self.task_cmd_builder.load_precision = '4bit'
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.ckpt_dir = '/tmp'
def test_target_modules(self):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id('llama3.1-8b-hf')
)
self.task_cmd_builder.target_modules = 'q_proj, v_proj, k_proj'
self.assertEqual(self.run_cmd(), 0)
if __name__ == '__main__':
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_throughput.py
+180
View File
@@ -0,0 +1,180 @@
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
"""Tests to check training throughput and GPU memory consumption."""
import os
import pathlib
from absl.testing import absltest
from absl.testing import parameterized
import instruct_lora_command_builder as task_cmd_builder
import test_util
class TrainerThroughputTest(test_util.TestBase):
_TEST_OUTPUT_DIR = os.path.expanduser('~/throughput_tests')
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.test_suite_output_dir = os.path.join(
cls._TEST_OUTPUT_DIR, os.path.splitext(os.path.basename(__file__))[0]
)
if not os.path.isdir(cls.test_suite_output_dir):
pathlib.Path(cls.test_suite_output_dir).mkdir(parents=True)
def setUp(self):
super().setUp()
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.gradient_accumulation_steps = 1
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.learning_rate = 5e-5
self.task_cmd_builder.warmup_ratio = 0.01
self.task_cmd_builder.max_steps = 10
self.task_cmd_builder.save_steps = 1000
self.task_cmd_builder.logging_steps = 1
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.example_packing = True
self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2-1k'
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'text'
self.task_cmd_builder.template = 'openassistant-guanaco'
self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
self.task_cmd_builder.logging_dir = '/tmp/logs'
def run_cmd_and_handle_failure(self):
ret = self.run_cmd()
if ret != 0:
with open(self.task_cmd_builder.benchmark_out_file, 'a') as f:
max_seq_length = self.task_cmd_builder.max_seq_length
f.write(f'{max_seq_length/1024.0:.1f}k | failed | n/a\n')
return ret
@parameterized.product(
model_name=[
'llama3-70b-hf',
'llama3.1-70b-hf',
'Mistral-7B-v0.1',
'Mixtral-8x7B-v0.1',
'Gemma2-9b-it',
],
precision=['4bit', '8bit', 'bfloat16'],
max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
)
def test_model_single_gpu(self, model_name, precision, max_seq_length):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.max_seq_length = max_seq_length
self.task_cmd_builder.load_precision = precision
self.task_cmd_builder.benchmark_out_file = os.path.join(
self.test_suite_output_dir, f'bm_{model_name}_{precision}.txt'
)
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
self.assertEqual(self.run_cmd_and_handle_failure(), 0)
@parameterized.product(
model_name=[
'llama3-70b-hf',
'llama3.1-70b-hf',
'Mistral-7B-v0.1',
'Mixtral-8x7B-v0.1',
'Gemma2-9b-it',
],
precision=['4bit', '8bit', 'bfloat16'],
max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
num_gpus=[8],
config=['deepspeed_zero2', 'deepspeed_zero3'],
)
def test_model_multi_gpu_deepspeed(
self, model_name, precision, max_seq_length, num_gpus, config
):
self.assertTrue(num_gpus == 4 or num_gpus == 8)
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.max_seq_length = max_seq_length
self.task_cmd_builder.load_precision = precision
self.task_cmd_builder.benchmark_out_file = os.path.join(
self.test_suite_output_dir,
f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
)
self.task_cmd_builder.config_file = (
f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
)
self.docker_builder.add_env_var(
'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
)
self.assertEqual(self.run_cmd_and_handle_failure(), 0)
@parameterized.product(
model_name=['llama3.1-70b-hf'],
precision=['4bit', '8bit', 'bfloat16'],
max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
num_gpus=[8],
)
def test_model_multi_gpu_fsdp_lora(
self, model_name, precision, max_seq_length, num_gpus
):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.max_seq_length = max_seq_length
self.task_cmd_builder.load_precision = precision
self.task_cmd_builder.benchmark_out_file = os.path.join(
self.test_suite_output_dir,
f'bm_fsdp_{num_gpus}gpu_{model_name}_{precision}.txt',
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.docker_builder.add_env_var(
'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
)
self.assertEqual(self.run_cmd_and_handle_failure(), 0)
@parameterized.product(
model_name=['llama3.1-70b-hf'],
precision=['bfloat16'],
max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
num_gpus=[8],
)
def test_model_multi_gpu_fsdp_full_finetuning(
self, model_name, precision, max_seq_length, num_gpus
):
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.max_seq_length = max_seq_length
self.task_cmd_builder.load_precision = precision
self.task_cmd_builder.benchmark_out_file = os.path.join(
self.test_suite_output_dir,
f'bm_fsdp_full_finetuning_{num_gpus}gpu_{model_name}_{precision}.txt',
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.task_cmd_builder.enable_peft = False
self.docker_builder.add_env_var(
'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
)
self.assertEqual(self.run_cmd_and_handle_failure(), 0)
if __name__ == '__main__':
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_trained_model_quality.py
+152
View File
@@ -0,0 +1,152 @@
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
"""Tests to make sure trained model achieves decent quality.
Right now, the metric is loss decreasing and we'll eyeball the TB graphs.
"""
import os
from absl.testing import absltest
from absl.testing import parameterized
import instruct_lora_command_builder as task_cmd_builder
import test_util
class TrainedModelQualityTest(test_util.TestBase):
_TEST_OUTPUT_DIR = os.path.expanduser('~/output')
@classmethod
def setUpClass(cls):
super().setUpClass()
cls.test_suite_output_dir = os.path.join(
cls._TEST_OUTPUT_DIR,
os.path.splitext(os.path.basename(__file__))[0],
cls.__class__.__name__,
)
def setUp(self):
super().setUp()
self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
self.task_cmd_builder.task = 'instruct-lora'
self.task_cmd_builder.eval_tasks = 'builtin_eval'
self.task_cmd_builder.eval_metric_name = 'loss'
self.task_cmd_builder.per_device_batch_size = 1
self.task_cmd_builder.gradient_accumulation_steps = 8
self.task_cmd_builder.lora_rank = 16
self.task_cmd_builder.lora_alpha = 32
self.task_cmd_builder.lora_dropout = 0.05
self.task_cmd_builder.learning_rate = 5e-5
self.task_cmd_builder.num_epochs = 2.0
self.task_cmd_builder.warmup_ratio = 0.01
self.task_cmd_builder.max_steps = -1
self.task_cmd_builder.save_steps = 10
self.task_cmd_builder.eval_steps = 10
self.task_cmd_builder.max_seq_length = 4096
self.task_cmd_builder.load_precision = '4bit'
self.task_cmd_builder.gradient_checkpointing = True
self.task_cmd_builder.completion_only = True
self.task_cmd_builder.attn_implementation = 'flash_attention_2'
self.task_cmd_builder.report_to = 'tensorboard'
def setup_output_dir(self, testcase_name: str):
testcase_output_dir = os.path.join(
self.test_suite_output_dir, testcase_name
)
self.task_cmd_builder.ckpt_dir = os.path.join(
testcase_output_dir, 'adapter'
)
self.task_cmd_builder.logging_dir = os.path.join(
testcase_output_dir, 'logs'
)
self.task_cmd_builder.merged_model_dir = os.path.join(
testcase_output_dir, 'merged'
)
@parameterized.named_parameters(
('llama3-8b', 'llama3-8b-hf'),
('llama3.1-8b', 'llama3.1-8b-hf'),
)
def test_8b_model_deepspeed(self, model_name):
self.setup_output_dir(f'test_deepspeed_{model_name}')
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
'peft_train_sample.jsonl'
)
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'input_text'
self.task_cmd_builder.template = 'llama3-text-bison'
self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
'peft_eval_sample.jsonl'
)
self.task_cmd_builder.eval_split_name = 'train'
self.task_cmd_builder.eval_instruct_column = (
self.task_cmd_builder.instruct_column
)
self.task_cmd_builder.eval_template = self.task_cmd_builder.template
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
self.assertEqual(self.run_cmd(), 0)
@parameterized.named_parameters(
('llama3-70b', 'llama3-70b-hf'),
('llama3.1-70b', 'llama3.1-70b-hf'),
)
def test_70b_model_deepspeed(self, model_name):
self.setup_output_dir(f'test_deepspeed_{model_name}')
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
)
self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'text'
self.task_cmd_builder.template = 'openassistant-guanaco'
self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
self.task_cmd_builder.eval_split_name = 'test'
self.task_cmd_builder.eval_instruct_column = (
self.task_cmd_builder.instruct_column
)
self.task_cmd_builder.eval_template = self.task_cmd_builder.template
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
self.assertEqual(self.run_cmd(), 0)
@parameterized.named_parameters(
('llama3-70b', 'llama3-70b-hf'),
('llama3.1-70b', 'llama3.1-70b-hf'),
)
def test_70b_model_fsdp(self, model_name):
self.setup_output_dir(f'test_fsdp_{model_name}')
self.task_cmd_builder.pretrained_model_id = (
test_util.get_pretrained_model_id(model_name)
)
self.task_cmd_builder.config_file = (
'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
)
self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
self.task_cmd_builder.train_split_name = 'train'
self.task_cmd_builder.instruct_column = 'text'
self.task_cmd_builder.template = 'openassistant-guanaco'
self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
self.task_cmd_builder.eval_split_name = 'test'
self.task_cmd_builder.eval_instruct_column = (
self.task_cmd_builder.instruct_column
)
self.task_cmd_builder.eval_template = self.task_cmd_builder.template
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
self.assertEqual(self.run_cmd(), 0)
if __name__ == '__main__':
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/test_quantize_model.py
+49
View File
@@ -0,0 +1,49 @@
# pylint: disable=missing-function-docstring
# pylint: disable=missing-class-docstring
"""Tests quantize model task in PEFT docker."""
import os
import time
from absl.testing import absltest
import quantize_model_command_builder as task_cmd_builder
import test_util
class QuantizeModelTest(test_util.TestBase):
def setUp(self):
super().setUp()
self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '')
self.docker_builder.add_mount_map(
os.path.expanduser('~'), os.path.expanduser('~')
)
self.task_cmd_builder = task_cmd_builder.QuantizeModelCommandBuilder()
self.task_cmd_builder.task = 'quantize-model'
self.task_cmd_builder.pretrained_model_id = (
'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf'
)
self.task_cmd_builder.quantization_method = 'awq'
self.task_cmd_builder.quantization_precision_mode = '4bit'
self.task_cmd_builder.quantization_dataset_name = 'pileval'
self.task_cmd_builder.text_column_in_quantization_dataset = 'text'
self.task_cmd_builder.quantization_output_dir = '~/llama3-8b-hf-quantized'
self.task_cmd_builder.device_map = None
self.task_cmd_builder.max_memory = None
self.task_cmd_builder.group_size = 128
self.task_cmd_builder.desc_act = False
self.task_cmd_builder.damp_percent = 0.1
self.task_cmd_builder.cache_examples_on_gpu = False
self.task_cmd_builder.awq_version = 'GEMM'
def test_llama3_8b_model_awq_quantization(self):
start_time = time.time()
self.assertEqual(self.run_cmd(), 0)
end_time = time.time()
self.assertLess(end_time - start_time, 1.5 * 60 * 60)
if __name__ == '__main__':
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/test_util.py
+134
View File
@@ -0,0 +1,134 @@
"""Test util class."""
import datetime
import os
import signal
import subprocess
import sys
from absl import flags
from absl import logging
from absl.testing import parameterized
import docker_command_builder as docker_cmd_builder
_DOCKER_URI = flags.DEFINE_string(
'docker_uri', None, 'docker image uri', required=True
)
_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://peft-docker-test',
'GCS directory that stores model checkpoint, dataset and etc.',
)
_GCS_OUTPUT_DIR = flags.DEFINE_string(
'gcs_output_dir',
'gs://peft-docker-test/output',
'GCS directory that stores test output.',
)
class TestBase(parameterized.TestCase):
"""Test base class that defines how to run commands."""
def setUp(self):
super().setUp()
self.docker_builder = docker_cmd_builder.DockerCommandBuilder(
_DOCKER_URI.value
)
self.docker_builder.add_mount_map(
os.path.expanduser('~'), os.path.expanduser('~')
)
self.docker_builder.add_mount_map(
self.local_input_dir(), self.local_input_dir()
)
self.task_cmd_builder = None
def cmd(self):
return self.docker_builder.build_cmd() + self.task_cmd_builder.build_cmd()
def run_cmd(self) -> int:
logging.info('running command: \n%s', ' \\\n'.join(self.cmd()))
if _DRY_RUN.value:
return 0
p = subprocess.Popen(self.cmd(), stdout=sys.stdout, stderr=sys.stderr)
try:
unused_output, unused_error = p.communicate()
return p.returncode
except KeyboardInterrupt:
p.send_signal(signal.SIGINT)
return 0
def gcs_output_dir(self):
return _GCS_OUTPUT_DIR.value
def local_output_dir(self):
return _LOCAL_OUTPUT_DIR.value
def local_input_dir(self):
return _LOCAL_INPUT_DIR.value
def get_timestamp():
return datetime.datetime.now(datetime.timezone.utc).strftime(
'%Y%m%d_%H%M%S%Z'
)
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.
"""
def _download_from_gcs(name):
if not os.path.exists(_LOCAL_INPUT_DIR.value):
os.mkdir(_LOCAL_INPUT_DIR.value)
subprocess.check_output([
'gsutil',
'-m',
'cp',
'-r',
os.path.join(_GCS_INPUT_DIR.value, name),
_LOCAL_INPUT_DIR.value,
])
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):
_download_from_gcs(name)
return local_data
def get_pretrained_model_id(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)
community-content/vertex_model_garden/model_oss/peft/train/tests/test_validate_dataset_with_template.py
+327
View File
@@ -0,0 +1,327 @@
"""Tests validate the dataset with template task in PEFT docker."""
from absl.testing import absltest
from absl.testing import parameterized
import test_util
import validate_dataset_with_template_command_builder as task_cmd_builder
class ValidateDatasetWithTemplateTest(test_util.TestBase):
"""Test the validate dataset with template task in PEFT docker."""
def setUp(self):
super().setUp()
self.task_cmd_builder = (
task_cmd_builder.ValidateDatasetWithTemplateCommandBuilder()
)
self.task_cmd_builder.task = "validate-dataset-with-template"
@parameterized.named_parameters(
dict(
testcase_name="valid_rows",
validate_top_k_rows=100,
expected_result=0,
),
dict(
testcase_name="negative_rows",
validate_top_k_rows=-10,
expected_result=0,
),
dict(
testcase_name="out_of_range_rows",
validate_top_k_rows=100000,
expected_result=1,
),
)
def test_validate_dataset_with_template_top_k_rows(
self,
validate_top_k_rows,
expected_result,
):
self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
self.task_cmd_builder.train_split_name = "train"
self.task_cmd_builder.instruct_column_in_dataset = "text"
self.task_cmd_builder.template = (
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
)
self.task_cmd_builder.validate_percentage_of_dataset = None
self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
self.task_cmd_builder.use_multiprocessing = True
result = self.run_cmd()
self.assertEqual(result, expected_result)
@parameterized.named_parameters(
dict(
testcase_name="valid_positive_x_percent",
validate_percentage_of_dataset=10,
expected_result=0,
),
dict(
testcase_name="valid_negative_x_percent",
validate_percentage_of_dataset=-10,
expected_result=0,
),
dict(
testcase_name="invalid_positive_x_percent",
validate_percentage_of_dataset=110,
expected_result=1,
),
dict(
testcase_name="invalid_negative_x_percent",
validate_percentage_of_dataset=-110,
expected_result=1,
),
)
def test_validate_dataset_with_template_x_percent(
self,
validate_percentage_of_dataset,
expected_result,
):
self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
self.task_cmd_builder.train_split_name = "train"
self.task_cmd_builder.instruct_column_in_dataset = "text"
self.task_cmd_builder.template = (
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
)
self.task_cmd_builder.validate_percentage_of_dataset = (
validate_percentage_of_dataset
)
self.task_cmd_builder.validate_k_rows_of_dataset = None
self.task_cmd_builder.use_multiprocessing = True
result = self.run_cmd()
self.assertEqual(result, expected_result)
@parameterized.named_parameters(
dict(
testcase_name="invalid_default_input_column",
dataset_name="timdettmers/openassistant-guanaco",
split="train",
input_column="",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="invalid_percentage",
dataset_name="timdettmers/openassistant-guanaco",
split="train",
input_column="text",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=110,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="negative_percentage",
dataset_name="timdettmers/openassistant-guanaco",
split="train",
input_column="text",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=-110,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="empty_dataset",
dataset_name="",
split="train",
input_column="text",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="empty_split",
dataset_name="timdettmers/openassistant-guanaco",
split="",
input_column="text",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="empty_template",
dataset_name="timdettmers/openassistant-guanaco",
split="train",
input_column="text",
template="",
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="wrong_gcs_template",
dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
split="train",
input_column="text",
template="gs://cloud-nas-260507-tmp-20240724/sample-template.json",
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
dict(
testcase_name="wrong_gcs_dataset_name",
dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft-train_sample.jsonl",
split="train",
input_column="text",
template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=1,
),
)
def test_validate_dataset_with_template_invalid_input(
self,
dataset_name,
split,
input_column,
template,
validate_percentage_of_dataset,
validate_top_k_rows,
use_multiprocessing,
expected_result,
):
self.task_cmd_builder.dataset_name = dataset_name
self.task_cmd_builder.train_split_name = split
self.task_cmd_builder.instruct_column_in_dataset = input_column
self.task_cmd_builder.template = template
self.task_cmd_builder.validate_percentage_of_dataset = (
validate_percentage_of_dataset
)
self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
self.task_cmd_builder.use_multiprocessing = use_multiprocessing
result = self.run_cmd()
self.assertEqual(result, expected_result)
@parameterized.named_parameters(
dict(
testcase_name="full_hf_dataset_with_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=0,
),
dict(
testcase_name="full_gcs_dataset_with_multiprocessing",
dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=0,
),
dict(
testcase_name="half_dataset_with_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=50,
validate_top_k_rows=None,
use_multiprocessing=True,
expected_result=0,
),
dict(
testcase_name="top_100_rows_with_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=100,
use_multiprocessing=True,
expected_result=0,
),
dict(
testcase_name="full_gcs_dataset_without_multiprocessing",
dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=False,
expected_result=0,
),
dict(
testcase_name="full_hf_dataset_without_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=None,
use_multiprocessing=False,
expected_result=0,
),
dict(
testcase_name="half_dataset_without_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=50,
validate_top_k_rows=None,
use_multiprocessing=False,
expected_result=0,
),
dict(
testcase_name="top_100_rows_without_multiprocessing",
dataset_name="timdettmers/openassistant-guanaco",
template=(
"gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
),
validate_percentage_of_dataset=None,
validate_top_k_rows=100,
use_multiprocessing=True,
expected_result=0,
),
)
def test_validate_dataset_with_template_multiprocessing_option(
self,
dataset_name,
template,
validate_percentage_of_dataset,
validate_top_k_rows,
use_multiprocessing,
expected_result,
):
self.task_cmd_builder.dataset_name = dataset_name
self.task_cmd_builder.train_split_name = "train"
self.task_cmd_builder.instruct_column_in_dataset = "text"
self.task_cmd_builder.template = template
self.task_cmd_builder.validate_percentage_of_dataset = (
validate_percentage_of_dataset
)
self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
self.task_cmd_builder.use_multiprocessing = use_multiprocessing
result = self.run_cmd()
self.assertEqual(result, expected_result)
if __name__ == "__main__":
absltest.main()
community-content/vertex_model_garden/model_oss/peft/train/tests/tools/command_builder_gen.py
+109
View File
@@ -0,0 +1,109 @@
"""Tools to generate CommandBuilder class.
See go/vmg-oss-peft-tests#commandbuilder-class-generation for details.
"""
import argparse
import dataclasses
from typing import List
_DO_NOT_MODIFY_WARNING = """
# DO NOT MODIFY: this file is auto-generated
# See go/vmg-oss-peft-tests#command-builder-genpy
"""
_GETTER_TMPL = """
@property
def {}(self):
return self._{}
"""
_SETTER_TMPL = """
@{}.setter
def {}(self, val: {}):
self._{} = val
"""
_INIT_NAME = """
def __init__(self):"""
_INIT_FIELDS = """
self._{} = None"""
_BUILD_CMD = r"""
def build_cmd(self) -> str:
cmd = []
for k, v in self.__dict__.items():
if v is not None:
cmd.append(f'--{k[1:]}={v}')
return cmd
"""
@dataclasses.dataclass
class FlagInfo:
api_name: str
impl_name: str
arg_type: str
def get_flag_info(line: str) -> FlagInfo:
api_name, impl_name, arg_type = [x.strip() for x in line.split(',')]
return FlagInfo(api_name, impl_name, arg_type)
def gen_getter(info: FlagInfo) -> str:
return _GETTER_TMPL.format(info.api_name, info.impl_name)
def gen_setter(info: FlagInfo) -> str:
return _SETTER_TMPL.format(
info.api_name, info.api_name, info.arg_type, info.impl_name
)
def gen_init(infos: List[FlagInfo]) -> str:
fields = [_INIT_FIELDS.format(i.impl_name) for i in infos]
return ''.join([_INIT_NAME] + fields)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--flags_def', required=True, help='file path contain flags definition.'
)
parser.add_argument(
'--generated_file',
required=True,
help='file path to the generated command builder.',
)
parser.add_argument(
'--class_name',
required=True,
help='class name for command build',
)
args = parser.parse_args()
flags_info = []
with open(args.flags_def, 'r') as flags_f:
for line in flags_f:
if not line.startswith('#'):
flags_info.append(get_flag_info(line))
with open(args.generated_file, 'w') as gen_f:
# Disables pylint messages.
# See https://stackoverflow.com/a/43510297
print('# pylint: disable=W,C,R', file=gen_f)
print(_DO_NOT_MODIFY_WARNING, file=gen_f)
print(f'class {args.class_name}:', file=gen_f)
print(gen_init(flags_info), file=gen_f)
for info in flags_info:
print(gen_getter(info), file=gen_f)
print(gen_setter(info), file=gen_f)
print(_BUILD_CMD, file=gen_f)
print(f'file generated at {args.generated_file}')
if __name__ == '__main__':
main()
community-content/vertex_model_garden/model_oss/peft/train/tests/tools/instruct_lora_flags.txt
+50
View File
@@ -0,0 +1,50 @@
# api_name, impl_name, value_type
# eval related and etc.
config_file, config_file, str
task, task, str
pretrained_model_id, pretrained_model_id, str
train_dataset, dataset_name, str
train_split_name, train_split_name, str
template, template, str
instruct_column, instruct_column_in_dataset, str
ckpt_dir, output_dir, str
merged_model_dir, merge_base_and_lora_output_dir, str
logging_dir, logging_output_dir, str
per_device_batch_size, per_device_train_batch_size, int
gradient_accumulation_steps, gradient_accumulation_steps, int
lora_rank, lora_rank, int
lora_alpha, lora_alpha, int
lora_dropout, lora_dropout, float
max_steps, max_steps, int
num_epochs, num_epochs, float
max_seq_length, max_seq_length, int
learning_rate, learning_rate, float
lr_scheduler_type, lr_scheduler_type, str
load_precision, precision_mode, str
train_precision, train_precision, str
gradient_checkpointing, enable_gradient_checkpointing, bool
example_packing, use_example_packing, bool
attn_implementation, attn_implementation, str
optimizer, optimizer, str
warmup_ratio, warmup_ratio, float
report_to, report_to, str
save_steps, save_steps, int
logging_steps, logging_steps, int
huggingface_access_token, huggingface_access_token, str
eval_dataset, eval_dataset_path, str
eval_instruct_column, eval_column, str
eval_template, eval_template, str
eval_split_name, eval_split, str
eval_steps, eval_steps, int
eval_tasks, eval_tasks, str
eval_metric_name, eval_metric_name, str
completion_only, completion_only, bool
max_grad_norm, max_grad_norm, float
logger_level, logger_level, str
benchmark_out_file, benchmark_out_file, str
tuning_data_stats_file, tuning_data_stats_file, str
enable_peft, enable_peft, bool
merge_model_precision_mode, merge_model_precision_mode, str
target_modules, target_modules, str
community-content/vertex_model_garden/model_oss/peft/train/tests/tools/quantize_model_flags.txt
+15
View File
@@ -0,0 +1,15 @@
# api_name, impl_name, value_type
task, task, str
pretrained_model_id, pretrained_model_id, str
quantization_method, quantization_method, str
quantization_precision_mode, quantization_precision_mode, str
quantization_dataset_name, quantization_dataset_name, str
text_column_in_quantization_dataset, text_column_in_quantization_dataset, str
quantization_output_dir, quantization_output_dir, str
device_map, device_map, str
max_memory, max_memory, str
group_size, group_size, int
desc_act, desc_act, bool
damp_percent, damp_percent, float
cache_examples_on_gpu, cache_examples_on_gpu, bool
awq_version, awq_version, str
community-content/vertex_model_garden/model_oss/peft/train/tests/tools/validate_dataset_with_template_flags.txt
+9
View File
@@ -0,0 +1,9 @@
# api_name, impl_name, value_type
task, task, str
template, template, str
dataset_name, dataset_name, str
train_split_name, train_split_name, str
instruct_column_in_dataset, instruct_column_in_dataset, str
use_multiprocessing, use_multiprocessing, bool
validate_k_rows_of_dataset, validate_k_rows_of_dataset, int
validate_percentage_of_dataset, validate_percentage_of_dataset, int
community-content/vertex_model_garden/model_oss/peft/train/tests/validate_dataset_with_template_command_builder.py
+88
View File
@@ -0,0 +1,88 @@
# pylint: disable=W,C,R
# DO NOT MODIFY: this file is auto-generated
# See go/vmg-oss-peft-tests#command-builder-genpy
class ValidateDatasetWithTemplateCommandBuilder:
def __init__(self):
self._task = None
self._template = None
self._dataset_name = None
self._train_split_name = None
self._instruct_column_in_dataset = None
self._use_multiprocessing = None
self._validate_k_rows_of_dataset = None
self._validate_percentage_of_dataset = None
@property
def task(self):
return self._task
@task.setter
def task(self, val: str):
self._task = val
@property
def template(self):
return self._template
@template.setter
def template(self, val: str):
self._template = val
@property
def dataset_name(self):
return self._dataset_name
@dataset_name.setter
def dataset_name(self, val: str):
self._dataset_name = val
@property
def train_split_name(self):
return self._train_split_name
@train_split_name.setter
def train_split_name(self, val: str):
self._train_split_name = val
@property
def instruct_column_in_dataset(self):
return self._instruct_column_in_dataset
@instruct_column_in_dataset.setter
def instruct_column_in_dataset(self, val: str):
self._instruct_column_in_dataset = val
@property
def use_multiprocessing(self):
return self._use_multiprocessing
@use_multiprocessing.setter
def use_multiprocessing(self, val: bool):
self._use_multiprocessing = val
@property
def validate_k_rows_of_dataset(self):
return self._validate_k_rows_of_dataset
@validate_k_rows_of_dataset.setter
def validate_k_rows_of_dataset(self, val: int):
self._validate_k_rows_of_dataset = val
@property
def validate_percentage_of_dataset(self):
return self._validate_percentage_of_dataset
@validate_percentage_of_dataset.setter
def validate_percentage_of_dataset(self, val: int):
self._validate_percentage_of_dataset = val
def build_cmd(self) -> str:
cmd = []
for k, v in self.__dict__.items():
if v is not None:
cmd.append(f'--{k[1:]}={v}')
return cmd
community-content/vertex_model_garden/model_oss/peft/train/vmg/callbacks.py
+86
View File
@@ -0,0 +1,86 @@
"""Different trainer callbacks for PEFT Trainer."""
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 vertex_vision_model_garden_peft.train.vmg import utils
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_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()
delta_t = float('nan')
else:
cur_time = time.time()
delta_t = cur_time - self._prev_time
self._prev_time = cur_time
self._avg_throughput += (delta_t - self._avg_throughput) / (
state.global_step - 1
)
gpu_stats = utils.gpu_stats()
self._peak_mem = max(gpu_stats.total_mem, self._peak_mem)
logging.info(
'on_step_end: %s, throughput: %.2f s/it',
utils.gpu_stats_str(gpu_stats),
delta_t,
)
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', utils.gpu_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
logging.info(
'training time %.2f s, throughput: %.2f s/it, peak_mem: %.2f GB',
train_time,
self._avg_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}k | {self._peak_mem:.2f} |'
f' {self._avg_throughput:.2f}\n'
)
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/custom_loglikelihood.yaml
+18
View File
@@ -0,0 +1,18 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/ddp_4gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/ddp_8gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero2_4gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero2_8gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero3_4gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero3_8gpu.yaml
+17
View File
@@ -0,0 +1,17 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_128gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_16gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: 2
num_processes: 16
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_24gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: 3
num_processes: 24
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_32gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: 4
num_processes: 32
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_8gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_16gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: HYBRID_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: 2
num_processes: 16
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_24gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: HYBRID_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: 3
num_processes: 24
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_32gpu.yaml
+28
View File
@@ -0,0 +1,28 @@
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: HYBRID_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: 4
num_processes: 32
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/zero2.json
+24
View File
@@ -0,0 +1,24 @@
{
"zero_optimization": {
"stage": 2,
"contiguous_gradients": false,
"overlap_comm": false
},
"bf16": {
"enabled": "auto"
},
"fp16": {
"enabled": "auto",
"auto_cast": false,
"loss_scale": 0,
"initial_scale_power": 32,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/zero3.json
+31
View File
@@ -0,0 +1,31 @@
{
"zero_optimization": {
"stage": 3,
"overlap_comm": false,
"contiguous_gradients": false,
"sub_group_size": 0,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 0,
"stage3_max_reuse_distance": 0,
"stage3_gather_16bit_weights_on_model_save": true
},
"bf16": {
"enabled": "auto"
},
"fp16": {
"enabled": "auto",
"auto_cast": false,
"loss_scale": 0,
"initial_scale_power": 32,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/merge_env.yaml
+43
View File
@@ -0,0 +1,43 @@
# Doc about format of conda environment file
# https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#create-env-file-manually
name: merge
channels:
- nodefaults
- conda-forge
dependencies:
- _libgcc_mutex=0.1=conda_forge
- _openmp_mutex=4.5=2_gnu
- bzip2=1.0.8=h4bc722e_7
- ca-certificates=2024.7.4=hbcca054_0
- ld_impl_linux-64=2.40=hf3520f5_7
- libffi=3.4.2=h7f98852_5
- libgcc-ng=14.1.0=h77fa898_0
- libgomp=14.1.0=h77fa898_0
- libnsl=2.0.1=hd590300_0
- libsqlite=3.46.0=hde9e2c9_0
- libuuid=2.38.1=h0b41bf4_0
- libxcrypt=4.4.36=hd590300_1
- libzlib=1.3.1=h4ab18f5_1
- ncurses=6.5=h59595ed_0
- openssl=3.3.1=h4bc722e_2
- pip=24.2=pyhd8ed1ab_0
- python=3.10.14=hd12c33a_0_cpython
- readline=8.2=h8228510_1
- setuptools=72.1.0=pyhd8ed1ab_0
- tk=8.6.13=noxft_h4845f30_101
- tzdata=2024a=h0c530f3_0
- wheel=0.44.0=pyhd8ed1ab_0
- xz=5.2.6=h166bdaf_0
- pip:
- --extra-index-url https://download.pytorch.org/whl/cu121
- absl-py==2.1.0
- accelerate==0.33.0 # Needed for fp8
- datasets==2.19.2
- fbgemm-gpu==0.8.0+cu121 # Needed for fp8
- kfp==2.5.0
- peft==0.12.0
- protobuf==3.20.3
- pynvml==11.5.3
- torch==2.4.0+cu121 # Needed for fp8
- transformers==4.43.1
- trl==0.9.6
community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/requirements.txt
+27
View File
@@ -0,0 +1,27 @@
# Doc about format of requirement file
# https://pip.pypa.io/en/stable/reference/requirements-file-format
--extra-index-url https://download.pytorch.org/whl/cu118
--extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
# keep sorted
accelerate==0.31.0
auto_gptq==0.7.1+cu118
autoawq==0.2.5
bitsandbytes==0.43.2
cloudml-hypertune==0.1.0.dev6
datasets==2.19.2
deepspeed==0.14.4
diffusers==0.25.1
fsspec==2024.3.1
gcsfs==2024.3.1
lm_eval==0.4.3
ninja==1.11.1 # Needed to avoid `ninja 1.11.1.1 is not supported on this platform` error
optimum==1.17.1
peft==0.12.0
pynvml==11.5.3
torch==2.2.2+cu118
torchvision==0.17.2+cu118
transformers==4.43.1
trl==0.9.6
wandb==0.17.1
community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/train.Dockerfile
+76
View File
@@ -0,0 +1,76 @@
# 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}
# Picked from https://cloud.google.com/deep-learning-containers/docs/choosing-container#pytorch
FROM us-docker.pkg.dev/deeplearning-platform-release/gcr.io/pytorch-cu121.2-2.py310:m123
RUN apt-get update && \
apt-get upgrade -y && \
apt-get install -y curl git wget software-properties-common vim libaio-dev && \
apt-get clean && \
rm -rf /var/lib/apt/lists*
# Copy license.
RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
# Install libraries.
ENV PIP_ROOT_USER_ACTION=ignore
RUN pip install --upgrade pip
# Prefer to install with requirement file as much as possible for reasons
# described in b/355034754.
COPY model_oss/peft/train/vmg/dockerfile/requirements.txt /tmp/requirements.txt
RUN pip install -r /tmp/requirements.txt
# flash-attn cannot be installed with the requirement file approach above
# because of the `no-build-isolation` requirement.
#
# It is OK to install it after other packages FOR NOW because it only has
# limited dependencies. And there's no concern about it overwriting previously
# installed packages.
# https://github.com/Dao-AILab/flash-attention/blob/v2.6.3/setup.py#L523
RUN pip install flash-attn==2.6.3 --no-build-isolation
# Install `diffusers` library as editable and in root folder (/) on purpose.
RUN git clone --depth 1 --branch v0.25.1 https://github.com/huggingface/diffusers.git
# Remove `diffusers` (NOTE that the dependency libraries are kept).
RUN pip uninstall -y diffusers
# Using `--no-deps` option to make sure previously installed packages are not
# overwritten.
RUN pip install --no-deps -e /diffusers
# Make sure there's no inconsistent pip libraries.
RUN pip check
# Install merge related packages in a separate env.
COPY model_oss/peft/train/vmg/dockerfile/merge_env.yaml /tmp/merge_env.yaml
RUN conda env create -n merge --yes --file /tmp/merge_env.yaml
RUN conda init
# Switch to diffusers examples folder.
WORKDIR /diffusers/examples
RUN mkdir -p ./vertex_vision_model_garden_peft/
COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
# custom `lm_eval` task.
ARG LM_EVAL_DIR=$(python -c 'import site; print(site.getsitepackages()[0])')/lm_eval
RUN mkdir -p $LM_EVAL_DIR/tasks/vertex && \
mv ./vertex_vision_model_garden_peft/custom_loglikelihood.yaml $LM_EVAL_DIR/tasks/vertex/
COPY model_oss/peft/train/vmg/*.py ./vertex_vision_model_garden_peft/train/vmg/
COPY model_oss/peft/train/vmg/templates /diffusers/examples/util/templates
COPY model_oss/util /diffusers/examples/util
COPY model_oss/notebook_util/dataset_validation_util.py /diffusers/examples/util
COPY model_oss/peft/train/tests/*.py ./vertex_vision_model_garden_peft/tests/
RUN chmod a+rwX -R /diffusers/examples/
ENV PYTHONPATH /diffusers/examples/
# Must disable torch XLA, otherwise runtime uses CPU even if GPU exists.
ENV USE_TORCH_XLA 0
ENTRYPOINT ["python3", "./vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py"]
community-content/vertex_model_garden/model_oss/peft/train/vmg/eval_lib.py
+183
View File
@@ -0,0 +1,183 @@
"""Library for running evaluations during training."""
import dataclasses
from typing import Any, Optional, Type
from absl import logging
import datasets
from lm_eval import evaluator
from lm_eval import tasks
from lm_eval import utils
from lm_eval.api import model as lm_model
from lm_eval.api import registry
from lm_eval.models import huggingface
from peft import peft_model
import transformers
from transformers import trainer
from util import dataset_validation_util
from util import constants
_DESCRIPTION_EVALUATION = "evaluation"
_BUILTIN_EVAL_TASK = "builtin_eval"
@dataclasses.dataclass(frozen=True)
class EvalConfig:
steps: int
tasks: list[str]
per_device_batch_size: int
num_fewshot: Optional[int]
limit: Optional[float]
metric_name: str
tokenize_dataset: bool
dataset_path: str = ""
split: str = "test"
template: str = ""
column: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET
class PeftCausalLMModel(huggingface.HFLM):
"""PeftCausalLMModel that supports loading an in-memory model."""
AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
def __init__(
self,
model: peft_model.PeftModelForCausalLM,
tokenizer: transformers.PreTrainedTokenizerBase,
batch_size_per_gpu: int,
):
lm_model.LM.__init__(self)
self._model = model
self.tokenizer = tokenizer
self.vocab_size = tokenizer.vocab_size
tokenizer.pad_token_id = tokenizer.eos_token_id
self._config = model.config
self.batch_size_per_gpu = batch_size_per_gpu
self._device = model.device
self._max_length = None # Will be automatically determined from config.
self._add_special_tokens = (
None # Will be automatically determined from AUTO_MODEL_CLASS.
)
def create_trainer(
cls: Type[transformers.Trainer],
eval_config: Optional[EvalConfig],
tokenizer: Optional[transformers.PreTrainedTokenizerBase],
args: trainer.TrainingArguments,
**kwargs,
) -> transformers.Trainer:
"""Creates a trainer. If eval config is provided, injects evaluation loop."""
if not eval_config:
return cls(args=args, **kwargs)
args.eval_strategy = "steps"
args.eval_steps = eval_config.steps
args.per_device_eval_batch_size = eval_config.per_device_batch_size
kwargs["tokenizer"] = tokenizer
if eval_config.tasks == [_BUILTIN_EVAL_TASK]:
try:
eval_dataset = dataset_validation_util.load_dataset_with_template(
dataset_name=eval_config.dataset_path,
split=eval_config.split,
input_column=eval_config.column,
template=eval_config.template,
tokenizer=tokenizer,
)
if eval_config.limit is not None:
if eval_config.limit >= 1:
limit = int(eval_config.limit)
else:
limit = int(eval_config.limit * len(eval_dataset))
eval_dataset = eval_dataset.select(range(limit))
if eval_config.tokenize_dataset:
eval_dataset = eval_dataset.map(
lambda samples: tokenizer(samples[eval_config.column])
)
kwargs["eval_dataset"] = eval_dataset
except (OSError, ValueError, IndexError) as e:
logging.warning(
"Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
eval_config.dataset_path,
e,
)
del args.evaluation_strategy
del args.eval_steps
del args.per_device_eval_batch_size
return cls(args=args, **kwargs)
class LMEvalTrainer(cls):
"""Trainer with lm_eval injected as the eval library."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
task_names = utils.pattern_match(eval_config.tasks, registry.ALL_TASKS)
logging.info("Selected Eval Tasks: %s", task_names)
task_args = {}
if eval_config.num_fewshot is not None:
task_args["num_fewshot"] = eval_config.num_fewshot
if eval_config.dataset_path:
task_args["dataset_path"] = "json"
task_args["dataset_kwargs"] = {
"data_files": {"test": eval_config.dataset_path},
}
self._eval_task_dict = tasks.get_task_dict(task_names, **task_args)
def evaluation_loop(
self,
dataloader: trainer.DataLoader,
description: str,
prediction_loss_only: Optional[bool] = None,
ignore_keys: Optional[list[str]] = None,
metric_key_prefix: str = "eval",
) -> trainer.EvalLoopOutput:
"""Custom evaluation loop that invokes lm_eval."""
if description.lower() != _DESCRIPTION_EVALUATION:
return super().evaluation_loop(
dataloader,
description,
prediction_loss_only,
ignore_keys,
metric_key_prefix,
)
model = self._wrap_model(self.model, training=False)
lm = PeftCausalLMModel(
model,
self.tokenizer or self.data_collator.tokenizer,
eval_config.per_device_batch_size,
)
results: dict[str, Any] = evaluator.evaluate(
lm=lm,
task_dict=self._eval_task_dict,
limit=eval_config.limit,
)["results"]
metric_name = eval_config.metric_name
# Compute average value if there are multiple tasks.
metric_values: list[float] = []
for result in results.values():
for key, value in result.items():
if key.split(",")[0] == metric_name:
metric_values.append(value)
if not metric_values:
raise ValueError(
f"Metric {metric_name} not found in eval response: {results}"
)
metric_average = sum(metric_values) / len(metric_values)
logging.info("%s value: %f\n%s", metric_name, metric_average, results)
return trainer.EvalLoopOutput(
# Only metrics field is set. Other fields are dummy values.
predictions=None,
label_ids=None,
metrics={f"{metric_key_prefix}_{metric_name}": metric_average},
num_samples=0,
)
# Use empty eval dataset as a placeholder.
return LMEvalTrainer(
args=args, eval_dataset=datasets.Dataset.from_dict({"test": []}), **kwargs
)
community-content/vertex_model_garden/model_oss/peft/train/vmg/instruct_lora.py
+794
View File
@@ -0,0 +1,794 @@
"""Instruct/Chat with LoRA models."""
import dataclasses
import datetime
import json
import os
from typing import Any, Dict, Optional, Sequence
import warnings
from absl import app
from absl import flags
from absl import logging
from accelerate import DistributedType
from accelerate import PartialState
import bitsandbytes as bnb
import hypertune
from peft import get_peft_model
from peft import LoraConfig
import torch
from transformers import AutoModelForCausalLM
from transformers import TrainingArguments
from trl import DataCollatorForCompletionOnlyLM
from trl import SFTTrainer
import wandb
from util import dataset_validation_util
from vertex_vision_model_garden_peft.train.vmg import callbacks
from vertex_vision_model_garden_peft.train.vmg import eval_lib
from vertex_vision_model_garden_peft.train.vmg import utils
from util import constants
from util import fileutils
_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. Note, there'
' might be different paddings for different models. This tool assumes the'
' pretrained_model_id contains model name, and then choose proper padding'
' methods. e.g. it must contain `llama` for `Llama2 models`.',
required=True,
)
_HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
'huggingface_access_token',
None,
'The access token for loading huggingface gated models.',
)
_DATASET_NAME = flags.DEFINE_string(
'dataset_name',
None,
'The dataset name in huggingface.',
)
_OUTPUT_DIR = flags.DEFINE_string(
'output_dir',
None,
'The output directory.',
)
_LOGGING_OUTPUT_DIR = flags.DEFINE_string(
'logging_output_dir',
'',
'The logging output directory, which defaults to same as output_dir.',
)
_PRECISION_MODE = flags.DEFINE_enum(
'precision_mode',
constants.PRECISION_MODE_16,
[
constants.PRECISION_MODE_4,
constants.PRECISION_MODE_8,
constants.PRECISION_MODE_16,
constants.PRECISION_MODE_16B,
constants.PRECISION_MODE_32,
],
'Precision to load model weights for finetuning.',
)
_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.',
)
_WEIGHT_DECAY = flags.DEFINE_float(
'weight_decay',
0.001,
'The weight decay in the learning rate scheduler.',
)
_NUM_EPOCHS = flags.DEFINE_float(
'num_epochs',
None,
'The number of training epochs. Only used for'
' "sequence-classification-lora" with an integer value and for'
' "instruct-lora" with a float value allowed.',
)
_MAX_STEPS = flags.DEFINE_integer(
'max_steps',
None,
'Total number of training steps. Overrides num_epochs if set. Only used for'
' "instruct-lora."',
)
_MAX_SEQ_LENGTH = flags.DEFINE_integer(
'max_seq_length',
512,
'The maximum sequence length.',
)
_LEARNING_RATE = flags.DEFINE_float(
'learning_rate',
2e-4,
'The learning rate after the potential warmup period.',
)
_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
'instruct_column_in_dataset',
constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
'The instruct column in dataset.',
)
_REPORT_TO = flags.DEFINE_string(
'report_to',
constants.REPORT_TO_NONE,
'Where logging is reported to, which can be tensorboard or none.',
)
_PER_DEVICE_TRAIN_BATCH_SIZE = flags.DEFINE_integer(
'per_device_train_batch_size',
4,
'The per device train batch size.',
)
_GRADIENT_ACCUMULATION_STEPS = flags.DEFINE_integer(
'gradient_accumulation_steps',
4,
'The gradient accumulation steps.',
)
_ENABLE_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
'enable_gradient_checkpointing',
False,
'Whether to enable gradient checkpointing.',
)
_ENABLE_PEFT = flags.DEFINE_boolean(
'enable_peft',
True,
'Whether to enable peft.',
)
_TEMPLATE = flags.DEFINE_string(
'template',
None,
'Template for formatting language model training data. Must be a filename'
' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
' Cloud Storage URI to a JSON file.',
)
_OPTIMIZER = flags.DEFINE_string(
'optimizer',
'adamw_torch',
'The optimizer.',
)
_LR_SCHEDULER_TYPE = flags.DEFINE_string(
'lr_scheduler_type',
'cosine',
'The learning rate scheduler type.',
)
_SAVE_STEPS = flags.DEFINE_integer(
'save_steps',
10,
'The save steps.',
)
_LOGGING_STEPS = flags.DEFINE_integer(
'logging_steps',
10,
'The logging steps.',
)
_EVAL_STEPS = flags.DEFINE_integer(
'eval_steps',
10,
'The number of training steps between evaluations.',
)
_TRAIN_SPLIT_NAME = flags.DEFINE_string(
'train_split_name',
'train',
'The train split name.',
)
_EVAL_TASKS = flags.DEFINE_list(
'eval_tasks',
None,
'List of eval task names (can have wildcards) as in'
' https://github.com/EleutherAI/lm-evaluation-harness. Will not run'
' evaluation if not set. Runs the built-in trainer evaluation loop if set'
' to `builtin_eval`.',
)
_EVAL_PER_DEVICE_BATCH_SIZE = flags.DEFINE_integer(
'eval_per_device_batch_size',
1,
'The per device batch size for model evaluation.',
)
_EVAL_NUM_FEWSHOT = flags.DEFINE_integer(
'eval_num_fewshot',
None,
'Run N-shot language model evaluation. Not implemented in `builtin_eval`.',
)
_EVAL_LIMIT = flags.DEFINE_float(
'eval_limit',
None,
'Limit the number of examples per task. If <1, limit is a percentage of the'
' total number of examples.',
)
_EVAL_METRIC_NAME = flags.DEFINE_string(
'eval_metric_name',
'acc',
'The metric name to aggregate during model evaluation.',
)
_EVAL_DATASET_PATH = flags.DEFINE_string(
'eval_dataset_path',
None,
'Overrides the default evaluation dataset path. In `builtin_eval` mode,'
' this can be any Hugging Face dataset name.',
)
# We set the default eval split as `test`, based on observation from
# https://huggingface.co/datasets/timdettmers/openassistant-guanaco/viewer/default/test.
_EVAL_SPLIT = flags.DEFINE_string(
'eval_split',
'test',
'Eval split name in the eval dataset for `builtin_eval`.',
)
_EVAL_TEMPLATE = flags.DEFINE_string(
'eval_template',
None,
'Template for formatting language model evaluation data for `builtin_eval`.'
' Must be a filename under `templates` folder, without `.json` extension,'
' e.g. `alpaca`, or a Cloud Storage URI to a JSON file.',
)
_EVAL_COLUMN = flags.DEFINE_string(
'eval_column',
None,
'Eval column name in the eval dataset for `builtin_eval`.',
)
_TRAIN_PRECISION = flags.DEFINE_enum(
'train_precision',
constants.PRECISION_MODE_16B,
[
constants.PRECISION_MODE_16,
constants.PRECISION_MODE_16B,
constants.PRECISION_MODE_32,
],
'Precision to train the model.',
)
_USE_EXAMPLE_PACKING = flags.DEFINE_boolean(
'use_example_packing',
False,
'Enables example packing during training, which uses '
'`ConstantLengthDataset` under the hood.',
)
_COMPLETION_ONLY = flags.DEFINE_boolean(
'completion_only',
False,
'If set, it uses DataCollatorForCompletionOnlyLM to train the model on the'
' generated prompts only, i.e., masking out the input',
)
_ATTN_IMPLEMENTATION = flags.DEFINE_string(
'attn_implementation',
None,
'Attention implementation, can be `eager`, `sdpa` or `flash_attention_2`',
)
_MAX_GRAD_NORM = flags.DEFINE_float(
'max_grad_norm',
0.3,
'Maximum gradient norm used for gradient clipping',
)
_WARNINGS_FILTER = flags.DEFINE_string(
'warnings_filter',
'ignore',
'Warning filter as defined in '
'https://docs.python.org/3/library/warnings.html#the-warnings-filter',
)
_LOGGER_LEVEL = flags.DEFINE_string(
'logger_level',
'passive',
'logging level passed to TrainingArguments. Note that this is for python'
' logging module, NOT the one from absl',
)
_BENCHMARK_OUT_FILE = flags.DEFINE_string(
'benchmark_out_file', None, 'file path for writing benchmark result'
)
_NCCL_TIMEOUT = flags.DEFINE_integer(
'nccl_timeout', 6000, 'nccl timeout in seconds'
)
_TUNING_DATA_STATS_FILE = flags.DEFINE_string(
'tuning_data_stats_file', None, 'file path for writing tuning data stats.'
)
_TARGET_MODULES = flags.DEFINE_list(
'target_modules', None, 'The names of the modules to apply LoRA adapter to.'
)
@flags.multi_flags_validator(
[
_COMPLETION_ONLY.name,
_USE_EXAMPLE_PACKING.name,
],
message=(
'`use_example_packing=True` does not work with `completion_only=True`'
),
)
def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
"""Check to make sure example packing is enabled properly.
Args:
flags_dict: Dictionary containing flags to check.
Returns:
If `use_example_packing` is set properly.
"""
if (
flags_dict[_COMPLETION_ONLY.name]
and flags_dict[_USE_EXAMPLE_PACKING.name]
):
return False
return True
@flags.multi_flags_validator(
[
_COMPLETION_ONLY.name,
_TEMPLATE.name,
],
message='`template` should be provided if using `completion_only=True`',
)
def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
"""Check to make sure completion_only is enabled properly.
Args:
flags_dict: Dictionary containing flags to check.
Returns:
If `completion_only` is set properly
"""
if flags_dict[_COMPLETION_ONLY.name] and flags_dict[_TEMPLATE.name] is None:
return False
return True
# References:
# Huggingface SFT trainer example:
# https://github.com/huggingface/trl/blob/main/examples/scripts/sft_trainer.py.
# Huggingface sagemaker example:
# https://github.com/huggingface/notebooks/blob/main/sagemaker/28_train_llms_with_qlora/scripts/run_clm.py.
# Copied from https://github.com/artidoro/qlora/blob/main/qlora.py.
def find_all_linear_names(
model: AutoModelForCausalLM, precision_mode: str
) -> list[str]:
"""Finds all linear module names."""
if precision_mode == constants.PRECISION_MODE_4:
cls = bnb.nn.Linear4bit
elif precision_mode == constants.PRECISION_MODE_8:
cls = bnb.nn.Linear8bitLt
else:
cls = torch.nn.Linear
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
def finetune_instruct(
pretrained_model_id: str,
dataset_name: str,
output_dir: str,
logging_output_dir: str,
lora_rank: int = 64,
lora_alpha: int = 16,
lora_dropout: float = 0.1,
warmup_ratio: int = 0.03,
num_epochs: Optional[float] = None,
max_steps: Optional[int] = None,
warmup_steps: int = 10,
max_seq_length: int = 512,
learning_rate: float = 2e-4,
precision_mode: str = None,
instruct_column_in_dataset: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
per_device_train_batch_size: int = 4,
gradient_accumulation_steps: int = 4,
optim: str = 'paged_adamw_32bit',
weight_decay: float = 0.001,
enable_gradient_checkpointing: bool = False,
enable_peft: bool = True,
template: str = None,
lr_scheduler_type: str = 'constant',
save_steps: int = 10,
logging_steps: int = 10,
train_split_name: str = 'train',
eval_config: Optional[eval_lib.EvalConfig] = None,
report_to: str = constants.REPORT_TO_NONE,
access_token: Optional[str] = None,
train_precision: str = constants.PRECISION_MODE_16B,
use_example_packing: bool = False,
attn_implementation: Optional[str] = None,
max_grad_norm: float = 0.3,
completion_only: bool = False,
logger_level: str = 'passive',
benchmark_out_file: Optional[str] = None,
tuning_data_stats_file: Optional[str] = None,
target_modules: Optional[str] = None,
) -> None:
"""Finetunes instruct."""
logging.info('on entering instruct_lora, %s', utils.gpu_stats_str())
gradient_checkpointing_kwargs = {}
# DDP provides limited support with the reentrant variant of gradient
# checkpoint [1]. Below is an indirect way of checking whether DDP will be
# used. It is "indirect" because there are complex logic under the hood of
# `SFTTrainer` and since those are not public API, they might change as we
# update the library.
if PartialState().distributed_type == DistributedType.MULTI_GPU:
gradient_checkpointing_kwargs['use_reentrant'] = False
tokenizer = utils.load_tokenizer(
pretrained_model_id,
'right',
access_token=access_token,
)
train_dataset = dataset_validation_util.load_dataset_with_template(
dataset_name,
split=train_split_name,
input_column=instruct_column_in_dataset,
template=template,
tokenizer=tokenizer,
)
if tuning_data_stats_file:
with PartialState().main_process_first():
effective_batch_size = (
per_device_train_batch_size
* gradient_accumulation_steps
* PartialState().num_processes
)
logging.info(
'getting tuning data stats with effective batch size %s',
effective_batch_size,
)
train_dataset_stats = utils.get_dataset_stats(
train_dataset,
tokenizer,
instruct_column_in_dataset,
effective_batch_size,
)
logging.info('stats: %s', train_dataset_stats)
tuning_data_stats_file = dataset_validation_util.force_gcs_fuse_path(
tuning_data_stats_file
)
with open(tuning_data_stats_file, 'w') as out_f:
json.dump(dataclasses.asdict(train_dataset_stats), out_f)
model = utils.load_model(
pretrained_model_id=pretrained_model_id,
tokenizer=tokenizer,
precision_mode=precision_mode,
enable_gradient_checkpointing=enable_gradient_checkpointing,
access_token=access_token,
attn_implementation=attn_implementation,
train_precision=train_precision,
)
if enable_peft:
if target_modules is None:
target_modules = find_all_linear_names(
model, precision_mode=precision_mode
)
logging.info('applying lora adapters to modules: %s', target_modules)
peft_config = LoraConfig(
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
r=lora_rank,
bias='none',
task_type='CAUSAL_LM',
target_modules=target_modules,
)
# If we pass in `peft_config` to SFTTrainer, it does a lot of magic under
# the hood, e.g., calling `prepare_model_for_kbit_training` before calling
# `get_peft_model`, which may revert other changes we did before. That's why
# we are calling `get_peft_model` explicitly here.
model = get_peft_model(model, peft_config)
# This is to work-around mix-precision training. This issue is not fixed as
# of transformers==4.41.2.
# See b/332760883#comment30 for more details.
if precision_mode in (
constants.PRECISION_MODE_16,
constants.PRECISION_MODE_16B,
):
for param in filter(lambda p: p.requires_grad, model.parameters()):
param.data = param.data.to(torch.float32)
if not logging_output_dir:
logging_output_dir = output_dir
# To use singleton PartialState() without re-initializing it. See
# b/357970482#comment3
accelerator_config = {'use_configured_state': True}
training_arguments = TrainingArguments(
report_to=report_to,
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,
save_strategy='steps',
save_total_limit=3,
logging_dir=os.path.join(logging_output_dir, 'logs'),
logging_steps=logging_steps,
learning_rate=learning_rate,
fp16=(train_precision == constants.PRECISION_MODE_16),
bf16=(train_precision == constants.PRECISION_MODE_16B),
max_grad_norm=max_grad_norm,
num_train_epochs=num_epochs if num_epochs else -1,
max_steps=max_steps if max_steps else -1,
warmup_ratio=warmup_ratio,
warmup_steps=warmup_steps,
group_by_length=False,
lr_scheduler_type=lr_scheduler_type,
gradient_checkpointing=enable_gradient_checkpointing,
gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
weight_decay=weight_decay,
log_level=logger_level,
accelerator_config=accelerator_config,
)
trainer_kwargs = {}
if completion_only and template:
template_json = dataset_validation_util.get_template(template_path=template)
instruction_sep = dataset_validation_util.get_instruction_separator(
template_json
)
response_sep = dataset_validation_util.get_response_separator(template_json)
if not response_sep:
raise ValueError(
'`response_separator` must be provided to use'
' `DataCollatorForCompletionOnlyLM`'
)
trainer_kwargs['data_collator'] = DataCollatorForCompletionOnlyLM(
instruction_template=instruction_sep,
response_template=response_sep,
tokenizer=tokenizer,
)
logging.info('using DataCollatorForCompletionOnlyLM')
trainer_stats_callback = callbacks.TrainerStatsCallback(
max_seq_length, benchmark_out_file
)
trainer = eval_lib.create_trainer(
cls=SFTTrainer,
eval_config=eval_config,
model=model,
train_dataset=train_dataset,
dataset_text_field=instruct_column_in_dataset,
max_seq_length=max_seq_length,
tokenizer=tokenizer,
args=training_arguments,
packing=use_example_packing,
callbacks=[trainer_stats_callback],
**trainer_kwargs,
)
# `eval_lib.create_trainer` might modify the training args. Printing here
# should capture what will be used by the trainer.
if PartialState().is_main_process:
logging.info('training args: %s', trainer.args)
if enable_peft:
trainer.model.print_trainable_parameters()
if trainer.is_fsdp_enabled:
logging.info('Trainer running with FSDP.')
elif trainer.is_deepspeed_enabled:
logging.info('Trainer running with DeepSpeed.')
else:
logging.info('Trainer running without parallelism.')
trainer.train()
# Always save the final checkpoint.
final_checkpoint = utils.get_final_checkpoint_path(output_dir)
logging.info('The final checkpoint is: %s.', final_checkpoint)
if trainer.is_fsdp_enabled:
trainer.accelerator.state.fsdp_plugin.set_state_dict_type('FULL_STATE_DICT')
# This method saves the sharded weights like `accelerator.save_state`, see
# https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
trainer.save_model(output_dir)
model = trainer.model.cpu() # Avoids GPU OOM
state_dict = trainer.accelerator.get_state_dict(model)
# To aggregate the weights from all the devices, we need to use
# `state_dict=state_dict`.
model.save_pretrained(
final_checkpoint,
state_dict=state_dict,
is_main_process=PartialState().is_main_process,
save_embedding_layers=False, # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
)
model.cuda() # Move back to GPU to do eval.
else:
trainer.model.save_pretrained(
final_checkpoint,
is_main_process=PartialState().is_main_process,
save_embedding_layers=False, # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
)
if eval_config is not None and trainer.eval_dataset is not None:
metrics = trainer.evaluate(metric_key_prefix='eval')
# Both `log_metrics` and `save_metrics` are multiple process safe.
# https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L911 #pylint: disable=line-too-long
# https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L1001 #pylint: disable=line-too-long
trainer.log_metrics('eval', metrics)
trainer.save_metrics('eval', metrics)
if PartialState().is_main_process:
hp_metric = metrics[f'eval_{eval_config.metric_name}']
hpt = hypertune.HyperTune()
hpt.report_hyperparameter_tuning_metric(
hyperparameter_metric_tag=constants.HP_METRIC_TAG,
metric_value=hp_metric,
)
logging.info('Send HP metric: %f to hyperparameter tuning.', hp_metric)
PartialState().wait_for_everyone()
if not enable_peft:
tokenizer.save_pretrained(
final_checkpoint, is_main_process=PartialState().is_main_process
)
def main(unused_argv: Sequence[str]) -> None:
# This needs to be called before any other PartialState() calls.
utils.init_partial_state(
timeout=datetime.timedelta(seconds=_NCCL_TIMEOUT.value)
)
utils.print_library_versions()
warnings.simplefilter(_WARNINGS_FILTER.value)
pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
if dataset_validation_util.is_gcs_path(pretrained_model_id):
pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
pretrained_model_id
)
output_dir = utils.GcsOrLocalDirectory(
_OUTPUT_DIR.value, check_empty=True, upload_from_all_nodes=True
)
# GCS Fuse does not sync flushed files if not closed. See b/361771727.
logging_output_dir = fileutils.force_gcs_path(_LOGGING_OUTPUT_DIR.value)
# Creates evaluation config.
if _EVAL_TASKS.value:
eval_config = eval_lib.EvalConfig(
tasks=_EVAL_TASKS.value,
per_device_batch_size=_EVAL_PER_DEVICE_BATCH_SIZE.value,
num_fewshot=_EVAL_NUM_FEWSHOT.value,
limit=_EVAL_LIMIT.value,
metric_name=_EVAL_METRIC_NAME.value,
steps=_EVAL_STEPS.value,
dataset_path=dataset_validation_util.force_gcs_fuse_path(
_EVAL_DATASET_PATH.value
),
split=_EVAL_SPLIT.value,
template=_EVAL_TEMPLATE.value,
column=_EVAL_COLUMN.value,
tokenize_dataset=False,
)
else:
eval_config = None
if _REPORT_TO.value == constants.REPORT_TO_WANDB:
wandb.login()
finetune_instruct(
pretrained_model_id=pretrained_model_id,
dataset_name=_DATASET_NAME.value,
output_dir=output_dir.local_dir,
logging_output_dir=logging_output_dir,
precision_mode=_PRECISION_MODE.value,
lora_rank=_LORA_RANK.value,
lora_alpha=_LORA_ALPHA.value,
lora_dropout=_LORA_DROPOUT.value,
warmup_ratio=_WARMUP_RATIO.value,
num_epochs=_NUM_EPOCHS.value,
warmup_steps=_WARMUP_STEPS.value,
max_steps=_MAX_STEPS.value,
max_seq_length=_MAX_SEQ_LENGTH.value,
learning_rate=_LEARNING_RATE.value,
instruct_column_in_dataset=_INSTRUCT_COLUMN_IN_DATASET.value,
per_device_train_batch_size=_PER_DEVICE_TRAIN_BATCH_SIZE.value,
optim=_OPTIMIZER.value,
weight_decay=_WEIGHT_DECAY.value,
gradient_accumulation_steps=_GRADIENT_ACCUMULATION_STEPS.value,
enable_gradient_checkpointing=_ENABLE_GRADIENT_CHECKPOINTING.value,
enable_peft=_ENABLE_PEFT.value,
template=_TEMPLATE.value,
lr_scheduler_type=_LR_SCHEDULER_TYPE.value,
save_steps=_SAVE_STEPS.value,
logging_steps=_LOGGING_STEPS.value,
train_split_name=_TRAIN_SPLIT_NAME.value,
eval_config=eval_config,
report_to=_REPORT_TO.value,
access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
train_precision=_TRAIN_PRECISION.value,
use_example_packing=_USE_EXAMPLE_PACKING.value,
attn_implementation=_ATTN_IMPLEMENTATION.value,
max_grad_norm=_MAX_GRAD_NORM.value,
completion_only=_COMPLETION_ONLY.value,
logger_level=_LOGGER_LEVEL.value,
benchmark_out_file=_BENCHMARK_OUT_FILE.value,
tuning_data_stats_file=_TUNING_DATA_STATS_FILE.value,
target_modules=_TARGET_MODULES.value,
)
# Frees the model from GPU.
utils.force_gc()
output_dir.upload_to_gcs(skip_if_exists=True)
if __name__ == '__main__':
app.run(main)
community-content/vertex_model_garden/model_oss/peft/train/vmg/merge_causal_language_model_lora.py
+132
View File
@@ -0,0 +1,132 @@
"""Script to merge PEFT adapter with base model."""
from typing import Any, Dict, Sequence
from absl import app
from absl import flags
from util import dataset_validation_util
from vertex_vision_model_garden_peft.train.vmg import utils
from util import constants
from util import fileutils
_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. Note, there'
' might be different paddings for different models. This tool assumes the'
' pretrained_model_id contains model name, and then choose proper padding'
' methods. e.g. it must contain `llama` for `Llama2 models`.',
required=True,
)
_MERGE_BASE_AND_LORA_OUTPUT_DIR = flags.DEFINE_string(
'merge_base_and_lora_output_dir',
None,
'The directory to store the merged model with the base and lora adapter.',
)
_MERGE_MODEL_PRECISION_MODE = flags.DEFINE_enum(
'merge_model_precision_mode',
constants.PRECISION_MODE_16,
[
constants.PRECISION_MODE_4,
constants.PRECISION_MODE_8,
constants.PRECISION_MODE_FP8,
constants.PRECISION_MODE_16,
constants.PRECISION_MODE_16B,
constants.PRECISION_MODE_32,
],
'Merging model precision mode.',
)
_FINETUNED_LORA_MODEL_DIR = flags.DEFINE_string(
'finetuned_lora_model_dir',
None,
'The directory storing finetuned LoRA model weights.',
)
_RESTRICT_MODEL_UPLOAD_DOCKER_URI = flags.DEFINE_string(
'restrict_model_upload_docker_uri',
'',
'If set, mark output model as only uploadable to Model Registry with the'
' specified Docker URI.',
)
_EXECUTOR_INPUT = flags.DEFINE_string(
'executor_input',
'',
'For internal use. Kubeflow pipeline context when running trainer as part'
' of an internal pipeline.',
)
_HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
'huggingface_access_token',
None,
'The access token for loading huggingface gated models.',
)
@flags.multi_flags_validator(
[
_PRETRAINED_MODEL_ID.name,
_FINETUNED_LORA_MODEL_DIR.name,
_MERGE_BASE_AND_LORA_OUTPUT_DIR.name,
],
)
def check_merge_lora_model_flags(flags_dict: Dict[str, Any]) -> bool:
"""Check if required flags are set on merge model LoRA task.
Args:
flags_dict: Dictionary containing task and flags to check.
Returns:
If required flags are not None.
"""
return all(map(lambda x: x is not None, flags_dict.values()))
def main(unused_argv: Sequence[str]) -> None:
pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
if dataset_validation_util.is_gcs_path(pretrained_model_id):
pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
pretrained_model_id
)
finetuned_lora_model_dir = utils.GcsOrLocalDirectory(
_FINETUNED_LORA_MODEL_DIR.value
)
merge_base_and_lora_output_dir = utils.GcsOrLocalDirectory(
_MERGE_BASE_AND_LORA_OUTPUT_DIR.value
)
utils.merge_causal_language_model_with_lora(
pretrained_model_id=pretrained_model_id,
precision_mode=_MERGE_MODEL_PRECISION_MODE.value,
finetuned_lora_model_dir=finetuned_lora_model_dir.local_dir,
merged_model_output_dir=merge_base_and_lora_output_dir.local_dir,
access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
)
if _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value:
utils.write_first_party_model_metadata(
merge_base_and_lora_output_dir.local_dir,
_RESTRICT_MODEL_UPLOAD_DOCKER_URI.value,
)
if _EXECUTOR_INPUT.value:
utils.write_kfp_outputs(
_EXECUTOR_INPUT.value,
{
'saved_model': _MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
},
)
merge_base_and_lora_output_dir.upload_to_gcs(skip_if_exists=True)
if __name__ == '__main__':
app.run(main)
community-content/vertex_model_garden/model_oss/peft/train/vmg/quantize_model.py
+349
View File
@@ -0,0 +1,349 @@
"""Quantizes the model."""
import json
import os
from typing import Any, Dict, List, Sequence, Union
from absl import app
from absl import flags
from absl import logging
from auto_gptq import AutoGPTQForCausalLM
from auto_gptq import BaseQuantizeConfig
from awq import AutoAWQForCausalLM
from optimum.gptq.data import get_dataset
from transformers import AutoTokenizer
from util import dataset_validation_util
from vertex_vision_model_garden_peft.train.vmg import utils
from util import constants
_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. Note, there'
' might be different paddings for different models. This tool assumes the'
' pretrained_model_id contains model name, and then choose proper padding'
' methods. e.g. it must contain `llama` for `Llama2 models`.',
)
_QUANTIZATION_METHOD = flags.DEFINE_enum(
'quantization_method',
None,
[constants.GPTQ, constants.AWQ],
'The quantization method. Choose from ["gtpq", "awq"].',
)
_QUANTIZATION_PRECISION_MODE = flags.DEFINE_enum(
'quantization_precision_mode',
constants.PRECISION_MODE_4,
[
constants.PRECISION_MODE_8,
constants.PRECISION_MODE_4,
constants.PRECISION_MODE_3,
constants.PRECISION_MODE_2,
],
'Quantization precision mode.',
)
_QUANTIZATION_DATASET_NAME = flags.DEFINE_string(
'quantization_dataset_name',
None,
'The dataset used for quantization. You can provide your own dataset in a'
' list of string or just use the original datasets used in GPTQ paper'
' ["wikitext2","c4","c4-new","ptb","ptb-new"] for GPTQ quantization. Using'
" a dataset more appropriate to the model's training can improve"
' quantisation accuracy. Note that the GPTQ dataset is not the same as the'
' dataset used to train the model.',
)
_TEXT_COLUMN_IN_QUANTIZATION_DATASET = flags.DEFINE_string(
'text_column_in_quantization_dataset',
constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
'The text column in quantization dataset.',
)
_QUANTIZATION_OUTPUT_DIR = flags.DEFINE_string(
'quantization_output_dir',
None,
'The directory to store the quantized model.',
)
_QUANTIZATION_DEVICE_MAP = flags.DEFINE_string(
'device_map', None, 'The device map.'
)
_QUANTIZATION_MAX_MEMORY = flags.DEFINE_string(
'max_memory', None, 'The maximum memory.'
)
_GROUP_SIZE = flags.DEFINE_integer(
'group_size',
None,
'The group size to use for quantization. Recommended value is 128 and -1'
' uses per-column quantization. Higher numbers use less VRAM, but have'
' lower quantisation accuracy. "None" is the lowest possible value.',
)
_DESC_ACT = flags.DEFINE_boolean(
'desc_act',
False,
'Whether to quantize columns in order of decreasing activation size.'
' Setting it to False can significantly speed up inference but the'
' perplexity may become slightly worse. Also known as act-order.',
)
_DAMP_PERCENT = flags.DEFINE_float(
'damp_percent',
0.1,
'The percent of the average Hessian diagonal to use for dampening.',
)
_CACHE_EXAMPLES_ON_GPU = flags.DEFINE_boolean(
'cache_examples_on_gpu',
True,
'Whether to cache the examples on GPU. Disabling will reduce VRAM usage,'
' but increase quantization time.',
)
_AWQ_VERSION = flags.DEFINE_enum(
'awq_version',
constants.GEMM,
[constants.GEMM, constants.GEMV],
'The version of the AWQ to use. It determines how matrix multiplication'
' runs under the hood. GEMV is 20% faster than GEMM, only at batch size 1'
' (not good for large contexts). GEMM is much faster than FP16 at batch'
' sizes below 8 (good with large contexts).',
)
@flags.multi_flags_validator(
[
_PRETRAINED_MODEL_ID.name,
_QUANTIZATION_METHOD.name,
_QUANTIZATION_PRECISION_MODE.name,
_QUANTIZATION_DATASET_NAME.name,
_QUANTIZATION_OUTPUT_DIR.name,
],
)
def check_quantization_flags(flags_dict: Dict[str, Any]) -> bool:
"""Check if required flags are set on quantization task.
Args:
flags_dict: Dictionary containing task and flags to check.
Returns:
If required flags are not None.
"""
required_flags = [
_QUANTIZATION_METHOD.name,
_PRETRAINED_MODEL_ID.name,
_QUANTIZATION_PRECISION_MODE.name,
_QUANTIZATION_DATASET_NAME.name,
_QUANTIZATION_OUTPUT_DIR.name,
]
return all(map(lambda x: flags_dict[x] is not None, required_flags))
def quantize_model(
quantization_method: str,
pretrained_model_id: str,
quantization_output_dir: str,
quantization_precision_mode: str = None,
quantization_dataset_name: Union[List[str]] = None,
text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
group_size: int = None,
desc_act: bool = True,
damp_percent: float = 0.1,
awq_version: str = 'GEMM',
device_map: str = None,
max_memory: Dict[Any, str] = None,
cache_examples_on_gpu: bool = True,
) -> None:
"""Quantizes the model using `quantization_method`."""
if quantization_method == constants.GPTQ:
gptq_quantize_model(
pretrained_model_id=pretrained_model_id,
gptq_output_dir=quantization_output_dir,
gptq_precision_mode=quantization_precision_mode,
gptq_dataset_name=quantization_dataset_name,
group_size=group_size,
desc_act=desc_act,
damp_percent=damp_percent,
cache_examples_on_gpu=cache_examples_on_gpu,
)
elif quantization_method == constants.AWQ:
awq_quantize_model(
pretrained_model_id=pretrained_model_id,
quantization_output_dir=quantization_output_dir,
quantization_precision_mode=quantization_precision_mode,
quantization_dataset_name=quantization_dataset_name,
text_column_in_quantization_dataset=text_column_in_quantization_dataset,
group_size=group_size,
awq_version=awq_version,
device_map=device_map,
max_memory=max_memory,
)
def awq_quantize_model(
pretrained_model_id: str,
quantization_output_dir: str,
quantization_precision_mode: str = None,
quantization_dataset_name: Union[List[str]] = None,
text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
group_size: int = None,
awq_version: str = 'GEMM',
device_map: str = None,
max_memory: Dict[Any, str] = None,
) -> None:
"""Quantizes the model using AWQ."""
if quantization_precision_mode != constants.PRECISION_MODE_4:
raise ValueError(
f'Invalid precision mode: {quantization_precision_mode} for AWQ. 4bit'
' quantization must be used.'
)
else:
bits = 4
if not group_size:
group_size = 128
if not device_map:
device_map = 'cpu'
if dataset_validation_util.is_gcs_path(quantization_dataset_name):
logging.info('Using custom dataset: %s', quantization_dataset_name)
with open(
dataset_validation_util.force_gcs_fuse_path(quantization_dataset_name),
'r',
) as f:
quantization_dataset = [line.rstrip('\n') for line in f]
else:
quantization_dataset = quantization_dataset_name
quant_config = {
'zero_point': True,
'q_group_size': group_size,
'w_bit': bits,
'version': awq_version,
}
logging.info('Quantization config: %s', quant_config)
model = AutoAWQForCausalLM.from_pretrained(
pretrained_model_id,
trust_remote_code=True,
device_map=device_map,
max_memory=max_memory,
low_cpu_mem_usage=True,
)
tokenizer = AutoTokenizer.from_pretrained(
pretrained_model_id, trust_remote_code=True
)
model.quantize(
tokenizer,
quant_config=quant_config,
calib_data=quantization_dataset,
text_column=text_column_in_quantization_dataset,
)
model.save_quantized(quantization_output_dir)
tokenizer.save_pretrained(quantization_output_dir)
def gptq_quantize_model(
pretrained_model_id: str,
gptq_output_dir: str,
gptq_precision_mode: str = None,
gptq_dataset_name: Union[List[str]] = None,
group_size: int = -1,
desc_act: bool = False,
damp_percent: float = 0.1,
cache_examples_on_gpu: bool = True,
) -> None:
"""Quantizes the model using GPTQ."""
logging.info(
'PYTORCH_CUDA_ALLOC_CONF: %s',
os.environ.get('PYTORCH_CUDA_ALLOC_CONF', ''),
)
if dataset_validation_util.is_gcs_path(gptq_dataset_name):
logging.info('Using custom dataset: %s', gptq_dataset_name)
with open(
dataset_validation_util.force_gcs_fuse_path(gptq_dataset_name), 'r'
) as f:
gptq_dataset = [line.rstrip('\n') for line in f]
else:
gptq_dataset = gptq_dataset_name
if gptq_precision_mode == constants.PRECISION_MODE_8:
bits = 8
elif gptq_precision_mode == constants.PRECISION_MODE_4:
bits = 4
elif gptq_precision_mode == constants.PRECISION_MODE_3:
bits = 3
elif gptq_precision_mode == constants.PRECISION_MODE_2:
bits = 2
else:
raise ValueError(f'Invalid precision mode: {gptq_precision_mode} for GPTQ.')
if not group_size:
group_size = -1
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_id)
gptq_dataset = get_dataset(gptq_dataset, tokenizer)
quantization_config = BaseQuantizeConfig(
bits=bits,
group_size=group_size,
damp_percent=damp_percent,
desc_act=desc_act,
)
logging.info('Quantization config: %s', quantization_config.to_dict())
model = AutoGPTQForCausalLM.from_pretrained(
pretrained_model_id,
quantization_config,
low_cpu_mem_usage=True,
torch_dtype='auto',
trust_remote_code=True,
)
model.quantize(
examples=gptq_dataset,
cache_examples_on_gpu=cache_examples_on_gpu,
)
if utils.should_add_pad_token(pretrained_model_id):
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))
model.save_pretrained(gptq_output_dir)
tokenizer.save_pretrained(gptq_output_dir)
def main(unused_argv: Sequence[str]) -> None:
pretrained_model_id = _PRETRAINED_MODEL_ID.value
if dataset_validation_util.is_gcs_path(pretrained_model_id):
pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
pretrained_model_id
)
pretrained_model_id = dataset_validation_util.force_gcs_fuse_path(
pretrained_model_id
)
if _QUANTIZATION_MAX_MEMORY.value:
max_memory = json.loads(_QUANTIZATION_MAX_MEMORY.value)
else:
max_memory = None
quantize_model(
quantization_method=_QUANTIZATION_METHOD.value,
pretrained_model_id=pretrained_model_id,
quantization_output_dir=_QUANTIZATION_OUTPUT_DIR.value,
quantization_precision_mode=_QUANTIZATION_PRECISION_MODE.value,
quantization_dataset_name=_QUANTIZATION_DATASET_NAME.value,
text_column_in_quantization_dataset=_TEXT_COLUMN_IN_QUANTIZATION_DATASET.value,
group_size=_GROUP_SIZE.value,
desc_act=_DESC_ACT.value,
damp_percent=_DAMP_PERCENT.value,
awq_version=_AWQ_VERSION.value,
device_map=_QUANTIZATION_DEVICE_MAP.value,
max_memory=max_memory,
cache_examples_on_gpu=_CACHE_EXAMPLES_ON_GPU.value,
)
if __name__ == '__main__':
app.run(main)
community-content/vertex_model_garden/model_oss/peft/train/vmg/scripts/build_train_docker.sh
Executable
+21
View File
@@ -0,0 +1,21 @@
#!/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/vmg/scripts/build_train_docker.sh
set -x
set -e
COPYBARA_DIR="/tmp/train_docker/"
pushd "${COPYBARA_DIR}"
PROJECT="cloud-nas-260507"
IMAGE_TAG="gcr.io/${PROJECT}/pytorch-peft-train:${USER}-test"
docker build -f model_oss/peft/train/vmg/dockerfile/train.Dockerfile . -t "${IMAGE_TAG}"
docker push "${IMAGE_TAG}"
popd
community-content/vertex_model_garden/model_oss/peft/sequence_classification_lora.py → community-content/vertex_model_garden/model_oss/peft/train/vmg/sequence_classification_lora.py
+97 -1
View File
@@ -1,7 +1,9 @@
"""Sequence classification with LoRA models."""
# pylint: disable=g-importing-member
from typing import Sequence
from absl import app
from absl import flags
from datasets import load_dataset
import evaluate
from peft import get_peft_model
@@ -14,6 +16,71 @@ from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
from transformers import get_linear_schedule_with_warmup
from util import dataset_validation_util
_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. Note, there"
" might be different paddings for different models. This tool assumes the"
" pretrained_model_id contains model name, and then choose proper padding"
" methods. e.g. it must contain `llama` for `Llama2 models`.",
)
_OUTPUT_DIR = flags.DEFINE_string(
"output_dir",
None,
"The output directory.",
)
_DATASET_NAME = flags.DEFINE_string(
"dataset_name",
None,
"The dataset name in huggingface.",
)
_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.",
)
_NUM_EPOCHS = flags.DEFINE_integer(
"num_epochs",
None,
"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 finetune_sequence_classification(
pretrained_model_id: str,
@@ -131,3 +198,32 @@ def finetune_sequence_classification(
print(f"epoch {epoch}:", eval_metric)
model.save_pretrained(output_dir)
def main(unused_argv: Sequence[str]) -> None:
if dataset_validation_util.is_gcs_path(_PRETRAINED_MODEL_ID.value):
pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
_PRETRAINED_MODEL_ID.value
)
else:
pretrained_model_id = _PRETRAINED_MODEL_ID.value
pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
pretrained_model_id
)
output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
finetune_sequence_classification(
pretrained_model_id=pretrained_model_path,
dataset_name=_DATASET_NAME.value,
output_dir=output_dir,
lora_rank=_LORA_RANK.value,
lora_alpha=_LORA_ALPHA.value,
lora_dropout=_LORA_DROPOUT.value,
num_epochs=int(_NUM_EPOCHS.value),
batch_size=_BATCH_SIZE.value,
learning_rate=_LEARNING_RATE.value,
)
if __name__ == "__main__":
app.run(main)
community-content/vertex_model_garden/model_oss/peft/templates/alpaca.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/alpaca.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/alpaca_short.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/alpaca_short.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/gemma.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/gemma.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/llama3-text-bison.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/llama3-text-bison.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/llama3.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/llama3.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/mistral.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/mistral.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/openai-completion.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/openai-completion.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/openassistant-guanaco.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/openassistant-guanaco.json
View File
community-content/vertex_model_garden/model_oss/peft/templates/vertex-chat.json → community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/vertex-chat.json
View File
community-content/vertex_model_garden/model_oss/peft/train/vmg/train_entrypoint.py
+242
View File
@@ -0,0 +1,242 @@
"""Entrypoint for peft train docker.
Dispatches to different scripts based on `task` type.
For task type in `_TASK_TO_SCRIPT`, if `--config_file` is specified, the script
will dispatch the call to `accelerate`, which is friendly for multi-GPU
environment. Otherwise, `python3` is used.
"""
import argparse
import json
import os
import subprocess
from typing import List, Optional, Sequence
from absl import app
from absl import flags
from absl import logging
from util import dataset_validation_util
from vertex_vision_model_garden_peft.train.vmg import utils
from util import constants
from util import hypertune_utils
_TEXT_TO_IMAGE_TASKS_SCRIPTS = {
constants.TEXT_TO_IMAGE: 'text_to_image/train_text_to_image.py',
constants.TEXT_TO_IMAGE_LORA: 'text_to_image/train_text_to_image_lora.py',
constants.TEXT_TO_IMAGE_DREAMBOOTH: 'dreambooth/train_dreambooth.py',
constants.TEXT_TO_IMAGE_DREAMBOOTH_LORA: (
'dreambooth/train_dreambooth_lora.py'
),
constants.TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL: (
'dreambooth/train_dreambooth_lora_sdxl.py'
),
}
_TASK_TO_SCRIPT = {
constants.INSTRUCT_LORA: (
'vertex_vision_model_garden_peft/train/vmg/instruct_lora.py'
),
constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: 'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py',
constants.QUANTIZE_MODEL: (
'vertex_vision_model_garden_peft/train/vmg/quantize_model.py'
),
constants.SEQUENCE_CLASSIFICATION_LORA: 'vertex_vision_model_garden_peft/train/vmg/sequence_classification_lora.py',
constants.VALIDATE_DATASET_WITH_TEMPLATE: 'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py',
}
def launch_script_cmd(
script: str,
config_file: Optional[str],
accelerate_args: argparse.Namespace = argparse.Namespace(),
) -> List[str]:
"""Returns the command to launch the script."""
if config_file:
cmd = [
'accelerate',
'launch',
'--config_file={}'.format(config_file),
]
else:
cmd = ['python3']
_append_args_to_command_in_place(accelerate_args, cmd)
cmd.append(script)
return cmd
def _get_accelerate_args() -> argparse.Namespace:
"""Returns the accelerate args."""
# For the format of the cluster spec, see
# https://cloud.google.com/vertex-ai/docs/training/distributed-training#cluster-spec-format # pylint: disable=line-too-long
cluster_spec = os.getenv('CLUSTER_SPEC', default=None)
if not cluster_spec:
return argparse.Namespace()
logging.info('CLUSTER_SPEC: %s', cluster_spec)
cluster_data = json.loads(cluster_spec)
if (
'workerpool1' not in cluster_data['cluster']
or not cluster_data['cluster']['workerpool1']
):
return argparse.Namespace()
# 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_worker_nodes = len(cluster_data['cluster']['workerpool1'])
num_nodes = num_worker_nodes + 1 # Add 1 for the primary node
logging.info('num nodes: %s', num_nodes)
accelerate_args = argparse.Namespace()
accelerate_args.machine_rank = node_rank
accelerate_args.num_machines = num_nodes
accelerate_args.main_process_ip = primary_node_addr
accelerate_args.main_process_port = primary_node_port
accelerate_args.max_restarts = 0
accelerate_args.monitor_interval = 120
return accelerate_args
def _append_args_to_command_in_place(
args: argparse.Namespace, command: List[str]
):
for key, value in vars(args).items():
# If not specified, skip.
if value is not None:
command.append(f'--{key}={value}')
def _get_train_cmd_and_maybe_merge_cmd(
task: str, config_file: str, unknown: Sequence[str]
) -> Sequence[Sequence[str]]:
"""Returns the training command and maybe the merge command if applicable."""
# Only populated when multi-node is used.
accelerate_args = _get_accelerate_args()
training_cmd = launch_script_cmd(
_TASK_TO_SCRIPT[task],
config_file,
accelerate_args=accelerate_args,
)
# Training only flag.
train_parser = argparse.ArgumentParser()
train_parser.add_argument('--output_dir', required=True)
training_args, unknown = train_parser.parse_known_args(unknown)
# Checks for `hypertune_utils._ENVIRONMENT_VARIABLE_FOR_TRIAL_ID` env var and
# appends the trial id if it exists.
training_args.output_dir = hypertune_utils.maybe_append_trial_id(
dataset_validation_util.force_gcs_fuse_path(training_args.output_dir)
)
# Merge only flags.
merge_parser = argparse.ArgumentParser()
merge_parser.add_argument('--merge_model_precision_mode')
merge_parser.add_argument('--executor_input')
merge_parser.add_argument('--restrict_model_upload_docker_uri')
merge_parser.add_argument('--merge_base_and_lora_output_dir')
merge_args, unknown = merge_parser.parse_known_args(unknown)
# Common flags shared by merging and training.
common_parser = argparse.ArgumentParser()
common_parser.add_argument('--pretrained_model_id', required=True)
common_parser.add_argument('--huggingface_access_token')
common_args, remaining = common_parser.parse_known_args(unknown)
# Add flags for training.
_append_args_to_command_in_place(training_args, training_cmd)
_append_args_to_command_in_place(common_args, training_cmd)
training_cmd.extend(remaining) # Remaining args are passed to training cmd.
commands = [training_cmd]
# Only the main node runs merging.
if (
merge_args.merge_base_and_lora_output_dir
and getattr(accelerate_args, 'machine_rank', 0) == 0
):
lora_dir = utils.get_final_checkpoint_path(training_args.output_dir)
merge_cmd = [
'WORLD_SIZE=1', # To ignore other nodes in multi-node setting.
'python3',
_TASK_TO_SCRIPT[constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA],
f'--finetuned_lora_model_dir={lora_dir}',
]
_append_args_to_command_in_place(merge_args, merge_cmd)
_append_args_to_command_in_place(common_args, merge_cmd)
# Run in a conda environment.
conda_run_cmd = [
'/bin/bash',
'-c',
f'conda run -n merge {" ".join(merge_cmd)}',
]
commands.append(conda_run_cmd)
return commands
def main(unused_argv: Sequence[str]) -> None:
parser = argparse.ArgumentParser()
parser.add_argument('--config_file')
parser.add_argument('--task')
args, unknown = parser.parse_known_args()
task = args.task
if task in _TEXT_TO_IMAGE_TASKS_SCRIPTS:
# Setup accelerate config before running trainer.
config_gen_cmd = [
'python',
'-c',
(
'from accelerate.utils import write_basic_config;'
' write_basic_config(mixed_precision="fp16")'
),
]
task_cmd = [
'accelerate',
'launch',
_TEXT_TO_IMAGE_TASKS_SCRIPTS[task],
] + list(map(dataset_validation_util.force_gcs_fuse_path, unknown))
commands = [config_gen_cmd, task_cmd]
elif task in [constants.INSTRUCT_LORA]:
commands = _get_train_cmd_and_maybe_merge_cmd(
task=task, config_file=args.config_file, unknown=unknown
)
else:
assert task in _TASK_TO_SCRIPT
cmd = launch_script_cmd(_TASK_TO_SCRIPT[task], args.config_file)
cmd.extend(unknown)
commands = [cmd]
for cmd in commands:
logging.info('launching task=%s with cmd: \n%s', task, ' \\\n'.join(cmd))
subprocess.run(cmd, check=True)
if __name__ == '__main__':
app.run(main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True))
community-content/vertex_model_garden/model_oss/peft/train/vmg/utils.py
+742
View File
@@ -0,0 +1,742 @@
"""Common libraries for PEFT."""
import dataclasses
import datetime
import gc
import multiprocessing as mp
import os
import subprocess
from typing import Any, Dict, Optional, Sequence
from absl import logging
import accelerate
from accelerate import DistributedType
from accelerate import PartialState
from google.protobuf import json_format
from kfp.pipeline_spec import pipeline_spec_pb2
import numpy as np
import peft
from peft import PeftModel
from peft import prepare_model_for_kbit_training
import pynvml
import torch
import transformers
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
from transformers import BitsAndBytesConfig
from transformers import FbgemmFp8Config
from transformers.integrations import is_deepspeed_zero3_enabled
import trl
from util import dataset_validation_util
from util import constants
from util import fileutils
_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
_LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
_LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
class GcsOrLocalDirectory(os.PathLike):
"""A class to represent a directory with upload support if GCS path is given.
This class is used to represent a directory. It can be used for a temporary
local directory and for uploading files to the GCS directory later if the
given path is a GCS directory. If the given path is a local directory, a call
to gcs_dir attribute will raise an error. This class has multi-node and
multi-process support with accelerate.
Attributes:
local_dir: The local directory to store the files.
gcs_dir: The path to the GCS directory.
"""
def __init__(
self,
path: str,
check_empty: bool = False,
upload_from_all_nodes: bool = False,
):
"""Initializes the GcsOrLocalDirectory.
Args:
path: The path to the directory.
check_empty: If True, check if the GCS directory is empty. No-op for local
directory.
upload_from_all_nodes: If True, upload the local directory to GCS from all
nodes.
"""
if len(path) > 1:
path = path.rstrip("/")
self._upload_from_all_nodes = upload_from_all_nodes
if path.startswith(constants.GCS_URI_PREFIX) or path.startswith(
constants.GCSFUSE_URI_PREFIX
):
self._is_gcs_path = True
self._local_dir = _get_local_dir_from_gcs_dir(path)
self._gcs_dir = fileutils.force_gcs_path(path)
os.makedirs(self.local_dir, exist_ok=True)
with PartialState().main_process_first():
if (
check_empty
and PartialState().is_main_process
and not _is_gcs_dir_empty(self._gcs_dir)
):
raise ValueError(f"{self._gcs_dir} needs to be empty.")
else:
self._is_gcs_path = False
self._local_dir = path
self._gcs_dir = path
def __fspath__(self) -> str:
return self.local_dir
@property
def local_dir(self) -> str:
return self._local_dir
@property
def gcs_dir(self) -> str:
"""Returns the GCS directory path.
Returns:
The GCS directory path.
Raises:
ValueError: If the path is not a GCS path.
"""
if not self._is_gcs_path:
raise ValueError(f"{self._gcs_dir} is not a GCS path.")
return self._gcs_dir
def upload_to_gcs(
self,
skip_if_exists: bool = True,
force_upload: bool = False,
):
"""Uploads the local directory to GCS."""
if not self._is_gcs_path:
logging.info(
"Not uploading to GCS since %s is not a GCS path.", self.local_dir
)
return
if not os.listdir(self.local_dir):
logging.info("Not uploading to GCS since %s is empty.", self.local_dir)
return
target = os.path.dirname(self.gcs_dir) + "/"
# Avoid race condition uploading the same file from multiple processes.
with PartialState().main_process_first():
if not PartialState().is_local_main_process:
# Non local main processes don't upload.
pass
elif self._upload_from_all_nodes or PartialState().is_main_process:
logging.info("Uploading %s to %s...", self.local_dir, target)
cmd = [
"gsutil",
"-m",
"cp",
"-r",
]
if skip_if_exists:
cmd.append("-n")
if force_upload:
cmd.append("-f")
cmd.extend([self.local_dir, target])
subprocess.check_output(cmd)
logging.info("%s uploaded.", self.local_dir)
def _get_local_dir_from_gcs_dir(path: str) -> str:
return os.path.join(
constants.LOCAL_OUTPUT_DIR,
dataset_validation_util.force_gcs_fuse_path(path)[1:],
)
def _is_gcs_dir_empty(path: str) -> bool:
"""Checks if a GCS directory is empty.
Args:
path: The GCS directory path.
Returns:
True if the directory is empty.
Raises:
subprocess.CalledProcessError: If the gsutil command failure reason is not
because the dir is empty.
"""
path = path.rstrip("/") + "/"
try:
subprocess.check_output(["gsutil", "ls", path], stderr=subprocess.STDOUT)
except subprocess.CalledProcessError as e:
if (
str(e.output, encoding="utf-8")
== "CommandException: One or more URLs matched no objects.\n"
):
return True
else:
logging.info(str(e.output, encoding="utf-8"))
raise
else:
return False
def load_tokenizer(
pretrained_model_id: str,
padding_side: Optional[str] = None,
access_token: Optional[str] = None,
) -> AutoTokenizer:
"""Loads tokenizer based on `pretrained_model_id`."""
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 PartialState().local_main_process_first():
tokenizer = 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 load_model(
pretrained_model_id: str,
tokenizer: AutoTokenizer,
precision_mode: str = None,
enable_gradient_checkpointing: bool = False,
gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
access_token: Optional[str] = None,
attn_implementation: Optional[str] = None,
train_precision: Optional[str] = None,
device_map: Optional[str] = None,
is_training: bool = True,
) -> AutoModelForCausalLM:
"""Loads models from the local dir if specified or from huggingface."""
# The `distributed_type` we got through `PartialState` is incorrect for FSDP.
# And that's why `Accelerator` is used here.
# See b/357138252 for more details.
accelerator = accelerate.Accelerator()
logging.info("using distributed_type %s", accelerator.distributed_type)
if device_map is None:
if accelerator.distributed_type == DistributedType.MULTI_GPU:
# https://github.com/artidoro/qlora/issues/186#issuecomment-1943045599
# and b/342038175.
device_map = {"": accelerator.process_index}
elif accelerator.distributed_type == DistributedType.DEEPSPEED:
# Deepspeed Zero3 does not allow setting device_map.
# https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/modeling_utils.py#L2941-L2943
device_map = None
elif accelerator.distributed_type == DistributedType.FSDP:
if precision_mode in [
constants.PRECISION_MODE_4,
constants.PRECISION_MODE_8,
]:
device_map = trl.get_kbit_device_map()
else:
device_map = None
elif (
accelerator.distributed_type == DistributedType.NO
and torch.cuda.device_count() > 1
):
# Setting device map to None to avoid using model parallelism (MP) when
# there are multiple GPUs, which can have very inefficient GPU utilization
# (b/342252819). This setting should trigger torch's nn.DataParallel
# instead, which has better GPU utilization.
device_map = None
else:
device_map = "auto"
logging.info("using device_map %s", device_map)
if train_precision == constants.PRECISION_MODE_32:
train_dtype = torch.float32
elif train_precision == constants.PRECISION_MODE_16:
train_dtype = torch.float16
elif train_precision == constants.PRECISION_MODE_16B:
train_dtype = torch.bfloat16
else:
train_dtype = "auto"
quantization_config = None
# Note: use_cache is False when enable gradient checkpointing.
if precision_mode == constants.PRECISION_MODE_32:
torch_dtype = torch.float32
elif precision_mode == constants.PRECISION_MODE_16:
torch_dtype = torch.float16
elif precision_mode == constants.PRECISION_MODE_16B:
torch_dtype = torch.bfloat16
elif precision_mode == constants.PRECISION_MODE_8:
quantization_config = BitsAndBytesConfig(
load_in_8bit=True, int8_threshold=0
)
torch_dtype = train_dtype
elif precision_mode == constants.PRECISION_MODE_4:
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=train_dtype,
)
# `bnb_4bit_quant_storage` must be set when using FSDP.
# https://huggingface.co/docs/bitsandbytes/main/en/fsdp_qlora
if accelerator.distributed_type == DistributedType.FSDP:
quantization_config.bnb_4bit_quant_storage = train_dtype
torch_dtype = train_dtype
else:
raise ValueError(f"Invalid precision mode: {precision_mode}")
logging.info("using torch_type=%s", torch_dtype)
model = AutoModelForCausalLM.from_pretrained(
pretrained_model_id,
use_cache=not enable_gradient_checkpointing,
device_map=device_map,
torch_dtype=torch_dtype,
quantization_config=quantization_config,
trust_remote_code=True,
token=access_token,
attn_implementation=attn_implementation,
)
if precision_mode in (constants.PRECISION_MODE_4, constants.PRECISION_MODE_8):
model = prepare_model_for_kbit_training(
model,
use_gradient_checkpointing=enable_gradient_checkpointing,
gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
)
if enable_gradient_checkpointing:
model.gradient_checkpointing_enable(
gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
)
# Flash attention only supports fp16 or bf16 [1].
# prepare_model_for_kbit_training will force cast some layers to float32 [2]
#
# [1]: https://github.com/Dao-AILab/flash-attention/issues/882
# [2]: https://github.com/huggingface/peft/blob/v0.10.0/src/peft/utils/other.py#L79-L81 # pylint: disable=line-too-long
if attn_implementation == "flash_attention_2" and precision_mode in (
constants.PRECISION_MODE_4,
constants.PRECISION_MODE_8,
):
for _, param in model.named_parameters():
if param.dtype == torch.float32:
param.data = param.data.to(torch_dtype)
if is_training:
# KV cache is useless during training
# https://stackoverflow.com/a/77408076
model.config.use_cache = False
if should_add_pad_token(pretrained_model_id):
model.resize_token_embeddings(len(tokenizer))
if is_training:
# The following is needed since we added a new token that needs to be
# learned.
# https://github.com/QwenLM/Qwen/issues/405#issuecomment-1751680291
model.enable_input_require_grads()
return model
def _merge_causal_language_model_with_lora_internal(
pretrained_model_id: str,
merge_precision_mode: str,
finetuned_lora_model_dir: str,
merged_model_output_dir: str,
access_token: Optional[str] = None,
) -> None:
"""Internal function to merges the base model with the lora adapter."""
logging.info("loading tokenizer...")
tokenizer = load_tokenizer(pretrained_model_id)
# Note: merging peft adapter requires loading model in 16 bits, so merging
# is done on CPU on purpose in case one GPU cannot hold the base model.
logging.info("loading model %s...", pretrained_model_id)
device_map = "cpu"
model = load_model(
pretrained_model_id=pretrained_model_id,
tokenizer=tokenizer,
precision_mode=merge_precision_mode,
access_token=access_token,
device_map=device_map,
is_training=False,
)
logging.info("loading LoRA model...")
model = PeftModel.from_pretrained(
model, finetuned_lora_model_dir, device_map=device_map
)
logging.info("merging base model with finetuned LoRA model...")
model = model.merge_and_unload()
logging.info("saving model to %s...", merged_model_output_dir)
model.save_pretrained(
merged_model_output_dir,
safe_serialization=False,
is_main_process=PartialState().is_main_process,
)
logging.info("saving tokenizer to %s...", merged_model_output_dir)
tokenizer.save_pretrained(
merged_model_output_dir,
is_main_process=PartialState().is_main_process,
)
def merge_causal_language_model_with_lora_fsdp(
pretrained_model_id: str,
merge_precision_mode: str,
finetuned_lora_model_dir: str,
merged_model_output_dir: str,
access_token: Optional[str] = None,
) -> None:
"""Merges the base model with the lora adapter for FSDP.
Only the main process should call this function.
Args:
pretrained_model_id: Predefined base model name or path to directory
containing model checkpoints.
merge_precision_mode: Precision mode for saving model weights.
finetuned_lora_model_dir: Path to directory containing PEFT-finetuned model
weights.
merged_model_output_dir: Path to directory to save the merged model.
access_token: Access token for accessing the model.
"""
assert PartialState().is_main_process
_merge_causal_language_model_with_lora_internal(
pretrained_model_id=pretrained_model_id,
merge_precision_mode=merge_precision_mode,
finetuned_lora_model_dir=finetuned_lora_model_dir,
merged_model_output_dir=merged_model_output_dir,
access_token=access_token,
)
def merge_causal_language_model_with_lora(
pretrained_model_id: str,
precision_mode: str,
finetuned_lora_model_dir: str,
merged_model_output_dir: str,
access_token: Optional[str] = None,
) -> None:
"""Merges the base model with the lora adapter."""
# Set merge related variables.
if precision_mode == constants.PRECISION_MODE_FP8:
# Merge as FP16. FP8 requires conversion after merge.
merge_precision_mode = constants.PRECISION_MODE_16
local_merged_model_dir = _LOCAL_MERGED_MODEL_DIR
else:
merge_precision_mode = precision_mode
local_merged_model_dir = merged_model_output_dir
if PartialState().is_main_process:
logging.info("Starting merging job...")
# When deepspeed Zero3 is enabled, users are not allowed to specify
# `device_map` when loading the model (even on CPU).
#
# To work-around this, we kick off another process (from the
# is_main_process) and set up the environment to avoid using Deepspeed when
# doing the merging.
if is_deepspeed_zero3_enabled():
ctx = mp.get_context("spawn")
os.environ["ACCELERATE_USE_DEEPSPEED"] = "false"
merge_job = ctx.Process(
target=_merge_causal_language_model_with_lora_internal,
args=(
pretrained_model_id,
merge_precision_mode,
finetuned_lora_model_dir,
local_merged_model_dir,
),
kwargs={
"access_token": access_token,
},
)
merge_job.start()
merge_job.join()
os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
else:
_merge_causal_language_model_with_lora_internal(
pretrained_model_id=pretrained_model_id,
merge_precision_mode=merge_precision_mode,
finetuned_lora_model_dir=finetuned_lora_model_dir,
merged_model_output_dir=local_merged_model_dir,
access_token=access_token,
)
logging.info("merging job is done")
# Wait for all processes to sync here.
PartialState().wait_for_everyone()
if precision_mode == constants.PRECISION_MODE_FP8:
convert_model_to_fp8(
pretrained_model_name_or_path=pretrained_model_id,
merged_model_output_dir=local_merged_model_dir,
quantized_model_output_dir=merged_model_output_dir,
access_token=access_token,
)
def convert_model_to_fp8(
pretrained_model_name_or_path: str,
merged_model_output_dir: str,
quantized_model_output_dir: str,
access_token: Optional[str] = None,
) -> None:
"""Converts the model to fp8.
Args:
pretrained_model_name_or_path: Original base model name or path.
merged_model_output_dir: Path to directory containing the merged model.
quantized_model_output_dir: Path to directory to save the quantized model.
access_token: Access token for accessing the model.
"""
if PartialState().is_main_process:
quantization_config = FbgemmFp8Config(
modules_to_not_convert=_maybe_get_modules_to_not_convert_by_model_id(
pretrained_model_name_or_path
)
)
quantized_model = AutoModelForCausalLM.from_pretrained(
merged_model_output_dir,
device_map="cpu",
quantization_config=quantization_config,
trust_remote_code=False,
token=access_token,
)
tokenizer = AutoTokenizer.from_pretrained(merged_model_output_dir)
quantized_model.save_pretrained(quantized_model_output_dir)
tokenizer.save_pretrained(quantized_model_output_dir)
PartialState().wait_for_everyone()
@dataclasses.dataclass
class TuningDataStats:
tuning_dataset_example_count: int
total_billable_token_count: int
tuning_step_count: int
def get_dataset_stats(
dataset: Any,
tokenizer: transformers.PreTrainedTokenizer,
column: str,
effective_batch_size: int,
) -> TuningDataStats:
"""Calculates dataset statistics, e.g., total number of tokens."""
tokenized_dataset = dataset.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]))
tuning_step_count = (
tuning_dataset_example_count + effective_batch_size - 1
) // effective_batch_size
return TuningDataStats(
tuning_dataset_example_count,
total_billable_token_count,
tuning_step_count,
)
def force_gc():
"""Collects garbage immediately to release unused CPU/GPU resources."""
gc.collect()
torch.cuda.empty_cache()
def should_add_pad_token(model_id: str) -> bool:
"""Returns whether the model requires adding a special pad token."""
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."""
return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
def write_kfp_outputs(
executor_input: str, output_artifacts: Dict[str, str]
) -> None:
"""Writes KFP outputs given a dict of output artifact names and URIs."""
# Only the main process writes to avoid race condition.
if PartialState().is_main_process:
executor_input = json_format.Parse(
executor_input, pipeline_spec_pb2.ExecutorInput()
)
outputs = executor_input.outputs
# set all artifacts
for name, uri in output_artifacts.items():
artifact_list = outputs.artifacts.get(name)
if not artifact_list or not artifact_list.artifacts:
raise ValueError(f"Artifact name={name} does not exist.")
artifact_list.artifacts[0].uri = uri
# write output file
executor_output = pipeline_spec_pb2.ExecutorOutput(
artifacts=outputs.artifacts
)
os.makedirs(os.path.dirname(outputs.output_file), exist_ok=True)
with open(outputs.output_file, "w") as f:
f.write(json_format.MessageToJson(executor_output, indent=None))
# Wait for the main process to finish before moving on to the next task.
PartialState().wait_for_everyone()
def upload_local_dir_to_gcs(local_dir: str, gcs_path: str):
"""Uploads local dir to GCS."""
if PartialState().is_main_process:
logging.info("uploading %s to %s...", local_dir, gcs_path)
subprocess.check_output([
"gsutil",
"-m",
"cp",
"-r",
local_dir,
gcs_path,
])
logging.info("%s uploaded.", local_dir)
PartialState().wait_for_everyone()
def write_first_party_model_metadata(output_dir: str, docker_uri: str) -> None:
"""Multi-process friendly version of fileutils.write_first_party_model_metadata."""
if PartialState().is_main_process:
fileutils.write_first_party_model_metadata(output_dir, docker_uri)
PartialState().wait_for_everyone()
@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
"""
# total memory
total_mem: 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
# total_mem, 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(mem_used_smi, occupied, unused, smi_diff, util)
def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
if stats is None:
stats = gpu_stats()
total, occupied, unused, smi_diff, util = stats
return (
f"GPU memory: {total:.2f}({occupied=:.2f}, {unused=:.2f},"
f" {smi_diff=:.2f}) GB. Utilization: {util:.2f}%"
)
def init_partial_state(
timeout: datetime.timedelta = datetime.timedelta(seconds=600),
) -> None:
"""Initializes the partial state with timeout."""
# This needs to be called before any other PartialState() calls, and
# TrainingArguments needs `use_configured_state`. See b/357970482#comment3
# for more details.
PartialState(timeout=timeout)
def print_library_versions():
if PartialState().is_main_process:
logging.info("======================")
logging.info("library versions")
logging.info("======================")
logging.info("accelerate: %s", accelerate.__version__)
logging.info("peft: %s", peft.__version__)
logging.info("transformers: %s", transformers.__version__)
logging.info("trl: %s", trl.__version__)
PartialState().wait_for_everyone()
def get_final_checkpoint_path(output_dir: str) -> str:
"""Returns the final checkpoint path."""
return os.path.join(output_dir, constants.FINAL_CHECKPOINT_DIRNAME)
def _maybe_get_modules_to_not_convert_by_model_id(
pretrained_model_name_or_path: str,
) -> Optional[Sequence[str]]:
"""Returns the modules to not convert for the model."""
if _LLAMA_3_1_405B_MODEL_ID in pretrained_model_name_or_path:
return _get_llama_3_1_405b_modules_to_not_convert()
else:
return None
def _get_llama_3_1_405b_modules_to_not_convert() -> Sequence[str]:
"""Returns the modules to not convert for Llama 3.1 405B model."""
modules_to_not_convert = ["lm_head"]
for idx in range(126):
for proj_name in ["k_proj", "o_proj", "q_proj", "v_proj"]:
modules_to_not_convert.append(f"model.layers.{idx}.self_attn.{proj_name}")
for proj_name in ["down_proj", "gate_proj", "up_proj"]:
modules_to_not_convert.append(f"model.layers.0.mlp.{proj_name}")
modules_to_not_convert.append(f"model.layers.125.mlp.{proj_name}")
return tuple(modules_to_not_convert)
community-content/vertex_model_garden/model_oss/peft/train/vmg/validate_dataset_with_template.py
+77
View File
@@ -0,0 +1,77 @@
"""Validate the dataset with the template."""
from typing import Sequence
from absl import app
from absl import flags
from util import dataset_validation_util
from vertex_vision_model_garden_peft.train.vmg import utils
from util import constants
_DATASET_NAME = flags.DEFINE_string(
'dataset_name',
None,
'The dataset name in huggingface.',
required=True,
)
_TRAIN_SPLIT_NAME = flags.DEFINE_string(
'train_split_name',
'train',
'The train split name.',
)
_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
'instruct_column_in_dataset',
constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
'The instruct column in dataset.',
)
_TEMPLATE = flags.DEFINE_string(
'template',
None,
'Template for formatting language model training data. Must be a filename'
' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
' Cloud Storage URI to a JSON file.',
required=True,
)
_VALIDATE_PERCENTAGE_OF_DATASET = flags.DEFINE_integer(
'validate_percentage_of_dataset',
None,
'The percentage of the dataset to validate with the template. If set to'
' -1, it loads the full dataset.',
)
_VALIDATE_K_ROWS_OF_DATASET = flags.DEFINE_integer(
'validate_k_rows_of_dataset',
None,
'The top k rows of the dataset to validate with the template. If set to -1,'
' it loads the full dataset.',
)
_USE_MULTIPROCESSING = flags.DEFINE_boolean(
'use_multiprocessing',
False,
'Whether to use multiprocessing for loading the dataset.',
)
def main(unused_argv: Sequence[str]) -> None:
utils.print_library_versions()
dataset_validation_util.validate_dataset_with_template(
dataset_name=_DATASET_NAME.value,
split=_TRAIN_SPLIT_NAME.value,
input_column=_INSTRUCT_COLUMN_IN_DATASET.value,
template=_TEMPLATE.value,
use_multiprocessing=_USE_MULTIPROCESSING.value,
validate_percentage_of_dataset=_VALIDATE_PERCENTAGE_OF_DATASET.value,
validate_k_rows_of_dataset=_VALIDATE_K_ROWS_OF_DATASET.value,
)
if __name__ == '__main__':
app.run(main)