PyTorch 2.13 Ships 12x Faster Attention and 4x Memory Savings for LLM Training

PyTorch 2.13 is out, and it's a release that pushes hard in three directions at once: hardware breadth (Apple Silicon, Arm, AMD, Intel), memory efficiency for large-vocabulary LLM training, and distributed training at scale. The release is composed of 3,328 commits from 526 contributors since PyTorch 2.12. That's a lot of surface area, so here's what actually matters.

FlexAttention finally lands on Apple Silicon

FlexAttention is PyTorch's unified API for expressing custom attention patterns -- things like sliding window attention, document masking, or ALiBi -- as plain Python functions that get compiled into fused kernels. Before 2.13, you needed CUDA to use it. Now it runs on Metal/MPS (Apple's GPU compute stack).

The MPS implementation provides hand-written Metal kernels for both the sparse prefill and decode paths, including GQA and captured buffers. Instead of writing a custom CUDA kernel for every attention variant, FlexAttention lets you write a two-line Python function and the compiler builds a fast kernel automatically.

The benchmark numbers for sparse patterns are striking. On long, sparse attention patterns, the speedups over SDPA are substantial -- on a 1x8x32768x64 shape with a 256-element sliding window (0.8% density), FlexAttention runs in ~35 ms vs ~431 ms for SDPA, a ~12.3x speedup; a smaller 8192-length / 64-window case achieves ~4.15x. Dense patterns still favor SDPA, as expected.

There's also a correctness fix for CUDA users: by default, the FlexAttention flash backend uses atomic operations in the backward pass for dQ accumulation, which introduces non-determinism -- repeated runs on the same input can produce slightly different gradients. The new deterministic backward path replaces atomics with a pre-computed write ordering that guarantees bit-for-bit reproducible gradients. The measured end-to-end overhead is well under 1% at longer sequence lengths -- just +0.2% at S=32768 -- making determinism effectively free for most production workloads. You opt in via the existing torch.use_deterministic_algorithms(True) flag.

The memory wall for large-vocab LLMs just got shorter

Anyone who has trained a language model with a large vocabulary (think 100K+ tokens) knows the pain: the final linear projection from hidden states to vocabulary logits produces a massive matrix that sits in GPU memory just long enough to compute the loss, then gets thrown away. For a model with a 128K-token vocabulary and a large batch, this can easily cost tens of gigabytes.

nn.LinearCrossEntropyLoss fuses the final linear projection and cross-entropy computation into a single module that processes the vocabulary dimension in chunks, never materializing the full logits matrix. This reduces peak memory by up to ~4x for large-vocabulary workloads while maintaining numerical equivalence with the unfused path.

The implementation supports label smoothing, weight tying, and z-loss regularization out of the box, and integrates with torch.compile for further optimization. As a drop-in replacement for separate nn.Linear + nn.CrossEntropyLoss, adoption requires no other code changes.

# Before
logits = self.lm_head(hidden_states)  # [B, T, vocab_size] -- huge!
loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
# After (2.13+)
from torch.nn import LinearCrossEntropyLoss
loss_fn = LinearCrossEntropyLoss(weight=self.lm_head.weight)
loss = loss_fn(hidden_states.view(-1, hidden_dim), targets.view(-1))

Distributed training gets a new backbone

Two major changes land for distributed training, both aimed at large-cluster workloads where the existing infrastructure starts to show cracks.

First, torchcomms is a new communications backend integrated into PyTorch Distributed's CI and device-mesh paths, providing improved fault tolerance (graceful timeout and partial-group recovery), better scalability across large clusters, and richer debuggability through structured logging and collective tracing. It serves as a modern alternative to the existing c10d backends while maintaining API compatibility.

Second, FSDP2 -- PyTorch's fully-sharded data-parallel training strategy, which shards model parameters across GPUs to reduce per-device memory -- gets a meaningful throughput improvement. In FSDP training, all-gather and reduce-scatter share a single NCCL communicator by default. Because NCCL serializes operations on the same communicator, these two collectives cannot overlap, leaving communication bandwidth underutilized. FSDPModule.set_separate_reduce_scatter_group(enable=True) gives reduce-scatter its own dedicated NCCL communicator, allowing it to progress concurrently with all-gather operations.

CuTeDSL: a second code-generation path for Inductor

PyTorch's Inductor compiler (the backend behind torch.compile) has always used Triton as its primary code-generation target for GPU kernels. 2.13 adds a second path: CuTeDSL.

CuTeDSL is a Python-native domain-specific language built on NVIDIA's CuTe (CUDA Templates) library. CuTeDSL enables authoring in Python what used to require CUTLASS C++, which makes JIT-style workflows more practical for FlexAttention. Inductor can now use CuTeDSL as an alternative code-generation backend alongside Triton -- specifically for matrix multiplication (GEMM) and normalization (RMSNorm), two of the most performance-critical operations in transformer training.

Kernel compilation has also moved from the thread pool to a subprocess pool, eliminating Python's GIL bottleneck and improving compile-time parallelism. This matters in practice: long compile times are one of the biggest friction points with torch.compile.

Everything else worth knowing

• Native safetensors loading: torch.load("foo.safetensors") now works natively, detecting the format automatically and returning tensors directly. The format is based on a simple JSON header and raw data buffers with no executable content -- loading a safetensors file cannot run arbitrary code on your machine.
• Python 3.15 wheels: PyTorch wheels now support Python 3.15, including the experimental free-threaded 3.15t build. Note: torch.compile is not yet supported on 3.15, and these wheels are only on Linux for now.
• Arm Armv9-A support: torch.compile on AArch64 now recognizes Armv9-A CPUs (e.g., Neoverse V2 used in AWS Graviton4), propagating the correct target triple and feature set through Inductor codegen.
• Intel XPU telemetry: New query APIs for Intel GPUs expose runtime device state: memory usage, utilization, power draw, clock rate, and temperature.
• ExecuTorch joins PyTorch Core: This release marks ExecuTorch's integration into PyTorch Core, making on-device inference a first-class capability of the framework.
• Named tensors removed: The deprecated named-tensor feature (Tensor.names and associated APIs) has been hard-removed to cut overhead and code bloat.

The bigger picture

Throughout the 2.x series, PyTorch has been evolving from a research-first framework into a unified, hardware-agnostic platform for production training and inference at scale. 2.13 makes that trajectory concrete: the same FlexAttention API now runs on CUDA, ROCm, CPU, and Apple Silicon. The same torch.compile pipeline now targets NVIDIA, AMD, Intel, and Arm hardware. And the same distributed training stack now has the fault-tolerance and observability primitives that production clusters actually need.

The full release notes are on GitHub. You can install 2.13 today via pip install torch, and a live Q&A with the core team is scheduled for July 22 at 11 a.m. PT.