06. Train a single-dataset Bolero model¶
The previous section turned ChromiumPBMC into the units a Bolero model trains on: a base-resolution 32 bp parquet store and a set of coverage-matched pseudobulks — each carrying a cell-state embedding, a coverage scale, and a sampling weight. This page trains the model on that single dataset — the simplest way to see the whole training loop end to end.
Single dataset vs. the full model. The model in the paper is trained jointly across the entire Bolero-10M corpus (dozens of datasets over several species). That multi-dataset recipe reuses everything here but adds per-dataset conditioning and a genome embedding; it is covered in the next two pages (multi-dataset ATAC, + gene). Start here to learn the mechanics on one dataset.
What the model is. Bolero is a frozen Flashzoi/Borzoi
genomics transformer wrapped in per-layer conditional LoRA adapters: the adapter's low-rank
weights are generated on the fly by small MLPs from a cell-state embedding, so every pseudobulk
gets its own effective network. The backbone weights stay frozen — training only fits the LoRA
MLPs and the output head. Here the head is the count / accessibility head (softplus +
Poisson-multinomial), and signal_model=True adds the base-resolution signal model on top.
How the pseudobulks condition it. Each training example is a pair (p0, p1) drawn by the
EnsemblePairedPseudobulker you validated on page 05: p1 is one coverage-matched pseudobulk (a
specific cell state), p0 is an on-the-fly ensemble background. Their 30-dim embeddings are
concatenated (30 + 30 → 60) into the reference→target vector every LoRA MLP reads.
In / out.
- Inputs — the 32 bp parquet store and the pseudobulk records from the previous section, plus a frozen Flashzoi fold-0 base checkpoint (see below).
- Outputs — LoRA checkpoints and epoch state under
model/(git-ignored).
This notebook runs in place and reads the artifacts you built earlier in the series via relative paths (
../meta_cell_adata_and_parquet/). It is a quick-demo configuration — a handful of tiny batches so the whole loop finishes in a few minutes. The production configuration for a real single-dataset run is given below.
Setup¶
In [1]:
Copied!
import os
from pathlib import Path
from bolero import init
from bolero.tl.model.borzoi.train import BorzoiLoRATrainer
from bolero.tl.pseudobulk.paired_pseudobulk import EnsemblePairedPseudobulker
# The trainer logs to Weights & Biases. For this demo we disable uploads so it runs without a
# login; for a real run, remove this line (or set WANDB_MODE="online") and use your own project.
os.environ["WANDB_MODE"] = "offline"
# Ray backs the streaming data pipeline (parquet -> CSR -> pseudobulk -> one-hot). num_cpus and
# object_store_memory_gb scale with the machine; these suit a single-GPU node with the data lake.
init(num_cpus=16, object_store_memory_gb=48)
import os
from pathlib import Path
from bolero import init
from bolero.tl.model.borzoi.train import BorzoiLoRATrainer
from bolero.tl.pseudobulk.paired_pseudobulk import EnsemblePairedPseudobulker
# The trainer logs to Weights & Biases. For this demo we disable uploads so it runs without a
# login; for a real run, remove this line (or set WANDB_MODE="online") and use your own project.
os.environ["WANDB_MODE"] = "offline"
# Ray backs the streaming data pipeline (parquet -> CSR -> pseudobulk -> one-hot). num_cpus and
# object_store_memory_gb scale with the machine; these suit a single-GPU node with the data lake.
init(num_cpus=16, object_store_memory_gb=48)
2026-07-13 11:26:54,784 INFO worker.py:1781 -- Started a local Ray instance.
Configure the run¶
Everything the trainer needs is a flat config dict. We first name the paths — the two inputs from the previous section, plus the frozen backbone — then derive the embedding widths from the pseudobulk records so the model is sized to match the data.
In [2]:
Copied!
# --- Configuration --------------------------------------------------------------
DATASET_NAME = "ChromiumPBMC"
DATA_DIR = Path(
"../meta_cell_adata_and_parquet"
) # artifacts built in the previous section
FOLD = 0 # chromosome cross-validation fold (hg38_splits)
TARGET_COV = 5_000_000 # must match the pseudobulk records built on page 05
# inputs from the previous section
PARQUET_DIR = DATA_DIR / f"{DATASET_NAME}-MetaCell-32bp" # page 04
PSEUDOBULK_RECORDS = (
DATA_DIR / f"pseudobulk_records_and_cond.cov{TARGET_COV}.joblib"
) # page 05
# frozen Flashzoi backbone (see the note below on where this comes from)
BASE_CHECKPOINT = (
"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt"
)
# run identity + where checkpoints land
SAVENAME = f"{DATASET_NAME}-atac-signal-fold{FOLD}"
WANDB_PROJECT = f"BorzoiSignal-{DATASET_NAME}"
assert (
PARQUET_DIR.exists() and PSEUDOBULK_RECORDS.exists()
), "run the previous section first"
# --- Configuration --------------------------------------------------------------
DATASET_NAME = "ChromiumPBMC"
DATA_DIR = Path(
"../meta_cell_adata_and_parquet"
) # artifacts built in the previous section
FOLD = 0 # chromosome cross-validation fold (hg38_splits)
TARGET_COV = 5_000_000 # must match the pseudobulk records built on page 05
# inputs from the previous section
PARQUET_DIR = DATA_DIR / f"{DATASET_NAME}-MetaCell-32bp" # page 04
PSEUDOBULK_RECORDS = (
DATA_DIR / f"pseudobulk_records_and_cond.cov{TARGET_COV}.joblib"
) # page 05
# frozen Flashzoi backbone (see the note below on where this comes from)
BASE_CHECKPOINT = (
"/home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt"
)
# run identity + where checkpoints land
SAVENAME = f"{DATASET_NAME}-atac-signal-fold{FOLD}"
WANDB_PROJECT = f"BorzoiSignal-{DATASET_NAME}"
assert (
PARQUET_DIR.exists() and PSEUDOBULK_RECORDS.exists()
), "run the previous section first"
The frozen Flashzoi backbone¶
BASE_CHECKPOINT points to the Flashzoi fold-0 weights that Bolero freezes and adapts. Bolero
uses the published Flashzoi replicate
checkpoints (four cross-validation folds, on the Hugging Face Hub) with their state-dict keys
renamed to match this repo's Borzoi module — hence the .renamed.pt file. Point
BASE_CHECKPOINT at your local copy of the renamed fold-0 checkpoint; the fold here matches
FOLD = 0, so the frozen backbone and the chromosome CV split agree.
In [3]:
Copied!
# The paired pseudobulker (page 05) both sizes the model and, at train time, samples the (p0, p1)
# pairs. Its metacell embedding is 30-dim; the reference->target concat doubles that to 60.
pseudobulker = EnsemblePairedPseudobulker(str(PSEUDOBULK_RECORDS))
emb_input_features = (
pseudobulker.meta_cell_emb.shape[1] * 2
) # 30 -> 60 (reference + target)
cond_emb_dim = dict(pseudobulker.cond_emb_dims) # {} for ChromiumPBMC (no conditions)
print("emb_input_features:", emb_input_features)
print("cond_emb_dim:", cond_emb_dim)
# The paired pseudobulker (page 05) both sizes the model and, at train time, samples the (p0, p1)
# pairs. Its metacell embedding is 30-dim; the reference->target concat doubles that to 60.
pseudobulker = EnsemblePairedPseudobulker(str(PSEUDOBULK_RECORDS))
emb_input_features = (
pseudobulker.meta_cell_emb.shape[1] * 2
) # 30 -> 60 (reference + target)
cond_emb_dim = dict(pseudobulker.cond_emb_dims) # {} for ChromiumPBMC (no conditions)
print("emb_input_features:", emb_input_features)
print("cond_emb_dim:", cond_emb_dim)
emb_input_features: 60
cond_emb_dim: {}
Build the trainer config¶
BorzoiLoRATrainer.make_config merges these keys onto the dataset / model / trainer defaults and
validates them. The block below is the quick-demo run: max_epochs, warmup_steps,
train_batches, and val_batches are cut to tiny values so the full loop finishes in a few
minutes. Everything else (batch size, learning rate, gradient accumulation, the count head, the
signal model) uses the shipped defaults — the same ones a real run uses.
In [4]:
Copied!
config = {
# run identity + outputs
"output_dir": "model",
"savename": SAVENAME,
"wandb_project": WANDB_PROJECT,
"wandb_job_type": "lora",
"wandb_name": SAVENAME,
# data (built in the previous section)
"dataset_path": str(PARQUET_DIR),
"fold_split_id": FOLD,
"pseudobulk_records": str(PSEUDOBULK_RECORDS),
# model
"base_checkpoint_path": BASE_CHECKPOINT,
"emb_input_features": emb_input_features,
"cond_emb_dim": cond_emb_dim,
"signal_model": True,
# ---- quick-demo overrides: a tiny run. See "The production configuration" below. ----
"max_epochs": 3,
"warmup_steps": 10,
"train_batches": 50,
"val_batches": 10,
# At the default "batch_size" of 4 training peaks at ~57 GB of GPU memory (fits an 80 GB
# H100). On a smaller GPU, lower it, e.g. "batch_size": 1.
}
config = BorzoiLoRATrainer.make_config(**config)
config = {
# run identity + outputs
"output_dir": "model",
"savename": SAVENAME,
"wandb_project": WANDB_PROJECT,
"wandb_job_type": "lora",
"wandb_name": SAVENAME,
# data (built in the previous section)
"dataset_path": str(PARQUET_DIR),
"fold_split_id": FOLD,
"pseudobulk_records": str(PSEUDOBULK_RECORDS),
# model
"base_checkpoint_path": BASE_CHECKPOINT,
"emb_input_features": emb_input_features,
"cond_emb_dim": cond_emb_dim,
"signal_model": True,
# ---- quick-demo overrides: a tiny run. See "The production configuration" below. ----
"max_epochs": 3,
"warmup_steps": 10,
"train_batches": 50,
"val_batches": 10,
# At the default "batch_size" of 4 training peaks at ~57 GB of GPU memory (fits an 80 GB
# H100). On a smaller GPU, lower it, e.g. "batch_size": 1.
}
config = BorzoiLoRATrainer.make_config(**config)
In [5]:
Copied!
# A few resolved values — the demo overrides plus the inherited defaults the real run keeps.
for k in [
"mode",
"fold_split_id",
"signal_model",
"emb_input_features",
"cond_emb_dim",
"lr",
"batch_size",
"accumulate_grad",
"max_epochs",
"warmup_steps",
"train_batches",
"val_batches",
]:
print(f"{k:20s} {config[k]}")
# A few resolved values — the demo overrides plus the inherited defaults the real run keeps.
for k in [
"mode",
"fold_split_id",
"signal_model",
"emb_input_features",
"cond_emb_dim",
"lr",
"batch_size",
"accumulate_grad",
"max_epochs",
"warmup_steps",
"train_batches",
"val_batches",
]:
print(f"{k:20s} {config[k]}")
mode lora
fold_split_id 0
signal_model True
emb_input_features 60
cond_emb_dim {}
lr 5e-05
batch_size 4
accumulate_grad 2
max_epochs 3
warmup_steps 10
train_batches 50
val_batches 10
Train¶
Constructing the trainer loads the frozen backbone, freezes everything except the output head,
swaps the Linear/Conv layers for conditional-LoRA variants, and prints the model plus its
trainable-parameter count. trainer.train() then runs the fit loop — building the Ray data
pipeline, training and validating per epoch, and writing 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 [6]:
Copied!
trainer = BorzoiLoRATrainer(config)
trainer.train()
trainer = BorzoiLoRATrainer(config)
trainer.train()
Using single GPU for training Create BorzoiDataset with config: dataset_path: ../meta_cell_adata_and_parquet/ChromiumPBMC-MetaCell-32bp dna_window: 524288 pos_resolution: 32 use_regions: borzoi train_region_step_sample: True n_pseudobulks: 10 paired_data: True paired_mode: ensemble reverse_complement: True rc_mode: random max_jitter: 3 shuffle_files: True batch_size: 4 min_cov: 0 normalize_cov: True reduce_resolution: False read_parquet_kwargs: None deg_list: None gene_data_path: None qtl_data_path: None _multihead: False
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 /large_storage/zhoulab/hanliu/pkg/liuhlab/bolero/docs/tutorials/model_training/wandb/offline-run-20260713_112657-erbg54ng 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: /home/hanliu/data/wmb/model/Borzoi/flashzoi_checkpoints/fold_0.renamed.pt
Condition embedding is not used.
Model does not have gene_count_output_head, skip Input shape: x: torch.Size([1, 4, 524288]) embedding: torch.Size([1, 60]) signal: torch.Size([1, 1, 16384])
=========================================================================================================================================================== Layer (type (var_name)) Input Shape Output Shape Param # =========================================================================================================================================================== BorzoiLoRA (BorzoiLoRA) -- [1, 1, 16352] 81,665 ├─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] 545,032 │ │ └─EmbeddingMLP (lora_B_module) -- [1, 512, 30] 4,029,952 │ │ └─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,147,232 │ └─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,516,000 │ └─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] 38,744,000 │ └─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,579,360 │ └─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,502,144 ├─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,183,360 ├─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, 60] [1, 30, 1280] 9,951,232 │ │ └─EmbeddingMLP (lora_B_module) [1, 60] [1, 1536, 30] 11,924,992 │ │ └─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, 60] [1, 30, 1536] 11,924,992 │ │ └─EmbeddingMLP (lora_B_module) [1, 60] [1, 1536, 30] 11,924,992 │ │ └─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,145,920 │ └─TransformerLayer (1) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (2) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (3) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (4) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (5) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (6) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 │ └─TransformerLayer (7) [1, 4096, 1536] [1, 4096, 1536] -- │ │ └─SequentialwithArgs (layers) [1, 4096, 1536] [1, 4096, 1536] 88,145,920 ├─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,210,816 │ └─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,035,162 │ │ └─ConditionalLoRAConv (1) [1, 1536, 8192] [1, 1536, 8192] 10,420,736 ├─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,210,816 │ └─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,035,162 │ │ └─ConditionalLoRAConv (1) [1, 1536, 16384] [1, 1536, 16384] 10,420,736 ├─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,407,424 │ └─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] 575,872 │ │ └─EmbeddingMLP (lora_B_module) -- [1, 1, 1] 82,689 │ │ └─Dropout (lora_dropout) [1, 1, 16352] [1, 1, 16352] -- │ └─Softplus (1) [1, 1, 16352] [1, 1, 16352] -- =========================================================================================================================================================== Total params: 1,189,177,183 Trainable params: 1,011,818,015 Non-trainable params: 177,359,168 Total mult-adds (Units.GIGABYTES): 36.48 =========================================================================================================================================================== Input size (MB): 8.45 Forward/backward pass size (MB): 2786.34 Params size (MB): 4039.50 Estimated Total Size (MB): 6834.30 =========================================================================================================================================================== Total model parameters 1189177183, trainable parameters 1011818015 EMA is not used. Current epoch: 1, max epochs: 3. Get dataloader with train mode
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': 300, 'drop_last': True, 'batch_size': 4}
2026-07-13 11:27:10,662 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:27:10,663 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(ReverseComplement)->MapBatches(_region_to_global_coords)] -> LimitOperator[limit=208]
- (Training) 0 [4/50] Ave Loss: 0.904; Ave Corr: 0.342 Last grad norm: 59.7296 Last batch Loss: 0.997
- (Training) 0 [9/50] Ave Loss: 0.799; Ave Corr: 0.429 Last grad norm: 2.4741 Last batch Loss: 0.818
/large_storage/zhoulab/hanliu/pkg/liuhlab/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/50] Ave Loss: 0.790; Ave Corr: 0.440 Last grad norm: 1.0000 Last batch Loss: 1.225
- (Training) 0 [19/50] Ave Loss: 0.761; Ave Corr: 0.478 Last grad norm: 0.4631 Last batch Loss: 0.904
- (Training) 0 [24/50] Ave Loss: 0.737; Ave Corr: 0.473 Last grad norm: 0.6218 Last batch Loss: 0.228
- (Training) 0 [29/50] Ave Loss: 0.747; Ave Corr: 0.540 Last grad norm: 0.4189 Last batch Loss: 1.123
- (Training) 0 [34/50] Ave Loss: 0.743; Ave Corr: 0.567 Last grad norm: 0.4239 Last batch Loss: 0.362
- (Training) 0 [39/50] Ave Loss: 0.721; Ave Corr: 0.571 Last grad norm: 0.1920 Last batch Loss: 0.247
- (Training) 0 [44/50] Ave Loss: 0.718; Ave Corr: 0.572 Last grad norm: 2.9624 Last batch Loss: 0.829
- (Training) 0 [49/50] Ave Loss: 0.719; Ave Corr: 0.594 Last grad norm: 0.3752 Last batch Loss: 0.736
Get dataloader with eval mode
2026-07-13 11:28:23,054 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:28:23,055 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(_region_to_global_coords)] -> LimitOperator[limit=132]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': None, 'drop_last': True, 'batch_size': 12}
- (Validation) 0 [4/10] Mean Loss: 0.177; Mean Track Corr: 0.761;
- (Validation) 0 [9/10] Mean Loss: 0.867; Mean Track Corr: 0.796;
- (Training) 1 Loss: 0.708; Mean Corr. : 0.603; Learning rate 4.996001119891201e-05. - (Validation) 1 Loss: 0.904; Mean Corr. : 0.795. Saving best checkpoint...
Current epoch: 2, max epochs: 3. Get dataloader with train mode
2026-07-13 11:29:34,125 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:29:34,125 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(ReverseComplement)->MapBatches(_region_to_global_coords)] -> LimitOperator[limit=208]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': 300, 'drop_last': True, 'batch_size': 4}
- (Training) 1 [4/50] Ave Loss: 1.345; Ave Corr: 0.705 Last grad norm: 0.4607 Last batch Loss: 1.383
- (Training) 1 [9/50] Ave Loss: 1.131; Ave Corr: 0.730 Last grad norm: 0.4595 Last batch Loss: 1.181
- (Training) 1 [14/50] Ave Loss: 1.112; Ave Corr: 0.745 Last grad norm: 0.2796 Last batch Loss: 0.846
- (Training) 1 [19/50] Ave Loss: 1.179; Ave Corr: 0.765 Last grad norm: 0.3771 Last batch Loss: 1.253
- (Training) 1 [24/50] Ave Loss: 1.267; Ave Corr: 0.771 Last grad norm: 0.8018 Last batch Loss: 1.763
- (Training) 1 [29/50] Ave Loss: 1.286; Ave Corr: 0.777 Last grad norm: 0.8005 Last batch Loss: 1.121
- (Training) 1 [34/50] Ave Loss: 1.262; Ave Corr: 0.781 Last grad norm: 0.1984 Last batch Loss: 1.478
- (Training) 1 [39/50] Ave Loss: 1.220; Ave Corr: 0.781 Last grad norm: 0.4207 Last batch Loss: 1.629
- (Training) 1 [44/50] Ave Loss: 1.211; Ave Corr: 0.786 Last grad norm: 0.1720 Last batch Loss: 1.671
- (Training) 1 [49/50] Ave Loss: 1.223; Ave Corr: 0.788 Last grad norm: 0.5341 Last batch Loss: 1.004
Get dataloader with eval mode
2026-07-13 11:30:44,417 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:30:44,417 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(_region_to_global_coords)] -> LimitOperator[limit=132]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': None, 'drop_last': True, 'batch_size': 12}
(raylet) [2026-07-13 11:31:28,634 E 2404099 2404099] (raylet) worker.cc:176: Failed to send wait complete: RpcError: RPC Error message: Socket closed; RPC Error details:
- (Validation) 1 [4/10] Mean Loss: 0.517; Mean Track Corr: 0.781;
- (Validation) 1 [9/10] Mean Loss: 0.547; Mean Track Corr: 0.857;
- (Training) 2 Loss: 1.229; Mean Corr. : 0.791; Learning rate 4.9908064012353764e-05. - (Validation) 2 Loss: 0.519; Mean Corr. : 0.858. Saving best checkpoint...
Current epoch: 3, max epochs: 3. Get dataloader with train mode
2026-07-13 11:32:01,307 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:32:01,308 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(ReverseComplement)->MapBatches(_region_to_global_coords)] -> LimitOperator[limit=208]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': 300, 'drop_last': True, 'batch_size': 4}
- (Training) 2 [4/50] Ave Loss: 0.415; Ave Corr: 0.838 Last grad norm: 0.1225 Last batch Loss: 0.047
- (Training) 2 [9/50] Ave Loss: 0.345; Ave Corr: 0.843 Last grad norm: 0.2270 Last batch Loss: 0.357
- (Training) 2 [14/50] Ave Loss: 0.303; Ave Corr: 0.853 Last grad norm: 0.0924 Last batch Loss: 0.401
- (Training) 2 [19/50] Ave Loss: 0.331; Ave Corr: 0.823 Last grad norm: 0.1433 Last batch Loss: 0.049
- (Training) 2 [24/50] Ave Loss: 0.306; Ave Corr: 0.813 Last grad norm: 0.0884 Last batch Loss: 0.350
- (Training) 2 [29/50] Ave Loss: 0.359; Ave Corr: 0.768 Last grad norm: 0.3019 Last batch Loss: 1.441
- (Training) 2 [34/50] Ave Loss: 0.356; Ave Corr: 0.762 Last grad norm: 0.1228 Last batch Loss: 0.219
- (Training) 2 [39/50] Ave Loss: 0.363; Ave Corr: 0.777 Last grad norm: 0.1721 Last batch Loss: 0.777
- (Training) 2 [44/50] Ave Loss: 0.376; Ave Corr: 0.781 Last grad norm: 0.2931 Last batch Loss: 0.089
- (Training) 2 [49/50] Ave Loss: 0.425; Ave Corr: 0.800 Last grad norm: 0.2538 Last batch Loss: 0.739
Get dataloader with eval mode
2026-07-13 11:33:11,206 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:33:11,206 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(_region_to_global_coords)] -> LimitOperator[limit=132]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': None, 'drop_last': True, 'batch_size': 12}
- (Validation) 2 [4/10] Mean Loss: 0.552; Mean Track Corr: 0.886;
- (Validation) 2 [9/10] Mean Loss: 0.624; Mean Track Corr: 0.868;
- (Training) 3 Loss: 0.418; Mean Corr. : 0.797; Learning rate 4.9856160045506116e-05. - (Validation) 3 Loss: 0.610; Mean Corr. : 0.869. Saving best checkpoint...
Load and update state dict from checkpoint file: /large_storage/zhoulab/hanliu/pkg/liuhlab/bolero/docs/tutorials/model_training/model/ChromiumPBMC-atac-signal-fold0.lora.best_checkpoint.pt
Get dataloader with eval mode
2026-07-13 11:34:47,553 INFO streaming_executor.py:108 -- Starting execution of Dataset. Full logs are in /tmp/ray/session_2026-07-13_11-26-52_473225_2403210/logs/ray-data
2026-07-13 11:34:47,554 INFO streaming_executor.py:109 -- Execution plan of Dataset: InputDataBuffer[Input] -> ActorPoolMapOperator[ReadParquet->Filter(region_length_filter)->Map(CompressedBytesToTensor)] -> ActorPoolMapOperator[FlatMap(GeneratePairedPseudobulk)] -> ActorPoolMapOperator[FlatMap(GenerateRegions)] -> ActorPoolMapOperator[MapBatches(MaskBlacklistAndClamp)] -> ActorPoolMapOperator[MapBatches(FetchRegionOneHot)] -> TaskPoolMapOperator[MapBatches(_region_to_global_coords)] -> LimitOperator[limit=372]
Data loader kwargs {'prefetch_batches': 3, 'local_shuffle_buffer_size': None, 'drop_last': True, 'batch_size': 12}
- (Validation) 3 [4/10] Mean Loss: 1.188; Mean Track Corr: 0.872;
- (Validation) 3 [9/10] Mean Loss: 1.591; Mean Track Corr: 0.839;
- (Validation) 3 [14/10] Mean Loss: 1.114; Mean Track Corr: 0.838;
- (Validation) 3 [19/10] Mean Loss: 1.002; Mean Track Corr: 0.850;
- (Validation) 3 [24/10] Mean Loss: 1.270; Mean Track Corr: 0.845;
- (Validation) 3 [29/10] Mean Loss: 1.247; Mean Track Corr: 0.848;
wandb: wandb: Run history: wandb: train/learning_rate ▁▆████████████████████████████ wandb: train/train_corr_ChromiumPBMC.MetaCell ▁▂▂▃▃▄▄▄▄▄▆▆▇▇▇▇▇▇▇▇████▇▇▇▇▇▇ wandb: train/train_loss ▅▄▆▄▂▅▂▂▄▄▆▆▄▆█▅▇▇█▅▁▂▂▁▂▇▂▄▁▄ wandb: train/train_loss_multinomial ▄▄▅▄▂▅▂▂▄▄▆▆▄▆█▅▇▇█▅▁▂▂▁▂▇▂▄▁▄ wandb: train/train_loss_poisson ▆▆▅▆█▄▇▇▅▅▂▃▅▄▁▃▂▂▁▄█▇▆█▇▃▇▅▇▅ wandb: train/train_total_grad_norm █▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ wandb: val/val_corr_ChromiumPBMC.MetaCell ▁▇█ wandb: val/val_loss █▁▃ wandb: val/val_loss_multinomial █▁▃ wandb: val/val_loss_poisson ▁█▆ wandb: wandb: Run summary: wandb: final_test_corr 0.84729 wandb: final_test_loss 1.26526 wandb: final_test_loss_multinomial 1.84384 wandb: final_test_loss_poisson -1.77229 wandb: final_valid_corr 0.86911 wandb: final_valid_loss 0.61025 wandb: final_valid_loss_multinomial 0.88474 wandb: final_valid_loss_poisson -0.83082 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 /large_storage/zhoulab/hanliu/pkg/liuhlab/bolero/docs/tutorials/model_training/wandb/offline-run-20260713_112657-erbg54ng
wandb: Find logs at: ./wandb/offline-run-20260713_112657-erbg54ng/logs
Inspect the model¶
Two objects are worth looking at: the frozen base backbone — the plain Flashzoi/Borzoi trunk
that Bolero loads from BASE_CHECKPOINT and never updates — and the full Bolero model the
trainer just built on top of it (the same trunk, with conditional-LoRA adapters swapped into its
Linear/Conv layers, plus the output head and the signal model). Comparing the two shows where the
trainable parameters actually live.
In [7]:
Copied!
from bolero.tl.model.borzoi.model import Borzoi, model_summary
device = trainer.device
# The frozen backbone on its own (~177M params): DNA stem -> residual conv tower + U-net ->
# transformer -> separable up-convs -> the 1920-channel trunk. This standalone copy is not frozen
# yet (torchinfo reports 0 non-trainable), but it is exactly the architecture BASE_CHECKPOINT
# fills, and these ~177M weights are the slice Bolero holds frozen inside the full model below.
base_model = Borzoi.create_from_config(config).to(device)
print(model_summary(base_model, depth=1))
from bolero.tl.model.borzoi.model import Borzoi, model_summary
device = trainer.device
# The frozen backbone on its own (~177M params): DNA stem -> residual conv tower + U-net ->
# transformer -> separable up-convs -> the 1920-channel trunk. This standalone copy is not frozen
# yet (torchinfo reports 0 non-trainable), but it is exactly the architecture BASE_CHECKPOINT
# fills, and these ~177M weights are the slice Bolero holds frozen inside the full model below.
base_model = Borzoi.create_from_config(config).to(device)
print(model_summary(base_model, depth=1))
========================================================================================================= Layer (type (var_name)) Param # ========================================================================================================= Borzoi (Borzoi) -- ├─ConvDna (conv_dna) 31,232 ├─MaxPool1d (_max_pool) -- ├─SequentialwithArgs (res_tower) 18,592,672 ├─SequentialwithArgs (unet1) 9,834,496 ├─SequentialwithArgs (transformer) 132,243,456 ├─ConvBlock (horizontal_conv0) 1,970,176 ├─ConvBlock (horizontal_conv1) 2,363,904 ├─Upsample (upsample) -- ├─SequentialwithArgs (upsampling_unet1) 2,363,904 ├─ConvBlock (separable1) 2,365,440 ├─SequentialwithArgs (upsampling_unet0) 2,363,904 ├─ConvBlock (separable0) 2,365,440 ├─TargetLengthCrop (crop) -- ├─SequentialwithArgs (final_joined_convs) 2,954,112 ========================================================================================================= Total params: 177,448,736 Trainable params: 177,448,736 Non-trainable params: 0 =========================================================================================================
In [8]:
Copied!
# The full model. On one GPU it is the module itself; multi-GPU training would wrap it in
# DataParallel, so unwrap defensively.
model = getattr(trainer.model, "module", trainer.model)
print(model_summary(model, depth=1))
# The full model. On one GPU it is the module itself; multi-GPU training would wrap it in
# DataParallel, so unwrap defensively.
model = getattr(trainer.model, "module", trainer.model)
print(model_summary(model, depth=1))
========================================================================================================= Layer (type (var_name)) Param # ========================================================================================================= BorzoiLoRA (BorzoiLoRA) 81,665 ├─ConvDna (conv_dna) 4,606,216 ├─SequentialwithArgs (res_tower) 201,496,352 ├─Sequential (signal_encoder) 2,168,832 ├─SequentialwithArgs (unet1) 71,185,920 ├─MaxPool1d (_max_pool) -- ├─ConvBlock (horizontal_conv0) 23,846,400 ├─ConvBlock (horizontal_conv1) 26,213,888 ├─SequentialwithArgs (transformer) 705,167,360 ├─SequentialwithArgs (upsampling_unet1) 26,213,888 ├─ConvBlock (separable1) 22,455,898 ├─SequentialwithArgs (upsampling_unet0) 26,213,888 ├─ConvBlock (separable0) 22,455,898 ├─SequentialwithArgs (final_joined_convs) 56,410,496 ├─TargetLengthCrop (crop) -- ├─OutputHead (final_output_head) 660,482 ========================================================================================================= Total params: 1,189,177,183 Trainable params: 1,011,818,015 Non-trainable params: 177,359,168 =========================================================================================================
In [9]:
Copied!
# Walk the top-level components and split each into total vs. trainable parameters. The backbone
# modules carry a large frozen weight plus a small conditional-LoRA adapter (the trainable part);
# the heads are fully trainable.
def count_params(module):
total = sum(p.numel() for p in module.parameters())
trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
return total, trainable
print(f"{'component':22s} {'class':26s} {'params':>14s} {'trainable':>13s}")
print("-" * 78)
for name, child in model.named_children():
total, trainable = count_params(child)
if total == 0:
continue # e.g. the crop layer holds no parameters
print(f"{name:22s} {type(child).__name__:26s} {total:14,d} {trainable:13,d}")
print("-" * 78)
print(
f"{'TOTAL':22s} {'':26s} {trainer.total_params:14,d} {trainer.trainable_params:13,d}"
)
print(f"\ntrainable fraction: {trainer.trainable_params / trainer.total_params:.2%}")
# Walk the top-level components and split each into total vs. trainable parameters. The backbone
# modules carry a large frozen weight plus a small conditional-LoRA adapter (the trainable part);
# the heads are fully trainable.
def count_params(module):
total = sum(p.numel() for p in module.parameters())
trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
return total, trainable
print(f"{'component':22s} {'class':26s} {'params':>14s} {'trainable':>13s}")
print("-" * 78)
for name, child in model.named_children():
total, trainable = count_params(child)
if total == 0:
continue # e.g. the crop layer holds no parameters
print(f"{name:22s} {type(child).__name__:26s} {total:14,d} {trainable:13,d}")
print("-" * 78)
print(
f"{'TOTAL':22s} {'':26s} {trainer.total_params:14,d} {trainer.trainable_params:13,d}"
)
print(f"\ntrainable fraction: {trainer.trainable_params / trainer.total_params:.2%}")
component class params trainable ------------------------------------------------------------------------------ conv_dna ConvDna 4,606,216 4,575,496 res_tower SequentialwithArgs 201,496,352 182,908,256 unet1 SequentialwithArgs 71,185,920 61,352,960 transformer SequentialwithArgs 705,167,360 572,997,632 horizontal_conv0 ConvBlock 23,846,400 21,877,760 horizontal_conv1 ConvBlock 26,213,888 23,851,520 upsampling_unet1 SequentialwithArgs 26,213,888 23,851,520 separable1 ConvBlock 22,455,898 20,091,994 upsampling_unet0 SequentialwithArgs 26,213,888 23,851,520 separable0 ConvBlock 22,455,898 20,091,994 final_joined_convs SequentialwithArgs 56,410,496 53,458,304 final_output_head OutputHead 660,482 658,562 signal_encoder Sequential 2,168,832 2,168,832 cond_emb_module ConditionEmbeddingModuleNoEffect 81,665 81,665 ------------------------------------------------------------------------------ TOTAL 1,189,177,183 1,011,818,015 trainable fraction: 85.09%
Reading the table top to bottom follows a sequence through the model:
conv_dna— the DNA stem: one-hot 524,288 bp in, first convolutional features out.res_tower,unet1,horizontal_conv0/1— the residual convolution tower and U-net that downsample the sequence to the 32 bp bin grid while building a multi-scale representation.transformer— the 8 flash-attention layers, the largest block in the network. In the base backbone it is ~132M params; in the full model it reports ~705M, because the conditional-LoRA generators that produce its per-layer adapter weights are nested inside it. Its trainable count (~573M) is those generators; the ~132M base attention / MLP weights are the frozen remainder.separable1,separable0,upsampling_unet0/1,final_joined_convs— separable up-convolutions that return to the 16,384-bin output grid, ending in the shared 1920-channel trunk. (cropcarries no parameters and is skipped.)final_output_head— the count / accessibility head (softplus); fully trainable.signal_encoder,cond_emb_module— the signal model, present becausesignal_model=True.cond_emb_moduleturns the paired pseudobulk embedding (and any conditions) into the vector every adapter generator reads; both are fully trainable. With no conditions, ChromiumPBMC's is a...NoEffectmodule.
The contrast between the two summaries is the point. The frozen part is the ~177M-parameter base
backbone — the Non-trainable params line in the full summary. Everything else — the
conditional-LoRA generator MLPs nested inside each backbone module, plus the heads and the signal
model — is trainable, and it is not small: adding it takes the model from ~177M to ~1.19B
parameters, so ~85% of the full network is trained. Bolero's conditional LoRA is a genuine
hypernetwork that regenerates the frozen backbone's effective weights for each cell state, not a thin
low-rank patch bolted on top.
The production configuration¶
The single-dataset run behind the paper is the same config with the four demo overrides removed — i.e. the full number of batches per epoch, over more epochs. Concretely:
config = {
# ... identical identity / data / model keys as above ...
"signal_model": True,
"max_epochs": 20, # the four quick-demo lines below are simply dropped:
# "warmup_steps": 500, # trainer default (demo used 10)
# "train_batches": 5000, # trainer default (demo used 50)
# "val_batches": 300, # trainer default (demo used 10)
}
Leaving warmup_steps, train_batches, and val_batches unset lets them fall back to their
defaults (500 / 5000 / 300). At those settings each epoch takes ~50 min on an H100 80 GB
GPU, so a full 20-epoch run is roughly ~17 hours. At the default batch_size of 4 the
model peaks at ~57 GB of GPU memory, so it fits comfortably on an 80 GB H100 — lower batch_size
if you have less. Everything else — learning rate, the count head, the signal model — is unchanged
from the demo.
For real runs, drop the WANDB_MODE="offline" line from the setup cell (or set it to "online")
and log to your own W&B project to track the training and validation curves.
You now have a trained, cell-state–conditioned Bolero model for ChromiumPBMC under model/:
the frozen Flashzoi backbone plus conditional-LoRA adapters and the accessibility head, fit on the
pseudobulks built in the previous section — the whole training loop on a single dataset.
The paper's final model scales this same recipe to the entire Bolero-10M corpus. The next two pages show that multi-dataset training and what changes at scale: joint training across dozens of datasets and several genomes with per-dataset conditioning (multi-dataset ATAC), then adding the gene-expression task on top (+ gene).