standard_transformers.py

# see https://github.com/pytorch/torchsnapshot/blob/main/benchmarks/fsdp/main.py
import torch
from torch import distributed as dist, nn

def create_model() -> nn.Module:
    # 7.8GB model, 1.9B parameters
    model = nn.Transformer(
        d_model=864,
        num_encoder_layers=1,
        num_decoder_layers=20,
        nhead=12,
        dim_feedforward=50257,
    )

    # 80GB 21B parameters
    # model = nn.Transformer(
    #     d_model=4000,
    #     num_encoder_layers=1,
    #     num_decoder_layers=40,
    #     nhead=40,
    #     dim_feedforward=50257,
    # )