⚙️ Требования

• Для работы по текущему гайду необходим драйвер NVIDIA, поддерживающий CUDA 12.4.

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

  /media/Data1/common/docker/deepspeed
  

  и все манипуляции будут происходить относительно него.

---

🐳 Docker и скрипты

Создаем Dockerfile и соответствующие скрипты для работы с образом.

---

🌐 Hostfile для DeepSpeed

Для работы DeepSpeed на нескольких машинах необходимо создать hostfile и положить его в volume папку (чтобы был доступ из контейнеров):

mlnode1_ds slots=1 env="PDSH_RCMD_TYPE=ssh,NCCL_DEBUG=INFO,NCCL_CROSS_NIC=0"
mlnode2_ds slots=1 env="PDSH_RCMD_TYPE=ssh,NCCL_DEBUG=INFO,NCCL_CROSS_NIC=0"

Здесь:

• mlnode1_ds и mlnode2_ds — это хосты, определённые в нашем ssh config.
• slots=1 — число GPU, выделяемых для этой ноды.
• env=... — дополнительные переменные для DeepSpeed и NCCL.

---

🚀 Запуск

При запуске мы должны передавать:

• hostfile
• файл скрипта (например, train.py)
• и deepspeed_config.json

Пример скрипта запуска — run.sh:

#!/bin/bash
deepspeed --hostfile hostfile train.py --deepspeed --deepspeed_config deepspeed_config.json

---

🛠 Рекомендуемые настройки перед запуском

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

---

🧪 Примеры для теста

Ниже приведены два минимальных примера для проверки работоспособности кластера:

• Легковесный (batch_size=4, hidden_dim=16).
• Потяжелее (batch_size=64, hidden_dim≈1024).

---

⚡ Пример 1: Лёгкий тест

deepspeed_config.json

{
  "train_batch_size": 4,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": false
  },
  "zero_optimization": {
    "stage": 0
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001
    }
  }
}

train.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset
import deepspeed
import argparse

class RandomDataset(Dataset):
    def __init__(self, num_samples=16, in_dim=10):
        self.x = torch.randn(num_samples, in_dim)
        self.y = torch.randn(num_samples, 1)
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]
    def __len__(self):
        return len(self.x)

class SimpleModel(nn.Module):
    def __init__(self, hidden_dim=16):
        super().__init__()
        self.linear = nn.Linear(10, hidden_dim)
        self.relu = nn.ReLU()
        self.output = nn.Linear(hidden_dim, 1)
    def forward(self, x):
        return self.output(self.relu(self.linear(x)))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=-1)
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()

    model = SimpleModel()
    model, _, trainloader, _ = deepspeed.initialize(
        args=args,
        model=model,
        training_data=RandomDataset(num_samples=20)
    )

    for epoch in range(10):
        for x, y in trainloader:
            inputs = x.to(model.device)
            targets = y.to(model.device)
            outputs = model(inputs)
            loss = nn.MSELoss()(outputs, targets)
            model.backward(loss)
            model.step()
        if model.global_rank == 0:
            print(f"Epoch {epoch:02d} - Loss: {loss.item():.4f}")

if __name__ == "__main__":
    main()

---

🚀 Пример 2: Более тяжёлый тест

deepspeed_config.json

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 32,
  "steps_per_print": 10,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 1e-3
    }
  },
  "fp16": {
    "enabled": false
  }
}

train.py

#!/usr/bin/env python3
import torch, torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import deepspeed
import argparse

class RandomDataset(Dataset):
    def __init__(self, num_samples=1024, in_dim=128):
        self.x = torch.randn(num_samples, in_dim)
        self.y = torch.randint(0, 2, (num_samples,))
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]
    def __len__(self):
        return len(self.x)

class TinyNet(nn.Module):
    def __init__(self, in_dim=128, n_classes=2):
        super().__init__()
        hidden_dim = 1024
        self.layers = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            *[nn.Linear(hidden_dim, hidden_dim) for _ in range(8)],
            nn.Linear(hidden_dim, n_classes)
        )
    def forward(self, x):
        return self.layers(x)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_epochs", type=int, default=30)
    parser.add_argument('--local_rank', type=int, default=-1)
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()

    model_engine, _, trainloader, _ = deepspeed.initialize(
        args=args,
        model=TinyNet(),
        training_data=RandomDataset(num_samples=2048)
    )

    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(args.num_epochs):
        for x, y in trainloader:
            x, y = x.to(model_engine.device), y.to(model_engine.device)
            y_hat = model_engine(x)
            loss = loss_fn(y_hat, y)
            model_engine.backward(loss)
            model_engine.step()
        if model_engine.global_rank == 0:
            print(f"Epoch {epoch:02d}: loss={loss.item():.4f}")

---