07. Train the multi-dataset Bolero model (ATAC)¶
The single-dataset page trained one Bolero model on ChromiumPBMC. The model in the paper is trained jointly across the entire Bolero-10M corpus — dozens of single-cell datasets spanning several mammalian genomes — so one network predicts cell-state–specific accessibility for every dataset at once. This page shows how that multi-dataset model is configured and what is different about it.
This notebook runs in the multi-dataset training directory, not in place under
docs/. Unlike the single-dataset tutorial (which runs on the small ChromiumPBMC artifacts), the multi-dataset model needs the full corpus (every dataset's parquet store + pseudobulk records) plus the intermediate checkpoints from the earlier training stages. All paths below are relative to that directory.
The full recipe is staged. Training the final model is done in stages, and the earlier stages are intentionally left out of these docs to keep them readable:
- init — register every dataset, build the shared model, train briefly so all datasets share one backbone. Produces the init config (which carries the whole dataset registry).
- dataset-specific — freeze the shared model and fine-tune each dataset's own conditioning encoder.
- combine — merge the per-dataset encoders into a single checkpoint
(
...PostDatasetFineTune.pt). - final — this page: load a frozen Flashzoi base + that combined checkpoint + the dataset registry, and train everything jointly.
Because a real final run is heavy (200 epochs, millions of steps, multi-GPU, days of compute) we
do not call trainer.train() here. Instead we build the model exactly as training would, then
load the already-trained final checkpoint and inspect what makes a multi-dataset model different.
Setup¶
import json
import os
import collections
import pandas as pd
from bolero import init
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer
os.environ["WANDB_MODE"] = (
"disabled" # demo only builds/loads a model; no training, no logging
)
init(num_cpus=16, object_store_memory_gb=48)
2026-07-13 12:09:53,684 INFO worker.py:1781 -- Started a local Ray instance.
# --- Configuration ---------------------------------------------------------------
MODEL_NAME = "Bolero10M"
FOLD = 0
# Stage-1 (init) config carries the dataset registry the whole model is built from.
INIT_CONFIG = f"model/250909-{MODEL_NAME}_atac_fold0-init.config.json"
# Combined per-dataset checkpoint (stage 3) — the starting weights for final training.
POST_DATASET_CKPT = f"250909-{MODEL_NAME}_atac_fold0-PostDatasetFineTune.pt"
# The trained final ATAC checkpoint we load for the demo (result of stage 4).
FINAL_CKPT = f"model/250909-{MODEL_NAME}_atac_fold0-final.lora.best_checkpoint.pt"
# Frozen Flashzoi backbone (same base checkpoint as the single-dataset model).
BASE_CKPT = (
f"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_{FOLD}.renamed.pt"
)
The dataset registry — what makes it multi-dataset¶
A single-dataset model takes one pseudobulk file. A multi-dataset model takes a dataset_records
registry: a dict mapping each dataset to the exact artifacts you built in the
data section — its 32 bp
data_path (parquet store) and its pseudobulk_path (coverage-matched pseudobulk records) — plus
the genome it maps to and a dataset_sample_weight. The registry was assembled during the
init stage from bolerodata.DATASETS, so we just read it back out of the init config.
with open(INIT_CONFIG) as f:
init_cfg = json.load(f)
dataset_records = init_cfg["dataset_records"]
# one row per dataset: which genome, which sampling weight
registry = pd.DataFrame(
{
k: {"genome": v["genome"], "weight": v["dataset_sample_weight"]}
for k, v in dataset_records.items()
}
).T
print(f"{len(dataset_records)} datasets across {registry['genome'].nunique()} genomes")
registry["genome"].value_counts()
41 datasets across 7 genomes
genome mm10 20 hg38 13 rheMac10 3 calJac4 2 mCalJac1.pat.X 1 panPan3 1 monDom5.split 1 Name: count, dtype: int64
# an example record: exactly a single dataset's parquet store + pseudobulk records + genome
example_key = next(iter(dataset_records))
print(example_key)
dataset_records[example_key]
Zu2023Nature+NonN
{'pseudobulk_path': '/large_storage/zhoulab/hanliu/wmb/standard/embedding/coverage_pseudobulk_with_condition/Zu2023Nature+NonN/pseudobulk_records_and_cond.cov5000000.joblib',
'data_path': '/large_storage/zhoulab/hanliu/wmb/standard/metacell_dataset/parquet/Zu2023Nature+NonN/Zu2023Nature-MetaCell-32bp',
'genome': 'mm10',
'dataset_sample_weight': 1}
Build the training config (real parameters)¶
The config is the single-dataset config with a few multi-dataset changes. The important keys:
dataset_records— the registry above (replaces the singlepseudobulk_records).base_checkpoint_path— the frozen Flashzoi fold-0 backbone (same as single-dataset).lora_checkpoint_path— the combined per-dataset checkpoint from stage 3; final training starts from it rather than from scratch.signal_model=True, and the real training schedule:max_epochs=200,lr_total_steps=4_000_000,save_state_every_n_epoch=25.
emb_input_features and cond_emb_dim are not set by hand here — for a multi-dataset model they
are computed per dataset when the dataset is built (each dataset has its own embedding width), and
merged into the config automatically. These are the real parameters used for the paper model; a
quick smoke-test would instead set max_epochs=3, warmup_steps=10, train_batches=50, val_batches=10
(as in the single-dataset page).
config = {
# run identity + outputs
"output_dir": "model",
"savename": f"250909-{MODEL_NAME}_atac_fold{FOLD}-final",
"wandb_project": "250909-BorzoiMulti",
"wandb_job_type": "lora",
"wandb_name": f"250909-{MODEL_NAME}_atac_fold{FOLD}-final",
# data: the whole registry instead of one pseudobulk file
"dataset_records": dataset_records,
"fold_split_id": FOLD,
# model
"base_checkpoint_path": BASE_CKPT,
"signal_model": True,
"lora_checkpoint_path": POST_DATASET_CKPT, # start from the combined per-dataset checkpoint
# real training schedule (see note above)
"max_epochs": 200,
"lr_total_steps": 4_000_000,
"save_state_every_n_epoch": 25,
"_predict_delta": False,
}
config = MultiBorzoiLoRATrainer.make_config(**config)
Build the model — then load the trained checkpoint (no training)¶
Constructing MultiBorzoiLoRATrainer builds the multi-dataset dataset eagerly: it opens all the
datasets' pseudobulk records, resolves each dataset's genome, and — crucially — computes the
per-dataset conditioning dimensions and merges them into the config (a single-dataset trainer has
nothing to compute here). _setup_model() then builds the model and loads the frozen base backbone
plus the combined per-dataset checkpoint.
At this point a real run would call trainer.train(). We skip that and instead load the trained final
checkpoint, giving us the finished paper model to inspect.
trainer = MultiBorzoiLoRATrainer(
config
) # builds the 41-dataset dataset; fills cond_module_kwargs
trainer._setup_model(
print_model=False
) # builds model; loads base + combined per-dataset ckpt
model = getattr(
trainer.model, "module", trainer.model
) # unwrap DataParallel if present
model.load_checkpoint_from_path(
FINAL_CKPT, strict=False
) # <- the trained final ATAC model
print("loaded trained checkpoint:", FINAL_CKPT)
Using single GPU for training Create BorzoiMultiDataset with config: dataset_records: 41 items <class 'dict'> dna_window: 524288 train_region_step_sample: True use_regions: borzoi n_pseudobulks: 10 paired_mode: ensemble reverse_complement: True max_jitter: 3 batch_size: 4 output_prefix: pseudobulk shared_data_paths: None _genomes: None
Updated cond_module_kwargs in config with dataset specific information.
Setting up model from config
Loading base model weights from: /home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt
Model does not have gene_count_output_head, skip Total model parameters 1209618856, trainable parameters 1032259688 Config contains lora_checkpoint_path, will load weights and perform fine-tuning. Loading LoRA model weights from: 250909-Bolero10M_atac_fold0-PostDatasetFineTune.pt
Loading LoRA model weights from: model/250909-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt
loaded trained checkpoint: model/250909-Bolero10M_atac_fold0-final.lora.best_checkpoint.pt
What's special #1 — per-dataset conditioning into one shared backbone¶
The single-dataset model had one cell-state encoder feeding its conditional-LoRA backbone. That
cannot work across datasets: each dataset was embedded by its own MULTIVI model, so its embedding
space (and its conditioning variables) differ from every other dataset's. The multi-dataset model
solves this with a ConditionEmbeddingModuleMulti: a per-dataset encoder
(cell_encoder_dict / cond_encoder_dict, one entry per dataset) that maps each dataset's own
embedding into a shared conditioning vector. That shared vector — identical in shape for every
dataset — is what drives the single, shared conditional-LoRA backbone.
So the backbone is trained on all datasets at once, while each dataset keeps a small private encoder to translate its embedding into the common space.
cem = model.cond_emb_module
print("cond_emb_module:", type(cem).__name__)
print(
"per-dataset encoders:",
len(cem.cell_encoder_dict),
"cell +",
sum(v is not None for v in cem.cond_encoder_dict.values()),
"with conditions",
)
print("shared conditioning vector width (output_dim):", cem.output_dim)
# per-dataset input dims: each dataset's own cell (and optional condition) embedding width
spec = pd.DataFrame(
{
ds: {
"cell_dim": d["cell"],
"cond_dims": (
sum(d["cond"].values()) if isinstance(d.get("cond"), dict) else 0
),
}
for ds, d in cem.dataset_specific_dims.items()
}
).T
print("\nheterogeneous per-dataset embedding widths (head):")
spec.head(8)
cond_emb_module: ConditionEmbeddingModuleMulti per-dataset encoders: 41 cell + 27 with conditions shared conditioning vector width (output_dim): 768 heterogeneous per-dataset embedding widths (head):
| cell_dim | cond_dims | |
|---|---|---|
| Zu2023Nature+NonN | 60 | 3 |
| Zu2023Nature+CNU | 60 | 3 |
| Zu2023Nature+IB | 60 | 3 |
| Zu2023Nature+CTXGlut | 60 | 3 |
| Zu2023Nature+CTXGaba | 60 | 3 |
| Zu2023Nature+CB | 60 | 3 |
| Zu2023Nature+HB | 60 | 3 |
| Zu2023Nature+MB | 60 | 3 |
What's special #2 — a genome (species) embedding¶
Because the corpus spans several genomes, the model also carries a shared genome embedding: a
one-hot code over the genomes in the registry, appended to the shared conditioning so the network
knows which species' DNA it is reading (the DNA one-hot for each example is fetched from that
dataset's own genome). This is the __genome__ entry in dataset_shared_dims.
dm = trainer.dataset.dm
print("dataset_shared_dims:", cem.dataset_shared_dims) # {'__genome__': n_genomes, ...}
genomes = collections.Counter(g.name for g in dm.genomes.values())
print("genomes in the registry:", dict(genomes))
print("genome one-hot width:", dm.dataset_idx_to_genome_emb[0].shape[-1])
dataset_shared_dims: {'__genome__': 7}
genomes in the registry: {'mm10': 20, 'hg38': 13, 'rheMac10': 3, 'calJac4': 2, 'mCalJac1.pat.X': 1, 'panPan3': 1, 'monDom5.split': 1}
genome one-hot width: 7
Model size¶
The full multi-dataset model: the frozen Flashzoi backbone, the shared conditional-LoRA generators, the per-dataset conditioning encoders, and the accessibility head. Notably it is only marginally larger than the single-dataset model — the 41 per-dataset encoders are lightweight, while the expensive shared conditional-LoRA backbone is amortized across every dataset.
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"total params : {total:,}")
print(f"trainable params: {trainable:,} ({trainable/total:.1%})")
total params : 1,209,618,856 trainable params: 1,032,259,688 (85.3%)
Real training vs. this demo¶
We loaded a finished checkpoint; producing it is a large job. The paper's final ATAC model trains for
200 epochs with lr_total_steps=4_000_000 across multiple GPUs — days of compute — starting from
the combined per-dataset checkpoint. To just watch the loop run, drop in the smoke-test overrides
(max_epochs=3, warmup_steps=10, train_batches=50, val_batches=10) and call trainer.train().
This is Bolero's final accessibility model. The next page extends it with the gene-expression task — adding a gene-count head and per-dataset RNA data on top of this exact model.