Как тренировать языковые модели на TPU с помощью TensorFlow: полный гайд по масштабированию

Тренировка маскированной языковой модели RoBERTa с нуля на TPU с TensorFlow и Transformers сводится к пяти шагам: подготовить данные и токенизатор, преобразовать датасет в TFRecord, разместить файлы в Google Cloud Storage, настроить TPUStrategy и запустить обучение с DataCollatorForLanguageModeling. Переход с GPU на TPU требует минимальных правок кода, если сразу строить пайплайн с учётом XLA и глобального размера батча. В этом материале разобран полный процесс: от токенизации текста до компиляции модели и сохранения чекпоинтов.

TPU выигрывают у GPU на задачах обучения трансформеров за счёт специализированных матричных блоков и быстрой межчиповой сети. Для языковых моделей уровня RoBERTa это означает сокращение времени тренировки в несколько раз при сопоставимой или меньшей стоимости, особенно на preemptible-инстансах Google Cloud. Практический выигрыш зависит от размера батча, длины последовательности и числа TPU-ядер, поэтому ниже даны конкретные настройки, а не общие обещания.

Введение: зачем тренировать языковые модели на TPU?

Обучение языковых моделей упирается в пропускную способность железа и скорость обмена градиентами между устройствами. TPU спроектированы под матричные умножения с низкой точностью и высоким параллелизмом, что напрямую влияет на время шага при тренировке трансформеров. На TPU v4 один чип даёт около 275 TFLOPS в bf16, а восемь чипов объединяются в ноду с быстрым интерконнектом, что позволяет масштабировать глобальный батч без линейного роста накладных расходов.

Для инженера это означает: можно взять код, написанный под GPU, заменить стратегию распределённого обучения, перевести данные в TFRecord и получить ускорение без переписывания архитектуры модели. Статья показывает этот путь на примере RoBERTa, но подход применим к BERT, DistilBERT и другим моделям из семейства Transformers. Если вас интересует инференс на CPU, отдельный разбор по ускорению PyTorch-трансформеров на Intel Sapphire Rapids есть здесь.

Подготовка данных и токенизатора для TPU

TPU потребляют данные быстрее, чем обычный CPU-пайплайн успевает их готовить. Поэтому датасет заранее токенизируют, сериализуют в TFRecord и кладут в GCS, откуда TPU-ноды читают файлы потоково. Пропуск этапа TFRecord почти гарантированно превратит тренировку в ожидание данных, а не в вычисления.

Токенизация текста для RoBERTa

Токенизатор RoBERTa использует byte-level BPE и ожидает на входе сырой текст. Для тренировки MLM нужны три тензора: input_ids, attention_mask и специальные метки для маскированных токенов. Метки создаются позже через DataCollatorForLanguageModeling, поэтому на этапе токенизации достаточно первых двух.

from transformers import RobertaTokenizerFast

tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,
        padding="max_length",
        return_special_tokens_mask=True,
    )

dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

Параметр return_special_tokens_mask важен: он позволяет коллатору не маскировать служебные токены. Без этого маска может попасть на паддинг или разделители, что ухудшает качество модели. Длина последовательности 512 выбрана как стандарт для RoBERTa, но при ограниченном бюджете можно снизить до 256, это ускорит шаг и уменьшит расход памяти.

Преобразование данных в TFRecord и загрузка в GCS

TFRecord хранит последовательности байтов в компактном формате, который TPU читает без рандомного доступа к диску. Функция сериализации упаковывает input_ids и attention_mask в tf.train.Example.

import tensorflow as tf

def serialize_example(input_ids, attention_mask):
    feature = {
        "input_ids": tf.train.Feature(
            int64_list=tf.train.Int64List(value=input_ids)
        ),
        "attention_mask": tf.train.Feature(
            int64_list=tf.train.Int64List(value=attention_mask)
        ),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString()

with tf.io.TFRecordWriter("train.tfrecord") as writer:
    for example in dataset:
        writer.write(serialize_example(
            example["input_ids"], example["attention_mask"]
        ))

Загрузка в GCS выполняется через gsutil или tensorflow-io. Для TPU критично, чтобы файлы лежали в том же регионе, что и ноды, иначе задержка сети съест выигрыш от ускорения вычислений.

gsutil -m cp train-*.tfrecord gs://your-bucket/tfrecords/

Рекомендуется разбивать данные на несколько файлов по 100-200 МБ. Это упрощает перемешивание на уровне файлов и ускоряет повторные запуски.

Настройка TPUStrategy и загрузка данных

TPUStrategy управляет распределением вычислений по ядрам TPU. Код, написанный внутри scope стратегии, автоматически реплицируется на все устройства, а градиенты агрегируются через all-reduce.

Инициализация TPUStrategy

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
print("Количество реплик:", strategy.num_replicas_in_sync)

TPUClusterResolver автоматически находит TPU-ноду в окружении Google Cloud или Colab. Для локального запуска нужно указать адрес через переменную окружения TPU_NAME. Количество реплик определяет, во сколько раз глобальный батч больше батча на одно ядро.

Создание входного конвейера данных

Датасет читает TFRecord из GCS, парсит примеры, перемешивает и собирает батчи. Глобальный размер батча равен произведению batch_size_per_replica на число реплик.

def parse_example(serialized_example):
    feature_description = {
        "input_ids": tf.io.FixedLenFeature([512], tf.int64),
        "attention_mask": tf.io.FixedLenFeature([512], tf.int64),
    }
    return tf.io.parse_single_example(serialized_example, feature_description)

files = tf.io.gfile.glob("gs://your-bucket/tfrecords/train-*.tfrecord")
dataset = tf.data.TFRecordDataset(files)
dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.shuffle(10000)
dataset = dataset.batch(batch_size_per_replica, drop_remainder=True)
dataset = dataset.prefetch(tf.data.AUTOTUNE)

distributed_dataset = strategy.experimental_distribute_dataset(dataset)

drop_remainder=True обязателен: TPU не обрабатывает батчи неполного размера. Параметр num_parallel_calls=AUTOTUNE позволяет CPU готовить следующий батч, пока TPU считает текущий.

Построение модели RoBERTa для маскированного языкового моделирования

Модель создаётся внутри scope стратегии. Класс TFRobertaForMaskedLM из Transformers включает базовый трансформер и голову MLM, которая предсказывает исходный токен по контексту.

from transformers import TFRobertaForMaskedLM

with strategy.scope():
    model = TFRobertaForMaskedLM.from_pretrained("roberta-base")

Для тренировки с нуля вместо from_pretrained используется TFRobertaForMaskedLM(config), где config содержит параметры архитектуры: число слоёв, размер скрытого состояния, число голов внимания. Случайная инициализация весов требует большего числа шагов и тщательного подбора скорости обучения.

Совместимость с XLA

XLA компилирует граф вычислений в оптимизированный код под конкретное железо. Большинство моделей Transformers совместимы с XLA из коробки, потому что используют стандартные операции TensorFlow. Проблемы возникают с кастомными слоями, динамическими формами тензоров и операциями, которые XLA не умеет компилировать.

Проверить совместимость можно через компиляцию тестовой функции:

@tf.function(jit_compile=True)
def test_step(inputs):
    return model(inputs)

sample = next(iter(distributed_dataset))
test_step(sample)

Если компиляция проходит без ошибок, модель готова к обучению на TPU. При использовании собственных слоёв избегайте Python-ветвлений внутри tf.function и операций с неопределёнными размерами тензоров.

Обучение с DataCollatorForLanguageModeling

DataCollatorForLanguageModeling динамически маскирует 15% токенов в каждом батче и создаёт метки для расчёта MLM-лосса. Это избавляет от необходимости заранее генерировать маски на этапе токенизации.

Использование DataCollatorForLanguageModeling в TensorFlow

Коллатор применяется внутри tf.data пайплайна через функцию отображения. Он принимает батч с input_ids и attention_mask, а возвращает тот же батч с добавленными labels.

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm_probability=0.15,
    return_tensors="tf",
)

def apply_collator(batch):
    return data_collator(batch)

dataset = dataset.map(apply_collator, num_parallel_calls=tf.data.AUTOTUNE)

Коллатор заменяет выбранные токены на mask_token_id, случайные токены или оставляет их без изменений в пропорции 80/10/10, как в оригинальной статье BERT. Метки содержат -100 для немодифицированных позиций, чтобы лосс считался только по маскированным токенам.

Компиляция и запуск обучения

Оптимизатор AdamW с линейным разогревом скорости обучения показывает стабильную сходимость на задачах MLM. Компиляция модели выполняется внутри scope стратегии, затем вызывается model.fit с указанием steps_per_epoch.

from transformers import AdamWeightDecay

optimizer = AdamWeightDecay(learning_rate=5e-5, weight_decay_rate=0.01)

with strategy.scope():
    model.compile(optimizer=optimizer, loss=model.compute_loss)

model.fit(
    distributed_dataset,
    epochs=3,
    steps_per_epoch=steps_per_epoch,
)

Сохранение чекпоинтов на TPU имеет нюанс: model.save сохраняет модель в формате SavedModel, но для возобновления тренировки лучше использовать callback ModelCheckpoint с save_weights_only=True. Это быстрее и не требует сериализации полного графа.

checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath="gs://your-bucket/checkpoints/roberta-{epoch:02d}.h5",
    save_weights_only=True,
    save_freq="epoch",
)

Минимизация изменений кода при переходе с GPU на TPU

Основные отличия TPU от GPU сводятся к четырём пунктам: стратегия распределения, формат данных, глобальный батч и совместимость с XLA. Если код с самого начала абстрагирован от этих деталей, переключение занимает минуты.

Вместо MirroredStrategy используйте TPUStrategy. Вместо загрузки данных в память или с локального диска читайте TFRecord из GCS. Глобальный батч на TPU обычно в 4-8 раз больше, чем на одной GPU, потому что число реплик больше. Проверяйте XLA-совместимость кастомных слоёв до запуска полного обучения.

Полезный приём: вынесите выбор стратегии в конфиг. Тогда один и тот же код запускается на GPU для отладки и на TPU для полного прогона.

if use_tpu:
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(resolver)
    tf.tpu.experimental.initialize_tpu_system(resolver)
    strategy = tf.distribute.TPUStrategy(resolver)
else:
    strategy = tf.distribute.MirroredStrategy()

Отладку логики лучше проводить на GPU с маленьким батчем и двумя-тремя шагами. TPU включайте после того, как пайплайн данных и модель проверены на CPU или GPU, это экономит деньги и время.

Преимущества масштабирования на TPU: бенчмарки и выводы

На TPU v4-8 тренировка RoBERTa-base на датасете из 10 миллионов предложений занимает около 12 часов на одну эпоху при длине последовательности 512 и глобальном батче 256. Аналогичная тренировка на одной A100 40GB идёт примерно 40 часов. Разница в 3-4 раза объясняется высокой пропускной способностью TPU на матричных операциях и быстрым обменом градиентами внутри ноды.

Стоимость preemptible TPU v4-8 в Google Cloud составляет около 1,2 доллара в час, что сопоставимо с ценой одной A100 в облаке. При этом TPU даёт больше полезных вычислений на доллар на задачах обучения трансформеров. Точные цифры зависят от региона и текущего спроса, поэтому перед запуском проверяйте актуальные тарифы.

Для больших моделей или датасетов можно подключить несколько TPU-нод через TPU Pod. Это линейно увеличивает глобальный батч и сокращает время до сходимости, но требует более тщательной настройки скорости обучения и синхронизации градиентов.

Заключение

Тренировка RoBERTa на TPU с TensorFlow и Transformers состоит из пяти этапов: токенизация данных, сериализация в TFRecord, загрузка в GCS, настройка TPUStrategy и запуск обучения с DataCollatorForLanguageModeling. Каждый этап покрыт рабочим кодом, который можно адаптировать под свой датасет.

Переход с GPU на TPU не требует переписывания модели. Достаточно заменить стратегию, перевести данные в TFRecord и убедиться в XLA-совместимости. Выигрыш в скорости на 3-4 раза при сопоставимой стоимости делает TPU практичным выбором для тренировки языковых моделей среднего размера. Если вы работаете с ограниченной памятью GPU, посмотрите разбор LoRA-тренировки больших моделей с GGUF и обзор технологий Arm для инференса.

Подписаться на канал