train_lm.py

if __name__ == "__main__":
    import argparse
    from pathlib import Path

    import msgpack
    import numpy as np
    from loguru import logger

    from cs336_basics import (
        AdamW,
        Tokenizer,
        TransformerLM,
        clip_grad_norm_,
        cross_entropy,
        get_batch,
        learning_rate_schedule,
        save_checkpoint,
    )

    parser = argparse.ArgumentParser()
    parser.add_argument("--tokenizer_dir", type=Path, required=True)
    parser.add_argument("--corpus", type=Path, required=True)
    parser.add_argument("--batch_size", default=64, type=int)
    parser.add_argument("--device", default="cuda", choices=["cpu", "cuda", "mps"], type=str)
    parser.add_argument("--context_length", default=128, type=int)
    parser.add_argument("--d_model", default=64, type=int)
    parser.add_argument("--num_layers", default=2, type=int)
    parser.add_argument("--num_heads", default=4, type=int)
    parser.add_argument("--d_ff", default=256, type=int)
    parser.add_argument("--rope_theta", default=10_000.0, type=float)
    args = parser.parse_args()

    logger.info(f"{vars(args)=}")

    with (args.tokenizer_dir / "vocab.msgpack").open("rb") as f:
        vocab = msgpack.unpack(f, strict_map_key=False)

    with (args.tokenizer_dir / "merges.msgpack").open("rb") as f:
        merges = [tuple(pair) for pair in msgpack.unpack(f)]

    with (args.tokenizer_dir / "special_tokens.msgpack").open("rb") as f:
        special_tokens = msgpack.unpack(f)

    tokenizer = Tokenizer(vocab=vocab, merges=merges, special_tokens=special_tokens)

    with open(args.corpus) as f:
        corpus = f.read()

    clm = TransformerLM(
        vocab_size=len(vocab),
        context_length=args.context_length,
        d_model=args.d_model,
        num_layers=args.num_layers,
        num_heads=args.num_heads,
        d_ff=args.d_ff,
        rope_theta=args.rope_theta,
    ).to(args.device)

    optimizer = AdamW(clm.parameters())

    dataset_np = np.array(tokenizer.encode(corpus))

    total_it = 200
    for it in range(1, total_it + 1):  # start from 1 to prevent 0.0 learning rate returned by scheduler
        # forward
        x, label = get_batch(
            dataset=dataset_np,
            batch_size=args.batch_size,
            context_length=args.context_length,
            device=args.device,
        )
        logits = clm(x)
        loss = cross_entropy(inputs=logits, targets=label)
        print(f"Loss: {loss.detach().cpu().item()}")

        # backward
        loss.backward()

        # clip grad
        clip_grad_norm_(clm.parameters(), max_l2_norm=1.0)

        lr = learning_rate_schedule(
            it=it,
            max_learning_rate=3e-4,
            min_learning_rate=3e-5,
            warmup_iters=int(total_it * 0.1),
            cosine_cycle_iters=total_it,
        )

        for group in optimizer.param_groups:
            group["lr"] = lr

        # update
        optimizer.step()
        optimizer.zero_grad()

        # save
        if it != 0 and it % 1_000 == 0:
            save_dir = f"./save/it{it:08d}.pt"
            save_checkpoint(
                model=clm,
                optimizer=optimizer,
                iteration=it,
                out=save_dir,
            )
            logger.info(f"Saved checkpoint to {save_dir}")