Torch.compile can be debugged now!

compiler
1

We are happy to announce that, with the newly introduced depyf package, we can debug torch.compile generated code now!

Example usage:

import torch

@torch.compile(backend="eager")
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

import depyf
with depyf.prepare_debug(toy_example, "./dump_src_debug_function"):
    for _ in range(100):
        toy_example(torch.randn(10), torch.randn(10))

with depyf.debug():
    toy_example(torch.randn(10), torch.randn(10))

General guideline for debugging torch.compile:

  • Use backend="eager" to debug Dynamo first, as this is the major place of black magic. If Dynamo works as you expect, it is straight-forward to debug the backend.
  • Run your code under with depyf.prepare_debug context for enough times, so that all branches of your code get compiled.
  • When leaving the with depyf.prepare_debug context, you will be prompted to set breakpoints for Dynamo generated code.
  • The following code under with depyf.debug(): context is debuggable.

And of course, remember to pip install depyf first.

A demo gif:

output

More information at GitHub - thuml/depyf: Decompile python bytecode, and understand PyTorch compiler! .

Hope it helps!

9 Likes
2

wow, much thank to team pytorch!
One quick question, what if I use inductor rather than eager backend.

3

Dynamo and Inductor can be separatedly debugged. By using eager backend, you are eseentially debugging Dynamo only.