PyTorch Fake Export

05-17-2026 05-17-2026 blog 11 minutes read (About 1634 words) visits

Introduction

PyTorch torch.export APIs produces a standardized, single-graph representation of a deep learning model designed for deployment in Python-less environments. Unlike other model export APIs and graph representations, such as torch.onnx.export APIs and ONNX, PyTorch exported program is not a pure device graph. In one PyTorch exported program, the graph can contain operators on different devices, and even explicit device transfer operators. Consequently, there is a difference between the CPU exported program and GPU exported program of the same PyTorch model.

Deep learning models are getting much larger and complex nowadays, and it is often common that a model cannot fit on a single GPU. On host, usually there is plenty of CPU memory to fit and run one large model. So a natural question is, how can the developer verify that the large model being developed can be successfully exported to a GPU exported program using torch.export APIs. Certainly, moving the model to CPU and run torch.export on CPU is an incorrect way to verify, even if the model can be run successfully on CPU, because of the difference between CPU and GPU exported programs. To address this problem, PyTorch provides a way to construct a fake model whose parameters are fake tensors on specific devices, such as CPU or GPU, that have no actual data. To verify the exportability of a large fake model, the developer can also use fake tensors as example inputs to run torch.export APIs for tracing.

In this blog post, I would like to discuss how to run PyTorch export for fake models with fake tensors for verifying the torch.export compatibility of a large model.

PyTorch Fake Export

In the following example, we define a simple MLP model with two linear layers and a GELU activation in between. The first linear layer and the activation are placed on fc1_device, and the second linear layer is placed on fc2_device. When executing the model, the output of the activation is transferred to the device of the second linear layer before being fed into it.

To instantiate a fake model, we create the model inside FakeTensorMode. To specify the device placement of the model, we can use torch.device context manager when constructing the linear layers, instead of using the PyTorch to API because the to move API would require accessing the actual data of tensors.

test_torch_fake_export.py

import torch
import torch.nn as nn
from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode


class MLP(nn.Module):
    """MLP configurable across CPU, GPU, or a CPU-GPU hybrid split.

    fc1 (+ GELU) is placed on *fc1_device*; fc2 is placed on *fc2_device*.
    When the two devices differ the forward pass inserts an explicit device
    transfer, preserved as an aten._to_copy node in the exported graph.
    When they are the same the transfer is a no-op.
    """

    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        out_features: int,
        fc1_device: torch.device = torch.device("cpu"),
        fc2_device: torch.device = torch.device("cpu")
    ) -> None:
        super().__init__()
        with torch.device(fc1_device):
            self.fc1 = nn.Linear(in_features, hidden_features)
            self.act = nn.GELU()
        with torch.device(fc2_device):
            self.fc2 = nn.Linear(hidden_features, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.act(self.fc1(x))
        
        h = h.to(self.fc2.weight.device)
        return self.fc2(h)


def fake_export(fc1_device: torch.device,
                fc2_device: torch.device) -> torch.export.ExportedProgram:
    """Export the MLP with fake tensors for the given device configuration.

    Both parameters and the example input are fake tensors created inside
    FakeTensorMode, so no real memory is allocated on either device.
    """
    with FakeTensorMode():
        model = MLP(in_features=128,
                    hidden_features=256,
                    out_features=10,
                    fc1_device=fc1_device,
                    fc2_device=fc2_device).eval()
        assert all(isinstance(p, FakeTensor) for p in model.parameters(
        )), "Model parameters were unexpectedly materialized (not FakeTensor)"
        example_input = torch.randn(4, 128, device=fc1_device)
    return torch.export.export(model, (example_input, ))


if __name__ == "__main__":

    cpu_device = torch.device("cpu")
    gpu_device = torch.device("cuda")

    
    ep_cpu = fake_export(fc1_device=cpu_device, fc2_device=cpu_device)
    print("MLP export (CPU) succeeded.")
    print(ep_cpu)

    ep_gpu = fake_export(fc1_device=gpu_device, fc2_device=gpu_device)
    print("MLP export (GPU) succeeded.")
    print(ep_gpu)

    ep_hybrid = fake_export(fc1_device=cpu_device, fc2_device=gpu_device)
    print("CPU-GPU hybrid export succeeded.")
    
    print(ep_hybrid)

Using NVIDIA NGC Docker container nvcr.io/nvidia/pytorch:26.04-py3, we could run the above script and see the successful export of the MLP fake model for CPU, GPU, and CPU-GPU hybrid device configurations. There will be no actual data allocated for the model parameters and the example input during the export, thanks to the use of fake tensors.

$ python test_torch_fake_export.py
MLP export (CPU) succeeded.
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_fc1_weight: "f32[256, 128]", p_fc1_bias: "f32[256]", p_fc2_weight: "f32[10, 256]", p_fc2_bias: "f32[10]", x: "f32[4, 128]"):
            
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            
            _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(gelu, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default = None
            to: "f32[4, 256]" = torch.ops.aten.to.dtype_layout(gelu, dtype = torch.float32, layout = torch.strided, device = device(type='cpu'));  gelu = None

            
            linear_1: "f32[4, 10]" = torch.ops.aten.linear.default(to, p_fc2_weight, p_fc2_bias);  to = p_fc2_weight = p_fc2_bias = None
            return (linear_1,)

Graph signature:
    
    p_fc1_weight: PARAMETER target='fc1.weight'
    p_fc1_bias: PARAMETER target='fc1.bias'
    p_fc2_weight: PARAMETER target='fc2.weight'
    p_fc2_bias: PARAMETER target='fc2.bias'
    x: USER_INPUT

    
    linear_1: USER_OUTPUT

Range constraints: {}

MLP export (GPU) succeeded.
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_fc1_weight: "f32[256, 128]", p_fc1_bias: "f32[256]", p_fc2_weight: "f32[10, 256]", p_fc2_bias: "f32[10]", x: "f32[4, 128]"):
            
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            
            _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(gelu, dtype = torch.float32, device = device(type='cuda', index=0), layout = torch.strided);  _assert_tensor_metadata_default = None
            to: "f32[4, 256]" = torch.ops.aten.to.dtype_layout(gelu, dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0));  gelu = None

            
            linear_1: "f32[4, 10]" = torch.ops.aten.linear.default(to, p_fc2_weight, p_fc2_bias);  to = p_fc2_weight = p_fc2_bias = None
            return (linear_1,)

Graph signature:
    
    p_fc1_weight: PARAMETER target='fc1.weight'
    p_fc1_bias: PARAMETER target='fc1.bias'
    p_fc2_weight: PARAMETER target='fc2.weight'
    p_fc2_bias: PARAMETER target='fc2.bias'
    x: USER_INPUT

    
    linear_1: USER_OUTPUT

Range constraints: {}

CPU-GPU hybrid export succeeded.
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_fc1_weight: "f32[256, 128]", p_fc1_bias: "f32[256]", p_fc2_weight: "f32[10, 256]", p_fc2_bias: "f32[10]", x: "f32[4, 128]"):
            
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            
            _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(gelu, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default = None
            to: "f32[4, 256]" = torch.ops.aten.to.dtype_layout(gelu, dtype = torch.float32, layout = torch.strided, device = device(type='cuda', index=0));  gelu = None

            
            linear_1: "f32[4, 10]" = torch.ops.aten.linear.default(to, p_fc2_weight, p_fc2_bias);  to = p_fc2_weight = p_fc2_bias = None
            return (linear_1,)

Graph signature:
    
    p_fc1_weight: PARAMETER target='fc1.weight'
    p_fc1_bias: PARAMETER target='fc1.bias'
    p_fc2_weight: PARAMETER target='fc2.weight'
    p_fc2_bias: PARAMETER target='fc2.bias'
    x: USER_INPUT

    
    linear_1: USER_OUTPUT

Range constraints: {}

推荐订阅源

Lei Mao's Log Book

Introduction

PyTorch Fake Export