Реализация DeepSpeed для масштабируемых трансформеров: продвинутое обучение с использованием градиентной контрольной точки и параллелизма

В этом продвинутом руководстве по DeepSpeed мы предлагаем практическое ознакомление с передовыми методами оптимизации для эффективного обучения больших языковых моделей. Сочетая оптимизацию ZeRO, обучение со смешанной точностью, накопление градиента и продвинутые конфигурации DeepSpeed, руководство демонстрирует, как максимально использовать память GPU, сократить накладные расходы на обучение и обеспечить масштабирование моделей трансформеров в средах с ограниченными ресурсами, таких как Colab.

Установка зависимостей

Мы устанавливаем PyTorch с поддержкой CUDA, DeepSpeed и необходимые библиотеки, такие как Transformers, Datasets, Accelerate и Weights & Biases. Это обеспечивает готовность к созданию и обучению моделей с использованием DeepSpeed.

“`python
import subprocess
import sys
import os
import json
import time
from pathlib import Path

def install_dependencies():
“””Install required packages for DeepSpeed in Colab”””
print(” Installing DeepSpeed and dependencies…”)
subprocess.check_call([
sys.executable, “-m”, “pip”, “install”,
“torch”, “torchvision”, “torchaudio”, “–index-url”,
“https://download.pytorch.org/whl/cu118”
])
subprocess.check_call([sys.executable, “-m”, “pip”, “install”, “deepspeed”])
subprocess.check_call([
sys.executable, “-m”, “pip”, “install”,
“transformers”, “datasets”, “accelerate”, “wandb”
])
print(” Installation complete!”)

install_dependencies()
“`

Создание синтетического набора данных

Мы создаём синтетический набор данных, генерируя случайные последовательности токенов для имитации реальных текстовых данных. Эти последовательности используются как входные данные и метки, что позволяет быстро тестировать обучение DeepSpeed без использования большого внешнего набора данных.

“`python
class SyntheticTextDataset(Dataset):
“””Synthetic dataset for demonstration purposes”””
def init(self, size: int = 1000, seqlength: int = 512, vocabsize: int = 50257):
self.size = size
self.seqlength = seqlength
self.vocabsize = vocabsize
self.data = torch.randint(0, vocabsize, (size, seqlength))
def len(self):
return self.size
def getitem(self, idx):
return {
‘input_ids’: self.data[idx],
‘labels’: self.data[idx].clone()
}
“`

Продвинутый тренер DeepSpeed

Мы создаём продвинутого тренера DeepSpeed с использованием нескольких методов оптимизации. Тренер создаёт модель GPT-2 для демонстрации, устанавливает конфигурацию DeepSpeed (ZeRO, FP16, AdamW, warmup scheduler, tensorboard) и инициализирует движок. Затем выполняются эффективные шаги обучения с логированием и статистикой памяти, сохраняются контрольные точки и демонстрируется вывод для проверки оптимизации и генерации в одном месте.

“`python
class AdvancedDeepSpeedTrainer:
“””Advanced DeepSpeed trainer with multiple optimization techniques”””
def init(self, modelconfig: Dict[str, Any], dsconfig: Dict[str, Any]):
self.modelconfig = modelconfig
self.dsconfig = dsconfig
self.model = None
self.engine = None
self.tokenizer = None
def create_model(self):
“””Create a GPT-2 style model for demonstration”””
print(” Creating model…”)
config = GPT2Config(
vocabsize=self.modelconfig[‘vocab_size’],
npositions=self.modelconfig[‘seq_length’],
nembd=self.modelconfig[‘hidden_size’],
nlayer=self.modelconfig[‘num_layers’],
nhead=self.modelconfig[‘num_heads’],
resid_pdrop=0.1,
embd_pdrop=0.1,
attn_pdrop=0.1,
)
self.model = GPT2LMHeadModel(config)
self.tokenizer = GPT2Tokenizer.from_pretrained(‘gpt2’)
self.tokenizer.padtoken = self.tokenizer.eostoken
print(f” Model parameters: {sum(p.numel() for p in self.model.parameters()):,}”)
return self.model
“`

Генерация конфигураций DeepSpeed

Мы генерируем повторно используемые конфигурации DeepSpeed, сравниваем стадии ZeRO для сравнения памяти и скорости и демонстрируем дополнительные функции, такие как динамическая масштабирование потерь и параллелизм конвейера/MoE.

“`python
def demonstratezerostages():
“””Demonstrate different ZeRO optimization stages”””
print(“\n ZeRO Optimization Stages Explained:”)
print(” Stage 0: Disabled (baseline)”)
print(” Stage 1: Optimizer state partitioning (~4x memory reduction)”)
print(” Stage 2: Gradient partitioning (~8x memory reduction)”)
print(” Stage 3: Parameter partitioning (~Nx memory reduction)”)
zero_configs = {
1: {“stage”: 1, “reducebucketsize”: 5e8},
2: {“stage”: 2, “allgatherpartitions”: True, “reducescatter”: True},
3: {“stage”: 3, “stage3prefetchbucketsize”: 5e8, “stage3parampersistencethreshold”: 1e6}
}
for stage, config in zero_configs.items():
estimatedmemoryreduction = [1, 4, 8, “Nx”][stage]
print(f” Stage {stage}: ~{estimatedmemoryreduction}x memory reduction”)
“`

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

1. Какие методы оптимизации используются в продвинутом тренере DeepSpeed для эффективного обучения больших языковых моделей?

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

2. Какие библиотеки и инструменты необходимы для установки зависимостей перед началом работы с DeepSpeed?

Для установки зависимостей необходимо установить PyTorch с поддержкой CUDA, DeepSpeed, а также библиотеки Transformers, Datasets, Accelerate и Weights & Biases.

3. Какие преимущества даёт использование синтетического набора данных при тестировании обучения DeepSpeed?

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

4. Какие стадии оптимизации ZeRO существуют и как они влияют на память и скорость обучения?

Существуют четыре стадии оптимизации ZeRO: Stage 0 (Disabled, baseline), Stage 1 (Optimizer state partitioning, ~4x memory reduction), Stage 2 (Gradient partitioning, ~8x memory reduction) и Stage 3 (Parameter partitioning, ~Nx memory reduction). Каждая стадия уменьшает использование памяти и может влиять на скорость обучения.

5. Какие дополнительные функции DeepSpeed демонстрируются в статье помимо основных методов оптимизации?

В статье демонстрируются дополнительные функции DeepSpeed, такие как динамическое масштабирование потерь и параллелизм конвейера/MoE.

Источник