Add full fine-tuning support to the tabular benchmark harness

benchmark/tabular/README.md

@@ -50,6 +50,17 @@ python -m benchmark.tabular.tabarena.main \
   --dataset blood-transfusion-service-center
 ```
 
+### Fine-tuning
+
+`--model tabiclv2-ft`, `--model kumo-tabular-large-ft`, `--model kumo-tabular-small-ft`, and `--model tabfm-ft` full fine-tune every parameter of the model on each dataset's training split before evaluating (see `benchmark/tabular/finetune.py`):
+
+```bash
+python -m benchmark.tabular.tabarena.main \
+  --model tabiclv2-ft \
+  --finetune_epochs 75 \
+  --finetune_lr 1e-6
+```
+
 ### Evaluate
 
 Evaluate all available model results with:
@@ -106,6 +117,14 @@ python -m benchmark.tabular.beyondarena.main \
 
 Available subset filters include problem types (`classification`, `regression`), size buckets (`tiny`, `small`, `medium`, `large`), split regimes (`iid`, `temporal`, `grouped`), feature groups (`low-dim`, `high-dim`, `text`, `high-cardinality`), and split selections (`core`, `lite`, `all`). Prefix a filter with `!` to negate it.
 
+### Fine-tuning
+
+Same `-ft` model variants and `--finetune_*` flags as TabArena above:
+
+```bash
+python -m benchmark.tabular.beyondarena.main --model kumo-tabular-small-ft --subset lite
+```
+
 ### Evaluate
 
 Evaluate all available model results with:

benchmark/tabular/beyondarena/main.py

@@ -12,7 +12,11 @@
 from tabarena.contexts import BeyondArenaContext
 from tabarena.utils.config_utils import ConfigGenerator
 
-from benchmark.tabular.model import MODEL_CONFIGS
+from benchmark.tabular.model import (
+    MODEL_CONFIGS,
+    add_finetune_args,
+    finetune_config_overrides,
+)
 
 parser = argparse.ArgumentParser(description=__doc__)
 parser.add_argument(
@@ -45,6 +49,7 @@
     type=int,
     help="Prediction batch size.",
 )
+add_finetune_args(parser)
 parser.add_argument(
     "--enable-kv-cache",
     action="store_true",
@@ -60,7 +65,6 @@
     / model_config.name
     / "outer_model"
 )
-result_dir.mkdir(parents=True, exist_ok=True)
 
 config = {
     "max_context_size": args.max_context_size,
@@ -69,6 +73,23 @@
 }
 if args.batch_size is not None:
     config["ag.max_batch_size"] = args.batch_size
+finetune_overrides = finetune_config_overrides(args)
+if finetune_overrides and not args.model.endswith("-ft"):
+    raise ValueError(
+        f"--finetune_* flags were passed but --model {args.model!r} isn't "
+        "a '-ft' variant, so fine-tuning is disabled and the flags would "
+        "silently have no effect; pass the matching '-ft' model instead."
+    )
+config.update(finetune_overrides)
+
+if finetune_overrides:
+    # tabarena's cache key is positional, not content-based; key by
+    # hyperparameters so different LRs don't silently collide.
+    variant = "_".join(
+        f"{key}={value}" for key, value in sorted(config.items())
+    )
+    result_dir = result_dir / variant
+result_dir.mkdir(parents=True, exist_ok=True)
 
 generator = ConfigGenerator(
     search_space={},

benchmark/tabular/finetune.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared full fine-tuning core for the tabular benchmark harness.
+
+Fine-tunes every parameter of an :class:`~sdm.models.ICLModel` with gradient
+descent, generic over ``x``/``y`` :class:`TableTensor` pools so
+:func:`full_finetune` drops in before any adapter's ``model.fit(...)`` call.
+"""
+
+import copy
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+
+import sdm
+
+_EPS = 1e-12
+
+# Shared by TabArena, BeyondArena, ScoringBench, and TALENT.
+FINETUNE_EPOCHS = 75
+FINETUNE_ITERS_PER_EPOCH = 10
+FINETUNE_LR = 1e-6
+FINETUNE_TRAIN_SIZE = 10_000
+FINETUNE_CONTEXT_FRAC = 0.8
+FINETUNE_VAL_FRAC = 0.2
+
+
+def cpu_state_dict(model: torch.nn.Module) -> dict[str, Any]:
+    """Copy weights to CPU while preserving extra state and metadata."""
+    state = model.state_dict()
+    for name, value in state.items():
+        state[name] = (
+            value.to(device="cpu", copy=True)
+            if isinstance(value, torch.Tensor)
+            else copy.deepcopy(value)
+        )
+    return state
+
+
+def kumo_small_binary_epochs(
+    epochs: int | None,
+    *,
+    is_kumo_small: bool,
+    is_binary: bool,
+) -> int:
+    """Respect explicit epochs; otherwise use the model-specific default.
+
+    Empirically found to need fewer epochs than the shared default to avoid
+    overfitting on binary classification.
+    """
+    if epochs is not None:
+        return epochs
+    return 50 if is_kumo_small and is_binary else FINETUNE_EPOCHS
+
+
+def evaluate(
+    model: sdm.models.ICLModel,
+    context_x: sdm.TableTensor,
+    context_y: sdm.TableTensor,
+    query_x: sdm.TableTensor,
+    query_y: sdm.TableTensor,
+    *,
+    task: str,
+    num_estimators: int,
+    generator: torch.Generator | None,
+) -> float:
+    """Zero-shot in-context evaluation via the model's own default recipe."""
+    model.eval()
+    with torch.inference_mode():
+        out = model(
+            x_context=context_x,
+            y_context=context_y,
+            x_query=query_x,
+            num_estimators=num_estimators,
+            generator=generator,
+        )
+        if task == "classification":
+            scores, indices = sdm.evaluation.to_class_indices(
+                out,
+                query_y,
+                missing_score=0.0,
+            )
+            return (scores.argmax(dim=-1) == indices).float().mean().item()
+
+        pred = out.numerical.mean(dim=-1)
+        target = query_y.numerical.squeeze(-1)
+        return (pred - target).pow(2).mean().sqrt().item()  # RMSE
+
+
+def full_finetune(
+    model: sdm.models.ICLModel,
+    x_pool: sdm.TableTensor,
+    y_pool: sdm.TableTensor,
+    *,
+    task: Literal["classification", "regression"],
+    max_epochs: int = FINETUNE_EPOCHS,
+    iters_per_epoch: int = FINETUNE_ITERS_PER_EPOCH,
+    train_size: int = FINETUNE_TRAIN_SIZE,
+    context_frac: float = FINETUNE_CONTEXT_FRAC,
+    val_frac: float = FINETUNE_VAL_FRAC,
+    lr: float = FINETUNE_LR,
+    num_estimators: int,
+    max_val_context_size: int | None = None,
+    generator: torch.Generator | None,
+) -> float:
+    """Full fine-tune every parameter of ``model`` in place.
+
+    Carves a ``val_frac`` validation split from ``x_pool``/``y_pool``; each
+    epoch trains on resampled context/query batches from the rest, then
+    checkpoints if validation improves. Leaves ``model`` in eval mode
+    holding its best-validation-metric weights and returns that metric
+    (accuracy for classification, RMSE for regression -- lower is
+    better for regression).
+    """
+    device = x_pool.device
+    n = x_pool.size(0)
+
+    perm = torch.randperm(n, generator=generator, device=device)
+    n_val = int(val_frac * n)
+    val_idx, train_idx = perm[:n_val], perm[n_val:]
+    x_train, y_train = x_pool[train_idx], y_pool[train_idx]
+    x_val, y_val = x_pool[val_idx], y_pool[val_idx]
+
+    val_context_x, val_context_y = x_train, y_train
+    if (
+        max_val_context_size is not None
+        and val_context_x.size(0) > max_val_context_size
+    ):
+        cap = torch.randperm(
+            val_context_x.size(0), generator=generator, device=device
+        )
+        cap = cap[:max_val_context_size]
+        val_context_x, val_context_y = val_context_x[cap], val_context_y[cap]
+
+    n_train = x_train.size(0)
+    epoch_size = min(train_size, n_train)
+    n_context = int(context_frac * epoch_size)
+
+    if n_val < 1:
+        raise ValueError(
+            f"val_frac={val_frac} leaves no validation rows out of "
+            f"{n} pool rows; pass a larger val_frac or pool."
+        )
+    if n_train < 2:
+        raise ValueError(
+            f"Only {n_train} training rows remain after the validation "
+            "split; need at least 2 to form a non-empty context and query."
+        )
+    if n_context < 1 or epoch_size - n_context < 1:
+        raise ValueError(
+            f"context_frac={context_frac} splits {epoch_size} sampled rows "
+            f"into an empty context or query ({n_context} / "
+            f"{epoch_size - n_context}); pick a context_frac strictly "
+            "between 0 and 1 relative to the sample size."
+        )
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    higher_is_better = task == "classification"
+
+    def current_metric() -> float:
+        return evaluate(
+            model,
+            val_context_x,
+            val_context_y,
+            x_val,
+            y_val,
+            task=task,
+            num_estimators=num_estimators,
+            generator=generator,
+        )
+
+    best_metric = current_metric()
+    best_state = cpu_state_dict(model)
+
+    for _epoch in range(max_epochs):
+        model.train()
+        for _iter in range(iters_per_epoch):
+            row = torch.randperm(n_train, generator=generator, device=device)
+            row = row[:epoch_size]
+            x_context, y_context = (
+                x_train[row[:n_context]],
+                y_train[row[:n_context]],
+            )
+            x_query, y_query = (
+                x_train[row[n_context:]],
+                y_train[row[n_context:]],
+            )
+
+            optimizer.zero_grad()
+            out = model(
+                x_context=x_context,
+                y_context=y_context,
+                x_query=x_query,
+                num_estimators=1,
+                generator=generator,
+            )
+            if task == "classification":
+                scores, indices = sdm.evaluation.to_class_indices(
+                    out,
+                    y_query,
+                    missing_score=0.0,
+                )
+                loss = F.nll_loss(scores.clamp_min(_EPS).log(), indices)
+            else:
+                # Pinball loss over quantile columns; TabFM's single
+                # point-estimate column has none, so this is squared error.
+                diff = y_query.numerical - out.numerical
+                n_quantiles = out.numerical.size(-1)
+                if n_quantiles > 1:
+                    levels = torch.linspace(
+                        0.001, 0.999, n_quantiles, device=device
+                    )
+                    loss = torch.maximum(
+                        levels * diff,
+                        (levels - 1) * diff,
+                    ).mean()
+                else:
+                    loss = diff.pow(2).mean()
+            loss.backward()
+            optimizer.step()
+
+        metric_value = current_metric()
+        improved = (
+            metric_value > best_metric
+            if higher_is_better
+            else metric_value < best_metric
+        )
+        if improved:
+            best_metric = metric_value
+            best_state = cpu_state_dict(model)
+
+    model.load_state_dict(best_state)
+    model.zero_grad(set_to_none=True)
+    model.eval()
+    return best_metric