Fine-tune and deploy LLaMA V2 models on AWS Trainium and AWS Inferentia based instances in SageMaker JumpStart

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

In this demo notebook, we demonstrate how to use the SageMaker Python SDK to deploy pre-trained Llama 2 model as well as fine-tune it for your dataset in domain adaptation or instruction tuning format on AWS Trainium and AWS Inferentia based instances.

AWS Neuron is an SDK with a compiler, runtime, and profiling tools that unlocks high-performance and cost-effective deep learning (DL) acceleration. It supports high-performance training on AWS Trainium-based Amazon Elastic Compute Cloud (Amazon EC2) Trn1 instances. For model deployment, it supports high-performance and low-latency inference on AWS Inferentia-based Amazon EC2 Inf1 instances and AWS Inferentia2-based Amazon EC2 Inf2 instances. For details, see Official documentation.

Model License information

To perform inference on these models, you need to pass 'accept_eula=True' as part of model.deploy() call. This means you have read and accept the end-user-license-agreement (EULA) of the model. EULA can be found in model card description or from https://ai.meta.com/resources/models-and-libraries/llama-downloads/. By default, this notebook sets 'accept_eula=False', so all inference requests will fail until you explicitly change this custom attribute.

Similarly, to perform fine-tuning on these models, you need pass environment variable '{"accept_eula": "true"}' to JumpStartEstimator class.

Set up

We begin by installing and upgrading necessary packages. Restart the kernel after executing the cell below for the first time.

Deploy Pre-trained Model

You can now deploy the model using SageMaker JumpStart through 2 options. Option 1 allows you to quickly deploy the endpoint with default setting in two lines of code. Option 2 allows you to have more customized configurations.

To deploy a model on AWS Trainium or AWS Inferentia based instances, we will firstly need call PyTorch Neuron (torch-neuronx) to compile the model into a Neuron specific graph. Then during runtime the graph is executed on the NeuronCores of the AWS Trainium or AWS Inferentia based instances. Compiling the graph involves running optimizations that can make use of the NeuronCores efficiently.

In SageMaker JumpStart, we pre-compile the neuron graphs for a varieity of configurations such that you do not spend time waiting for compiling the graph during endpoint deployment, as long as the deployment parameters (environmental variables) matches one of configurations listed as below. Otherwise, the compilation will be triggered during endpoint deployment, which will take a slightly longer time to deploy a model.

Invoke the endpoint

Next, we invoke the endpoint with some sample queries. Later, in this notebook, we will fine-tune this model with a custom dataset and carry out inference using the fine-tuned model. We will also show comparison between results obtained via the pre-trained and the fine-tuned models.

Dataset preparation for fine-tuning

You can fine-tune on the dataset with domain adaptation format or instruction tuning format. Below are the instructions for how the training data should be formatted for input to the model.

• Input: A train directory containing either a JSON lines (.jsonl) or text (.txt) formatted file.
  • For JSON lines (JSONL) file, each line is a dictionary, repsentating a dictionary. The key in dictionary (each line) has to be 'text'.
  • The number of files under train directory should equal to one.
• Output: A trained model that can be deployed for inference.

In this demo, we will use a subset of Dolly dataset in an instruction tuning format. Dolly dataset contains roughly 15,000 instruction following records for various categories such as question answering, summarization, information extraction etc. It is available under Apache 2.0 license. We will select the summarization examples for fine-tuning.

For demonstration of using text file as input, please see Appendix 2

Next, we use a prompt template for preprocessing the data in an instruction / input format for the training job, and also for inferencing the deployed endpoint.

Upload dataset to S3

We will upload the prepared dataset to S3 which will be used for fine-tuning.

Train the model

Next, we fine-tune the LLaMA v2 model on the summarization dataset from Dolly on AWS Trainium instance. You have two options: ml.trn1.32xlarge (default) and ml.trn1n.32xlarge. Finetuning scripts are based on scripts provided by Neuronx-Nemo-Megatron. For a list of supported hyper-parameters and their default values, please see supported hyperparameters for fine-tuning.

Overwrite some of the hyperparameters

Validate hyperparameters

Studio Kernel idle issue: If your studio kernel goes idle and you lose reference to the estimator object, please see section 4. Studio Kernel Dead/Creating JumpStart Model from the training Job on how to deploy endpoint using the training job name and the model id.

Deploy the fine-tuned model

Next, we deploy the fine-tuned model. We will compare the performance of fine-tuned and pre-trained model.

Evaluate the pre-trained and fine-tuned model

Next, we use the test data to evaluate the performance of the fine-tuned model and compare it with the pre-trained model.

Clean up resources

Appendix

1. Supported Inference Parameters

This model supports the following inference payload parameters:

• max_length: Model generates text until the output length (which includes the input context length) reaches max_length. If specified, it must be a positive integer.
• max_new_tokens: Model generates text until the output length (excluding the input context length) reaches max_new_tokens. If specified, it must be a positive integer.
• num_beams: Number of beams used in the greedy search. If specified, it must be integer greater than or equal to num_return_sequences.
• no_repeat_ngram_size: Model ensures that a sequence of words of no_repeat_ngram_size is not repeated in the output sequence. If specified, it must be a positive integer greater than 1.
• temperature: Controls the randomness in the output. Higher temperature results in output sequence with low-probability words and lower temperature results in output sequence with high-probability words. If temperature -> 0, it results in greedy decoding. If specified, it must be a positive float.
• early_stopping: If True, text generation is finished when all beam hypotheses reach the end of sentence token. If specified, it must be boolean.
• do_sample: If True, sample the next word as per the likelihood. If specified, it must be boolean.
• top_k: In each step of text generation, sample from only the top_k most likely words. If specified, it must be a positive integer.
• top_p: In each step of text generation, sample from the smallest possible set of words with cumulative probability top_p. If specified, it must be a float between 0 and 1.
• stop: If specified, it must be a list of strings. Text generation stops if any one of the specified strings is generated.

We may specify any subset of the parameters mentioned above while invoking an endpoint.

2. Use text file as input to fine-tune LLaMA-2

3. Supported Hyper-parameters for fine-tuning

• max_input_length: Maximum total input sequence length after tokenization. Sequences longer than this will be truncated. Default: 2048.
• learning_rate: The rate at which the model weights are updated after working through each batch of training examples. Must be a positive float greater than 0. Default: 6e-6.
• min_learning_rate: The learning rate at the last step of learning rate scheduler 'CosineAnnealing'. Default: 1e-06.
• global_train_batch_size: The global batch size for training. Based on global_train_batch_size, the gradient accumulation is calculated as global_train_batch_size / (data_parallel_degree * per_device_train_batch_size), where data_parallel_degree is calculated as total number of neuron cores / (tensor_parallel_degree * pipeline_parallel_degree). Default: 256.
• per_device_train_batch_size: The batch size per Neuron core for training. Default: 1
• layer_norm_epilson: During layer normalization, a value added to the denominator for numerical stability. See documentation. Default: 0.00001.
• preprocessing_num_workers: The number of processors to use for the preprocessing. If None, all of workers (number of vCPUs) are used for preprocessing. Default: "None"
• weight_decay: The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in AdamW optimizer. Default: 0.1.
• lr_scheduler_type: Learning rate scheduler type. Default: 'CosineAnnealing' (currently we only support 'CosineAnnealing' scheduler type).
• warmup_steps: Linear warmup over warmup steps. Default: 10.
• constant_steps: The number of steps for learning rate to be constant after warmup_steps in 'CosineAnnealing' scheduler type. Default: 0.
• adam_beta1: The beta1 hyperparameter (exponential decay rate for the first moment estimates) for the AdamW optimizer. Default: 0.9.
• adam_beta2: The beta2 hyperparameter (exponential decay rate for the first moment estimates) for the AdamW optimizer. Default: 0.95.
• mixed_precision: Whether to use mixed precision. If mixed_precision to be 'True', it means that master weights and optimizer states are stored in fp32, and model weights are saved in bf16. For details, see reference. Default: 'True'.
• tensor_parallel_degree: The number of neuron cores which specific model weights, gradients, and optimizer states are split across. For details, see reference. Default: "8" (currently we only support parallel degree as 8).
• pipeline_parallel_degree: The number of neuron cores which the layers of a model are partitioned across. Default: "1" (currently we only support "1" for LLaMA-2 7B and "4" for LLaMA-2 13B).
• append_eod: Whether to append an <eod> token to the end of each example. By setting it to 'True', the fine-tuned model tends to generate succinct output. Default: 'False'.

4. Studio Kernel goes idle/Creating JumpStart Model from the training Job

Training job may take several hours due to setting of hyperparameters and the studio kernel may be in idle stage during the training phase. However, during this time, training is still running in SageMaker. If this happens, you can still deploy the endpoint using the training job name with the following code:

How to find the training job name? Go to Console -> SageMaker -> Training -> Training Jobs -> Identify the training job name and substitute in the following cell.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.