diff --git a/llm-litgpt-finetuning/.assets/model.png b/llm-litgpt-finetuning/.assets/model.png new file mode 100644 index 00000000..2201c2e9 Binary files /dev/null and b/llm-litgpt-finetuning/.assets/model.png differ diff --git a/llm-lora-finetuning/.copier-answers.yml b/llm-litgpt-finetuning/.copier-answers.yml similarity index 100% rename from llm-lora-finetuning/.copier-answers.yml rename to llm-litgpt-finetuning/.copier-answers.yml diff --git a/llm-litgpt-finetuning/.dockerignore b/llm-litgpt-finetuning/.dockerignore new file mode 100644 index 00000000..496552c8 --- /dev/null +++ b/llm-litgpt-finetuning/.dockerignore @@ -0,0 +1,9 @@ +* +!/pipelines/** +!/steps/** +!/materializers/** +!/evaluate/** +!/finetune/** +!/generate/** +!/lit_gpt/** +!/scripts/** diff --git a/llm-litgpt-finetuning/LICENSE b/llm-litgpt-finetuning/LICENSE new file mode 100644 index 00000000..75d01fb4 --- /dev/null +++ b/llm-litgpt-finetuning/LICENSE @@ -0,0 +1,15 @@ +Apache Software License 2.0 + +Copyright (c) ZenML GmbH 2024. All rights reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/llm-litgpt-finetuning/README.md b/llm-litgpt-finetuning/README.md new file mode 100644 index 00000000..993fa744 --- /dev/null +++ b/llm-litgpt-finetuning/README.md @@ -0,0 +1,128 @@ +# ☮️ Fine-tuning open source LLMs using MLOps pipelines with LitGPT + +Welcome to your newly generated "ZenML LLM LitGPT Finetuning project" project! This is +a great way to get hands-on with ZenML using production-like template. +The project contains a collection of ZenML steps, pipelines and other artifacts +and useful resources that can serve as a solid starting point for finetuning open-source LLMs using ZenML. + +Using these pipelines, we can run the data-preparation and model finetuning with a single command while using YAML files for [configuration](https://docs.zenml.io/user-guide/production-guide/configure-pipeline) and letting ZenML take care of tracking our metadata and [containerizing our pipelines](https://docs.zenml.io/user-guide/advanced-guide/infrastructure-management/containerize-your-pipeline). + +
+
+ + Model version metadata + +
+
+ +## :earth_americas: Inspiration and Credit + +This project heavily relies on the [Lit-GPT project](https://github.com/Lightning-AI/litgpt) of the amazing people at Lightning AI. We used [this blogpost](https://lightning.ai/pages/community/lora-insights/#toc14) to get started with LoRA and QLoRA and modified the commands they recommend to make them work using ZenML. + +## 🏃 How to run + +In this project we provide a few predefined configuration files for finetuning models on the [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) dataset. Before we're able to run any pipeline, we need to set up our environment as follows: + +```bash +# Set up a Python virtual environment, if you haven't already +python3 -m venv .venv +source .venv/bin/activate + +# Install requirements +pip install -r requirements.txt +``` + +### Combined feature engineering and finetuning pipeline + +The easiest way to get started with just a single command is to run the finetuning pipeline with the `finetune-alpaca.yaml` configuration file, which will do both feature engineering and finetuning: + +```shell +python run.py --finetuning-pipeline --config finetune-alpaca.yaml +``` + +When running the pipeline like this, the trained adapter will be stored in the ZenML artifact store. You can optionally upload the adapter, the merged model or both by specifying the `adapter_output_repo` and `merged_output_repo` parameters in the configuration file. + + +### Evaluation pipeline + +Before running this pipeline, you will need to fill in the `adapter_repo` in the `eval.yaml` configuration file. This should point to a huggingface repository that contains the finetuned adapter you got by running the finetuning pipeline. + +```shell +python run.py --eval-pipeline --config eval.yaml +``` + +### Merging pipeline + +In case you have trained an adapter using the finetuning pipeline, you can merge it with the base model by filling in the `adapter_repo` and `output_repo` parameters in the `merge.yaml` file, and then running: + +```shell +python run.py --merge-pipeline --config merge.yaml +``` + +### Feature Engineering followed by Finetuning + +If you want to finetune your model on a different dataset, you can do so by running the feature engineering pipeline followed by the finetuning pipeline. To define your dataset, take a look at the `scripts/prepare_*` scripts and set the dataset name in the `feature-alpaca.yaml` config file. + +```shell +python run.py --feature-pipeline --config feature-alpaca.yaml +python run.py --finetuning-pipeline --config finetune-from-dataset.yaml +``` + +## ☁️ Running with a remote stack + +To finetune an LLM on remote infrastructure, you can either use a remote orchestrator or a remote step operator. Follow these steps to set up a complete remote stack: +- Register the [orchestrator](https://docs.zenml.io/stacks-and-components/component-guide/orchestrators) (or [step operator](https://docs.zenml.io/stacks-and-components/component-guide/step-operators)) and make sure to configure it in a way so that the finetuning step has access to a GPU with at least 24GB of VRAM. Check out our docs for more [details](https://docs.zenml.io/stacks-and-components/component-guide). + - To access GPUs with this amount of VRAM, you might need to increase your GPU quota ([AWS](https://docs.aws.amazon.com/servicequotas/latest/userguide/request-quota-increase.html), [GCP](https://console.cloud.google.com/iam-admin/quotas), [Azure](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-manage-quotas?view=azureml-api-2#request-quota-and-limit-increases)). + - The GPU instance that your finetuning will be running on will have CUDA drivers of a specific version installed. If that CUDA version is not compatible with the one provided by the default Docker image of the finetuning pipeline, you will need to modify it in the configuration file. See [here](https://hub.docker.com/r/pytorch/pytorch/tags) for a list of available PyTorch images. + - If you're running out of memory, you can experiment with quantized LoRA (QLoRA) by setting a different value for the `quantize` parameter in the configuration, or reduce the `global_batch_size`/`micro_batch_size`. +- Register a remote [artifact store](https://docs.zenml.io/stacks-and-components/component-guide/artifact-stores) and [container registry](https://docs.zenml.io/stacks-and-components/component-guide/container-registries). +- Register a stack with all these components + ```shell + zenml stack register llm-finetuning-stack -o \ + -a \ + -c \ + [-s ] + ``` + +## 💾 Running with custom data + +To finetune a model with your custom data, you will need to convert it to a CSV file with the columns described +[here](https://github.com/Lightning-AI/litgpt/blob/main/tutorials/prepare_dataset.md#preparing-custom-datasets-from-a-csv-file). + +Next, update the `configs/feature-custom.yaml` file and set the value of the `csv_path` parameter to that CSV file. +With all that in place, you can now run the feature engineering pipeline to convert your CSV into the correct format for training and then run the finetuning pipeline as follows: +```shell +python run.py --feature-pipeline --config feature-custom.yaml +python run.py --finetuning-pipeline --config finetune-from-dataset.yaml +``` + +## 📜 Project Structure + +The project loosely follows [the recommended ZenML project structure](https://docs.zenml.io/user-guide/starter-guide/follow-best-practices): + +``` +. +├── configs # pipeline configuration files +│ ├── eval.yaml # configuration for the evaluation pipeline +│ ├── feature-alpaca.yaml # configuration for the feature engineering pipeline +│ ├── feature-custom.yaml # configuration for the feature engineering pipeline +│ ├── finetune-alpaca.yaml # configuration for the finetuning pipeline +│ ├── finetune-from-dataset.yaml # configuration for the finetuning pipeline +│ └── merge.yaml # configuration for the merging pipeline +├── pipelines # `zenml.pipeline` implementations +│ ├── evaluate.py # Evaluation pipeline +│ ├── feature_engineering.py # Feature engineering pipeline +│ ├── finetuning.py # Finetuning pipeline +│ └── merge.py # Merging pipeline +├── steps # logically grouped `zenml.steps` implementations +│ ├── evaluate.py # evaluate model performance +│ ├── feature_engineering.py # preprocess data +│ ├── finetune.py # finetune a model +│ ├── merge.py # merge model and adapter +│ ├── params.py # shared parameters for steps +│ └── utils.py # utility functions +├── .dockerignore +├── README.md # this file +├── requirements.txt # extra Python dependencies +└── run.py # CLI tool to run pipelines on ZenML Stack +``` diff --git a/llm-lora-finetuning/configs/eval.yaml b/llm-litgpt-finetuning/configs/eval.yaml similarity index 100% rename from llm-lora-finetuning/configs/eval.yaml rename to llm-litgpt-finetuning/configs/eval.yaml diff --git a/llm-lora-finetuning/configs/feature-alpaca.yaml b/llm-litgpt-finetuning/configs/feature-alpaca.yaml similarity index 100% rename from llm-lora-finetuning/configs/feature-alpaca.yaml rename to llm-litgpt-finetuning/configs/feature-alpaca.yaml diff --git a/llm-lora-finetuning/configs/feature-custom.yaml b/llm-litgpt-finetuning/configs/feature-custom.yaml similarity index 100% rename from llm-lora-finetuning/configs/feature-custom.yaml rename to llm-litgpt-finetuning/configs/feature-custom.yaml diff --git a/llm-lora-finetuning/configs/finetune-alpaca.yaml b/llm-litgpt-finetuning/configs/finetune-alpaca.yaml similarity index 100% rename from llm-lora-finetuning/configs/finetune-alpaca.yaml rename to llm-litgpt-finetuning/configs/finetune-alpaca.yaml diff --git a/llm-lora-finetuning/configs/finetune-from-dataset.yaml b/llm-litgpt-finetuning/configs/finetune-from-dataset.yaml similarity index 100% rename from llm-lora-finetuning/configs/finetune-from-dataset.yaml rename to llm-litgpt-finetuning/configs/finetune-from-dataset.yaml diff --git a/llm-lora-finetuning/configs/merge.yaml b/llm-litgpt-finetuning/configs/merge.yaml similarity index 100% rename from llm-lora-finetuning/configs/merge.yaml rename to llm-litgpt-finetuning/configs/merge.yaml diff --git a/llm-lora-finetuning/evaluate/lm_eval_harness.py b/llm-litgpt-finetuning/evaluate/lm_eval_harness.py similarity index 100% rename from llm-lora-finetuning/evaluate/lm_eval_harness.py rename to llm-litgpt-finetuning/evaluate/lm_eval_harness.py diff --git a/llm-lora-finetuning/finetune/adapter.py b/llm-litgpt-finetuning/finetune/adapter.py similarity index 100% rename from llm-lora-finetuning/finetune/adapter.py rename to llm-litgpt-finetuning/finetune/adapter.py diff --git a/llm-lora-finetuning/finetune/adapter_v2.py b/llm-litgpt-finetuning/finetune/adapter_v2.py similarity index 100% rename from llm-lora-finetuning/finetune/adapter_v2.py rename to llm-litgpt-finetuning/finetune/adapter_v2.py diff --git a/llm-lora-finetuning/finetune/full.py b/llm-litgpt-finetuning/finetune/full.py similarity index 100% rename from llm-lora-finetuning/finetune/full.py rename to llm-litgpt-finetuning/finetune/full.py diff --git a/llm-lora-finetuning/finetune/lora.py b/llm-litgpt-finetuning/finetune/lora.py similarity index 100% rename from llm-lora-finetuning/finetune/lora.py rename to llm-litgpt-finetuning/finetune/lora.py diff --git a/llm-lora-finetuning/generate/adapter.py b/llm-litgpt-finetuning/generate/adapter.py similarity index 100% rename from llm-lora-finetuning/generate/adapter.py rename to llm-litgpt-finetuning/generate/adapter.py diff --git a/llm-lora-finetuning/generate/adapter_v2.py b/llm-litgpt-finetuning/generate/adapter_v2.py similarity index 100% rename from llm-lora-finetuning/generate/adapter_v2.py rename to llm-litgpt-finetuning/generate/adapter_v2.py diff --git a/llm-lora-finetuning/generate/base.py b/llm-litgpt-finetuning/generate/base.py similarity index 100% rename from llm-lora-finetuning/generate/base.py rename to llm-litgpt-finetuning/generate/base.py diff --git a/llm-lora-finetuning/generate/full.py b/llm-litgpt-finetuning/generate/full.py similarity index 100% rename from llm-lora-finetuning/generate/full.py rename to llm-litgpt-finetuning/generate/full.py diff --git a/llm-lora-finetuning/generate/lora.py b/llm-litgpt-finetuning/generate/lora.py similarity index 100% rename from llm-lora-finetuning/generate/lora.py rename to llm-litgpt-finetuning/generate/lora.py diff --git a/llm-lora-finetuning/generate/sequentially.py b/llm-litgpt-finetuning/generate/sequentially.py similarity index 100% rename from llm-lora-finetuning/generate/sequentially.py rename to llm-litgpt-finetuning/generate/sequentially.py diff --git a/llm-lora-finetuning/generate/tp.py b/llm-litgpt-finetuning/generate/tp.py similarity index 100% rename from llm-lora-finetuning/generate/tp.py rename to llm-litgpt-finetuning/generate/tp.py diff --git a/llm-lora-finetuning/lit_gpt/__init__.py b/llm-litgpt-finetuning/lit_gpt/__init__.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/__init__.py rename to llm-litgpt-finetuning/lit_gpt/__init__.py diff --git a/llm-lora-finetuning/lit_gpt/adapter.py b/llm-litgpt-finetuning/lit_gpt/adapter.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/adapter.py rename to llm-litgpt-finetuning/lit_gpt/adapter.py diff --git a/llm-lora-finetuning/lit_gpt/adapter_v2.py b/llm-litgpt-finetuning/lit_gpt/adapter_v2.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/adapter_v2.py rename to llm-litgpt-finetuning/lit_gpt/adapter_v2.py diff --git a/llm-lora-finetuning/lit_gpt/args.py b/llm-litgpt-finetuning/lit_gpt/args.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/args.py rename to llm-litgpt-finetuning/lit_gpt/args.py diff --git a/llm-lora-finetuning/lit_gpt/config.py b/llm-litgpt-finetuning/lit_gpt/config.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/config.py rename to llm-litgpt-finetuning/lit_gpt/config.py diff --git a/llm-lora-finetuning/lit_gpt/lora.py b/llm-litgpt-finetuning/lit_gpt/lora.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/lora.py rename to llm-litgpt-finetuning/lit_gpt/lora.py diff --git a/llm-lora-finetuning/lit_gpt/model.py b/llm-litgpt-finetuning/lit_gpt/model.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/model.py rename to llm-litgpt-finetuning/lit_gpt/model.py diff --git a/llm-lora-finetuning/lit_gpt/packed_dataset.py b/llm-litgpt-finetuning/lit_gpt/packed_dataset.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/packed_dataset.py rename to llm-litgpt-finetuning/lit_gpt/packed_dataset.py diff --git a/llm-lora-finetuning/lit_gpt/rmsnorm.py b/llm-litgpt-finetuning/lit_gpt/rmsnorm.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/rmsnorm.py rename to llm-litgpt-finetuning/lit_gpt/rmsnorm.py diff --git a/llm-lora-finetuning/lit_gpt/tokenizer.py b/llm-litgpt-finetuning/lit_gpt/tokenizer.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/tokenizer.py rename to llm-litgpt-finetuning/lit_gpt/tokenizer.py diff --git a/llm-lora-finetuning/lit_gpt/utils.py b/llm-litgpt-finetuning/lit_gpt/utils.py similarity index 100% rename from llm-lora-finetuning/lit_gpt/utils.py rename to llm-litgpt-finetuning/lit_gpt/utils.py diff --git a/llm-litgpt-finetuning/materializers/__init__.py b/llm-litgpt-finetuning/materializers/__init__.py new file mode 100644 index 00000000..757bd841 --- /dev/null +++ b/llm-litgpt-finetuning/materializers/__init__.py @@ -0,0 +1,16 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/llm-litgpt-finetuning/materializers/directory_materializer.py b/llm-litgpt-finetuning/materializers/directory_materializer.py new file mode 100644 index 00000000..4adc7b4a --- /dev/null +++ b/llm-litgpt-finetuning/materializers/directory_materializer.py @@ -0,0 +1,71 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +from pathlib import Path +from tempfile import mkdtemp +from typing import Any, ClassVar, Tuple, Type + +from zenml.enums import ArtifactType +from zenml.io import fileio +from zenml.materializers.base_materializer import BaseMaterializer + + +class DirectoryMaterializer(BaseMaterializer): + """Materializer to store local directories in the artifact store.""" + + ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (Path,) + ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA + + def load(self, data_type: Type[Any]) -> Any: + """Copy the artifact files to a local temp directory. + + Args: + data_type: Unused. + + Returns: + Path to the local directory that contains the artifact files. + """ + directory = mkdtemp(prefix="zenml-artifact") + self._copy_directory(src=self.uri, dst=directory) + return Path(directory) + + def save(self, data: Any) -> None: + """Store the directory in the artifact store. + + Args: + data: Path to a local directory to store. + """ + assert isinstance(data, Path) + self._copy_directory(src=str(data), dst=self.uri) + + @staticmethod + def _copy_directory(src: str, dst: str) -> None: + """Recursively copy a directory. + + Args: + src: The directory to copy. + dst: Where to copy the directory to. + """ + for src_dir, _, files in fileio.walk(src): + dst_dir = os.path.join(dst, os.path.relpath(src_dir, src)) + fileio.makedirs(dst_dir) + + for file in files: + src_file = os.path.join(src_dir, file) + dst_file = os.path.join(dst_dir, file) + fileio.copy(src_file, dst_file) diff --git a/llm-lora-finetuning/pipelines/__init__.py b/llm-litgpt-finetuning/pipelines/__init__.py similarity index 100% rename from llm-lora-finetuning/pipelines/__init__.py rename to llm-litgpt-finetuning/pipelines/__init__.py diff --git a/llm-lora-finetuning/pipelines/evaluate.py b/llm-litgpt-finetuning/pipelines/evaluate.py similarity index 100% rename from llm-lora-finetuning/pipelines/evaluate.py rename to llm-litgpt-finetuning/pipelines/evaluate.py diff --git a/llm-lora-finetuning/pipelines/feature_engineering.py b/llm-litgpt-finetuning/pipelines/feature_engineering.py similarity index 100% rename from llm-lora-finetuning/pipelines/feature_engineering.py rename to llm-litgpt-finetuning/pipelines/feature_engineering.py diff --git a/llm-lora-finetuning/pipelines/finetuning.py b/llm-litgpt-finetuning/pipelines/finetuning.py similarity index 100% rename from llm-lora-finetuning/pipelines/finetuning.py rename to llm-litgpt-finetuning/pipelines/finetuning.py diff --git a/llm-lora-finetuning/pipelines/merge.py b/llm-litgpt-finetuning/pipelines/merge.py similarity index 100% rename from llm-lora-finetuning/pipelines/merge.py rename to llm-litgpt-finetuning/pipelines/merge.py diff --git a/llm-litgpt-finetuning/requirements.txt b/llm-litgpt-finetuning/requirements.txt new file mode 100644 index 00000000..ad19fe96 --- /dev/null +++ b/llm-litgpt-finetuning/requirements.txt @@ -0,0 +1,17 @@ +zenml +torch>=2.2.0 +lightning @ git+https://github.com/Lightning-AI/lightning@ed367ca675861cdf40dbad2e4d66f7eee2ec50af +jsonargparse[signatures] # CLI +bitsandbytes==0.41.0 # quantization +scipy # required by bitsandbytes +sentencepiece # llama-based models +tokenizers # pythia, falcon, redpajama +datasets # eval +requests # scripts/prepare_* +zstandard # scripts/prepare_redpajama.py, scripts/prepare_starcoder.py +pandas # scripts/prepare_csv.py, scripts/prepare_starcoder.py +pyarrow # scripts/prepare_starcoder.py +# eval +git+https://github.com/EleutherAI/lm-evaluation-harness.git@115206dc89dad67b8beaa90051fb52db77f0a529 +# scripts/prepare_slimpajama.py, scripts/prepare_starcoder.py, pretrain/tinyllama.py +lightning[data] @ git+https://github.com/Lightning-AI/lightning@ed367ca675861cdf40dbad2e4d66f7eee2ec50af diff --git a/llm-litgpt-finetuning/run.py b/llm-litgpt-finetuning/run.py new file mode 100644 index 00000000..c8faf2a8 --- /dev/null +++ b/llm-litgpt-finetuning/run.py @@ -0,0 +1,131 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +from typing import Optional + +import click +from pipelines import ( + llm_lora_evaluation, + llm_lora_feature_engineering, + llm_lora_finetuning, + llm_lora_merging, +) +from zenml.logger import get_logger + +logger = get_logger(__name__) + + +@click.command( + help=""" +ZenML LLM Finetuning project CLI v0.1.0. + +Run the ZenML LLM Finetuning project LLM LoRA finetuning pipelines. + +Examples: + + \b + # Run the feature feature engineering pipeline + python run.py --feature-pipeline + + \b + # Run the finetuning pipeline + python run.py --finetuning-pipeline + + \b + # Run the merging pipeline + python run.py --merging-pipeline + + \b + # Run the evaluation pipeline + python run.py --eval-pipeline +""" +) +@click.option( + "--config", + type=str, + default=None, + help="Path to the YAML config file.", +) +@click.option( + "--feature-pipeline", + is_flag=True, + default=False, + help="Whether to run the pipeline that creates the dataset.", +) +@click.option( + "--finetuning-pipeline", + is_flag=True, + default=False, + help="Whether to run the pipeline that finetunes the model.", +) +@click.option( + "--merging-pipeline", + is_flag=True, + default=False, + help="Whether to run the pipeline that merges the model and adapter.", +) +@click.option( + "--eval-pipeline", + is_flag=True, + default=False, + help="Whether to run the pipeline that evaluates the model.", +) +@click.option( + "--no-cache", + is_flag=True, + default=False, + help="Disable caching for the pipeline run.", +) +def main( + config: Optional[str] = None, + feature_pipeline: bool = False, + finetuning_pipeline: bool = False, + merging_pipeline: bool = False, + eval_pipeline: bool = False, + no_cache: bool = False, +): + """Main entry point for the pipeline execution. + + Args: + no_cache: If `True` cache will be disabled. + """ + config_folder = os.path.join( + os.path.dirname(os.path.realpath(__file__)), + "configs", + ) + pipeline_args = {"enable_cache": not no_cache} + if not config: + raise RuntimeError("Config file is required to run a pipeline.") + + pipeline_args["config_path"] = os.path.join(config_folder, config) + + if feature_pipeline: + llm_lora_feature_engineering.with_options(**pipeline_args)() + + if finetuning_pipeline: + llm_lora_finetuning.with_options(**pipeline_args)() + + if merging_pipeline: + llm_lora_merging.with_options(**pipeline_args)() + + if eval_pipeline: + llm_lora_evaluation.with_options(**pipeline_args)() + + +if __name__ == "__main__": + main() diff --git a/llm-lora-finetuning/scripts/convert_hf_checkpoint.py b/llm-litgpt-finetuning/scripts/convert_hf_checkpoint.py similarity index 100% rename from llm-lora-finetuning/scripts/convert_hf_checkpoint.py rename to llm-litgpt-finetuning/scripts/convert_hf_checkpoint.py diff --git a/llm-lora-finetuning/scripts/convert_lit_checkpoint.py b/llm-litgpt-finetuning/scripts/convert_lit_checkpoint.py similarity index 100% rename from llm-lora-finetuning/scripts/convert_lit_checkpoint.py rename to llm-litgpt-finetuning/scripts/convert_lit_checkpoint.py diff --git a/llm-lora-finetuning/scripts/convert_pretrained_checkpoint.py b/llm-litgpt-finetuning/scripts/convert_pretrained_checkpoint.py similarity index 100% rename from llm-lora-finetuning/scripts/convert_pretrained_checkpoint.py rename to llm-litgpt-finetuning/scripts/convert_pretrained_checkpoint.py diff --git a/llm-lora-finetuning/scripts/download.py b/llm-litgpt-finetuning/scripts/download.py similarity index 100% rename from llm-lora-finetuning/scripts/download.py rename to llm-litgpt-finetuning/scripts/download.py diff --git a/llm-lora-finetuning/scripts/merge_lora.py b/llm-litgpt-finetuning/scripts/merge_lora.py similarity index 100% rename from llm-lora-finetuning/scripts/merge_lora.py rename to llm-litgpt-finetuning/scripts/merge_lora.py diff --git a/llm-lora-finetuning/scripts/prepare_alpaca.py b/llm-litgpt-finetuning/scripts/prepare_alpaca.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_alpaca.py rename to llm-litgpt-finetuning/scripts/prepare_alpaca.py diff --git a/llm-lora-finetuning/scripts/prepare_csv.py b/llm-litgpt-finetuning/scripts/prepare_csv.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_csv.py rename to llm-litgpt-finetuning/scripts/prepare_csv.py diff --git a/llm-lora-finetuning/scripts/prepare_dolly.py b/llm-litgpt-finetuning/scripts/prepare_dolly.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_dolly.py rename to llm-litgpt-finetuning/scripts/prepare_dolly.py diff --git a/llm-lora-finetuning/scripts/prepare_flan.py b/llm-litgpt-finetuning/scripts/prepare_flan.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_flan.py rename to llm-litgpt-finetuning/scripts/prepare_flan.py diff --git a/llm-lora-finetuning/scripts/prepare_lima.py b/llm-litgpt-finetuning/scripts/prepare_lima.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_lima.py rename to llm-litgpt-finetuning/scripts/prepare_lima.py diff --git a/llm-lora-finetuning/scripts/prepare_longform.py b/llm-litgpt-finetuning/scripts/prepare_longform.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_longform.py rename to llm-litgpt-finetuning/scripts/prepare_longform.py diff --git a/llm-lora-finetuning/scripts/prepare_openwebtext.py b/llm-litgpt-finetuning/scripts/prepare_openwebtext.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_openwebtext.py rename to llm-litgpt-finetuning/scripts/prepare_openwebtext.py diff --git a/llm-lora-finetuning/scripts/prepare_redpajama.py b/llm-litgpt-finetuning/scripts/prepare_redpajama.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_redpajama.py rename to llm-litgpt-finetuning/scripts/prepare_redpajama.py diff --git a/llm-lora-finetuning/scripts/prepare_slimpajama.py b/llm-litgpt-finetuning/scripts/prepare_slimpajama.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_slimpajama.py rename to llm-litgpt-finetuning/scripts/prepare_slimpajama.py diff --git a/llm-lora-finetuning/scripts/prepare_starcoder.py b/llm-litgpt-finetuning/scripts/prepare_starcoder.py similarity index 100% rename from llm-lora-finetuning/scripts/prepare_starcoder.py rename to llm-litgpt-finetuning/scripts/prepare_starcoder.py diff --git a/llm-litgpt-finetuning/steps/__init__.py b/llm-litgpt-finetuning/steps/__init__.py new file mode 100644 index 00000000..c9630597 --- /dev/null +++ b/llm-litgpt-finetuning/steps/__init__.py @@ -0,0 +1,21 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from steps.evaluate import evaluate +from steps.feature_engineering import feature_engineering +from steps.finetune import finetune +from steps.merge import merge diff --git a/llm-lora-finetuning/steps/evaluate.py b/llm-litgpt-finetuning/steps/evaluate.py similarity index 100% rename from llm-lora-finetuning/steps/evaluate.py rename to llm-litgpt-finetuning/steps/evaluate.py diff --git a/llm-lora-finetuning/steps/feature_engineering.py b/llm-litgpt-finetuning/steps/feature_engineering.py similarity index 100% rename from llm-lora-finetuning/steps/feature_engineering.py rename to llm-litgpt-finetuning/steps/feature_engineering.py diff --git a/llm-litgpt-finetuning/steps/finetune.py b/llm-litgpt-finetuning/steps/finetune.py new file mode 100644 index 00000000..6168c263 --- /dev/null +++ b/llm-litgpt-finetuning/steps/finetune.py @@ -0,0 +1,249 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import shutil +from pathlib import Path +from typing import Literal, Optional + +import torch +from finetune.lora import setup +from huggingface_hub import upload_folder +from lit_gpt.args import EvalArgs, IOArgs, TrainArgs +from materializers.directory_materializer import DirectoryMaterializer +from pydantic import BaseModel +from typing_extensions import Annotated +from zenml import get_step_context, log_model_metadata, step +from zenml.logger import get_logger +from zenml.materializers import BuiltInMaterializer + +from scripts.convert_lit_checkpoint import convert_lit_checkpoint +from scripts.download import download_from_hub +from scripts.merge_lora import merge_lora +from scripts.prepare_alpaca import prepare +from steps.params import LoraParameters +from steps.utils import ( + convert_to_lit_checkpoint_if_necessary, + get_huggingface_access_token, +) + +logger = get_logger(__file__) + + +class DataParameters(BaseModel): + """Data preprocessing parameters.""" + + seed: int = 42 + test_split_fraction: float = 0.03865 + mask_inputs: bool = False + ignore_index: int = -1 + max_seq_length: Optional[int] = None + + +class TrainingParameters(BaseModel): + """Training parameters.""" + + save_interval: int = 1000 + log_interval: int = 1 + global_batch_size: int = 64 + micro_batch_size: int = 4 + lr_warmup_steps: int = 100 + epochs: Optional[int] = None + epoch_size: Optional[int] = None + max_tokens: Optional[int] = None + max_seq_length: Optional[int] = None + + learning_rate: float = 1e-3 + weight_decay: float = 0.02 + beta1: float = 0.9 + beta2: float = 0.95 + max_norm: Optional[float] = None + min_lr: float = 6e-5 + + +class EvalParameters(BaseModel): + """Mid-training evaluation parameters.""" + + interval: int = 100 + max_new_tokens: int = 100 + max_iters: int = 100 + + +class FinetuningParameters(BaseModel): + """Parameters for the finetuning step.""" + + base_model_repo: str + from_safetensors: bool = False + + adapter_output_repo: Optional[str] = None + merged_output_repo: Optional[str] = None + convert_to_hf_checkpoint: bool = False + + precision: Optional[str] = None + quantize: Optional[ + Literal[ + "bnb.nf4", + "bnb.nf4-dq", + "bnb.fp4", + "bnb.fp4-dq", + "bnb.int8-training", + ] + ] = None + + data: DataParameters = DataParameters() + training: TrainingParameters = TrainingParameters() + eval: EvalParameters = EvalParameters() + lora: LoraParameters = LoraParameters() + + +@step(output_materializers=[DirectoryMaterializer, BuiltInMaterializer]) +def finetune( + config: FinetuningParameters, dataset_directory: Optional[Path] = None +) -> Annotated[Optional[Path], "adapter"]: + """Finetune model using LoRA. + + Args: + config: Configuration for this step. + """ + torch.set_float32_matmul_precision("high") + + access_token = get_huggingface_access_token() + + checkpoint_root_dir = Path("checkpoints") + checkpoint_dir = checkpoint_root_dir / config.base_model_repo + + if checkpoint_dir.exists(): + logger.info( + "Checkpoint directory already exists, skipping download..." + ) + else: + download_from_hub( + repo_id=config.base_model_repo, + from_safetensors=config.from_safetensors, + checkpoint_dir=checkpoint_root_dir, + access_token=access_token, + ) + + convert_to_lit_checkpoint_if_necessary(checkpoint_dir=checkpoint_dir) + + if dataset_directory: + try: + dataset_name = ( + get_step_context() + .inputs["dataset_directory"] + .run_metadata["dataset_name"] + .value + ) + except KeyError: + dataset_name = "unknown_dataset" + else: + dataset_directory = Path("data/alpaca") + dataset_name = dataset_directory.name + prepare( + destination_path=dataset_directory, + checkpoint_dir=checkpoint_dir, + test_split_fraction=config.data.test_split_fraction, + seed=config.data.seed, + mask_inputs=config.data.mask_inputs, + ignore_index=config.data.ignore_index, + max_seq_length=config.data.max_seq_length, + ) + + model_name = checkpoint_dir.name + + log_model_metadata( + metadata={"model_name": model_name, "dataset_name": dataset_name} + ) + adapter_output_dir = Path("output/lora") / dataset_name / model_name + + io_args = IOArgs( + train_data_dir=dataset_directory, + val_data_dir=dataset_directory, + checkpoint_dir=checkpoint_dir, + out_dir=adapter_output_dir, + ) + train_args = TrainArgs(**config.training.dict()) + eval_args = EvalArgs(**config.eval.dict()) + setup( + devices=1, + io=io_args, + train=train_args, + eval=eval_args, + precision=config.precision, + quantize=config.quantize, + **config.lora.dict(), + ) + + if config.merged_output_repo: + lora_path = adapter_output_dir / "lit_model_lora_finetuned.pth" + + merge_output_dir = ( + Path("output/lora_merged") / dataset_name / model_name + ) + merge_lora( + lora_path=lora_path, + checkpoint_dir=checkpoint_dir, + out_dir=merge_output_dir, + precision=config.precision, + **config.lora.dict(), + ) + + for path in Path(checkpoint_dir).glob("*.json"): + destination = Path(merge_output_dir) / path.name + shutil.copy(src=path, dst=destination) + + if config.convert_to_hf_checkpoint: + upload_dir = ( + Path("output/lora_merged_hf") / dataset_name / model_name + ) + upload_dir.mkdir(parents=True, exist_ok=True) + convert_lit_checkpoint( + checkpoint_path=config.merged_output_repo / "lit_model.pth", + config_path=config.merged_output_repo / "lit_config.json", + output_path=upload_dir / "pytorch_model", + ) + else: + upload_dir = merge_output_dir + + commit = upload_folder( + repo_id=config.merged_output_repo, + folder_path=upload_dir, + token=access_token, + ) + log_model_metadata( + metadata={ + "merged_model_huggingface_commit_hash": commit.oid, + "merged_model_huggingface_commit_url": commit.commit_url, + } + ) + + if config.adapter_output_repo: + commit = upload_folder( + repo_id=config.adapter_output_repo, + folder_path=adapter_output_dir, + token=access_token, + ) + log_model_metadata( + metadata={ + "adapter_huggingface_commit_hash": commit.oid, + "adapter_huggingface_commit_url": commit.commit_url, + } + ) + return None + else: + # If the adapter should not be uploaded to the HF Hub, we store it + # in the artifact store + return adapter_output_dir diff --git a/llm-lora-finetuning/steps/merge.py b/llm-litgpt-finetuning/steps/merge.py similarity index 100% rename from llm-lora-finetuning/steps/merge.py rename to llm-litgpt-finetuning/steps/merge.py diff --git a/llm-lora-finetuning/steps/params.py b/llm-litgpt-finetuning/steps/params.py similarity index 100% rename from llm-lora-finetuning/steps/params.py rename to llm-litgpt-finetuning/steps/params.py diff --git a/llm-lora-finetuning/steps/utils.py b/llm-litgpt-finetuning/steps/utils.py similarity index 100% rename from llm-lora-finetuning/steps/utils.py rename to llm-litgpt-finetuning/steps/utils.py diff --git a/llm-lora-finetuning/.assets/model.png b/llm-lora-finetuning/.assets/model.png index 2201c2e9..c58dd211 100644 Binary files a/llm-lora-finetuning/.assets/model.png and b/llm-lora-finetuning/.assets/model.png differ diff --git a/llm-lora-finetuning/.assets/pipeline.png b/llm-lora-finetuning/.assets/pipeline.png new file mode 100644 index 00000000..6e49ee7a Binary files /dev/null and b/llm-lora-finetuning/.assets/pipeline.png differ diff --git a/llm-lora-finetuning/.dockerignore b/llm-lora-finetuning/.dockerignore index 496552c8..c43c482f 100644 --- a/llm-lora-finetuning/.dockerignore +++ b/llm-lora-finetuning/.dockerignore @@ -1,9 +1,5 @@ * +!/materializers/** !/pipelines/** !/steps/** -!/materializers/** -!/evaluate/** -!/finetune/** -!/generate/** -!/lit_gpt/** -!/scripts/** +!/utils/** diff --git a/llm-lora-finetuning/README.md b/llm-lora-finetuning/README.md index 94267385..b4567b3b 100644 --- a/llm-lora-finetuning/README.md +++ b/llm-lora-finetuning/README.md @@ -1,6 +1,6 @@ -# ☮️ Fine-tuning open source LLMs using MLOps pipelines +# ☮️ Fine-tuning open source LLMs using MLOps pipelines with PEFT -Welcome to your newly generated "ZenML LLM Finetuning project" project! This is +Welcome to your newly generated "ZenML LLM PEFT Finetuning project" project! This is a great way to get hands-on with ZenML using production-like template. The project contains a collection of ZenML steps, pipelines and other artifacts and useful resources that can serve as a solid starting point for finetuning open-source LLMs using ZenML. @@ -15,13 +15,13 @@ Using these pipelines, we can run the data-preparation and model finetuning with
-## :earth_americas: Inspiration and Credit +## 🌎 Inspiration and Credit -This project heavily relies on the [Lit-GPT project](https://github.com/Lightning-AI/litgpt) of the amazing people at Lightning AI. We used [this blogpost](https://lightning.ai/pages/community/lora-insights/#toc14) to get started with LoRA and QLoRA and modified the commands they recommend to make them work using ZenML. +This project heavily relies on the [PEFT project](https://huggingface.co/docs/peft/en/index) by the amazing people at Huggingface and the [`microsoft/phi-2`](https://huggingface.co/microsoft/phi-2) model from the amazing people at Microsoft. ## 🏃 How to run -In this project we provide a few predefined configuration files for finetuning models on the [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) dataset. Before we're able to run any pipeline, we need to set up our environment as follows: +In this project, we provide a predefined configuration file to finetune models on the [Viggio](https://huggingface.co/datasets/GEM/viggo) dataset. Before we're able to run any pipeline, we need to set up our environment as follows: ```bash # Set up a Python virtual environment, if you haven't already @@ -32,41 +32,23 @@ source .venv/bin/activate pip install -r requirements.txt ``` -### Combined feature engineering and finetuning pipeline +### 👷 Combined feature engineering and finetuning pipeline -The easiest way to get started with just a single command is to run the finetuning pipeline with the `finetune-alpaca.yaml` configuration file, which will do both feature engineering and finetuning: +The easiest way to get started with just a single command is to run the finetuning pipeline with the `orchestrator_finetune.yaml` configuration file, which will do data preparation, model finetuning, evaluation with [Rouge](https://huggingface.co/spaces/evaluate-metric/rouge) and promotion: ```shell -python run.py --finetuning-pipeline --config finetune-alpaca.yaml +python run.py --config orchestrator_finetune.yaml ``` -When running the pipeline like this, the trained adapter will be stored in the ZenML artifact store. You can optionally upload the adapter, the merged model or both by specifying the `adapter_output_repo` and `merged_output_repo` parameters in the configuration file. +When running the pipeline like this, the trained model will be stored in the ZenML artifact store. - -### Evaluation pipeline - -Before running this pipeline, you will need to fill in the `adapter_repo` in the `eval.yaml` configuration file. This should point to a huggingface repository that contains the finetuned adapter you got by running the finetuning pipeline. - -```shell -python run.py --eval-pipeline --config eval.yaml -``` - -### Merging pipeline - -In case you have trained an adapter using the finetuning pipeline, you can merge it with the base model by filling in the `adapter_repo` and `output_repo` parameters in the `merge.yaml` file, and then running: - -```shell -python run.py --merge-pipeline --config merge.yaml -``` - -### Feature Engineering followed by Finetuning - -If you want to finetune your model on a different dataset, you can do so by running the feature engineering pipeline followed by the finetuning pipeline. To define your dataset, take a look at the `scripts/prepare_*` scripts and set the dataset name in the `feature-alpaca.yaml` config file. - -```shell -python run.py --feature-pipeline --config feature-alpaca.yaml -python run.py --finetuning-pipeline --config finetune-from-dataset.yaml -``` +
+
+ + Model version metadata + +
+
## ☁️ Running with a remote stack @@ -74,7 +56,6 @@ To finetune an LLM on remote infrastructure, you can either use a remote orchest - Register the [orchestrator](https://docs.zenml.io/stacks-and-components/component-guide/orchestrators) (or [step operator](https://docs.zenml.io/stacks-and-components/component-guide/step-operators)) and make sure to configure it in a way so that the finetuning step has access to a GPU with at least 24GB of VRAM. Check out our docs for more [details](https://docs.zenml.io/stacks-and-components/component-guide). - To access GPUs with this amount of VRAM, you might need to increase your GPU quota ([AWS](https://docs.aws.amazon.com/servicequotas/latest/userguide/request-quota-increase.html), [GCP](https://console.cloud.google.com/iam-admin/quotas), [Azure](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-manage-quotas?view=azureml-api-2#request-quota-and-limit-increases)). - The GPU instance that your finetuning will be running on will have CUDA drivers of a specific version installed. If that CUDA version is not compatible with the one provided by the default Docker image of the finetuning pipeline, you will need to modify it in the configuration file. See [here](https://hub.docker.com/r/pytorch/pytorch/tags) for a list of available PyTorch images. - - If you're running out of memory, you can experiment with quantized LoRA (QLoRA) by setting a different value for the `quantize` parameter in the configuration, or reduce the `global_batch_size`/`micro_batch_size`. - Register a remote [artifact store](https://docs.zenml.io/stacks-and-components/component-guide/artifact-stores) and [container registry](https://docs.zenml.io/stacks-and-components/component-guide/container-registries). - Register a stack with all these components ```shell @@ -84,18 +65,6 @@ To finetune an LLM on remote infrastructure, you can either use a remote orchest [-s ] ``` -## 💾 Running with custom data - -To finetune a model with your custom data, you will need to convert it to a CSV file with the columns described -[here](https://github.com/Lightning-AI/litgpt/blob/main/tutorials/prepare_dataset.md#preparing-custom-datasets-from-a-csv-file). - -Next, update the `configs/feature-custom.yaml` file and set the value of the `csv_path` parameter to that CSV file. -With all that in place, you can now run the feature engineering pipeline to convert your CSV into the correct format for training and then run the finetuning pipeline as follows: -```shell -python run.py --feature-pipeline --config feature-custom.yaml -python run.py --finetuning-pipeline --config finetune-from-dataset.yaml -``` - ## 📜 Project Structure The project loosely follows [the recommended ZenML project structure](https://docs.zenml.io/user-guide/starter-guide/follow-best-practices): @@ -103,24 +72,23 @@ The project loosely follows [the recommended ZenML project structure](https://do ``` . ├── configs # pipeline configuration files -│ ├── eval.yaml # configuration for the evaluation pipeline -│ ├── feature-alpaca.yaml # configuration for the feature engineering pipeline -│ ├── feature-custom.yaml # configuration for the feature engineering pipeline -│ ├── finetune-alpaca.yaml # configuration for the finetuning pipeline -│ ├── finetune-from-dataset.yaml # configuration for the finetuning pipeline -│ └── merge.yaml # configuration for the merging pipeline +│ ├── orchestrator_finetune.yaml # default local or remote orchestrator +│ └── remote_finetune.yaml # default step operator configuration +├── materializers +│ └── directory_materializer.py # custom materializer to push whole directories to the artifact store and back ├── pipelines # `zenml.pipeline` implementations -│ ├── evaluate.py # Evaluation pipeline -│ ├── feature_engineering.py # Feature engineering pipeline -│ ├── finetuning.py # Finetuning pipeline -│ └── merge.py # Merging pipeline +│ └── train.py # Finetuning and evaluation pipeline ├── steps # logically grouped `zenml.steps` implementations -│ ├── evaluate.py # evaluate model performance -│ ├── feature_engineering.py # preprocess data -│ ├── finetune.py # finetune a model -│ ├── merge.py # merge model and adapter -│ ├── params.py # shared parameters for steps -│ └── utils.py # utility functions +│ ├── evaluate_model.py # evaluate base and finetuned models using Rouge metrics +│ ├── finetune.py # finetune the base model +│ ├── prepare_datasets.py # load and tokenize dataset +│ └── promote.py # promote good models to target environment +├── utils # utility functions +│ ├── callbacks.py # custom callbacks +│ ├── cuda.py # helpers for CUDA +│ ├── loaders.py # loaders for models and data +│ ├── logging.py # logging helpers +│ └── tokenizer.py # load and tokenize ├── .dockerignore ├── README.md # this file ├── requirements.txt # extra Python dependencies diff --git a/llm-lora-finetuning/configs/orchestrator_finetune.yaml b/llm-lora-finetuning/configs/orchestrator_finetune.yaml new file mode 100644 index 00000000..9bd271f6 --- /dev/null +++ b/llm-lora-finetuning/configs/orchestrator_finetune.yaml @@ -0,0 +1,62 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +model: + name: llm-peft-microsoft-phi-2 + description: "Fine-tune `microsoft/phi-2`." + tags: + - llm + - peft + - microsoft/phi-2 + version: 300_steps + +settings: + docker: + parent_image: pytorch/pytorch:2.2.2-cuda11.8-cudnn8-runtime + requirements: requirements.txt + python_package_installer: uv + environment: + PJRT_DEVICE: CUDA + USE_TORCH_XLA: "false" + MKL_SERVICE_FORCE_INTEL: "1" + +parameters: + base_model_id: microsoft/phi-2 + use_fast: False + load_in_4bit: True + system_prompt: | + Given a target sentence construct the underlying meaning representation of the input sentence as a single function with attributes and attribute values. + This function should describe the target string accurately and the function must be one of the following ['inform', 'request', 'give_opinion', 'confirm', 'verify_attribute', 'suggest', 'request_explanation', 'recommend', 'request_attribute']. + The attributes must be one of the following: ['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating', 'genres', 'player_perspective', 'has_multiplayer', 'platforms', 'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier'] + + +steps: + prepare_data: + parameters: + dataset_name: gem/viggo + + finetune: + enable_step_logs: False + parameters: + max_steps: 300 + eval_steps: 30 + bf16: True + + promote: + parameters: + metric: rouge2 + target_stage: staging diff --git a/llm-lora-finetuning/configs/remote_finetune.yaml b/llm-lora-finetuning/configs/remote_finetune.yaml new file mode 100644 index 00000000..6fda3072 --- /dev/null +++ b/llm-lora-finetuning/configs/remote_finetune.yaml @@ -0,0 +1,69 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +model: + name: llm-peft-microsoft-phi-2 + description: "Fine-tune `microsoft/phi-2`." + tags: + - llm + - peft + - microsoft/phi-2 + version: 300_steps + +settings: + docker: + parent_image: pytorch/pytorch:2.2.2-cuda11.8-cudnn8-runtime + requirements: requirements.txt + python_package_installer: uv + environment: + PJRT_DEVICE: CUDA + USE_TORCH_XLA: "false" + MKL_SERVICE_FORCE_INTEL: "1" + +parameters: + base_model_id: microsoft/phi-2 + use_fast: False + load_in_4bit: True + system_prompt: | + Given a target sentence construct the underlying meaning representation of the input sentence as a single function with attributes and attribute values. + This function should describe the target string accurately and the function must be one of the following ['inform', 'request', 'give_opinion', 'confirm', 'verify_attribute', 'suggest', 'request_explanation', 'recommend', 'request_attribute']. + The attributes must be one of the following: ['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating', 'genres', 'player_perspective', 'has_multiplayer', 'platforms', 'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier'] + + +steps: + prepare_data: + parameters: + dataset_name: gem/viggo + + finetune: + enable_step_logs: False + step_operator: gcp_a100 + parameters: + max_steps: 300 + eval_steps: 30 + bf16: True + + evaluate_finetuned: + step_operator: gcp_a100 + + evaluate_base: + step_operator: gcp_a100 + + promote: + parameters: + metric: rouge2 + target_stage: staging diff --git a/llm-lora-finetuning/materializers/directory_materializer.py b/llm-lora-finetuning/materializers/directory_materializer.py index 4adc7b4a..05e90e43 100644 --- a/llm-lora-finetuning/materializers/directory_materializer.py +++ b/llm-lora-finetuning/materializers/directory_materializer.py @@ -16,6 +16,7 @@ # import os +import shutil from pathlib import Path from tempfile import mkdtemp from typing import Any, ClassVar, Tuple, Type @@ -52,6 +53,7 @@ def save(self, data: Any) -> None: """ assert isinstance(data, Path) self._copy_directory(src=str(data), dst=self.uri) + shutil.rmtree(data) # clean-up locally stored data @staticmethod def _copy_directory(src: str, dst: str) -> None: diff --git a/llm-lora-finetuning/pipelines/train.py b/llm-lora-finetuning/pipelines/train.py new file mode 100644 index 00000000..e6614416 --- /dev/null +++ b/llm-lora-finetuning/pipelines/train.py @@ -0,0 +1,86 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +from steps import evaluate_model, finetune, prepare_data, promote +from zenml import logging as zenml_logging +from zenml import pipeline + +zenml_logging.STEP_LOGS_STORAGE_MAX_MESSAGES = ( + 10000 # workaround for https://github.com/zenml-io/zenml/issues/2252 +) + + +@pipeline +def llm_peft_full_finetune( + system_prompt: str, + base_model_id: str, + use_fast: bool = True, + load_in_8bit: bool = False, + load_in_4bit: bool = False, +): + """Pipeline for finetuning an LLM with peft. + + It will run the following steps: + + - prepare_data: prepare the datasets and tokenize them + - finetune: finetune the model + - evaluate_model: evaluate the base and finetuned model + - promote: promote the model to the target stage, if evaluation was successful + """ + if not load_in_8bit and not load_in_4bit: + raise ValueError( + "At least one of `load_in_8bit` and `load_in_4bit` must be True." + ) + if load_in_4bit and load_in_8bit: + raise ValueError( + "Only one of `load_in_8bit` and `load_in_4bit` can be True." + ) + + datasets_dir = prepare_data( + base_model_id=base_model_id, + system_prompt=system_prompt, + use_fast=use_fast, + ) + ft_model_dir = finetune( + base_model_id, + datasets_dir, + use_fast=use_fast, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + ) + evaluate_model( + base_model_id, + system_prompt, + datasets_dir, + ft_model_dir, + use_fast=use_fast, + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + id="evaluate_finetuned", + ) + evaluate_model( + base_model_id, + system_prompt, + datasets_dir, + None, + use_fast=use_fast, + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + id="evaluate_base", + ) + promote(after=["evaluate_finetuned", "evaluate_base"]) diff --git a/llm-lora-finetuning/requirements.txt b/llm-lora-finetuning/requirements.txt index ad19fe96..d270781a 100644 --- a/llm-lora-finetuning/requirements.txt +++ b/llm-lora-finetuning/requirements.txt @@ -1,17 +1,11 @@ zenml torch>=2.2.0 -lightning @ git+https://github.com/Lightning-AI/lightning@ed367ca675861cdf40dbad2e4d66f7eee2ec50af -jsonargparse[signatures] # CLI -bitsandbytes==0.41.0 # quantization -scipy # required by bitsandbytes -sentencepiece # llama-based models -tokenizers # pythia, falcon, redpajama -datasets # eval -requests # scripts/prepare_* -zstandard # scripts/prepare_redpajama.py, scripts/prepare_starcoder.py -pandas # scripts/prepare_csv.py, scripts/prepare_starcoder.py -pyarrow # scripts/prepare_starcoder.py -# eval -git+https://github.com/EleutherAI/lm-evaluation-harness.git@115206dc89dad67b8beaa90051fb52db77f0a529 -# scripts/prepare_slimpajama.py, scripts/prepare_starcoder.py, pretrain/tinyllama.py -lightning[data] @ git+https://github.com/Lightning-AI/lightning@ed367ca675861cdf40dbad2e4d66f7eee2ec50af +datasets +transformers +peft +bitsandbytes>=0.41.3 +scipy +evaluate +rouge_score +nltk +accelerate \ No newline at end of file diff --git a/llm-lora-finetuning/run.py b/llm-lora-finetuning/run.py index c8faf2a8..9d7aa67a 100644 --- a/llm-lora-finetuning/run.py +++ b/llm-lora-finetuning/run.py @@ -19,72 +19,32 @@ from typing import Optional import click -from pipelines import ( - llm_lora_evaluation, - llm_lora_feature_engineering, - llm_lora_finetuning, - llm_lora_merging, -) -from zenml.logger import get_logger - -logger = get_logger(__name__) +from pipelines.train import llm_peft_full_finetune @click.command( help=""" -ZenML LLM Finetuning project CLI v0.1.0. +ZenML LLM Finetuning project CLI v0.2.0. -Run the ZenML LLM Finetuning project LLM LoRA finetuning pipelines. +Run the ZenML LLM Finetuning project LLM PEFT finetuning pipelines. Examples: \b - # Run the feature feature engineering pipeline - python run.py --feature-pipeline + # Run the pipeline + python run.py \b - # Run the finetuning pipeline - python run.py --finetuning-pipeline - - \b - # Run the merging pipeline - python run.py --merging-pipeline - - \b - # Run the evaluation pipeline - python run.py --eval-pipeline + # Run the pipeline with custom config + python run.py --config custom_finetune.yaml """ ) @click.option( "--config", type=str, - default=None, + default="default_finetune.yaml", help="Path to the YAML config file.", ) -@click.option( - "--feature-pipeline", - is_flag=True, - default=False, - help="Whether to run the pipeline that creates the dataset.", -) -@click.option( - "--finetuning-pipeline", - is_flag=True, - default=False, - help="Whether to run the pipeline that finetunes the model.", -) -@click.option( - "--merging-pipeline", - is_flag=True, - default=False, - help="Whether to run the pipeline that merges the model and adapter.", -) -@click.option( - "--eval-pipeline", - is_flag=True, - default=False, - help="Whether to run the pipeline that evaluates the model.", -) @click.option( "--no-cache", is_flag=True, @@ -93,10 +53,6 @@ ) def main( config: Optional[str] = None, - feature_pipeline: bool = False, - finetuning_pipeline: bool = False, - merging_pipeline: bool = False, - eval_pipeline: bool = False, no_cache: bool = False, ): """Main entry point for the pipeline execution. @@ -114,17 +70,7 @@ def main( pipeline_args["config_path"] = os.path.join(config_folder, config) - if feature_pipeline: - llm_lora_feature_engineering.with_options(**pipeline_args)() - - if finetuning_pipeline: - llm_lora_finetuning.with_options(**pipeline_args)() - - if merging_pipeline: - llm_lora_merging.with_options(**pipeline_args)() - - if eval_pipeline: - llm_lora_evaluation.with_options(**pipeline_args)() + llm_peft_full_finetune.with_options(**pipeline_args)() if __name__ == "__main__": diff --git a/llm-lora-finetuning/steps/__init__.py b/llm-lora-finetuning/steps/__init__.py index c9630597..b8bfdaa3 100644 --- a/llm-lora-finetuning/steps/__init__.py +++ b/llm-lora-finetuning/steps/__init__.py @@ -15,7 +15,7 @@ # limitations under the License. # -from steps.evaluate import evaluate -from steps.feature_engineering import feature_engineering -from steps.finetune import finetune -from steps.merge import merge +from .evaluate_model import evaluate_model +from .finetune import finetune +from .prepare_datasets import prepare_data +from .promote import promote diff --git a/llm-lora-finetuning/steps/evaluate_model.py b/llm-lora-finetuning/steps/evaluate_model.py new file mode 100644 index 00000000..dc504bd8 --- /dev/null +++ b/llm-lora-finetuning/steps/evaluate_model.py @@ -0,0 +1,112 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pathlib import Path +from typing import Optional + +import evaluate +import torch +from datasets import load_from_disk +from utils.loaders import ( + load_base_model, + load_pretrained_model, +) +from utils.tokenizer import load_tokenizer, tokenize_for_eval +from zenml import log_model_metadata, save_artifact, step +from zenml.logger import get_logger + +logger = get_logger(__name__) + + +@step +def evaluate_model( + base_model_id: str, + system_prompt: str, + datasets_dir: Path, + ft_model_dir: Optional[Path], + use_fast: bool = True, + load_in_4bit: bool = False, + load_in_8bit: bool = False, +) -> None: + """Evaluate the model with ROUGE metrics. + + Args: + base_model_id: The base model id to use. + system_prompt: The system prompt to use. + datasets_dir: The path to the datasets directory. + ft_model_dir: The path to the finetuned model directory. If None, the + base model will be used. + use_fast: Whether to use the fast tokenizer. + load_in_4bit: Whether to load the model in 4bit mode. + load_in_8bit: Whether to load the model in 8bit mode. + """ + logger.info("Evaluating model...") + + logger.info("Loading dataset...") + tokenizer = load_tokenizer( + base_model_id, + is_eval=True, + use_fast=use_fast, + ) + test_dataset = load_from_disk(datasets_dir / "test_raw") + test_dataset = test_dataset[:50] + ground_truths = test_dataset["meaning_representation"] + tokenized_train_dataset = tokenize_for_eval( + test_dataset, tokenizer, system_prompt + ) + + if ft_model_dir is None: + logger.info("Generating using base model...") + model = load_base_model( + base_model_id, + is_training=False, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + ) + else: + logger.info("Generating using finetuned model...") + model = load_pretrained_model( + ft_model_dir, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + ) + + model.eval() + with torch.no_grad(): + predictions = model.generate( + input_ids=tokenized_train_dataset["input_ids"], + attention_mask=tokenized_train_dataset["attention_mask"], + max_new_tokens=100, + pad_token_id=2, + ) + predictions = tokenizer.batch_decode( + predictions[:, tokenized_train_dataset["input_ids"].shape[1] :], + skip_special_tokens=True, + ) + + logger.info("Computing ROUGE metrics...") + prefix = "base_model_" if ft_model_dir is None else "finetuned_model_" + rouge = evaluate.load("rouge") + rouge_metrics = rouge.compute( + predictions=predictions, references=ground_truths + ) + metadata = {prefix + k: float(v) for k, v in rouge_metrics.items()} + + log_model_metadata(metadata) + logger.info("Computed metrics: " + str(metadata)) + + save_artifact(rouge_metrics, prefix + "rouge_metrics") diff --git a/llm-lora-finetuning/steps/finetune.py b/llm-lora-finetuning/steps/finetune.py index 6168c263..01171f90 100644 --- a/llm-lora-finetuning/steps/finetune.py +++ b/llm-lora-finetuning/steps/finetune.py @@ -15,235 +15,127 @@ # limitations under the License. # -import shutil from pathlib import Path -from typing import Literal, Optional -import torch -from finetune.lora import setup -from huggingface_hub import upload_folder -from lit_gpt.args import EvalArgs, IOArgs, TrainArgs +import transformers +from datasets import load_from_disk from materializers.directory_materializer import DirectoryMaterializer -from pydantic import BaseModel from typing_extensions import Annotated -from zenml import get_step_context, log_model_metadata, step +from utils.callbacks import ZenMLCallback +from utils.loaders import load_base_model +from utils.tokenizer import load_tokenizer +from zenml import logging as zenml_logging +from zenml import step from zenml.logger import get_logger from zenml.materializers import BuiltInMaterializer -from scripts.convert_lit_checkpoint import convert_lit_checkpoint -from scripts.download import download_from_hub -from scripts.merge_lora import merge_lora -from scripts.prepare_alpaca import prepare -from steps.params import LoraParameters -from steps.utils import ( - convert_to_lit_checkpoint_if_necessary, - get_huggingface_access_token, +logger = get_logger(__name__) +zenml_logging.STEP_LOGS_STORAGE_MAX_MESSAGES = ( + 10000 # workaround for https://github.com/zenml-io/zenml/issues/2252 ) -logger = get_logger(__file__) - - -class DataParameters(BaseModel): - """Data preprocessing parameters.""" - - seed: int = 42 - test_split_fraction: float = 0.03865 - mask_inputs: bool = False - ignore_index: int = -1 - max_seq_length: Optional[int] = None - - -class TrainingParameters(BaseModel): - """Training parameters.""" - - save_interval: int = 1000 - log_interval: int = 1 - global_batch_size: int = 64 - micro_batch_size: int = 4 - lr_warmup_steps: int = 100 - epochs: Optional[int] = None - epoch_size: Optional[int] = None - max_tokens: Optional[int] = None - max_seq_length: Optional[int] = None - - learning_rate: float = 1e-3 - weight_decay: float = 0.02 - beta1: float = 0.9 - beta2: float = 0.95 - max_norm: Optional[float] = None - min_lr: float = 6e-5 - - -class EvalParameters(BaseModel): - """Mid-training evaluation parameters.""" - - interval: int = 100 - max_new_tokens: int = 100 - max_iters: int = 100 - - -class FinetuningParameters(BaseModel): - """Parameters for the finetuning step.""" - - base_model_repo: str - from_safetensors: bool = False - - adapter_output_repo: Optional[str] = None - merged_output_repo: Optional[str] = None - convert_to_hf_checkpoint: bool = False - - precision: Optional[str] = None - quantize: Optional[ - Literal[ - "bnb.nf4", - "bnb.nf4-dq", - "bnb.fp4", - "bnb.fp4-dq", - "bnb.int8-training", - ] - ] = None - - data: DataParameters = DataParameters() - training: TrainingParameters = TrainingParameters() - eval: EvalParameters = EvalParameters() - lora: LoraParameters = LoraParameters() - @step(output_materializers=[DirectoryMaterializer, BuiltInMaterializer]) def finetune( - config: FinetuningParameters, dataset_directory: Optional[Path] = None -) -> Annotated[Optional[Path], "adapter"]: - """Finetune model using LoRA. + base_model_id: str, + dataset_dir: Path, + max_steps: int = 1000, + logging_steps: int = 50, + eval_steps: int = 50, + save_steps: int = 50, + optimizer: str = "paged_adamw_8bit", + lr: float = 2.5e-5, + per_device_train_batch_size: int = 2, + gradient_accumulation_steps: int = 4, + warmup_steps: int = 5, + bf16: bool = True, + use_fast: bool = True, + load_in_4bit: bool = False, + load_in_8bit: bool = False, +) -> Annotated[Path, "ft_model_dir"]: + """Finetune the model using PEFT. + + Base model will be derived from configure step and finetuned model will + be saved to the output directory. + + Finetuning parameters can be found here: https://github.com/huggingface/peft#fine-tuning Args: - config: Configuration for this step. + base_model_id: The base model id to use. + dataset_dir: The path to the dataset directory. + max_steps: The maximum number of steps to train for. + logging_steps: The number of steps to log at. + eval_steps: The number of steps to evaluate at. + save_steps: The number of steps to save at. + optimizer: The optimizer to use. + lr: The learning rate to use. + per_device_train_batch_size: The batch size to use for training. + gradient_accumulation_steps: The number of gradient accumulation steps. + warmup_steps: The number of warmup steps. + bf16: Whether to use bf16. + use_fast: Whether to use the fast tokenizer. + load_in_4bit: Whether to load the model in 4bit mode. + load_in_8bit: Whether to load the model in 8bit mode. + + Returns: + The path to the finetuned model directory. """ - torch.set_float32_matmul_precision("high") - - access_token = get_huggingface_access_token() - - checkpoint_root_dir = Path("checkpoints") - checkpoint_dir = checkpoint_root_dir / config.base_model_repo - - if checkpoint_dir.exists(): - logger.info( - "Checkpoint directory already exists, skipping download..." - ) - else: - download_from_hub( - repo_id=config.base_model_repo, - from_safetensors=config.from_safetensors, - checkpoint_dir=checkpoint_root_dir, - access_token=access_token, - ) - - convert_to_lit_checkpoint_if_necessary(checkpoint_dir=checkpoint_dir) - - if dataset_directory: - try: - dataset_name = ( - get_step_context() - .inputs["dataset_directory"] - .run_metadata["dataset_name"] - .value - ) - except KeyError: - dataset_name = "unknown_dataset" - else: - dataset_directory = Path("data/alpaca") - dataset_name = dataset_directory.name - prepare( - destination_path=dataset_directory, - checkpoint_dir=checkpoint_dir, - test_split_fraction=config.data.test_split_fraction, - seed=config.data.seed, - mask_inputs=config.data.mask_inputs, - ignore_index=config.data.ignore_index, - max_seq_length=config.data.max_seq_length, - ) - - model_name = checkpoint_dir.name - - log_model_metadata( - metadata={"model_name": model_name, "dataset_name": dataset_name} + project = "zenml-finetune" + base_model_name = "mistral" + run_name = base_model_name + "-" + project + output_dir = "./" + run_name + + logger.info("Loading datasets...") + tokenizer = load_tokenizer(base_model_id, use_fast=use_fast) + tokenized_train_dataset = load_from_disk(dataset_dir / "train") + tokenized_val_dataset = load_from_disk(dataset_dir / "val") + + logger.info("Loading base model...") + model = load_base_model( + base_model_id, + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, ) - adapter_output_dir = Path("output/lora") / dataset_name / model_name - io_args = IOArgs( - train_data_dir=dataset_directory, - val_data_dir=dataset_directory, - checkpoint_dir=checkpoint_dir, - out_dir=adapter_output_dir, + trainer = transformers.Trainer( + model=model, + train_dataset=tokenized_train_dataset, + eval_dataset=tokenized_val_dataset, + args=transformers.TrainingArguments( + output_dir=output_dir, + warmup_steps=warmup_steps, + per_device_train_batch_size=per_device_train_batch_size, + gradient_checkpointing=True, + gradient_accumulation_steps=gradient_accumulation_steps, + max_steps=max_steps, + learning_rate=lr, + logging_steps=logging_steps, + bf16=bf16, + optim=optimizer, + logging_dir="./logs", + save_strategy="steps", + save_steps=save_steps, + evaluation_strategy="steps", + eval_steps=eval_steps, + do_eval=True, + ), + data_collator=transformers.DataCollatorForLanguageModeling( + tokenizer, mlm=False + ), + callbacks=[ZenMLCallback()], ) - train_args = TrainArgs(**config.training.dict()) - eval_args = EvalArgs(**config.eval.dict()) - setup( - devices=1, - io=io_args, - train=train_args, - eval=eval_args, - precision=config.precision, - quantize=config.quantize, - **config.lora.dict(), + + model.config.use_cache = ( + False # silence the warnings. Please re-enable for inference! ) - if config.merged_output_repo: - lora_path = adapter_output_dir / "lit_model_lora_finetuned.pth" - - merge_output_dir = ( - Path("output/lora_merged") / dataset_name / model_name - ) - merge_lora( - lora_path=lora_path, - checkpoint_dir=checkpoint_dir, - out_dir=merge_output_dir, - precision=config.precision, - **config.lora.dict(), - ) - - for path in Path(checkpoint_dir).glob("*.json"): - destination = Path(merge_output_dir) / path.name - shutil.copy(src=path, dst=destination) - - if config.convert_to_hf_checkpoint: - upload_dir = ( - Path("output/lora_merged_hf") / dataset_name / model_name - ) - upload_dir.mkdir(parents=True, exist_ok=True) - convert_lit_checkpoint( - checkpoint_path=config.merged_output_repo / "lit_model.pth", - config_path=config.merged_output_repo / "lit_config.json", - output_path=upload_dir / "pytorch_model", - ) - else: - upload_dir = merge_output_dir - - commit = upload_folder( - repo_id=config.merged_output_repo, - folder_path=upload_dir, - token=access_token, - ) - log_model_metadata( - metadata={ - "merged_model_huggingface_commit_hash": commit.oid, - "merged_model_huggingface_commit_url": commit.commit_url, - } - ) - - if config.adapter_output_repo: - commit = upload_folder( - repo_id=config.adapter_output_repo, - folder_path=adapter_output_dir, - token=access_token, - ) - log_model_metadata( - metadata={ - "adapter_huggingface_commit_hash": commit.oid, - "adapter_huggingface_commit_url": commit.commit_url, - } - ) - return None - else: - # If the adapter should not be uploaded to the HF Hub, we store it - # in the artifact store - return adapter_output_dir + logger.info("Training model...") + trainer.train() + + logger.info("Saving model...") + model.config.use_cache = True + ft_model_dir = Path("model_dir") + ft_model_dir.mkdir(parents=True, exist_ok=True) + trainer.save_model(ft_model_dir) + + return ft_model_dir diff --git a/llm-lora-finetuning/steps/prepare_datasets.py b/llm-lora-finetuning/steps/prepare_datasets.py new file mode 100644 index 00000000..b0378a70 --- /dev/null +++ b/llm-lora-finetuning/steps/prepare_datasets.py @@ -0,0 +1,70 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from functools import partial +from pathlib import Path + +from materializers.directory_materializer import DirectoryMaterializer +from typing_extensions import Annotated +from utils.tokenizer import generate_and_tokenize_prompt, load_tokenizer +from zenml import log_model_metadata, step +from zenml.materializers import BuiltInMaterializer + + +@step(output_materializers=[DirectoryMaterializer, BuiltInMaterializer]) +def prepare_data( + base_model_id: str, + system_prompt: str, + dataset_name: str = "gem/viggo", + use_fast: bool = True, +) -> Annotated[Path, "datasets_dir"]: + """Prepare the datasets for finetuning. + + Args: + base_model_id: The base model id to use. + system_prompt: The system prompt to use. + dataset_name: The name of the dataset to use. + use_fast: Whether to use the fast tokenizer. + + Returns: + The path to the datasets directory. + """ + from datasets import load_dataset + + log_model_metadata( + {"system_prompt": system_prompt, "base_model_id": base_model_id} + ) + + tokenizer = load_tokenizer(base_model_id, False, use_fast) + gen_and_tokenize = partial( + generate_and_tokenize_prompt, + tokenizer=tokenizer, + system_prompt=system_prompt, + ) + + train_dataset = load_dataset(dataset_name, split="train") + tokenized_train_dataset = train_dataset.map(gen_and_tokenize) + eval_dataset = load_dataset(dataset_name, split="validation") + tokenized_val_dataset = eval_dataset.map(gen_and_tokenize) + test_dataset = load_dataset(dataset_name, split="test") + + datasets_path = Path("datasets") + tokenized_train_dataset.save_to_disk(str((datasets_path / "train").absolute())) + tokenized_val_dataset.save_to_disk(str((datasets_path / "val").absolute())) + test_dataset.save_to_disk(str((datasets_path / "test_raw").absolute())) + + return datasets_path diff --git a/llm-lora-finetuning/steps/promote.py b/llm-lora-finetuning/steps/promote.py new file mode 100644 index 00000000..53549d43 --- /dev/null +++ b/llm-lora-finetuning/steps/promote.py @@ -0,0 +1,76 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from zenml import get_step_context, step +from zenml.client import Client +from zenml.logger import get_logger + +logger = get_logger(__name__) + + +@step +def promote( + metric: str = "rouge1", + target_stage: str = "staging", +) -> None: + """Promote the model to the target stage. + + If the model does not outperform the base model, it will be skipped. + If the model does not outperform the model in the target stage, it will be skipped. + + Args: + metric: The metric to use for promotion. + target_stage: The target stage to promote to. + """ + context_model = get_step_context().model + base_metrics = context_model.load_artifact("base_model_rouge_metrics") + ft_metrics = context_model.load_artifact("finetuned_model_rouge_metrics") + staging_metrics = None + try: + staging_model = Client().get_model_version( + context_model.name, target_stage + ) + staging_metrics = staging_model.get_artifact( + "finetuned_model_rouge_metrics" + ).load() + except KeyError: + pass + + msg = ( + f"`{metric}` values to compare:\n" + f"base={base_metrics[metric]*100:.2f}%\n" + f"finetuned={ft_metrics[metric]*100:.2f}%" + ) + if staging_metrics: + msg += f"\nstaging={staging_metrics[metric]*100:.2f}%" + logger.info(msg) + + if base_metrics[metric] <= ft_metrics[metric]: + if staging_metrics is not None and ( + staging_metrics[metric] > ft_metrics[metric] + ): + logger.info( + "Skipping promotion: model does not " + f"outperform the current model in `{target_stage}`." + ) + else: + logger.info(f"Promoting model to `{target_stage}`") + get_step_context().model.set_stage(target_stage, True) + else: + logger.info( + "Skipping promotion: model does not outperform the base model." + ) diff --git a/llm-lora-finetuning/utils/callbacks.py b/llm-lora-finetuning/utils/callbacks.py new file mode 100644 index 00000000..fa374564 --- /dev/null +++ b/llm-lora-finetuning/utils/callbacks.py @@ -0,0 +1,67 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Dict + +from transformers import ( + TrainerCallback, + TrainerControl, + TrainerState, + TrainingArguments, +) +from zenml import get_step_context + + +class ZenMLCallback(TrainerCallback): + """Callback that logs metrics to ZenML.""" + + def on_evaluate( + self, + args: TrainingArguments, + state: TrainerState, + control: TrainerControl, + metrics: Dict[str, float], + **kwargs, + ): + """Log metrics to the ZenML Model version as metadata. + + Args: + args: The training arguments. + state: The trainer state. + control: The trainer control. + metrics: The metrics to log. + """ + try: + context = get_step_context() + context.model.log_metadata( + { + f"step_{state.global_step}_eval_metrics": metrics, + } + ) + except RuntimeError: + # If we can't get the context, silently pass + return + + def on_save( + self, + args: TrainingArguments, + state: TrainerState, + control: TrainerControl, + **kwargs, + ): + # TODO: add ability to save model checkpoints here, will likely get redundant with Mounts + pass diff --git a/llm-lora-finetuning/utils/cuda.py b/llm-lora-finetuning/utils/cuda.py new file mode 100644 index 00000000..35d78d4a --- /dev/null +++ b/llm-lora-finetuning/utils/cuda.py @@ -0,0 +1,26 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import gc + +import torch + + +def cleanup_memory() -> None: + """Clean up GPU memory.""" + while gc.collect(): + torch.cuda.empty_cache() diff --git a/llm-lora-finetuning/utils/loaders.py b/llm-lora-finetuning/utils/loaders.py new file mode 100644 index 00000000..9301b06a --- /dev/null +++ b/llm-lora-finetuning/utils/loaders.py @@ -0,0 +1,111 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pathlib import Path +from typing import Any + +import torch +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import AutoModelForCausalLM, BitsAndBytesConfig +from utils.logging import print_trainable_parameters + + +def load_base_model( + base_model_id: str, + is_training: bool = True, + load_in_8bit: bool = False, + load_in_4bit: bool = False, +) -> Any: + """Load the base model. + + Args: + base_model_id: The base model id to use. + is_training: Whether the model should be prepared for training or not. + If True, the Lora parameters will be enabled and PEFT will be + applied. + load_in_8bit: Whether to load the model in 8-bit mode. + load_in_4bit: Whether to load the model in 4-bit mode. + + Returns: + The base model. + """ + bnb_config = BitsAndBytesConfig( + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + + model = AutoModelForCausalLM.from_pretrained( + base_model_id, quantization_config=bnb_config, device_map="auto" + ) + + if is_training: + model.gradient_checkpointing_enable() + model = prepare_model_for_kbit_training(model) + + config = LoraConfig( + r=8, + lora_alpha=16, + target_modules=[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + "lm_head", + ], + bias="none", + lora_dropout=0.05, # Conventional + task_type="CAUSAL_LM", + ) + + model = get_peft_model(model, config) + print_trainable_parameters(model) + + return model + + +def load_pretrained_model( + ft_model_dir: Path, + load_in_4bit: bool = False, + load_in_8bit: bool = False, +) -> AutoModelForCausalLM: + """Load the finetuned model saved in the output directory. + + Args: + ft_model_dir: The path to the finetuned model directory. + load_in_4bit: Whether to load the model in 4-bit mode. + load_in_8bit: Whether to load the model in 8-bit mode. + + Returns: + The finetuned model. + """ + bnb_config = BitsAndBytesConfig( + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + model = AutoModelForCausalLM.from_pretrained( + ft_model_dir, quantization_config=bnb_config, device_map="auto" + ) + return model diff --git a/llm-lora-finetuning/utils/logging.py b/llm-lora-finetuning/utils/logging.py new file mode 100644 index 00000000..0fd2df78 --- /dev/null +++ b/llm-lora-finetuning/utils/logging.py @@ -0,0 +1,35 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Any + +from zenml.logger import get_logger + +logger = get_logger(__name__) + + +def print_trainable_parameters(model: Any): + """Prints the number of trainable parameters in the model.""" + trainable_params = 0 + all_param = 0 + for _, param in model.named_parameters(): + all_param += param.numel() + if param.requires_grad: + trainable_params += param.numel() + logger.info( + f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" + ) diff --git a/llm-lora-finetuning/utils/tokenizer.py b/llm-lora-finetuning/utils/tokenizer.py new file mode 100644 index 00000000..22f963ff --- /dev/null +++ b/llm-lora-finetuning/utils/tokenizer.py @@ -0,0 +1,138 @@ +# Apache Software License 2.0 +# +# Copyright (c) ZenML GmbH 2024. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +from transformers import AutoTokenizer + + +def load_tokenizer( + base_model_id: str, + is_eval: bool = False, + use_fast: bool = True, +) -> AutoTokenizer: + """Loads the tokenizer for the given base model id. + + Args: + base_model_id: The base model id to use. + is_eval: Whether to load the tokenizer for evaluation. + use_fast: Whether to use the fast tokenizer. + + Returns: + The tokenizer. + """ + if is_eval: + tokenizer = AutoTokenizer.from_pretrained( + base_model_id, + add_bos_token=True, + device_map="auto", + use_fast=use_fast, + ) + tokenizer.pad_token_id = 0 + else: + tokenizer = AutoTokenizer.from_pretrained( + base_model_id, + model_max_length=512, + padding_side="left", + add_eos_token=True, + device_map="auto", + use_fast=use_fast, + ) + tokenizer.pad_token = tokenizer.eos_token + return tokenizer + + +def tokenize( + prompt: str, + tokenizer: AutoTokenizer, +) -> dict: + """Tokenizes the prompt for single entry. + + Args: + prompt: The prompt to tokenize. + tokenizer: The tokenizer to use. + + Returns: + The tokenized prompt. + """ + result = tokenizer( + prompt, + truncation=True, + max_length=512, + padding="max_length", + ) + result["labels"] = result["input_ids"].copy() + return result + + +def generate_and_tokenize_prompt( + data_point: dict, + tokenizer: AutoTokenizer, + system_prompt: str, +): + """Generates and tokenizes the prompt for single entry. + + To be used in map function of the dataset. + + Args: + data_point: The data point to generate and tokenize. + tokenizer: The tokenizer to use. + system_prompt: The system prompt to use. + + Returns: + The tokenized prompt. + """ + full_prompt = f"""{system_prompt} + +### Target sentence: +{data_point["target"]} + +### Meaning representation: +{data_point["meaning_representation"]} +""" + return tokenize(full_prompt, tokenizer) + + +def tokenize_for_eval( + data_points: dict, + tokenizer: AutoTokenizer, + system_prompt: str, +): + """Tokenizes the prompts for evaluation. + + This runs for the whole test dataset at once. + + Args: + data_points: The data points to tokenize. + tokenizer: The tokenizer to use. + system_prompt: The system prompt to use. + + Returns: + The tokenized prompt. + """ + eval_prompts = [ + f"""{system_prompt} + +### Target sentence: +{data_point} + +### Meaning representation: +""" + for data_point in data_points["target"] + ] + return tokenizer(eval_prompts, padding="longest", return_tensors="pt").to( + "cuda" + )