Training Module¶

def train_epoch(
    self, progress: Progress | None = None, task_id: TaskID | None = None
) -> tuple[float, float]:
    """
    Train for one epoch.

    Args:
        progress: Rich Progress instance (optional)
        task_id: Progress task ID (optional)

    Returns:
        Tuple of (average_loss, accuracy)
    """
    self.model.train()
    if self.adversarial_head is not None:
        self.adversarial_head.train()
    if self.cl_regression_head is not None:
        self.cl_regression_head.train()

    # Training metrics reduce to a mean loss and an accuracy, so nothing
    # here needs the predictions themselves -- a running correct-count is
    # enough and keeps the per-epoch prediction arrays (one boxed
    # np.float32 per prediction) from ever existing. Totals stay on the
    # device until the epoch ends; see _DeviceTally.
    tally = _DeviceTally()
    seen = 0

    for batch in self.train_loader:
        self.optimizer.zero_grad()

        # Gradient accumulation: N sub-batches, each scaled by 1/N, one
        # optimizer step per full batch. N=1 leaves the graph untouched.
        sub_batches = split_batch(batch, self.grad_accum_split)
        num_splits = len(sub_batches)

        for micro_step, sub_batch in enumerate(sub_batches):
            # no_sync on every micro-step but the last: N sub-batches
            # otherwise pay N allreduces for one step's arithmetic.
            with self._accumulating(micro_step == num_splits - 1):
                if self.use_mixed_precision:
                    with torch.amp.autocast("cuda"):
                        logits, labels, main_loss, loss, adv, cl_loss = (
                            self._compute_batch_loss(sub_batch)
                        )
                    if num_splits > 1:
                        loss = loss / num_splits
                    self.scaler.scale(loss).backward()
                else:
                    logits, labels, main_loss, loss, adv, cl_loss = self._compute_batch_loss(
                        sub_batch
                    )
                    if num_splits > 1:
                        loss = loss / num_splits
                    loss.backward()

                # Track metrics (sub-batch losses are averaged back to a
                # per-batch mean so history is comparable across split settings)
                tally.add("loss", main_loss, num_splits)
                if adv is not None:
                    adv_loss, adv_preds, adv_labels = adv
                    tally.add("adv_loss", adv_loss, num_splits)
                    # `mask.any()` would sync; the masked matches and the
                    # masked count are all the accuracy needs, and both are
                    # correct when the mask is empty.
                    mask = adv_labels != -1
                    tally.add("adv_correct", ((adv_preds == adv_labels) & mask).sum())
                    tally.add("adv_seen", mask.sum())
                if cl_loss is not None:
                    tally.add("cl_loss", cl_loss, num_splits)

                labels_flat = labels.detach().flatten()
                if self.loss_type == "cross_entropy":
                    # argmax(logits) is exactly the softmax > 0.5 threshold
                    # (softmax is monotonic per-row in the logits), for
                    # both binary and multi-class CE -- so this reads the
                    # predicted class straight off the logits and never
                    # computes the fp16 softmax value autocast would
                    # otherwise hand back for a probability nothing here
                    # needs (#264).
                    preds = torch.argmax(logits, dim=-1).detach().flatten()
                else:
                    # sigmoid(logits) > 0.5 is exactly logits > 0; comparing
                    # the raw logit sidesteps autocast's fp16 sigmoid
                    # output entirely instead of quantizing a probability
                    # this branch only ever thresholds (#264).
                    preds = logits.detach().flatten() > 0
                matches = preds.to(torch.float64) == labels_flat.to(torch.float64)
                tally.add("correct", matches.sum())
                seen += labels_flat.numel()

        # Gradient clipping + optimizer step (once per full batch)
        needs_clip = self.clip_grad_fn is not None or self.max_grad_norm > 0
        if needs_clip:
            if self.use_mixed_precision:
                self.scaler.unscale_(self.optimizer)
            if self.clip_grad_fn is not None:
                self.clip_grad_fn(self.model.parameters())
            else:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

        if self.use_mixed_precision:
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()

        # Update progress if provided
        if progress is not None and task_id is not None:
            progress.update(task_id, advance=1)

    # Compute metrics. accuracy_score normalizes to matches / samples, so
    # the running counts reproduce it exactly for both the multi-class
    # (argmax) and the binary (threshold at 0.5) case.
    #
    # Under DDP each rank has tallied only its own shard, so the sums are
    # reduced BEFORE they are divided -- and the divisor grows with the
    # world, since every rank ran the same number of batches. Averaging the
    # per-rank means instead would be the same number here and a different
    # one the moment the shards stop being equal.
    totals = {
        "loss": tally.value("loss"),
        "correct": tally.value("correct"),
        "seen": float(seen),
    }
    if self.adversarial_head is not None:
        totals["adv_loss"] = tally.value("adv_loss")
        totals["adv_correct"] = tally.value("adv_correct")
        totals["adv_seen"] = tally.value("adv_seen")
    if self.cl_regression_head is not None:
        totals["cl_loss"] = tally.value("cl_loss")
    totals = all_reduce_sum(totals, self.dist, self.device)
    num_batches = len(self.train_loader) * self.dist.world_size

    avg_loss = totals["loss"] / num_batches
    accuracy = totals["correct"] / totals["seen"] if totals["seen"] else 0.0

    # Adversarial metrics
    if self.adversarial_head is not None:
        avg_adv_loss = totals["adv_loss"] / max(1, num_batches)
        adv_seen = totals["adv_seen"]
        adv_acc = totals["adv_correct"] / adv_seen if adv_seen else 0.0
        self.history["train_adv_loss"].append(avg_adv_loss)
        self.history["train_adv_acc"].append(adv_acc)

    # CL regression metrics
    if self.cl_regression_head is not None:
        avg_cl_loss = totals["cl_loss"] / max(1, num_batches)
        self.history["train_cl_loss"].append(avg_cl_loss)

    return avg_loss, accuracy

def validate(
    self, progress: Progress | None = None, task_id: TaskID | None = None
) -> tuple[float, float, float, float]:
    """
    Validate model.

    Args:
        progress: Rich Progress instance (optional)
        task_id: Progress task ID (optional)

    Returns:
        Tuple of (average_loss, accuracy, roc_auc, f1)
    """
    if self.val_loader is None:
        return 0.0, 0.0, 0.0, 0.0

    self.model.eval()
    if self.cl_regression_head is not None:
        self.cl_regression_head.eval()
    # AUROC needs the probabilities themselves, so validation keeps them --
    # but as one array per batch, never as a list of boxed numpy scalars
    # (~40 bytes each while alive, and unboxed one at a time at the end).
    tally = _DeviceTally()
    all_preds: list[np.ndarray] = []
    all_labels: list[np.ndarray] = []
    # For multiclass: keep full softmax probabilities so we can compute
    # macro one-vs-rest AUROC (argmax-only preds lose that information).
    all_probs_mc: list[np.ndarray] = []
    val_seen = 0

    with torch.inference_mode():
        for batch in self.val_loader:
            # Move labels to device
            labels = batch["label"].to(self.device, non_blocking=True)

            # Adapt labels for CrossEntropyLoss
            if self.loss_type == "cross_entropy":
                ce_labels = labels.squeeze(-1).long()
            else:
                ce_labels = None

            # Forward pass (wrapper handles moving tensors and calling model correctly)
            if self.use_mixed_precision:
                with torch.amp.autocast("cuda"):
                    logits = self.model_wrapper.forward_batch(batch, self.device)
                    if ce_labels is not None:
                        loss = self.criterion(logits, ce_labels)
                    elif self.loss_type == "noise_corrected_bce":
                        loss = self.criterion(
                            logits, labels, noise_rate=self._batch_noise_rate(batch)
                        )
                    else:
                        loss = self.criterion(logits, labels)
            else:
                logits = self.model_wrapper.forward_batch(batch, self.device)
                if ce_labels is not None:
                    loss = self.criterion(logits, ce_labels)
                elif self.loss_type == "noise_corrected_bce":
                    loss = self.criterion(
                        logits, labels, noise_rate=self._batch_noise_rate(batch)
                    )
                else:
                    loss = self.criterion(logits, labels)

            # The metrics below need real probability resolution: under
            # autocast the final Linear emits fp16, and unlike softmax
            # (which autocast promotes to fp32 on its own), torch.sigmoid
            # is not on that list -- computing it on the raw logits would
            # quantize the AUROC ranking and val_auc, the checkpoint
            # -selection criterion (#264). The loss above is deliberately
            # computed on the pre-cast logits -- BCE/CE are already
            # autocast's fp32-safe, so this cast must not move earlier.
            logits = logits.float()

            # CL regression validation loss
            if (
                self.cl_regression_head is not None
                and "cl_target" in batch
                and self.model_wrapper.captured_repr is not None
            ):
                cl_targets = batch["cl_target"].to(self.device, non_blocking=True)
                cl_mask = cl_targets >= 0
                if cl_mask.any():
                    cl_preds = self.cl_regression_head(
                        self.model_wrapper.captured_repr[cl_mask]
                    )
                    tally.add(
                        "cl_loss",
                        nn.functional.mse_loss(cl_preds, cl_targets[cl_mask]),
                    )

            # Track metrics. The sample-weighted sum is what the
            # distributed reduction divides; the unweighted per-batch mean
            # below is kept for the single-GPU number, which must not move.
            tally.add("loss", loss)
            tally.add("loss_sum", loss * labels.shape[0])
            val_seen += int(labels.shape[0])
            if self.loss_type == "cross_entropy" and self._num_out > 2:
                probs_mc = torch.softmax(logits, dim=-1).cpu().numpy()
                all_probs_mc.append(probs_mc)
                all_preds.append(probs_mc.argmax(axis=-1).ravel())
            elif self.loss_type == "cross_entropy":
                all_preds.append(torch.softmax(logits, dim=-1)[:, 1].cpu().numpy().ravel())
            else:
                all_preds.append(torch.sigmoid(logits).cpu().numpy().ravel())
            all_labels.append(labels.cpu().numpy().ravel())

            # Update progress if provided
            if progress is not None and task_id is not None:
                progress.update(task_id, advance=1)

    # Compute metrics
    avg_loss = tally.value("loss") / len(self.val_loader)
    preds_arr = np.concatenate(all_preds) if all_preds else np.zeros(0, dtype=np.float32)
    labels_arr = np.concatenate(all_labels) if all_labels else np.zeros(0, dtype=np.float32)

    if self.dist.enabled:
        # Every rank scored a different slice, and AUROC does not decompose
        # over slices -- so the predictions themselves are gathered and the
        # metrics computed once, globally. That makes the 2-GPU numbers
        # equal to the 1-GPU ones rather than merely close.
        preds_arr = np.concatenate(gather_arrays(preds_arr, self.dist))
        labels_arr = np.concatenate(gather_arrays(labels_arr, self.dist))
        if all_probs_mc:
            all_probs_mc = [
                np.concatenate(
                    gather_arrays(np.concatenate(all_probs_mc, axis=0), self.dist), axis=0
                )
            ]
        # The loss is the one metric that cannot be recovered exactly: a
        # per-batch mean depends on where the batch boundaries fell, and
        # sharding moves them. The sample-weighted mean is reported
        # instead, which differs from the single-GPU number only by the
        # weight of the tail batch.
        reduced = all_reduce_sum(
            {
                "loss_sum": tally.value("loss_sum"),
                "n": float(val_seen),
                "cl_loss": tally.value("cl_loss"),
                "batches": float(len(self.val_loader)),
            },
            self.dist,
            self.device,
        )
        avg_loss = reduced["loss_sum"] / reduced["n"] if reduced["n"] else 0.0
    if self._num_out > 2:
        # Multi-class: preds are class indices; probs are needed for AUROC.
        accuracy = accuracy_score(labels_arr, preds_arr)
        f1 = f1_score(labels_arr, preds_arr, average="macro", zero_division=0.0)
        present = np.unique(labels_arr)
        # Macro one-vs-rest AUROC computed per-class. We avoid passing the
        # sliced probs to roc_auc_score(multi_class="ovr") because sklearn
        # rejects scores that don't sum to 1.0 across columns — which fires
        # whenever the val set is missing any of the model's output classes
        # (silent ValueError → AUC reported as 0.0). Manual OvR sidesteps
        # the check and lets us skip classes that are degenerate (zero
        # positives or zero negatives) in this validation set.
        if len(present) > 1 and all_probs_mc:
            probs_all = np.concatenate(all_probs_mc, axis=0)
            per_class_aucs: list[float] = []
            for c in present:
                c_int = int(c)
                if c_int >= probs_all.shape[1]:
                    continue
                y_bin = (labels_arr == c).astype(int)
                pos = int(y_bin.sum())
                if pos == 0 or pos == len(y_bin):
                    continue
                try:
                    per_class_aucs.append(float(roc_auc_score(y_bin, probs_all[:, c_int])))
                except ValueError as e:
                    logger.warning(f"AUROC failed for class {c_int}: {e}")
            auc = float(np.mean(per_class_aucs)) if per_class_aucs else 0.0
        else:
            auc = 0.0
    else:
        all_preds_binary = (preds_arr > 0.5).astype(int)
        accuracy = accuracy_score(labels_arr, all_preds_binary)
        auc = roc_auc_score(labels_arr, preds_arr) if len(np.unique(labels_arr)) > 1 else 0.0
        f1 = f1_score(labels_arr, all_preds_binary, zero_division=0.0)
        # Stashed for _current_selection_value: the parametric metrics
        # need the raw gathered (labels, probs), not a scalar already
        # reduced here. Post-gather under DDP -- the same population
        # val_acc/val_auc/val_f1 above were computed from -- so a
        # parametric metric is "computed once on the gathered
        # predictions" exactly like the others (see CLAUDE.md).
        self._last_val_labels = labels_arr
        self._last_val_probs = preds_arr

    # CL regression validation metrics
    if self.cl_regression_head is not None:
        if self.dist.enabled:
            avg_cl_loss = reduced["cl_loss"] / max(1.0, reduced["batches"])
        else:
            avg_cl_loss = tally.value("cl_loss") / max(1, len(self.val_loader))
        self.history["val_cl_loss"].append(avg_cl_loss)

    return avg_loss, accuracy, auc, f1

def train(self, epochs: int, early_stopping_patience: int = 10) -> dict[str, Any]:
    """
    Train model for multiple epochs.

    Args:
        epochs: Number of epochs to train
        early_stopping_patience: Stop if validation accuracy doesn't improve for N epochs (0 to disable)

    Returns:
        Training history dictionary
    """
    # NOT reset here -- self.patience_counter carries over from a resumed
    # checkpoint (or starts at 0 from __init__ on a fresh Trainer), so
    # early stopping continues counting rather than getting a fresh
    # `early_stopping_patience` epochs on every retry of an interrupted run.
    total_epochs = max(0, epochs - self.start_epoch + 1)
    last_epoch = self.start_epoch - 1

    # If training already completed (resume past final epoch), save best and exit
    if self.start_epoch > epochs:
        logger.info(
            f"Training already complete (resumed at epoch {self.start_epoch - 1}, "
            f"requested {epochs}). Saving checkpoints and exiting."
        )
        self._finalize(last_epoch=self.start_epoch - 1)
        barrier(self.dist)
        return self.history

    # Create progress bars
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
        console=console,
        disable=not self._is_main,
    ) as progress:
        epoch_task = progress.add_task("[cyan]Training epochs...", total=total_epochs)

        for epoch in range(self.start_epoch, epochs + 1):
            last_epoch = epoch
            # Re-seed the shared draw. Skipping this trains every epoch on
            # the same shard -- which converges, and reports metrics from a
            # world_size-th of the corpus.
            if self.dist.enabled:
                epoch_sampler = getattr(self.train_loader, "sampler", None)
                if hasattr(epoch_sampler, "set_epoch"):
                    epoch_sampler.set_epoch(epoch)
            # LR warmup. The cosine schedule folds warmup into its own
            # LambdaLR multiplier, so only the plateau/none paths need the
            # manual ramp here.
            if (
                self.scheduler_type != "cosine"
                and self.warmup_epochs > 0
                and epoch <= self.warmup_epochs
            ):
                warmup_lr = self.base_lr * (epoch / self.warmup_epochs)
                for param_group in self.optimizer.param_groups:
                    param_group["lr"] = warmup_lr

            # Adversarial lambda annealing (linear ramp)
            if self.adversarial_head is not None and self.adversarial_anneal_epochs > 0:
                progress_frac = min(1.0, epoch / self.adversarial_anneal_epochs)
                current_lambda = self.adversarial_lambda * progress_frac
                self.adversarial_head.set_lambda(current_lambda)

            # Create tasks for this epoch
            train_task = progress.add_task(
                f"[green]Epoch {epoch}/{epochs} - Training", total=len(self.train_loader)
            )

            # Train
            train_loss, train_acc = self.train_epoch(progress, train_task)
            self.history["train_loss"].append(train_loss)
            self.history["train_acc"].append(train_acc)

            progress.remove_task(train_task)

            # Validate
            if self.val_loader is not None:
                val_task = progress.add_task(
                    f"[yellow]Epoch {epoch}/{epochs} - Validation", total=len(self.val_loader)
                )
                val_loss, val_acc, val_auc, val_f1 = self.validate(progress, val_task)
                self.history["val_loss"].append(val_loss)
                self.history["val_acc"].append(val_acc)
                self.history["val_auc"].append(val_auc)
                self.history["val_f1"].append(val_f1)
                current_value = self._current_selection_value(val_acc, val_auc, val_f1)
                self.history["val_selection"].append(current_value)
                is_parametric_metric = self._checkpoint_metric not in (
                    "val_acc",
                    "val_f1",
                    "val_auc",
                )

                progress.remove_task(val_task)

                # LR scheduler step (only after warmup). For reduce_on_plateau
                # we step on the *selection metric*, not val_loss, so the
                # plateau detection matches what early-stopping and best-
                # model checkpointing react to. If val_loss were used while
                # selection runs off F1/AUC, LR can drop because loss
                # plateaued on easy examples even though the ranking metric
                # is still improving — burning the schedule prematurely.
                # The cosine LambdaLR owns warmup itself, so it steps every
                # epoch; plateau keeps the original post-warmup guard.
                if self.scheduler is not None and (
                    self.scheduler_type == "cosine" or epoch > self.warmup_epochs
                ):
                    old_lr = self.optimizer.param_groups[0]["lr"]
                    if self.scheduler_type == "reduce_on_plateau":
                        self.scheduler.step(-current_value)
                    else:
                        # Epoch-based schedulers (cosine, etc.)
                        self.scheduler.step()
                    new_lr = self.optimizer.param_groups[0]["lr"]
                    if new_lr != old_lr:
                        logger.info(f"LR reduced: {old_lr:.6f} -> {new_lr:.6f}")

                # Display metrics — [*] marks the checkpoint criterion
                lr_str = ""
                current_lr = self.optimizer.param_groups[0]["lr"]
                if current_lr != self.base_lr:
                    lr_str = f" LR: {current_lr:.6f}"
                acc_label = "Acc[*]" if self._checkpoint_metric == "val_acc" else "Acc"
                f1_label = "F1[*]" if self._checkpoint_metric == "val_f1" else "F1"
                auc_label = "AUC[*]" if self._checkpoint_metric == "val_auc" else "AUC"
                sel_str = (
                    f" | {self._checkpoint_metric}[*]: {current_value:.4f}"
                    if is_parametric_metric
                    else ""
                )
                adv_str = ""
                if self.adversarial_head is not None and self.history["train_adv_loss"]:
                    _adv_l = self.history["train_adv_loss"][-1]
                    _adv_a = self.history["train_adv_acc"][-1]
                    _lam = self.adversarial_head.grl.lambda_
                    adv_str = f" | Adv: L={_adv_l:.3f} Acc={_adv_a:.3f} λ={_lam:.3f}"
                cl_str = ""
                if self.cl_regression_head is not None and self.history["train_cl_loss"]:
                    _tcl = self.history["train_cl_loss"][-1]
                    _vcl = (
                        self.history["val_cl_loss"][-1] if self.history["val_cl_loss"] else 0.0
                    )
                    cl_str = f" | CL: train={_tcl:.4f} val={_vcl:.4f}"
                self._print(
                    f"[cyan]Epoch {epoch}/{epochs}[/cyan] | "
                    f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
                    f"Val Loss: {val_loss:.4f} {acc_label}: {val_acc:.4f} "
                    f"{f1_label}: {val_f1:.4f} {auc_label}: {val_auc:.4f}"
                    f"{sel_str}{lr_str}{adv_str}{cl_str}"
                )

                # Save best model per configured checkpoint metric
                improved = current_value > self.best_val_selection
                if improved:
                    self.best_val_acc = val_acc
                    self.best_val_f1 = val_f1
                    self.best_val_auc = val_auc
                    self.best_val_selection = current_value
                    self.best_epoch = epoch
                    self._best_model_state = copy.deepcopy(self.model.state_dict())
                    self.patience_counter = 0

                    if self.output_dir:
                        self.save_checkpoint("model_best.pt", epoch=epoch)
                        sel_note = (
                            f", {self._checkpoint_metric}: {current_value:.4f}"
                            if is_parametric_metric
                            else ""
                        )
                        self._print(
                            f"[bold green]✓ Saved best model "
                            f"(val_acc: {val_acc:.4f}, val_f1: {val_f1:.4f}, "
                            f"val_auc: {val_auc:.4f}{sel_note})[/bold green]"
                        )
                else:
                    self.patience_counter += 1

                # Early stopping (disabled if patience is 0)
                if (
                    early_stopping_patience > 0
                    and self.patience_counter >= early_stopping_patience
                ):
                    self._print(f"[yellow]Early stopping at epoch {epoch}[/yellow]")
                    break
            else:
                self._print(
                    f"[cyan]Epoch {epoch}/{epochs}[/cyan] | "
                    f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f}"
                )

            # Rolling resume checkpoint (issue #330): written every epoch,
            # after everything above (validation, best-model tracking, the
            # scheduler step) has updated for it, so it always reflects the
            # state the *next* epoch should start from. Deliberately not a
            # declared Snakemake output -- see `_finalize` -- so a rule
            # retry after a walltime kill still finds it on disk.
            if self.output_dir:
                self.save_checkpoint("model_resume.pt", epoch=epoch)

            progress.update(epoch_task, advance=1)

    # Save final model. The writers are no-ops off rank 0; the barrier is
    # what keeps a rank from tearing down its process group while rank 0 is
    # still inside a collective-free write the others would outlive.
    self._finalize(last_epoch=last_epoch)
    barrier(self.dist)

    return self.history