AWS adds GPU partitioning to SageMaker for cheaper AI inference

Sharing the hardware

Amazon Web Services (AWS) announced that it is bringing NVIDIA Multi-Instance GPU (MIG) support to Amazon SageMaker Inference. This update allows developers to host multiple hardware-isolated models on a single accelerator rather than assigning an entire GPU to every individual deployment. By sharing the hardware, companies can better manage their resources when they have several models running on real-time endpoints.

Before this update, customers were often limited by the total number of GPUs available on an instance. Even if a model only required a small amount of memory or compute power, it still occupied a whole device. This often led to wasted capacity on expensive hardware, especially when using newer, high-performance GPUs that were too large for small models to justify.

The new feature addresses these constraints by carving a single GPU into smaller, fixed partitions. Each partition gets its own dedicated block of memory and a specific share of compute units. This hardware-level isolation ensures that one model's activity does not negatively impact the performance of another model sharing the same physical chip.

How partitioning works

The technology relies on SageMaker's inference components, which are the constructs used to deploy multiple models to a single endpoint. Developers can now use two optional API fields to define how the hardware is divided. One field declares the specific layout of the partitions, while the other selects the profile used by each individual inference component.

When a layout is defined, it applies uniformly across every accelerator on every instance in a production variant. For example, AWS showed how an ml.p4d.24xlarge instance with eight A100 GPUs could be divided into 24 different partitions of varying sizes. Each copy of an inference component then occupies exactly one of these requested partitions.

Because the models remain resident on the device, this arrangement avoids the typical delays seen in multi-model endpoints. In those older setups, a model might have to be loaded onto the shared device the first time it is called. With MIG, the model is already in its assigned partition, which helps maintain predictable performance and faster response times for users.

Performance and costs

To validate the new feature, AWS conducted tests comparing an H100 partition against a full A100 GPU. They used the Llama-3.1-8B-Instruct model to measure performance across different traffic shapes. The results showed that the H100 partition provided 53% to 89% more output tokens per second at the specific load levels where the cost benefits were most significant.

The testing also revealed that the H100 partition had lower inter-token latency and a faster time to first token compared to the full A100. At the highest concurrency levels, the A100 failed to meet serving deadlines on up to half of the requests, while the H100 partition successfully met every deadline. This suggests that partitioning can be more reliable for heavy, long-context workloads.

However, the cost benefits are not universal and depend on the specific workload. AWS found that the H100 partition resulted in an 18% to 34% lower cost per million output tokens once traffic passed a certain break-even point. For very low-concurrency or short-input prompts, the older whole-GPU approach remained the cheaper option for the user.

Limitations and requirements

While partitioning increases the number of models an instance can host, it does not change the underlying billing structure. Customers continue to pay for the full GPU instance regardless of how many partitions they create. The goal of the feature is to increase usable capacity and efficiency rather than to lower the hourly price of the hardware itself.

There are several technical limitations to keep in mind for developers. The feature is currently unavailable for single-model endpoints, serverless endpoints, and heterogeneous instance pools. It also does not support asynchronous workloads because those types of calls cannot address specific inference components in the same way.

Additionally, the layout of a partitioned endpoint cannot be changed once it is running. If a developer needs to modify the partition sizes or the layout, they must create a new endpoint and migrate the traffic. Furthermore, partitioned endpoints do not currently publish GPU, CPU, or memory utilization metrics to CloudWatch, though they still provide standard invocation and latency metrics.

Deployment and regions

AWS has made this feature available across several major regions, including parts of the US, Asia Pacific, and Europe. Availability depends on the specific regional instance quotas and the availability of the supported hardware families. The launch covers four specific instance families, with the exact partition counts determined by the type of accelerator used.

For those using P4d instances, specific driver requirements apply. The mig-1g.10gb profile requires NVIDIA’s R525 driver or newer, which means AWS recommends using a specific newer inference AMI. The default images provided by AWS may not include the necessary drivers to support these specific partition sizes.

Salesforce helped define the requirements for this launch and validated the initial implementation. Richard Chang, a software engineering architect at Salesforce, noted that the company expected the feature to consolidate workloads safely. They are currently preparing to validate the software with representative production workloads to see how it performs in real-world scenarios.