Optimizing an AI Model

We developed a fairly large model for an internal AI service. Used as-is, it filled an entire 96 GB GPU. I first examined memory and speed optimizations, but one requirement mattered more: neither improvement could come at the cost of model quality. Most quantization research accepts some trade-off, and in an environment where large-scale post-training was unavailable, few methods preserved quality. A method whose paper reported little degradation behaved differently on our service model, so we chose not to deploy it.

Splitting the Model

The model happened to divide cleanly into two functionally independent parts, each occupying about 48 GB. Micromanaging the resulting system taught me several lessons.

1. Loading Can Temporarily Double Parameter Memory

If new tensors are allocated on the GPU before old parameters or checkpoint buffers are released, identical values occupy separate memory regions.

It is like reserving space both for moving boxes and for those same boxes inside the new house. Common inefficient patterns include loading both the model and checkpoint directly to CUDA:

device = 'cuda'
model = MyModel().to(device)
ckpt = torch.load('model.pt', map_location=device)
model.load_state_dict(ckpt['model'])

model = MyModel().to('cuda')
state_dict = torch.load('model.pt')
model.load_state_dict(state_dict)

It is safer to load on CPU, apply the state dict, and move the completed model afterward:

ckpt = torch.load(path, map_location='cpu')
model.load_state_dict(ckpt)
model.to('cuda')

Surprisingly many codebases contain variants of this inefficiency, so refactoring before production use matters.

2. Preprocessing Efficiency Matters as Much as Model Efficiency

Our media model required video preprocessing, and CPU load often exceeded GPU load. I moved preprocessing to CUDA decoding and rewrote CPU-based open-source components for the GPU. A fast model does not help when decoding remains the bottleneck.

3. Normalization and Error Fallbacks Are Essential

Video encodings vary enormously. CUDA decoding could hang and freeze the API, so I isolated it in a process and fell back to CPU after a timeout. An exception can be caught with try/except; a hang cannot. An earlier thread-based attempt merely accumulated zombie work.

FlashAttention Is Remarkably Effective

FlashAttention improved speed and sharply reduced memory use. A model that had approached 48 GB under larger batches stayed below 24 GB after the change. Docker images built around it still require careful compatibility checks among CUDA, PyTorch, and each server's driver.

Losing CUDA Visibility after NVML Environment Changes

One production failure was a container that initially passed torch.cuda.is_available() and initialized NVML, but appeared to lose its CUDA devices after several days. Already loaded inference continued because its CUDA Graph remained connected to GPU memory, while new CUDA contexts and NVML initialization failed.

The suspected cause was disabled cgroup support in /etc/nvidia-container-runtime/config.toml:

--no-cgroups=true

NVIDIA Container Runtime uses cgroups to track GPU allocation, manage NVML and visibility, retain device IDs, and isolate access permissions. Without it, a long-running container had no reliable way to preserve and revalidate the device. Enabling cgroups helped, but the issue later returned on a particular server; granting the container privileged = true ultimately resolved that environment. The broader lesson was that long-lived CUDA containers require active maintenance.

Is MIG Essential?

After reducing inference memory, we could run several containers on one GPU. I tested whether NVIDIA MIG—isolating one GPU into instances with separate memory bandwidth, cache, and compute—would improve throughput and stability. The answer depended heavily on the workload.

Throughput without MIG

Four containers shared one GPU. Each measurement includes network request, analysis, and response time and is averaged over three runs.

Throughput with MIG

First, nvidia-smi shows whether a GPU supports MIG.

N/A means unsupported; Disabled means available but inactive. Enable it with sudo nvidia-smi -i {gpu_id} -mig 1.

Instance sizes are chosen from predefined profiles rather than arbitrary allocations. nvidia-smi mig -lgip lists them; Blackwell GPUs commonly divide evenly into $2^n$ instances.

For four instances I used profile 14:

sudo nvidia-smi mig -i {gpu_id} -cgi 14,14,14,14 -C

nvidia-smi -L exposes each instance UUID for container assignment. The resulting throughput was:

MIG May Not Be the Best Choice

MIG improves isolation and may reduce memory, NVML, and CUDA interference between services. But when workloads overlap only modestly, hard partitioning can lower speed. A single non-MIG request improved from $\times7.26$ on the legacy server to $\times15.19$, whereas one MIG instance reached only $\times6.99$. A difference that looks small compounds to roughly 53 hours over 10,000 hours of data and nearly three weeks over 100,000 hours.

Our service had two APIs: a long-running batch workload and a short-response endpoint. Because they rarely contended directly, leaving the GPU unpartitioned was more effective.

Conclusion

AI efficiency is a systems problem, not merely a model problem. FlashAttention and functional model separation halved memory without degrading quality. Loading order and preprocessing, though easy to overlook, also had major effects. CUDA visibility failures and decoder hangs showed how tightly NVML, Docker runtime, cgroups, and drivers are coupled.

MIG demonstrated that more isolation is not automatically better. It improves stability, but real throughput can fall depending on CPU/GPU preprocessing ratios, API behavior, and the mix of batch and interactive requests. Infrastructure has no universal answer; it must fit the service, data, model, and operations.

AI optimization is end-to-end systems engineering across the model, code, runtime environment, and GPU configuration.