From 25d8cafb1c8509437fef0da155223f48794a746e Mon Sep 17 00:00:00 2001 From: Simon Veitner Date: Fri, 29 Mar 2024 18:46:25 +0100 Subject: [PATCH] Introduction --- _posts/2024-29-03-multi-chip-performance.md | 8 ++++++++ assets/images/LayerMatMul.png | Bin 0 -> 35993 bytes 2 files changed, 8 insertions(+) create mode 100644 _posts/2024-29-03-multi-chip-performance.md create mode 100644 assets/images/LayerMatMul.png diff --git a/_posts/2024-29-03-multi-chip-performance.md b/_posts/2024-29-03-multi-chip-performance.md new file mode 100644 index 0000000..99716a2 --- /dev/null +++ b/_posts/2024-29-03-multi-chip-performance.md @@ -0,0 +1,8 @@ +# Multi chip performance in JAX + +The larger the models we use get the more it becomes necessary to be able to perform training of machine learning models over multiple chips. +In this blog post we will explain how to efficently use Googles TPU. TPUs are especially convenient as they are designed especially for machine learning and easily deployable on Google cloud. For an introduction on how to deploy your own TPU with Google cloud [see this excellent documentation](https://github.com/ayaka14732/tpu-starter?tab=readme-ov-file#2-introduction-to-tpu). + +In this tutorial we will take a simple layerwise matrix multiplication of activations with weights as our running example. The workload may be visualized like this: + +![Layer-wise Matrix Multiplication](assets/LayerMatMul.png) \ No newline at end of file diff --git a/assets/images/LayerMatMul.png b/assets/images/LayerMatMul.png new file mode 100644 index 0000000000000000000000000000000000000000..b9784ccbab0a1e3241688e6319142435b156b972 GIT binary patch literal 35993 zcmeHP2V9fa+IMfo$}PA6wMDCpfXEi2fKo+8r3#1;8DSJikQE@1II31r(5gfQL8VF+ zAu<%1!HE(fB0@wN0x|?-g#ZZ|-}43{0llsF``!EdZu;rZRygl@&w0+XpZ_`MJ-T&^ z`OInar;QmiX6A+;&9;vjGrnoen6Y!GOoA5Ut4D0djM>)izs|hnD;<4&)LlFr-FzLr{nT+-e`tc|-gp-eR}ZYqV4sy5D_5y$XsNB# z|4CCb&6OIO8t_*Wn$~;xVqNj@+Ims1!RfI6 zj&6gbt<}N%YdSfpxoW7X;ccCM{%PeF@4-HSSYJO6JZ@+OtJGJj!*Iyj{DXb4gRL&u zKo4gu(yFyiSz|3cL;h-+4)Fti!8v*l_1y|fR2v$>V~8&yW*E$m8}vi`d^Y?0?Aorg z$=6hE^l0KSy@Sks9evz4<6W>`FtAI|APJhfzzGpqF2RE>x;i?84Q{^Td4Md0aHPS? zBV>pmO>t|0hYQwkXtcqO{&>8XzlYDrMrS+@hjktj($GR2eSPtSk=f-|o z9~^yXLZrdU+0jc($KaeutF?!Vzx!ZcJ#C5JKVm)H+{H74WC!1Ha&-1|^Th|?hUQaV zspTsAGx*rsQ9RaAe?NCe7d%0N6iJ~CvRTY*UpyYBlaQ(8V{s?%AXBUtB8yRM237ml z-~Gre-;ahcME3t`xSzDNbo2?{jyl>_`T@>vCTlnP?o!hPB|!8b(9tVki1~vJe*VE? zX%O5!{IOO(j-b|BW(3eUFf(_5Z!h=_u|h{bpFtgP^$0?&8_ZHPBG%U*E7r)tCLhP) zm53Ji=j~e;(AT1P%qe zP`M{<68KnpiAyo*l>oR-QROOSweTkr5QhcnvjHhBhClH_(kp>4sD=SF6QP}c0p9Dp zyZ{pfT*$-+>+1m+4b~S>A%MV;)~&-$CVzn;4dZr}0sdYdI0R`LM0iC8`5%RnG{vz` z?7+k!Mf%L=3o2?DuLx3vS3V1kAs?Q;Ve~$D?QiKq)74HH279WwymE!5d>Lrek;+Y2dF!UCp z?tcwK{>79Z>wteYrN5)E8`l4SF{PHN2><-Z>p$nrP;U_o`d2-_3@z0Ysf6^w)t^X2 zYVm44DGkwDVy`Z?3F(!<7li1)2)F+yTQ;BAfYc^QdS4J{|0^Xh+6JUo0#XWDe_{jT zWI!Ti5er2!Wf8ZDbB6!>eM>zl#jHPx1JWx2DFv-Z+rY=)rOyVW6ty001JWx2{UP*M z+CbA)VuvoyMy1vZ=}ReWy-Lh+>6d_%*2JGkKx`9I7YFpEv?l(<3rVj8zMv-lo59wr z!~llafr*Pd=`)`%D5s?ew$}e!z}DJYL#dTmy~OcST!Q}lVQZby8a;YSsfDfer4+9C zL?U7dNuQ6AQn*4#9FoOmD7_MpQn=z18xTuG`fNZ-;fhaeKx!rM@t%^j;T4}4fjBfv zoe_v`03y5WQYQkj){7&76pKJg5sT4QAicP$riPS47NcWhO1FpK}jU1VM4^bz7g z{|1ie3y_EZN)d>$p1*+zq_i~F9W4TJc$UJ9NJ?wtPaKgHi{J~2Ur<;}5rD0!A@a$8eqwj<2=agP!D~J7*G$9;CRVur9D=PUp5mWPDYXc;rkI(3 zRyKcD$NgZ<{{O<2gD<^6Anod@CXQA^y7KW*&VNolM^E|j;+S|!DOS(V)w4QY@Wmbz zSG?#t#_$JiboK@ypYjnsQ1cUg{0R9@la`JT+-QLdEaKXqzZ&n?Z|XfIohaTG7=gKYw4mC)O12g-5Pj;qb_vDHHf8m)EGK zK`N1X|8rmCQTyl~#EwsV=|v)#i<>0RwTffDSONd<&$Vi5enDCO)zL}bK$86VGZ&yn z`MHmGCB-uT_irSRzR~x|lv3>HK8INRoBZ6muYT^MDvwHRhrT){c40$aPJH_Gzre|f zuG;;RPt=Jo?G1k`P%?KCJr#eQaPYy<3&SjwD-r^?(9*x0m~X?f|~k2znar(&PK*)8Esq0S_O zmyR=JC#7a4Rh-%8@4ED)#x;{OYWwd_+OS;ZZ06Ll1$Q>pb`XUFz1CTGYC4Pla~QI<*4GqL0L-B4vU!4^z+DuDKk{}=qc%qb%-5aBeK}J z#bqNKX3w?B_c8VvmlL|_t&*TtsnAl`V_6`_p*8PF+?}}Qve81@tmOQQh5w;=~>rAVhjPAeZ`&5~7Tw&QT`(R;eAV{@;Bx#-l4hA#F zXTHA{U3it^@LGHKLJ49fm;iyZ_58c%&kHA|pSnAs0xAK37 z0Z+@_5tzZhvJJQOB<3=jI=7g3pm*hI5T51ZvKtco>QRA7m~Rrbf!n&1Zc8rgAaHlG zW8ak|U_;XPNv0BWorEPXX6FV1lNFby59`3#aZ{o5)ZgCD&YHf^AO={BQm03$wWhy~ zQnM?JP%t!jG;N854XrVQ0ZL>I?;;$2x9<+FYL}o>MyQeci_NgKCrt4Nc@vyu&#VKb zQssQHMa8zE<;Xqdv z2~+a_7PR!NtZvjAg+gj>@9S*GW?@zV6kbi%xpe>CNMWOK;8f4@8&1~~RgE@uP zd8z3WU3+ZVRji-8dP*%R->U39bK26l;-qnW<{h#+cvN3DD6_gL-^Q%(o3xa0Vw zuG^l>sL6^!dwOWx{xtMJTK6^GT#q(wx8-|kt5}DO$_OQ`^{<>G)vPgLw+S!YmRIx+ z$CW)(!S@3(%u=7phn&n9$fs_!G4J7_%4>_kF`Db4C4IZ{+hj4>gQe>R$ss zmL4kY<@9f^%E9$b^~2%W{5^P6oi&AM!t zfm5TB*oky=bYTwB}-?=(psa)#J(lO^0BJZGuz(TnSwc8feLIZy@|*vB|%+QJb7OnH{QuQs$= zC0`&x_@Xs?d>}TwUG6hxScxEBYIK0Of~jK$M&vU4?*YTU68x3J!;1CqQ;Bo=0>X#a zJKD25@}}ObbbVTtO^B4+l^y6EJ#pEsLN2?N{`ka#Ww%n5r*6vFCz!*C6U@;~zwb4z z8P%23f`UBpS;oFvo(88T)32$S<`CNTh-I&!KoAblgaZx5h8?O4biA7H-gkL`%X)t; z)#=f(RFBeYPxH1Gl*zZ|ixbm)(Pw5QX9CXO?g zdqN^Lqi#_0dw|T8N^D>Fhb^o2lOWkSqQ@i@)d&TB6cni|yn)+FMNu3=JkaOQUABH& z9aFWMu{EWluyQKAcE)Y+wR3BFSr7OYsaX^Y|p%!=jVXthQ9VlIGaoC z+hVF{r4S!(r>aoFcgmLCylCl3KF^1>kWO_UX9_fE^4mx$RY#XTkAwjhHF1+U*){z3 z?hkl!_Ug|3@mNlGL%fg{&{|GP)%=e7Tdsa~aO^3Qv=f>Cc#d)MV6@A>x#&JlbAMaY zkpYSgX`Sr3N#H%!+TQ56z$DPog04uVermOm679HsUs{!(Wc;I`xlYc)G?4bH-Awk^sV-=Aq5^)bGj=D};6Xsg4J5Rx9pi*()qF6i!>q=DmO8obiX>JN-GpfA-0fi2 z?}y9?u|DW}HVRR0T&(Z(InA5l@wU%fS*1^P&KCPhr{5)4@_Bngy1q!&b<<@OS>e`> z%LUujY>^Z(xE2tx2UHmgildp=!|jf*l32+P{x?n!*FG`5G z44=-rwalD2wPz_j{zFp4 z!-j=uV1hbM^M&L}zAEcy-NjdB{nwPWDks=4vniHXctP9q0=3HN95qQ)~`{TZM#y6@AuS$5H{OW2V z3H&8zR4;sQa~nH>_He}yL=wn@NQFTUq$u*M^t!4 zbsUEKf>J2p573<>jdZwa#t!g%R9T{Rq+oUZaV>);Bg7O85&~C@L(J?7ZAS1CEE;_{ z-L}xvt*;c*L%I60>-X5xeBQa2c>@kikwN%XSXSHS+%@WYD2oMW7cZJ6fSe6(t=x1(}iynXz%pTLe3rh z`z-SJ75v8Qw7cY|ggcPvb{`EeT!OejyM~3zD+Zn#TF7S@5qr-hCTHEK45~kYb$pdf zjt%X8YF^X|8ME>HI{CvKx_b7_8~|(7JfX@-S4?eU`sBG!OW;mLOlJ8y^1Swc&`G2j z%hJ-Vi`x@u%_xh+cTZQV3~)5h+xO9KhzzMEH7+O3XD^ZaK7kI1ZXvN|pt}1`MsAYV zn>*e&2{f54Az@7Z1<&MpNrZ;8))D1~g8uqQ8~F;BryicPF=3YmV$=+jJpXD<>9PS% z3H5P5HLA*zT!MK_Z3+CQf)N>Ls1SIy(?>9!);S#uV4TJ4N61!D#H3%XbLzfra3eLk z6)vdmJRCo3ti>*W6H9}}bIyf+3?tpkWp2cZwq+gZn96ntF5O%vEiZ;F;KZ-xrk8H< z+`cu^nA_~WM)=Il*b0}&Y3v3Zt@~E1D=mF^u5Mxbm?U|FDITm+|*+bqcqL;#fD{P?QogcRW2?>*p6}g9>h^3 zpJA|y8F`6`*2Py>8f`E?n!_%)Gk-wl{}IY(JvZ;=%`aZIUCu+U2x0`i#9|%!eB1Ri zph0cz^@)Q@6?B$>SqxN!xn(Yip4C*KNc6hi5uw!Hg%+~V4be5cMgqTVxh3M{6k!Sph^%IB(qGoeA3^HleIbkJ(zdIP@)|c~Z z8%H|^wfLSD?+oJBRULM7>2>Sv<4*B;%~{mEy3~R9sBr58g4Da)m~NnUS?wxzyGv_B zORsHiu!`fp{HCuz2_Kfoqtkd6neADvg$2gFDdUwEU^0EN6}dzLC$D25Z&tt* zV=Bo%5CGrZAGdZeCh@UIcrMk6X+_oV$PuiJWu_b1k&j}Z@hvK;0=ZbO^Pzc4pt8qZ z?IS#H)#i{YcJY@a0vRVF25wWsygAM0=+{^o*`t~<1h*pKKxvU5e`hZg=4T0ctc#OL zk3$`OVqfkfHFMc+W;KG|3<@f=c6vOa9E`Kbug{M#e@Pm%CY?>;PE6_uVP^MGl;voG zRI@*phh2c=sPZ0(Edax1SiyguJL8OA`8ce_ur4qc_Rr;6@?Tn>WsrK)JYr@xRaw>6 zl~*ONS8TH;_w@r*Y}>NTVisxB&uh>xvh417Emzrjb8gszpG8r%1%;#`(bsBsljqT1 zvngFg*Or^WvgxpFhGCQD1KVE}{q=;LI3cI|sU^wxoeeFRixPIDO3%$N&Q7@hC@H(@ z-UANomhE!rO{JVQR}6l4i>98$-bYGjw@}v_@Ow(>4`k50TNsgL=KPW~$DzcXm-eFl z9OQQm2jT=Tinfv8H|nTugB-|;3zPaNYQ-FCzIT|#P3a9u!{E8C^(~21b$l@0?!|~X z!G-T^p)RhzHDFKnE+BBftA`9GXb3<;|u0C1U)?~iNGxrzQzUU}IF`I&G z_*t(*JlS-}_jJwqSEJwDTsKoEhPgjpu+TG16_Xdb>UY{ZWz>dxPrZ)VY4REEfzuyw zbo07S__;*yqq zc|Vt}c))Smpnt!{4UIOyC1%`$+*JdVxr$y^=}M(@`kK^xv0{lFw~kcZ7ZYAV-U_+A21zw_-e*=2~>2`<@Kx)^4AB6W?z_ z)T?fg8i4KGDzXARISm*=gEC5|{gP2Tt95^9v#HX2{?0?JVnc>c-n(b3g-5|`BHpy9 zsb<%#n$g6y=<^L=RaW8VFM+f&>ISGl)mHL*oUX;G%|}>!9o4g%uasni_FI<&B6g?->^4sewN*o3$Nn*qJ^95SQm z3wgiLUyphcwjS7SHZ^^DeP<13y9Fa&9Z_og6TA~O{U^Gu$!zO*6sEP?#;enztwWwz zgzTwcxK}YNI)Vxe<%lhd$Q^Y$Q;2EuRRBCSm(F9_6YMkl>%wWC_E7jI(H?^HI#=_& zgiCkPf}%Z?E&45-O}VOw8T7A{NqEkwS(iu$>!HBWpoyIK zvOcz-1&VOAYBeDv@f5{A(3Zn%LSqCUY6M9v6Z84`@dtyAF3nt^BZ>eyO-3E>Lqku` z@ZTH#d;2?Pbp{@z4?vnZ&G?qUbsyJm^(@5~!2DT5HP!-yUF7p#li)Wf4_FHfLQ({& zFY8;w^w9nJEmYT#zU-Q?w`hJTns22)(3VCCX_F0$1+h>Ooz_mBwM5vZB;1As;^!tu zX4R=pVuh@s?DpMY3$cNdKiPP;LEeUm*>^W@yaPZb&6r##8~1}Hn(uF$;`FE)n0VM9 zvi`{c+(Nm*TWQvcI!!+yn3Dzt4r$&MsLrB7LJ!VrLVPB&CuFM_(tNq3E3Z9(s-jyH z%Gw?oc7>r8Lue8gryKeAX~!caE7N&0_@2!P=w4^^PC$K4N!)bFBtizZ&dBoUSNANU<3}To{_D>9Qs`eHr#1lwNDyAmvwCr&n zf`^~7SGZ}nx`lC=c`s*Hvl?kj&a9@mfy$;rI!4efcZ3;wf*Y?z4SJ+fXiejmi3;SW z%i>gp{2HOX@jwORjt%Y_Qx9>*-$kzW-4n^yJ@C>mUhp7U@Is9ZelX+E55RxQ@?B() zxS2ru4TSV{O3AdBmN7Wk*(-UlFjSAKFmiLSK4k`pS@Gz{69>wSey(TO->w~lDc*C^>F*qK04D z#sQD=u&2HR!c(Wa3V?*jCWEejj}sqzy+F6w2idocFeYVboTT;MrYi8KU7R8FAh=J} zg{J?cn%6>Ia=l|#IqFa=_uRXT+f2;W-ov^%mwCN-VIcECwEON}o+}!?^g@iXs)MkM z6J>{+$3Itk&y-{0K1hY?aK`1;Ng+KH9B#6({;=3T z&pZVo+M>)7eDxnR|BUFXPpXfw++u@Wdj_Igi}*PNdndQ=jl;;i)20R&H8YjU{ft@-0hF;&KYu;+3S?Jy>{3kD zf)-TOhii+s!;~8&rzEUJoTmj}k9Y^;iPh`)tw@M|j3hn>YLKj;B*CJk;EAPP?Jw)( z+g7?W*A*Nk)*W58UMY=Ai_2D|)P|O%FYKKPK@gYSf?5GVuyXf%R8?!h{BAn}RY2B7+(f%25?a zT?Jm`0(PH$EV6%M4@u)T`Gs!z#+4=(C#baUTY_U5T$!KQ^o42FJiwtK)k}wL`uV3x zmk0(c+UOso7x(S2rhSMrvPIIK$udo_#r3|~+z7I|QzaeWM$RY?S155B&y9~X-n9kI z@9FM*AVZNM_dkeln9&q|T~W}H7N=AdScds2kmiw@`2v$w)|UQG*OT0Str~L%e7Zv? zFj^lNUFoV;=tKJDIAoEOEF`=iJi-d`W0z;--CoUkcN*DW6DF^dU+_nP(XsfX3@klr zJGrB0hqZk}f!yOmM(;Sl&$qox6=SRL7jVfvlrDO}g>FMZyBX}i&DVQnNuF1;9scu@ z`)+Wb;REb57`^IvpAe_sTArlR5!z%(Z!C(IQLPNoz}WQrw3gT0*ouDPKr?HBDpoh9 z;6-#b>Td4QBdsXSy5sx?*l;6q?#t&v2;Xja4AQ$*ujr2d$yj$aEoRJSqd2 z_VMV1+D68sB#ZKdW$OV)M1l#T_if;KR87rWP;BGI0y3|lil%xWLQBgqR*9@Frlv@@ zW_3s(ubYlS(Rw{;J8?s??bUPzrp~QxV6nXqwr`4tS?~4cT+aerk@C*bxrnA^l>UYHKQRAt+&t zvJq_OVPej}G|-}UL;ANKBIxxAAPTE6l|QGpFmPqH9Pni*Jn_w&hV{G8_ty3&;nwNtNvr?lJ1 z-`pY)C9X>M{9se9#;FA6eH207C8!obNQnR;dC-EBrzX}M*fse*B%`_UV{=K*_<$ER zpP?^{4Z5=1uSh$dot)EMSQ5gvtBCA5S+nKYI{B@7%Jv z!ZL)%TzdbGD)l9RF5nU#B)f>#nGfq!-CZg7+Z{JXYza*c&;o=Ux(H|vph@R>!5T@zAcL4^FC1*iwim3z zxqia#s$*&X{9kvx2sIa>&5SX~{-HdL+y0Y6y8Ph}T(#CjYpVKma`|^5Z7Br^DksCs z)Og0+Pl_Q_WiP~Ra>;R*&*;jz_5oslREj!iqEL&8abi=9vgn+WrdZb$K@WgS3_thz z4s;P8McZC-Bn0J))%}(Eok7qZG=l3XxrEoZ^6D(6#5{_oD0>lcu~!66*KxH8w43w7 zb}g@uUq{kvPoz2*w(H6;b@}Va?^{C|Tmunw{n}r&K}RQrvV2%`5^uea*hx;|Vf&x~ zAb&xsc_Pi`aOwaSZ49tCS&ZUii%`hE?m*X&29f5^u4)&t-V^zY(7T6qR}3ZOUl z;|_->3E6R|hy?h5C%GcN45MmV!B15};C}D{z=O-kQ-Hk7YQzY>EeqK^1X2L}=h79E z1Wxk`+bR5h0Q&*Z077ScB-B-m-(SEWyO_Q80R5L?dfnOb^sSI*6hQ8cLjEhN`;T7nC_;Tr zAUoZ_e1QB%ulTEIM(WVtQ#APMoXkqVgt0{zw2PRyl()Hk)poEscsCVSl_f{l#yB%# zjT)eb8YHNmtlM^~HYl8UcQCx+0(EccX5F*MsmjgKI@*}*M+c?{zRAGjQz|QyE<>XeeP` zKa4mvqGWOW>sge~oCf;3z5lOm?|(gu^7Y`9G#s2_b+~vyrO|+tItekph0U~ue$O-; z;M|GIK<@_?rUTJCv+;uGwZ(DEMd$Vo|MgKQfkG4P15@DURpg`|E)CA@2lTipJhL=- z|9kgYYUtIt4`9K`#t)U5X*|F-T#-Zk!OfSCtQWEa zs6=~CPXgMWRJV+Dr4lxY|CNJ{=qDUNnbc*B`}c5Qx6`3K+1@U&quL(s+vsrJ2jGM| z_vw=vIUA9KXTg(>J(FwP31{T48cUp4lQ?-d8GAjC=zuO1G{;9O_EQls(gybPA~&0= zmo7w3qkj9>r{{RHo0K@Gl$K)#UNhH^58?khfJENn)vRa#%2CwCwg&xGtP=}#xwfoX zJKE~X$9}kU@~_23n%ktryQic|>}{KuOKf*5DnLCHojdK<|90)7a5x`HH2($99#%3`9f5`FD{nYD z%G>kJ1lQfbP-Pnk%%2=JPCVE+E<5nUioPlWh819^WyBDWF@g@GgQEQ;@KiLcm+E)nd>gKGk;i{$K) zmj@=mm4XQye5Q=NXix*KoMKv)#QnE5`__!xFO)HN%-{XTEXC?E@MpuiEoR9kdk*{$ DMryG~ literal 0 HcmV?d00001