v4.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.checkpoint import checkpoint

# ---- config ----
N, D, H, C, L = 4096, 512, 1024, 10, 12
batch_size = 64
accum_steps = 4 #&
epochs = 5
lr = 1e-2
device = "cuda" if torch.cuda.is_available() else "cpu"
amp_dtype = torch.float16 # or torch.bfloat16 if your device support it

# ---- data ----
class DemoDataset(Dataset):
    def __init__(self, n, d, c):
        """
        this method is mostly for setting the path to data folder, split IDs, augemntations, ...

        reading data (images, text files, ...) can happen here if data is small enough to comfortably fit in RAM (an when using few workers)
        each worker will create an instance of this class, meaning it will run this method and create / allocate memory for whatever its implementation do
        """
        self.x = torch.randn(n, d)
        self.y = torch.randint(0, c, (n,))

    # this is used to determin how many batches / samples in an epoch
    def __len__(self):
        return self.x.size(0)

    # this is what's called by a single worker to callect samples for the batch
    def __getitem__(self, idx):
        """
        it's most common and recommended to read / preprocess a data sample in here on the fly
        """
        return self.x[idx], self.y[idx]


dataset = DemoDataset(N, D, C)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# ---- model ----
class MLPBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fc = nn.Linear(dim, dim)

    def forward(self, x):
        return x + torch.relu(self.fc(self.norm(x)))

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, n_layers):
        super().__init__()
        self.in_proj = nn.Linear(in_dim, hidden_dim)

        self.blocks = nn.ModuleList(
            [MLPBlock(hidden_dim) for _ in range(n_layers - 1)]
        )

        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.in_proj(x) 

        for block in self.blocks:
            
            # checkpoint each whole block
            x = checkpoint(block, x, use_reentrant=False) 
            # the "use_reentrant=False" is just an implementation detail that favor a modern implementation over a legacy, potentially problematic one

        return self.out(x)

model = MLP(D, H, C, L).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
loss_fn = nn.CrossEntropyLoss()
scaler = torch.amp.GradScaler(device, enabled=(amp_dtype == torch.float16))

# ---- train ----
model.train()
for epoch in range(epochs):
    total_loss, seen = 0.0, 0
    for step, (xb, yb) in enumerate(loader):
        xb, yb = xb.to(device), yb.to(device)


        with torch.autocast(device_type=device, dtype=amp_dtype):
            logits = model(xb)
            loss = loss_fn(logits, yb) / accum_steps #&
        
        # we still run the forward and backward passes each step
        scaler.scale(loss).backward()

        total_loss += loss.item() * accum_steps * xb.size(0)
        seen += xb.size(0)


        # we only compute the grad from the (accum) computed grads ones per "accum_steps"
        # and lastly we zero out existing grads to start computing the next (batch_size * accum_steps) grads 
        if (step + 1) % accum_steps == 0: #&
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

    print(f"epoch {epoch}: loss {total_loss / seen:.4f}")