The rise of Large Language Models (LLMs) has fundamentally shifted infrastructure requirements. Unlike traditional microservices, LLM inference workloads are resource-intensive and highly sensitive to hardware capabilities. As organizations move from development to production, they encounter a fragmented hardware landscape: GPU clusters for training and high-throughput inference, CPU-only instances for edge or cost-optimized deployments, and varying CPU architectures (amd64 vs. arm64) across cloud providers. Ensuring your LLM inference containers run reliably across this spectrum requires sophisticated build strategies.
The Challenge of Heterogeneous Hardware
A standard Docker image built on a development laptop with an NVIDIA GPU will fail on a CPU-only production node, and vice versa. Furthermore, mixing CUDA libraries into a base image creates bloat and dependency conflicts if the target machine lacks NVIDIA drivers. In LLMOps, we must separate concerns: the application logic, the model weights, and the runtime dependencies (CUDA/CPU) must be modular.
Strategy 1: Separate Base Images by Architecture
Instead of one monolithic "universal" image, maintain distinct base images for different hardware targets.
- GPU Image: Based on
nvidia/cuda:11.8-runtime-ubuntu22.04. Includes CUDA toolkits and GPU-specific libraries likeflash-attn. - CPU Image: Based on
python:3.10-slim. Includes MKL or OpenBLAS for optimized CPU math, but excludes all CUDA dependencies.
Strategy 2: Modular Dependency Management
Use Python’s conditional dependency mechanism to install packages only when relevant hardware is detected. This prevents ImportError issues on incompatible nodes.
# requirements-gpu.txt
torch==2.0.1+cu118
transformers==4.36.0
flash-attn==2.3.6 # GPU only
# requirements-cpu.txt
torch==2.0.1
transformers==4.36.0
In your Dockerfile, use build arguments to select the correct requirements file:
ARG BUILD_VARIANT=cpu
RUN if [ "$BUILD_VARIANT" = "gpu" ]; then \
pip install -r requirements-gpu.txt; \
else \
pip install -r requirements-cpu.txt; \
fi
Strategy 3: Leveraging BuildX for Multi-Arch Builds
For CPU targets, you must support both amd64 and arm64 (common in AWS Graviton and Apple Silicon). Use Docker Buildx to create multi-architecture manifests without needing actual hardware for every arch.
docker buildx create --use
docker buildx build --platform linux/amd64,linux/arm64 \
--build-arg BUILD_VARIANT=cpu \
-t my-llm-app:cpu-mixed .
Strategy 4: Model Weight Separation
Never bake large model weights (7B-13B parameters) into your Docker image. Images become multi-GB in size, slowing CI/CD pipelines and registry pulls. Instead, use volume mounts or pull weights from a registry (like Hugging Face Hub) at container startup. This ensures your image remains lightweight and versionable, while models can be updated independently.
CI/CD Implementation Best Practices
1. Cache Layers Aggressively: Keep pip install for core dependencies in early layers. Only copy application code in later layers.
2. Health Checks: Implement CMD HEALTHCHECK that verifies GPU availability (using nvidia-smi) or CPU core count before declaring the service ready.
3. Tagging Convention: Use tags like my-llm:v1.0-gpu and my-llm:v1.0-cpu to prevent accidental deployment mismatches.
Conclusion
Containerizing LLM inference is not about creating one size-fits-all image. It’s about building a portfolio of optimized, hardware-specific artifacts. By decoupling runtime dependencies, leveraging build arguments, and separating model weights from application code, you can achieve reliable, scalable, and cost-efficient LLM deployments across any hardware topology. Start with modular Dockerfiles early in your LLMOps pipeline to avoid technical debt as your infrastructure scales.