Fine-tune Llama3.1 405B Base and Instruct Variants on 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 fine-tune Llama3.1 and further deploy the fine-tuned model on SageMaker JumpStart.

Below is the content of the notebook.

1. Setup
2. Deploy model
3. Fine-tune model
4. Evaluate the pre-trained and fine-tuned model

The notebook requires users to specify following variables to start with.

• Specify model_id (default value: meta-textgeneration-llama-3-1-405b-fp8) (other option: meta-textgeneration-llama-3-1-405b-instruct-fp8)
• Specify accept_eula argument to be True in model.deploy() to accept the end-user license agreement (EULA) before deployment the model in an endpoint, given Llama model is gated.
• Sepcify {"accept_eula": "true"} in argument environment of JumpStartEstimator to accept the end-user license agreement (EULA) before fine-tuning.

The fine-tuning job will require more than 5 hours for the Llama3.1 405b model due to the large number of model parameters.

1. Setup

First, upgrade to the latest sagemaker SDK to ensure all available models are deployable.

Select the desired model to deploy. The provided dropdown filters all text generation models available in SageMaker JumpStart.

2. Deploy model

Create a JumpStartModel object, which initializes default model configurations conditioned on the selected instance type. JumpStart already sets a default instance type, but you can deploy the model on other instance types by passing instance_type to the JumpStartModel class. The deployment requires single ml.p5.48xlarge EC2 instance.

You can now deploy the model using SageMaker JumpStart. If the selected model is gated, you will need to accept the end-user license agreement (EULA) prior to deployment. This is accomplished by providing the accept_eula=True argument to the deploy method. The deployment might take few minutes.

Invoke the endpoint

This section demonstrates how to invoke the endpoint using example payloads that are retrieved programmatically from the JumpStartModel object. You can replace these example payloads with your own payloads. Inference payload parameters are provided as below.

• 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.
• 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.
• 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.

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

JumpStart stores model-specific default example payloads in its SDK. You can retrieve and view them using following code.

3. Fine-tune model

You can fine-tune on the dataset with domain adaptation format or fine-tuning with chat / conversational dataset. Please find more details in the section Dataset formatting instruction for training. In this demo, we will use a subset of OpenAssistant's TOP-1 Conversation Threads as an example dataset for chat fine-tuning. It can be downloaded from here. It contains roughly 13,000 samples of conversations between the Assistant and the user.

Training data should be formatted in JSON lines (.jsonl) format, where each line is a dictionary representing a set of conversations. Here is an example of a line in the training file. The key has to be one of dialog, messages, and conversations to be compatible with chatml format in HuggingFace SFTTrainer.

{"dialog": [{"content":"what is the height of the empire state building","role":"user"},{"content":"381 meters, or 1,250 feet, is the height of the Empire State Building. If you also account for the antenna, it brings up the total height to 443 meters, or 1,454 feet","role":"assistant"},{"content":"Some people need to pilot an aircraft above it and need to know.\nSo what is the answer in feet?","role":"user"},{"content":"1454 feet","role":"assistant"}]}

Upload dataset to S3

Retrieve and customize hyperparameters

Underlying the training scripts, JumpStart leverages HuggingFace SFTTrainer with QLoRA.

• For chat fine-tuning, specify chat_dataset to be True
• For domain adaptation fine-tuning, specify chat_dataset to be False

Note. The fine-tuning model artifacts for 405B fine-tuning are in its original precision bf16. After QLoRA fine-tuning, we conducted fp8 quantization (code source) on the trained model artifacts in bf16 to make it deployable on single ml.p5.48xlarge instance.

Deploy the fine-tuned model

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

4. 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 the endpoint

Don't forget to clean up resources when finished to avoid unnecessary charges.

Appendix

1. Dataset formatting instruction for training

Fine-tune the Model on a New Dataset

We currently offer two types of fine-tuning: chat fine-tuning and domain adaption fine-tuning. You can easily switch to one of the training methods by specifying parameter chat_dataset being 'True' or 'False'.

1.1. Chat fine-tuning

The Text generation model can be fine-tuned on the chat dataset, provided that the data is in the expected format. The resulting chat model can be further deployed for inference. Below are the instructions for how the training data should be formatted for input to the model.

Below are the instructions for how the training data should be formatted for input to the model.

• Input: A train and an optional validation directory. Train and validation directories should contain one or multiple JSON lines (.jsonl) formatted files. All training data must be in a single folder, however it can be saved in multiple jsonl files. The .jsonl file extension is mandatory.
  • The training data must be formatted in a JSON lines (.jsonl) format, where each line is a dictionary representing a single data sample. Each line in the file is a list of conversations between the user and the assistant model. This model only supports 'system', 'user' and 'assistant' roles, starting with 'system', then 'user' and alternating (u/a/u/a/u...).
• Output: A trained model that can be deployed for inference.

The best model is selected according to the validation loss, calculated at the end of each epoch. If a validation set is not given, an (adjustable) percentage of the training data is automatically split and used for validation.The training data must be formatted in a JSON lines (.jsonl) format, where each line is a dictionary representing a single data sample.

Here is an example of a line in the training file:

{"dialog": [{"content":"what is the height of the empire state building","role":"user"},{"content":"381 meters, or 1,250 feet, is the height of the Empire State Building. If you also account for the antenna, it brings up the total height to 443 meters, or 1,454 feet","role":"assistant"},{"content":"Some people need to pilot an aircraft above it and need to know.\nSo what is the answer in feet?","role":"user"},{"content":"1454 feet","role":"assistant"}]}

1.2. Domain adaptation fine-tuning

The Text Generation model can also be fine-tuned on any domain specific dataset. After being fine-tuned on the domain specific dataset, the model is expected to generate domain specific text and solve various NLP tasks in that specific domain with few shot prompting.

Below are the instructions for how the training data should be formatted for input to the model.

• Input: A train and an optional validation directory. Each directory contains a CSV/JSON/TXT file.
  • For CSV/JSON files, the train or validation data is used from the column called 'text' or the first column if no column called 'text' is found.
  • The number of files under train and validation (if provided) should equal to one, respectively.
• Output: A trained model that can be deployed for inference.

Below is an example of a TXT file for fine-tuning the Text Generation model. The TXT file is SEC filings of Amazon from year 2021 to 2022.

This report includes estimates, projections, statements relating to our
business plans, objectives, and expected operating results that are “forward-
looking statements” within the meaning of the Private Securities Litigation
Reform Act of 1995, Section 27A of the Securities Act of 1933, and Section 21E
of the Securities Exchange Act of 1934. Forward-looking statements may appear
throughout this report, including the following sections: “Business” (Part I,
Item 1 of this Form 10-K), “Risk Factors” (Part I, Item 1A of this Form 10-K),
and “Management’s Discussion and Analysis of Financial Condition and Results
of Operations” (Part II, Item 7 of this Form 10-K). These forward-looking
statements generally are identified by the words “believe,” “project,”
“expect,” “anticipate,” “estimate,” “intend,” “strategy,” “future,”
“opportunity,” “plan,” “may,” “should,” “will,” “would,” “will be,” “will
continue,” “will likely result,” and similar expressions. Forward-looking
statements are based on current expectations and assumptions that are subject
to risks and uncertainties that may cause actual results to differ materially.
We describe risks and uncertainties that could cause actual results and events
to differ materially in “Risk Factors,” “Management’s Discussion and Analysis
of Financial Condition and Results of Operations,” and “Quantitative and
Qualitative Disclosures about Market Risk” (Part II, Item 7A of this Form
10-K). Readers are cautioned not to place undue reliance on forward-looking
statements, which speak only as of the date they are made. We undertake no
obligation to update or revise publicly any forward-looking statements,
whether because of new information, future events, or otherwise.
GENERAL
Embracing Our Future ...

License: Apache 2.0

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.