Train, Deploy and Optimize AI models with Amazon SageMaker JumpStart | Amazon Web Services

Key Takeaways

Hi, I'm Jazeppi Zappia, a IML specialist solution architect at AWS. In this video, I'm going to show you how to access a library of over 250 foundation models that you can quickly deploy, optimize, and even fine-tune in just a few clicks. Let's get into it. The process of model selection involves a few different steps. Usually, you have a use case in mind, possibly with some sample prompts related to it. Next, you'll do an investigation into model capabilities and requirements to scope it down to a manageable short list. Once you have a short list, you can either scope it down further through baseline prompts or begin deeper evaluation on a full data set. Afterwards, you can pick the model with the best balance of quality, speed, and cost. After looking at the standard model's performance, you can still leverage techniques like quantization or specular decoding to improve throughput. However, implementing these techniques may require specialized knowledge and time. In this example, you'll see how SageMaker Jumpstart can take the most of the heavy lifting out of this process. First, we need to access SageMaker Studio either through your IM identity center or the SageMaker AI console. I'm going to access through identity center. Once we're in, we're head over to the Jumpstart section in the left navigation. Jumpstart has models from dozens of providers as you see here. And in this example, we'll focus on models from Meta. Once we're into the model provider section, you can see all the available jump start models. Let's find a Llama 3.170B for this example. Here in the model page, you can see a full model card with background information on the model, how it was trained, and even evaluation results on standard benchmarks. You can use all of this data to reduce your list of model candidates. There's a few things you can do here. First, we can view any sample notebooks for this model. This particular model has two notebooks. One that will load the model and invoke it with sample data as well as show all of the supported payload parameters and another that shows how to fine-tune and deploy it. This could be useful as a first step of building your own notebooks or scripts. To use the notebook, you can click the open in Jupyter lab and it will allow you to use an existing studio space or create a new one for experimentation. Let's take a look at how we can optimize a model next. Clicking on optimize takes us to the inference optimization job console. You can see here that there are a few different techniques that we can leverage. Quantization reduces the precision of the data types of the weights and/or the activations of the model resulting in an increase in speed and a reduction in memory footprint with a slight penalty to model quality. You have three options for quantization. One is INT4 activation aware weight quantization which is a quantization techniques for LLMs that is efficient, accurate, low bit and weight only. FP8 is 8bit floating point low precision format for floatingoint numbers. It balances memory efficiency and model accuracy by representing values with fewer bits than the standard FP16 floatingoint format. Full FP8 quantization requires hardware support found on the NVIDIA L40, L40S, and H100/H200 GPUs. N8 smooth quant is an 8-bit data format which is a mixed precision quantization method that scales activations and weights jointly by balancing their dynamic ranges. Then you have speculative decoding. Speculative decoding is a technique that uses a smaller draft model to generate tokens that are then verified by the larger model which corrects any misgenerated tokens. This results in a small increase in memory footprint while providing an increase in throughput with no loss of model quality. Fast model loading is an optimization that prepares the model artifact so that it can be directly streamed to the GPU on the endpoint, bypassing the need of staging it to disk. This also breaks the model into equal chunks to match the number of GPUs so that this doesn't have to be done on the host machine at runtime. This will benefit larger models more than smaller ones and benefits both the initial model load as well as autoscaling. Optimizing the model artifact for a GPU power instance will compile it using NVIDIA Tensor RT LLM or using the Neuron SDK for AWS Tranium and AWS Inferentia instances. The model optimization job will run on an ephemeral instance and output the final model artifacts to S3. From there, you can deploy them via the console or the SDKs. For a simple deployment of a model to an endpoint without specifying any customization, we can click on deploy. This takes us to the model deployment console where we can adjust some of the parameters for model hosting. The llama 3.17b model here is an inference optimized model in jumpstart which automatically leverages a sagemaker draft model to improve throughput. You can also see some performance benchmarking information related to the selected instance type and control the inference configuration based on the expected concurrency by selecting it from the table. Note that higher user concurrency has an impact on latency as well as the tokens per second per user metrics and that the metrics are per instance. Finally, for situations where you're looking to customize the base model, you can select train to create a fine-tuning job. You can bring your data set to further train the model for your use case. And the training job will output your model artifacts to S3, allowing you to host it on SageMaker, import it into Bedrock if it's a supported model, or anywhere else. In this example, Jumpstart is providing a chat assistant data set as training data focused on making the model better at turn-based interactions. You can also perform instruction based fine-tuning to enhance the model's ability to perform a specific task by bringing it an appropriate data set. You can find more information on training data set formats in the jump start documentation. In the next section, you have all the available hyperparameters that will be submitted as part of the training job. You can adjust everything from the number of epochs to the learning rate to Laura configuration and even quantization. These give you control over the quality of your model outputs, training time and the amount of resources necessary to train the model. So you can optimize it based on your needs. With your hyperparameters all set up, you can select the infrastructure type for training which will be dependent on the model being used. The smaller models will have lower resource requirements and techniques like quantization can help reduce them further. Next, provide an S3 location for the training job to export the model artifacts for training. When the job completes, you'll be able to access the model here. After that is the infrastructure configuration for the training job. Here you can configure things like BPC access, encryption keys, and IM permissions. This ensures that you have the flexibility to provide access to resources and control over how they are accessed. If you'd like to customize the name of the job, you can do that in the last section. Then hit submit to start training. In this video, you learned how to use SageMaker Jumpstart to train, deploy, and optimize Foundation models for your generative AI applications. Jumpstart provides you over 250 models and a variety of techniques to allow you to get started quickly while still maintaining a high level of control over the entire process. That's all for this video. Thanks for watching and happy experimenting with SageMaker Jumpstart.