A simple problem for which Julia is a better choice than Python | Hacker News Reader

A simple problem for which Julia is a better choice than Python

3 pointscs7024 comments

If we try this with PyTorch:

  import torch
  x = torch.randn(4000, 4000, device='cuda:0')
  y = torch.randn(4000, 4000, device='cuda:0')
  z = torch.einsum('ij,jk->ijk', x, y).relu().sum(dim=1)

we will likely run out of memory, because the tensor product of x and y consists of 4000 x 4000 x 4000 = 64 billion floating-point values, and we must evaluate the relu() function on each of them before summing them over index j (dim 1).

And we can't solve this problem with a nice for-loop that iterates over 64 billion elements in Python -- it would be waaaaay too slow. Python and PyTorch force us to drop down to a fast, low-level, compiled language.

The same is true for TensorFlow and other similar "vectorized" frameworks. Our choices are to store 64 billion floating-point values simultaneously in memory, or drop down to a low-level language. No fun.

But solving this in Julia is trivial: We just write the loops.

  import torch
  x = torch.randn(4000, 4000, device='cuda:0')
  y = torch.randn(4000, 4000, device='cuda:0')

  import time
  torch.cuda.synchronize()
  t0 = time.perf_counter()
  z = torch.zeros(4000, 4000, device='cuda:0')
  for j in range(4000):
      z += (x[:, j] * z[j, :]).relu()
  torch.cuda.synchronize()
  t1 = time.perf_counter()
  print(t1-t0)

seem like 640ms-ish for me (on a rtx3090). How much faster is Julia, then? Is it double or single digit ms?

Of course, the backward will be more work because you need to recompute which bits were 0.

I can appreciate that Python is slow and that Julia has very nice properties dissolving the boundaries of operators that are awefully hard in PyTorch & Co, but this is one of these examples where the story seems rather flawed (you don't need the for loop over 64 billion elements, just over the 4000 j-coordinates, and the Python overhead here isn't too bad, either). As far as I can tell, the large overhead here is from materializing the intermediate results in global memory. Personally, I would hope that the fusers might be able to fuse in the not to far future, so the entire thing can run on the GPU in one kernel, but until then, if Julia can do that, it'll have a huge advantage here.

Shameless plug: In my courses about efficient PyTorch, I always give the rule of thumb "avoid operations on individual elements, but things operating on hundreds+ elements should be fine". And the reason isn't as much just "Python is slow" but the overhead of the tensor objects.

>I always give the rule of thumb "avoid operations on individual elements, but things operating on hundreds+ elements should be fine". And the reason isn't as much just "Python is slow" but the overhead of the tensor objects.

That's just not possible in a lot of cases, like when defining a large nonlinear ODE which is precisely described by the actions of individual elements. Indeed, it's not the ML case, but it is a case where PyTorch's assumptions are too strong.

Show the code!

This is more a complaint about pytorch library than python language itself

Reply on news.ycombinator.com