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.

A demo on two datasets. The paper model trains on all 42 published datasets, and its full training data is the whole corpus: about 1.6 TB of 32 bp parquet stores. This page keeps the published Borzoi10M model (its config, which carries the dataset registry, and its trained checkpoint) and the frozen Flashzoi backbone, all fetched through bolerodata. It then trains a short demo on two small datasets from that registry, ChromiumPBMC and Li2023Science+Inh (both hg38; about 9.4 GB of downloads).

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:

1. 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).
2. dataset-specific — freeze the shared model and fine-tune each dataset's own conditioning encoder.
3. combine — merge the per-dataset encoders into a single checkpoint.
4. 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 (100 epochs on all 42 datasets, millions of steps, multi-GPU, days of compute), this page builds the model exactly as final training would, loads the published final checkpoint (Borzoi10M), inspects what makes a multi-dataset model different, and then runs a few training steps on the two demo datasets. The final checkpoint holds every weight the stage-3 combined checkpoint sets, so this page skips the stage-3 load and starts from the published weights.

Setup¶

In [1]:

import collections
import copy
import os

import pandas as pd
import torch

from bolero import init
from bolero.tl.model.borzoi.model_lora import filter_dataset_encoders
from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer
from bolerodata import MODELS, get_flashzoi_checkpoint

# The trainer logs to Weights & Biases. As on the single-dataset page, log offline so the demo
# runs without a login; for a real run, remove this line and use your own project.
os.environ["WANDB_MODE"] = "offline"
init(num_cpus=16, object_store_memory_gb=48)

import collections import copy import os import pandas as pd import torch from bolero import init from bolero.tl.model.borzoi.model_lora import filter_dataset_encoders from bolero.tl.model.borzoi.train import MultiBorzoiLoRATrainer from bolerodata import MODELS, get_flashzoi_checkpoint # The trainer logs to Weights & Biases. As on the single-dataset page, log offline so the demo # runs without a login; for a real run, remove this line and use your own project. os.environ["WANDB_MODE"] = "offline" init(num_cpus=16, object_store_memory_gb=48)

2026-10-27 16:22:17,535	INFO worker.py:1781 -- Started a local Ray instance.

In [2]:

# --- Configuration ---------------------------------------------------------------
MODEL_KEY = "Borzoi10M"  # the published multi-dataset ATAC model
FOLD = 0
# the two small hg38 datasets from the registry that the demo trains on
DEMO_DATASETS = ["ChromiumPBMC", "Li2023Science+Inh"]

published = MODELS[MODEL_KEY]
# The published model's config carries the dataset registry the whole model is built from.
MODEL_CONFIG = published.model_config
# The trained final ATAC checkpoint (result of stage 4); downloaded once.
FINAL_CKPT = published.ckpt_path
# Frozen Flashzoi backbone (same base checkpoint as the single-dataset model).
BASE_CKPT = str(get_flashzoi_checkpoint(FOLD))

# --- Configuration --------------------------------------------------------------- MODEL_KEY = "Borzoi10M" # the published multi-dataset ATAC model FOLD = 0 # the two small hg38 datasets from the registry that the demo trains on DEMO_DATASETS = ["ChromiumPBMC", "Li2023Science+Inh"] published = MODELS[MODEL_KEY] # The published model's config carries the dataset registry the whole model is built from. MODEL_CONFIG = published.model_config # The trained final ATAC checkpoint (result of stage 4); downloaded once. FINAL_CKPT = published.ckpt_path # Frozen Flashzoi backbone (same base checkpoint as the single-dataset model). BASE_CKPT = str(get_flashzoi_checkpoint(FOLD))

2026-10-27 16:22:20 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M/config.json "HTTP/1.1 307 Temporary Redirect"

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

2026-10-27 16:22:20 - WARNING - Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

2026-10-27 16:22:20 - INFO - HTTP Request: HEAD https://huggingface.co/api/resolve-cache/models/arcinstitute/bolero-models/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS%2FBorzoi10M%2Fconfig.json "HTTP/1.1 200 OK"

2026-10-27 16:22:20 - INFO - HTTP Request: HEAD https://huggingface.co/arcinstitute/bolero-models/resolve/main/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt "HTTP/1.1 302 Found"

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; the published model's config keeps it, so we read it back from there.

In [3]:

full_records = MODEL_CONFIG["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 full_records.items()
    }
).T
print(f"{len(full_records)} datasets across {registry['genome'].nunique()} genomes")
registry["genome"].value_counts()

full_records = MODEL_CONFIG["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 full_records.items() } ).T print(f"{len(full_records)} datasets across {registry['genome'].nunique()} genomes") registry["genome"].value_counts()

42 datasets across 7 genomes

Out[3]:

genome
mm10              19
hg38              15
rheMac10           3
calJac4            2
mCalJac1.pat.X     1
panPan3            1
monDom5.split      1
Name: count, dtype: int64

In [4]:

# an example record: exactly a single dataset's parquet store + pseudobulk records + genome
example_key = next(iter(full_records))
print(example_key)
full_records[example_key]

# an example record: exactly a single dataset's parquet store + pseudobulk records + genome example_key = next(iter(full_records)) print(example_key) full_records[example_key]

Zu2023Nature+NonN

Out[4]:

{'pseudobulk_path': 'standard/embedding/coverage_pseudobulk_with_condition/Zu2023Nature+NonN/pseudobulk_records_and_cond.cov5000000.joblib',
 'data_path': 'standard/metacell_dataset/parquet/Zu2023Nature+NonN/Zu2023Nature-MetaCell-32bp',
 'genome': 'mm10',
 'dataset_sample_weight': 1}

Keep two datasets for the demo¶

Training reads every registered dataset's 32 bp parquet store and pseudobulk records, and for the full registry that is the whole corpus. The demo keeps two datasets. Two things stay as in the full model:

• the genome list: _genomes lists every genome of the full registry, in order, so the genome one-hot keeps the width the checkpoint was trained with;
• each kept dataset's conditioning encoder: encoders are keyed by dataset name, so the kept datasets' encoders load from the checkpoint unchanged and the others are dropped.

Both kept datasets sample with weight 1. bolerodata downloads their files the first time the trainer opens them.

In [5]:

# every genome of the full registry, in the order the model's genome embedding uses
GENOMES = sorted({v["genome"] for v in full_records.values()})
dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS}
for rec in dataset_records.values():
    rec["dataset_sample_weight"] = 1

print(f"{len(full_records)} datasets in the corpus, {len(GENOMES)} genomes: {GENOMES}")
pd.DataFrame(dataset_records).T[["genome", "data_path"]]

# every genome of the full registry, in the order the model's genome embedding uses GENOMES = sorted({v["genome"] for v in full_records.values()}) dataset_records = {k: copy.deepcopy(full_records[k]) for k in DEMO_DATASETS} for rec in dataset_records.values(): rec["dataset_sample_weight"] = 1 print(f"{len(full_records)} datasets in the corpus, {len(GENOMES)} genomes: {GENOMES}") pd.DataFrame(dataset_records).T[["genome", "data_path"]]

42 datasets in the corpus, 7 genomes: ['calJac4', 'hg38', 'mCalJac1.pat.X', 'mm10', 'monDom5.split', 'panPan3', 'rheMac10']

Out[5]:

	genome	data_path
ChromiumPBMC	hg38	standard/metacell_dataset/parquet/ChromiumPBMC...
Li2023Science+Inh	hg38	standard/metacell_dataset/parquet/Li2023Scienc...

Build the training config¶

The config is the single-dataset config with a few multi-dataset changes. The important keys:

• dataset_records — the registry (replaces the single pseudobulk_records); here, the two demo datasets.
• _genomes — the full genome list, so the genome embedding matches the checkpoint.
• base_checkpoint_path — the frozen Flashzoi fold-0 backbone (same as single-dataset).
• lora_checkpoint_path — the weights training starts from. A real final run starts from the combined per-dataset checkpoint of stage 3; the demo starts from the published final weights.
• signal_model=True, and a demo schedule: one short epoch of 20 training steps at batch size 1. The real run used max_epochs=100, lr_total_steps=2_000_000, save_state_every_n_epoch=10, n_pseudobulks=50 and the default batch_size of 4; the demo lines are simply dropped for it.

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.

In [6]:

config = {
    # run identity + outputs
    "output_dir": "model",
    "savename": f"{MODEL_KEY}_atac_fold{FOLD}-demo",
    "wandb_project": "BorzoiMulti",
    "wandb_job_type": "lora",
    "wandb_name": f"{MODEL_KEY}_atac_fold{FOLD}-demo",
    # data: a registry instead of one pseudobulk file
    "dataset_records": dataset_records,
    "fold_split_id": FOLD,
    "_genomes": GENOMES,  # keep the genome embedding of the full model
    # model
    "base_checkpoint_path": BASE_CKPT,
    "signal_model": True,
    "lora_checkpoint_path": str(FINAL_CKPT),  # a real final run: the stage-3 checkpoint
    "n_pseudobulks": 50,
    "_predict_delta": False,
    # schedule: real run max_epochs=100, lr_total_steps=2_000_000, save_state_every_n_epoch=10
    "max_epochs": 1,
    "lr_total_steps": 2_000_000,
    "save_state_every_n_epoch": 10,
    # demo only: a few small steps (drop these four lines for a real run)
    "warmup_steps": 10,
    "train_batches": 20,
    "val_batches": 5,
    "batch_size": 1,
}
config = MultiBorzoiLoRATrainer.make_config(**config)

config = { # run identity + outputs "output_dir": "model", "savename": f"{MODEL_KEY}_atac_fold{FOLD}-demo", "wandb_project": "BorzoiMulti", "wandb_job_type": "lora", "wandb_name": f"{MODEL_KEY}_atac_fold{FOLD}-demo", # data: a registry instead of one pseudobulk file "dataset_records": dataset_records, "fold_split_id": FOLD, "_genomes": GENOMES, # keep the genome embedding of the full model # model "base_checkpoint_path": BASE_CKPT, "signal_model": True, "lora_checkpoint_path": str(FINAL_CKPT), # a real final run: the stage-3 checkpoint "n_pseudobulks": 50, "_predict_delta": False, # schedule: real run max_epochs=100, lr_total_steps=2_000_000, save_state_every_n_epoch=10 "max_epochs": 1, "lr_total_steps": 2_000_000, "save_state_every_n_epoch": 10, # demo only: a few small steps (drop these four lines for a real run) "warmup_steps": 10, "train_batches": 20, "val_batches": 5, "batch_size": 1, } config = MultiBorzoiLoRATrainer.make_config(**config)

Build the model and load the published weights¶

Constructing MultiBorzoiLoRATrainer builds the multi-dataset dataset eagerly: it opens the datasets' pseudobulk records and 32 bp stores, 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, loads the frozen base backbone and loads lora_checkpoint_path.

We then load the published checkpoint once more with strict=True, after dropping the encoders of the datasets the demo leaves out. A strict load fails if any weight of the model is missing from the checkpoint, so it confirms that every weight of the demo model comes from the published model.

In [7]:

trainer = MultiBorzoiLoRATrainer(config)  # builds the dataset; fills cond_module_kwargs
trainer._setup_model(print_model=False)  # builds model; loads base + published weights

model = getattr(
    trainer.model, "module", trainer.model
)  # unwrap DataParallel if present
state = torch.load(FINAL_CKPT, map_location="cpu", weights_only=False)["state_dict"]
state = filter_dataset_encoders(
    state, dataset_records
)  # keep the demo datasets' encoders
model.load_state_dict(
    state, strict=True
)  # every model weight comes from the checkpoint
del state
print("loaded trained checkpoint:", MODEL_KEY)

trainer = MultiBorzoiLoRATrainer(config) # builds the dataset; fills cond_module_kwargs trainer._setup_model(print_model=False) # builds model; loads base + published weights model = getattr( trainer.model, "module", trainer.model ) # unwrap DataParallel if present state = torch.load(FINAL_CKPT, map_location="cpu", weights_only=False)["state_dict"] state = filter_dataset_encoders( state, dataset_records ) # keep the demo datasets' encoders model.load_state_dict( state, strict=True ) # every model weight comes from the checkpoint del state print("loaded trained checkpoint:", MODEL_KEY)

2026-10-27 16:22:21 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/metacell/parquet/ChromiumPBMC/ChromiumPBMC-MetaCell-32bp.tar "HTTP/1.1 302 Found"

2026-10-27 16:22:21 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/metacell/parquet/Li2023Science%2BInh/Li2023Science-MetaCell-32bp.tar "HTTP/1.1 302 Found"

Using single GPU for training
Create BorzoiMultiDataset with config:
  dataset_records: 2 items <class 'dict'>
  dna_window: 524288
  train_region_step_sample: True
  use_regions: borzoi
  n_pseudobulks: 50
  paired_mode: ensemble
  reverse_complement: True
  max_jitter: 3
  batch_size: 1
  output_prefix: pseudobulk
  shared_data_paths: None
  _genomes: 7 items <class 'list'>

2026-10-27 16:22:21 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/pseudobulk/condition/ChromiumPBMC/pseudobulk_records_and_cond.cov5000000.joblib "HTTP/1.1 302 Found"

2026-10-27 16:22:21 - INFO - HTTP Request: HEAD https://huggingface.co/datasets/arcinstitute/bolero-data/resolve/main/pseudobulk/condition/Li2023Science%2BInh/pseudobulk_records_and_cond.cov5000000.joblib "HTTP/1.1 302 Found"

Updated cond_module_kwargs in config with dataset specific information.

Setting up model from config

Loading base model weights from: /mnt/work/bdhome/flashzoi/fold_0.renamed.pt

Model does not have gene_count_output_head, skip
Total model parameters 1201714785, trainable parameters 1024355617
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: /mnt/work/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt

Skipping checkpoint encoders of datasets not in the config: ['AIBSDev', 'CBDevEvoBonobo', 'CBDevEvoHuman', 'CBDevEvoMacaque', 'CBDevEvoMarmoset', 'CBDevEvoMouse', 'CBDevEvoOpossum', 'CZIBodyAtlas', 'EE+Control', 'EncodeHumanAtlas', 'Fear+Control', 'HMBABasalGangliaHuman', 'HMBABasalGangliaMacaque', 'HMBABasalGangliaMarmoset', 'HYSocialRebound', 'Harris2025Nature', 'HumanBodyDev', 'HumanBrainAging', 'HumanBrainDev', 'HumanCBDev', 'HumanFirstTriBrain', 'KA+Control', 'Li2023Science+Exc', 'Li2023Science+NonN', 'MouseAging', 'MouseBodyAtlas', 'MultiSpeciesHuman', 'MultiSpeciesMacaque', 'MultiSpeciesMarmoset', 'MultiSpeciesMouse', 'MultipleSclerosis+Control', 'NeurIPSBMMC', 'Stress+Control', 'Torpor+Control', 'Zhang2021Cell', 'Zu2023Nature+CB', 'Zu2023Nature+CNU', 'Zu2023Nature+CTXGaba', 'Zu2023Nature+CTXGlut', 'Zu2023Nature+HB', 'Zu2023Nature+IB', 'Zu2023Nature+MB', 'Zu2023Nature+NonN']

loaded trained checkpoint: Borzoi10M

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.

In [8]:

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)

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: 2 cell + 0 with conditions
shared conditioning vector width (output_dim): 768

heterogeneous per-dataset embedding widths (head):

Out[8]:

	cell_dim	cond_dims
ChromiumPBMC	60	0
Li2023Science+Inh	60	0

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.

The demo registry holds only hg38 datasets, but the one-hot keeps the full genome list, so its width matches the checkpoint.

In [9]:

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 demo registry:", dict(genomes))
print("genome one-hot width:", dm.dataset_idx_to_genome_emb[0].shape[-1])

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 demo 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 demo registry: {'hg38': 2}
genome one-hot width: 7

Model size¶

The multi-dataset model: the frozen Flashzoi backbone, the shared conditional-LoRA generators, the per-dataset conditioning encoders, and the accessibility head. It is only marginally larger than the single-dataset model: the per-dataset encoders are lightweight, while the expensive shared conditional-LoRA backbone is amortized across every dataset.

In [10]:

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 = 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,201,714,785
trainable params: 1,024,355,617  (85.2%)

Train on the demo datasets¶

trainer.train() runs the fit loop on the two datasets: it builds the model again from the config (base + published weights), trains on the demo schedule, validates and writes checkpoints under model/. A ...success.flag file is dropped on completion, so re-running the cell is a no-op (delete the flag to retrain).

In [11]:

# train() builds its own copy of the model; free the one inspected above first
del model
trainer.model = None
torch.cuda.empty_cache()

trainer.train()

# train() builds its own copy of the model; free the one inspected above first del model trainer.model = None torch.cuda.empty_cache() trainer.train()

wandb: Tracking run with wandb version 0.28.0

wandb: W&B syncing is set to `offline` in this directory. Run `wandb online` or set WANDB_MODE=online to enable cloud syncing.
wandb: Run data is saved locally in /mnt/work/runs/wandb/offline-run-20261027_162238-nsumgvax
wandb: View this run in the terminal with `wandb leet`

wandb: WARNING URL not available in offline run

Setting up model from config

Loading base model weights from: /mnt/work/bdhome/flashzoi/fold_0.renamed.pt

Model does not have gene_count_output_head, skip
Total model parameters 1201714785, trainable parameters 1024355617
Config contains lora_checkpoint_path, will load weights and perform fine-tuning.
Loading LoRA model weights from: /mnt/work/bdhome/hf/models--arcinstitute--bolero-models/snapshots/eabdbb2e61072357ad0441d49eb66708fb858e33/BorzoiSignalMultiDS/Borzoi10M/checkpoint.pt

Input shape:
x: torch.Size([1, 4, 524288])
embedding: torch.Size([1, 768])
signal: torch.Size([1, 1, 16384])

===========================================================================================================================================================
Layer (type (var_name))                                                          Input Shape               Output Shape              Param #
===========================================================================================================================================================
BorzoiLoRAMulti (BorzoiLoRAMulti)                                                --                        [1, 1, 16352]             656,899
├─ConvDna (conv_dna)                                                             [1, 4, 524288]            [1, 512, 262144]          --
│    └─ConditionalLoRAConv (conv_layer)                                          [1, 4, 524288]            [1, 512, 524288]          31,232
│    │    └─EmbeddingMLP (lora_A_module)                                         --                        [1, 30, 60]               726,280
│    │    └─EmbeddingMLP (lora_B_module)                                         --                        [1, 512, 30]              4,211,200
│    │    └─Dropout (lora_dropout)                                               [1, 512, 524288]          [1, 512, 524288]          --
│    └─MaxPool1d (max_pool)                                                      [1, 512, 524288]          [1, 512, 262144]          --
├─SequentialwithArgs (res_tower)                                                 [1, 512, 262144]          [1, 1280, 16384]          --
│    └─ConvBlock (0)                                                             [1, 512, 262144]          [1, 608, 262144]          --
│    │    └─BatchNorm1d (norm)                                                   [1, 512, 262144]          [1, 512, 262144]          (1,024)
│    │    └─GELU (activation)                                                    [1, 512, 262144]          [1, 512, 262144]          --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 512, 262144]          [1, 608, 262144]          26,509,728
│    └─MaxPool1d (1)                                                             [1, 608, 262144]          [1, 608, 131072]          --
│    └─ConvBlock (2)                                                             [1, 608, 131072]          [1, 736, 131072]          --
│    │    └─BatchNorm1d (norm)                                                   [1, 608, 131072]          [1, 608, 131072]          (1,216)
│    │    └─GELU (activation)                                                    [1, 608, 131072]          [1, 608, 131072]          --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 608, 131072]          [1, 736, 131072]          31,878,496
│    └─MaxPool1d (3)                                                             [1, 736, 131072]          [1, 736, 65536]           --
│    └─ConvBlock (4)                                                             [1, 736, 65536]           [1, 896, 65536]           --
│    │    └─BatchNorm1d (norm)                                                   [1, 736, 65536]           [1, 736, 65536]           (1,472)
│    │    └─GELU (activation)                                                    [1, 736, 65536]           [1, 736, 65536]           --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 736, 65536]           [1, 896, 65536]           39,106,496
│    └─MaxPool1d (5)                                                             [1, 896, 65536]           [1, 896, 32768]           --
│    └─ConvBlock (6)                                                             [1, 896, 32768]           [1, 1056, 32768]          --
│    │    └─BatchNorm1d (norm)                                                   [1, 896, 32768]           [1, 896, 32768]           (1,792)
│    │    └─GELU (activation)                                                    [1, 896, 32768]           [1, 896, 32768]           --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 896, 32768]           [1, 1056, 32768]          47,941,856
│    └─MaxPool1d (7)                                                             [1, 1056, 32768]          [1, 1056, 16384]          --
│    └─ConvBlock (8)                                                             [1, 1056, 16384]          [1, 1280, 16384]          --
│    │    └─BatchNorm1d (norm)                                                   [1, 1056, 16384]          [1, 1056, 16384]          (2,112)
│    │    └─GELU (activation)                                                    [1, 1056, 16384]          [1, 1056, 16384]          --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 1056, 16384]          [1, 1280, 16384]          57,864,640
├─Sequential (signal_encoder)                                                    [1, 1, 16384]             [1, 1280, 16384]          --
│    └─Conv1d (0)                                                                [1, 1, 16384]             [1, 128, 16384]           512
│    └─GroupNorm (1)                                                             [1, 128, 16384]           [1, 128, 16384]           256
│    └─SiLU (2)                                                                  [1, 128, 16384]           [1, 128, 16384]           --
│    └─Conv1d (3)                                                                [1, 128, 16384]           [1, 512, 16384]           197,120
│    └─GroupNorm (4)                                                             [1, 512, 16384]           [1, 512, 16384]           1,024
│    └─SiLU (5)                                                                  [1, 512, 16384]           [1, 512, 16384]           --
│    └─Conv1d (6)                                                                [1, 512, 16384]           [1, 1280, 16384]          1,967,360
│    └─GroupNorm (7)                                                             [1, 1280, 16384]          [1, 1280, 16384]          2,560
├─SequentialwithArgs (unet1)                                                     [1, 1280, 16384]          [1, 1536, 8192]           --
│    └─MaxPool1d (0)                                                             [1, 1280, 16384]          [1, 1280, 8192]           --
│    └─ConvBlock (1)                                                             [1, 1280, 8192]           [1, 1536, 8192]           --
│    │    └─BatchNorm1d (norm)                                                   [1, 1280, 8192]           [1, 1280, 8192]           (2,560)
│    │    └─GELU (activation)                                                    [1, 1280, 8192]           [1, 1280, 8192]           --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 1280, 8192]           [1, 1536, 8192]           71,545,856
├─MaxPool1d (_max_pool)                                                          [1, 1536, 8192]           [1, 1536, 4096]           --
├─ConvBlock (horizontal_conv0)                                                   [1, 1280, 16384]          [1, 1536, 16384]          --
│    └─BatchNorm1d (norm)                                                        [1, 1280, 16384]          [1, 1280, 16384]          (2,560)
│    └─GELU (activation)                                                         [1, 1280, 16384]          [1, 1280, 16384]          --
│    └─ConditionalLoRAConv (conv_layer)                                          [1, 1280, 16384]          [1, 1536, 16384]          1,967,616
│    │    └─EmbeddingMLP (lora_A_module)                                         [1, 768]                  [1, 30, 1280]             10,132,480
│    │    └─EmbeddingMLP (lora_B_module)                                         [1, 768]                  [1, 1536, 30]             12,106,240
│    │    └─Dropout (lora_dropout)                                               [1, 1536, 16384]          [1, 1536, 16384]          --
├─ConvBlock (horizontal_conv1)                                                   [1, 1536, 8192]           [1, 1536, 8192]           --
│    └─BatchNorm1d (norm)                                                        [1, 1536, 8192]           [1, 1536, 8192]           (3,072)
│    └─GELU (activation)                                                         [1, 1536, 8192]           [1, 1536, 8192]           --
│    └─ConditionalLoRAConv (conv_layer)                                          [1, 1536, 8192]           [1, 1536, 8192]           2,360,832
│    │    └─EmbeddingMLP (lora_A_module)                                         [1, 768]                  [1, 30, 1536]             12,106,240
│    │    └─EmbeddingMLP (lora_B_module)                                         [1, 768]                  [1, 1536, 30]             12,106,240
│    │    └─Dropout (lora_dropout)                                               [1, 1536, 8192]           [1, 1536, 8192]           --
├─SequentialwithArgs (transformer)                                               [1, 4096, 1536]           [1, 4096, 1536]           --
│    └─TransformerLayer (0)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (1)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (2)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (3)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (4)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (5)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (6)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
│    └─TransformerLayer (7)                                                      [1, 4096, 1536]           [1, 4096, 1536]           --
│    │    └─SequentialwithArgs (layers)                                          [1, 4096, 1536]           [1, 4096, 1536]           88,870,912
├─SequentialwithArgs (upsampling_unet1)                                          [1, 1536, 4096]           [1, 1536, 8192]           --
│    └─ConvBlock (0)                                                             [1, 1536, 4096]           [1, 1536, 4096]           --
│    │    └─BatchNorm1d (norm)                                                   [1, 1536, 4096]           [1, 1536, 4096]           (3,072)
│    │    └─GELU (activation)                                                    [1, 1536, 4096]           [1, 1536, 4096]           --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 1536, 4096]           [1, 1536, 4096]           26,573,312
│    └─Upsample (1)                                                              [1, 1536, 4096]           [1, 1536, 8192]           --
├─ConvBlock (separable1)                                                         [1, 1536, 8192]           [1, 1536, 8192]           --
│    └─Identity (norm)                                                           [1, 1536, 8192]           [1, 1536, 8192]           --
│    └─Identity (activation)                                                     [1, 1536, 8192]           [1, 1536, 8192]           --
│    └─Sequential (conv_layer)                                                   --                        --                        --
│    │    └─ConditionalLoRAConv (0)                                              [1, 1536, 8192]           [1, 1536, 8192]           12,397,658
│    │    └─ConditionalLoRAConv (1)                                              [1, 1536, 8192]           [1, 1536, 8192]           10,783,232
├─SequentialwithArgs (upsampling_unet0)                                          [1, 1536, 8192]           [1, 1536, 16384]          --
│    └─ConvBlock (0)                                                             [1, 1536, 8192]           [1, 1536, 8192]           --
│    │    └─BatchNorm1d (norm)                                                   [1, 1536, 8192]           [1, 1536, 8192]           (3,072)
│    │    └─GELU (activation)                                                    [1, 1536, 8192]           [1, 1536, 8192]           --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 1536, 8192]           [1, 1536, 8192]           26,573,312
│    └─Upsample (1)                                                              [1, 1536, 8192]           [1, 1536, 16384]          --
├─ConvBlock (separable0)                                                         [1, 1536, 16384]          [1, 1536, 16384]          --
│    └─Identity (norm)                                                           [1, 1536, 16384]          [1, 1536, 16384]          --
│    └─Identity (activation)                                                     [1, 1536, 16384]          [1, 1536, 16384]          --
│    └─Sequential (conv_layer)                                                   --                        --                        --
│    │    └─ConditionalLoRAConv (0)                                              [1, 1536, 16384]          [1, 1536, 16384]          12,397,658
│    │    └─ConditionalLoRAConv (1)                                              [1, 1536, 16384]          [1, 1536, 16384]          10,783,232
├─SequentialwithArgs (final_joined_convs)                                        [1, 1536, 16384]          [1, 1920, 16384]          --
│    └─ConvBlock (0)                                                             [1, 1536, 16384]          [1, 1920, 16384]          --
│    │    └─BatchNorm1d (norm)                                                   [1, 1536, 16384]          [1, 1536, 16384]          (3,072)
│    │    └─GELU (activation)                                                    [1, 1536, 16384]          [1, 1536, 16384]          --
│    │    └─ConditionalLoRAConv (conv_layer)                                     [1, 1536, 16384]          [1, 1920, 16384]          56,769,920
│    └─Dropout (1)                                                               [1, 1920, 16384]          [1, 1920, 16384]          --
│    └─GELU (2)                                                                  [1, 1920, 16384]          [1, 1920, 16384]          --
├─TargetLengthCrop (crop)                                                        [1, 1920, 16384]          [1, 1920, 16352]          --
├─OutputHead (final_output_head)                                                 [1, 1920, 16352]          [1, 1, 16352]             --
│    └─ConditionalLoRAConv (0)                                                   [1, 1920, 16352]          [1, 1, 16352]             1,921
│    │    └─EmbeddingMLP (lora_A_module)                                         --                        [1, 1, 1920]              757,120
│    │    └─EmbeddingMLP (lora_B_module)                                         --                        [1, 1, 1]                 263,937
│    │    └─Dropout (lora_dropout)                                               [1, 1, 16352]             [1, 1, 16352]             --
│    └─Softplus (1)                                                              [1, 1, 16352]             [1, 1, 16352]             --
===========================================================================================================================================================
Total params: 1,201,714,785
Trainable params: 1,024,355,617
Non-trainable params: 177,359,168
Total mult-adds (Units.GIGABYTES): 36.49
===========================================================================================================================================================
Input size (MB): 8.46
Forward/backward pass size (MB): 2786.34
Params size (MB): 4087.35
Estimated Total Size (MB): 6882.15
===========================================================================================================================================================
EMA is not used.
Current epoch: 1, max epochs: 1.
Get dataloader with train mode

2026-10-27 16:22:51,879	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:22:51,880	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]

2026-10-27 16:22:52,894	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:22:52,894	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> TaskPoolMapOperator[MapBatches(ReverseComplement)] -> LimitOperator[limit=22]

Data loader kwargs {'batch_size': 1, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}

 - (Training) 0 [4/20] Ave Loss: 0.338; Ave Corr: 0.648 Last grad norm: 0.5991 Last batch Loss: 0.267 

 - (Training) 0 [9/20] Ave Loss: 0.353; Ave Corr: 0.716 Last grad norm: 0.3308 Last batch Loss: 0.283 

/mnt/work/bolero/.pixi/envs/dev/lib/python3.11/site-packages/torch/optim/lr_scheduler.py:232: UserWarning: The epoch parameter in `scheduler.step()` was not necessary and is being deprecated where possible. Please use `scheduler.step()` to step the scheduler. During the deprecation, if epoch is different from None, the closed form is used instead of the new chainable form, where available. Please open an issue if you are unable to replicate your use case: https://github.com/pytorch/pytorch/issues/new/choose.
  warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)

 - (Training) 0 [14/20] Ave Loss: 0.364; Ave Corr: 0.715 Last grad norm: 0.0882 Last batch Loss: 0.277 

 - (Training) 0 [19/20] Ave Loss: 0.354; Ave Corr: 0.716 Last grad norm: 0.1433 Last batch Loss: 0.243 

2026-10-27 16:23:22,705	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:23:22,705	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]

2026-10-27 16:23:22,873	WARNING plan.py:454 -- Warning: The Ray cluster currently does not have any available CPUs. The Dataset job will hang unless more CPUs are freed up. A common reason is that cluster resources are used by Actors or Tune trials; see the following link for more details: https://docs.ray.io/en/latest/data/data-internals.html#ray-data-and-tune

Get dataloader with eval mode

2026-10-27 16:23:22,889	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:23:22,889	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> LimitOperator[limit=6]

Data loader kwargs {'batch_size': 3, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}

 - (Training)   1 Loss: 0.351; Mean Corr. : 0.726; Learning rate 4.999875000625004e-05.
 - (Validation) 1 Loss: 0.829; Mean Corr. : 0.904.
Saving best checkpoint...

Load and update state dict from checkpoint file: /mnt/work/bolero/docs/tutorials/model_training/model/Borzoi10M_atac_fold0-demo.lora.best_checkpoint.pt

2026-10-27 16:24:20,861	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:24:20,862	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> AllToAllOperator[Repartition]

Get dataloader with eval mode

2026-10-27 16:24:21,637	INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-10-27_16-22-15_178396_3600807/logs/ray-data

2026-10-27 16:24:21,637	INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> ActorPoolMapOperator[FlatMap(GenerateMultiGenomeParquetAndPseudobulk)] -> LimitOperator[limit=16]

Data loader kwargs {'batch_size': 3, 'prefetch_batches': 3, 'drop_last': True, 'local_shuffle_buffer_size': 300}

 - (Validation) 1 [4/5] Mean Loss: 0.087; Mean Track Corr: 0.553; 

wandb: 
wandb: Run history:
wandb:          train/learning_rate ▁▆██
wandb:        train/train_corr_data ▁███
wandb:             train/train_loss ▅█▇▁
wandb: train/train_loss_multinomial ▅██▁
wandb:     train/train_loss_poisson ▆▂▁█
wandb:  train/train_total_grad_norm █▄▁▂
wandb:            val/val_corr_data ▁
wandb:                 val/val_loss ▁
wandb:     val/val_loss_multinomial ▁
wandb:         val/val_loss_poisson ▁
wandb: 
wandb: Run summary:
wandb:              final_test_corr 0.55309
wandb:              final_test_loss 0.08714
wandb:  final_test_loss_multinomial 0.11609
wandb:      final_test_loss_poisson -0.06487
wandb:             final_valid_corr 0.90374
wandb:             final_valid_loss 0.82923
wandb: final_valid_loss_multinomial 1.18761
wandb:     final_valid_loss_poisson -1.05226
wandb:                      success True
wandb:          train/learning_rate 5e-05
wandb:                           +9 ...
wandb: 

wandb: You can sync this run to the cloud by running:
wandb: wandb sync /mnt/work/runs/wandb/offline-run-20261027_162238-nsumgvax

wandb: Find logs at: /mnt/work/runs/wandb/offline-run-20261027_162238-nsumgvax/logs

Real training vs. this demo¶

The demo trained a few steps on two datasets. The paper's final ATAC model trains for 100 epochs with lr_total_steps=2_000_000 on all 42 datasets across multiple GPUs — days of compute — starting from the combined per-dataset checkpoint. For that run, keep the full registry (full_records), set lora_checkpoint_path to the stage-3 checkpoint, and drop the demo lines.

---

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.