train_tensor_parallel.py

import os
import time
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.distributed as dist

from torch.utils.data import DataLoader

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
)

from src.data import RandomTextDataset


@dataclass
class Config:
    vocab_size: int = 32000
    seq_len: int = 256
    d_model: int = 768
    hidden_dim: int = 3072
    num_layers: int = 6
    num_classes: int = 10

    num_samples: int = 10000

    # IMPORTANT FOR PURE TENSOR PARALLELISM:
    # This is the global batch processed by the tensor-parallel group.
    # All TP ranks receive the same batch.
    #
    # It is NOT per-GPU batch size like DDP.
    batch_size: int = 32

    num_workers: int = 4
    num_epochs: int = 1
    max_steps: int = 100
    lr: float = 1e-4


class TensorParallelMLPBlock(nn.Module):
    """
    Transformer-style feed-forward block:

        x -> Linear(d_model, hidden_dim)
          -> GELU
          -> Linear(hidden_dim, d_model)
          -> residual + LayerNorm

    fc1 is column-parallel:
        output hidden dimension is split across GPUs.

    fc2 is row-parallel:
        input hidden dimension is split across GPUs.
    """

    def __init__(self, d_model: int, hidden_dim: int):
        super().__init__()

        self.ln = nn.LayerNorm(d_model)
        self.fc1 = nn.Linear(d_model, hidden_dim, bias=False)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_dim, d_model, bias=False)

    def forward(self, x):
        residual = x

        x = self.ln(x)
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)

        return x + residual


class TensorParallelToyTransformer(nn.Module):
    """
    Toy Transformer-like classifier.

    For learning tensor parallelism, we keep attention out and focus on MLP blocks.

    Real LLM tensor parallelism commonly shards:
        - attention Q/K/V/O projections
        - MLP up/gate/down projections
    """

    def __init__(self, cfg: Config):
        super().__init__()

        self.cfg = cfg

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

        self.layers = nn.ModuleList(
            [
                TensorParallelMLPBlock(
                    d_model=cfg.d_model,
                    hidden_dim=cfg.hidden_dim,
                )
                for _ in range(cfg.num_layers)
            ]
        )

        self.final_ln = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.num_classes)

    def forward(self, input_ids):
        batch_size, seq_len = input_ids.shape

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

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

        for layer in self.layers:
            x = layer(x)

        # CLS-style pooling
        x = self.final_ln(x[:, 0])
        logits = self.head(x)

        return logits


def setup_distributed():
    dist.init_process_group(backend="nccl")

    local_rank = int(os.environ["LOCAL_RANK"])
    global_rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    return local_rank, global_rank, world_size, device


def cleanup_distributed():
    dist.destroy_process_group()


def is_main_process():
    # Gives global rank
    return dist.get_rank() == 0


def parallelize_model(model: nn.Module, world_size: int):
    """
    Apply tensor parallelism to each MLP block.

    IMPORTANT:
    hidden_dim must be divisible by world_size.

    With 3 GPUs:
        hidden_dim = 3072
        each rank gets 1024 hidden features
    """

    cfg = model.cfg

    if cfg.hidden_dim % world_size != 0:
        raise ValueError(
            f"hidden_dim={cfg.hidden_dim} must be divisible by world_size={world_size}"
        )

    # Creates a 1D tensor-parallel mesh over CUDA devices.
    #
    # With:
    #     torchrun --standalone --nproc_per_node=3 train_tensor_parallel.py
    #
    # mesh is conceptually:
    #     [cuda:0, cuda:1, cuda:2]
    tp_mesh = init_device_mesh(
        device_type="cuda",
        mesh_shape=(world_size,),
    )

    for layer_id, layer in enumerate(model.layers):
        parallelize_module(
            module=layer,
            device_mesh=tp_mesh,
            parallelize_plan={
                # fc1 weight shape:
                #     [hidden_dim, d_model]
                #
                # ColwiseParallel splits output features.
                #
                # Conceptually:
                #     GPU 0 computes first 1/3 of hidden features
                #     GPU 1 computes second 1/3
                #     GPU 2 computes third 1/3
                "fc1": ColwiseParallel(),

                # fc2 weight shape:
                #     [d_model, hidden_dim]
                #
                # RowwiseParallel splits input features.
                #
                # Each GPU consumes its hidden-feature shard,
                # then outputs are reduced/combined.
                "fc2": RowwiseParallel(),
            },
        )

        if is_main_process():
            print(f"Tensor parallelized MLP layer {layer_id}")

    return model


def print_parameter_summary(model: nn.Module):
    if not is_main_process():
        return

    print("\nParameter summary:")
    for name, param in model.named_parameters():
        print(f"{name:40s} shape={tuple(param.shape)} type={type(param)}")
    print()


def broadcast_batch_from_rank0(
    input_ids,
    labels,
    cfg: Config,
    device: torch.device,
):
    """
    Pure tensor parallelism rule:

        All TP ranks must process the SAME batch.

    Rank 0 loads the batch from the DataLoader.
    Then rank 0 broadcasts input_ids and labels to every other TP rank.

    This is correct for pure TP learning/demo code.

    In real large-scale training, TP is usually combined with DP:
        - ranks inside same TP group get the same batch
        - different DP groups get different batches
    """

    if is_main_process():
        input_ids = input_ids.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
    else:
        input_ids = torch.empty(
            (cfg.batch_size, cfg.seq_len),
            dtype=torch.long,
            device=device,
        )
        labels = torch.empty(
            (cfg.batch_size,),
            dtype=torch.long,
            device=device,
        )

    dist.broadcast(input_ids, src=0)
    dist.broadcast(labels, src=0)

    return input_ids, labels


def main():
    cfg = Config()

    local_rank, global_rank, world_size, device = setup_distributed()

    # IMPORTANT:
    # For pure tensor parallelism, all ranks are cooperating on one logical model.
    # So model initialization should be consistent across ranks.
    #
    # Do NOT use:
    #     torch.manual_seed(1234 + global_rank)
    #
    # for pure TP model initialization.
    torch.manual_seed(1234)

    if is_main_process():
        print("Pure Tensor Parallel training started")
        print(f"world_size / TP size: {world_size}")
        print(f"global TP batch size: {cfg.batch_size}")
        print("effective global batch size is NOT batch_size * world_size for pure TP")
        print(f"d_model: {cfg.d_model}")
        print(f"hidden_dim: {cfg.hidden_dim}")
        print(f"hidden_dim per GPU: {cfg.hidden_dim // world_size}")

    # Only rank 0 needs to load/generate the batch.
    #
    # Other ranks receive the same batch through dist.broadcast().
    if is_main_process():
        dataset = RandomTextDataset(
            num_samples=cfg.num_samples,
            seq_len=cfg.seq_len,
            vocab_size=cfg.vocab_size,
            num_classes=cfg.num_classes,
        )

        loader = DataLoader(
            dataset,
            batch_size=cfg.batch_size,
            shuffle=True,
            num_workers=cfg.num_workers,
            pin_memory=True,
            drop_last=True,
        )
    else:
        loader = None

    # Build model on each rank's GPU.
    model = TensorParallelToyTransformer(cfg).to(device)

    # Apply tensor parallelism to selected submodules.
    model = parallelize_model(model, world_size)

    print_parameter_summary(model)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=cfg.lr,
        foreach=False,
    )
    criterion = nn.CrossEntropyLoss()

    model.train()

    start = time.time()
    global_step = 0

    for epoch in range(cfg.num_epochs):
        if is_main_process():
            data_iter = iter(loader)

        # Since only rank 0 has a DataLoader, all ranks follow the same
        # step count controlled by cfg.max_steps.
        while global_step < cfg.max_steps:
            if is_main_process():
                try:
                    input_ids, labels = next(data_iter)
                except StopIteration:
                    # Restart loader if max_steps is larger than one epoch.
                    data_iter = iter(loader)
                    input_ids, labels = next(data_iter)
            else:
                input_ids, labels = None, None

            input_ids, labels = broadcast_batch_from_rank0(
                input_ids=input_ids,
                labels=labels,
                cfg=cfg,
                device=device,
            )

            optimizer.zero_grad(set_to_none=True)

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

            loss.backward()
            optimizer.step()

            if is_main_process() and global_step % 10 == 0:
                elapsed = time.time() - start

                allocated = torch.cuda.memory_allocated(device) / 1024**3
                reserved = torch.cuda.memory_reserved(device) / 1024**3

                print(
                    f"step={global_step}, "
                    f"loss={loss.item():.4f}, "
                    f"gpu_mem_allocated={allocated:.2f} GB, "
                    f"gpu_mem_reserved={reserved:.2f} GB, "
                    f"elapsed={elapsed:.2f}s"
                )

            global_step += 1

        if global_step >= cfg.max_steps:
            break

    dist.barrier()

    if is_main_process():
        print("Pure Tensor Parallel training finished.")

    cleanup_distributed()


if __name__ == "__main__":
    main()