Fixing PyTorch Tensor Corruption Bug

Introduction

In the ever-evolving world of deep learning, PyTorch stands out as a powerful and flexible framework. Its ability to handle complex tensor operations is a cornerstone of modern AI development. However, even the most robust frameworks can encounter unexpected issues. One such issue, recently highlighted, involves the potential for tensor corruption when attempting to resize storage that cannot be resized, particularly when a tensor shares its storage with a non-resizable buffer, such as a NumPy array injected via set_(). This problem can lead to what's colloquially termed a "Zombie Tensor" – a tensor that appears to have a valid shape but is fundamentally broken, often resulting in crashes like segmentation faults or internal runtime errors when accessed. This article delves into the intricacies of this bug, explains why it happens, and discusses the implications for developers using PyTorch.

The "Zombie Tensor" Bug Explained

Let's dive deep into the core of the issue. When you're working with PyTorch tensors, you're essentially dealing with two key pieces of information: the tensor's shape (how it's organized in dimensions) and its storage (where the actual data resides in memory). Normally, these two are in sync. However, a bug has been identified where this synchronization breaks down, leading to a very problematic state. The bug occurs when you try to resize a tensor's storage using resize_() that is intrinsically not resizable. A common scenario for this is when a tensor is created from or shares its underlying storage with a NumPy array that has been embedded into PyTorch using set_(). NumPy arrays, by their nature, have fixed-size storage once created, and PyTorch respects this. When resize_() is called on such a tensor, PyTorch correctly identifies that the storage cannot be resized and raises a RuntimeError, specifically stating: Trying to resize storage that is not resizable. This is the expected and safe behavior.

However, the problem lies in the fact that the operation is not entirely exception-safe. Before the actual check for resizable storage fails, PyTorch has already gone ahead and updated the tensor's shape and stride metadata to reflect the new target size you requested. So, imagine you have a small, empty tensor derived from a non-resizable buffer, and you try to resize it to a large (5, 5, 5) shape. The RuntimeError is raised, but your tensor's shape attribute now incorrectly reports torch.Size([5, 5, 5]). Meanwhile, the actual underlying storage remains stubbornly at 0 bytes because it couldn't be resized. This creates a dangerous inconsistency. The tensor thinks it's large, but its storage is empty. This is what we mean by a "Zombie Tensor": it looks alive with a specific shape, but it's dead inside, lacking the data it claims to hold. Accessing such a tensor – for example, by trying to print it or perform any operation that requires reading its data – can lead to severe issues. The program might encounter a Segmentation Fault (a low-level error indicating an attempt to access memory that doesn't belong to the process) or another internal RuntimeError because the code is trying to read data from a non-existent memory location. This is a critical bug because it can manifest in unpredictable ways, potentially corrupting your workflow and leading to hard-to-debug crashes, especially in complex, long-running training loops where the source of the error might be buried deep within operations.

Illustrating the Corruption: A Minimal Reproduction

To truly understand the severity and nature of this bug, it's essential to see it in action. The PyTorch team has provided a minimal, reproducible example that clearly demonstrates the problem. Let's walk through it. The setup begins by creating a piece of non-resizable storage that is effectively 0 bytes in size. This is achieved by creating a NumPy array with no elements (np.array([], dtype=np.int32)) and then converting its underlying storage to a PyTorch untyped storage (untyped_storage()). This locked_storage is now a PyTorch object that represents memory which cannot be altered in size. Next, a fresh PyTorch tensor, t, is created, also empty and with the same integer data type (dtype=torch.int32). The crucial step here is t.set_(locked_storage). This command tells the tensor t to use the locked_storage as its data backing. At this point, t correctly reflects its empty state: its shape is torch.Size([0]) and its storage size is 0 bytes.

The problem arises when we attempt to resize this tensor. The code then calls t.resize_((5, 5, 5)). Based on our understanding, this operation should fail gracefully. The expectation is that because locked_storage is not resizable, the resize_() method should raise a RuntimeError, and importantly, the tensor's metadata (its shape and stride) should remain unchanged, sticking to the original torch.Size([0]). This is what's known as a strong exception guarantee: if an operation fails, the object should be left in a state as if the operation never occurred.

However, the actual behavior is different and problematic. The try...except RuntimeError: pass block demonstrates this. The resize_((5, 5, 5)) call does trigger a RuntimeError as expected because the storage is locked. But, critically, before the error is raised, the tensor's internal metadata is updated to the new target shape of (5, 5, 5). So, when the exception is caught and execution continues, the tensor t is left in this corrupted state. If you were to then print(f"Shape: {t.shape}"), it would output Shape: torch.Size([5, 5, 5]). Simultaneously, print(f"Storage: {t.untyped_storage().nbytes()}") would correctly report Storage: 0, because the underlying storage was never actually resized. The stark mismatch between the reported shape (5x5x5, which implies 125 elements) and the actual available storage (0 bytes) is the root cause of the subsequent issues. The final line, print(t), is where the crash often occurs. Trying to display a tensor that claims to have data but has no memory allocated for it will inevitably lead to a crash, either a Segmentation Fault or another RuntimeError, depending on the specific context and how the memory access is handled by the system and PyTorch internals. This minimal example vividly illustrates how easily a tensor can become a "Zombie" – appearing valid in shape but critically corrupted in its data backing, posing a significant risk to program stability.

Why This Matters: Implications for Developers

The discovery and understanding of this