Customizing foundation models with Amazon SageMaker AI | Amazon Web Services

Key Takeaways

Hi, I'm Jazeppi Zappia, a IML specialist solution architect at AWS. The topic of foundation model customization spans a wide gamut of possibilities from prompt engineering and retrieval augmented generation all the way to supervised fine-tuning, preference optimization, continued pre-training, and even building models from scratch. In this video, I'm going to show how to customize your own foundation models using Amazon SageMaker AI. Let's get into it. Fine-tuning is used as a customization technique when you're trying to get a foundation model to perform better or in a very specific way when presented with a given task. This could reduce the amount of input or output tokens, correct inconsistency in model outputs, and even provide additional capabilities that didn't exist before, such as tool calling or advanced reasoning. When fine-tuning a model, the approaches fall in a few different buckets. First is quick experimentation with a small data set and a small number of epochs. Then you could do single instance training with bigger data sets and possibly with multiple GPUs. Then you have large-scale distributed training workloads with large models and data sets spanning multiple instances. SageMaker Studio spaces support GPU powered instances and could be used for smaller scale model customization where you're focused on quick iterations in a notebook environment but don't allow you to scale to multi-node training and results in a tightly coupled development environment. SageMaker managed training jobs provide ephemeral scalable resources for both single instance and distributed training workloads at scale. Managed training jobs support bringing your own custom containers as well as using managed containers where you can bring your own training script. This variety of options allows you to choose the one that gives you the flexibility without sacrificing time to value. SageMaker HyperPod provides long-running clusters that can be orchestrated via Slurm or Kubernetes, optimized for large-scale distributed training and provides direct access to the underlying instances via SSH. This provides the highest level of control while still taking advantage of managed infrastructure. SageMaker training jobs are tailored for organizations that want to focus on model development rather than infrastructure management and prefer ease of use with a managed experience. SageMaker training jobs feature a user-friendly interface, simplified setup and scaling, automatic handling of distributed training tasks, built-in synchronization, checkpointing, fault tolerance, and abstraction of infrastructure complexities. For organizations that require granular control over training infrastructure and extensive customization options, SageMaker Hyperpod is the ideal choice. HyperPod offers custom network configurations, flexible parallelism strategies, and support for custom orchestration techniques. It integrates seamlessly with tools such as Slurm, Amazon EKS, Nvidia's Nroot, and Pixus, and provides SSH access for in-depth debugging and custom configurations. Let's take a look at some code. In this example, I'll be using a SageMaker Studio code editor space and walk you through training models both through HyperPod and through manage training jobs. We'll start off with HyperPod. Here I have a SageMaker HyperPod cluster that is a mix of CPU and GPU powered instances using Amazon EKS as the orchestrator. This allows my Kubernetes cluster to see the HyperPod cluster as a set of nodes that I can directly deploy containers or pods to. The HyperPod nodes can additionally leverage shared persistent file systems that are available to the Kubernetes cluster. In this example, I've attached an Amazon FSX for Luster file system to the cluster which will provide high performance disk access for models and data. I've also taken advantage of a feature in SageMaker Studio that allows me to mount volumes to my code editor space. This way, I can load larger resources directly onto the FSX for Lustre file system without having to continuously shuffle data into my training job, accelerating my development cycle. For simplicity, I've already staged my training data and script on FSX. To submit your training job, you can use the cube cuddle command and a file with the configuration information. This job is for fine-tuning a DeepSeek distilled quen 7B model. Jumping right into the PyTorch job, you can see that the FSX volume here is tied to the persistent volume claim, allowing us to use the FSX for luster file system. The container for this job is one of the SageMaker managed PyTorch containers, but you could also have created your own. The invidia.com/GPU resource will ensure that this job gets allocated one GPU. Then you have some environment configuration followed by the entry point of the container. This command will install dependencies from a requirements.ext file. Then run our training script with a set of training arguments. The paths for these files correspond to the mount point at the end here where the container will mount the fsx for luster file system on /data. Running cubecuddle apply-fpod_finetuning.yaml will schedule this pod for deployment and fire off the training job. The training script will write the train model artifacts back to the fsx volume so it can be prepared for deployment. This will take a few minutes to complete. Running cube cuddle get pods will show the currently running pods on the cluster. You can see that the worker pod for training job has started. Running cube cuddle logs in the name of the pod will allow you to view the standard output from the training container. Right now it is installing dependencies and will eventually show a progress indicator once training begins. We'll come back to this once training has completed. The logs now show that the training job is finished and running cubecuddle get pods shows the training worker pod in a completed state. The train model artifacts are now located on the fsx for luster volume that we set up in the training configuration and is ready for deployment. You can deploy this model through SageMaker Hyperbot inference on the SageMaker managed endpoint, Amazon EKS, Amazon EC2 or anywhere else that you'd like. Let's go through a similar exercise with a manage training job. I've already installed my dependencies and downloaded a medical reasoning data set from hugging face for this example. I prepared a prompt template that I'll apply to all of my training samples. Then upload the files to S3 so that they could be downloaded by the training job. You can also provide an FSX for luster file system to the training job similar to the hyperpod example but I wanted to show a different method here. Next is the training configuration where I've defined all the locations for where the training script can expect to find data training hyperparameters and where to output the final model artifact. The training job will package these files in this folder as a tar file and upload it to S3 upon completion of the job. I'm leveraging PyTorch's fully sharded data parallel capability in the configuration here, which will take advantage of all available GPUs on my training instance to speed up the process. I'll upload my training configuration to S3, then set up my model trainer object. The SageMaker model trainer SDK construct simplifies the steps necessary to build SageMaker training jobs. All I need to do is specify my training container infrastructure configuration, the training parameters, and any custom code I want to use for the job. Then the path from where my training data is stored. After that, I call train and a SageMaker training job will be created for me. This will take a little while to complete. After the training job finishes, I can fetch the last training job from the model trainer, which gives me the path to the fine-tuned model artifacts. From there, I could use SageMaker's large model inference container, supply my model artifacts and inference configuration, and deploy the model for inference. Amazon SageMaker AI provides several manage hosting options for your models. Realtime inference is a model deployment option designed for use cases requiring immediate low latency responses and high throughput. It provides a persistent fully managed endpoint that can handle continuous inference requests and deliver responses in real time. Serverless inference is a pay-peruse option that provides ondemand inference without configuring or managing the underlying infrastructure. This is ideal for CPUbased inference workloads that have idle periods between spurts of traffic and can tolerate cold starts. Asynchronous inference is optimal for nearreal-time use cases. Inference that requires long processing times or have requests with large payloads. Asynchronous inference cues incoming requests and processes them on a fleet of instances, scaling the instance count to zero when there are no requests to process. Batch transform is ideal for longunning offline batch inference jobs on large data sets that don't require a persistent real-time endpoint. Batch transform handles the distribution of your data sets and models, removing the undifferiated heavy lifting of building your own job orchestration. Through these options for model hosting, SageMaker AI allows you to choose the one that best aligns with your application's specific requirements. Now that the HyperPod job has completed, you can see the fine-tuned model artifact stored in the FSX file system I mounted to the code editor space. From here, I could package the model artifact up myself or as part of the training script and store it in S3 for deployment to a managed endpoint. Once the SageMaker real-time endpoint has fully deployed, the model is ready to serve inference. I've created a predictor object for my endpoint using the SageMaker Python SDK that I can use to invoke the model. I've taken a sample data point from the medical data set we used for training earlier. I'll format it for my prompt and call predict on the predictor with the desired inference parameters to get a response. The inference response latency will be dependent on the size of the model and the type of instance serving the model. Finally, you can see the output from my request. The model is reasoning through the scenario I presented and giving an answer. Last, you can terminate your endpoint by running the delete model and delete endpoint APIs which will stop any of the running costs for hosting. To recap, in this video, you learned how you can use either SageMaker fully managed training jobs or SageMaker Hyperpod to fine-tune a foundation model on custom data set. Thanks for watching and happy experimenting with fine-tuning.