| """Deterministic Stage 1 training, validation, checkpointing, and teacher export.""" |
|
|
| import json |
| import math |
| import platform |
| import random |
| import subprocess |
| import time |
| from collections import Counter |
| from copy import deepcopy |
| from hashlib import sha256 |
| from importlib.metadata import PackageNotFoundError, version |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, Subset |
|
|
| from causalcelljepa.model import CellEncoder, CellJEPA, ema_momentum, jepa_loss, mask_gene_tokens |
| from causalcelljepa.resources import file_sha256, load_gmt_gene_indices |
|
|
|
|
| def stage1_split(cell_ids, validation_fraction, seed): |
| """Hash-split only admitted cells into fixed Stage 1 train/validation subsets.""" |
| threshold = int(validation_fraction * (1 << 256)) |
| validation = np.asarray( |
| [ |
| int.from_bytes(sha256(f"{seed}\0stage1-validation\0{cell_id}".encode()).digest(), "big") |
| < threshold |
| for cell_id in cell_ids |
| ], |
| dtype=bool, |
| ) |
| if not validation.any() or validation.all(): |
| raise ValueError("Stage 1 split must contain both training and validation cells") |
| return np.flatnonzero(~validation), np.flatnonzero(validation) |
|
|
|
|
| def stage1_split_report(dataset, validation_fraction, seed, replogle_manifest): |
| """Summarize and hash the exact admitted-cell split without storing cell identifiers.""" |
| cell_ids = [sample[2] for sample in dataset.samples] |
| train_indices, validation_indices = stage1_split(cell_ids, validation_fraction, seed) |
| splits = {} |
| for name, indices in (("train", train_indices), ("validation", validation_indices)): |
| selected = [dataset.samples[index] for index in indices] |
| ids = sorted(sample[2] for sample in selected) |
| splits[name] = { |
| "cells": len(selected), |
| "cell_ids_sha256": sha256("\n".join(ids).encode()).hexdigest(), |
| "context_counts": dict(sorted(Counter(sample[0] for sample in selected).items())), |
| "role_counts": dict(sorted(Counter(sample[3] for sample in selected).items())), |
| } |
| return { |
| "seed": seed, |
| "validation_fraction": validation_fraction, |
| "admission_policy": dataset.config["stage1"]["admission_policy"], |
| "admitted_cells": len(dataset), |
| "replogle_manifest_sha256": replogle_manifest["manifest_sha256"], |
| "hvg_sha256": replogle_manifest["genes"]["hvg_sha256"], |
| "splits": splits, |
| } |
|
|
|
|
| def write_stage1_split_manifest(path, dataset, stage1_config, replogle_manifest): |
| payload = stage1_split_report( |
| dataset, stage1_config["validation"]["fraction"], stage1_config["seed"], replogle_manifest |
| ) |
| payload["config_sha256"] = file_sha256("configs/stage1.yaml") |
| payload["manifest_sha256"] = sha256( |
| json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| return payload |
|
|
|
|
| def epoch_order(indices, seed, epoch): |
| """Return a replayable per-epoch order independent of global RNG state.""" |
| generator = np.random.default_rng(np.random.SeedSequence([seed, epoch])) |
| return generator.permutation(np.asarray(indices)).tolist() |
|
|
|
|
| def seed_everything(seed, deterministic_algorithms=True): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = deterministic_algorithms |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| torch.use_deterministic_algorithms(deterministic_algorithms) |
|
|
|
|
| def batch_views( |
| batch, |
| programs, |
| epoch, |
| seed, |
| padding_id, |
| device, |
| mask_range=(0.30, 0.60), |
| biological_probability=0.25, |
| minimum_program_genes=3, |
| ): |
| """Build deterministic per-cell masks and copy tensor-only views to a device.""" |
| masked_rows = [] |
| for row, cell_id in enumerate(batch["cell_id"]): |
| generator = torch.Generator().manual_seed( |
| int.from_bytes(sha256(f"{seed}\0{epoch}\0{cell_id}".encode()).digest()[:8], "little") |
| ) |
| masked_rows.append( |
| mask_gene_tokens( |
| batch["gene_ids"][row : row + 1], |
| batch["values"][row : row + 1], |
| batch["padding_mask"][row : row + 1], |
| programs=programs, |
| mask_range=mask_range, |
| biological_probability=biological_probability, |
| minimum_program_genes=minimum_program_genes, |
| padding_id=padding_id, |
| generator=generator, |
| ) |
| ) |
| keys = ("gene_ids", "values", "padding_mask") |
| non_blocking = device.type == "cuda" |
| student = { |
| key: torch.cat([masked[key] for masked in masked_rows]).to( |
| device, non_blocking=non_blocking |
| ) |
| for key in keys |
| } |
| teacher = {key: batch[key].to(device, non_blocking=non_blocking) for key in keys} |
| return student, teacher |
|
|
|
|
| def build_stage1_model(config): |
| stage1 = config["stage1"] |
| encoder = CellEncoder( |
| vocab_size=config["data"]["hvg_count"], |
| token_dim=stage1["token_dim"], |
| latent_queries=stage1["latent_queries"], |
| blocks=stage1["blocks"], |
| heads=stage1["heads"], |
| ffn_dim=stage1["ffn_dim"], |
| dropout=stage1["dropout"], |
| cell_dim=stage1["cell_dim"], |
| ) |
| return CellJEPA(encoder, stage1["predictor_hidden"]) |
|
|
|
|
| def learning_rate_schedule(total_steps, warmup_fraction, minimum_fraction): |
| warmup_steps = max(1, round(total_steps * warmup_fraction)) |
|
|
| def multiplier(step): |
| if step < warmup_steps: |
| return (step + 1) / warmup_steps |
| progress = min(1.0, (step - warmup_steps) / max(total_steps - warmup_steps, 1)) |
| cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) |
| return minimum_fraction + (1.0 - minimum_fraction) * cosine |
|
|
| return multiplier |
|
|
|
|
| def capture_rng_state(): |
| return { |
| "python": random.getstate(), |
| "numpy": np.random.get_state(), |
| "torch": torch.get_rng_state(), |
| "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [], |
| } |
|
|
|
|
| def restore_rng_state(state): |
| random.setstate(state["python"]) |
| np.random.set_state(state["numpy"]) |
| torch.set_rng_state(state["torch"].cpu()) |
| if state["cuda"] and torch.cuda.is_available(): |
| torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda"]]) |
|
|
|
|
| def _declared_manifest_hash(path): |
| payload = json.loads(Path(path).read_text()) |
| declared = payload.pop("manifest_sha256") |
| actual = sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() |
| if actual != declared: |
| raise ValueError(f"Manifest integrity failure for {path}: {actual} != {declared}") |
| return declared |
|
|
|
|
| def _git_state(): |
| try: |
| commit = subprocess.run( |
| ["git", "rev-parse", "HEAD"], check=True, capture_output=True, text=True |
| ).stdout.strip() |
| dirty = bool( |
| subprocess.run( |
| ["git", "status", "--porcelain"], check=True, capture_output=True, text=True |
| ).stdout |
| ) |
| return {"commit": commit, "dirty": dirty} |
| except (OSError, subprocess.CalledProcessError): |
| return {"commit": None, "dirty": None} |
|
|
|
|
| def _runtime_source_hash(): |
| paths = [ |
| *Path("src").rglob("*.py"), |
| *Path("scripts").glob("*.py"), |
| Path("pyproject.toml"), |
| Path("uv.lock"), |
| ] |
| digest = sha256() |
| for path in sorted(paths): |
| digest.update(str(path).encode()) |
| digest.update(b"\0") |
| digest.update(path.read_bytes()) |
| digest.update(b"\0") |
| return digest.hexdigest() |
|
|
|
|
| def _runtime_environment(): |
| packages = {} |
| for distribution in ("anndata", "h5py", "numpy", "PyYAML", "scipy", "torch"): |
| try: |
| packages[distribution] = version(distribution) |
| except PackageNotFoundError: |
| packages[distribution] = None |
| runtime = { |
| "python": platform.python_version(), |
| "platform": platform.platform(), |
| "packages": packages, |
| "torch_cuda": torch.version.cuda, |
| } |
| if torch.cuda.is_available(): |
| properties = torch.cuda.get_device_properties(torch.cuda.current_device()) |
| runtime["cuda_device"] = { |
| "name": properties.name, |
| "capability": [properties.major, properties.minor], |
| "total_memory_bytes": properties.total_memory, |
| } |
| else: |
| runtime["cuda_device"] = None |
| return runtime |
|
|
|
|
| def build_provenance( |
| replogle_config_path="configs/replogle.yaml", |
| stage1_config_path="configs/stage1.yaml", |
| replogle_manifest_path="manifests/replogle_v1.json", |
| go_manifest_path="manifests/go_bp_2026-08-05.json", |
| stage1_manifest_path="manifests/stage1_v1.json", |
| ): |
| go_manifest = json.loads(Path(go_manifest_path).read_text()) |
| actual_gmt_hash = file_sha256(go_manifest["output"]["gmt_path"]) |
| if actual_gmt_hash != go_manifest["output"]["gmt_sha256"]: |
| raise ValueError("GO BP GMT does not match its frozen manifest") |
| return { |
| "config_sha256": { |
| "replogle": file_sha256(replogle_config_path), |
| "stage1": file_sha256(stage1_config_path), |
| }, |
| "manifest_sha256": { |
| "replogle": _declared_manifest_hash(replogle_manifest_path), |
| "go_bp": _declared_manifest_hash(go_manifest_path), |
| "stage1_split": _declared_manifest_hash(stage1_manifest_path), |
| }, |
| "hvg_sha256": json.loads(Path(replogle_manifest_path).read_text())["genes"]["hvg_sha256"], |
| "go_bp_gmt_sha256": go_manifest["output"]["gmt_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
|
|
|
|
| def initial_train_state(): |
| return { |
| "epoch": 0, |
| "batch_in_epoch": 0, |
| "global_step": 0, |
| "best_validation_loss": float("inf"), |
| "best_validation_epoch": None, |
| "epochs_without_improvement": 0, |
| "complete": False, |
| "completion_reason": None, |
| } |
|
|
|
|
| def save_checkpoint(path, model, optimizer, scheduler, state, configuration, provenance): |
| """Atomically save every state required for exact optimizer-step resumption.""" |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| payload = { |
| "format_version": 1, |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict(), |
| "scheduler": scheduler.state_dict(), |
| "state": deepcopy(state), |
| "rng": capture_rng_state(), |
| "configuration": deepcopy(configuration), |
| "provenance": deepcopy(provenance), |
| } |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| torch.save(payload, temporary) |
| temporary.replace(path) |
| return path |
|
|
|
|
| def load_checkpoint(path, model, optimizer, scheduler, expected_provenance, device="cpu"): |
| payload = torch.load(path, map_location=device, weights_only=False) |
| if payload["format_version"] != 1: |
| raise ValueError(f"Unsupported checkpoint format {payload['format_version']}") |
| if payload["provenance"] != expected_provenance: |
| raise ValueError("Checkpoint provenance does not match the current frozen inputs/code") |
| model.load_state_dict(payload["model"]) |
| optimizer.load_state_dict(payload["optimizer"]) |
| scheduler.load_state_dict(payload["scheduler"]) |
| restore_rng_state(payload["rng"]) |
| return payload["state"], payload["configuration"] |
|
|
|
|
| def _configuration_for_resume(configuration): |
| normalized = deepcopy(configuration) |
| normalized["stage1"]["training"]["resume_from"] = None |
| return normalized |
|
|
|
|
| def export_canonical_teacher(path, model, state, replogle_config, provenance): |
| """Export only a completed EMA teacher as the frozen canonical state encoder.""" |
| if not state["complete"]: |
| raise ValueError("Canonical teacher export is forbidden before Stage 1 completion") |
| path = Path(path) |
| payload = { |
| "format_version": 1, |
| "frozen": True, |
| "cell_dim": replogle_config["stage1"]["cell_dim"], |
| "encoder_configuration": deepcopy(replogle_config["stage1"]), |
| "teacher": model.teacher.state_dict(), |
| "training_state": deepcopy(state), |
| "provenance": deepcopy(provenance), |
| } |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| torch.save(payload, temporary) |
| temporary.replace(path) |
| return path |
|
|
|
|
| def _json_log(path, event): |
| with Path(path).open("a") as handle: |
| handle.write(json.dumps(event, sort_keys=True) + "\n") |
|
|
|
|
| def _data_loader( |
| dataset, |
| indices, |
| batch_size, |
| num_workers, |
| seed, |
| pin_memory=False, |
| persistent_workers=False, |
| prefetch_factor=None, |
| ): |
| generator = torch.Generator().manual_seed(seed) |
| options = {} |
| if num_workers: |
| options["persistent_workers"] = persistent_workers |
| options["prefetch_factor"] = prefetch_factor |
| return DataLoader( |
| Subset(dataset, list(indices)), |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=num_workers, |
| pin_memory=pin_memory, |
| generator=generator, |
| **options, |
| ) |
|
|
|
|
| @torch.no_grad() |
| def validate(model, loader, programs, config, epoch, seed, device): |
| model.eval() |
| totals, cells = {}, 0 |
| stage1 = config["stage1"] |
| for batch in loader: |
| student, teacher = batch_views( |
| batch, |
| programs, |
| epoch, |
| seed, |
| model.online.padding_id, |
| device, |
| stage1["random_mask_range"], |
| stage1["biological_mask_probability"], |
| stage1["minimum_program_genes"], |
| ) |
| prediction, target, context = model(student, teacher) |
| losses = jepa_loss( |
| prediction, |
| target, |
| context, |
| stage1["variance_weight"], |
| stage1["covariance_weight"], |
| ) |
| batch_size = len(batch["cell_id"]) |
| cells += batch_size |
| for name, value in losses.items(): |
| totals[name] = totals.get(name, 0.0) + float(value) * batch_size |
| model.train() |
| return {name: value / cells for name, value in totals.items()} |
|
|
|
|
| def train_stage1( |
| dataset, |
| replogle_config, |
| stage1_config, |
| replogle_manifest, |
| go_manifest, |
| device, |
| max_steps=None, |
| ): |
| """Train or exactly resume Stage 1; a max_steps cap is reserved for CPU smoke checks.""" |
| training = stage1_config["training"] |
| seed = stage1_config["seed"] |
| if training["optimizer"] != "AdamW": |
| raise ValueError(f"Unsupported Stage 1 optimizer {training['optimizer']}") |
| if replogle_config["stage1"]["teacher_mode"] != "eval": |
| raise ValueError("The EMA teacher must remain in eval mode") |
| if replogle_config["stage1"]["prediction_loss"] != "normalized_smooth_l1": |
| raise ValueError("Unsupported Stage 1 prediction loss") |
| seed_everything(seed, training["deterministic_algorithms"]) |
| train_indices, validation_indices = stage1_split( |
| [sample[2] for sample in dataset.samples], stage1_config["validation"]["fraction"], seed |
| ) |
| expected_split = json.loads(Path(stage1_config["validation"]["manifest_path"]).read_text()) |
| actual_split = stage1_split_report( |
| dataset, stage1_config["validation"]["fraction"], seed, replogle_manifest |
| ) |
| if actual_split != { |
| key: value |
| for key, value in expected_split.items() |
| if key not in {"config_sha256", "manifest_sha256"} |
| }: |
| raise ValueError("Runtime Stage 1 admission/split does not match its frozen manifest") |
| programs = load_gmt_gene_indices( |
| go_manifest["output"]["gmt_path"], replogle_manifest["genes"]["hvg_gene_names"] |
| ) |
| model = build_stage1_model(replogle_config).to(device) |
| parameters = [*model.online.parameters(), *model.predictor.parameters()] |
| optimizer = torch.optim.AdamW( |
| parameters, lr=training["learning_rate"], weight_decay=training["weight_decay"] |
| ) |
| steps_per_epoch = math.ceil(len(train_indices) / training["batch_size"]) |
| total_steps = training["epochs"] * steps_per_epoch |
| scheduler = torch.optim.lr_scheduler.LambdaLR( |
| optimizer, |
| learning_rate_schedule( |
| total_steps, |
| training["warmup_fraction"], |
| training["minimum_learning_rate_fraction"], |
| ), |
| ) |
| provenance = build_provenance() |
| configuration = {"replogle": deepcopy(replogle_config), "stage1": deepcopy(stage1_config)} |
| output = Path(training["output_directory"]) |
| output.mkdir(parents=True, exist_ok=True) |
| latest = output / "latest.pt" |
| best = output / "best.pt" |
| state = initial_train_state() |
| if training["resume_from"]: |
| state, saved_configuration = load_checkpoint( |
| training["resume_from"], model, optimizer, scheduler, provenance, device |
| ) |
| if _configuration_for_resume(saved_configuration) != _configuration_for_resume( |
| configuration |
| ): |
| raise ValueError("Checkpoint configuration does not match the current run") |
| validation_loader = _data_loader( |
| dataset, |
| validation_indices, |
| training["batch_size"], |
| training["num_workers"], |
| seed + 1, |
| training["pin_memory"] and device.type == "cuda", |
| training["persistent_workers"], |
| training["prefetch_factor"], |
| ) |
| started = time.perf_counter() |
| model.train() |
| for epoch in range(state["epoch"], training["epochs"]): |
| order = epoch_order(train_indices, seed, epoch) |
| start = state["batch_in_epoch"] * training["batch_size"] |
| loader = _data_loader( |
| dataset, |
| order[start:], |
| training["batch_size"], |
| training["num_workers"], |
| seed + 2 + epoch, |
| training["pin_memory"] and device.type == "cuda", |
| training["persistent_workers"], |
| training["prefetch_factor"], |
| ) |
| for offset, batch in enumerate(loader, start=state["batch_in_epoch"]): |
| student, teacher = batch_views( |
| batch, |
| programs, |
| epoch, |
| seed, |
| model.online.padding_id, |
| device, |
| replogle_config["stage1"]["random_mask_range"], |
| replogle_config["stage1"]["biological_mask_probability"], |
| replogle_config["stage1"]["minimum_program_genes"], |
| ) |
| optimizer.zero_grad(set_to_none=True) |
| prediction, target, context = model(student, teacher) |
| losses = jepa_loss( |
| prediction, |
| target, |
| context, |
| replogle_config["stage1"]["variance_weight"], |
| replogle_config["stage1"]["covariance_weight"], |
| ) |
| if not all(torch.isfinite(value) for value in losses.values()): |
| raise FloatingPointError( |
| f"Non-finite Stage 1 loss at optimizer step {state['global_step']}" |
| ) |
| losses["loss"].backward() |
| gradient_norm = torch.nn.utils.clip_grad_norm_( |
| parameters, training["gradient_clip_norm"], error_if_nonfinite=True |
| ) |
| optimizer.step() |
| ema_start, ema_end = replogle_config["stage1"]["ema_momentum"] |
| momentum = ema_momentum(state["global_step"], total_steps, start=ema_start, end=ema_end) |
| model.update_teacher(momentum) |
| scheduler.step() |
| state["global_step"] += 1 |
| state["batch_in_epoch"] = offset + 1 |
| event = { |
| "event": "train_step", |
| "epoch": epoch, |
| "global_step": state["global_step"], |
| "cells": len(batch["cell_id"]), |
| "loss": float(losses["loss"].detach()), |
| "prediction": float(losses["prediction"].detach()), |
| "variance": float(losses["variance"].detach()), |
| "covariance": float(losses["covariance"].detach()), |
| "gradient_norm": float(gradient_norm), |
| "ema_momentum": momentum, |
| "learning_rate": scheduler.get_last_lr()[0], |
| } |
| _json_log(output / "training.jsonl", event) |
| if state["global_step"] % training["checkpoint_every_steps"] == 0: |
| save_checkpoint( |
| latest, model, optimizer, scheduler, state, configuration, provenance |
| ) |
| if max_steps is not None and state["global_step"] >= max_steps: |
| save_checkpoint( |
| latest, model, optimizer, scheduler, state, configuration, provenance |
| ) |
| return ( |
| model, |
| state, |
| { |
| "train_cells": len(train_indices), |
| "validation_cells": len(validation_indices), |
| "programs": len(programs), |
| "elapsed_seconds": time.perf_counter() - started, |
| "checkpoint": str(latest), |
| }, |
| ) |
| validation = validate( |
| model, |
| validation_loader, |
| programs, |
| replogle_config, |
| stage1_config["validation"]["mask_epoch"], |
| seed, |
| device, |
| ) |
| _json_log( |
| output / "training.jsonl", |
| { |
| "event": "validation", |
| "epoch": epoch, |
| "global_step": state["global_step"], |
| **validation, |
| }, |
| ) |
| improved = validation["loss"] < state["best_validation_loss"] |
| state["best_validation_loss"] = min(state["best_validation_loss"], validation["loss"]) |
| if improved: |
| state["best_validation_epoch"] = epoch |
| state["epochs_without_improvement"] = ( |
| 0 if improved else state["epochs_without_improvement"] + 1 |
| ) |
| state["epoch"], state["batch_in_epoch"] = epoch + 1, 0 |
| save_checkpoint(latest, model, optimizer, scheduler, state, configuration, provenance) |
| if improved: |
| save_checkpoint(best, model, optimizer, scheduler, state, configuration, provenance) |
| if state["epochs_without_improvement"] >= training["early_stopping_patience"]: |
| state["complete"], state["completion_reason"] = True, "early_stopping" |
| break |
| else: |
| state["complete"], state["completion_reason"] = True, "configured_epochs" |
| save_checkpoint(latest, model, optimizer, scheduler, state, configuration, provenance) |
| best_payload = torch.load(best, map_location=device, weights_only=False) |
| best_teacher = { |
| key.removeprefix("teacher."): value |
| for key, value in best_payload["model"].items() |
| if key.startswith("teacher.") |
| } |
| model.teacher.load_state_dict(best_teacher) |
| export_canonical_teacher( |
| output / "canonical_ema_teacher.pt", model, state, replogle_config, provenance |
| ) |
| return ( |
| model, |
| state, |
| { |
| "train_cells": len(train_indices), |
| "validation_cells": len(validation_indices), |
| "programs": len(programs), |
| "elapsed_seconds": time.perf_counter() - started, |
| "checkpoint": str(latest), |
| }, |
| ) |
|
|
