Ускорение spaCy-трансформеров
через TensorRT и ONNX Runtime

Проблема

spaCy — стандарт для production NLP-пайплайнов. Модели типа en_core_web_trf используют трансформеры (RoBERTa, BERT) как feature-extractor: на каждый батч текстов трансформер делает forward pass, результат идёт в компоненты NER, POS-теггера и других.

PyTorch сам по себе не оптимален для инференса. Он проектировался для обучения: граф вычислений строится динамически, каждый оператор запускается отдельно, автодифференцирование тянет за собой граф для градиентов, которые при инференсе не нужны. На больших нагрузках это заметно.

Как устроен spaCy изнутри

Архитектура пайплайна

Любой spaCy-пайплайн — это последовательность компонентов, применяемых к тексту по цепочке. Трансформер занимает центральное место: он вычисляет векторные представления, которые все последующие компоненты читают из Doc._.trf_data.

Внутренняя иерархия объектов

spaCy использует библиотеку thinc как слой абстракции над ML-фреймворками. Это создаёт вложенную иерархию: от компонента до конкретного nn.Module.

Подход к ускорению

PyTorch нужен для обучения, для инференса он избыточен. Идея: один раз экспортировать модель в статичный граф (ONNX), запустить его через специализированный движок (TensorRT или ORT), подменить shim._model на прокси-объект с тем же интерфейсом.

# Вся оптимизация — один вызов
nlp = spacy_accelerate.optimize(nlp, provider="tensorrt", precision="fp16")

# После этого nlp работает как обычно
doc = nlp("Apple Inc. was founded by Steve Jobs in Cupertino.")
docs = list(nlp.pipe(texts, batch_size=128))

Внутри optimize() выполняется 9 шагов:

1. 01
   Validate pipeline Проверяем, что в пайплайне есть компонент transformer
2. 02
   Find transformer shim Рекурсивно обходим thinc-иерархию в поисках PyTorchShim
3. 03
   Detect architecture Читаем state_dict, определяем num_layers и hidden_size
4. 04
   Build OptimizeConfig Собираем конфигурацию runtime с валидацией всех параметров
5. 05
   Check provider Проверяем, что TensorrtExecutionProvider доступен в ORT
6. 06
   Get or export ONNX model Cache hit → берём готовый файл. Cache miss → конвертируем через export_to_onnx()
7. 07
   Create proxy Инстанцируем IOBindingProxy / ORTProxy / CPUProxy
8. 08
   Patch pipeline shim._original_model = shim._model, затем shim._model = proxy
9. 09
   Warmup inference Прогреваем движок, заранее компилируем TRT-ядра для каждого bucket

Конвертация весов

spaCy использует curated-transformers со своим форматом хранения весов, который отличается от HuggingFace. Чтобы экспортировать в ONNX, нужна HF-модель — значит, нужно переложить веса. Это самая нетривиальная часть.

Главная проблема: матрицы Q, K, V

В curated-transformers Q, K, V объединены в одну матрицу. В HuggingFace это три отдельных тензора.

Полная карта переименований (слой i)

После маппинга создаём HF-модель с load_state_dict(strict=False) и проверяем parity: прогоняем dummy input через обе модели, сравниваем max absolute diff — должно быть < 1e-4.

ONNX export и FP16

Export

torch.onnx.export(
    ONNXWrapper(hf_model),    # возвращает только last_hidden_state
    (dummy_input_ids, dummy_mask),
    onnx_path,
    input_names=["input_ids", "attention_mask"],
    output_names=["last_hidden_state"],
    dynamic_axes={              # динамические оси для батча и длины
        "input_ids":         {0: "batch", 1: "seq"},
        "attention_mask":    {0: "batch", 1: "seq"},
        "last_hidden_state": {0: "batch", 1: "seq"},
    },
    opset_version=17,
)

FP16 конвертация

optimizer.optimize_model(
    onnx_path,
    model_type="bert",
    optimization_options=FusionOptions(
        enable_bias_gelu=False  # TensorRT не поддерживает этот кастомный ORT-оператор
    )
)
model.convert_float_to_float16(keep_io_types=True)
# keep_io_types=True: int64 входы и float32 выход сохраняются,
# внутренние операции переводятся в FP16

Фьюжны, которые применяются: LayerNorm, Multi-Head Attention, GELU. BiasGelu отключён — TensorRT не имеет для него custom op.

Кеширование

Экспорт занимает 20–60 секунд. Повторять при каждом запуске неприемлемо.

Ключ кеша

cache_key = SHA256({
    "name":               nlp.meta["name"],        # "en_core_web_trf"
    "version":            nlp.meta["version"],     # "3.8.0"
    "lang":               nlp.meta["lang"],         # "en"
    "precision":          "fp16",
    "fixed_batch_size":   None,
    "structure_hash":     SHA256(sorted first 20 state_dict keys)[:8],
    "first_weight_shape": str(state_dict[first_key].shape),
})[:16]

structure_hash и first_weight_shape — страховка от модификации весов в runtime. Если модель изменилась, ключ изменится и экспорт произойдёт заново.

Структура директории

~/.cache/spacy-accelerate/
├── a3f1b2c4d5e6f7a8/
│   ├── model.onnx
│   └── model_fp16.onnx
├── 9d8e7f6a5b4c3d2e/
│   └── model_fp16.onnx
└── trt_engines/
    └── TensorrtExecutionProvider_cache_...   # скомпилированные GPU-ядра

TensorRT кеширует скомпилированные ядра в trt_engines/. При совпадении формы входа TRT переиспользует готовый engine — без компиляции.

Runtime: IO Binding

Два способа запускать ONNX Runtime с GPU. Разница — в том, как данные попадают в движок.

IO Binding передаёт в ONNX Runtime не данные, а GPU-указатели через data_ptr(). Движок читает входы и пишет выходы напрямую в GPU-буферы. Для TensorRT FP16 это устраняет основной bottleneck — memory bandwidth между GPU и CPU.

Batch bucketing

TensorRT компилирует отдельное GPU-ядро для каждой уникальной формы входа. В production батчи разных размеров — норма, и каждый новый размер без подготовки вызывает компиляцию длиной 100–500 мс. Это неприемлемо.

Решение: заранее задать набор bucket-размеров и паддить каждый батч до ближайшего сверху. Все ядра компилируются один раз при старте, в рантайме компиляций нет.

Аналогично работает fixed_seq_length: если известна максимальная длина последовательности (например, 144 токена для en_core_web_trf), паддим до неё на GPU — TRT дополнительно оптимизирует под фиксированную форму.

# Оптимальная конфигурация для en_core_web_trf
nlp = spacy_accelerate.optimize(
    nlp,
    provider="tensorrt",
    precision="fp16",
    batch_buckets=[8, 16, 64, 96, 128, 160, 192, 256],
    fixed_seq_length=144,       # 128 window + 16 overhead thinc
    trt_timing_cache=True,
)

Совместимость выходов

ONNX экспортирует только last_hidden_state — финальный слой трансформера. Но spaCy и разные версии spacy-transformers обращаются к результату по разным именам атрибутов. Создаём объект-обёртку, который под всеми этими именами возвращает один и тот же тензор:

class UniversalTransformerOutput:
    def __init__(self, hidden_state, num_layers):
        # Все известные имена атрибутов — один тензор
        self.embedding_output         = hidden_state
        self.last_hidden_layer_state   = hidden_state
        self.all_hidden_layer_states   = [hidden_state] * num_layers
        self.layer_hidden_states       = self.all_hidden_layer_states
        self.last_hidden_layer_states  = self.all_hidden_layer_states
        self.all_outputs               = self.all_hidden_layer_states
        self.num_layers                = num_layers

Упрощение: все "промежуточные слои" одинаковы — это дубликат финального. Для NER и POS-теггера это не проблема — они используют только последний слой. Для probing по промежуточным слоям потребовался бы полный экспорт всех hidden states.

Режимы работы

Первый запуск TensorRT компилирует GPU-ядра для каждого bucket-размера. Скомпилированные ядра кешируются и переиспользуются во всех последующих запусках.

Поддерживаемые модели

Архитектура определяется автоматически по ключам state_dict:

"curated_encoder" в ключах  →  curated-transformers (spaCy native)
"roberta"          в ключах  →  RoBERTa
"distilbert"       в ключах  →  DistilBERT
"xlm"              в ключах  →  XLM-RoBERTa
иначе                        →  BERT

Ограничения

Только inference

Граф заморожен после экспорта. Fine-tuning через spacy-accelerate невозможен — в ONNX Runtime нет автодифференцирования.

Только финальный hidden state

ONNX экспортирует один выход. Если downstream-компонент обращается к промежуточным слоям, они все будут одинаковы (дубликат финального).

Формы батчей для TensorRT

TRT требует заранее известных форм. Для нетипичных размеров батчей нужно добавить их в batch_buckets. Иначе при первом появлении нового размера будет пауза на компиляцию.

CUDA 12.x

Весь стек (PyTorch, TensorRT, cupy, onnxruntime-gpu) собран под CUDA 12. CUDA 11 не поддерживается.

Один GPU

Мульти-GPU не реализован. device_id выбирает устройство, параллелизм по нескольким GPU отсутствует.

Итог

Пять решений, которые вместе дают результат: