Skip to content

2.3x speed-up by switching to pytorch torch.linalg.qr #593

Description

@RoelantStegmann

Context

We noticed pygam's PIRLS loop spends the large majority of its wall-clock time in a single
numpy.linalg.qr(WB.toarray()) call (pygam.py:783). Profiling a real workload — LogisticGAM
with a monotonic-constrained TensorTerm, fit under thread-pinned BLAS (OMP_NUM_THREADS=1 to have
one process per core in a multi-process worker pool) — showed this call dominating fit time.

What we tried

We monkeypatched this call site to route through torch.linalg.qr instead of
numpy.linalg.qr (converting the dense array to a CPU tensor and back), keeping everything
else about the fit identical. On real data (10 representative fits, various data sizes),
this consistently cut per-fit wall time by roughly 2.3x (mean 8.72s -> 3.74s across the
sample), with the effect holding across every fit we tested, not just on average.

We verified correctness by comparing predicted values between the numpy-QR and torch-QR runs:
max absolute difference across ~1000 column comparisons was ~1.3e-13.

Why this might matter for other users

This seems like it could generalize beyond our specific use case: any model with enough
observations/spline terms that WB.toarray() produces a reasonably large dense matrix would hit
the same QR-dominated cost, and any user running many concurrent fits under thread-pinned BLAS
(a common setup to avoid oversubscription in multiprocessing pools) would likely see a similar
win, since torch's linear algebra backend appears to handle this workload more efficiently in
that configuration.

Proposal

pygam already has precedent for exactly this kind of optional, auto-detected accelerated
backend -- pygam/utils.py's cholesky() function checks if SKSPIMPORT and transparently
uses sksparse.cholmod when available, falling back to scipy.linalg.cholesky otherwise. A
similar pattern could apply here: if torch is importable, use torch.linalg.qr for the PIRLS
QR step; otherwise fall back to the current numpy.linalg.qr path. This would keep torch as a
fully optional dependency (no behavior change for anyone who doesn't have it installed) while
giving users who do have it a meaningful, opt-in speedup.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions