| |
| """How much does full BPTT through the latent chain actually cost at depth 15/20? |
| |
| `backprop_depth: null` (every config) means full BPTT: gradients flow through all |
| D latent recurrence steps. coconut.Coconut also implements a truncated window -- |
| if backprop_depth=W, passes older than W steps are detached, capping backward |
| compute and activation memory at O(W) instead of O(D). |
| |
| This times one forward+backward at batch_size_training on real prompts, sweeping |
| the number of latents and the BPTT window, so the depth cost can be separated |
| from the dataset-size cost. |
| |
| Usage: |
| PYTHONPATH=. python scripts/bench_bptt.py --train data/star_2arm_L20_100k_train_fo_bfs.json |
| """ |
| import argparse |
| import json |
| import time |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoConfig |
|
|
| from stokenizer import STokenizer |
| from coconut import Coconut |
|
|
|
|
| def build_batch(data, tok, latent_id, n_latent, bs, device): |
| seqs = [] |
| for s in data[:bs]: |
| q = "<eos> " + "|".join([f" {e[0]} {e[1]} " for e in s["edges"]]).strip() + " [Q] " |
| q += f'{s["target"]} {s["neg_target"]} [R] {s["root"]}' |
| ids = tok.encode(q, add_special_tokens=False) + [latent_id] * n_latent |
| |
| hop = min(max(1, n_latent + 1), max(int(k) for k in s["neighbor_k"])) |
| tgt = s["neighbor_k"][str(hop)][0] |
| ids = ids + tok.encode(str(tgt), add_special_tokens=False) |
| seqs.append(ids) |
| n = min(len(x) for x in seqs) |
| seqs = [x[:n] for x in seqs] |
| ids = torch.tensor(seqs, device=device) |
| labels = ids.clone() |
| labels[:, :-1] = -100 |
| return ids, labels |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--train", required=True) |
| ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json") |
| ap.add_argument("--bs", type=int, default=128) |
| ap.add_argument("--device", default="cuda:0") |
| ap.add_argument("--iters", type=int, default=8) |
| args = ap.parse_args() |
|
|
| tok = STokenizer() |
| latent_id = tok.convert_tokens_to_ids("<|latent|>") |
| data = json.load(open(args.train))[: args.bs] |
|
|
| print(f"data: {args.train} edges/graph: {len(data[0]['edges'])} bs={args.bs}\n") |
| print(f"{'n_latent':>9} {'bptt_W':>7} {'fwd+bwd ms':>12} {'peak MiB':>10}") |
|
|
| for n_latent in (0, 5, 10, 15, 20): |
| for W in (None, 1, 2, 5): |
| if W is not None and n_latent <= W: |
| continue |
| base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(args.model_id)) |
| model = Coconut(base, latent_id, |
| tok.convert_tokens_to_ids("<|start-latent|>"), |
| tok.convert_tokens_to_ids("<|end-latent|>"), |
| tok.eos_token_id, backprop_depth=W).to(args.device) |
| ids, labels = build_batch(data, tok, latent_id, n_latent, args.bs, args.device) |
| pos = torch.arange(ids.shape[1], device=args.device).unsqueeze(0).expand(len(ids), -1) |
| attn = torch.ones_like(ids) |
| torch.cuda.reset_peak_memory_stats(args.device) |
| for i in range(args.iters): |
| if i == 2: |
| torch.cuda.synchronize(); t0 = time.time() |
| out = model.forward(ids, attn, labels, pos) |
| out.loss.backward() |
| model.zero_grad(set_to_none=True) |
| torch.cuda.synchronize() |
| ms = (time.time() - t0) / (args.iters - 2) * 1e3 |
| peak = torch.cuda.max_memory_allocated(args.device) / 2**20 |
| print(f"{n_latent:>9} {str(W):>7} {ms:>12.1f} {peak:>10.0f}") |
| del model, ids, labels |
| torch.cuda.empty_cache() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
