train_model_parallel.py

import time

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from src.data import RandomTextDataset
from src.runtime import require_cuda

MAX_STEPS = 100


class ManualModelParallelTransformer(nn.Module):
    def __init__(
        self,
        vocab_size: int = 32000,
        seq_len: int = 256,
        d_model: int = 512,
        nhead: int = 8,
        num_classes: int = 10,
        dim_feedforward: int = 2048,
    ):
        super().__init__()

        self.gpu0 = torch.device("cuda:0")
        self.gpu1 = torch.device("cuda:1")
        self.gpu2 = torch.device("cuda:2")

        self.seq_len = seq_len

        self.token_emb = nn.Embedding(vocab_size, d_model).to(self.gpu0)
        self.pos_emb = nn.Embedding(seq_len, d_model).to(self.gpu0)

        def make_layer():
            return nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=dim_feedforward,
                batch_first=True,
                activation="gelu",
            )

        self.block0 = nn.Sequential(make_layer(), make_layer()).to(self.gpu0)
        self.block1 = nn.Sequential(make_layer(), make_layer()).to(self.gpu1)
        self.block2 = nn.Sequential(make_layer(), make_layer()).to(self.gpu2)

        self.norm = nn.LayerNorm(d_model).to(self.gpu2)
        self.head = nn.Linear(d_model, num_classes).to(self.gpu2)

    def forward(self, input_ids):

        batch_size, seq_len = input_ids.shape

        positions = torch.arange(seq_len, device=self.gpu0)
        positions = positions.unsqueeze(0).expand(batch_size, seq_len)

        x = self.token_emb(input_ids) + self.pos_emb(positions)
        x = self.block0(x)

        x = x.to(self.gpu1)
        x = self.block1(x)

        x = x.to(self.gpu2)
        x = self.block2(x)

        x = self.norm(x[:, 0])
        logits = self.head(x)

        return logits


def main():
    require_cuda(min_devices=3)

    dataset = RandomTextDataset(num_samples=5000)
    loader = DataLoader(
        dataset,
        batch_size=64,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )

    model = ManualModelParallelTransformer()

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss().to("cuda:2")

    model.train()
    start = time.time()

    for step, (input_ids, labels) in enumerate(loader):
        if step >= MAX_STEPS:
            break

        input_ids = input_ids.to("cuda:0", non_blocking=True)
        labels = labels.to("cuda:2", non_blocking=True)

        optimizer.zero_grad(set_to_none=True)

        logits = model(input_ids)
        loss = criterion(logits, labels)

        loss.backward()
        optimizer.step()

        if step % 20 == 0:
            elapsed = time.time() - start
            print(f"step={step}, loss={loss.item():.4f}, elapsed={elapsed:.2f}s")


if __name__ == "__main__":
    main()