Buckets:
| """Profile one training step (forward / backward / loss) to find the bottleneck.""" | |
| import sys, time, random, torch | |
| sys.path.insert(0, "/workspace/think-decide/src") | |
| from common import DATA, S1_PATH, Renderer, load_canonical, tokenizer_and_mapping | |
| from model import ThinkDecide, load_backbone, collate, item_logits | |
| from train import micro_batches, decision_losses | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| train, _ = load_canonical(DATA + "/train.jsonl", Renderer(tok, mid, mt, 32), 2048, 4000) | |
| dev = torch.device("cuda") | |
| model = ThinkDecide(load_backbone(S1_PATH, torch.float32, dev), mp["Yes"], mp["No"], n_think=32).to(dev).train() | |
| opt = torch.optim.AdamW(model.param_groups(1e-5, 1e-4, 0.01), fused=True) | |
| rng = random.Random(0) | |
| if len(sys.argv) > 3 and sys.argv[3] == "compile": model.compile_layers() | |
| K, budget = int(sys.argv[1]), int(sys.argv[2]) | |
| def step(prof=False): | |
| idx = rng.sample(range(len(train)), 32) | |
| tf = tb = tl = 0 | |
| for mb in micro_batches([train[i] for i in idx], K, budget, 64): | |
| torch.cuda.synchronize(); t = time.time() | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| batch = collate(mb, dev, tok.pad_token_id) | |
| s = model(batch, K=K)[-1] | |
| torch.cuda.synchronize(); t1 = time.time() | |
| from shared_yesno_supervised import supervised_terms | |
| terms = [supervised_terms(lg, torch.tensor([x["target"]], device=dev), x["primitive"], 0.25) for lg, x in zip(item_logits(s, batch), mb)] | |
| loss = sum(t["loss"] for t in terms) / 32 | |
| torch.cuda.synchronize(); t2 = time.time() | |
| loss.backward(); torch.cuda.synchronize(); t3 = time.time() | |
| tf += t1 - t; tl += t2 - t1; tb += t3 - t2 | |
| torch.cuda.synchronize(); t = time.time() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); opt.zero_grad(set_to_none=True) | |
| torch.cuda.synchronize(); to = time.time() - t | |
| return tf, tl, tb, to, sum(len(train[i]["ids"]) for i in idx) * K | |
| for _ in range(6): step() | |
| tot = [0]*5 | |
| for _ in range(5): | |
| r = step(); tot = [a + b for a, b in zip(tot, r)] | |
| print("K=%d budget=%d fwd %.3f loss %.3f bwd %.3f opt %.3f s/step %.3f pass-tok/s %.0f" % (K, budget, *[x/5 for x in tot[:4]], sum(tot[:4])/5, tot[4]/sum(tot[:4]))) | |
| from torch.profiler import profile, ProfilerActivity | |
| with profile(activities=[ProfilerActivity.CUDA]) as p: | |
| step() | |
| print(p.key_averages().table(sort_by="cuda_time_total", row_limit=18)) | |
Xet Storage Details
- Size:
- 2.45 kB
- Xet hash:
- 5284ec5f8b5c22547d1349baf5b34dd829f949aa90f04070144f669b776557c4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.