Edward Z. Yang ezyang

TSAN verification of the CachingHostAllocator snapshot data races

This report documents a ThreadSanitizer (TSAN) reproduction of two data races in aten/src/ATen/core/CachingHostAllocator.h that were introduced/exposed by the "Support memory snapshot for CPU pinned memory" change, and confirms the committed fixes eliminate them.

Background

The host (pinned-memory) caching allocator uses fine-grained per-block locking

Ref: https://x.com/ezyang/status/2048485559576789083

I think one way to think about fine-grained precision APIs is that we are exposing a little about the underlying memory hierarchy to the user. For the single node ops, the most important thing is memory or not. For collectives, what the comms are actually done in is another dimension.

PyTorch generally has these rules:

If hardware is involved (e.g., tensor cores), defer to the hardware
Always do accumulation in fp32 (this is formalized as acc_dtype, which is not exposed to users but is hard coded per dtype)
- We don't always do accumulation in fp32 for matmuls. This is generally controlled by global knobs: torch.backends.cuda.matmul.fp32_precision, torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp
Output dtype defaults to the same as the input. out_dtype to control how to convert the accumulator

	#!/usr/bin/env bash
	# Claude Code status line - shows context window remaining percentage

	input=$(cat)
	remaining=$(echo "$input" \| jq -r '.context_window.remaining_percentage // empty')

	if [ -n "$remaining" ]; then
	printf 'Context: %s%% remaining' "$remaining"
	fi

	import torch
	import torch.nn.functional as F


	def topk_softmax(logits, k):
	vals, idx = torch.topk(logits, k, dim=-1)
	return idx, F.softmax(vals, dim=-1)


	def full_softmax_then_renorm(logits, k):

	#!/usr/bin/env bash
	# Claude Code status line - shows context window remaining percentage

	input=$(cat)
	remaining=$(echo "$input" \| jq -r '.context_window.remaining_percentage // empty')

	if [ -n "$remaining" ]; then
	printf 'Context: %s%% remaining' "$remaining"
	fi

	https://github.com/pytorch/pytorch/issues/163449

	https://github.com/pytorch/pytorch/issues/163457

	https://github.com/pytorch/pytorch/issues/163420

	https://github.com/pytorch/pytorch/issues/163300

	https://github.com/pytorch/pytorch/issues/162723

	import torch
	import unittest
	from torch import Tensor
	from torch.distributed.tensor import (
	DTensor,
	DeviceMesh,
	distribute_tensor,
	init_device_mesh,
	Partial,
	Replicate,

	x = DTensor.from_local(arange_nd(15), mesh["m", "n", "k"], [R, R, R])
	# Eliminate M
	x = DTensor.from_local(x.redistribute(placements=[R, R, S(0)]).to_local(), mesh["m", "n"]) # shard K
	x = DTensor.from_local(x.redistribute(placements=[R, S(0)]).to_local(), mesh["m"]) # shard N
	x = x.redistribute(placements=[S(0)]).to_local() # shard M
	x = DTensor.from_local(x, mesh["n"], [S(0)]).redistribute(placements=[R]) # unshard N
	x = DTensor.from_local(x.to_local(), mesh["n", "k"], [R, S(0)]).redistribute(placements=[R, R]) # unshard K
	# Eliminate N
	x = DTensor.from_local(x.redistribute(placements=[R, S(0)]).to_local(), mesh["n"]) # shard K
	x = x.redistribute(placements=[S(0)]).to_local() # shard N


	def render(tensor, cell_width=None):
	"""
	Print a tensor following the matrix-of-matrices algorithm.

	Args:
	tensor: A tensor-like object with .shape attribute and indexing
	cell_width: Width for each cell (calculated globally if None)

	Returns: