what matrix multiplication teaches us
Matrix multiplication, and why it’s so onerous
Matrices are vital to machine studying programs as a easy technique to symbolize knowledge effectively — akin to enter knowledge (e.g., units of pixels in a picture) or the inner workings between the mannequin layers. Because of this, multiplying these matrices collectively makes up a big portion of the overall computations in deep studying fashions. The truth is, matmuls make up roughly 45-60% of the overall runtime of many common transformer fashions like BERT, CLIP, and even ChatGPT, is a well-liked method to computing the convolution operation that kinds the inspiration of most pc imaginative and prescient fashions, and makes up the spine of many high-performance computing (”HPC”) purposes.
Given its significance, there was in depth analysis on writing environment friendly matmul algorithms. Papers from the 60s, 70s, 80s, 90s, 2000s, to the current day exist that attempt to clear up the issue with the {hardware} of that period.
However the conceptual matmul algorithm isn’t what makes it troublesome. As an alternative, the problem is writing a matmul quick sufficient to attain state-of-the-art efficiency throughout all of the {hardware}, mannequin, and knowledge variety within the AI business. It’s even more durable to make it compose with all different AI operators.
{Hardware}
Every system used to run AI fashions has totally different traits, together with totally different memory hierarchies and totally different multiply and accumulate units (MAC).
For instance, CPUs make use of a hierarchy of reminiscence from gradual RAM to more and more quicker caches — Degree-3, Degree-2, Degree-1, and CPU registers. The dimensions of the reminiscence is inversely proportional to its pace — for instance, L1 cache entry is usually on the order of single nanoseconds, whereas RAM entry is on the order of 100 nanoseconds. To get the best efficiency matmul, the algorithm itself must be carried out to effectively work with the totally different reminiscence ranges and sizes. Uncooked matrices are too large to suit into the registers or the quickest reminiscence cache at one time, so the problem is figuring out decompose them into the suitable sized blocks or “tiles” that maximize utilization of the quickest reminiscence.
Furthermore, the precise shapes of the models that deal with the core matrix performance differ throughout {hardware}. CPUs have historically been scalar machines, that means they course of instructions one step at a time, however all CPU distributors have added vector models (SIMD) during the last 20 years. GPUs execute single operations across multiple threads (SIMT) to maximise the effectivity of extremely parallel, repetitive operations like matmuls. And extra specialised {hardware} takes this additional by working on 2D matrices. Google TPUs are probably the most well-known, however Apple and Intel have added their very own matrix multiplication options referred to as AMX. However whereas extra superior MAC models have led to improved efficiency, they’ve additionally created a better want for versatile algorithms that work throughout scalar, vector, and matrix processors.
Fashions
AI fashions are additionally fairly numerous. Whereas matmuls kind the inspiration of many fashions, the matrix sizes utilized in these matmuls can fluctuate considerably. For instance, fashions have totally different enter shapes (like various sequence lengths), totally different inner shapes (i.e., the matrices which might be being multiplied collectively as a part of the hidden layers of the mannequin), and even totally different batch sizes (vital for coaching and inference efficiencies). Because of this, matmuls are available a whole lot of various shapes in manufacturing, which makes decomposing them into blocks that maximize reminiscence effectivity difficult.
Knowledge
Lastly, your knowledge may also be numerous. Most readers might be acquainted with knowledge variety by way of structured and unstructured knowledge, however on this case, we’re extra involved with knowledge sort (”dtype”). Knowledge in AI fashions is often of dtype FP32, however the business can also be adopting decrease precision knowledge varieties like Bfloat16, Int8, and the extra unique FP4 and Int4, to scale back mannequin dimension and enhance efficiency. Matmul algorithms, then, want to have the ability to function on knowledge of many various precisions, relying on the use case.
