Spaces:
No application file
No application file
| import os | |
| import argparse | |
| import json | |
| # We WANT the ZeroGPU shim active here, so do NOT pop SPACES_ZERO_GPU. | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # MUST come before torch | |
| import torch | |
| # Importing train.py pops SPACES_ZERO_GPU, but spaces + torch are already | |
| # initialised with the GPU shim above, so the decoration below still routes | |
| # through the scheduler. Reuse its model / data helpers. | |
| from train import TinyLM, build_vocab, make_batch | |
| def run(data_path, epochs, batch_size, block_size, n_embed, n_layer, n_head, lr): | |
| device = "cuda" | |
| text = open(data_path, encoding="utf-8", errors="ignore").read() | |
| stoi, itos = build_vocab(text) | |
| model = TinyLM(len(stoi), n_embed, n_layer, n_head, block_size).to(device) | |
| opt = torch.optim.AdamW(model.parameters(), lr=lr) | |
| model.train() | |
| last = None | |
| for epoch in range(epochs): | |
| x, y = make_batch(stoi, text, block_size, batch_size, device) | |
| logits, loss = model(x, tgt=y) | |
| opt.zero_grad() | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| opt.step() | |
| last = float(loss.item()) | |
| if epoch % 25 == 0 or epoch == epochs - 1: | |
| print(f"epoch {epoch}/{epochs} loss {last:.4f}", flush=True) | |
| torch.save({"model": model.cpu().state_dict(), "vocab_size": len(stoi), | |
| "itos": {str(k): v for k, v in itos.items()}, | |
| "cfg": {"n_embed": n_embed, "n_layer": n_layer, "n_head": n_head, | |
| "block_size": block_size}}, | |
| "cyber_model.pt") | |
| print("saved cyber_model.pt") | |
| return f"done loss={last:.4f}" | |
| if __name__ == "__main__": | |
| p = argparse.ArgumentParser(description="Train TinyLM on ZeroGPU.") | |
| p.add_argument("--data", required=True) | |
| p.add_argument("--epochs", type=int, default=500) | |
| p.add_argument("--batch_size", type=int, default=32) | |
| p.add_argument("--block_size", type=int, default=128) | |
| p.add_argument("--n_embed", type=int, default=128) | |
| p.add_argument("--n_layer", type=int, default=4) | |
| p.add_argument("--n_head", type=int, default=4) | |
| p.add_argument("--lr", type=float, default=1e-3) | |
| a = p.parse_args() | |
| print(run(a.data, a.epochs, a.batch_size, a.block_size, | |
| a.n_embed, a.n_layer, a.n_head, a.lr)) |