agk4444/bdh-shakespeare

BDH (Dragon Hatchling) — tiny Shakespeare

Pathway's BDH ("Dragon Hatchling", arXiv:2509.26507) trained from scratch on the tiny-Shakespeare corpus — the same toy task the official repo ships, but run to completion on GPU. Code: github.com/pathwaycom/bdh.

This is a toy model for studying the BDH architecture, not a general text model.

Model

• Architecture: BDH at the repo-default size — 25.3M parameters, byte-level (vocab 256), with BDH's wide latent state (8,192-wide per head)
• Config: n_layer=6, n_embd=256, n_head=4, vocab_size=256, block size 512
• Corpus: tiny-Shakespeare (karpathy/char-rnn input.txt)

Training

• 3,000 steps (the repo-default full toy run) on a Kaggle T4 GPU, 2026-10-23
• Micro-batch 8 with 4 gradient-accumulation steps (effective batch 32 — a flat batch of 32 OOMs on a 16 GB T4 because of the wide BDH latent activations)
• lr 1e-3, weight decay 0.1, bf16/fp16 autocast, seed 1337
• Loss: 5.650 → 0.1147 over 40,068s (~11.1h)
• Samples at steps 1000/2000/3000 get progressively more coherent as Shakespeare pastiche

Checkpoint

bdh_tinyshakespeare.pt holds the trained weights plus its config:

{"model": state_dict, "config": {...}, "iters": 3000, "final_loss": 0.1147}

Usage

Clone the official repo first so bdh is importable, then (from the training notebook's reload check):

import torch
import bdh  # from github.com/pathwaycom/bdh

ckpt = torch.load("bdh_tinyshakespeare.pt", map_location="cuda", weights_only=True)
model = bdh.BDH(bdh.BDHConfig(**ckpt["config"])).to("cuda")
model.load_state_dict(ckpt["model"])
model.eval()

prompt = torch.tensor(bytearray("To be or ", "utf-8"), dtype=torch.long).unsqueeze(0).to("cuda")
out = model.generate(prompt, max_new_tokens=300, top_k=3)
print(bytes(out.to(torch.uint8).cpu().squeeze(0)).decode(errors="backslashreplace"))

Limitations

• Trained only on tiny-Shakespeare: output is Shakespeare-style pastiche and nothing else. Expect repetition and nonsense over long generations.
• No instruction tuning, no safety tuning, no real-world knowledge.