train on pytorch lightning

[selmanozleyen]

Stacked on #149 (which is stacked on #148) — review those first. This PR's own commit is
3f220fe; everything below is relative to feat/simplify_methods.

Note: GitHub keeps this PR in stack #151 (148 → 150) and refuses a base-branch change
while it is a stack member, so the base still reads feat/flatten_backends_module and the
diff shown here includes #149's commit. The branch itself is correctly based on
feat/simplify_methods.

Replaces the hand-rolled training loop with PyTorch Lightning. Net effect on src/ is a wash
in line count (+662 / -663) — the interesting part is what stopped existing.

What Lightning now owns

Was	Now
Trainer.train() — manual step loop, tqdm bar, log accumulation, valid_freq modulo	pl.Trainer(max_steps=…, val_check_interval=…)
BaseCallback / ComputationalCallback / LoggingCallback / TrainingCallbacks	pl.Callback + self.log()
WandBLogger	lightning.pytorch.loggers.WandbLogger
OptimizationManager	configure_optimizers()
set_train_mode, _device_id/_dtype, Model.to_device's optimizer-state walk	Lightning's device/mode management

And it brings things that simply did not exist before: checkpointing and resume, AMP, gradient
clipping and accumulation, early stopping, LR monitoring, deterministic seeding, multi-device.

Node-steps

The train sampler used to hand back a tuple of n_nodes nodes per call, and the loop took
one optimizer step per node inside that tuple. So the gradient schedule was already
per-node; only the counting was per-round.

TrainSampler.__iter__ now yields one node at a time, and a node is one batch and one
optimizer step. Rounds are still drawn n_nodes at a time so replace_nodes=False keeps
meaning "distinct nodes within a round". The stream is unbounded and re-iterable —
max_iter_steps now defaults to None and the trainer's n_train_steps governs run length.

Migration: a run that used n_train_steps=N with n_nodes=k now needs
n_train_steps=N * k. With the default n_nodes=1 nothing changes.

Samplers are passed to fit() as plain iterables — Lightning accepts them directly, so there
is no DataLoader layer and data/ stays torch-free.

Following scvi-tools, and where it doesn't

scvi-tools wraps its module in a separate TrainingPlan(LightningModule). Their reason is that
one module gets trained by many plans — 13+ plan classes, selected via _training_plan_cls,
swapped at runtime for TOTALANVI, plus Pyro's SVI which needs a different optimization scheme
entirely. sckitflow has exactly one training recipe and every method is
one-loss-one-optimizer-backward, so BaseMethod is the LightningModule and there is no
plan layer. If a second recipe ever shows up, extracting a plan is mechanical.

The two scvi layers that do earn their keep are kept:

• Trainer(pl.Trainer) — a thin subclass with sckitflow defaults (node-step semantics,
  check_val_every_n_epoch=None, no checkpointing, log_every_n_steps=1), mirroring their
  Trainer.
• DataFrameLogger — an in-memory pl.Logger, mirroring their SimpleLogger, so
  get_train_logs_df() / get_val_logs_df() keep working with zero configuration. Pass
  logger= to also ship to CSV/W&B; it is composed alongside, not replaced.

Two things worth a second look

Frozen dataclasses. Nodes are trees of frozen dataclasses wrapping numpy arrays, and
Lightning's default recursive device transfer refuses to walk them (ValueError: A frozen dataclass was passed to apply_to_collection). transfer_batch_to_device is overridden to
leave the node on the host; extract_step_data already does the tensor conversion and
placement per step.

Per-validation-set metric state. MetricsCallback gives each validation dataloader its own
copy of every metric. The copies are cloned from the pristine template rather than from another
dataloader's — cloning from an already-updated metric silently pooled the two sets' state.
There is a regression test for it.

Breaking

The method contract renamed in #149 (compute_loss, infer) is unchanged here. On top of that:

• The callback hierarchy is gone. Custom callbacks become pl.Callback subclasses.
• _runtime.py lost its backend switch.
• device_id defaults to None (leave in place) instead of eagerly selecting CUDA — the
  trainer's accelerator decides. Pass accelerator= to Model.train.
• Model.train no longer forwards *args/**kwargs to train_step; extra kwargs go to the
  trainer. Validation-time prediction arguments move to val_predict_kwargs (e.g. n_samples
  for noise-generating methods).
• Model.save no longer pickles the trainer, so logs and optimizer state do not survive a
  round trip. Use ModelCheckpoint for those.
• OptimConfig lost its unused plan_kwargs, gained lr_scheduler_monitor, and now resolves
  itself into the mapping configure_optimizers returns. lr_scheduler_step still accepts
  "train_step".

Verification

• Full suite green: 1253 passed, 30 skipped, 15 xfailed.
• All 4 notebooks execute via .run_notebooks.sh.
• pre-commit run --all-files clean.
• End-to-end run of real CFM through Model.train with two validation sets and
  EnergyDistance/MaximumMeanDiscrepancy: per-step train losses, per-val-id metric frames,
  and predict all check out.
• .test_durations regenerated and remapped onto flatten backends module #148's tests/core/ layout; all 51 entries
  resolve to files that exist, and --splits 2 verified.

New tests cover the parts that carry the semantics: one-node-one-optimizer-step, validation
cadence in node-steps, per-node sampler iteration (lazy, unbounded, re-iterable),
DataFrameLogger train/val routing, and the metric-isolation regression.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[lorenzo-consoli]

I would stack this on top of #149 too :)