gistfile1.txt

import torch
from torch._dynamo.testing import AotEagerAndRecordGraphs
import torch.fx.traceback as fx_traceback

def forward(x):
    with fx_traceback.annotate({"pp_stage": 0}):
        with fx_traceback.annotate({"fdsp_bucket": 0}):
            sin = torch.sin(x)
        sub = sin - 2
    with fx_traceback.annotate({"cuda_stream": 2, "fsdp_bucket": 1}):
        mul = sub * 2
        div = mul / 3
    return div, mul

backend = AotEagerAndRecordGraphs()
opt_m = torch.compile(forward, backend=backend, fullgraph=True)
x = torch.randn(10, requires_grad=True)
opt_m(x)[1].sum().backward()

and if I run your gist with TORCH_LOGS="aot_graphs", I do see the same annotation

P2023694204