Hardware Acceleration

PyTorch-BSF can train with PyTorch Lightning (solver="trainer"), which provides seamless support for hardware acceleration (GPUs, TPUs, etc.) and distributed training (multi-GPU, multi-node).

This page explains how to leverage these features to speed up your Bézier simplex (and simploid) fitting tasks.

Installation

Lightning is an optional dependency, which only solver="trainer" needs. Install it with the lightning extra:

pip install "pytorch-bsf[lightning]"

Without Lightning, solver="trainer" (and --solver=trainer on the command line) raises an ImportError that tells you to run this command; the default solver="lstsq" and torch_bsf.fit_sgd() do not need it. See Installation for all extras and for upgrading from a version that installed Lightning by default.

Note

By default, fit() solves the least-squares problem exactly on the CPU in double precision (solver="lstsq"), which is usually the fastest and most accurate choice, so the options on this page do not apply. They take effect when you train with the Lightning trainer instead: pass solver="trainer" to fit(), or --solver=trainer on the command line. Passing a trainer option without solver="trainer" raises a ValueError in the Python API; the command line ignores the trainer options unless --solver=trainer is given.

CLI / MLflow

When using the command-line interface or MLflow with --solver=trainer, you can control hardware usage via several flags. These flags are passed directly to the underlying PyTorch Lightning Trainer.

Accelerator and Devices

The --accelerator and --devices flags are the primary way to specify your hardware.

--accelerator: Choose the hardware backend. Supported values include cpu, gpu, tpu, hpu, mps (for Apple Silicon), or auto.
--devices: Specify which devices to use. You can provide an integer for the number of devices (e.g., 1, 2), a list of device IDs (e.g., 0,1), or -1 / auto to use all available devices.

Example: Use all available GPUs .. code-block:: bash

python -m torch_bsf –params data.csv –values results.csv –degree 3 –solver trainer –accelerator gpu –devices -1

Precision

You can reduce memory usage and increase training speed by using lower-precision floating-point arithmetic. This is particularly effective on modern GPUs (e.g., NVIDIA A100, RTX 30/40 series).

--precision: Supports 16-mixed, bf16-mixed, 32-true (default), or 64-true.

Example: Mixed precision training .. code-block:: bash

python -m torch_bsf –params data.csv –values results.csv –degree 3 –solver trainer –accelerator gpu –precision 16-mixed

Multi-Node Training

To scale training across multiple machines, use the --num_nodes flag.

--num_nodes: The number of machines (nodes) to use.

Example: Training on 4 nodes with 4 GPUs each .. code-block:: bash

# Run this on each node (usually handled by a cluster manager like SLURM) python -m torch_bsf –params data.csv –values results.csv –degree 3 –solver trainer –accelerator gpu –devices 4 –num_nodes 4

Python API

Pass solver="trainer" together with any PyTorch Lightning Trainer arguments as keyword arguments to the fit() function.

import torch_bsf

# Use multi-GPU and mixed precision via API
bs = torch_bsf.fit(
    params=ts,
    values=xs,
    degree=3,
    solver="trainer",
    accelerator="gpu",
    devices=2,
    precision="16-mixed",
)

The same holds for a Bézier simploid (see Simploidal Sets):

bs = torch_bsf.fit(ts, xs, degree=3, simploid_type=(2, 1), solver="trainer", accelerator="gpu", devices=2)

The returned model can be moved to any device for prediction, like any PyTorch module (e.g., bs.to("cuda")).

Detailed Use Cases

Large-Scale Pareto Front Approximation

When dealing with thousands of Pareto optimal points and high-degree Bézier simplices (e.g., degree 10+), the number of control points increases significantly. The exact "lstsq" solver then solves a large dense linear system on the CPU; if it becomes too slow or too large for memory, train with solver="trainer" and:

Use GPUs: Fitting is highly parallelizable.
Vectorized Forward: PyTorch-BSF uses a highly efficient, fully vectorized forward pass. This means that instead of looping through each control point, it uses matrix operations, which significantly speeds up computation on GPUs.
Use Mixed Precision: Set precision="16-mixed" to fit larger models into GPU memory.
Increase Batch Size: Use --batch_size to optimize GPU throughput.

Distributed Training on Clusters

If your dataset is massive or you are performing an extensive search over hyperparameters, you can use multi-node training. PyTorch-BSF supports all distributed strategies provided by Lightning (DDP, FSDP, etc.).

Tip

For most cases, strategy="auto" is sufficient. If you encounter issues on specialized clusters, you might need to specify strategy="ddp".