Buckets:
| """Compiled vs eager (bf16 autocast) against an fp32 reference: loss and gradient relative errors; K=8 peak memory.""" | |
| import sys, 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, 2000) | |
| dev = torch.device("cuda") | |
| model = ThinkDecide(load_backbone(S1_PATH, torch.float32, dev), mp["Yes"], mp["No"], n_think=32).to(dev).train() | |
| with torch.no_grad(): # open the gate so the feedback path is exercised | |
| model.gate.fill_(0.05) | |
| for p in model.head.o.parameters(): p.normal_(0, 0.01) | |
| mb = sorted(random.Random(1).sample(train, 6), key=lambda x: len(x["ids"])) | |
| def run(autocast): | |
| model.zero_grad(set_to_none=True) | |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): | |
| terms = decision_losses(model, mb, 2, dev, tok.pad_token_id, 0.25) | |
| loss = sum(t["loss"] for t in terms) / len(mb) | |
| loss.backward() | |
| g = torch.cat([p.grad.flatten() for n, p in model.named_parameters() if p.grad is not None and "embed" not in n]) | |
| return loss.item(), g.clone() | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| l32, g32 = run(False) | |
| le, ge = run(True) | |
| model.compile_layers() | |
| lc, gc = run(True); lc, gc = run(True) | |
| rel = lambda g: ((g - g32).norm() / g32.norm()).item() | |
| print("loss fp32 %.6f eager-bf16 %.6f (err %.2e) compiled-bf16 %.6f (err %.2e)" % (l32, le, abs(le-l32)/l32, lc, abs(lc-l32)/l32)) | |
| print("grad rel err vs fp32: eager-bf16 %.4f compiled-bf16 %.4f" % (rel(ge), rel(gc))) | |
| # K=8 memory at the training budget, longest items | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| for budget in (32768, 49152): | |
| torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() | |
| longest = sorted(train, key=lambda x: -len(x["ids"]))[:32] | |
| for m in micro_batches(longest, 8, budget, 64): | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| terms = decision_losses(model, m, 8, dev, tok.pad_token_id, 0.25) | |
| (sum(t["loss"] for t in terms) / 32).backward() | |
| print("budget %d K=8 longest-32: peak %.1f GiB" % (budget, torch.cuda.max_memory_allocated() / 2**30)) | |
Xet Storage Details
- Size:
- 2.37 kB
- Xet hash:
- b566ede0befa06a6e0d2e0aa305c66d8268b1334c6b7a8ef5791001659b5ccda
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.