deepspeed_setup/README.md

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

* Для работы по текущему гайду необходим **драйвер NVIDIA**, поддерживающий **CUDA 12.4**.
* Кроме того, для работы DeepSpeed необходимо, чтобы **пути к файлам совпадали на всех машинах**.
  В данном гайде основной путь будет:

  ```
  /media/Data1/common/docker/deepspeed
  ```

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

---

## 🐳 Docker и скрипты

[Создаем Dockerfile и соответствующие скрипты для работы с образом](Dockerfile_and_scripts.md). Однако перед эти убеждаемся, что рядом с нашим Dockerfile'ом лежит `dsync.py` (Если нет, то копируем из папки `dsync`, также можем ознакомиться с [инструкцией](dsync/dsync.md)).

---

## 🌐 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`:

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

---

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

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

---

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

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

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

---

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

#### `deepspeed_config.json`

```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`

```python
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`

```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`

```python
#!/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}")
```

---