Portable train.py: CPU/GPU/Colab support, env-var budgets

f3fdcef verified about 1 month ago

22.8 kB

"""
tinystories-10m : a ~10M parameter GPT-2 trained from scratch on TinyStories.
One script, end to end:
  1. download the TinyStories train/valid raw text
  2. train a byte-level BPE tokenizer (vocab 8192) on a sample of the corpus
  3. tokenize the corpus into flat uint16 arrays
  4. train a GPT-2 (256 dim / 10 layers / 8 heads / tied embeddings ~ 10.06M params)
  5. evaluate val loss + bits-per-byte, and compare against the published TinyStories baselines
  6. sample generations, plot the loss curve, push everything to the Hub
Runs on CUDA (fast path, bfloat16) or on CPU (slow path, fp32) with no code changes.
Environment variables (all optional):
  TS_MAX_STEPS     training steps, default 12000      (12000 x 65536 tok = ~786M tokens)
  TS_MAX_MINUTES   wall-clock cap on training, default 110
  TS_DEVICE        'cuda' or 'cpu', default auto-detect
  TS_WORK          scratch directory, default /tmp/tswork
  HF_TOKEN         needed only for pushing to the Hub; skips upload if unset
Examples
  HF Jobs (GPU):
    hf jobs uv run --flavor a10g-small --timeout 3h --secrets HF_TOKEN \\
      -w torch -w transformers -w tokenizers -w matplotlib -w huggingface_hub \\
      -c "from huggingface_hub import hf_hub_download; import runpy;
          runpy.run_path(hf_hub_download('kunal10tripathi/tinystories-10m','train.py'), run_name='__main__')"
  Quick CPU smoke run (roughly 40M tokens):
    TS_DEVICE=cpu TS_MAX_STEPS=600 TS_MAX_MINUTES=90 python train.py
"""

import os, json, math, time, shutil, contextlib
import numpy as np
import torch
import torch.nn.functional as F
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
from transformers import (
    GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast,
    AutoTokenizer, AutoModelForCausalLM,
)
from huggingface_hub import hf_hub_download, HfApi

# ----------------------------------------------------------------------------- config
REPO_ID = 'kunal10tripathi/tinystories-10m'
DATASET = 'roneneldan/TinyStories'
EOS = '<|endoftext|>'
WORK = os.environ.get('TS_WORK', '/tmp/tswork')
os.makedirs(WORK, exist_ok=True)
T0 = time.time()


def log(*a):
    msg = ' '.join(str(x) for x in a)
    print('[' + str(round(time.time() - T0, 1)).rjust(8) + 's] ' + msg, flush=True)


HF_TOKEN = os.environ.get('HF_TOKEN')
api = HfApi(token=HF_TOKEN)
if HF_TOKEN:
    try:
        log('logged in as', api.whoami()['name'])
    except Exception as e:
        log('WARN whoami failed:', repr(e))
else:
    log('no HF_TOKEN set: results will be kept local and not uploaded')

DEVICE = os.environ.get('TS_DEVICE') or ('cuda' if torch.cuda.is_available() else 'cpu')
IS_CUDA = DEVICE.startswith('cuda')
log('device', DEVICE, torch.cuda.get_device_name(0) if IS_CUDA else '')


def amp():
    """bfloat16 autocast on GPU, plain fp32 on CPU."""
    if IS_CUDA:
        return torch.autocast('cuda', dtype=torch.bfloat16)
    return contextlib.nullcontext()


# ----------------------------------------------------------------------------- data
log('downloading TinyStories...')
train_txt = hf_hub_download(DATASET, 'TinyStories-train.txt', repo_type='dataset')
valid_txt = hf_hub_download(DATASET, 'TinyStories-valid.txt', repo_type='dataset')
log('train bytes', os.path.getsize(train_txt), '| valid bytes', os.path.getsize(valid_txt))

# ----------------------------------------------------------------------------- tokenizer
TOK_PATH = os.path.join(WORK, 'tokenizer.json')
if os.path.exists(TOK_PATH):
    tok = Tokenizer.from_file(TOK_PATH)
else:
    def sample_text():
        n = 0
        with open(train_txt, 'r', encoding='utf-8', errors='ignore') as f:
            while n < 80_000_000:
                c = f.read(1_000_000)
                if not c:
                    break
                n += len(c)
                yield c

    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    tr = trainers.BpeTrainer(
        vocab_size=8192, min_frequency=2, special_tokens=[EOS],
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False,
    )
    log('training BPE tokenizer (vocab 8192) on an 80MB sample...')
    tok.train_from_iterator(sample_text(), trainer=tr)
    tok.save(TOK_PATH)

VOCAB = tok.get_vocab_size()
EOS_ID = tok.token_to_id(EOS)
log('tokenizer ready | vocab', VOCAB, '| eos_id', EOS_ID)


def tokenize_to_bin(in_path, out_path):
    total = os.path.getsize(in_path)
    read = last = n_tok = 0
    buf = ''
    t = time.time()
    with open(in_path, 'r', encoding='utf-8', errors='ignore') as fin, open(out_path, 'wb') as fout:
        while True:
            c = fin.read(32 << 20)
            if not c:
                break
            read += len(c)
            buf += c
            parts = buf.split(EOS)
            buf = parts.pop()
            docs = [p.strip(chr(10)) for p in parts]
            docs = [d for d in docs if d.strip()]
            for i in range(0, len(docs), 2000):
                ids = []
                for e in tok.encode_batch(docs[i:i + 2000]):
                    ids.extend(e.ids)
                    ids.append(EOS_ID)
                arr = np.asarray(ids, dtype=np.uint16)
                fout.write(arr.tobytes())
                n_tok += int(arr.size)
            if read - last > 300 << 20:
                last = read
                log('  tokenized', round(read / 1e9, 2), '/', round(total / 1e9, 2), 'GB |',
                    round(n_tok / 1e6, 1), 'M tokens |',
                    int(n_tok / max(1.0, time.time() - t) / 1000), 'k tok/s')
        tail = buf.strip(chr(10))
        if tail.strip():
            ids = []
            for e in tok.encode_batch([tail]):
                ids.extend(e.ids)
                ids.append(EOS_ID)
            arr = np.asarray(ids, dtype=np.uint16)
            fout.write(arr.tobytes())
            n_tok += int(arr.size)
    return n_tok


TRAIN_BIN = os.path.join(WORK, 'train.bin')
VAL_BIN = os.path.join(WORK, 'val.bin')
if os.path.exists(TRAIN_BIN):
    n_train = os.path.getsize(TRAIN_BIN) // 2
else:
    log('tokenizing train corpus (1.9 GB)...')
    n_train = tokenize_to_bin(train_txt, TRAIN_BIN)
log('train tokens', n_train)
if os.path.exists(VAL_BIN):
    n_val = os.path.getsize(VAL_BIN) // 2
else:
    log('tokenizing valid corpus...')
    n_val = tokenize_to_bin(valid_txt, VAL_BIN)
log('valid tokens', n_val)
log('corpus tokens', round(n_train / 1e6, 1), 'M')

train_data = np.memmap(TRAIN_BIN, dtype=np.uint16, mode='r')
val_data = np.memmap(VAL_BIN, dtype=np.uint16, mode='r')

# ----------------------------------------------------------------------------- model
CTX = 256
torch.manual_seed(1337)
cfg = GPT2Config(
    vocab_size=VOCAB, n_positions=CTX, n_ctx=CTX,
    n_embd=256, n_layer=10, n_head=8, n_inner=1024,
    activation_function='gelu_new',
    resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0,
    layer_norm_epsilon=1e-5, initializer_range=0.02,
    tie_word_embeddings=True, bos_token_id=EOS_ID, eos_token_id=EOS_ID,
)
model = GPT2LMHeadModel(cfg)
N_PARAMS = model.num_parameters()
N_NONEMB = N_PARAMS - model.transformer.wte.weight.numel()
log('params: {:,} total ({:.2f}M), {:.2f}M non-embedding'.format(
    N_PARAMS, N_PARAMS / 1e6, N_NONEMB / 1e6))
log('config: d_model=256 layers=10 heads=8 ctx={} vocab={} tied-embeddings'.format(CTX, VOCAB))

model.to(DEVICE)
torch.set_float32_matmul_precision('high')
if IS_CUDA:
    torch.backends.cuda.matmul.allow_tf32 = True

# ----------------------------------------------------------------------------- training
MICRO, ACCUM = 32, 8
BATCH_TOKENS = MICRO * ACCUM * CTX
MAX_STEPS = int(os.environ.get('TS_MAX_STEPS', 12000))
MAX_TRAIN_SECONDS = float(os.environ.get('TS_MAX_MINUTES', 110)) * 60
WARMUP, LR, MIN_LR, WD = 300, 1e-3, 1e-4, 0.1
EVAL_EVERY, CKPT_EVERY = 500, 2000

opt = torch.optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95), weight_decay=WD)


def lr_lambda(s):
    if s < WARMUP:
        return (s + 1) / WARMUP
    p = min(1.0, max(0.0, (s - WARMUP) / max(1, MAX_STEPS - WARMUP)))
    return MIN_LR / LR + (1 - MIN_LR / LR) * 0.5 * (1 + math.cos(math.pi * p))


sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
drng = np.random.default_rng(1234)


def make_batch(bs):
    ix = drng.integers(0, n_train - CTX - 1, size=bs)
    x = np.empty((bs, CTX), dtype=np.int64)
    y = np.empty((bs, CTX), dtype=np.int64)
    for j, i in enumerate(ix):
        x[j] = train_data[i:i + CTX]
        y[j] = train_data[i + 1:i + CTX + 1]
    return torch.from_numpy(x).to(DEVICE), torch.from_numpy(y).to(DEVICE)


@torch.no_grad()
def eval_val(nb=32, bs=32, seed=99):
    model.eval()
    r = np.random.default_rng(seed)
    ix = r.integers(0, n_val - CTX - 1, size=nb * bs)
    tot, ntok = 0.0, 0
    for k in range(0, len(ix), bs):
        sl = ix[k:k + bs]
        x = np.stack([val_data[i:i + CTX] for i in sl]).astype(np.int64)
        y = np.stack([val_data[i + 1:i + CTX + 1] for i in sl]).astype(np.int64)
        x = torch.from_numpy(x).to(DEVICE)
        y = torch.from_numpy(y).to(DEVICE)
        with amp():
            logits = model(x).logits
        tot += F.cross_entropy(logits.float().view(-1, VOCAB), y.view(-1),
                               reduction='sum').item()
        ntok += y.numel()
    model.train()
    return tot / ntok


def build_fast_tokenizer():
    return PreTrainedTokenizerFast(
        tokenizer_object=tok, eos_token=EOS, bos_token=EOS, pad_token=EOS,
        model_max_length=CTX,
    )


def push_ckpt(step):
    if not HF_TOKEN:
        return
    d = os.path.join(WORK, 'ckpt')
    os.makedirs(d, exist_ok=True)
    model.save_pretrained(d, safe_serialization=True)
    build_fast_tokenizer().save_pretrained(d)
    api.upload_folder(repo_id=REPO_ID, folder_path=d, token=HF_TOKEN,
                      commit_message='checkpoint at step ' + str(step))


hist, evals = [], []
t_start = time.time()
model.train()
log('training | {} tokens/step | up to {} steps | wall-clock cap {} min'.format(
    BATCH_TOKENS, MAX_STEPS, int(MAX_TRAIN_SECONDS / 60)))
for step in range(MAX_STEPS):
    if time.time() - t_start > MAX_TRAIN_SECONDS:
        log('wall-clock budget reached, stopping early at step', step)
        break
    opt.zero_grad(set_to_none=True)
    acc = 0.0
    for _ in range(ACCUM):
        x, y = make_batch(MICRO)
        with amp():
            out = model(x, labels=y)
            loss = out.loss / ACCUM
        loss.backward()
        acc += out.loss.item() / ACCUM
    gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    sched.step()
    hist.append((step, acc, opt.param_groups[0]['lr']))
    if step % 50 == 0:
        tps = (step + 1) * BATCH_TOKENS / (time.time() - t_start)
        log('step {:6d} | loss {:.4f} | lr {:.2e} | gnorm {:.2f} | {} k tok/s'.format(
            step, acc, opt.param_groups[0]['lr'], float(gn), int(tps / 1000)))
    if step > 0 and step % EVAL_EVERY == 0:
        vl = eval_val()
        evals.append((step, vl))
        log('  >> val loss {:.4f} at step {}'.format(vl, step))
    if step > 0 and step % CKPT_EVERY == 0:
        try:
            push_ckpt(step)
            log('  >> pushed checkpoint at step', step)
        except Exception as e:
            log('  WARN checkpoint push failed:', repr(e))

STEPS_DONE = (hist[-1][0] + 1) if hist else 0
TOKENS_SEEN = STEPS_DONE * BATCH_TOKENS
TRAIN_SECONDS = time.time() - t_start
log('training done | steps', STEPS_DONE, '| tokens seen', round(TOKENS_SEEN / 1e6, 1), 'M |',
    round(TOKENS_SEEN / n_train, 2), 'epochs |', round(TRAIN_SECONDS / 60, 1), 'min')

log('final val loss...')
FINAL_VAL = eval_val(nb=64)
log('final val loss {:.4f}'.format(FINAL_VAL))

# ----------------------------------------------------------------------------- bits-per-byte
NCHARS = 1_000_000 if IS_CUDA else 300_000
with open(valid_txt, 'r', encoding='utf-8', errors='ignore') as f:
    VTEXT = f.read(NCHARS)
NBYTES = len(VTEXT.encode('utf-8'))
log('comparison text:', NBYTES, 'bytes')


@torch.no_grad()
def nll_stats(m, ids, ctx=256):
    m.eval()
    tot, ntok = 0.0, 0
    for i in range(0, len(ids) - 1, ctx):
        ch = ids[i:i + ctx + 1]
        if len(ch) < 2:
            break
        x = torch.tensor(ch[:-1], dtype=torch.long, device=DEVICE).unsqueeze(0)
        y = torch.tensor(ch[1:], dtype=torch.long, device=DEVICE).unsqueeze(0)
        with amp():
            logits = m(x).logits
        tot += F.cross_entropy(logits.float().view(-1, logits.shape[-1]), y.view(-1),
                               reduction='sum').item()
        ntok += y.numel()
    return tot / ntok, ntok


results = {}


def scored(name, m, ids):
    l, nt = nll_stats(m, ids)
    bpt = NBYTES / nt
    results[name] = {'params': int(m.num_parameters()), 'loss_per_token': l,
                     'bits_per_byte': l / (bpt * math.log(2)), 'tokens': int(nt)}
    log('  {} | {:.2f}M params | loss/token {:.4f} | bits/byte {:.4f} | {:.3f} bytes/token'.format(
        name, results[name]['params'] / 1e6, l, results[name]['bits_per_byte'], bpt))


log('evaluating our model...')
scored('tinystories-10m (ours)', model, tok.encode(VTEXT, add_special_tokens=False).ids)

if IS_CUDA:
    neo_tok = None
    for rid in ['roneneldan/TinyStories-8M', 'roneneldan/TinyStories-28M']:
        try:
            log('loading baseline', rid)
            if neo_tok is None:
                neo_tok = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-125M')
            bm = AutoModelForCausalLM.from_pretrained(rid, torch_dtype=torch.float32).to(DEVICE).eval()
            scored(rid, bm, neo_tok.encode(VTEXT))
            del bm
            torch.cuda.empty_cache()
        except Exception as e:
            log('WARN baseline failed', rid, repr(e))
else:
    log('skipping published baselines on CPU')

# ----------------------------------------------------------------------------- samples
log('generating samples...')
gens = []


def sample(prompt, max_new=200, seed=0):
    torch.manual_seed(seed)
    ids = tok.encode(prompt, add_special_tokens=False).ids
    x = torch.tensor([ids], device=DEVICE)
    with torch.no_grad():
        out = model.generate(x, max_new_tokens=max_new, do_sample=True, top_k=50,
                             temperature=0.8, repetition_penalty=1.1, pad_token_id=EOS_ID)
    return tok.decode(out[0].tolist()).strip()


model.eval()
for p, s in [('Once upon a time', 0), ('One day, a little boy named Tom', 1),
             ('The cat wanted to', 2), ('Lily and her friend', 3)]:
    try:
        txt = sample(p, seed=s)
        gens.append((p, txt))
        log('--- sample', repr(p), '---')
        print(txt, flush=True)
    except Exception as e:
        log('WARN generation failed:', repr(e))

# ----------------------------------------------------------------------------- save + upload
FINAL = os.path.join(WORK, 'final')
os.makedirs(FINAL, exist_ok=True)
model.save_pretrained(FINAL, safe_serialization=True)
build_fast_tokenizer().save_pretrained(FINAL)
try:
    shutil.copy(os.path.abspath(__file__), os.path.join(FINAL, 'train.py'))
except Exception:
    pass

summary = {
    'repo_id': REPO_ID, 'dataset': DATASET,
    'params_total': N_PARAMS, 'params_non_embedding': N_NONEMB,
    'config': {'n_embd': 256, 'n_layer': 10, 'n_head': 8, 'n_inner': 1024, 'n_ctx': CTX,
               'vocab_size': VOCAB, 'tie_word_embeddings': True,
               'tokenizer': 'byte-level BPE, 8192 tokens, trained on TinyStories'},
    'corpus_tokens': int(n_train), 'valid_tokens': int(n_val),
    'steps_done': STEPS_DONE, 'batch_tokens': BATCH_TOKENS, 'tokens_seen': int(TOKENS_SEEN),
    'epochs': TOKENS_SEEN / n_train, 'train_seconds': TRAIN_SECONDS, 'device': DEVICE,
    'optimizer': {'name': 'AdamW', 'lr': LR, 'min_lr': MIN_LR, 'betas': [0.9, 0.95],
                  'weight_decay': WD, 'warmup_steps': WARMUP, 'schedule': 'cosine',
                  'grad_clip': 1.0, 'precision': 'bfloat16 autocast' if IS_CUDA else 'fp32'},
    'final_val_loss': FINAL_VAL, 'val_loss_curve': evals, 'train_loss_curve': hist,
    'comparison': results,
}

with open(os.path.join(FINAL, 'training_summary.json'), 'w') as f:
    json.dump(summary, f, indent=2)

try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    xs = [h[0] for h in hist]
    ys = [h[1] for h in hist]
    k = min(25, max(1, len(ys) // 10))
    ys_s = np.convolve(ys, np.ones(k) / k, mode='valid')
    plt.figure(figsize=(8, 4.5))
    plt.plot(xs[k - 1:], ys_s, lw=1.4, label='train loss (smoothed)')
    if evals:
        plt.plot([e[0] for e in evals], [e[1] for e in evals], 'o-', ms=3, lw=1.2, label='val loss')
    plt.xlabel('step')
    plt.ylabel('cross-entropy (nats / token)')
    plt.title('tinystories-10m | {}M tokens seen'.format(round(TOKENS_SEEN / 1e6)))
    plt.grid(alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(FINAL, 'loss_curve.png'), dpi=130)
    plt.close()
except Exception as e:
    log('WARN plot failed:', repr(e))

rows = []
for k_, v in results.items():
    rows.append('| ' + k_ + ' | ' + ('%.2fM' % (v['params'] / 1e6)) + ' | ' +
                ('%.4f' % v['loss_per_token']) + ' | ' + ('%.4f' % v['bits_per_byte']) + ' |')
rows = chr(10).join(rows)

gen_md = chr(10).join(['**' + p + '**' + chr(10) + chr(10) + '> ' + t.replace(chr(10), ' ') + chr(10)
                       for p, t in gens])

card = '---\n' + \
    'language:\n- en\n' + \
    'license: cdla-sharing-1.0\n' + \
    'datasets:\n- roneneldan/TinyStories\n' + \
    'pipeline_tag: text-generation\n' + \
    'library_name: transformers\n' + \
    'tags:\n- gpt2\n- tinystories\n- tiny-llm\n- from-scratch\n' + \
    '---\n\n' + \
    '# tinystories-10m\n\n' + \
    'A **10.06M parameter** GPT-2 style language model trained **from scratch** on\n' + \
    '[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)\n' + \
    '(Eldan & Li, 2023, [arXiv:2305.07759](https://arxiv.org/abs/2305.07759)) — a synthetic corpus of\n' + \
    'simple English children stories built as a sandbox for very small language models.\n\n' + \
    'One run built all of it: the tokenizer was trained on the corpus, the corpus was tokenized, and\n' + \
    'the model was trained and evaluated. No pretrained weights were used.\n\n' + \
    '## Model\n\n' + \
    '| | |\n|---|---|\n' + \
    '| Architecture | GPT-2 (pre-LN, learned absolute positions, GELU) |\n' + \
    '| Parameters | **' + format(N_PARAMS, ',') + '** (' + ('%.2fM' % (N_PARAMS / 1e6)) + ' total, ' + \
    ('%.2fM' % (N_NONEMB / 1e6)) + ' non-embedding) |\n' + \
    '| Layers / heads / hidden | 10 / 8 / 256 |\n' + \
    '| Context length | 256 tokens |\n' + \
    '| MLP inner size | 1024 |\n' + \
    '| Vocab | 8192, byte-level BPE trained on TinyStories |\n' + \
    '| Embeddings | tied input/output |\n' + \
    '| Precision | ' + ('bfloat16 autocast' if IS_CUDA else 'fp32') + ' |\n\n' + \
    '## Training\n\n' + \
    '| | |\n|---|---|\n' + \
    '| Data | TinyStories train split (' + format(n_train, ',') + ' tokens, ' + \
    format(n_val, ',') + ' in validation) |\n' + \
    '| Tokens seen | ' + format(int(TOKENS_SEEN), ',') + ' (' + \
    ('%.2f' % (TOKENS_SEEN / n_train)) + ' epochs over the corpus) |\n' + \
    '| Steps | ' + format(STEPS_DONE, ',') + ' at ' + format(BATCH_TOKENS, ',') + ' tokens/step |\n' + \
    '| Optimizer | AdamW, lr 1e-3 cosine to 1e-4, 300 step warmup, wd 0.1, grad clip 1.0 |\n' + \
    '| Hardware / time | ' + DEVICE + ', ' + ('%.1f' % (TRAIN_SECONDS / 60)) + ' min |\n' + \
    '| Final val loss | **' + ('%.4f' % FINAL_VAL) + '** nats/token |\n\n' + \
    '## Results vs published TinyStories baselines\n\n' + \
    'Losses are computed on the same slice of the TinyStories validation set. The tokenizers differ,\n' + \
    'so per-token loss is not comparable across rows — **bits per byte** is tokenizer-independent, so\n' + \
    'use that column to compare.\n\n' + \
    '| Model | Params | Val loss (nats/token) | Bits per byte |\n|---|---|---|---|\n' + rows + '\n\n' + \
    '## Samples\n\n' + gen_md + '\n' + \
    '## Usage\n\n' + \
    '```python\n' + \
    'from transformers import AutoModelForCausalLM, AutoTokenizer\n\n' + \
    "model = AutoModelForCausalLM.from_pretrained('kunal10tripathi/tinystories-10m')\n" + \
    "tok = AutoTokenizer.from_pretrained('kunal10tripathi/tinystories-10m')\n\n" + \
    "ids = tok('Once upon a time', return_tensors='pt').input_ids\n" + \
    'out = model.generate(ids, max_new_tokens=200, do_sample=True, top_k=50, temperature=0.8)\n' + \
    'print(tok.decode(out[0], skip_special_tokens=True))\n' + \
    '```\n\n' + \
    'The tokenizer uses `<|endoftext|>` (id 0) as BOS/EOS/PAD and as the story separator.\n\n' + \
    '## Reproducing\n\n' + \
    '```bash\n' + \
    'hf jobs uv run --flavor a10g-small --timeout 3h --secrets HF_TOKEN \\\n' + \
    '  -w torch -w transformers -w tokenizers -w matplotlib -w huggingface_hub \\\n' + \
    '  -c "from huggingface_hub import hf_hub_download; import runpy;\n' + \
    "      runpy.run_path(hf_hub_download('kunal10tripathi/tinystories-10m','train.py'), run_name='__main__')\"\n" + \
    '```\n\n' + \
    '## Files\n\n' + \
    '- `train.py` — the exact script used, start to finish\n' + \
    '- `training_summary.json` — full config, loss curves, comparison numbers\n' + \
    '- `loss_curve.png` — train and validation loss\n\n' + \
    '## Notes\n\n' + \
    'TinyStories models are conventionally named by their *non-embedding* parameter count, and the\n' + \
    'published baselines use a 50257 token GPT-2/GPT-Neo vocabulary whose embedding table dominates\n' + \
    'their parameter budget. This model uses an 8192 token vocabulary trained on the corpus itself and\n' + \
    'is named by its ~10M *total* parameter count.\n\n' + \
    'Trained from scratch for fun. Facts will be invented; grammar mostly will not.\n'

with open(os.path.join(FINAL, 'README.md'), 'w') as f:
    f.write(card)

if HF_TOKEN:
    log('uploading to', REPO_ID)
    try:
        api.upload_folder(
            repo_id=REPO_ID, folder_path=FINAL, token=HF_TOKEN,
            commit_message='tinystories-10m: {:.2f}M params, val loss {:.4f}'.format(
                N_PARAMS / 1e6, FINAL_VAL))
        log('uploaded')
    except Exception as e:
        log('WARN upload failed:', repr(e), '| artifacts are in', FINAL)
else:
    log('no HF_TOKEN: artifacts are in', FINAL)

log('DONE. wall clock', round((time.time() - T0) / 60, 1), 'min')
log(json.dumps({'params': N_PARAMS, 'steps': STEPS_DONE, 'tokens_seen': int(TOKENS_SEEN),
                'final_val_loss': FINAL_VAL, 'comparison': results}, indent=2))

	"""
	tinystories-10m : a ~10M parameter GPT-2 trained from scratch on TinyStories.

	One script, end to end:
	1. download the TinyStories train/valid raw text
	2. train a byte-level BPE tokenizer (vocab 8192) on a sample of the corpus
	3. tokenize the corpus into flat uint16 arrays
	4. train a GPT-2 (256 dim / 10 layers / 8 heads / tied embeddings ~ 10.06M params)
	5. evaluate val loss + bits-per-byte, and compare against the published TinyStories baselines
	6. sample generations, plot the loss curve, push everything to the Hub

	Runs on CUDA (fast path, bfloat16) or on CPU (slow path, fp32) with no code changes.

	Environment variables (all optional):
	TS_MAX_STEPS training steps, default 12000 (12000 x 65536 tok = ~786M tokens)
	TS_MAX_MINUTES wall-clock cap on training, default 110
	TS_DEVICE 'cuda' or 'cpu', default auto-detect
	TS_WORK scratch directory, default /tmp/tswork
	HF_TOKEN needed only for pushing to the Hub; skips upload if unset

	Examples
	HF Jobs (GPU):
	hf jobs uv run --flavor a10g-small --timeout 3h --secrets HF_TOKEN \\
	-w torch -w transformers -w tokenizers -w matplotlib -w huggingface_hub \\
	-c "from huggingface_hub import hf_hub_download; import runpy;
	runpy.run_path(hf_hub_download('kunal10tripathi/tinystories-10m','train.py'), run_name='__main__')"

	Quick CPU smoke run (roughly 40M tokens):
	TS_DEVICE=cpu TS_MAX_STEPS=600 TS_MAX_MINUTES=90 python train.py
	"""

	import os, json, math, time, shutil, contextlib
	import numpy as np
	import torch
	import torch.nn.functional as F
	from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
	from transformers import (
	GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast,
	AutoTokenizer, AutoModelForCausalLM,
	)
	from huggingface_hub import hf_hub_download, HfApi

	# ----------------------------------------------------------------------------- config
	REPO_ID = 'kunal10tripathi/tinystories-10m'
	DATASET = 'roneneldan/TinyStories'
	EOS = '<\|endoftext\|>'
	WORK = os.environ.get('TS_WORK', '/tmp/tswork')
	os.makedirs(WORK, exist_ok=True)
	T0 = time.time()


	def log(*a):
	msg = ' '.join(str(x) for x in a)
	print('[' + str(round(time.time() - T0, 1)).rjust(8) + 's] ' + msg, flush=True)


	HF_TOKEN = os.environ.get('HF_TOKEN')
	api = HfApi(token=HF_TOKEN)
	if HF_TOKEN:
	try:
	log('logged in as', api.whoami()['name'])
	except Exception as e:
	log('WARN whoami failed:', repr(e))
	else:
	log('no HF_TOKEN set: results will be kept local and not uploaded')

	DEVICE = os.environ.get('TS_DEVICE') or ('cuda' if torch.cuda.is_available() else 'cpu')
	IS_CUDA = DEVICE.startswith('cuda')
	log('device', DEVICE, torch.cuda.get_device_name(0) if IS_CUDA else '')


	def amp():
	"""bfloat16 autocast on GPU, plain fp32 on CPU."""
	if IS_CUDA:
	return torch.autocast('cuda', dtype=torch.bfloat16)
	return contextlib.nullcontext()


	# ----------------------------------------------------------------------------- data
	log('downloading TinyStories...')
	train_txt = hf_hub_download(DATASET, 'TinyStories-train.txt', repo_type='dataset')
	valid_txt = hf_hub_download(DATASET, 'TinyStories-valid.txt', repo_type='dataset')
	log('train bytes', os.path.getsize(train_txt), '\| valid bytes', os.path.getsize(valid_txt))

	# ----------------------------------------------------------------------------- tokenizer
	TOK_PATH = os.path.join(WORK, 'tokenizer.json')
	if os.path.exists(TOK_PATH):
	tok = Tokenizer.from_file(TOK_PATH)
	else:
	def sample_text():
	n = 0
	with open(train_txt, 'r', encoding='utf-8', errors='ignore') as f:
	while n < 80_000_000:
	c = f.read(1_000_000)
	if not c:
	break
	n += len(c)
	yield c

	tok = Tokenizer(models.BPE(unk_token=None))
	tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
	tok.decoder = decoders.ByteLevel()
	tr = trainers.BpeTrainer(
	vocab_size=8192, min_frequency=2, special_tokens=[EOS],
	initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False,
	)
	log('training BPE tokenizer (vocab 8192) on an 80MB sample...')
	tok.train_from_iterator(sample_text(), trainer=tr)
	tok.save(TOK_PATH)

	VOCAB = tok.get_vocab_size()
	EOS_ID = tok.token_to_id(EOS)
	log('tokenizer ready \| vocab', VOCAB, '\| eos_id', EOS_ID)


	def tokenize_to_bin(in_path, out_path):
	total = os.path.getsize(in_path)
	read = last = n_tok = 0
	buf = ''
	t = time.time()
	with open(in_path, 'r', encoding='utf-8', errors='ignore') as fin, open(out_path, 'wb') as fout:
	while True:
	c = fin.read(32 << 20)
	if not c:
	break
	read += len(c)
	buf += c
	parts = buf.split(EOS)
	buf = parts.pop()
	docs = [p.strip(chr(10)) for p in parts]
	docs = [d for d in docs if d.strip()]
	for i in range(0, len(docs), 2000):
	ids = []
	for e in tok.encode_batch(docs[i:i + 2000]):
	ids.extend(e.ids)
	ids.append(EOS_ID)
	arr = np.asarray(ids, dtype=np.uint16)
	fout.write(arr.tobytes())
	n_tok += int(arr.size)
	if read - last > 300 << 20:
	last = read
	log(' tokenized', round(read / 1e9, 2), '/', round(total / 1e9, 2), 'GB \|',
	round(n_tok / 1e6, 1), 'M tokens \|',
	int(n_tok / max(1.0, time.time() - t) / 1000), 'k tok/s')
	tail = buf.strip(chr(10))
	if tail.strip():
	ids = []
	for e in tok.encode_batch([tail]):
	ids.extend(e.ids)
	ids.append(EOS_ID)
	arr = np.asarray(ids, dtype=np.uint16)
	fout.write(arr.tobytes())
	n_tok += int(arr.size)
	return n_tok


	TRAIN_BIN = os.path.join(WORK, 'train.bin')
	VAL_BIN = os.path.join(WORK, 'val.bin')
	if os.path.exists(TRAIN_BIN):
	n_train = os.path.getsize(TRAIN_BIN) // 2
	else:
	log('tokenizing train corpus (1.9 GB)...')
	n_train = tokenize_to_bin(train_txt, TRAIN_BIN)
	log('train tokens', n_train)
	if os.path.exists(VAL_BIN):
	n_val = os.path.getsize(VAL_BIN) // 2
	else:
	log('tokenizing valid corpus...')
	n_val = tokenize_to_bin(valid_txt, VAL_BIN)
	log('valid tokens', n_val)
	log('corpus tokens', round(n_train / 1e6, 1), 'M')

	train_data = np.memmap(TRAIN_BIN, dtype=np.uint16, mode='r')
	val_data = np.memmap(VAL_BIN, dtype=np.uint16, mode='r')

	# ----------------------------------------------------------------------------- model
	CTX = 256
	torch.manual_seed(1337)
	cfg = GPT2Config(
	vocab_size=VOCAB, n_positions=CTX, n_ctx=CTX,
	n_embd=256, n_layer=10, n_head=8, n_inner=1024,
	activation_function='gelu_new',
	resid_pdrop=0.0, embd_pdrop=0.0, attn_pdrop=0.0,
	layer_norm_epsilon=1e-5, initializer_range=0.02,
	tie_word_embeddings=True, bos_token_id=EOS_ID, eos_token_id=EOS_ID,
	)
	model = GPT2LMHeadModel(cfg)
	N_PARAMS = model.num_parameters()
	N_NONEMB = N_PARAMS - model.transformer.wte.weight.numel()
	log('params: {:,} total ({:.2f}M), {:.2f}M non-embedding'.format(
	N_PARAMS, N_PARAMS / 1e6, N_NONEMB / 1e6))
	log('config: d_model=256 layers=10 heads=8 ctx={} vocab={} tied-embeddings'.format(CTX, VOCAB))

	model.to(DEVICE)
	torch.set_float32_matmul_precision('high')
	if IS_CUDA:
	torch.backends.cuda.matmul.allow_tf32 = True

	# ----------------------------------------------------------------------------- training
	MICRO, ACCUM = 32, 8
	BATCH_TOKENS = MICRO * ACCUM * CTX
	MAX_STEPS = int(os.environ.get('TS_MAX_STEPS', 12000))
	MAX_TRAIN_SECONDS = float(os.environ.get('TS_MAX_MINUTES', 110)) * 60
	WARMUP, LR, MIN_LR, WD = 300, 1e-3, 1e-4, 0.1
	EVAL_EVERY, CKPT_EVERY = 500, 2000

	opt = torch.optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95), weight_decay=WD)


	def lr_lambda(s):
	if s < WARMUP:
	return (s + 1) / WARMUP
	p = min(1.0, max(0.0, (s - WARMUP) / max(1, MAX_STEPS - WARMUP)))
	return MIN_LR / LR + (1 - MIN_LR / LR) * 0.5 * (1 + math.cos(math.pi * p))


	sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
	drng = np.random.default_rng(1234)


	def make_batch(bs):
	ix = drng.integers(0, n_train - CTX - 1, size=bs)
	x = np.empty((bs, CTX), dtype=np.int64)
	y = np.empty((bs, CTX), dtype=np.int64)
	for j, i in enumerate(ix):
	x[j] = train_data[i:i + CTX]
	y[j] = train_data[i + 1:i + CTX + 1]
	return torch.from_numpy(x).to(DEVICE), torch.from_numpy(y).to(DEVICE)


	@torch.no_grad()
	def eval_val(nb=32, bs=32, seed=99):
	model.eval()
	r = np.random.default_rng(seed)
	ix = r.integers(0, n_val - CTX - 1, size=nb * bs)
	tot, ntok = 0.0, 0
	for k in range(0, len(ix), bs):
	sl = ix[k:k + bs]
	x = np.stack([val_data[i:i + CTX] for i in sl]).astype(np.int64)
	y = np.stack([val_data[i + 1:i + CTX + 1] for i in sl]).astype(np.int64)
	x = torch.from_numpy(x).to(DEVICE)
	y = torch.from_numpy(y).to(DEVICE)
	with amp():
	logits = model(x).logits
	tot += F.cross_entropy(logits.float().view(-1, VOCAB), y.view(-1),
	reduction='sum').item()
	ntok += y.numel()
	model.train()
	return tot / ntok


	def build_fast_tokenizer():
	return PreTrainedTokenizerFast(
	tokenizer_object=tok, eos_token=EOS, bos_token=EOS, pad_token=EOS,
	model_max_length=CTX,
	)


	def push_ckpt(step):
	if not HF_TOKEN:
	return
	d = os.path.join(WORK, 'ckpt')
	os.makedirs(d, exist_ok=True)
	model.save_pretrained(d, safe_serialization=True)
	build_fast_tokenizer().save_pretrained(d)
	api.upload_folder(repo_id=REPO_ID, folder_path=d, token=HF_TOKEN,
	commit_message='checkpoint at step ' + str(step))


	hist, evals = [], []
	t_start = time.time()
	model.train()
	log('training \| {} tokens/step \| up to {} steps \| wall-clock cap {} min'.format(
	BATCH_TOKENS, MAX_STEPS, int(MAX_TRAIN_SECONDS / 60)))
	for step in range(MAX_STEPS):
	if time.time() - t_start > MAX_TRAIN_SECONDS:
	log('wall-clock budget reached, stopping early at step', step)
	break
	opt.zero_grad(set_to_none=True)
	acc = 0.0
	for _ in range(ACCUM):
	x, y = make_batch(MICRO)
	with amp():
	out = model(x, labels=y)
	loss = out.loss / ACCUM
	loss.backward()
	acc += out.loss.item() / ACCUM
	gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	opt.step()
	sched.step()
	hist.append((step, acc, opt.param_groups[0]['lr']))
	if step % 50 == 0:
	tps = (step + 1) * BATCH_TOKENS / (time.time() - t_start)
	log('step {:6d} \| loss {:.4f} \| lr {:.2e} \| gnorm {:.2f} \| {} k tok/s'.format(
	step, acc, opt.param_groups[0]['lr'], float(gn), int(tps / 1000)))
	if step > 0 and step % EVAL_EVERY == 0:
	vl = eval_val()
	evals.append((step, vl))
	log(' >> val loss {:.4f} at step {}'.format(vl, step))
	if step > 0 and step % CKPT_EVERY == 0:
	try:
	push_ckpt(step)
	log(' >> pushed checkpoint at step', step)
	except Exception as e:
	log(' WARN checkpoint push failed:', repr(e))

	STEPS_DONE = (hist[-1][0] + 1) if hist else 0
	TOKENS_SEEN = STEPS_DONE * BATCH_TOKENS
	TRAIN_SECONDS = time.time() - t_start
	log('training done \| steps', STEPS_DONE, '\| tokens seen', round(TOKENS_SEEN / 1e6, 1), 'M \|',
	round(TOKENS_SEEN / n_train, 2), 'epochs \|', round(TRAIN_SECONDS / 60, 1), 'min')

	log('final val loss...')
	FINAL_VAL = eval_val(nb=64)
	log('final val loss {:.4f}'.format(FINAL_VAL))

	# ----------------------------------------------------------------------------- bits-per-byte
	NCHARS = 1_000_000 if IS_CUDA else 300_000
	with open(valid_txt, 'r', encoding='utf-8', errors='ignore') as f:
	VTEXT = f.read(NCHARS)
	NBYTES = len(VTEXT.encode('utf-8'))
	log('comparison text:', NBYTES, 'bytes')


	@torch.no_grad()
	def nll_stats(m, ids, ctx=256):
	m.eval()
	tot, ntok = 0.0, 0
	for i in range(0, len(ids) - 1, ctx):
	ch = ids[i:i + ctx + 1]
	if len(ch) < 2:
	break
	x = torch.tensor(ch[:-1], dtype=torch.long, device=DEVICE).unsqueeze(0)
	y = torch.tensor(ch[1:], dtype=torch.long, device=DEVICE).unsqueeze(0)
	with amp():
	logits = m(x).logits
	tot += F.cross_entropy(logits.float().view(-1, logits.shape[-1]), y.view(-1),
	reduction='sum').item()
	ntok += y.numel()
	return tot / ntok, ntok


	results = {}


	def scored(name, m, ids):
	l, nt = nll_stats(m, ids)
	bpt = NBYTES / nt
	results[name] = {'params': int(m.num_parameters()), 'loss_per_token': l,
	'bits_per_byte': l / (bpt * math.log(2)), 'tokens': int(nt)}
	log(' {} \| {:.2f}M params \| loss/token {:.4f} \| bits/byte {:.4f} \| {:.3f} bytes/token'.format(
	name, results[name]['params'] / 1e6, l, results[name]['bits_per_byte'], bpt))


	log('evaluating our model...')
	scored('tinystories-10m (ours)', model, tok.encode(VTEXT, add_special_tokens=False).ids)

	if IS_CUDA:
	neo_tok = None
	for rid in ['roneneldan/TinyStories-8M', 'roneneldan/TinyStories-28M']:
	try:
	log('loading baseline', rid)
	if neo_tok is None:
	neo_tok = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-125M')
	bm = AutoModelForCausalLM.from_pretrained(rid, torch_dtype=torch.float32).to(DEVICE).eval()
	scored(rid, bm, neo_tok.encode(VTEXT))
	del bm
	torch.cuda.empty_cache()
	except Exception as e:
	log('WARN baseline failed', rid, repr(e))
	else:
	log('skipping published baselines on CPU')

	# ----------------------------------------------------------------------------- samples
	log('generating samples...')
	gens = []


	def sample(prompt, max_new=200, seed=0):
	torch.manual_seed(seed)
	ids = tok.encode(prompt, add_special_tokens=False).ids
	x = torch.tensor([ids], device=DEVICE)
	with torch.no_grad():
	out = model.generate(x, max_new_tokens=max_new, do_sample=True, top_k=50,
	temperature=0.8, repetition_penalty=1.1, pad_token_id=EOS_ID)
	return tok.decode(out[0].tolist()).strip()


	model.eval()
	for p, s in [('Once upon a time', 0), ('One day, a little boy named Tom', 1),
	('The cat wanted to', 2), ('Lily and her friend', 3)]:
	try:
	txt = sample(p, seed=s)
	gens.append((p, txt))
	log('--- sample', repr(p), '---')
	print(txt, flush=True)
	except Exception as e:
	log('WARN generation failed:', repr(e))

	# ----------------------------------------------------------------------------- save + upload
	FINAL = os.path.join(WORK, 'final')
	os.makedirs(FINAL, exist_ok=True)
	model.save_pretrained(FINAL, safe_serialization=True)
	build_fast_tokenizer().save_pretrained(FINAL)
	try:
	shutil.copy(os.path.abspath(__file__), os.path.join(FINAL, 'train.py'))
	except Exception:
	pass

	summary = {
	'repo_id': REPO_ID, 'dataset': DATASET,
	'params_total': N_PARAMS, 'params_non_embedding': N_NONEMB,
	'config': {'n_embd': 256, 'n_layer': 10, 'n_head': 8, 'n_inner': 1024, 'n_ctx': CTX,
	'vocab_size': VOCAB, 'tie_word_embeddings': True,
	'tokenizer': 'byte-level BPE, 8192 tokens, trained on TinyStories'},
	'corpus_tokens': int(n_train), 'valid_tokens': int(n_val),
	'steps_done': STEPS_DONE, 'batch_tokens': BATCH_TOKENS, 'tokens_seen': int(TOKENS_SEEN),
	'epochs': TOKENS_SEEN / n_train, 'train_seconds': TRAIN_SECONDS, 'device': DEVICE,
	'optimizer': {'name': 'AdamW', 'lr': LR, 'min_lr': MIN_LR, 'betas': [0.9, 0.95],
	'weight_decay': WD, 'warmup_steps': WARMUP, 'schedule': 'cosine',
	'grad_clip': 1.0, 'precision': 'bfloat16 autocast' if IS_CUDA else 'fp32'},
	'final_val_loss': FINAL_VAL, 'val_loss_curve': evals, 'train_loss_curve': hist,
	'comparison': results,
	}

	with open(os.path.join(FINAL, 'training_summary.json'), 'w') as f:
	json.dump(summary, f, indent=2)

	try:
	import matplotlib
	matplotlib.use('Agg')
	import matplotlib.pyplot as plt
	xs = [h[0] for h in hist]
	ys = [h[1] for h in hist]
	k = min(25, max(1, len(ys) // 10))
	ys_s = np.convolve(ys, np.ones(k) / k, mode='valid')
	plt.figure(figsize=(8, 4.5))
	plt.plot(xs[k - 1:], ys_s, lw=1.4, label='train loss (smoothed)')
	if evals:
	plt.plot([e[0] for e in evals], [e[1] for e in evals], 'o-', ms=3, lw=1.2, label='val loss')
	plt.xlabel('step')
	plt.ylabel('cross-entropy (nats / token)')
	plt.title('tinystories-10m \| {}M tokens seen'.format(round(TOKENS_SEEN / 1e6)))
	plt.grid(alpha=0.3)
	plt.legend()
	plt.tight_layout()
	plt.savefig(os.path.join(FINAL, 'loss_curve.png'), dpi=130)
	plt.close()
	except Exception as e:
	log('WARN plot failed:', repr(e))

	rows = []
	for k_, v in results.items():
	rows.append('\| ' + k_ + ' \| ' + ('%.2fM' % (v['params'] / 1e6)) + ' \| ' +
	('%.4f' % v['loss_per_token']) + ' \| ' + ('%.4f' % v['bits_per_byte']) + ' \|')
	rows = chr(10).join(rows)

	gen_md = chr(10).join(['' + p + '' + chr(10) + chr(10) + '> ' + t.replace(chr(10), ' ') + chr(10)
	for p, t in gens])

	card = '---\n' + \
	'language:\n- en\n' + \
	'license: cdla-sharing-1.0\n' + \
	'datasets:\n- roneneldan/TinyStories\n' + \
	'pipeline_tag: text-generation\n' + \
	'library_name: transformers\n' + \
	'tags:\n- gpt2\n- tinystories\n- tiny-llm\n- from-scratch\n' + \
	'---\n\n' + \
	'# tinystories-10m\n\n' + \
	'A 10.06M parameter GPT-2 style language model trained from scratch on\n' + \
	'[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)\n' + \
	'(Eldan & Li, 2023, [arXiv:2305.07759](https://arxiv.org/abs/2305.07759)) — a synthetic corpus of\n' + \
	'simple English children stories built as a sandbox for very small language models.\n\n' + \
	'One run built all of it: the tokenizer was trained on the corpus, the corpus was tokenized, and\n' + \
	'the model was trained and evaluated. No pretrained weights were used.\n\n' + \
	'## Model\n\n' + \
	'\| \| \|\n\|---\|---\|\n' + \
	'\| Architecture \| GPT-2 (pre-LN, learned absolute positions, GELU) \|\n' + \
	'\| Parameters \| ' + format(N_PARAMS, ',') + ' (' + ('%.2fM' % (N_PARAMS / 1e6)) + ' total, ' + \
	('%.2fM' % (N_NONEMB / 1e6)) + ' non-embedding) \|\n' + \
	'\| Layers / heads / hidden \| 10 / 8 / 256 \|\n' + \
	'\| Context length \| 256 tokens \|\n' + \
	'\| MLP inner size \| 1024 \|\n' + \
	'\| Vocab \| 8192, byte-level BPE trained on TinyStories \|\n' + \
	'\| Embeddings \| tied input/output \|\n' + \
	'\| Precision \| ' + ('bfloat16 autocast' if IS_CUDA else 'fp32') + ' \|\n\n' + \
	'## Training\n\n' + \
	'\| \| \|\n\|---\|---\|\n' + \
	'\| Data \| TinyStories train split (' + format(n_train, ',') + ' tokens, ' + \
	format(n_val, ',') + ' in validation) \|\n' + \
	'\| Tokens seen \| ' + format(int(TOKENS_SEEN), ',') + ' (' + \
	('%.2f' % (TOKENS_SEEN / n_train)) + ' epochs over the corpus) \|\n' + \
	'\| Steps \| ' + format(STEPS_DONE, ',') + ' at ' + format(BATCH_TOKENS, ',') + ' tokens/step \|\n' + \
	'\| Optimizer \| AdamW, lr 1e-3 cosine to 1e-4, 300 step warmup, wd 0.1, grad clip 1.0 \|\n' + \
	'\| Hardware / time \| ' + DEVICE + ', ' + ('%.1f' % (TRAIN_SECONDS / 60)) + ' min \|\n' + \
	'\| Final val loss \| ' + ('%.4f' % FINAL_VAL) + ' nats/token \|\n\n' + \
	'## Results vs published TinyStories baselines\n\n' + \
	'Losses are computed on the same slice of the TinyStories validation set. The tokenizers differ,\n' + \
	'so per-token loss is not comparable across rows — bits per byte is tokenizer-independent, so\n' + \
	'use that column to compare.\n\n' + \
	'\| Model \| Params \| Val loss (nats/token) \| Bits per byte \|\n\|---\|---\|---\|---\|\n' + rows + '\n\n' + \
	'## Samples\n\n' + gen_md + '\n' + \
	'## Usage\n\n' + \
	'```python\n' + \
	'from transformers import AutoModelForCausalLM, AutoTokenizer\n\n' + \
	"model = AutoModelForCausalLM.from_pretrained('kunal10tripathi/tinystories-10m')\n" + \
	"tok = AutoTokenizer.from_pretrained('kunal10tripathi/tinystories-10m')\n\n" + \
	"ids = tok('Once upon a time', return_tensors='pt').input_ids\n" + \
	'out = model.generate(ids, max_new_tokens=200, do_sample=True, top_k=50, temperature=0.8)\n' + \
	'print(tok.decode(out[0], skip_special_tokens=True))\n' + \
	'```\n\n' + \
	'The tokenizer uses `<\|endoftext\|>` (id 0) as BOS/EOS/PAD and as the story separator.\n\n' + \
	'## Reproducing\n\n' + \
	'```bash\n' + \
	'hf jobs uv run --flavor a10g-small --timeout 3h --secrets HF_TOKEN \\\n' + \
	' -w torch -w transformers -w tokenizers -w matplotlib -w huggingface_hub \\\n' + \
	' -c "from huggingface_hub import hf_hub_download; import runpy;\n' + \
	" runpy.run_path(hf_hub_download('kunal10tripathi/tinystories-10m','train.py'), run_name='__main__')\"\n" + \
	'```\n\n' + \
	'## Files\n\n' + \
	'- `train.py` — the exact script used, start to finish\n' + \
	'- `training_summary.json` — full config, loss curves, comparison numbers\n' + \
	'- `loss_curve.png` — train and validation loss\n\n' + \
	'## Notes\n\n' + \
	'TinyStories models are conventionally named by their non-embedding parameter count, and the\n' + \
	'published baselines use a 50257 token GPT-2/GPT-Neo vocabulary whose embedding table dominates\n' + \
	'their parameter budget. This model uses an 8192 token vocabulary trained on the corpus itself and\n' + \
	'is named by its ~10M total parameter count.\n\n' + \
	'Trained from scratch for fun. Facts will be invented; grammar mostly will not.\n'

	with open(os.path.join(FINAL, 'README.md'), 'w') as f:
	f.write(card)

	if HF_TOKEN:
	log('uploading to', REPO_ID)
	try:
	api.upload_folder(
	repo_id=REPO_ID, folder_path=FINAL, token=HF_TOKEN,
	commit_message='tinystories-10m: {:.2f}M params, val loss {:.4f}'.format(
	N_PARAMS / 1e6, FINAL_VAL))
	log('uploaded')
	except Exception as e:
	log('WARN upload failed:', repr(e), '\| artifacts are in', FINAL)
	else:
	log('no HF_TOKEN: artifacts are in', FINAL)

	log('DONE. wall clock', round((time.time() - T0) / 60, 1), 'min')
	log(json.dumps({'params': N_PARAMS, 'steps': STEPS_DONE, 'tokens_seen': int(TOKENS_SEEN),
	'final_val_loss': FINAL_VAL, 'comparison': results}, indent=2))