Справочник API torch.compiler

torch.compiler.precompile

torch.compiler.precompile(fn, *example_inputs, backend='inductor', tracer='make_fx', decompositions=None)

Предварительно компилирует fn до выполнения, используя примеры входных данных, и возвращает самодостаточную исполняемую исходную строку Python вместе с кэшем ускорения в виде (python_code, cache). fn — это вычисление целиком, принимающее модель (или модели) в качестве явных аргументов, например lambda model, x: model(x) или шаг обучения. У аргументов nn.Module параметры и буферы выносятся во входы графа, поэтому веса не встраиваются в артефакт: во время выполнения вы снова передаёте модель вызываемому объекту после загрузки. Загрузка выполняется с помощью torch.compiler.precompile.load (см. ниже).

Примечание

При использовании трассировщика make_fx по умолчанию захват выполняется в нережимном режиме. Поток управления специализируется под примеры входных данных, а формы статичны — каждый размер встраивается в артефакт. Исключение — измерение тензора, явно помеченное как не имеющее привязки (только для бэкенда inductor) с помощью torch._dynamo.decorators.mark_unbacked на входных данных до вызова; такое измерение захватывается как не имеющий привязки symint, поэтому один артефакт подходит для любого его размера во время выполнения, а захват завершается ошибкой, если графу требуется установить для него защитную проверку. Тип данных и устройство каждого входа также специализируются (несовпадение во время выполнения отклоняется), а бэкенд inductor дополнительно специализируется на формате памяти входных данных. См. Note [precompile programming model] в torch/_precompile.py. torch.compiler.precompile отличается от torch._dynamo.config.caching_precompile (режима кэширования torch.compile).

Если fn выполняет обратный проход, артефакт повторно запускает прямой и обратный проходы целиком и записывает полученные градиенты параметров в поля parameters() .grad модели, переданной во время выполнения, накапливая их (p.grad += g) так же, как при обычном выполнении .backward(). Поэтому используйте привычный цикл zero_grad() / optimizer.step(). Параметры, получающие градиент, фиксируются во время захвата (замороженные параметры и параметры, не вносящие вклад, остаются .grad = None). Артефакт возвращает собственный результат fn (None для отдельного шага .backward()), а не градиенты.

Параметры:

• fn – Вычисление целиком, которое нужно захватить; принимает модели и входные данные времени выполнения в качестве позиционных аргументов.
• example_inputs – Примеры позиционных аргументов для fn; аргументы nn.Module выносятся, остальные являются входными данными времени выполнения.
• backend – "inductor" (по умолчанию) выполняет понижение через AOTAutograd + Inductor; "eager" сохраняет захваченный граф ATen (гибкий к формату памяти, без ядер; формы всё равно специализируются под примеры).
• tracer – Фронтенд захвата. "make_fx" (по умолчанию) — это трассировка make_fx в нережимном режиме и единственный реализованный на данный момент трассировщик; "dynamo" запланирован и пока вызывает NotImplementedError.
• decompositions – Необязательная таблица декомпозиций (dict из OpOverload в функцию декомпозиции), передаваемая в make_fx; по умолчанию — None.

Возвращает:

(python_code, cache) — самодостаточную исходную строку Python (единственный источник истины для соглашения о вызове) и бинарный кэш ускорения (без весов и метаданных соглашения о вызове; содержит небольшую метку целостности формата/версии/бэкенда/code_hash, которую проверяет load).

Вызывает исключение:

PrecompileError — если захват, понижение или вызов во время выполнения нарушает условия (см. описание исключения ниже).

Пример:

python_code, cache = torch.compiler.precompile(lambda m, x: m(x), model, x)
f = torch.compiler.precompile.load(python_code, cache)
out = f(model, x)   # pass the model again at runtime

precompile.load(python_code, cache) [исходный код]

Восстанавливает исполняемый объект из пары (python_code, cache), возвращённой precompile. Соглашение о вызове считывается из python_code (единственного источника истины); cache лишь ускоряет загрузку — он содержит только скомпилированный артефакт бэкенда (пакет Inductor для backend="inductor"; пустой для backend="eager") и не содержит весов. Во время выполнения вы снова передаёте модели.

Предупреждение

load выполняет артефакт как код: он исполняет python_code (через exec) и для бэкенда inductor предварительно заполняет кэши ядер из cache. Считайте (python_code, cache) доверенным исполняемым вводом — загружайте только пару, созданную вами или иным образом заслуживающую доверия, точно так же, как и любой другой код, который вы собираетесь запустить (см. Note [precompile programming model], инвариант 7). Перед запуском load также выводит предупреждение при каждом вызове.

Параметры:

• python_code – Самодостаточная исходная строка Python, возвращённая precompile.
• cache – Бинарный кэш ускорения, возвращённый precompile.

Возвращает:

Исполняемый вызываемый объект с тем же соглашением о вызове, что и у захваченного fn. Аргументы сопоставляются позиционно как во время захвата, так и во время загрузки; соглашения о вызове с именованными аргументами не поддерживаются.

Вызывает исключение:

PrecompileError — если python_code не является допустимым артефактом предварительной компиляции (не удаётся выполнить разбор или отсутствуют метаданные соглашения о вызове), если cache сопоставлен с другим python_code (не совпадает метка backend или code_hash) либо если вызов во время выполнения нарушает условия предварительной компиляции.

exception torch.compiler.PrecompileError

Тип ошибки, вызываемой torch.compiler.precompile и его артефактами.

Вызывается, когда захват, понижение, load или вызов во время выполнения нарушает условия предварительной компиляции — например, тензор встроен как константа (инвариант 1), используется неподдерживаемая операция или операция с побочными эффектами, выводится нетензорное значение, которое бэкенд inductor не может понизить, либо форма или формат памяти входных данных во время выполнения отличается от примера (инварианты 3 и 6). Полное описание условий см. в Note [precompile programming model] этого модуля.