Slower processing using Metal on PyTorch versus CPU for specific calculations

basavyr: <p>I am performing some training procedure for a classification task, and I am running this on a M4 Pro chip. Before the actual training stage, I need to apply some transformations to some embeddings.</p>
<p>During my benchmarks, I have noticed some strange (or rather <em>unexpected</em>) behavior: <strong>the processing of those embeddings happens faster on the <code>CPU</code> rather than <code>MPS</code>.</strong></p>
<p>The processing of embeddings is as follows (omitting some irrelevant stuff):</p>
<pre data-code-wrap="python"><code class="lang-python">def algorithm(embeddings: torch.Tensor, labels: torch.Tensor):
    ...
    for label in torch.unique(labels):
        idx = (labels == label).nonzero(as_tuple=True)[0]
        emb = embeddings[idx]
        dot = torch.mm(emb, emb.t())
        sq = emb.pow(2).sum(1, keepdim=True)
        dists = sq + sq.t() - 2 * dot
        dists = torch.sqrt(torch.clamp(dists, min=0))
        dists_min = dists.clone()
        dists_min.fill_diagonal_(float('inf'))
        min_dist = dists_min.min().item()
        label_radii[int(label.item())] = min_dist
    ...
</code></pre>
<p>It is the <code>algorithm()</code> execution where I’ve noticed decreased performance when both tensors are loaded onto <code>device=torch.device(“mps”)</code>. The differences are not dramatic, however, they are noticeable:</p>
<pre data-code-wrap="bash"><code class="lang-bash">Process duration on [cpu] : 2.9765 s
</code></pre>
<pre data-code-wrap="bash"><code class="lang-bash">Process duration on [mps] : 3.4318 s
</code></pre>
<p>(the numbers are consistent across multiple runs)</p>
<p><em><strong>Q: Any ideas why the CPU is faster here?</strong></em></p>
<p>What I can do next is to actually measure each individual tensor operation from <code>algorithm()</code> and see which might be the “culprit”… However, at a first glance, it should be expected that since everything is preloaded into GPU/CPU memory beforehand, the GPU should be faster.</p>

basavyr: <p>I did just perform some <a href="https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html" rel="noopener nofollow ugc">profiling</a> and saw that when I execute <code>algorithm()</code> on <code>MPS</code>, we see a lot of CPU time spent performing <code>aten::_local_scalar_dense</code> and <code>aten::item</code> (roughly 97%).</p>
<p>On the other hand, when trying on the <code>CPU</code>, it is indeed the expected matrix multiplication: <code>aten::mm</code> (77.18%). So I guess the <code>item()</code> is the “culprit”.</p>

ptrblck: <aside class="quote no-group" data-username="basavyr" data-post="2" data-topic="222548">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/basavyr/48/73507_2.png" class="avatar"> basavyr:</div>
<blockquote>
<p>we see a lot of CPU time spent performing <code>aten::_local_scalar_dense</code> and <code>aten::item</code> (roughly 97%).</p>
</blockquote>
</aside>
<p>Assuming MPS executes their workloads asynchronously as GPUs using CUDA, this output is misleading since <code>item</code> will synchronize the code w.r.t. the host and thus accumulate the runtime of the GPU kernels.</p>
<p>You would need to use a proper profiler which is able to report the actual device execution time.</p>