Phi-3 and Llama-3 tutorial

llama3-finetune/README.md

@@ -0,0 +1,59 @@
+# Llama-3 Model Fine-tuning Demo
+
+This demo will show how to use ACPT (Azure Container for PyTorch) along with accelerators such as onnxruntime training (through ORTModule) and DeepSpeed to fine-tune Llama-3 model.
+
+## Background
+
+[Llama-3-8b](https://huggingface.co/blog/llama3) the latest large language model from Meta, is built on the architecture of Llama 2 and comes in two sizes (8B and 70B parameters) with pre-trained and instruction-tuned versions. Here we show fine-tuning on 8b model.
+
+## Set up
+
+### AzureML
+The easiest option to run the demo will be using AzureML as the environment details are already included, there is another option to run directly on the machine which is provided later. For AzureML, please complete the following prerequisites:
+
+#### Local environment
+Follow [Install Azure CLI](https://learn.microsoft.com/en-us/cli/azure/install-azure-cli-linux?pivots=apt#install-azure-cli) to install Azure CLI.
+Set up your local environment with az-cli and azureml dependency for script submission:
+
+```
+az login
+pip install azure-ai-ml azure-identity
+```
+
+#### AzureML Workspace
+- An AzureML workspace is required to run this demo. Download the config.json file ([How to get config.json file from Azure Portal](https://docs.microsoft.com/en-us/azure/machine-learning/how-to-configure-environment#workspace)) for your workspace. Make sure to put this config file in this folder and name it ws_config.json.
+- The workspace should have a gpu cluster. This demo was tested with GPU cluster of SKU [Standard_ND96asr_A100_v4](https://learn.microsoft.com/en-us/azure/virtual-machines/nda100-v4-series). [Standard_ND40rs_v2](https://docs.microsoft.com/en-us/azure/virtual-machines/ndv2-series) should also work with reduced block size. See this document for [creating gpu cluster](https://docs.microsoft.com/en-us/azure/machine-learning/how-to-create-attach-compute-cluster?tabs=python). We do not recommend running this demo on `NC` series VMs which uses old architecture (K80).
+- Additionally, you'll need to create a [Custom Curated Environment ACPT](https://learn.microsoft.com/en-us/azure/machine-learning/resource-curated-environments) with PyTorch >=2.2.0 and the steps in the Dockerfile.
+
+## Run Experiments
+The demo is ready to be run.
+
+#### `aml_submit_clm.py` submits an training job to AML for both Pytorch+DeepSpeedStage2+LoRA and ORT+DeepSpeedStage2+LoRA. This job builds the training environment and runs the fine-tuning script in it.
+
+```bash
+python aml_submit_clm.py
+```
+
+The above script will generate two URLs, one for Pytorch and another for ONNX Runtime training.
+
+We observe **12% speedup** for Llama-3 trained leveraging ONNX Runtime Training with 8 V100 GPUs with 32GB memory with batch size of 1.
+
+### Run directly on your compute
+
+If you are using CLI by directly logging into your machine then you can follow the below instructions. The below steps assume you have the required packages like Pytorch, ONNX Runtime training, Transformers and more already installed in your system. For easier setup, you can look at the environment folder.
+
+```bash
+cd finetune-clm
+
+# To run the model using Pytorch
+torchrun --nproc_per_node 8 run_clm.py --model_name_or_path meta-llama/Meta-Llama-3-8B --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --do_train --save_strategy 'no' --fp16 --block_size 2048 --max_steps -1 --per_device_train_batch_size 1 --num_train_epochs 2 --output_dir output_dir --overwrite_output_dir --deepspeed zero_stage_2.json --evaluation_strategy no --remove_unused_columns False
+
+# To run the model using ONNX Runtime training, you need to export couple of variables and run the same command above, overall these would be your steps:
+export APPLY_ORT="True"
+export ORTMODULE_FALLBACK_POLICY="FALLBACK_DISABLE"
+export ORTMODULE_DEEPCOPY_BEFORE_MODEL_EXPORT=0
+# Optionally you can enable/disable Triton, for faster performance it is turned on
+export ORTMODULE_USE_TRITON=1
+torchrun --nproc_per_node 8 run_clm.py --model_name_or_path meta-llama/Meta-Llama-3-8B --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --do_train --save_strategy 'no' --fp16 --block_size 2048 --max_steps -1 --per_device_train_batch_size 1 --num_train_epochs 2 --output_dir output_dir --overwrite_output_dir --deepspeed zero_stage_2.json --evaluation_strategy no --remove_unused_columns False
+```
+

llama3-finetune/aml_submit_clm.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2023 Microsoft Corp. 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
+
+import argparse
+from pathlib import Path
+import json
+import os
+
+from azure.ai.ml import MLClient, command
+from azure.ai.ml.entities import Environment, BuildContext 
+from azure.identity import AzureCliCredential
+
+# run test on automode workspace
+ws_config = json.load(open("ws_config.json"))
+subscription_id = ws_config["subscription_id"]
+resource_group = ws_config["resource_group"]
+workspace_name = ws_config["workspace_name"]
+compute = ws_config["compute"]
+nproc_per_node = ws_config["nproc_per_node"]
+
+hf_token = os.getenv("HUGGINGFACE_TOKEN", None)
+
+
+def get_args(raw_args=None):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--experiment_name", default="Llama-3-ORT-CLM-Stage2-Experiment", help="Experiment name for AML Workspace")
+
+    args = parser.parse_args(raw_args)
+    return args
+
+def main(raw_args=None):
+    args = get_args(raw_args)
+
+    ml_client = MLClient(
+        AzureCliCredential(), subscription_id, resource_group, workspace_name
+    )
+
+    root_dir = Path(__file__).resolve().parent
+    environment_dir = root_dir / "environment"
+    code_dir = root_dir / "finetune-clm"
+
+    model = "meta-llama/Meta-Llama-3-8B"
+
+    num_train_epochs = 15
+    bsz = 1
+    max_steps = -1
+    block_size = 2048
+
+    dataset_name = "wikitext"
+    dataset_config_name = "wikitext-2-raw-v1"
+    text_column_name = "text"
+    label_column_name = "label"
+
+    # Pytorch eager
+    pytorch_job = command(
+        code=code_dir,  # local path where the code is stored
+        command=f"gmon torchrun --nproc_per_node {nproc_per_node} run_clm.py \
+            --model_name_or_path {model} \
+            --dataset_name {dataset_name} \
+            --dataset_config_name {dataset_config_name} \
+            --do_train \
+            --save_strategy 'no' \
+            --per_device_train_batch_size {bsz} \
+            --num_train_epochs {num_train_epochs} \
+            --output_dir ./outputs --overwrite_output_dir \
+            --fp16 --max_steps {max_steps} \
+            --block_size {block_size} \
+            --deepspeed zero_stage_2.json \
+            --evaluation_strategy epoch --remove_unused_columns False \
+            --token $HUGGINGFACE_TOKEN",
+        environment=Environment(build=BuildContext(path=environment_dir)),
+        environment_variables={
+            "HUGGINGFACE_TOKEN": hf_token
+        },
+        experiment_name="Llama-3-Pytorch-CLM-LORA-Stage2-Experiment",
+        compute=compute,
+        display_name=model.replace(
+            "meta-llama/Meta-Llama-3",
+            f"TorchEager+DS2_lora-{bsz}-{block_size}-{num_train_epochs}epoch"
+        ),
+        description=f"Finetune HuggingFace's Llama-3 using PyTorch.",
+        tags={"model": model,
+            "bsz": bsz,
+            "dataset_name": dataset_name,
+            "block_size": block_size},
+        shm_size="16g",
+        job_tier="Premium", #https://learn.microsoft.com/en-us/python/api/azure-ai-ml/azure.ai.ml?view=azure-python
+        priority="High"
+    )
+
+    print("submitting PyTorch job for " + model)
+    pytorch_returned_job = ml_client.create_or_update(pytorch_job)
+    print("submitted job")
+
+    pytorch_aml_url = pytorch_returned_job.studio_url
+    print("job link:", pytorch_aml_url)
+
+    # ORT perf
+    ort_job = command(
+        code=code_dir,  # local path where the code is stored
+        command=f"gmon torchrun --nproc_per_node {nproc_per_node} run_clm.py \
+            --model_name_or_path {model} \
+            --dataset_name {dataset_name} \
+            --dataset_config_name {dataset_config_name} \
+            --do_train \
+            --save_strategy 'no' \
+            --per_device_train_batch_size {bsz} \
+            --num_train_epochs {num_train_epochs} \
+            --output_dir ./outputs --overwrite_output_dir \
+            --fp16 --max_steps {max_steps} \
+            --block_size {block_size} \
+            --deepspeed zero_stage_2.json \
+            --evaluation_strategy epoch --remove_unused_columns False \
+            --token $HUGGINGFACE_TOKEN",
+        environment=Environment(build=BuildContext(path=environment_dir)),
+        environment_variables={
+            "HUGGINGFACE_TOKEN": hf_token,
+            "APPLY_ORT": "True",
+            "ORTMODULE_FALLBACK_POLICY": "FALLBACK_DISABLE",
+            "ORTMODULE_DEEPCOPY_BEFORE_MODEL_EXPORT": "0"},
+        experiment_name="Llama-3-ORT-CLM-Stage2-Experiment",
+        compute=compute,
+        display_name=model.replace(
+            "meta-llama/Meta-Llama-3",
+            f"ORT+DS2_lora-{bsz}-{block_size}-{num_train_epochs}epoch"
+        ),
+        description=f"Finetune HuggingFace's Llama-3 using ONNX Runtime.",
+        tags={"model": model,
+            "bsz": bsz,
+            "dataset_name": dataset_name,
+            "block_size": block_size,
+            "epoch": num_train_epochs},
+        shm_size="16g",
+        job_tier="Premium",
+        priority="High"
+    )
+
+    print("submitting ORT job for " + model)
+    ort_returned_job = ml_client.create_or_update(ort_job)
+    print("submitted job")
+
+    ort_aml_url = ort_returned_job.studio_url
+    print("job link:", ort_aml_url)
+
+
+if __name__ == "__main__":
+    main()

llama3-finetune/environment/Dockerfile

@@ -0,0 +1,24 @@
+
+FROM mcr.microsoft.com/aifx/acpt/stable-ubuntu2004-cu118-py310-torch222
+
+# RUN pip install -U datasets evaluate accelerate scikit-learn transformers==4.40.1
+RUN pip install -U datasets evaluate accelerate scikit-learn
+RUN pip install git+https://github.com/huggingface/transformers.git
+RUN pip install git+https://github.com/huggingface/optimum.git
+
+
+# Insall numpy 1.26.4 because numpy 2.0.0 are not supported with many libraries
+RUN pip uninstall numpy -y && \
+    pip install numpy==1.26.4
+
+# Error with deepspeed=0.14.2
+RUN pip uninstall deepspeed -y && \
+    pip install deepspeed==0.13.1
+
+RUN pip install einops
+RUN pip install --upgrade pytest
+RUN pip install peft
+
+
+RUN pip list
+