Initial perf results from compiling NestedTensor Python Subclass via desugaring

– with @ezyang, @bdhirsh, @jbschlosser, @cpuhrsch

Recently NestedTensor became compilable (stay tuned for a different post about this and our plans here), and just to get a sense of what the numbers should look like, I did some measuring:


NestedTensor allows one to handle raggedness in data efficiently in native PyTorch. Historically, NestedTensor (and most subclasses in general) are written in C++ to ensure fast execution. More recently however, writing tensor subclasses in Python has become more appealing as PT2 stack gains more capabilities to potentially remove that overhead.

In particular, the approach featured here today is possible thanks to @bdhirsh’s work on subclass compilation via desugaring, e.g. converting the subclass into its underlying dense tensors before it’s processed by the compiler backend. (Take note that this is not the only way we plan to support compilation for NestedTensor, see future work).

In this post we compare the performance of (1) the original NestedTensor in C++, (2) a new NestedTensor written as a Python subclass, (3) and the new python subclass torch.compile via desugaring. Results were obtained on PyTorch main branch with a single A100 GPU. We use variations of dummy program that consists of sin, cos, and linear ops to demonstrate the effect of varying input and model size (number of ops) and operator fusion, etc on the runtime.


Overall, the results are directionally as expected, with the degree of eager python overhead being the biggest surprise.

  • Varying input size: Compiled Python NT outperforms eager cpp across all input sizes. In the overhead-bound regime this is due to reduced overhead. In the compute-bound regime this is due to operator fusion. Non-compiled Python subclass matches eager cpp on large input sizes, but underperforms on smaller inputs due to Python subclass overhead.
  • Varying program size (num operators): Speedups from compilation over eager cpp are more pronounced for larger programs (with more ops) as guard execution overhead etc is amortized away (as long as the larger program doesn’t induce proportionally more guards).

Future work:

  • Overhead in the eager mode should be investigated.
  • We plan to preserve the nested structure all the way to inductor so that we can rely on a generic codegen rather custom kernels. (@jbschlosser is working on this)
  • There is also interest to expose codegened inductor kernels into eager to ensure fast eager performance.


How does the size of the input tensor affect things?

Example 1: A moderate sized program (~60 ops in forward) with fusion opportunities

Compiled NestedTensor subclass is measured to be both faster than eager cpp at small input sizes due to reduction in overhead, and faster than cpp at large input sizes due operator fusion speeding up
bandwidth-bound operations. Python subclass w/o compile matches cpp at large input sizes since both are compute bound and launch the same kernels, but significantly slower at small inputs sizes due to overhead.*

Input size cpp w/o compile python w/o compile python w/ compile
[[20, 30, 40], 2^14] 8.812ms** 37.774ms** 4.27ms**
[[20, 30, 40], 2^16] 11.955ms 38.983ms** 10.559ms
[[20, 30, 40], 2^18] 43.27ms 43.861ms 37.73ms
[[20, 30, 40], 2^20] 168.7ms 167.665ms 145.7ms

Table 1: ** indicates overhead bound

Example 2: Slightly smaller sized program (~40 ops), and with fewer fusion opportunities

When there are fewer fusion opportunities the gap between compile and eager cpp narrows at large input sizes.

Compiled NestedTensor subclass is measured to be faster than the C++ NT in the overhead-bound regime, though the gap has narrowed compared to example 1 because we have fewer ops in forward/backward and hence less overhead to remove; see row 1 of Table 2 below. On the same row, we see that without torch.compile, Python NT subclass is much slower than C++. All 3 columns have identical runtime in the compute-bound regime as expected since their implementations do not differ significantly (or at all?) in terms of the underlying kernels called.

Input size cpp w/o compile python w/o compile python w/ compile
[[20, 30, 40], 2^14] 5.93ms** 28.9ms** 4.27ms**
[[20, 30, 40], 2^16] 10.5ms 29.6ms** 10.5ms
[[20, 30, 40], 2^18] 37.8ms 38.3ms 37.4ms
[[20, 30, 40], 2^20] 148.8ms 147.6ms 145.2ms

Table 2: ** indicates overhead bound

How does the size of the program (number of ops) affect things?

Result: Smaller programs means that the overhead of guard execution is more visible, so larger programs achieve higher multiples in speedup.

The program used for measurement in this section is the same program used in example 1 of the previous section, except now we vary the number of times we execute unit, and test for both overhead-bound and compute-bound regimes. Each execution of unit corresponds to 6 aten ops.

Example 1: Overhead-bound (smaller tensor): [[20, 30, 40], 2^14]

range(N) cpp w/o compile python w/o compile python w/ compile
2 2.5ms 7.9ms 2.0ms
5 4.7ms 18.34ms 2.9ms
10 8.5ms 38ms 4.5ms
20 18ms 73ms 7ms

Table 3

Compute-bound (smaller tensor): [[20, 30, 40], 2^14]

Python w/o compile matches cpp w/o compile at all graph sizes as long as input size is large. Compile is faster due to operator fusion. (See appendix for some profiler traces showing this.)

range(N) cpp w/o compile python w/o compile python w/ compile
2 34ms 34ms 30ms
5 84ms 85ms 75ms
10 169ms 170ms 148ms
20 334ms 334ms 329ms

Table 4

Future work

Reduce eager mode overhead

As pointed out above in the results, Python subclasses seem to have a significant amount of overhead. Eager performance of tensor subclasses can be important as well for users who cannot compile their model for various reasons.

aten::sin call for Python NT without compile

aten::sin call for C++ NT without compile

Compile support without desugaring Inductor codegen

Currently, with the current torch compile desugaring solution, we still must rely on custom kernels to implement operations that need to reduce or broadcast over the ragged dimension. These custom kernels can take significant effort to write. Preserving the nested tensor structure all the way to inductor and using inductor codegen would a way to cleanly support these operations without the need for custom kernels.

Eager mode performance for inductor generated kernels

Following up from the previous item, if you rely on inductor generated kernels during compile, but still rely on a unbind-based fallback during eager, you would observe poor performance for eager subclasses (C++ or Python) even in a compute-bound regime. Hence, it may be necessary to have some way of running those inductor generated kernels in eager.


Code to reproduce

Profiler traces: Operator fusion speed up

Unfused backward for sin/cos

Fused backward for sin/cos

Profiler traces: compute bound vs overhead bound execution

Compute-bound ([20, 30, 40], 2^18])

Python NT no compile (38.3ms)

C++ NT no compile (37.8ms)

Python NT compile (37.4ms)

Overhead-bound ([[20, 30, 40], 2^14] case)

Python NT no compile (29.0ms)

C++ NT no compile (5.93ms)

C++ NT no compile (4.27ms)