Setup#

Setting up is three things: an environment created with uv, a .env holding the paths and switches every entry point reads, and a data root pointing at the build you downloaded in the dataset quick start.

Nothing here needs credentials. The brainsets on S3 are public, and W&B logging is off until you turn it on.

Requirements#

  • Python 3.10 or newer.

  • uv, which creates the environment and installs the dependencies.

  • A CUDA GPU for anything past the smallest baselines.

Installing#

The training code is meant to be edited, not imported from a Python library (at least for now), so you work inside the clone and install it editable. To version and push your own progress, fork the repository first and clone that instead. Make sure uv is installed.

# or your own fork: git clone https://github.com/<you>/ibl-bwb.git
git clone https://github.com/brainbench-org/ibl-bwb.git
cd ibl-bwb

uv venv .venv -p python3.10
source .venv/bin/activate
uv pip install -e ".[train]"

Configuring .env#

Environment variables and paths for the benchmark are set in a .env at the repo root, which is gitignored. Edit it and the change applies on the next run: every training entry point loads the file, so there is nothing to re-source.

Copy the annotated template:

cp .env.example .env

Every variable in it starts commented out, so uncomment and fill only the ones you need.

The ones you are most likely to edit:

.env#
BWB_CKPT_DIR=./ckpt                # checkpoints, and what a relative ckpt.load_from resolves against
BWB_PREDICTIONS_DIR=./predictions  # saved predictions, for scoring
BWB_EMBS_DIR=./embs                # TS3 unit embeddings
BWB_CACHE_DIR=./cache              # TS3 waveform/ACG cache, built on first use
WANDB_MODE=disabled                # online, offline or disabled (the default)

The four directories take relative paths, resolved against wherever you launch from, and the values above are what you get if you leave them unset. Set them if you launch from outside the repo root.

Model settings such as hyperparameters live elsewhere, in the Hydra .yaml configs (Where the code lives). There you will often see the pattern ${oc.env:VAR,default}, which is what pulls a variable out of this file: the command line beats .env, which beats the default. Leave a variable commented out rather than empty, since an exported empty string beats the default.

Setting the data root#

data_root is where a build you downloaded (dataset quick start) lives: everything above the brainset directory itself.

/abs/path/to/data/all_units/ibl_brain_wide_bench_2026/   <- the build you downloaded
/abs/path/to/data/all_units                              <- data_root

The different builds carry that same directory name ibl_brain_wide_bench_2026, so they cannot share a root.

Set the ones you have downloaded in your .env, as absolute paths:

.env#
BWB_DATA_ROOT_ALL_UNITS=/abs/path/to/data/all_units            # TS1
BWB_DATA_ROOT_SELECTED_UNITS=/abs/path/to/data/selected_units  # TS2, TS3
BWB_DATA_ROOT_PRETRAIN=/abs/path/to/data/all_units             # pretraining, either build

Each falls back to BWB_DATA_ROOT when unset, so one root covers a single build. The exception is trainer=nuclr_pretrain and nemo_pretrain: they score TS3 as they train, so they read BWB_DATA_ROOT_SELECTED_UNITS rather than BWB_DATA_ROOT_PRETRAIN, and need their pretraining sessions under the selected_units root next to the eval ones.

Where the code lives#

The benchmark codebase is organized as follows, colored by role:

shared machinery   pretraining   the three evaluation suites   the scoring contract

src/
├── core/           # shared machinery: trainer, datasets, samplers, checkpointing
├── pretrain/       # pretraining: train.py and one directory per model
├── ts1/            # decoding behavior
├── ts2/            # neural activity
├── ts3/            # brain region
├── ibl_bwb_eval/   # the published contract: tasks, metrics, scoring
├── hydra_plugins/  # config auto-discovery
└── tests/

Take NDT Stitch as an example: pretrained and then evaluated, it shows how a model is organized in the codebase.

src/pretrain/models/ndt_stitch/               # pretraining
├── ndt_stitch.py                             # the model class
├── ndt_stitch_pretrain.py                    # its trainer
└── configs/
    ├── model/ndt_stitch_10M.yaml             # reused by every suite
    └── trainer/ndt_stitch_pretrain.yaml

src/ts1/models/pretrained/ndt_stitch/         # the same model, on TS1
├── ndt_stitch_eval_trainer.py                # a trainer only
└── configs/trainer/ndt_stitch_finetune.yaml

Each entry point is a train.py you call with Hydra overrides.

python src/pretrain/train.py trainer=ndt_stitch_pretrain
  • trainer=ndt_stitch_pretrain names the trainer config pretrain/models/ndt_stitch/configs/trainer/ndt_stitch_pretrain.yaml.
  • That file asks for /model: ndt_stitch_10M, which is the model config pretrain/models/ndt_stitch/configs/model/ndt_stitch_10M.yaml.
python src/ts1/train.py trainer=ndt_stitch_finetune ckpt.load_from=<checkpoint>.pt
  • TS1 works the same way, from ts1/models/pretrained/ndt_stitch/configs/trainer/ndt_stitch_finetune.yaml, except that its /model is still the pretraining one: a suite can use the pretraining model config, and adds only a trainer.

Configs are found by name, wherever they sit, because src/hydra_plugins/ discovers every configs/ in the tree. That is why adding a model edits no central file.

Three things to know before editing:

  • Every model is one directory. As in the NDT Stitch example above, the model class, its trainer and their configs sit together, and nothing outside that directory knows about them.

  • The task suites are self-contained, on purpose. ts1/, ts2/ and ts3/ repeat each other so that someone who cares about one suite can read one directory.

  • Some parts are fixed, so that two submissions stay comparable. ibl_bwb_eval/ holds the shared contract: tasks, metrics, evaluation seeds and the prediction format. Do not edit it, nor the code that evaluates the held-out test split: the suites’ test mixins, and ts3/protocol.py, the unit table every TS3 metric is computed against.

Weights & Biases#

Weights & Biases tracks a run’s metrics and config in a dashboard you can compare runs in.

It is off until you turn it on in the .env, and every run logs from then on.

.env#
WANDB_MODE=online         # online, offline or disabled (the default, records nothing)
WANDB_API_KEY=<your key>  # online needs it, from https://wandb.ai/authorize
WANDB_ENTITY=<team>       # default: your own entity
WANDB_DIR=wandb/          # where offline keeps runs

offline is the middle ground if you want your own metrics but no account: it keeps the full run locally, ready to push later with wandb sync wandb/offline-run-*.

The project a run logs to is set in the configs, named after the entry point and the model: pretrain-ndt2, ts1-poyo, ts2-lfads, ts3-lolcat.

Checking the install#

One short run exercises the whole path, from the loader to the metrics:

python src/ts1/train.py trainer=mlp task=choice debug=true

debug=true cuts the run to one batch per split and one epoch, with W&B, checkpointing and dataloader workers off.

Optional extras#

xformers, for POYO+. POYO+ attends over chained tokens with a stock-torch kernel by default (model.attn_impl=nested). xformers provides a leaner one – chiefly a much lower peak memory, so a larger batch fits on one GPU, plus a smaller speedup – but pins a wheel to an exact torch build, so it is kept out of train:

uv pip install -e ".[train,xformers]"

Then pass model.attn_impl=xformers. Both backends compute the same attention, so checkpoints move between them freely.

Ray, for tuning. tune.py and tune.sh need a running cluster. Start a local one with ray start --head, or point RAY_ADDRESS at an existing one. BWB_RAY_RESULTS_DIR and BWB_RAY_OUTPUTS_DIR hold trial storage and result CSVs, and the first has to be on a shared filesystem for a multi-node cluster. RAY_TMPDIR moves Ray’s scratch space off /tmp when that is small.