Table of Contents
Shortcuts

# Getting Started

Before you read this section, make sure to read the torch.compiler.

Let’s start by looking at a simple `torch.compile` example that demonstrates how to use `torch.compile` for inference. This example demonstrates the `torch.cos()` and `torch.sin()` features which are examples of pointwise operators as they operate element by element on a vector. This example might not show significant performance gains but should help you form an intuitive understanding of how you can use `torch.compile` in your own programs.

Note

To run this script, you need to have at least one GPU on your machine. If you do not have a GPU, you can remove the `cuda()` code in the snippet below and it will run on CPU.

    import torch
    def fn(x, y):
        a = torch.cos(x).cuda()
        b = torch.sin(y).cuda()
        return a + b
    new_fn = torch.compile(fn, backend="inductor")
    input_tensor = torch.randn(10000).to(device="cuda:0")
    a = new_fn(input_tensor, input_tensor)

A more famous pointwise operator you might want to use would be something like `torch.relu()`. Pointwise ops in eager mode are suboptimal because each one would need to read a tensor from the memory, make some changes, and then write back those changes. The single most important optimization that inductor performs is fusion. In the example above we can turn 2 reads and 2 writes into 1 read and 1 write which is crucial especially for newer GPUs where the bottleneck is memory bandwidth (how quickly you can send data to a GPU) rather than compute (how quickly your GPU can crunch floating point operations).

Another major optimization that inductor provides is automatic support for CUDA graphs. CUDA graphs help eliminate the overhead from launching individual kernels from a Python program which is especially relevant for newer GPUs.

TorchDynamo supports many different backends, but TorchInductor specifically works by generating Triton kernels. Let’s save our example above into a file called `example.py`. We can inspect the code generated Triton kernels by running `TORCH_COMPILE_DEBUG=1 python example.py`. As the script executes, you should see `DEBUG` messages printed to the terminal. Closer to the end of the log, you should see a path to a folder that contains `torchinductor_<your_username>`. In that folder, you can find the `output_code.py` file that contains the generated kernel code similar to the following:

    @pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
    @triton.jit
    def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
        xnumel = 10000
        xoffset = tl.program_id(0) * XBLOCK
        xindex = xoffset + tl.arange(0, XBLOCK)[:]
        xmask = xindex < xnumel
        x0 = xindex
        tmp0 = tl.load(in_ptr0 + (x0), xmask)
        tmp1 = tl.cos(tmp0)
        tmp2 = tl.sin(tmp0)
        tmp3 = tmp1 + tmp2
        tl.store(out_ptr0 + (x0), tmp3, xmask)

Note

The above code snippet is an example. Depending on your hardware, you might see different code generated.

And you can verify that fusing the `cos` and `sin` did actually occur because the `cos` and `sin` operations occur within a single Triton kernel and the temporary variables are held in registers with very fast access.

Read more on Triton’s performance here. Because the code is written in Python, it’s fairly easy to understand even if you have not written all that many CUDA kernels.

Next, let’s try a real model like resnet50 from the PyTorch hub.

    import torch
    model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
    opt_model = torch.compile(model, backend="inductor")
    opt_model(torch.randn(1,3,64,64))

And that is not the only available backend, you can run in a REPL `torch.compiler.list_backends()` to see all the available backends. Try out the `cudagraphs` next as inspiration.

## Using a pretrained model

PyTorch users frequently leverage pretrained models from transformers or TIMM and one of the design goals is TorchDynamo and TorchInductor is to work out of the box with any model that people would like to author.

Let’s download a pretrained model directly from the HuggingFace hub and optimize it:

    import torch
    from transformers import BertTokenizer, BertModel
    # Copy pasted from here https://huggingface.co/bert-base-uncased
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
    model = torch.compile(model, backend="inductor") # This is the only line of code that we changed
    text = "Replace me by any text you'd like."
    encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
    output = model(**encoded_input)

If you remove the `to(device="cuda:0")` from the model and `encoded_input`, then Triton will generate C++ kernels that will be optimized for running on your CPU. You can inspect both Triton or C++ kernels for BERT. They are more complex than the trigonometry example we tried above but you can similarly skim through it and see if you understand how PyTorch works.

Similarly, let’s try out a TIMM example:

    import timm
    import torch
    model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
    opt_model = torch.compile(model, backend="inductor")
    opt_model(torch.randn(64,3,7,7))

## Next Steps

In this section, we have reviewed a few inference examples and developed a basic understanding of how torch.compile works. Here is what you check out next:

- torch.compile tutorial on training

- torch.compiler API reference

- TorchDynamo APIs for fine-grained tracing

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.