FFCV

Large-Scale Linear Regression

In this example, we’ll see how to run large-scale regularized linear regression with FFCV (by “large-scale” here we mean a dataset that definitely doesn’t fit in GPU memory, and may barely even fit in RAM).

See here for the script corresponding to this tutorial.

Setup: Generating a fake dataset

Let’s start by generating a fake dataset on which we’ll run linear regression. We’ll generate the independent variables (also known as the covariates or inputs) as random uniform vectors, and the dependent variable (also known as the responses or outputs) as the noised product of the dependent variable and a fixed weight vector:

import numpy as np

# 1,000,000 inputs each of dimension 10,000 = 40GB of data
N, D = 1000000, 10000
X = np.random.rand(N, D).astype('float32')
# Ground-truth vector
W, b = np.random.rand(D).astype('float32'), np.random.rand()
# Response variables
Y = X @ W + np.random.randn(N)
# Save the dataset:
pkl.dump((X, W, b, Y), open('/tmp/linreg_data.pkl', 'wb'))

Basic code template

Our goal is to, given X and Y, recover the true parameter W. We will accomplish this via SGD on the squared-loss:

import torch as ch
from tqdm import tqdm
import time

train_loader = None # TODO!

# Calculate data mean and variance for normalization
def calculate_stats(loader, N):
    mean, stdev = 0., 0.
    for x_batch, _ in tqdm(loader):
        mean += x_batch.sum(0) / N
        stdev += x_batch.pow(2).sum(0) / N
    return mean, ch.sqrt(stdev - mean.pow(2))

start_time = time.time()
mean, stdev = calculate_stats(train_loader, N)
mean, stdev = mean.cuda(), stdev.cuda()
w_est, b_est = ch.zeros(D).cuda(), ch.zeros(1).cuda() # Initial guess for W
num_epochs = 10 # Number of full passes over the data to do

lr = 5e-2
for _ in range(num_epochs):
    total_loss, num_examples = 0., 0.
    start_time = time.time()
    for (x_batch, y_batch) in tqdm(train_loader):
        x_batch = x_batch.cuda()
        y_batch = y_batch.cuda()
        # Normalize the data for stability
        x_batch = (x_batch - mean) / stdev
        residual = x_batch @ w_est + b_est - y_batch
        # Gradients
        w_grad = x_batch.T @ residual / x_batch.shape[0]
        b_grad = ch.mean(residual, dim=0)
        # import ipdb; ipdb.set_trace()
        w_est = w_est - lr * w_grad
        b_est = b_est - lr * b_grad
        total_loss += residual.pow(2).sum()
        num_examples += x_batch.shape[0]
    print('Epoch time:', time.time() - start_time)
    print(f'Average loss: {total_loss / num_examples:.3f} | ',
        f'Norm diff', ch.norm(w_est / stdev - ch.tensor(W).cuda()).cpu().item())

print(f'Total script running time: {time.time() - start_time:.2f}s')

Note

Note that in general, using vanilla gradient descent to solve regularized linear regression is typically a very bad idea, and users are better served implementing an algorithm based on conjugate gradients or variance-reduced gradient methods. That said, the exact same principles here apply to any algorithm, so we use gradient descent to keep the code as clean as possible.

Naive approach: PyTorch TensorDataset

The only thing that remains unspecified in our implementation above is the train_loader. The standard way of making a loader here would be to use PyTorch’s built-in TensorDataset class, as follows:

from torch.utils.data import TensorDataset, DataLoader

X, W, b, Y = pkl.load(open('/tmp/linreg_data.pkl', 'rb'))
dataset = TensorDataset(ch.tensor(X), ch.tensor(Y))
train_loader = DataLoader(dataset, num_workers=8, shuffle=True)
# ... rest of code as above

The resulting code is runnable and correct. It will use 40GB of memory, since the entire tensor X will be kept in RAM. Running our script in an environment with a single A100 GPU and 8 CPU cores takes 16 seconds per epoch.

Speeding things up with FFCV

We’ll now try to improve on these results by replacing the standard PyTorch data loading pipeline with FFCV. The first step is to rewrite X and Y as a FFCV dataset (as detailed in the Writing a dataset to FFCV format guide):

from ffcv.fields import NDArrayField, FloatField

class LinearRegressionDataset:
    def __getitem__(self, idx):
        return (X[idx], np.array(Y[idx]).astype('float32'))

    def __len__(self):
        return len(X)

writer = DatasetWriter('/tmp/linreg_data.beton', {
    'covariate': NDArrayField(shape=(D,), dtype=np.dtype('float32')),
    'label': NDArrayField(shape=(1,), dtype=np.dtype('float32')),
}, num_workers=16)

writer.from_indexed_dataset(LinearRegressionDataset())

This allows us to replace the TensorDataset from the previous section with an FFCV data loader:

from ffcv.loader import Loader, OrderOption
from ffcv.fields.decoders import NDArrayDecoder
from ffcv.transforms import ToTensor, Squeeze, ToDevice

train_loader = Loader('/tmp/linreg_data.beton', batch_size=2048,
            num_workers=8, order=OrderOption.RANDOM,
            pipelines={
                'covariate': [NDArrayDecoder(), ToTensor(), ToDevice(ch.device('cuda:0'))],
                'label': [NDArrayDecoder(), ToTensor(), Squeeze(), ToDevice(ch.device('cuda:0'))]
            })

With just this simple substitution, our code goes from 16 seconds per epoch on an A100 GPU to 6 seconds.

As expected, GPU utilization also increases dramatically since data loading is no longer a bottleneck—this allows us to make optimizations elsewhere and make the code even faster!

More speed, less memory

We conclude this guide by suggesting a few ways to make our linear regression program even faster, and to reduce its memory footprint:

  • In our example above, FFCV caches the entire dataset in-memory: which means that, in the event of insufficient RAM, the program will not error our (unlike the TensorDataset example, which will raise a Segmentation Fault), but it will become significantly slower. An alternative discussed in the Tuning Guide that we didn’t explore here is to initialize the loader with os_cache=False and order=OrderOption.QUASI_RANDOM—this will disable caching of the full dataset (and thus can operate with very little memory!), and will read examples in an order which is nearly random but still minimizes underlying disk reads.

  • We can also optimize the main loop itself: for example, the gradient updates should be performed as in-place operations, as should the normalization. Since data loading is no longer the main bottleneck, such optimizations will result in improved performance.