Devices, Batching and Precision in Pipelines

Making pipelines fast enough

Older releases always ran a pipeline on the CPU in 32-bit floats, and defaults have shifted between versions, so set device and precision explicitly. Three arguments control this. device chooses where the model runs: device=0 or device="cuda" for the first GPU, "mps" on Apple silicon. For big models, device_map="auto" (which needs the accelerate package) spreads layers across available GPUs and CPU memory automatically. dtype sets precision: torch.float16 or torch.bfloat16 halves memory and usually speeds up GPU inference with little quality loss, and "auto" uses the dtype stored in the checkpoint; older releases called this argument torch_dtype. Finally, pass a list of inputs (or a dataset) with batch_size so several examples share one forward pass. Batching helps on GPUs, but it is not free: inputs are padded to the longest example in the batch, so mixing very short and very long texts wastes compute. Measure with your own data before choosing a batch size.

GPU, half precision and batching

Pass a list and a batch size; the pipeline groups inputs into batches internally.

import torch
from transformers import pipeline

clf = pipeline(
    "text-classification",
    model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
    device=0 if torch.cuda.is_available() else -1,  # -1 means CPU
    dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)

reviews = ["Great phone"] * 64 + ["Battery died in a day"] * 64
preds = clf(reviews, batch_size=32)
print(len(preds), preds[0]["label"], preds[-1]["label"])

float16 on CPU is a trap

Half precision is meant for GPUs. On most CPUs float16 matrix maths is slow or unsupported, so keep float32 (or try bfloat16 on CPUs that support it) when no GPU is present.

त्वरित जाँच: What does `device_map="auto"` do when loading a large model?

• Places model layers across the available GPUs and CPU memory automatically
• Converts the model to 4-bit precision
• Downloads the smallest variant of the model
• Runs the model on the Hub's servers instead of locally