Шаги по отладке медленных прогнозов BERT (и других LLM) на персональном компьютере

Все началось, когда я экспериментировал с моделями BERT и получил зловещее сообщение, которого все специалисты по данным надеются избежать:

Страшное сообщение «Kernel Died» 💀

Это случилось со мной, когда я запускал свою модель TensorFlow BERT на своем ноутбуке Jupyter. Обучение больших языковых моделей (LLM), как известно, требует больших объемов данных и вычислений, поэтому мой сравнительно слабенький ноутбук может иметь смысл зависнуть здесь…

… за исключением того, что этот сбой произошел во время прогнозирования, а не обучения, что было странно, учитывая мое предположение, что во время обучения использовалось больше памяти, чем при прогнозировании.

Предоставленная ошибка «Kernel Died», к сожалению, не очень информативна, и построчная отладка через TensorFlow звучала как сложное упражнение.

Несколько быстрых поисков в Stack Overflow также не полностью ответили на мои нерешенные вопросы. Но мне все еще нужен был путь вперед.

Это мое исследование проблемы умирания ядра и то, как я нашел решение. 🚀

Копать глубже

Учитывая, что единственное, что я знал о своей проблеме, это то, что ядро ​​умерло, мне пришлось собрать больше контекста. Из нескольких других потоков стало ясно, что причиной смерти ядра было то, что предсказание моей модели требовало больше оперативной памяти, чем мог обеспечить мой процессор (8 ГБ), даже во время предсказания.

Теперь очень прямое решение (которое большинство предположило бы) — просто получить или арендовать графический процессор через Google Colab или что-то в этом роде. И я думаю, что это, безусловно, жизнеспособное решение.

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

Размер партии

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

Сначала я написал три упрощенные версии BERT, изменив только размер пакетов, используемых моделью. Я запускал три версии:

  • FULL: BERT прогнозирует сразу все входные данные.
  • SINGLE: прогнозирование BERT для одного входа за раз
  • BATCH (100): прогнозирование BERT пакетами по 100 входных данных за раз.

Код для этого ниже:

from transformers import BertTokenizer, BertForSequenceClassification, TFBertForSequenceClassification
import tensorflow as tf

class BERT_model_full:
    """
    BERT model predicting on all inputs at once
    """

    def __init__(self):

        self.model = TFBertForSequenceClassification.from_pretrained("bert-base-uncased")
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

    def predict(self,inputs):

        tf_batch = self.tokenizer(inputs, max_length=128, padding=True, truncation=True, return_tensors='tf')
        tf_outputs = self.model(tf_batch)

        return(tf_outputs.logits.numpy())
    
class BERT_model_batch:
    """
    BERT model predicting on batches of 100 inputs at a time
    """

    def __init__(self):

        self.model = TFBertForSequenceClassification.from_pretrained("bert-base-uncased")
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

    def predict(self,inputs):

        # Pred by batchsize
        i = 0
        batch_size = 100
        og_preds = []
        int_preds = []

        while i < len(inputs):

            j = min([len(inputs),i+batch_size])
            tf_batch = self.tokenizer(inputs[i:j], max_length=128, padding=True, truncation=True, return_tensors='tf')
            tf_outputs = self.model(tf_batch)
            
            i = j


        return(True)
    
class BERT_model_single:
    """
    BERT model predicting on a single input at a time
    """

    def __init__(self):

        self.model = TFBertForSequenceClassification.from_pretrained("bert-base-uncased")
        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

    def predict(self,inputs):

        for i in inputs:
            tf_batch = self.tokenizer([i], max_length=128, padding=True, truncation=True, return_tensors='tf')
            tf_outputs = self.model(tf_batch)

        return(tf_outputs.logits.numpy())

Затем я прогнал каждую из этих моделей через одни и те же тестовые примеры с увеличением размера входных данных. Для этого я использовал классический набор данных imdb.

size_list = [1,10,100,1000,2000,4000,6000,8000]
single_time_list = []
batch_time_list = []
full_time_list = []

BERT = BERT_model_single()
print("BERT Single Input:")
for s in size_list:
    data = imdb_data.sample(s)['DATA_COLUMN']
    
    start = time.time()
    
    _ = BERT.predict(data)
    
    end = time.time()
    
    single_time_list.append(end-start)
    print(f"{s} samples: {(end-start)/60:.2f} minutes")
    
BERT = BERT_model_batch()
print("\nBERT Small Batch:")
for s in size_list:
    data = list(imdb_data.sample(s)['DATA_COLUMN'])
    
    start = time.time()
    
    _ = BERT.predict(data)
    
    end = time.time()
    
    batch_time_list.append(end-start)
    print(f"{s} samples: {(end-start)/60:.2f} minutes")
    
    
BERT = BERT_model_full()
print("\nBERT Full Batch:")
for s in size_list:
    data = list(imdb_data.sample(s)['DATA_COLUMN'])
    
    start = time.time()
    
    _ = BERT.predict(data)
    
    end = time.time()
    
    full_time_list.append(end-start)
    print(f"{s} samples: {(end-start)/60:.2f} minutes")

И построение графика вывода показало интересную тенденцию:

BATCH превосходит SINGLE по производительности, потому что большинство моделей и пакетов машинного обучения, таких как Tensorflow, предназначены для использования преимуществ векторизации.

Но что было удивительно, так это насколько FULL хуже по сравнению с BATCH.

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

FULL на самом деле работал хуже, чем обработка одного ввода за раз без какой-либо векторизации на больших входных данных. 🤯

Примерно при 2000 примерах эти требования к ОЗУ начинают сказываться на моем процессоре. И что удивительно, так это то, что до достижения этих 2 КБ разница между BATCH и FULL не так уж отличается.

Основываясь на приведенной выше диаграмме, я предположил, что использование пакета размером 2000 даст наилучшие результаты.

Я был неправ.

Похоже, что лучший размер пакета ближе к 1 КБ, потому что время прогнозирования начинает увеличиваться, если мы используем размер пакета 2 КБ:

Токенизатор

Следующим фрагментом кода, который я исследовал, был Tokenizer. Учитывая, сколько гиперпараметров содержится в строке, я подумал, что это также место для оптимизации:

tf_batch = self.tokenizer(inputs, max_length=128, 
                          padding=True, truncation=True, 
                          return_tensors='tf')

Однако, когда я проверял производительность своей ПОЛНОЙ модели как на входных данных 1 КБ, где она работала на уровне BATCH, так и на 4 КБ, где она работала значительно хуже, время работы токенизатора составляло несущественную часть общего времени:

1000 samples:
Tokenizer Time: 0.06 minutes
Predictionn Time: 1.97 minutes
Tokenizer takes up 3.06%  of prediction time

4000 samples:
Tokenizer Time: 0.29 minutes
Predictionn Time: 27.25 minutes
Tokenizer takes up 1.06%  of prediction time

Хотя увеличение времени Tokenizer немного опережало увеличение размера входных данных (учетверение размера входных данных привело к увеличению времени Tokenizer в 4,8 раза), время предсказания увеличилось в поразительные 13,8x!

Очевидно, что проблема в .predict() части моей воронки продаж.

Версия тензорного потока

Основываясь на уже упомянутом выше потоке Stack Overflow, наиболее популярным решением было понизить Tensorflow, чтобы ускорить прогнозирование.

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

Перейдя на страницу Tensorflow Pypi, мы можем увидеть более старые версии пакета. Выбирая пакеты, выпущенные примерно через год, мы получаем следующие версии пакетов:

  • 2.10.0, выпущен в сентябре 2022 г.
  • 2.6.1, выпущено в ноябре 2021 г.
  • 1.15.4, выпущен в сентябре 2020 г.
  • 1.15.0, выпущено в октябре 2019 г.

Чтобы итеративно устанавливать разные версии одного и того же пакета, нам нужно использовать пакет os, позволяющий запускать команды терминала из кода Python:

import os
data = list(imdb_data.sample(4000)['DATA_COLUMN'])
full_time_list = []
versions = ["2.10.0","2.6.1","1.15.4","1.15.0"]
for version in versions:
    print(version,":")
    os.system(f"pip install tensorflow=={version}")
    
    try:
        from transformers import BertTokenizer, BertForSequenceClassification, TFBertForSequenceClassification
        import tensorflow as tf
    except:
        print("Cannot import relevant packages")
        continue
    
    BERT = BERT_model_full()
    
    start = time.time()
    
    _ = BERT.predict(data)
    
    end = time.time()
    
    minutes = (end-start)/60
    full_time_list.append(minutes)
    print(f"{s} batch size: {minutes:.2f} minutes")
  • Предложение try/except здесь, потому что мы не знаем, существовали ли эти функции в более ранних версиях пакета.К счастью, все они есть
  • Операторы import в цикле выглядят неправильно, но необходимы, поскольку нам нужно повторно импортировать функции после установки правильной версии пакета.

Перебрав каждую версию, мы обнаруживаем, что понижение версии tensorflow улучшает время выполнения на целых 15 %!

Моя любимая теория о том, почему это так, заключается в том, что новые версии tensorflow создаются с учетом интенсивного использования графического процессора, что означает, что он оптимизирован для этого конкретного варианта использования за счет локальной производительности процессора.

Если у кого-то есть ответ на вопрос, почему старые версии TF работают быстрее, дайте мне знать!

Заключение и все вместе

Со следующими сведениями о среде выполнения Tensorflow:

  • Оптимальный размер пакета прогноза составляет около 1000.
  • Параметры токенизатора действительно играют большую роль во времени предсказания
  • Tensorflow 1.X.X увеличивает время прогнозирования на 15%

Мы можем собрать все это вместе и посмотреть, как это работает по сравнению с нашим исходным экспериментом с размером партии:

В самом большом испытанном случае наш оптимальный запуск превосходит пакетный (100) на 20% и одиночный на 57%!

В целом, это упражнение было простым и приятным выражением того, что значит быть Data Scientist. Вам нужно определить проблему, выдвинуть гипотезу, разработать строгий тест и проанализировать свои результаты. В данном случае это была моя среда выполнения Tensorflow. В будущем, я уверен, вы столкнетесь со сложными данными/вопросами/проблемами в своей работе.

И в следующий раз, я надеюсь, вместо того, чтобы проверять Stack Overflow и сдаваться, если ответа нет, вы закатаете рукава и сами исследуете проблемное пространство. Вы никогда не знаете, что вы можете узнать 💡

Надеюсь, это было полезно при отладке ваших проблем со временем прогнозирования Tensorflow! 🎉

Все изображения, если не указано иное, принадлежат автору