ML//PyTorch

PyTorch is a Python library for numerical computing on tensors (multidimensional arrays) with automatic differentiation, and it is the framework in which most neural networks, transformers included, are written, trained and studied. A model is ordinary Python code that multiplies tensors; PyTorch records the operations, computes the gradients for gradient descent and runs every operation on whichever device holds the tensors.

PyTorch is a Python library for numerical computing on tensors (multidimensional arrays) with automatic differentiation, and it is the framework in which most neural networks, transformers included, are written, trained and studied. A model is ordinary Python code that multiplies tensors; PyTorch records the operations, computes the gradients for gradient descent and runs every operation on whichever device holds the tensors.

The device is a property of each tensor, chosen with .to("cpu"), .to("cuda") or .to("mps") (Apple), and the code of the model does not change between them. What changes underneath is the backend that executes each operation: optimised CPU kernels on the processor, CUDA libraries on an NVIDIA GPU, ROCm on a supported AMD GPU.

On a CPU everything works, only slower: the same transformer code trains a small model such as GPT-2 small on a laptop, and the arithmetic is identical. The limits are the processor's compute and the memory bandwidth, not the correctness.

Much of the surrounding ecosystem assumes NVIDIA. Libraries that ship their own GPU kernels (quantised training, fast serving engines) are often written for CUDA first, so on another vendor's hardware PyTorch itself may run while those extensions do not.