File size: 3,777 Bytes
8f46582 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | #!/usr/bin/env python3
"""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
# supervise a single node token, as dataset.expand_data does
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()
|