Начало работы

Прежде чем читать этот раздел, убедитесь, что вы ознакомились с torch.compiler.

Начнем с простого torch.compile примера, который демонстрирует, как использовать torch.compile для вывода. Этот пример демонстрирует torch.cos() и torch.sin() функции, которые являются примерами точечных операторов, поскольку они работают по элементу над вектором. Этот пример может не показать значительного прироста производительности, но должен помочь вам получить интуитивное представление о том, как вы можете использовать torch.compile в собственных программах.

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)

Более известным точечным оператором, который вы можете использовать, будет что-то вроде torch.relu(). Точечные операции в режиме eager неэффективны, поскольку каждый из них должен считывать тензор из памяти, внести некоторые изменения и затем записать эти изменения обратно. Самая важная оптимизация, которую выполняет inductor, — это слияние. В примере выше мы можем преобразовать 2 чтения и 2 записи в 1 чтение и 1 запись, что имеет решающее значение, особенно для новых графических процессоров, где узким местом является пропускная способность памяти (скорость передачи данных на GPU), а не вычислительная мощность (скорость, с которой ваш GPU может обрабатывать операции с плавающей точкой).

Еще одной важной оптимизацией, предоставляемой inductor, является автоматическая поддержка CUDA-графиков. CUDA-графики помогают устранить накладные расходы, связанные с запуском отдельных ядер из программы Python, что особенно актуально для новых графических процессоров.

TorchDynamo поддерживает множество различных бэкэндов, но TorchInductor специально работает путем генерации ядер Triton. Сохраним наш пример выше в файле с именем example.py. Мы можем просмотреть сгенерированные ядра Triton, выполнив TORCH_COMPILE_DEBUG=1 python example.py. При выполнении скрипта на терминале должны отображаться сообщения DEBUG. Ближе к концу лога вы должны увидеть путь к папке, которая содержит torchinductor_<your_username>. В этой папке вы можете найти файл output_code.py, содержащий сгенерированный код ядра, аналогичный следующему:

@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)

И вы можете проверить, что слияние cos и sin действительно произошло, потому что операции cos и sin происходят внутри одного ядра Triton, а временные переменные хранятся в регистрах с очень быстрым доступом.

Дополнительную информацию о производительности Triton можно найти здесь. Поскольку код написан на Python, его довольно легко понять, даже если вы не писали много ядер CUDA.

Далее попробуем реальную модель, например, resnet50 из 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))

И это не единственный доступный бэкэнд; вы можете запустить REPL torch.compiler.list_backends() , чтобы увидеть все доступные бэкэнды. Попробуйте cudagraphs в качестве вдохновения.

Использование предварительно обученной модели

Пользователи PyTorch часто используют предварительно обученные модели из transformers или TIMM, и одной из целей разработки TorchDynamo и TorchInductor является работа «из коробки» с любой моделью, которую пользователи хотят создать.

Загрузим предварительно обученную модель напрямую с платформы HuggingFace и оптимизируем ее:

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)

Если вы удалите to(device="cuda:0") из модели и encoded_input, Triton сгенерирует ядра C++, которые будут оптимизированы для работы на вашем процессоре. Вы можете просмотреть как ядра Triton, так и ядра C++ для BERT. Они сложнее, чем пример с тригонометрией, который мы пробовали выше, но вы можете аналогичным образом просмотреть их и увидеть, понимаете ли вы, как работает PyTorch.

Аналогичным образом, давайте попробуем пример TIMM:

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))

Следующие шаги

В этом разделе мы рассмотрели несколько примеров вывода и получили базовое понимание того, как работает torch.compile. Вот что вам следует изучить дальше:

• Учебник по torch.compile по обучению
• Справочник по API torch.compiler
• API TorchDynamo для точного отслеживания