Высокопроизводительный ускоритель машинного обучения, созданный AWS

Об этой статье

В предыдущих сериях мы представили популярную платформу машинного обучения PyTorch. Часть первая: основы. Часть вторая: конкретный пример использования PyTorch. Часть третья: как компилятор машинного обучения ускоряет процесс машинного обучения Это часть 4, в которой мы познакомим вас с некоторыми аппаратными знаниями о том, как недавно выпущенный чип AWS может помочь ускорить ваш путь машинного обучения.

Что такое AWS Trainium и Inferentia?

Тип экземпляра Trainium — это одно из семейств микросхем AWS. Он специально разработан, чтобы помочь вам ускорить обучение вашей модели. Точно так же Inferentia также специально разработана, чтобы помочь вам сократить время логического вывода и сократить расходы.

Очень хорошей документацией является AWS Neutron, которая включает всю информацию о выпуске новых чипов, поддерживаемых AWS. AWS также предоставляет пакет SDK для нейтронов, который поможет вам лучше адаптироваться к использованию этих двух чипов.

Архитектура Trainium

Вот схема чипов Trainium. Trn1 состоит из 16 устройств Trainium 1. Каждое устройство Trainium 1 имеет два нейрона для обработки запроса.

Немного статистики с веб-сайта AWS нейтрона.

  • Вычисления: 2 ядра NeuronCore-v2, обеспечивающие 420 INT8 TOPS, 210 FP16/BF16/cFP8/TF32 TFLOPS и 52,5 FP32 TFLOPS.
  • Память устройства: 32 ГБ памяти устройства (для хранения состояния модели) с пропускной способностью 820 ГБ/с.
  • Перемещение данных: пропускная способность прямого доступа к памяти 1 ТБ/с с встроенным сжатием/распаковкой памяти.
  • NeuronLink: NeuronLink-v2 для соединения устройств друг с другом обеспечивает эффективное обучение масштабированию, а также объединение памяти между различными устройствами Trainium.

Пример кода использования Trainium

Есть много вариантов, когда вы пытаетесь использовать преимущества Trainium. Вы можете использовать AWS EKS/ECS со своим собственным контейнером. Вы можете использовать пакетный сервис AWS для выполнения задания HPC. Вы также можете выбрать управляемый сервис машинного обучения AWS SageMaker. Один пример кода от AWS SageMaker выглядит так, как показано ниже, в котором используются как PyTorch, так и Trainium для ускорения вашего пути к машинному обучению.

from sagemaker.pytorch import PyTorch

pt_estimator = PyTorch(
    entry_point="train_torch_distributed.py",
    role="SageMakerRole",
    framework_version="1.11.0",
    py_version="py38",
    instance_count=1,
    instance_type="ml.trn1.2xlarge",
    distribution={
        "torch_distributed": {
            "enabled": True
        }
    }
)

pt_estimator.fit("s3://bucket/path/to/training/data")

Сравнение цен между разными экземплярами

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

В итоге

В этой статье мы кратко представили два чипа из семейства AWS, чтобы ускорить обучение и выводы для вашего пути к машинному обучению. Это надстройка для популярных платформ машинного обучения, таких как PyTorch. В следующей статье мы углубимся в автообучение, которое представляет собой рекламный ход PyTorch о том, как он ускоряет ваше обучение.

Пожалуйста, не переключайтесь.