Post-Training Memory Reduction Techniques for Model Inference

When a model leaves the training phase, the engineering priorities change completely. You are no longer optimizing for gradient flow or convergence. You are optimizing for the VRAM budget on a target GPU, the latency contract of a serving system, and the per-request cost on a cloud bill. A 13B-parameter model in FP32 occupies around 52 GB of weights alone, which does not fit on a single A100 40 GB. The same model in INT4 occupies roughly 7 GB, which fits comfortably on a consumer-grade RTX 4090 with headroom for the KV cache. The techniques in this post are post-training only: they assume you already have a trained checkpoint and want to shrink its memory footprint at inference time, with minimal engineering effort and no retraining.

1. Half-Precision (FP16 / BF16)

The simplest and highest-impact change is to cast the model weights from FP32 to a 16-bit floating point format. This halves the memory footprint of both the weights and the activations, and on modern accelerators (Ampere, Hopper, Ada), it also unlocks Tensor Core kernels, which are 2× to 8× faster than FP32 paths.

The two main options are FP16 (5-bit exponent, 10-bit mantissa) and BF16 (8-bit exponent, 7-bit mantissa). BF16 has the same dynamic range as FP32, which makes it considerably more robust against overflow and underflow in attention logits, softmax denominators, and layer norm statistics. For inference of transformer-based models, BF16 should be your default. FP16 remains useful on older hardware (Volta, Turing) that does not support BF16 natively.

import torch
from transformers import AutoModelForCausalLM

# Option A: load directly in BF16 (preferred for transformers)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

# Option B: cast an existing FP32 model
model = model.to(dtype=torch.bfloat16, device="cuda")
model.eval()

# Inference
with torch.inference_mode():
    output = model.generate(input_ids, max_new_tokens=128)

If you cannot cast the full model, for example, because a specific layer is numerically sensitive, you can apply mixed precision only inside an autocast region, while keeping the master weights in FP32:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    with torch.inference_mode():
        output = model(input_tensor)

The Trade-off with BF16, the accuracy degradation is usually negligible (well below 0.1% on most benchmarks). With FP16, you may observe instabilities in attention softmax or in models trained originally in BF16, particularly large language models, where logits can exceed the FP16 representable range and produce NaNs. Always validate end-to-end before shipping.

2. Post-Training Quantization to INT8

The next step down is INT8 quantization, which represents weights (and optionally activations) with 8-bit integers instead of 16-bit floats. This gives a further 2× reduction over FP16 and a 4× reduction over FP32. There are two flavours worth knowing.

Dynamic quantization quantizes only the weights ahead of time, and quantizes activations on the fly at every forward pass. It requires no calibration data and applies most naturally to linear and LSTM layers. It is the cheapest path to deploy.

import torch
import torch.nn as nn

model_fp32 = MyModel()
model_fp32.load_state_dict(torch.load("checkpoint.pt"))
model_fp32.eval()

model_int8 = torch.quantization.quantize_dynamic(
    model_fp32,
    {nn.Linear, nn.LSTM},  # which layer types to quantize
    dtype=torch.qint8,
)

# Same inference API
with torch.inference_mode():
    output = model_int8(input_tensor)

For modern transformer models, the practical default has shifted tobitsandbytes, which integrates cleanly with Hugging Face and applies LLM.int8() a mixed decomposition that keeps outlier activation channels in FP16 and quantizes the rest to INT8. This preserves accuracy much better than naive INT8 on large models.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,  # outlier detection threshold
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    quantization_config=bnb_config,
    device_map="auto",
)

Trade-off. Expect a 0.5–2% accuracy drop on downstream tasks for naive INT8, and near-zero drop for LLM.int8(). The compute path also matters: INT8 GEMM kernels are fast on Tensor Cores from Turing onward, but on CPUs and older GPUs, the dequantization overhead can erase the latency benefit, even if the memory benefit remains.

3. Weight-Only Quantization to INT4 (NF4 / GPTQ / AWQ)

For large language models, the dominant memory cost is the weight matrix itself, not the activations. Weight-only quantization exploits this by storing weights at very low precision (typically 4 bits) while keeping the compute path in BF16 or FP16. The weights are dequantized on the fly inside the matmul kernel. The reduction is substantial: a 70B model that requires 140 GB in BF16 can be reduced to roughly 35 GB in INT4.

The two paths most commonly used today are NF4 (via bitsandbytes) and AWQ / GPTQ (which apply a more careful calibration to minimize quantization error on important weights).

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # normalized float 4
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,      # quantize the quantization constants too
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-70B",
    quantization_config=bnb_config,
    device_map="auto",
)

For higher throughput at inference time, AWQ is usually preferable, because it produces a kernel-friendly layout that fuses dequantization directly into the GEMM:

# Loading a pre-quantized AWQ checkpoint
from transformers import AutoModelForCausalLM, AwqConfig

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-3.1-8B-AWQ",
    device_map="cuda",
)

Trade-off. Accuracy degradation is typically 1–3% on language modelling perplexity, and noticeably larger on reasoning-heavy tasks. Calibration-based methods (AWQ, GPTQ) are more accurate than data-free methods (NF4), but they require a small calibration dataset and a one-time offline quantization step. Also, be aware that not every kernel supports every quantization scheme on every GPU architecture. Check compatibility before committing.

4. Graph Compilation and Runtime Export (ONNX Runtime / TensorRT)

The previous techniques reduce the size of the tensors. This one reduces the overhead of operator dispatch, intermediate tensor allocations, and layout conversions by exporting the model to an optimized inference runtime. ONNX Runtime is the portable option; TensorRT is the NVIDIA-specific option, which usually delivers the highest throughput and lowest latency on NVIDIA hardware.

First, export the model to ONNX:

import torch

model.eval()
dummy_input = torch.randn(1, 3, 224, 224, device="cuda")

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch"},
        "output": {0: "batch"},
    },
    opset_version=17,
    do_constant_folding=True,
)

Then run with ONNX Runtime, which performs graph-level fusions (e.g. fused attention, fused LayerNorm) and removes the Python dispatch overhead:

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

inputs = {"input": np.random.randn(1, 3, 224, 224).astype(np.float32)}
outputs = session.run(None, inputs)

For maximum performance on NVIDIA GPUs, convert the ONNX graph further to a TensorRT engine, which performs kernel autotuning and precision calibration at build time:

trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --fp16 \
        --memPoolSize=workspace:4096 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:32x3x224x224

Note that these runtimes compose with the techniques above: you can export an INT8-quantized graph to ONNX, or build a TensorRT engine with INT8 calibration and a representative calibration dataset, and combine all the savings.

Trade-off: Export increases engineering complexity. Dynamic shapes, control flow, and custom ops often do not export cleanly, so you may need to refactor the forward pass to make it export-friendly. TensorRT engines are also hardware-specific — an engine built on an A100 will not run optimally (or sometimes at all) on a different GPU architecture, so you have to maintain a build matrix.

Applying These Techniques to a Custom PyTorch Model

The examples above used Hugging Face checkpoints because they are convenient, but every technique generalizes to a hand-written one nn.Module. The friction points are different for each one, so let me go through them in turn.

Casting Custom Modules to BF16 / FP16

model.to(dtype=torch.bfloat16) recursively casts every parameter and buffer registered through nn.Module. The places where this fails are almost always the same: hardcoded dtypes insideforward, tensors created on the fly with the default FP32 dtype, and numerically sensitive blocks that you want to keep in FP32.

import torch
import torch.nn as nn

class MyBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.proj = nn.Linear(dim, dim)
        self.norm = nn.LayerNorm(dim)
        # Register constants as buffers so .to(dtype=...) catches them
        self.register_buffer("scale", torch.tensor(dim ** -0.5))

    def forward(self, x):
        # BAD: forces FP32 regardless of model dtype
        # mask = torch.zeros(x.shape[0], device=x.device)
        # GOOD: inherit dtype from the input
        mask = torch.zeros(x.shape[0], device=x.device, dtype=x.dtype)

        x = self.proj(x) * self.scale

        # Keep LayerNorm in FP32 for stability, then cast back
        x_fp32 = x.float()
        x_fp32 = self.norm(x_fp32)
        return x_fp32.to(x.dtype)

model = MyBlock(1024).to(device="cuda", dtype=torch.bfloat16)

For diffusion-style models with sensitive paths (timestep embeddings, attention softmax denominators, classifier-free guidance combinations), the common pattern is to keep the trunk in BF16 but wrap the sensitive blocks with an explicit .float() cast and a .to(orig_dtype) at the end. This is the same trick that diffusers is used internally for the UNet and the noise scheduler.

Replacing Linear Layers with bitsandbytes Equivalents

For any custom architecture where the dominant cost is linear projections, which includes all transformer and DiT variants, essentially, you can apply weight-only INT8 or INT4 by simply swapping nn.Linear for the bitsandbytes equivalent. No graph rewriting, no calibration step, no Hugging Face dependency.

import torch
import torch.nn as nn
import bitsandbytes as bnb

class MyTransformerBlock(nn.Module):
    def __init__(self, dim, heads, use_4bit=False):
        super().__init__()
        Linear = bnb.nn.Linear4bit if use_4bit else nn.Linear

        if use_4bit:
            self.qkv = bnb.nn.Linear4bit(
                dim, 3 * dim,
                bias=False,
                quant_type="nf4",
                compute_dtype=torch.bfloat16,
            )
            self.out = bnb.nn.Linear4bit(
                dim, dim,
                quant_type="nf4",
                compute_dtype=torch.bfloat16,
            )
        else:
            self.qkv = nn.Linear(dim, 3 * dim, bias=False)
            self.out = nn.Linear(dim, dim)

        self.norm = nn.LayerNorm(dim)

To convert an already-trained model without rewriting its definition, walk the module tree and substitute layers in place:

def replace_linear_with_4bit(module, skip_names=("lm_head",)):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear) and name not in skip_names:
            new_layer = bnb.nn.Linear4bit(
                child.in_features,
                child.out_features,
                bias=child.bias is not None,
                quant_type="nf4",
                compute_dtype=torch.bfloat16,
            )
            # Move quantized weights over
            new_layer.weight = bnb.nn.Params4bit(
                child.weight.data, requires_grad=False, quant_type="nf4"
            ).cuda()
            if child.bias is not None:
                new_layer.bias = nn.Parameter(child.bias.data)
            setattr(module, name, new_layer)
        else:
            replace_linear_with_4bit(child, skip_names)

replace_linear_with_4bit(model)
model = model.cuda()

Note the skip_names argument — for generative models, the final output projection (lm_head, the noise prediction head, etc.) is often worth leaving in BF16, because quantization error there directly distorts the output distribution.

Native PyTorch Quantization with torch.ao.quantization

If you cannot use bitsandbytes (for example, when targeting CPU or non-NVIDIA hardware), The native FX graph-mode quantization API is the right tool. It traces the model into a graph, inserts observers, runs calibration, and then converts observed modules into quantized counterparts.

import torch
from torch.ao.quantization import quantize_fx, get_default_qconfig_mapping

model.eval()
qconfig_mapping = get_default_qconfig_mapping("fbgemm")   # use "qnnpack" on ARM

# Example input that matches the actual forward signature
example_inputs = (torch.randn(1, 3, 224, 224),)

prepared = quantize_fx.prepare_fx(model, qconfig_mapping, example_inputs)

# Calibrate with representative data (a few hundred batches is usually enough)
with torch.inference_mode():
    for batch in calibration_loader:
        prepared(batch)

quantized = quantize_fx.convert_fx(prepared)

# Save the quantized model
torch.save(quantized.state_dict(), "model_int8.pt")

The most common failure mode is that FX tracing cannot handle your forward. Tracing supports static control flow but breaks on data-dependent branches, Python list mutations, and dictionary outputs. The workaround is either to restructure the forward to be trace-friendly, or to mark non-traceable submodules as leaf modules so FX treats them as opaque:

from torch.fx import wrap

# Make a free function a leaf so FX does not trace inside it
@wrap
def custom_op(x, scale):
    return x * scale + scale.sin()

Exporting a Custom Model to ONNX

Custom modules export cleanly as long as every operator used inside is forward has an ONNX symbolic implementation. The new torch.onnx.export(..., dynamo=True) path (built on torch.export) handles more dynamic shapes and Python control flow than the legacy tracer.

import torch

class MyModel(nn.Module):
    def forward(self, x, mask):
        # data-dependent shapes — handled by dynamo export
        return (x * mask.unsqueeze(-1)).sum(dim=1)

model = MyModel().eval()
x = torch.randn(2, 16, 64)
mask = torch.ones(2, 16)

torch.onnx.export(
    model,
    (x, mask),
    "custom.onnx",
    input_names=["x", "mask"],
    output_names=["out"],
    dynamic_axes={
        "x":    {0: "batch", 1: "seq"},
        "mask": {0: "batch", 1: "seq"},
    },
    opset_version=17,
    dynamo=True,           # use the newer exporter
)

When an operator is missing — typical for attention variants, custom CUDA kernels, or 3D-rendering operations — you have three options: rewrite the op in terms of ONNX-supported primitives, register a symbolic function for the existing op, or implement a custom ONNX Runtime operator. For research code that uses unusual ops (for example, the rasterization step in Gaussian splatting), the rewrite path is usually the most pragmatic, even if it produces a less efficient graph.

from torch.onnx import register_custom_op_symbolic

def my_symbolic(g, x, alpha):
    # Express the op in terms of standard ONNX nodes
    return g.op("Mul", x, g.op("Constant", value_t=torch.tensor(alpha)))

register_custom_op_symbolic("mylib::my_custom_op", my_symbolic, opset_version=17)

torch.compile for Memory-Efficient Inference

torch.compile is mainly a latency tool, but it also reduces peak memory in inference by fusing element-wise operations and eliminating intermediate allocations. It works on any custom nn.Module without modification, which makes it the lowest-effort technique on this list.

model = MyModel().to("cuda", dtype=torch.bfloat16).eval()
model = torch.compile(model, mode="reduce-overhead", fullgraph=False)

with torch.inference_mode():
    output = model(input_tensor)

The mode="reduce-overhead" setting uses CUDA graphs under the hood, which removes per-kernel-launch overhead but requires static input shapes. For variable-shape inputs (variable sequence lengths, variable batch sizes), use the default mode and accept slightly higher overhead. Setting fullgraph=True forces the compiler to fail loudly on a graph break, which is useful during development to make sure no Python-level fallback path is silently slowing things down.

A Practical Order of Operations

For a custom model, the order I would recommend is:

1. Move to BF16 and fix any hardcoded FP32 paths inside the forward. Validate accuracy.
2. Wrap the model with torch.compile and confirm there are no graph breaks on the hot path.
3. If the model is weight-heavy, replace nn.Linear instances with bnb.nn.Linear4bit (or Linear8bitLt) using the in-place tree walk above. Re-validate accuracy.
4. If targeting CPU or non-NVIDIA hardware, apply FX graph-mode INT8 quantization instead of bitsandbytes.
5. Only after the model is functionally stable, export to ONNX or TensorRT for deployment. The export step is the most brittle, so it should be the last thing you change.

Each of these steps is independently revertible, which matters when you are debugging an accuracy regression in production, you want to know exactly which transformation is responsible for the drop.

Putting It Together

In practice, the techniques compose, and a reasonable production pipeline looks like this:

1. Cast the model to BF16 as the baseline, almost free, with almost no accuracy cost.
2. Apply weight-only INT4 (AWQ or NF4) if the model is large and weight-bound — this is the single largest lever for modern LLMs.
3. Apply INT8 dynamic quantization or LLM.int8() for medium-sized models where weight-only INT4 is too aggressive.
4. Export to ONNX Runtime or TensorRT once the model is functionally validated, to remove framework overhead and unlock kernel-level fusions.