Performance of Right-Facing vs Left-Facing matmuls

Neel_Nanda: <p>I’m writing an implementation of a transformer to pre-train from scratch, and wrote my matrices to be left multiplying (eg, in a MLP layer with 4000 neurons and a 1000 dimensional residual stream, <code>W_in.shape==[4000, 1000]</code> and <code>neuron_pre_act = einsum("nm,bm-&gt;bn", W_in, residual_stream)</code>). I notice that in most implementations, eg <code>nn.Linear</code>, matrices are right facing, and this lets you use notation like <code>neuron_pre_act = residual_stream @ W_in</code>.</p>
<p>I’m curious if anyone knows why this convention is there, and how much it matters for eg performance? Based on some preliminary tests, right-facing is a performance improvement, but seems to vary a ton in ways I don’t fully understand, eg in some contexts it’s a minor 0.5% speed improvement, in others it’s a major 5-10%?!</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/neel_nanda">@Neel_Nanda</a>,</p>
<p>Can you share a minimal reproducible example?</p>
<ol>
<li>
<p>When you compare <code>torch.einsum</code> with <code>nn.Linear</code> make sure you have <code>nn.Linear(bias=False)</code> otherwise the operations aren’t equivalent.</p>
</li>
<li>
<p>When you are measuring times for code snippets, make sure that you synchronize torch (<code>torch.cuda.synchronize</code>) before calling <code>time.time()</code> otherwise you won’t record the <em>whole</em> runtime of an operation but only its call (especially if ran on the GPU). There’s more info about it <a href="https://discuss.pytorch.org/t/how-does-torch-cuda-synchronize-behave/147049">here</a>.</p>
</li>
</ol>

Neel_Nanda: <p>Minimal counter-example</p>
<pre><code class="lang-auto">batch = 3 * 10**4
residual = 3 * 10**4
d_mlp = 3 * 10**4

input_mat = torch.randn(batch, residual, device='cuda')
right_mat = torch.randn(residual, d_mlp, device='cuda')
left_mat = torch.randn(d_mlp, residual, device='cuda')

print("Cuda warmup")
_ = (input_mat @ right_mat)
torch.cuda.synchronize()

start_time = time.time()
_ = (input_mat @ right_mat)
torch.cuda.synchronize()
print("Time taken right mul: ", time.time() - start_time)

start_time = time.time()
_ = (input_mat @ left_mat.T)
torch.cuda.synchronize()
print("Time taken left mul transpose: ", time.time() - start_time)

start_time = time.time()
_ = torch.einsum("rm,br-&gt;bm", right_mat, input_mat)
torch.cuda.synchronize()
print("Time taken right mul einsum: ", time.time() - start_time)

start_time = time.time()
_ = torch.einsum("mr,br-&gt;bm", left_mat, input_mat)
torch.cuda.synchronize()
print("Time taken left mul einsum: ", time.time() - start_time)
</code></pre>
<p>On my A100 80GB GPU this returns:</p>
<pre><code class="lang-auto">Cuda warmup
Time taken right mul:  3.2466354370117188
Time taken left mul transpose:  3.315218448638916
Time taken right mul einsum:  3.061365842819214
Time taken left mul einsum:  3.315138339996338
</code></pre>

Neel_Nanda: <p>Some more experiments:</p>
<p>In float16:</p>
<pre><code class="lang-auto">Cuda warmup
Time taken right mul:  0.23980021476745605
Time taken left mul transpose:  0.2591671943664551
Time taken right mul einsum:  0.2865767478942871
Time taken left mul einsum:  0.27009034156799316
</code></pre>
<p>In bfloat16</p>
<pre><code class="lang-auto">Cuda warmup
Time taken right mul:  0.23795533180236816
Time taken left mul transpose:  0.25300073623657227
Time taken right mul einsum:  0.2702064514160156
Time taken left mul einsum:  0.2489306926727295
</code></pre>

Neel_Nanda: <p>With 10**4 rather than 3*10**4</p>
<pre><code class="lang-auto">Cuda warmup
Time taken right mul:  0.11016678810119629
Time taken left mul transpose:  0.11061906814575195
Time taken right mul einsum:  0.11274075508117676
Time taken left mul einsum:  0.11332011222839355
</code></pre>

KFrank: <p>Hi Neel!</p>
<p>The short story:  I speculate that pytorch follows a <code>Linear.weight</code> “transpose”<br>
convention for improved performance on “low-end” hardware.</p>
<aside class="quote no-group" data-username="Neel_Nanda" data-post="1" data-topic="162582" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neel_nanda/48/53641_2.png" class="avatar"> Neel_Nanda:</div>
<blockquote>
<p>I notice that in most implementations, eg <code>nn.Linear</code>, matrices are right facing</p>
</blockquote>
</aside>
<p>I’m not sure what you mean by “right facing,” but I comment below on what<br>
I believe are the likely roots of pytorch’s choice of <code>Linear</code> implementation.</p>
<aside class="quote no-group">
<blockquote>
<p>this lets you use notation like <code>neuron_pre_act = residual_stream @ W_in</code>.</p>
</blockquote>
</aside>
<p>As an aside, I  believe this should be <code>residual_stream @ W_in.T</code>.  Based<br>
on your shapes (and your <code>einsum()</code> expression), your version (without the<br>
transpose) will fail because the second dimension of <code>residual_stream</code><br>
doesn’t match the first dimension of <code>W_in</code>.</p>
<p>As for the <code>weight</code> of <code>Linear</code>, I believe that pytorch chooses to have the<br>
<code>weight</code> multiply the <code>input</code> to <code>Linear</code> from the right so that <code>Linear</code> can<br>
accept <code>input</code>s with any number of leading “batch” dimensions, e.g., no<br>
batch dimension, a single batch dimension, or a batch and a channel<br>
dimension.</p>
<p>Thus shapes <code>[10]</code>, <code>[5, 10]</code>, and <code>[5, 3, 10]</code> would all be valid inputs to<br>
<code>Linear (in_features = 10, out_features = 2)</code>.</p>
<p>However, <code>torch.nn.Linear (in_features = 10, out_features = 2)</code><br>
has a <code>weight</code> of shape <code>[2, 10]</code> and you can see by inspection that <code>Linear</code><br>
multiplies its <code>input</code> from the right by the transpose  of its <code>weight</code>, that is, e.g.,<br>
<code>torch.randn (5, 10) @ torch.nn.Linear (10, 2).weight.T</code>.</p>
<p>Why is the matrix that is stored as <code>weight</code> the transpose of the of the matrix<br>
that is used in the actual matrix-matrix multiplication (or, more generally,<br>
the tensor-matrix multiplication)?</p>
<p>Although there have been historical exceptions, <em>row-major</em> is the de facto<br>
standard for matrix storage*.  That is to say that as you move sequentially<br>
though the (linear) address space in which a matrix is stored, the column<br>
index varies the most rapidly.  Equivalently, individual rows are stored<br>
sequentially (and then chained together one after the other).</p>
<p>It is therefore true that performing matrix-matrix-transpose multiplication<br>
has greater <em>memory locality</em> than does matrix-matrix multiplication because<br>
as you perform the row-column vector-vector dot products of which matrix<br>
multiplication is composed, moving sequentially through a column of the<br>
transposed matrix translates to moving sequentially through the corresponding<br>
row of the untransposed matrix, and therefore the data in transposed matrix’s<br>
column is accessed sequentially in memory.</p>
<p>Whether this memory locality affects performance is highly dependent on<br>
the details of the hardware – size and page size of cache memory, relative<br>
cost of cache hits vs. misses, structure and number of floating-point pipelines,<br>
cpu vs. gpu, and so on.  On “small” systems where a row of a “large” matrix<br>
might fit into cache memory, but the whole matrix doesn’t, matrix-matrix<br>
multiplication could be dramatically slower than matrix-matrix-transpose<br>
multiplication, and significant algorithmic effort was made in deciding whether<br>
to store a matrix or its transpose (or even sometimes both).</p>
<p>With modern gpus, this is probably much less likely to matter.  I’ll also note<br>
that gpus (as well as most modern cpu pipelines) are optimized to be able<br>
to stride efficiently through memory (that is, with strides greater than one,<br>
rather than purely sequentially), which is what you are doing when you access<br>
the column of a matrix (that is stored in row-major form).</p>
<p>Having said all of that, here are some real-life timings for matrix-matrix vs.<br>
matrix-matrix-transpose multiplication on one particular system.</p>
<p>Here is the test script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
print (torch.version.cuda)
print (torch.cuda.get_device_name())

import math
import statistics
import time

_ = torch.manual_seed (2022)

device = 'cpu'
print ('device =', device)

nWarm = 100
nStat = 10
nTime = 100
print ('nStat =', nStat, ', nTime =', nTime)

b = 100
m = 10000
n = 10000
print ('b =', b, ', m =', m, ', n =', n)

a = torch.randn (b, m, device = device)
mr = torch.randn (m, n, device = device)
ml = torch.randn (n, m, device = device)

# warm up
for i in range (nWarm):
    p = a @ mr
    p = a @ ml.T

mr_timing = []
mlT_timing = []

for  i in range (nStat):
    torch.cuda.synchronize()
    tStart = time.time()
    for  i in range (nTime):
        p = a @ mr
    torch.cuda.synchronize()
    tEnd = time.time()
    mr_timing.append (tEnd - tStart)
    
    torch.cuda.synchronize()
    tStart = time.time()
    for  i in range (nTime):
        p = a @ ml.T
    torch.cuda.synchronize()
    tEnd = time.time()
    mlT_timing.append (tEnd - tStart)

mr_mean = statistics.mean (mr_timing)
mr_stderr = statistics.stdev (mr_timing) / math.sqrt (nStat - 1)
mlT_mean = statistics.mean (mlT_timing)
mlT_stderr = statistics.stdev (mlT_timing) / math.sqrt (nStat - 1)

print ('a @ mr  timings:', mr_mean, '+-', mr_stderr)
print ('a @ ml.T timing:', mlT_mean, '+-', mlT_stderr)
</code></pre>
<p>Here is its cpu-version output:</p>
<pre><code class="lang-plaintext">1.12.0
11.6
GeForce GTX 1050 Ti
device = cpu
nStat = 10 , nTime = 100
b = 100 , m = 10000 , n = 10000
a @ mr  timings: 12.70658016204834 +- 0.007860937265143848
a @ ml.T timing: 12.256337642669678 +- 0.007302343186920837
</code></pre>
<p>And here is its companion gpu-version output:</p>
<pre><code class="lang-plaintext">1.12.0
11.6
GeForce GTX 1050 Ti
device = cuda
nStat = 10 , nTime = 100
b = 100 , m = 10000 , n = 10000
a @ mr  timings: 1.354783058166504 +- 0.0005118169535293032
a @ ml.T timing: 1.3442912817001342 +- 0.0014081416198141895
</code></pre>
<p>For this particular system, the performance differences are quite small, and<br>
not necessarily statistically significant.  (If significant, they could perhaps<br>
be meaningful at the margins – maybe a percent or so, but certainly not<br>
something like a factor of two).</p>
<p>It’s interesting to note – but I wouldn’t read too much into it – that<br>
matrix-matrix-transpose multiplication outperforms matrix-matrix multiplication<br>
on both the cpu and gpu by a little bit.  (Of course, not surprisingly, the gpu<br>
outperforms the cpu by a factor of ten – which one would hope because this<br>
is precisely the kind of thing that should be fast on a gpu.)</p>
<p>Again, any such differences will depend on the hardware, and will likely also<br>
depend on the size and shape of the tensors involved.</p>
<p>So, why does pytorch choose to store its multiplying-from-the-right <code>weight</code><br>
matrix as its transpose?</p>
<p>It could just be historical habit from the time when the matrix-transpose<br>
optimization almost always mattered for large matrices.</p>
<p>But it could be that on low-end hardware (think inference on an embedded<br>
system) the matrix-transpose optimization really does pay off, while on<br>
“fancy” hardware it doesn’t really matter, one way or the other.</p>
<p>*)  Note, it is possible to have a <code>Linear</code> store its <code>weight</code> in column-major<br>
form by setting its <code>weight</code> to the transpose view of a row-major matrix:</p>
<pre data-code-wrap="python"><code class="lang-python">lin = torch.nn.Linear (10, 2, bias = False)
lin.weight.shape
wt_column_major = torch.randn (10, 2)
wt_column_major.shape
lin.weight = torch.nn.Parameter (wt_column_major.T)
lin.weight.shape
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Neel_Nanda: <p>Thanks so much, that’s a really helpful answer! And interesting that it mostly doesn’t matter on larger/more modern systems.<br>
I also was completely wrong about nn.Linear and hadn’t noticed the tranpose, so thanks for correcting me on that.</p>