Skip to content

Preloading models to high-speed local storage to accelerate LLM startup

10 minute read
Content level: Advanced
0

Instead of streaming models from network-based storage, this article describes how to preload model weights onto high-speed local storage to reduce startup times for large language models (LLMs).

Introduction

LLMs have become foundational to modern Generative AI (GenAI) applications, powering capabilities from natural language understanding to content creation and conversational AI. As enterprises scale their AI initiatives, the ability to efficiently deploy and serve these massive models directly affects performance, user experience, and operational costs. Because LLMs often exceed hundreds of gigabytes, the time it takes to load these models into GPU memory is a critical bottleneck that affects the speed and responsiveness of GenAI services.

Modern LLMs are massive in size and often exceed 900 GB in size. It can take several minutes to directly load these giants from remote storage, especially in distributed environments with multiple nodes. When you run expensive GPU instances such as P4 or P5, every second of loading time directly translates to wasted resources and higher costs.

Recently, AWS Support worked with a global enterprise customer that was facing significant performance bottlenecks in their LLM operations. The customer's AI team struggled with models that exceeded their storage and time limits.

The AWS Enterprise Support team collaborated closely with the customer through their Technical Account Manager (TAM). To resolve this issue, the customer's TAM recommended a clever preloading step that copied the model to lightning-fast local Non-Volatile Memory Express (NVMe) storage during pod initialization. Combined with Elastic Fabric Adapter (EFA) for accelerated data transfer, this technique reduces model load latency and boosts the customer's efficiency, scalability, and bottom line.

AWS Enterprise Support interaction

The AWS Enterprise Support team collaborated closely with the customer through their TAM. After analyzing the customer's workload, the Specialist STAM recommended a preloading step that copies models to fast local NVMe storage during pod initialization. This approach significantly reduces model load latency while enhancing efficiency, scalability, and cost-effectiveness across their operations.

Solution overview

To implement the solution, the team had to complete the following tasks:

  1. Publish the machine learning (ML) model to a shared Amazon FSx for Lustre volume. This volume is EFA-enabled to maximize network throughput during model transfer.

  2. Configure an init container to perform 32 concurrent copy operations to transfer the model from FSx to the instance's local ephemeral NVMe storage. This fully uses the aggregate EFA bandwidth.

To implement this solution, the team undertook several critical tasks designed to optimize the model loading process and minimize inference startup latency. First, the team published the LLM model to a shared FSx for Lustre file system that's configured with EFA support. The team chose FSx for Lustre because of its high throughput and parallel file access capabilities, which made it well suited to handle multi-hundred-gigabyte models. By turning on EFA on the FSx for Lustre volume, the solution benefits from significantly enhanced network bandwidth and reduced latency. These benefits are essential for accelerating large data transfers between storage and compute nodes.

Next, the solution uses Kubernetes init container functionality to perform a highly parallelized model preloading step during pod initialization. The init container executes a copy operation that runs 32 concurrent file transfers from the FSx for Lustre volume to the instance's local ephemeral NVMe storage. This parallelism maximizes use of the aggregated EFA network bandwidth and the instance's CPU and disk I/O resources. This use reduces a multi-minute copy operation to a fraction of the time. By preloading the model onto the local NVMe storage, the system avoids the costly overhead of streaming model data over the network at inference runtime.

This preloading step is seamlessly integrated into the vLLM deployment workflow. After the init container copies the model, the main vLLM container loads the model directly from the high-speed local NVMe volume. This shift from network-based streaming to local disk access significantly improves startup times, and allows inference workloads to begin serving requests almost immediately after pod startup. This optimization not only enhances user experience by reducing latency, but also minimizes the idle time of expensive GPU instances to improve resource efficiency.

The use of EFA is pivotal in this architecture. EFA provides a high-bandwidth, low-latency networking interface that accelerates data transfer between FSx for Lustre and Amazon Elastic Compute Cloud (Amazon EC2) instances. This enhanced networking capability makes sure that network constraints don't bottleneck the parallel copy operations in the init container. This capability also allows for rapid model preloading even for very large models. As a result, the time from pod launch to inference readiness significantly decreases, improving overall system throughput and cost-effectiveness.

For a complete Kubernetes deployment manifest that demonstrates end-to-end setup, see the Accelerated model loading in vLLM section of this article.

Enter image description here

Solution overview

Prerequisites:

  • An Amazon Elastic Kubernetes Service (Amazon EKS) cluster

  • GPU-based instances with local NVMe storage, such as g6 or p5

  • NVIDIA device plugin for Kubernetes. For more information, see NVIDIA device plugin for Kubernetes on the GitHub website.

  • (Optional) FSx for Lustre with EFA support

Turning on FSx for Lustre with EFA

Because the customer used storage with EFA turned on, the team first set up the necessary components in the Kubernetes cluster. To do this, the team installed the FSx CSI Driver and EFA Device Plugin Helm charts to use FSx for Lustre storage with EFA networking:

  1.  Add the Helm repositories:
helm repo add aws-fsx-csi-driver https://kubernetes-sigs.github.io/aws-fsx-csi-driver
helm repo add eks https://aws.github.io/eks-charts
helm repo update
  1. Install the FSx driver:
helm upgrade -i aws-fsx-csi-driver aws-fsx-csi-driver/aws-fsx-csi-driver \
  --namespace kube-system \
  ---version 1.10.0
  1.  Install the EFA Device plugin:
helm upgrade -i aws-efa-k8s-device-plugin eks/aws-efa-k8s-device-plugin \
  --namespace kube-system
  1.  If your base Amazon Machine Image (AMI) doesn't include Lustre Client, then install Lustre Client:
dnf install lustre-client -y
  1.  Create an FSx storage class with EFA turned on:
kind: StorageClass
apiVersion: storage.k8s.io/v1
metadata:
  name: fsx-sc-efa
provisioner: fsx.csi.aws.com
allowVolumeExpansion: true
reclaimPolicy: Delete
parameters:
  subnetId: subnet-example
  securityGroupIds: sg-example
  deploymentType: PERSISTENT_2
  perUnitStorageThroughput: "1000"
  fileSystemTypeVersion: "2.15"
  metadataConfigurationMode: "AUTOMATIC"
  efaEnabled: "true"
mountOptions:
  - flock
  1.  Create a persistent volume claim:
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: fsx-claim-efa
spec:
  accessModes:
    - ReadWriteMany
  storageClassName: fsx-sc-efa
  resources:
    requests:
      storage: 14400Gi

Use an init container to preload the model

The key part of this solution is to use an init container to preload the LLM onto the high-speed local NVMe storage before the main vLLM container starts. The init container is a specialized container that runs before the app containers in the pod and performs one-time setup tasks.

Add an init container to your pod spec that copies the model from FSx for Lustre to local NVMe-backed storage:

      initContainers:
        - name: init
          image: docker.io/busybox:1.37.0
          command: ["/bin/sh"]
          args:
          - "-c"
          - "find /fsx/models/large-model -type f | xargs -P32 -I{} cp {} /local/"
          volumeMounts:
          - name: local
            mountPath: "/local"
          - name: fsx
            mountPath: "/fsx"
      volumes:
        - name: local
          emptyDir:
            sizeLimit: 1024Gi
        - name: fsx
          persistentVolumeClaim:
            claimName: fsx-claim-efa

The preceding example uses Bash script to recursively find all the files in the /fsx/models/large-model directory and copy them in parallel (-P32) to the /local directory.

The script then mounts the /local volume into the parent vLLM container that the model subsequently loads onto:

      containers:
        - name: vllm
          image: docker.io/vllm/vllm-openai:v0.8.5.post1
          args:
            - '--model=/local/'
          volumeMounts:
            - name: local
              mountPath: /local
      volumes:
        - name: local
          emptyDir:
            sizeLimit: 1024Gi

Accelerated model loading in vLLM

After you copy the model to local storage, the vLLM server starts and loads it directly from NVMe. This avoids performance penalties of network-based model loading during startup.

The following Kubernetes deployment example shows a complete end-to-end setup that preloads a large model onto high-speed local NVMe storage before launching the vLLM inference container.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: large-model
spec:
  selector:
    matchLabels:
      model: large-model
  template:
    metadata:
      labels:
        model: large-model
    spec:
      initContainers:
        - name: init
          image: docker.io/busybox:1.37.0
          command: ["/bin/sh"]
          args:
          - "-c"
          - "find /fsx/models/large-model -type f | xargs -P32 -I{} cp {} /local/"
          volumeMounts:
          - name: local
            mountPath: "/local"
          - name: fsx
            mountPath: "/fsx"
      containers:
        - name: vllm
          image: docker.io/vllm/vllm-openai:v0.8.5.post1
          imagePullPolicy: IfNotPresent
          args:
            - '--model=/local/'
          volumeMounts:
            - name: local
              mountPath: /local
      volumes:
        - name: local
          emptyDir:
            sizeLimit: 1024Gi
        - name: fsx
          persistentVolumeClaim:
            claimName: fsx-claim-efa

After you deploy the init container, the init container automatically runs on pod startup and copies the model files from the FSx volume to the local NVMe storage. This process makes sure that the vLLM container runs with the model that's already present on a local NVMe volume, minimizing startup latency.

Optimizing your performance

To optimize the solution, review the following best practices:

  • Adjust your parallel copying (xargs -P32): Adjust the parallelism level based on the CPU capabilities of the instance. If you have too many concurrent copies, then you might experience I/O contention.

  • Optimize your EFA bandwidth: To optimize your transfer speeds, make sure that EFA is turned on for both the FSx file system and Amazon EC2 instances.

  • Review your emptyDir storage limits: The emptyDir volume is backed by instance storage. Make sure that your node type, such as p5.48xlarge, provides sufficient local NVMe capacity to hold the full model.

Cleanup

To avoid unnecessary costs, make sure to clean up all resources after you implement the solution:

  • Delete the vLLM deployment: kubectl delete deployment large-model

  • Remove persistent volume claims: kubectl delete pvc fsx-claim-efa

  • If you don't need the plugins, then uninstall the Helm charts for FSx CSI and EFA plugins:

helm uninstall aws-fsx-csi-driver -n kube-system  helm uninstall aws-efa-k8s-device-plugin -n kube-system

Conclusion

This solution demonstrates a highly effective strategy for reducing startup times for LLMs by preloading the LLMS onto high-speed local storage. Through EFA, you can accelerate data transfer from network-attached storage. The preloading step can then complete rapidly, and the inference framework can initialize with minimal delay. This approach improves operational efficiency and cost-effectiveness, particularly in environments with large GPU fleets where it's critical to minimize idle time.

To get the most out of your AWS environment, contact your TAMs. They can help you with general guidance, best practices, troubleshooting, and operational support on AWS. To learn more about our plans and offerings, see AWS Support. To learn more about the suggested solution, contact your TAM or AWS account team.

About the authors

Enter image description here

Nick Zverev Nick Zverev is a STAM at AWS. He helps Enterprise customers use containers and cloud-native technologies to support their digital transformation and AI/ML initiatives.

Enter image description here

Alma Mohapatra Alma Mohapatra is an Enterprise Support Manager who supports strategic AI/ML customers running on high performance computing (HPC) environments. She leads specialized technical teams that help organizations optimize their large-scale ML workloads across distributed GPU clusters. With extensive experience in enterprise ML operations, Alma guides customers through complex performance challenges and infrastructure optimization for training LLMs. She excels at translating technical requirements into practical solutions that enhance reliability and efficiency in production AI environments. Alma is known for her collaborative approach to customer success, and works closely with TAMs to make sure that AI/ML initiatives meet critical business objectives.

Enter image description here

Alan Halcyon Alan Halcyon is a Senior STAM at AWS in the Containers domain. He helps AWS customers optimize and troubleshoot large-scale Kubernetes workloads, and loves weird and difficult problems.