Merge pull request #125 from zenml-io/project/addsimplefinetune

Add a simple finetunining example to motivate the lightning ai studio integration

Author: htahir1

README.md

@@ -75,7 +75,8 @@ A list of updated and maintained projects by the ZenML team and the community:
 | [Huggingface to Sagemaker](huggingface-sagemaker)                      | NLP                                   | `pytorch` `mlflow` `huggingface` `aws` `s3` `kubeflow` `slack` `github`  |
 | [Complete Guide to LLMs (from RAG to finetuning)](llm-complete-guide)               | NLP, LLMs, embeddings, finetuning                           | `openai` `supabase` `huggingface` `argilla`  |
 | [LLM LoRA Finetuning (Phi3 and Llama 3.1)](llm-lora-finetuning)               | NLP, LLMs                           | `gcp`  |
-| [ECP Price Prediction with GCP Cloud Composer](airflow-cloud-composer-etl-feature-train/)               | Regression, Airflow                           | `cloud-composer` `airflow` |
+| [ECP Price Prediction with GCP Cloud Composer](airflow-cloud-composer-etl-feature-train/README.md)               | Regression, Airflow                           | `cloud-composer` `airflow` |
+| [Simple LLM finetuning with Lightning Studio](simple-llm-finetuning/README.md)               | Lightning AI Studio, LLMs                           | `cloud-composer` `airflow` |
 
 # 💻 System Requirements
 

llm-finetuning-simple/.dockerignore

@@ -0,0 +1,5 @@
+*
+!/materializers/**
+!/pipelines/**
+!/steps/**
+!/utils/**

llm-finetuning-simple/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.

llm-finetuning-simple/README.md

@@ -0,0 +1,180 @@
+# README for LLM Fine-Tuning with ZenML and Lightning AI Studios
+
+## Overview
+
+In the fast-paced world of AI, the ability to efficiently fine-tune Large Language Models (LLMs) for specific tasks is crucial. This project combines ZenML with Lightning AI Studios to streamline and automate the LLM fine-tuning process, enabling rapid iteration and deployment of task-specific models. This is a toy showcase only but can be extended for full production use.
+
+## Table of Contents
+
+1. [Introduction](#introduction)
+2. [Installation](#installation)
+3. [Running the Pipeline](#running-the-pipeline)
+4. [Configuration](#configuration)
+5. [Accelerated Fine-Tuning](#accelerated-fine-tuning)
+6. [Running with Remote Stack](#running-with-remote-stack)
+7. [Customizing Data Preparation](#customizing-data-preparation)
+8. [Project Structure](#project-structure)
+9. [Benefits & Future](#benefits--future)
+10. [Credits](#credits)
+
+## Introduction
+
+As LLMs such as GPT-4, Llama 3.1, and Mistral become more accessible, companies aim to adapt these models for specialized tasks like customer service chatbots, content generation, and specialized data analysis. This project addresses the challenge of scaling fine-tuning and managing numerous LLM variants by combining Lightning AI Studios with the automation capabilities of ZenML.
+
+### Key Benefits
+
+- **Efficient Fine-Tuning:** Fine-tune models with minimal computational resources.
+- **Ease of Management:** Store and distribute adapter weights efficiently.
+- **Scalability:** Serve thousands of fine-tuned variants from a single base model.
+
+## Installation
+
+To set up your environment, follow these steps:
+
+```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
+
+# Install ZenML and Lightning integrations
+pip install zenml
+zenml integration install lightning s3 aws -y
+
+# Initialize and connect to a deployed ZenML server
+zenml init
+zenml connect --url <MYZENMLSERVERURL>
+```
+
+## Running the Pipeline
+
+To run the fine-tuning pipeline, use the `run.py` script with the appropriate configuration file:
+
+```shell
+python run.py --config configs/config_large_gpu.yaml
+```
+
+## Configuration
+
+The fine-tuning process can be configured using YAML files located in the `configs` directory. Here are examples:
+
+### Example `config_large_gpu.yaml`
+
+```yaml
+model:
+  name: llm-finetuning-gpt2-large
+  description: "Fine-tune GPT-2 on larger GPU."
+  tags:
+    - llm
+    - finetuning
+    - gpt2-large
+
+parameters:
+  base_model_id: gpt2-large
+
+steps:
+  prepare_data:
+    parameters:
+      dataset_name: squad
+      dataset_size: 1000
+      max_length: 512
+
+  finetune:
+    parameters:
+      num_train_epochs: 3
+      per_device_train_batch_size: 8
+
+    settings:
+      orchestrator.lightning:
+        machine_type: A10G
+```
+
+### Example `config_small_cpu.yaml`
+
+```yaml
+model:
+  name: llm-finetuning-distilgpt2-small
+  description: "Fine-tune DistilGPT-2 on smaller computer."
+  tags:
+    - llm
+    - finetuning
+    - distilgpt2
+
+parameters:
+  base_model_id: distilgpt2
+
+steps:
+  prepare_data:
+    parameters:
+      dataset_name: squad
+      dataset_size: 100
+      max_length: 128
+
+  finetune:
+    parameters:
+      num_train_epochs: 1
+      per_device_train_batch_size: 4
+```
+
+## Running with Remote Stack
+
+Set up a remote lightning stack with ZenML for fine tuning on remote infrastructure:
+
+1. **Register Orchestrator and Artifact Store:**
+
+    ```shell
+    zenml integration install lightning s3
+    zenml orchestrator register lightning_orchestrator --flavor=lightning --machine_type=CPU --user_id=<YOUR_LIGHTNING_USER_ID> --api_key=<YOUR_LIGHTNING_API_KEY> --username=<YOUR_LIGHTNING_USERNAME>
+    zenml artifact-store register s3_store --flavor=s3 --path=s3://yourpath
+    ```
+
+2. **Set up and Register the Stack:**
+
+    ```shell
+    zenml stack register lightning_stack -o lightning_orchestrator -a s3_store
+    zenml stack set lightning_stack
+    ```
+
+## Customizing Data Preparation
+
+Customize the `prepare_data` step for different datasets by modifying loading logic or tokenization patterns. Update the relevant YAML configuration parameters to fit your dataset and requirements.
+
+## Project Structure
+
+The project follows a structured layout for easy navigation and management:
+
+```
+.
+├── configs                                          # Configuration files for the pipeline
+│   ├── config_large_gpu.yaml                        # Config for large GPU setup
+│   ├── config_small_cpu.yaml                        # Config for small CPU setup
+├── .dockerignore                                    # Docker ignore file
+├── LICENSE                                          # License file
+├── README.md                                        # This file
+├── requirements.txt                                 # Python dependencies
+├── run.py                                           # CLI tool to run pipelines on ZenML Stack
+```
+
+## Benefits & Future
+
+Using smaller, task-specific models is more efficient and cost-effective than relying on large general-purpose models. This strategy allows for:
+
+- **Cost-Effectiveness:** Less computational resources reduce operational costs.
+- **Improved Performance:** Models fine-tuned on specific data often outperform general models on specialized tasks.
+- **Faster Iteration:** Quicker experimentation and iteration cycles.
+- **Data Privacy:** Control over training data, crucial for industries with strict privacy requirements.
+
+## Credits
+
+This project relies on several tools and libraries:
+
+- [Hugging Face Transformers](https://huggingface.co/transformers/)
+- [Hugging Face Datasets](https://huggingface.co/datasets)
+- [ZenML](https://zenml.io/)
+- [Lightning AI Studios](https://www.lightning.ai/)
+
+With these tools, you can efficiently manage the lifecycle of multiple fine-tuned LLM variants, benefiting from the robust infrastructure provided by ZenML and the scalable resources of Lightning AI Studios.
+
+For more details, consult the [ZenML documentation](https://docs.zenml.io) and the [Lightning AI Studio documentation](https://lightning.ai).

llm-finetuning-simple/configs/config_large_gpu.yaml

@@ -0,0 +1,44 @@
+settings:
+  docker:
+    requirements: requirements.txt
+    python_package_installer: uv
+    apt_packages: 
+      - git
+    environment:
+      PJRT_DEVICE: CUDA
+      USE_TORCH_XLA: "false"
+      MKL_SERVICE_FORCE_INTEL: "1"
+      PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
+  orchestrator.lightning:
+      machine_type: CPU
+      user_id: USERID
+      api_key: APIKEY
+      username: USERNAME
+      teamspace: TEAMSPACE
+
+model:
+  name: llm-finetuning-gpt2-large
+  description: "Fine-tune GPT-2 on larger GPU."
+  tags:
+    - llm
+    - finetuning
+    - gpt2-large
+
+parameters:
+  base_model_id: gpt2-large
+
+steps:
+  prepare_data:
+    parameters:
+      dataset_name: squad
+      dataset_size: 1000
+      max_length: 256
+
+  finetune:
+    parameters:
+      num_train_epochs: 3
+      per_device_train_batch_size: 4
+
+    settings:
+        orchestrator.lightning:
+          machine_type: A10G

llm-finetuning-simple/configs/config_small_cpu.yaml

@@ -0,0 +1,22 @@
+model:
+  name: llm-finetuning-distilgpt2-small
+  description: "Fine-tune DistilGPT-2 on smaller computer."
+  tags:
+    - llm
+    - finetuning
+    - distilgpt2
+
+parameters:
+  base_model_id: distilgpt2
+
+steps:
+  prepare_data:
+    parameters:
+      dataset_name: squad
+      dataset_size: 100
+      max_length: 128
+
+  finetune:
+    parameters:
+      num_train_epochs: 1
+      per_device_train_batch_size: 4

llm-finetuning-simple/requirements.txt

@@ -0,0 +1,14 @@
+datasets>=2.19.1
+transformers>=4.43.1
+peft
+bitsandbytes>=0.41.3
+scipy
+evaluate
+rouge_score
+nltk
+accelerate>=0.30.0
+urllib3<2
+zenml>=0.62.0
+torch>=2.2.0
+sentencepiece
+huggingface_hub

llm-finetuning-simple/run.py

@@ -0,0 +1,68 @@
+import torch
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from zenml import pipeline, step, log_model_metadata
+from typing_extensions import Annotated
+import argparse
+from zenml.integrations.huggingface.materializers.huggingface_datasets_materializer import HFDatasetMaterializer
+
+@step(output_materializers=HFDatasetMaterializer)
+def prepare_data(base_model_id: str, dataset_name: str, dataset_size: int, max_length: int) -> Annotated[Dataset, "tokenized_dataset"]:
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+    dataset = load_dataset(dataset_name, split=f"train[:{dataset_size}]")
+    
+    def tokenize_function(example):
+        prompt = f"Question: {example['question']}\nAnswer: {example['answers']['text'][0]}"
+        return tokenizer(prompt, truncation=True, padding="max_length", max_length=max_length)
+
+    tokenized_data = dataset.map(tokenize_function, remove_columns=dataset.column_names)
+    log_model_metadata(metadata={"dataset_size": len(tokenized_data), "max_length": max_length})
+    return tokenized_data
+
+@step
+def finetune(base_model_id: str, tokenized_dataset: Dataset, num_train_epochs: int, per_device_train_batch_size: int) -> None:
+    torch.cuda.empty_cache()
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        device_map="auto",
+        torch_dtype=torch.float32,  # Changed from float16 to float32
+        low_cpu_mem_usage=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = tokenizer.pad_token_id
+
+    training_args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=num_train_epochs,
+        per_device_train_batch_size=per_device_train_batch_size,
+        gradient_accumulation_steps=8,
+        logging_steps=10,
+        save_strategy="epoch",
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        optim="adamw_torch",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
+    )
+    
+    train_result = trainer.train()
+    log_model_metadata(metadata={"metrics": {"train_loss": train_result.metrics.get("train_loss")}})
+    trainer.save_model("finetuned_model")
+
+@pipeline
+def llm_finetune_pipeline(base_model_id: str):
+    tokenized_dataset = prepare_data(base_model_id)
+    finetune(base_model_id, tokenized_dataset)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True, help='Path to the YAML config file')
+    args = parser.parse_args()
+    llm_finetune_pipeline.with_options(config_path=args.config)()