train.py

"""
Run with: torchrun --nproc_per_node=2 train.py
"""

import os

import torch
import torch.distributed as dist

from transformers import LlamaConfig, LlamaForCausalLM

from src.fsdp import fully_shard, FSDPModule


def main():
    # Get distributed settings from environment (set by torchrun)
    local_rank = int(os.environ.get("LOCAL_RANK"))
    rank = int(os.environ.get("RANK"))
    world_size = int(os.environ.get("WORLD_SIZE"))

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

    dist.init_process_group(backend, rank=rank, world_size=world_size, device_id=device)

    print(
        f"[Rank {rank}] Local rank: {local_rank}, Global rank: {rank}, World size: {world_size}"
    )

    torch.manual_seed(42)

    config = LlamaConfig()
    model = LlamaForCausalLM(config)

    base_model = model.model
    for layer in base_model.layers:
        fully_shard(layer)
    fully_shard(base_model.embed_tokens)
    fully_shard(base_model.norm)
    fully_shard(base_model.rotary_emb)
    fully_shard(model.lm_head)

    model.to(device=device)
    optim = torch.optim.Adam(model.parameters(), lr=1e-3)

    print(f"✓ Training started")
    for step in range(5):
        x = torch.randint(0, config.vocab_size, (1, 512), device=device)

        optim.zero_grad()

        output = model(x).logits
        loss = output.sum()
        loss.backward()

        optim.step()

    print(f"✓ Training completed, final output shape: {output.shape}")

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()