train.py

import gzip
import random

import numpy as np
import torch
import tqdm
from torch.utils.data import DataLoader, Dataset
from zeta.optim import StableAdamWUnfused

from zeta.structs import AutoregressiveWrapper
from griffin_torch import Griffin

# constants

NUM_BATCHES = int(1e5)
BATCH_SIZE = 4
GRADIENT_ACCUMULATE_EVERY = 4
LEARNING_RATE = 2e-4
VALIDATE_EVERY = 100
GENERATE_EVERY = 500
GENERATE_LENGTH = 512
SEQ_LEN = 1024


# helpers
def cycle(loader):
    while True:
        yield from loader


def decode_token(token):
    return str(chr(max(32, token)))


def decode_tokens(tokens):
    return "".join(list(map(decode_token, tokens)))


# instantiate GPT-like decoder model
model = Griffin(
    num_tokens=256,
    dim=512,
    depth=8,
    seq_len=SEQ_LEN,
    mlp_mult=4,
    heads=8,
    dropout=0.1,
)
model = AutoregressiveWrapper(model, speculative=True)


# Use all available GPUs
if torch.cuda.device_count() > 1:
    print("Using", torch.cuda.device_count(), "GPUs!")
    model = torch.nn.DataParallel(model)


model.cuda()

# prepare enwik8 data
with gzip.open("./data/enwik8.gz") as file:
    X = np.fromstring(file.read(int(95e6)), dtype=np.uint8)
    trX, vaX = np.split(X, [int(90e6)])
    data_train, data_val = torch.from_numpy(trX), torch.from_numpy(
        vaX
    )


class TextSamplerDataset(Dataset):
    def __init__(self, data, seq_len):
        super().__init__()
        self.data = data
        self.seq_len = seq_len

    def __getitem__(self, index):
        rand_start = torch.randint(
            0, self.data.size(0) - self.seq_len, (1,)
        )
        full_seq = self.data[
            rand_start : rand_start + self.seq_len + 1
        ].long()
        return full_seq.cuda()

    def __len__(self):
        return self.data.size(0) // self.seq_len


train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
val_dataset = TextSamplerDataset(data_val, SEQ_LEN)
train_loader = cycle(DataLoader(train_dataset, batch_size=BATCH_SIZE))
val_loader = cycle(DataLoader(val_dataset, batch_size=BATCH_SIZE))

# optimizer
optim = StableAdamWUnfused(
    model.parameters(),
    lr=LEARNING_RATE,
)

# training
for i in tqdm.tqdm(
    range(NUM_BATCHES), mininterval=10.0, desc="training"
):
    model.train()

    for __ in range(GRADIENT_ACCUMULATE_EVERY):
        loss = model(next(train_loader))
        loss.mean().backward()

    print(f"training loss: {loss.mean().item()}")
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optim.step()
    optim.zero_grad()

    if i % VALIDATE_EVERY == 0:
        model.eval()
        with torch.no_grad():
            loss = model(next(val_loader))
            print(f"validation loss: {loss.mean().item()}")

    if i % GENERATE_EVERY == 0:
        model.eval()
        inp = random.choice(val_dataset)[:-1]
        prime = decode_tokens(inp)
        print("%s \n\n %s", (prime, "*" * 100))

        sample = model.module.generate(
            inp[None, ...], GENERATE_LENGTH
        )
        output_str = decode_tokens(sample[0])
        print(output_str)