| |
| """Fair, held-out pilot for Laya-Bio original-vocabulary vs M2. |
| |
| This pilot is still too small to be a paper result. It exists to validate the |
| comparison protocol after the GPU smoke tests: |
| |
| * train and validation are cleaned from the two-task audit policy; |
| * M1 and M2 each start from the same untouched Laya checkpoint; |
| * the same natural-language/sequence serialization is used in both conditions; |
| * the M2 vocabulary uses modality-scoped tokens (`▶` DNA, `◆` protein) built |
| from the existing source BPE tokenizers, preventing bare biological tokens |
| from matching English prompt text; |
| * a held-out validation accuracy/NLL is computed only as a pilot diagnostic; |
| * finite-loss/gradient checks and checkpoint reload are explicit. |
| |
| No CPT data is read. The raw JSONL files remain unchanged. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import random |
| import re |
| import time |
| from collections import Counter, defaultdict |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
| from tokenizers import AddedToken, Tokenizer |
| from transformers import AutoTokenizer |
|
|
| from laya_gpu_validation import ( |
| TASKS, |
| collate, |
| extract_context_and_sequence, |
| expand_tokenizer_and_model, |
| make_items, |
| move_batch, |
| ) |
|
|
|
|
| def args() -> argparse.Namespace: |
| p = argparse.ArgumentParser() |
| p.add_argument("--model-dir", default="artifacts/laya_model") |
| p.add_argument("--output-dir", default="artifacts/laya_fair_pilot") |
| p.add_argument("--laya-repo", default="vendor/laya") |
| p.add_argument("--train-per-class", type=int, default=100) |
| p.add_argument("--val-per-class", type=int, default=50) |
| p.add_argument("--steps", type=int, default=96) |
| p.add_argument("--micro-batch", type=int, default=4) |
| p.add_argument("--grad-accum", type=int, default=8) |
| p.add_argument("--expanded-tokens", type=int, default=64) |
| p.add_argument("--max-len", type=int, default=1024) |
| p.add_argument("--head-max-len", type=int, default=256) |
| p.add_argument("--seed", type=int, default=20260922) |
| return p.parse_args() |
|
|
|
|
| def load_laya(repo: str): |
| import sys |
|
|
| sys.path.insert(0, repo) |
| from laya.common import QTYPES, build_model, build_sequence |
|
|
| return QTYPES, build_model, build_sequence |
|
|
|
|
| def raw_rows(path: Path, split: str) -> list[dict[str, Any]]: |
| out = [] |
| with path.open(encoding="utf-8") as fh: |
| for line in fh: |
| rec = json.loads(line) |
| if rec.get("split") == split: |
| out.append(rec) |
| return out |
|
|
|
|
| def seq_key(rec: dict[str, Any]) -> str: |
| return extract_context_and_sequence(rec)[1].strip().upper() |
|
|
|
|
| def clean_train_val(path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: |
| train_raw, val_raw = raw_rows(path, "train"), raw_rows(path, "val") |
| grouped: dict[str, set[str]] = defaultdict(set) |
| for rec in train_raw + val_raw: |
| grouped[seq_key(rec)].add(str(rec["answer_short"])) |
| conflict = {k for k, labels in grouped.items() if len(labels) > 1} |
| train, train_seen = [], set() |
| for rec in train_raw: |
| key = seq_key(rec) |
| if key in conflict or key in train_seen: |
| continue |
| train_seen.add(key) |
| train.append(rec) |
| val, val_seen = [], set() |
| for rec in val_raw: |
| key = seq_key(rec) |
| if key in conflict or key in train_seen or key in val_seen: |
| continue |
| val_seen.add(key) |
| val.append(rec) |
| return train, val, { |
| "raw_train": len(train_raw), |
| "raw_val": len(val_raw), |
| "clean_train": len(train), |
| "clean_val": len(val), |
| "conflict_sequences_removed": len(conflict), |
| "train_val_exact_overlap_after_cleanup": len(train_seen & val_seen), |
| } |
|
|
|
|
| def balanced(rows: list[dict[str, Any]], per_class: int, seed: int) -> list[dict[str, Any]]: |
| groups: dict[str, list[dict[str, Any]]] = defaultdict(list) |
| for rec in rows: |
| groups[str(rec["answer_short"])].append(rec) |
| rng = random.Random(seed) |
| out = [] |
| for label in sorted(groups): |
| group = list(groups[label]) |
| rng.shuffle(group) |
| out.extend(group[:per_class]) |
| rng.shuffle(out) |
| return out |
|
|
|
|
| def modality(rec: dict[str, Any]) -> str: |
| return "dna" if "dna" in rec.get("modality", []) else "protein" |
|
|
|
|
| def source_serializers() -> dict[str, Tokenizer]: |
| return { |
| "dna": Tokenizer.from_file("data/02_vocab/dna_bpe_20k.json"), |
| "protein": Tokenizer.from_file("data/02_vocab/protein_bpe_8k.json"), |
| } |
|
|
|
|
| def source_state_fn(source: dict[str, Tokenizer]): |
| def state_fn(rec: dict[str, Any]) -> str: |
| context, sequence = extract_context_and_sequence(rec) |
| kind = modality(rec) |
| pieces = source[kind].encode(sequence).tokens |
| prefix = "▶" if kind == "dna" else "◆" |
| |
| |
| |
| serialized = prefix + prefix.join(pieces) |
| return f"Task context: {context}\nSequence: {serialized}" |
|
|
| return state_fn |
|
|
|
|
| def build_candidates(rows, source: dict[str, Tokenizer], base_tok, n_new: int): |
| |
| |
| counts = {"dna": Counter(), "protein": Counter()} |
| prompt_texts = [] |
| for rec in rows: |
| context, _ = extract_context_and_sequence(rec) |
| prompt_texts.append(context) |
| prompt_texts.extend(str(x) for x in rec["choices"]) |
| kind = modality(rec) |
| prefix = "▶" if kind == "dna" else "◆" |
| for piece in source[kind].encode(extract_context_and_sequence(rec)[1]).tokens: |
| counts[kind][prefix + piece] += 1 |
| ranked = {} |
| for kind, counter in counts.items(): |
| ranked[kind] = [ |
| token for token, count in sorted( |
| counter.items(), key=lambda kv: (-kv[1], -len(kv[0]), kv[0]) |
| ) |
| if len(token) >= 4 and not any(token in text for text in prompt_texts) |
| ] |
| n_dna = n_new // 2 |
| n_protein = n_new - n_dna |
| candidates = [] |
| dna, protein = ranked["dna"][: n_dna * 4], ranked["protein"][: n_protein * 4] |
| for i in range(max(len(dna), len(protein))): |
| if i < len(dna): |
| candidates.append(dna[i]) |
| if i < len(protein): |
| candidates.append(protein[i]) |
| |
| |
| return candidates, ranked |
|
|
|
|
| def load_model(model_dir, cfg, build_model, device): |
| model = build_model(cfg, encoder_dir=str(model_dir / "encoder")) |
| model.load_state_dict(load_file(str(model_dir / "model.safetensors")), strict=True) |
| model.encoder.gradient_checkpointing_enable( |
| gradient_checkpointing_kwargs={"use_reentrant": False} |
| ) |
| for name, parameter in model.named_parameters(): |
| if name.startswith("act_head."): |
| parameter.requires_grad_(False) |
| model.to(device) |
| return model |
|
|
|
|
| def train(model, items, tokenizer, device, steps, micro_batch, grad_accum, seed, tag): |
| if steps % grad_accum: |
| raise ValueError("--steps must be divisible by --grad-accum for this pilot") |
| model.train() |
| optimizer = torch.optim.AdamW( |
| [p for p in model.parameters() if p.requires_grad], lr=2e-5, weight_decay=0.01 |
| ) |
| bf16 = torch.cuda.is_bf16_supported() |
| amp_dtype = torch.bfloat16 if bf16 else torch.float16 |
| scaler = torch.amp.GradScaler("cuda", enabled=not bf16) |
| order = list(range(len(items))) |
| rng = random.Random(seed) |
| cursor = 0 |
| losses, updates, examples = [], 0, 0 |
| optimizer.zero_grad(set_to_none=True) |
| torch.cuda.reset_peak_memory_stats() |
| start_time = time.time() |
| for step in range(steps): |
| if cursor + micro_batch > len(order): |
| rng.shuffle(order) |
| cursor = 0 |
| chunk = [items[i] for i in order[cursor: cursor + micro_batch]] |
| cursor += micro_batch |
| batch = move_batch(collate(chunk, tokenizer.pad_token_id), device) |
| with torch.autocast("cuda", dtype=amp_dtype): |
| logits, _ = model( |
| batch["input_ids"], batch["attention_mask"], batch["marker_pos"], |
| batch["marker_mask"], batch["qtype"], |
| ) |
| loss = torch.nn.functional.cross_entropy( |
| logits.masked_fill(~batch["marker_mask"], -1e4), batch["label"] |
| ) |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f"{tag}: non-finite loss at step {step + 1}") |
| scaler.scale(loss / grad_accum).backward() |
| if (step + 1) % grad_accum == 0: |
| scaler.unscale_(optimizer) |
| grad_norm = torch.nn.utils.clip_grad_norm_( |
| [p for p in model.parameters() if p.requires_grad], 1.0, |
| error_if_nonfinite=True, |
| ) |
| if not torch.isfinite(grad_norm): |
| raise FloatingPointError(f"{tag}: non-finite gradient at step {step + 1}") |
| scaler.step(optimizer) |
| scaler.update() |
| optimizer.zero_grad(set_to_none=True) |
| updates += 1 |
| losses.append(float(loss.detach().cpu())) |
| examples += len(chunk) |
| return { |
| "tag": tag, |
| "steps": steps, |
| "optimizer_updates": updates, |
| "examples_consumed": examples, |
| "loss_first": losses[0], |
| "loss_last": losses[-1], |
| "loss_min": min(losses), |
| "seconds": time.time() - start_time, |
| "max_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, |
| "max_reserved_gib": torch.cuda.max_memory_reserved() / 2**30, |
| "amp_dtype": str(amp_dtype), |
| "finite": True, |
| } |
|
|
|
|
| def evaluate(model, items, tokenizer, device, batch_size=8): |
| model.eval() |
| records = [] |
| with torch.no_grad(): |
| for start in range(0, len(items), batch_size): |
| chunk = items[start: start + batch_size] |
| batch = move_batch(collate(chunk, tokenizer.pad_token_id), device) |
| with torch.autocast("cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16): |
| logits, _ = model( |
| batch["input_ids"], batch["attention_mask"], batch["marker_pos"], |
| batch["marker_mask"], batch["qtype"], |
| ) |
| probs = torch.softmax(logits.masked_fill(~batch["marker_mask"], -1e4).float(), -1) |
| for i, item in enumerate(chunk): |
| k = len(item["markers"]) |
| p = probs[i, :k] |
| gold = int(item["label"]) |
| records.append({ |
| "task": item["task"], |
| "correct": int(int(p.argmax()) == gold), |
| "nll": float(-torch.log(p[gold].clamp_min(1e-9))), |
| }) |
| by_task = defaultdict(list) |
| for record in records: |
| by_task[record["task"]].append(record) |
| result = {} |
| for task, values in sorted(by_task.items()): |
| result[task] = { |
| "n": len(values), |
| "accuracy": sum(x["correct"] for x in values) / len(values), |
| "nll": sum(x["nll"] for x in values) / len(values), |
| } |
| result["all"] = { |
| "n": len(records), |
| "accuracy": sum(x["correct"] for x in records) / len(records), |
| "nll": sum(x["nll"] for x in records) / len(records), |
| } |
| return result |
|
|
|
|
| def grad_by_task(model, items, tokenizer, device, new_ids): |
| out = {} |
| model.train() |
| for task in sorted({x["task"] for x in items}): |
| covered = [x for x in items if x["task"] == task and any(i in new_ids for i in x["ids"])] |
| if not covered: |
| out[task] = None |
| continue |
| model.zero_grad(set_to_none=True) |
| batch = move_batch(collate(covered[:4], tokenizer.pad_token_id), device) |
| with torch.autocast("cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16): |
| logits, _ = model( |
| batch["input_ids"], batch["attention_mask"], batch["marker_pos"], |
| batch["marker_mask"], batch["qtype"], |
| ) |
| loss = torch.nn.functional.cross_entropy( |
| logits.masked_fill(~batch["marker_mask"], -1e4), batch["label"] |
| ) |
| loss.backward() |
| grad = model.encoder.get_input_embeddings().weight.grad |
| out[task] = None if grad is None else float(grad[new_ids].norm().detach().cpu()) |
| model.zero_grad(set_to_none=True) |
| return out |
|
|
|
|
| def candidate_text_ids(tokenizer, rows): |
| texts = [] |
| for rec in rows: |
| context, _ = extract_context_and_sequence(rec) |
| texts.extend([context, *map(str, rec["choices"]), "Choose the correct biological label for the sequence."]) |
| return [tokenizer(x, add_special_tokens=False)["input_ids"] for x in texts] |
|
|
|
|
| def main(): |
| a = args() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA is not available") |
| random.seed(a.seed) |
| torch.manual_seed(a.seed) |
| device = torch.device("cuda:0") |
| output = Path(a.output_dir) |
| output.mkdir(parents=True, exist_ok=True) |
| model_dir = Path(a.model_dir) |
| qtypes, build_model, build_sequence = load_laya(a.laya_repo) |
| source = source_serializers() |
| state_fn = source_state_fn(source) |
|
|
| train_rows, val_rows, cleanup = {}, {}, {} |
| for task, spec in TASKS.items(): |
| tr, va, stats = clean_train_val(Path(spec["file"])) |
| train_rows[task], val_rows[task], cleanup[task] = tr, va, stats |
| train_selected = sum( |
| (balanced(train_rows[t], a.train_per_class, a.seed + i) for i, t in enumerate(TASKS)), [] |
| ) |
| val_selected = sum( |
| (balanced(val_rows[t], a.val_per_class, a.seed + 100 + i) for i, t in enumerate(TASKS)), [] |
| ) |
|
|
| base_tok = AutoTokenizer.from_pretrained(model_dir / "tokenizer") |
| with (model_dir / "rl_agent_config.json").open() as fh: |
| cfg = json.load(fh) |
| cfg.update({"max_len": a.max_len, "head_max_len": a.head_max_len, "gradient_checkpointing": True}) |
| base_weights = load_file(str(model_dir / "model.safetensors")) |
|
|
| def fresh_model(): |
| model = build_model(cfg, encoder_dir=str(model_dir / "encoder")) |
| model.load_state_dict(base_weights, strict=True) |
| model.encoder.gradient_checkpointing_enable( |
| gradient_checkpointing_kwargs={"use_reentrant": False} |
| ) |
| for name, parameter in model.named_parameters(): |
| if name.startswith("act_head."): |
| parameter.requires_grad_(False) |
| model.to(device) |
| return model |
|
|
| |
| m1_tok = AutoTokenizer.from_pretrained(model_dir / "tokenizer") |
| m1_train, _ = make_items(train_selected, m1_tok, build_sequence, qtypes, a.max_len, a.head_max_len, state_fn) |
| m1_val, _ = make_items(val_selected, m1_tok, build_sequence, qtypes, a.max_len, a.head_max_len, state_fn) |
| m1 = fresh_model() |
| m1_train_result = train(m1, m1_train, m1_tok, device, a.steps, a.micro_batch, a.grad_accum, a.seed, "M1-original") |
| m1_eval = evaluate(m1, m1_val, m1_tok, device) |
| del m1 |
| torch.cuda.empty_cache() |
| gc.collect() |
|
|
| |
| m2_tok = AutoTokenizer.from_pretrained(model_dir / "tokenizer") |
| candidates, ranked = build_candidates(train_selected, source, m2_tok, a.expanded_tokens) |
| m2 = fresh_model() |
| chosen, old_size = expand_tokenizer_and_model(m2_tok, m2, candidates, a.expanded_tokens) |
| m2_train, m2_train_lengths = make_items(train_selected, m2_tok, build_sequence, qtypes, a.max_len, a.head_max_len, state_fn) |
| m2_val, m2_val_lengths = make_items(val_selected, m2_tok, build_sequence, qtypes, a.max_len, a.head_max_len, state_fn) |
| m1_train_lengths = [x["n_tokens"] for x in m1_train] |
| m1_val_lengths = [x["n_tokens"] for x in m1_val] |
| m1_prompt_ids = candidate_text_ids(base_tok, val_selected) |
| m2_prompt_ids = candidate_text_ids(m2_tok, val_selected) |
| prompt_changed = sum(x != y for x, y in zip(m1_prompt_ids, m2_prompt_ids)) |
| m2_train_result = train(m2, m2_train, m2_tok, device, a.steps, a.micro_batch, a.grad_accum, a.seed, "M2-expanded") |
| m2_eval = evaluate(m2, m2_val, m2_tok, device) |
| new_ids = list(range(old_size, len(m2_tok))) |
| modality_grad = grad_by_task(m2, m2_train, m2_tok, device, new_ids) |
|
|
| |
| |
| m2.eval() |
| checkpoint = output / "m2_checkpoint" |
| checkpoint.mkdir(parents=True, exist_ok=True) |
| save_file({k: v.detach().contiguous().cpu() for k, v in m2.state_dict().items()}, str(checkpoint / "model.safetensors")) |
| m2.encoder.config.save_pretrained(checkpoint / "encoder") |
| m2_tok.save_pretrained(checkpoint / "tokenizer") |
| (checkpoint / "rl_agent_config.json").write_text(json.dumps(cfg, indent=2)) |
| disk_tok = AutoTokenizer.from_pretrained(checkpoint / "tokenizer") |
| disk_val, _ = make_items(val_selected, disk_tok, build_sequence, qtypes, a.max_len, a.head_max_len, state_fn) |
| reload_ids_match = all(x["ids"] == y["ids"] for x, y in zip(m2_val, disk_val)) |
| reloaded = load_laya(a.laya_repo)[1](cfg, encoder_dir=str(checkpoint / "encoder")) |
| reloaded.load_state_dict(load_file(str(checkpoint / "model.safetensors")), strict=True) |
| reloaded.to(device).eval() |
| reload_eval = evaluate(reloaded, disk_val, disk_tok, device) |
| del reloaded |
|
|
| summary = { |
| "pilot_only": True, |
| "no_cpt": True, |
| "gpu": torch.cuda.get_device_name(device), |
| "vram_gib": torch.cuda.get_device_properties(device).total_memory / 2**30, |
| "cuda_version": torch.version.cuda, |
| "transformers_version": __import__("transformers").__version__, |
| "laya_source_revision": "573e5b62696ba441230cd6be71d593331b5d23af", |
| "model_revision": "f9ab0b228f0fc0f14d873dbc99038f135c2da1b2", |
| "cleanup": cleanup, |
| "n_train": len(train_selected), |
| "n_val": len(val_selected), |
| "train_tasks": sorted({x["task_id"] for x in train_selected}), |
| "conditions_share_initial_checkpoint": True, |
| "steps": a.steps, |
| "micro_batch": a.micro_batch, |
| "grad_accum": a.grad_accum, |
| "m1_training": m1_train_result, |
| "m1_val": m1_eval, |
| "m2_training": m2_train_result, |
| "m2_val": m2_eval, |
| "m2_reload_val": reload_eval, |
| "reload_tokenizer_ids_match": reload_ids_match, |
| "base_vocab_size": len(base_tok), |
| "expanded_vocab_size": len(m2_tok), |
| "new_tokens": chosen, |
| "candidate_counts_by_source": {k: len(v) for k, v in ranked.items()}, |
| "prompt_only_tokenization_changed_examples": prompt_changed, |
| "token_lengths": { |
| "m1_train_p50": sorted(m1_train_lengths)[len(m1_train_lengths) // 2], |
| "m2_train_p50": sorted(m2_train_lengths)[len(m2_train_lengths) // 2], |
| "m1_val_p50": sorted(m1_val_lengths)[len(m1_val_lengths) // 2], |
| "m2_val_p50": sorted(m2_val_lengths)[len(m2_val_lengths) // 2], |
| "m1_train_max": max(m1_train_lengths), |
| "m2_train_max": max(m2_train_lengths), |
| "m1_val_max": max(m1_val_lengths), |
| "m2_val_max": max(m2_val_lengths), |
| }, |
| "new_embedding_grad_norm_by_task": modality_grad, |
| "source_tokenizer_isolation": "modality-prefixed source BPE pieces serialized with ▶/◆", |
| "formal_paper_result": False, |
| } |
| (output / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2)) |
| print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
