PyTorch on Apple M1 MAX GPUs with SHARK – faster than TensorFlow-Metal | Hacker News Reader

PyTorch on Apple M1 MAX GPUs with SHARK – faster than TensorFlow-Metal

nod.ai

154 pointsyayr30 comments

I have questions!

What is the size of the matrix elements? My understanding is that float32 is super-easy to represent in SPIR-V, and that float16 is possible but the path is not as smoothly paved. Of course, on the Metal side "half" is fine.

Does the M1 silicon have anything like cooperative matrices[1]? This is a huge performance bump, but currently requires fiddly vendor-specific extensions. Hopefully a vendor-neutral standard will emerge before long. Conversely, are cooperative matrices used in the A100 comparison?

Are you using MoltenVK to run the dispatches? I've had ok results with this but ultimately built my own GPU API abstraction layer, so my code is calling into Metal. There are certain features, such as access to "memoryless" buffers, that may be difficult to access through MoltenVK, as I believe the concept doesn't really exist in Vulkan yet. Is that something that might be useful in a machine learning context?

Do you support operations like prefix sums? If so, are these using the decoupled look-back algorithm, or multiple dispatches?

I'm excited to see more portable use of compute shaders, I think to a large extent it's the future.

[1]: https://developer.nvidia.com/blog/machine-learning-accelerat...

Here are the matmul sizes for the MiniLM model used for inference: https://github.com/mmperf/mmperf/blob/main/benchmark_sizes/b...

These are the matmul sizes for the BERT training workload https://github.com/mmperf/mmperf/blob/main/benchmark_sizes/b...

Yes we use the latest MoltenVK (1.3.204.0) installed in the system.

I will let @noxa and other IREE devs chime in on the SPIR-V path but we do support prefix sums etc in the GPU path.

//part of nod.ai team.

Thanks for the matmul sizes, but the question I am more interested in is precision. Matrix multiply throughput can be dramatically higher for tensor cores[1] than normal shader ALU, specifically for reduced precision arithmetic. I'm wondering to what extent that's accessible on the M1, and to what extent IREE can address them.

Regarding prefix sum, the specific question I'm interested in is that SPIRV OpControlBarrier with device scope gets translated into threadgroup_barrier(mem_device) [2]. That's insufficient to make decoupled look-back work. Conversely, if you're not using decoupled look-back, you're not getting the full throughput on GPUs that do support that barrier. I'm wondering how your infrastructure deals with that.

[1]: https://developer.nvidia.com/blog/programming-tensor-cores-c...

[2]: https://github.com/linebender/piet-gpu/blob/d81e5cb4ee145abd...

So with Tensorcores you use TF32 which is more like FP19-ish and the marketing makes you think you get 8x the performance. But if you want actual FP32 precision you will need something like [1] but then your performance in the Tensorcore path is _only_ 2X faster than the SIMT path.

I'll leave the prefix sum for other devs who know more :D

https://github.com/NVIDIA/cutlass/blob/master/examples/27_am...

//part of nod.ai/shark team

I think we're talking past each other to some extent. Putting aside the question of how misleading it is to market a 16 bit multiply as a "TF32" operation, this is all about tradeoffs. The specific tradeoff that these tensor cores make is that in exchange for reduced precision (and a programming model which is even more of a pain than ordinary compute shaders, an astonishing achievement in and of itself), you get a lot more throughput. For certain AI workloads, particularly inference, that tradeoff is well worth it.

Reading between the lines a little, it sounds like your infrastructure is potentially able to exploit a good deal of the available throughput for FP32 workloads. That's great, and I'm happy to see it! However, for workloads that don't need that much precision, the tradeoff might be a lot less advantageous to M1. That may change again if and when Apple opens up lower-level APIs to their hardware, or reverse engineering delivers usable results.

tf32 and fp16 tensor cores are completely different, and tf32 is not 16 bit multiplication.

RE cooperative matrix: they have functional units (AMX/ANE) that could hopefully be exposed in MSL via something shaped like cooperative matrix and I'm pretty sure it'd be fantastic. Everything today is locked behind CoreML and Accelerate and those are poor targets for modern compiler-based approaches :( On the Vulkan side there's been rumblings of a vendor-agnostic extension for cooperative matrix and support from major vendors - at which point I'm hoping that leads to Apple wanting to show off their own HW features.

RE MoltenVK: we were surprised how robust it is nowadays - we're definitely going to build out our own Metal backend for our abstraction layer but wanted to see what we could hit with the zero-code option and it's proven very useful for that! A MoltenVK build is on the order of ~12MB (last I looked) while the entire IREE runtime is 50-100KB so it's a hard pill to swallow, all other issues (security/memory consumption/startup time/etc) notwithstanding :)

AFAIK the memoryless storage is only for textures and mostly useful in render passes - Vulkan has this via VK_MEMORY_PROPERTY_LAZILY_ALLOCATED_BIT. In a way what we do in our compiler is fuse dispatches such that they never hit memory at all so it hasn't been something we've needed yet. For heavy compute workloads more control over virtual memory would be useful, though, ala https://developer.nvidia.com/blog/introducing-low-level-gpu-... - some of the biggest performance hurdles when it comes to highly dynamic ML is around memory management (being able to zero-copy resize buffers when data-dependent shapes change would be killer).

// IREE dev

Great, thanks. That answers my questions. I'll read up on the lazily allocated bit; I wasn't aware that this provided similar functionality as dispatchThreadsPerTile[1], but perhaps it's something I'm misunderstanding. I'm excited about that as a way to stitch 2D graphics rendering operations together without having to hit main memory, but from your explanation I can see that functionality might not be very useful in AI workloads.

Amen on more control over dynamic memory access patterns. It's something I'm struggling with too, and I have a feeling that whatever solution I come up with is going to be a compromise.

Keep up the good work, these are exciting times!

[1]: https://developer.apple.com/documentation/metal/mtlrendercom...

Oh nice, I hadn't seen dispatchThreadsPerTile! I don't believe that's possible today in Vulkan :(

Our goal (though still WIP) is to have the interaction between user applications and our compiled code happen at the command buffer boundary - you would submit some work, pass in a VkSemaphore/MTLSharedEvent/cuEvent/futex/etc, we would use that when submitting our own work, and then we'd pass you back a VkSemaphore/etc you can continue chaining with. So one level of granularity coarser than mid-pass interleaving but still hopefully all pipelined properly with no host/device synchronization required. There will be programs that this doesn't work well with (heavily data-dependent stuff) but at least making it work turns it into an optimization problem vs today's representation problem!

The approach we're exploring for such integration is having the host pass a VkCommandBuffer/MTLCommandBuffer to the guest, which then records its dispatches onto that (the guest also having a VkDevice/MTLDevice handle for resource creation). This has a downside of requiring a single CPU thread at the integration point (and also not exploiting async compute on GPUs that legitimately run more than one compute queue in parallel), but on the other hand it avoids cross command buffer synchronization overhead (which I haven't measured, but if it's like GPU-to-CPU latency, it'd be very much worth avoiding).

Going direct to Metal, instead of through MoltenVK, would seem very appealing to me. Metal is pretty easy, and the idea of having an extremely light runtime that Just Works could be really nice for the user experience.

If you're planning to open-source the runtime, I'd be happy to chat about my experience building piet-gpu-hal. I haven't actually tried running any machine learning workloads on it, but I think it could be adapted for that without too much pain.

Exactly the kind of things we've been talking about! A fun and challenging tradeoff space and it's always great to connect with others!

Ahh linebender - I hadn't connected the name with your github account - piet-gpu is great, as is your blog! Also, for anyone skimming the comments this talk is fantastic and I share it with anyone new to the GPGPU space: https://www.youtube.com/watch?v=DZRn_jNZjbw

We waffled a bit with the API granularity in the beginning and it's taken building out most of the rest of the project in order to nail it down (the big refactor still pending). The biggest issue is that in simple models we'll end up emitting a single command buffer but anything with control flow (that we can't predicate), data dependencies (sparsity, thresholding, etc), or CPU work in the middle (IO, custom user code, etc) can break that up. We also hit cases where we need to flush work - such as if we run out of usable memory and need to defragment or resize our pools. We want to be able to (but aren't yet) reuse command buffers (CUDA graphs, etc) and that requires being able to both cache them and recreate them on demand (if we resize a pool we have to invalidate all cached command buffers using those resources, as update-after-bind is not universally available and if shapes change there's big ripples). Since most models beyond simple vision ones are ~thousands of dispatches it also lets us better integrate into multithreaded applications like you mention as apps can record commands for themselves in parallel without synchronization. It still would be nice to have certain operations inlined, though, and for that we want to allow custom hooks that we call into to add commands to the command buffers, turning things inside-out to make small amounts of work like image transformations in-between model layers possible (I'm really hoping we can avoid modeling the entire graphics pipeline in the compiler and this would be a way around that :). We haven't yet started on scheduling across queues but that's also very interesting especially in multi-GPU cases (with x4/x8 GPUs being common in datacenters, or NUMA CPU clusters that can be scheduled similarly).

We're fully open source (https://github.com/google/iree) but have been operating quietly while we get the groundwork in place - it's taken some time but now we're finally starting to stumble into success on certain problem categories (like transformers as in the post). Right now it's mostly just organized as a systems/compiler nerd honeypot for people looking for an ML/number crunching framework that (purposefully) doesn't look like any of the existing ones :)

Would love to chat more - even if just to commiserate over GPU APIs and such - everyone is welcome on the discord where a bunch of us nerds have gathered or we could grab virtual coffee (realized just now that this hn acct is ancient - I'm benvanik@google.com :)

Has anyone else actually installed and demonstrated this? (might be the perfect excuse to buy an M1)

"All the code the recreate the tests are here and here. However you will need to install PyTorch torchvision from source since torchvision doesn’t have support for M1 yet. You will also need to build SHARK from the apple-m1-max-support branch from the SHARK repository."

So I bet you won't get it running within an hour, possibly 24 hours.

I'm testing FFTs on GPU. I got cupy running in jupyter today after a day of pain. The error said: "One method of fixing this is to repeatedly uninstall numpy until none is found" -- I had to manually delete directories and reboot to get cupy working -- now I have very fast fourier tranforms, but the internet says the PyTorch implementation is even better on GPU.... so I have to try it...

I created a new venv and tried to install pytorch with GPU support but torch.cuda.is_available() returns False.

Why is package management such a disaster?

This is not part of regular pytorch install.

If you can build torch-mlir and SHARK from src you can use it. So hopefully soon we can make pip installable packages but for now the interfaces are in constant development so you will have to build from source.

//part of nod.ai / SHARK team.

Thanks for your comment, I'm betting that the M1 example demoed here should be easy to build, and I really do want the M1 for my work. I'm just lamenting how annoying it is to install the packaged versions of PyTorch for GPU since that's what I'm fighting with at this very moment!

hear your pain and we really want to make it easy (after we make it work).

//part of nod.ai / SHARK team.

there is also some hint about performance/watt vs an NVIDIA A100 using BERT models:

> The A100 draws 131W peak during the training run, the M1MAX GPU draws a maximum of 15.4 W. The A100 runs an iteration at 7ms vs 104ms (with SHARK) and 196ms (with TF-Metal)

So the A100 is 75% more efficient despite being in a high power variation of a larger node. Seems Apple still haw a lot of work to do in the GPU/ML side of things.

My reading implied that they weren't able to use the apple hardware acceleration for AI, which given the massive perf delta I suspect the A100 is using. If they weren't using the A100's ML hardware I'd be both amazed by Nvidia's engineers, and also unhappy as it would be unfair for the test to not use the features of the hardware (this is a common refrain in some benchmarking circles where you disable features to get "fairness" in comparisons).

Honestly though I'm more concerned about the other aspects of apple's software: "The OSX Window Server crashes when all GPUs are used to the maximum"

Correct! We're currently compiling to SPIR-V and then going through MoltenVK/spirv-cross to get to Metal - and our SPIR-V needs quite a bit of work (currently tuned for Mali, which is a very different architecture). The biggest limiter though is as you note Apple doesn't have a way to use the fancy hardware instructions whereas our CUDA version is using the Tensor Cores.

The good news (for us) is that this is all effectively -O1 today; there's still potential for 2-8x speedups over these numbers and that's not even factoring in the inaccessible HW features that maybe one day Apple will expose :crossed-fingers: :)

// IREE dev

Any chance of using the Neural Engine (otherwise exposed by CoreML) ?

unlikely since the interface from ANE is not public and it may change between hardware versions.

yeah, you could also say the performance is only 57% in this particular test for the M1 vs A100. Still not bad for a laptop though with a first gen chip.

The M1 is (supposedly) specialized for ML calculations and on a decidedly denser node, though. I guess it's a testament to either CUDA's power or machine learning GPGPU optimization.

I like this blog to understand NVidia perf stats (search for "Normalized GPU deep learning performance"):

https://timdettmers.com/2020/09/07/which-gpu-for-deep-learni...

It seems that their consumer GPUs also have a gap

> The M1 is (supposedly) specialized for ML calculations

It is, but they’re not using any of that. Only the M1 GPU at this time.

Yeah the ANE and AMX on cpu are wrapped behind Accelerate Framework and CoreML. So you will have to use CoreML (which wasn't able to compile the latest TF BERT). ANE is also inference only. So if you want training you will have to use the GPU with Apple's Tensorflow-Metal.

//part of nod.ai / SHARK team.

This is a big deal. Congrats to the team who pulled this off.

Thanks to: LLVM/MLIR --> For the awesome compiler infrastructure

IREE --> For the awesome backend to MLIR

SHARK/nod.ai --> For adapting IREE for use on various hardware and fine tuning for target hardware.

//part of nod.ai / SHARK team

How do your benchmarks compare against CoreML?

I updated the blog with the reference. Basically it crashes to compile the model with https://github.com/NodLabs/shark-samples/blob/main/examples/.... The coremltools converter is very version specific (like all vendor conversion kits) and still on a version of TF I couldn't get on conda. Also it doesn't allow for training and only FP16 for inference with ANE. All our tests were with FP32.

//part of nod.ai/shark team.

Reply on news.ycombinator.com