PyTorch

PyTorch is an open-source machine learning framework developed by Meta. It provides tensor computation with GPU acceleration and automatic differentiation, making it the primary framework for training and running many local AI models. Operators encounter PyTorch when loading models in Hugging Face Transformers or when using tools like vLLM that depend on PyTorch for model inference. PyTorch's dynamic computation graph allows flexible model building, but its memory usage can be high—a typical 7B model in FP16 requires ~14 GB of VRAM, often necessitating quantization for consumer GPUs.

PyTorch is the backbone of most modern local AI workflows. It manages tensors (multi-dimensional arrays) on CPU or GPU and records operations for automatic gradient computation, which is essential for training. For inference, PyTorch executes the model's forward pass. Operators often use PyTorch indirectly: Hugging Face Transformers loads models into PyTorch by default, and vLLM uses PyTorch for its inference engine. PyTorch's memory footprint is a key concern—model weights, activations, and optimizer states all consume VRAM. Quantization (e.g., using bitsandbytes or GPTQ) reduces memory by converting weights to lower precision (INT4/INT8), trading off some accuracy for speed and fit. PyTorch also supports torch.compile, which can optimize inference speed by fusing operations, though it may increase compilation time.

A rig with an RTX 3090 (24 GB VRAM) can run Llama 3.1 8B in FP16 (16 GB) with a 2K context. But running the same model in FP32 (32 GB) would exceed VRAM. Using PyTorch's built-in quantization (torch.quantization) or loading a 4-bit quantized version via bitsandbytes reduces memory to ~5 GB, leaving room for larger contexts or other tasks.

When you run from transformers import AutoModelForCausalLM; model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", torch_dtype=torch.float16, device_map="auto"), Hugging Face loads the model into PyTorch tensors. The device_map="auto" tells PyTorch to split layers across GPU and CPU if VRAM is insufficient. You'll see PyTorch allocate memory and, if VRAM is tight, offload some layers to system RAM—slowing inference from ~40 tok/s to ~5 tok/s.