Task Statement 3.3: Describe the training and fine-tuning process for foundation models
Objectives:
- Describe the key elements of training an FM (for example, pre-training, fine-tuning, continuous pre-training, distillation).
- Define methods for fine-tuning an FM (for example, instruction tuning, adapting models for specific domains, transfer learning, continuous pre-training).
- Describe how to prepare data to fine-tune an FM (for example, data curation, governance, size, labeling, representativeness, reinforcement learning from human feedback [RLHF]).
Lesson 1
Let's get started with the third task statement from Domain 3, which is to describe the training and fine-tuning process for foundation models. This task statement is split into two lessons.
What are the key elements of training a foundation model? They include pre-training, fine-tuning, and continuous pre-training. Let's start with pre-training, which is a complex process. It requires millions of graphic processing units, GPUs, compute hours, terabytes and petabytes of data, trillions of tokens, trial and error, and time. During the pre-training, generative AI models learn its capabilities. What is the difference between pre-training and fine-tuning? With pre-training, you train the LLM by using huge amounts of unstructured data with self-supervised learning.
Even though foundation models have been trained on millions of documents, videos, images, files, audio files and more, and have learned the fundamentals of the human language, they might need additional training or instructions for your domain and dataset, or they might need to learn how to perform human tasks and reasoning. Fine-tuning is a process that extends the training of the model to improve the generation of completions for a specific task. It is a supervised learning process and you use a dataset of labeled examples to update the weights of the LLM. Fine-tuning helps to adapt foundation models to your custom datasets and use cases. Instruction based fine-tuning uses labeled examples for performance improvements on specific tasks. LLMs can perform many different tasks within a single model. However, what if your application needs to perform a single task? Well, you can fine-tune a pre-trained model to improve performance for the tasks specific to your use cases. However, there is a possible limitation to fine-tuning on a single task, which is catastrophic forgetting. Catastrophic forgetting happens when the whole fine-tuning process modifies the weights of the original LLM. This can improve the performance of the single task fine-tuning, but it can degrade performance on other tasks. And another consideration is to decide whether catastrophic forgetting will impact your use case. If you need reliable performance on a single task, then you might not be concerned that the model can't generalize to other tasks.
During full fine-tuning, every parameter in the model is updated through supervised learning, but there are parameter efficient methods to adapt your model. When you train and tune a foundation model, you load the model parameters and add memory for the optimizer, gradients, forward activations, and temporal memory. These components can add additional bytes of GPU memory for each model parameter, and they can also increase the compute costs of the fine-tuning process. Parameter-efficient fine-tuning, PEFT, is a process and set of techniques that freeze or preserve the parameters and weights of the original LLM and fine-tune or train a small number of task-specific adaptor layers and parameters. PEFT reduces the compute and memory that's needed because it's fine-tuning a small set of model parameters. Low-rank adaptation or LoRA, is a popular PEFT technique that also preserves or freezes the original weights of the foundation model and creates new trainable low-rank matrices into each layer of a transformer architecture. PEFT and LoRA modify the weights of your model, but not the representations.
Representations encode semantic information similar to embeddings. Representation fine-tuning, ReFT, is a fine-tuning process that freezes the base model and learns task-specific interventions on hidden representations. The linear representation hypothesis says that concepts are encoded in linear subspaces of representation in a neural network. Multitask fine-tuning is an extension of fine-tuning a single task. Multitask fine-tuning requires a lot of data. For this process, the training dataset has examples of inputs and outputs for multiple tasks. For example, the dataset might contain examples that instruct a model to complete multiple tasks such as reviews or ratings, summarization, translating code, and more. This produces an instruction tuned model that has learned how to complete many different tasks simultaneously. Additionally, you can calculate losses from the examples to update the weights of the model. This helps mitigate and avoid catastrophic forgetting. Here is a question.
What fine-tuning process modifies the weights of the model to adapt to domain-specific data? Domain adaptation fine-tuning gives you the ability to use the pre-trained foundation models and adapt them to specific tasks by using limited domain-specific data. You can use domain adaptation fine-tuning to help your model work with domain-specific language such as industry jargon, technical terms, or other specialized data. Amazon SageMaker JumpStart provides the capability to fine-tune a large language model, particularly a text generation model on the domain- specific dataset. You can fine-tune models with their custom dataset to improve performance in specific domains. You can also fine-tune reinforcement learning from human feedback, or RLHF. In this way you can improve the performance of your model and help it better understand human-like prompts to generate human-like responses. RLHF uses reinforcement learning to fine-tune the LLM with human feedback data to better align the model with human preferences. I am going to pause this lesson here, and in the next lesson we will continue with task statement 3.3.
Lesson 2
Let's continue with task statement 3.3 and talk about pre-training, fine-tuning, and continuous pre-training, but with a focus on preparing the data to fine-tune the foundation model. The first step is to prepare your training data. Many publicly available datasets and prompt template libraries have been used to train large language models. The prompt template libraries include many templates for different tasks and datasets. After your instruction dataset is ready, you can divide the dataset into training validation, and test splits. During fine-tuning, you select prompts from your training dataset and pass them to the LLM to generate completions. Then, compare the distribution of completions, and the training label, to calculate a loss between the two token distributions. You can use the calculated loss to update your model's weights.
After many batches of prompt completion pairs, update the weights so the model's performance on the task improves. As in standard supervised learning, you can define separate evaluation steps to measure your LLM's performance, by using the holdout validation dataset. You will get the validation accuracy, and after you've completed your fine-tuning, you can perform a final performance evaluation by using the holdout test dataset. This last result will give you the test accuracy. How can you prepare your data for the fine-tuning process in AWS? Well, let's pause for fundamentals and talk about data preparation. In machine learning, data preparation is the collecting pre-processing, and organizing of your raw data for your model. I have added flashcards to dive deeper, but for now, let's quickly review a few options of data preparation in AWS.
If you have low-code data preparation, you can use Amazon SageMaker Canvas to create data flows that define your ML data pre-processing. These data flows feature engineering workflows that use little to no coding. If you have data preparation that needs to scale, you can use open source frameworks such as Apache Spark, Apache Hive, or Presto. Amazon SageMaker Studio Classic provides a built-in integration with Amazon EMR. Also, if you need serverless, you can use the Apache Spark based serverless engine from AWS Glue interactive sessions. They will aggregate, transform, and prepare data from multiple sources in SageMaker Studio Classic.
If you need to use structured query language, or SQL, in SageMaker Studio for data preparation, you can use Jupyter Lab. If you have data preparation for feature discovery and storage, you can use Amazon SageMaker Feature Store to search, discover, and retrieve features for model training. You can also use it to provide a centralized repository to store feature data in a standardized format. Suppose that you have data preparation that needs to detect bias in your data. You can use Amazon SageMaker Clarify to analyze your data and detect potential biases across multiple facets. For example, SageMaker Clarify can help you detect whether your training data contains imbalanced representations or labeling biases between groups such as gender, race, or age. If you have data that needs to be labeled, you can use SageMaker Ground Truth to manage the data labeling workflows for your training datasets. Continuous pre-training is also important with generative AI for a few reasons. The output of generative AI models is non-deterministic, which makes validation more difficult.
It's important to choose metrics, benchmarks, and datasets that help evaluate the capabilities of the models and ensure that it doesn't create harmful outputs. When you continually pre-train models on data across different topics, genres, and context, over time, the models become more powerful. They learn how to use out of domain data better by accumulating wider knowledge and adaptability. This capability creates value for your organization. For example, continued pre-training in Amazon Bedrock helps you to train Amazon Titan Text Express, and Amazon Titan Text Lite FMs. You can customize them by using your own unlabeled data in a secure and managed environment. Alright, let's get started with task statement 3.4 in the next lesson.