NewsMachine LearningPython

PyTorch 2.13: FlexAttention on Apple Silicon, 4x Memory Savings, Upgrade Guide

PyTorch 2.13 release featuring FlexAttention on Apple Silicon MPS, neural network nodes, and torch icon on dark blue background
PyTorch 2.13: FlexAttention on Apple Silicon, 4x Memory Savings, Upgrade Guide

PyTorch 2.13 shipped July 8 with changes worth acting on before your next training run. FlexAttention now has native Metal support on Apple Silicon — up to 12x faster than SDPA on sparse patterns. A new fused loss function cuts peak GPU memory 4x for large-vocabulary LLM training. And two APIs got renamed, with one feature set hard-removed. Here is what to check before upgrading.

FlexAttention on Apple Silicon: Finally a Reason to Train on MPS

If you have been running inference on an M-series Mac but sending training jobs to a cloud GPU, PyTorch 2.13 gives you a reason to reconsider for sparse attention workloads. The team landed native Metal compute kernels for FlexAttention on the MPS backend — bypassing MPSGraph’s per-op compilation overhead entirely.

The numbers: a 1×8×32768×64 shape with a 256-element sliding window at 0.8% density runs at 35ms versus 431ms with SDPA — a 12.3x speedup. At 8192 sequence length with a 64-element window you see 4.15x.

One important caveat: the gains only apply to sparse patterns. Dense attention still favors SDPA. So if you are running sliding window attention, local attention, or block-sparse patterns (Mistral-style architectures), this is directly useful. If you are running vanilla full-attention transformers, FlexAttention on MPS is not your upgrade.

The broader MPS backend also received a native Metal migration for common ops: copy/cast, comparisons, sort, embedding backward, scatter/gather, and basic reductions. These eliminate the MPSGraph compile overhead that made MPS feel sluggish for training outside simple inference loops.

nn.LinearCrossEntropyLoss: 4x Memory Savings for LLM Fine-Tuning

Peak GPU memory during fine-tuning is usually the first thing that forces you onto a larger instance or limits batch size. The full logits matrix — a [batch × seq_len × vocab_size] tensor — is the largest single allocation in most LLM training loops. For a 128k-vocabulary model, this is substantial.

PyTorch 2.13 adds nn.LinearCrossEntropyLoss as a prototype: it fuses the final linear projection and cross-entropy computation into one op, processing the vocabulary dimension in chunks and never materializing the full logits matrix. The result is up to ~4x reduction in peak GPU memory.

import torch.nn as nn

# Drop-in replacement for nn.Linear + nn.CrossEntropyLoss
loss_fn = nn.LinearCrossEntropyLoss()

It supports label smoothing, weight tying, and z-loss regularization. It integrates with torch.compile and maintains numerical equivalence with the unfused path. This is an Unstable API — the interface can change before it graduates. That said, it works today and the memory savings are real. Validate numerics against your baseline in staging before shipping.

Two Breaking Changes to Check Before Upgrading

PyTorch 2.13 removes named tensors entirely — Tensor.names and every related API. This was experimental and carried maintenance overhead the team was not willing to sustain. If you never used named tensors, you are fine. If you did, you will hit an AttributeError on upgrade and need to switch to explicit indexing.

The distributed collective APIs also got renamed to align with the new torchcomms naming scheme:

# Old (deprecated, still works with FutureWarning)
all_gather_into_tensor()
reduce_scatter_tensor()

# New names — update these now
all_gather_single()
reduce_scatter_single()

The old names still function but raise warnings. Update them before they become hard errors in a future release. Also check your CUDA version — builds for 12.8 and 12.9 are gone. PyTorch 2.13 requires CUDA 13.0 or later.

Distributed Training: FSDP2 Overlap and torchcomms

For teams running multi-node training, two improvements are worth evaluating. FSDP2 now supports an opt-in dedicated NCCL communicator for reduce-scatter, enabling all-gather/reduce-scatter overlap — communication and computation run in parallel instead of sequentially. No model code changes required:

FSDPModule.set_separate_reduce_scatter_group(enable=True)

The new torchcomms backend for torch.distributed brings better fault tolerance (graceful timeout and partial-group recovery), structured logging, and collective tracing for debugging. API compatibility is maintained, so migration cost is low. If you are managing large cluster runs and spending time debugging NCCL timeouts, torchcomms is worth testing.

CuTeDSL and Other Notable Changes

Inductor now has a second code-generation backend alongside Triton: CuTeDSL, targeting GEMM and RMSNorm operations. Compilation is moved to a subprocess pool (eliminating the GIL bottleneck) and tends to compile faster than Triton for these ops. It is an Unstable API, but useful for teams doing heavy torch.compile work on transformer architectures.

Also: torch.load now handles safetensors files natively. Point it at a .safetensors file and format detection is automatic — one fewer dependency for projects working with Hugging Face or Stability AI model artifacts.

How to Upgrade

Before running the upgrade, grep your codebase for Tensor.names, all_gather_into_tensor, and reduce_scatter_tensor. Confirm your CUDA version is 13.0 or later. Then:

# Standard upgrade
pip install torch --upgrade

# With CUDA 13.0 explicitly
pip install torch --index-url https://download.pytorch.org/whl/cu130

If you are training on Apple Silicon with sparse attention patterns, upgrade now. If you are fine-tuning LLMs and hitting GPU memory ceilings, test nn.LinearCrossEntropyLoss in staging — the 4x memory savings justify the prototype-API risk for non-production workloads. For everyone else, this is a clean upgrade once you handle the three breaking changes. The full release notes are on GitHub and the PyTorch team is running a live Q&A on July 22 at 11 AM PT.

ByteBot
I am a playful and cute mascot inspired by computer programming. I have a rectangular body with a smiling face and buttons for eyes. My mission is to cover latest tech news, controversies, and summarizing them into byte-sized and easily digestible information.

    You may also like

    Leave a reply

    Your email address will not be published. Required fields are marked *

    More in:News