Где применять torch.compile?

Создано: 28 июля 2025 г. | Последнее обновление: 27 июля 2026 г.

Рекомендуем применять torch.compile к функции самого высокого уровня, которая не вызывает чрезмерных проблем. Обычно это:

• шаг train или eval с оптимизатором, но без цикла,
• ваш nn.Module верхнего уровня
• или некоторые под-nn.Module.

torch.compile не очень хорошо справляется с модулями-обёртками для распределённых вычислений, такими как DDP или FSDP, поэтому рассмотрите возможность применения torch.compile к внутреннему модулю, передаваемому обёртке.

# inference
model = ...
model.compile()

for _ in range(N_ITERS):
    inp = ...
    out = model(inp)

# training
model = ...
opt = torch.optim.Adam(model.parameters())

@torch.compile
def train(mod, data):
    opt.zero_grad(True)
    pred = mod(data[0])
    loss = torch.nn.CrossEntropyLoss()(pred, data[1])
    loss.backward()
    opt.step()

for _ in range(N_ITERS):
    inp = ...
    train(model, inp)

# DistributedDataParallel
model = ...
model.compile()
model_ddp = DistributedDataParallel(model, ...)

for _ in range(N_ITERS):
    inp = ...
    out = model_ddp(inp)

# DO NOT DO THIS
model = MyModel()
model = torch.compile(model)
model(inp)

# DO THIS
model = MyModel()
model.compile()
model(inp)

# this is also acceptable
@torch.compile
def fn(model, inp):
    return model(inp)
model = MyModel()
fn(model, inp)