PyTorch Fake Export

   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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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))
        # Transfer to fc2's device (no-op when fc1 and fc2 share the same device).
        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")

    # PyTorch export specializes the graph on the device configuration.
    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.")
    # The graph contains both cpu and cuda ops plus an aten._to_copy transfer node.
    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.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
$ 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]"):
            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/activation.py:816 in forward, code: return F.gelu(input, approximate=self.approximate)
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            # File: /mnt/test_torch_fake_export.py:33 in forward, code: h = h.to(self.fc2.weight.device)
            _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

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            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:
    # inputs
    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

    # outputs
    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]"):
            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/activation.py:816 in forward, code: return F.gelu(input, approximate=self.approximate)
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            # File: /mnt/test_torch_fake_export.py:33 in forward, code: h = h.to(self.fc2.weight.device)
            _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

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            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:
    # inputs
    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

    # outputs
    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]"):
            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[4, 256]" = torch.ops.aten.linear.default(x, p_fc1_weight, p_fc1_bias);  x = p_fc1_weight = p_fc1_bias = None

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/activation.py:816 in forward, code: return F.gelu(input, approximate=self.approximate)
            gelu: "f32[4, 256]" = torch.ops.aten.gelu.default(linear);  linear = None

            # File: /mnt/test_torch_fake_export.py:33 in forward, code: h = h.to(self.fc2.weight.device)
            _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

            # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:134 in forward, code: return F.linear(input, self.weight, self.bias)
            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:
    # inputs
    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

    # outputs
    linear_1: USER_OUTPUT

Range constraints: {}

PyTorch Fake Export

https://leimao.github.io/blog/PyTorch-Fake-Export/

Author

Lei Mao

Posted on

05-17-2026

Updated on

05-17-2026

Licensed under

 

Like this article? Support the author with

Comments

  Follow   Sponsor

Advertisement


Catalogue

×