Add estimator sharding for DDP fine-tuning

[eliott-kalfon]

Issue + Motivation

Running into OOMs doing DDP fine-tuning for 8 estimators, 4 M rows, 1M susbample, 200 features, 20k and 10k chunks, on A100 and H100.

• In the current set up, each GPU holds the activation of all estimators, generating OOMs
• The change I want to propose is for each GPU to process one estimator/divide the number of estimators among them

This will enable us to run finetuning on larger data and smaller GPU. It now works

Public API Changes

• No Public API changes
• Yes, Public API changes (Details below)

---

How Has This Been Tested?

Locally and with GPUs on the cluster. Ran some parity tests to ensure we got the same outputs.

---

Checklist

• The changes have been tested locally.
• Documentation has been updated (if the public API or usage changes).
• A changelog entry has been added (see changelog/README.md), or "no changelog needed" label requested.
• The code follows the project's style guidelines.
• I have considered the impact of these changes on the public API.

---

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit dca2280. Configure here.}

[anuragg1209]

Hi @eliott-kalfon, did you get the chance to test the accuracy by testing the same model twice with and without sharding.

[eliott-kalfon]

Hi @eliott-kalfon, did you get the chance to test the accuracy by testing the same model twice with and without sharding.

Hi @anuragg1209, yes I did with this script. Getting this results

{
"status": "ok",
"gpu_name": "NVIDIA RTX PRO 6000 Blackwell Server Edition",
"world_size": 4,
"train_rows": 2000,
"test_rows": 400,
"epochs": 2,
"n_estimators": 4,
"unsharded_accuracy": 0.965,
"sharded_accuracy": 0.9675,
"prediction_agreement": 0.9975,
"probability_max_abs_difference": 0.011978328227996826,
"probability_mean_abs_difference": 0.0003741188265848905,
"weight_max_abs_difference": 5.485536530613899e-05
}

#!/usr/bin/env python3
"""Compare unsharded single-GPU and estimator-sharded DDP fine-tuning."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch
from tabpfn.finetuning.finetuned_classifier import FinetunedTabPFNClassifier


ROOT = ""
BASELINE_PATH = ROOT / "results_4m/parity_unsharded.pt"
RESULT_PATH = ROOT / "results_4m/sharded_parity.json"


def make_data() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    rng = np.random.default_rng(20260814)
    X = rng.standard_normal((2_400, 20), dtype=np.float32)
    margin = X[:, :5].sum(axis=1) + 0.25 * X[:, 5] * X[:, 6]
    y = (margin > 0).astype(np.int64)
    return X[:2_000], y[:2_000], X[2_000:], y[2_000:]


def make_model(*, sharded: bool) -> FinetunedTabPFNClassifier:
    return FinetunedTabPFNClassifier(
        device="cuda",
        epochs=2,
        validation_split_ratio=None,
        n_finetune_ctx_plus_query_samples=1_000,
        finetune_ctx_query_split_ratio=0.2,
        random_state=17,
        early_stopping=False,
        use_lr_scheduler=False,
        n_estimators_finetune=4,
        n_estimators_validation=4,
        n_estimators_final_inference=4,
        shard_estimators_across_gpus=sharded,
        save_checkpoint_interval=None,
        use_fixed_preprocessing_seed=True,
    )


def cpu_state_dict(model: FinetunedTabPFNClassifier) -> dict[str, torch.Tensor]:
    return {
        key: value.detach().cpu().clone()
        for key, value in model.finetuned_estimator_.model_.state_dict().items()
    }


def run_baseline() -> None:
    torch.manual_seed(17)
    np.random.seed(17)
    X_train, y_train, X_test, y_test = make_data()
    model = make_model(sharded=False).fit(X_train, y_train)
    probabilities = model.predict_proba(X_test)
    predictions = probabilities.argmax(axis=1)
    torch.save(
        {
            "state_dict": cpu_state_dict(model),
            "probabilities": probabilities,
            "predictions": predictions,
            "accuracy": float(np.mean(predictions == y_test)),
        },
        BASELINE_PATH,
    )


def run_sharded() -> None:
    rank = int(os.environ["RANK"])
    torch.manual_seed(17)
    np.random.seed(17)
    X_train, y_train, X_test, y_test = make_data()
    model = make_model(sharded=True).fit(X_train, y_train)
    if rank != 0:
        return

    probabilities = model.predict_proba(X_test)
    predictions = probabilities.argmax(axis=1)
    baseline = torch.load(BASELINE_PATH, map_location="cpu", weights_only=False)
    sharded_state = cpu_state_dict(model)
    parameter_differences = [
        (sharded_state[key] - baseline["state_dict"][key]).abs().max().item()
        for key in sharded_state
        if sharded_state[key].is_floating_point()
    ]
    probability_delta = np.abs(probabilities - baseline["probabilities"])
    result = {
        "status": "ok",
        "gpu_name": torch.cuda.get_device_name(0),
        "world_size": int(os.environ["WORLD_SIZE"]),
        "train_rows": len(X_train),
        "test_rows": len(X_test),
        "epochs": 2,
        "n_estimators": 4,
        "unsharded_accuracy": baseline["accuracy"],
        "sharded_accuracy": float(np.mean(predictions == y_test)),
        "prediction_agreement": float(np.mean(predictions == baseline["predictions"])),
        "probability_max_abs_difference": float(probability_delta.max()),
        "probability_mean_abs_difference": float(probability_delta.mean()),
        "weight_max_abs_difference": max(parameter_differences),
    }
    RESULT_PATH.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result), flush=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=("baseline", "sharded"))
    args = parser.parse_args()
    if args.mode == "baseline":
        run_baseline()
    else:
        run_sharded()

[anuragg1209]

Overall, the PR looks solid to me! We are almost ready to merge it!

[anuragg1209]

Solid PR! Thanks, @eliott-kalfon, for adding this feature!