10 ms·
Doing small network scientific machine learning in Julia faster than PyTorch
- sjkelly 5y agoThis is one of the benefits of a capable JIT for numerical computing. Even including compile time, the overall execution time is way lower. Chris Elrod is an AVX2/512 whisperer.
- gnufx 5y agoJIT-compiled code is way faster than linear algebra kernels which are completely vectorized by GCC et al? How does that work?
- adgjlsfhk1 5y agoif you only vectorize the linear algebra, you leave performance on the table. Vectorizing fused operations reduces the number of memory passes. Also knowing the sizes (which are chosen at runtime) is necessary to make optimal decisions.
- gnufx 5y agoI was responding to a blanket statement alluding to AVX512 (which is rather a can of worms). I don't understand what this is on about, especially with reference to the article. Of course you account for matrix dimensions at run time, like to decide whether or not to do packing in GEMM, whether to call out to GEMM for matmul, etc.
- celrod 5y agoSimpleChains does not call out to gemm. All layers are currently implemented as naive for loops. It uses LoopVectorization.jl to compile them, which does a good job leveraging AVX512 -- much better than llvm or gcc. It doesn't pack, which is why column major A' * B is slow: https://juliasimd.github.io/LoopVectorization.jl/latest/examples/matrix_multiplication/#Matrix-Multiplication https://juliasimd.github.io/LoopVectorization.jl/latest/exam... And also why SimpleChains wouldn't scale to large matrix multiplies. But, with 1 MiB L2 cache (e.g. Skylake-X), a 512x256 dense layer of Float32 is still small enough to not really need packing, so I haven't yet needed to implement it (but I will eventually, in a future version of a rewritten LoopVectorization that also actually adds dependency analysis). For an ML library, I'd implement packing via changing the data layout of the parameter vector to tile major, to just skip any runtime packing altogether (i.e., the data would be pre-packed). Only extremely large arrays benefit from a second packing level, so that I don't think it's worthwhile; smaller batch sizes would avoid the need. The benchmark plots I linked above used dynamically sized arrays. LoopVectorization.jl performed better at small sizes than MKL, and much better than everything else. Compared to that application, LV can also specialize on compile time sizes, fuse the addition of the bias vector, and fuse the activation function.
- a1371 5y ago"small" appears to be key here: > When we get to larger matrix-matrix operations, such as 100x100 * 100x100, we can effectively write off any overheads due to memory allocations. But we definitely see that there is a potential for some fairly significant performance gains in the lower end! Notice too that these gains are realized by using the pure-Julia LoopVectorization.jl as the standard BLAS tools tend to have extra threading overhead in this region (again, not optimizing as much in this region). > But, if you have been riding the GPU gospel without looking into the details then this plot may be a shocker! However, GPUs are designed as dumb slow chips with many cores, and thus they are only effective on very parallel operations, such as large matrix-matrix multiplications. It is from this point that assumption (2) is derived for large newtork operations. But again, in the case of small networks such GPU kernels will be outperformed by well-designed CPU kernels due to the lack of parallel opportunities.
- adgjlsfhk1 5y agoyeah, SimpleChains definitely won't be useful for the next large language model, but there's a lot of very small models in the world from physics, to pharma, to high frequency trading. Having libraries that are optimized for these usecases is a major benefit for a fairly large niche of users.
- dleslie 5y agoAnd video games, too. Small models may be useful for procedural content and dynamic environment and actor responses. The barrier, so far, had been the performance overhead of using them.
- dunefox 5y agoThe question is if these users don't already have all of their code in Python/C++/etc. and would even consider using Julia.
- liuliu 5y agoNeural network renaissance definitely renewed the interests to use neural networks as function approximators in many areas. For the one I am interested in recently (dynamics control), many networks are small feed-forward ones or two-layer LSTMs. The importance there is about balance robust controller v.s. update frequency. If you can run the same controller at higher frequency, it is basically free performance improvement (the robot will be more stable, more responsive etc.).
- machinekob 5y agoCool that tiny models have better performance on some Julia frameworks then on pytorch but till Julia get some framework that allows for seamless multi-gpu training a lot of DL people just wont use it sadly. Its for sure easier to just optimize your small model on device exporting it from pytorch then write your own multi-gpu bridge using Nvidia software to train it :/
- WithinReason 5y agoIf your network is small, use a larger batch size to take advantage of GPUs
- Sukera 5y agoThe article addresses that: > This problem is far too small to saturate the GPU, even with such a large batch size. Time is dominated by moving batches from the CPU to the GPU. Unfortunately, as the batch sizes get larger, we need more epochs to reach the same accuracy, so we can hit a limit in terms of maximizing accuracy/time. The batch size in their example is 2048.
- ChrisRackauckas 5y agoYes, we discuss this in the post, and cases where this is not possible. For example, in the adjoints of universal differential equations you have to take a Jacobian-vector product. You can make that a matrix by batching, but then you are making a trade-off because you're effectively solving multiple ODEs simultaneously. When doing this, you might be able to fill the GPU kernels more, but you also have to take more time steps! This is because adaptive time stepping is only correct if you take the minimum of all of the steps (average would make some above error tolerance, or worse unstable). If you have some ODEs stiff, or worse have stochasticity that can cause random times to shrink the time steps at different points for different solvers, making a solver of the union of equations something that is extremely wasteful. And that becomes a dilemma: we started trying to solve multiple ODEs at the same time to fill a matrix-matrix multiplication kernel, but by doing so we now have to take 100x more matrix-matrix multiplications! So that ended up nowhere, and we focused on matrix-vector performance via direct SIMD and effective use of mutation (as the blog post describes) and got very good results. We are also redoing some of the GPU ODE solvers and seeing about 100x improvement by avoiding batching in a similar way, but that's a post for another time. All in all, some applications can just slam a big net. Others can't. The two cases need different optimizations.
- Sukera 5y agoI think this is a great opportunity - actually being able to show real-world (though small) nets in an educational context is great. This means students can try & tweak the net without requiring access to a big cluster or tons of beefy GPUs.
- deleted 5y ago[deleted]
- pbowyer 5y agoI thought that - like numpy - pytorch was a thin wrapper over a C++ library. So why the slowdown?
- w-m 5y ago...a C++ library with a CUDA backend. But these high-performance building blocks might only be saturating the GPU fully if the data is large enough. I haven't looked at implementing these things, but I imagine uf you have smaller networks and thus less data, the large building blocks may not be optimal. You may for example want to fuse some operations to reduce memory latency from repeated memory access. In PyTorch world, there are approaches for small networks as well, there is https://github.com/NVlabs/tiny-cuda-nn https://github.com/NVlabs/tiny-cuda-nn - as far as I understand from the first link in the README, it makes clever use of the CUDA shared memory, which can hold all the weights of a tiny network (but not larger ones).
- rejuvyesh 5y agoNot related to the authors and don't have the same machine, but on a V100, tiny-cuda-nn performance for the blog post matrix power example: Warning: FullyFusedMLP is not supported for the selected architecture 70. Falling back to CutlassMLP. For maximum performance, raise the target GPU architecture to 75+. Warning: FullyFusedMLP is not supported for the selected architecture 70. Falling back to CutlassMLP. For maximum performance, raise the target GPU architecture to 75+. Initial Train Loss: 5.7188 Initial Test Loss: 5.2812 Took: 11.41 seconds Train Loss: 0.0354 Test Loss: 0.0514 Took: 11.58 seconds Train Loss: 0.0327 Test Loss: 0.0511 Took: 11.42 seconds Train Loss: 0.0316 Test Loss: 0.0505 I think almost of the time here is python overhead because if we increase the batch size 10x, it still takes the same time: Warning: FullyFusedMLP is not supported for the selected architecture 70. Falling back to CutlassMLP. For maximum performance, raise the target GPU architecture to 75+. Warning: FullyFusedMLP is not supported for the selected architecture 70. Falling back to CutlassMLP. For maximum performance, raise the target GPU architecture to 75+. Initial Train Loss: 5.5391 Initial Test Loss: 5.5938 Took: 11.03 seconds Train Loss: 0.0444 Test Loss: 0.0545 Took: 11.16 seconds Train Loss: 0.0388 Test Loss: 0.0496 Took: 11.01 seconds Train Loss: 0.0384 Test Loss: 0.0490 See [gist](https://gist.github.com/rejuvyesh/6c428ea12154edbb36cd4359fa75e4c9 https://gist.github.com/rejuvyesh/6c428ea12154edbb36cd4359fa...) for the implementation.
- jstx1 5y agoThe 5x speedup does nothing for me - I want to use a stable, widely adopted, well documented tool, with convenient deployment and lots of available pretrained models for transfer learning. Julia is very far from it right now.
- amelius 5y agoHell, the CPU is probably faster than GPU (... wait for it: for small tensor sizes).
- mochomocha 5y agoA library I designed a few years ago (https://github.com/Netflix/vectorflow https://github.com/Netflix/vectorflow) is also much faster than pytorch/tensorflow in these cases. In "small" or "very sparse" setups, you're memory bound, not compute bound. TF and Pytorch are bad at that because they assume memory movements are worth it and do very little in-place operations. Different tools for different jobs.
- teleforce 5y agoThanks for designing this tool, somehow it's missed my radar. Interestingly it is kind of given that D language performance is always fast. It is a shame that it's not being used more on data science since personally I think it's more natural and better fit to replace Python and C eco-system that's very popular now for data science. Heck, it even beats Fortran in its own turf for the BLAS benchmark several years back [1]. Just wondering why did you choose D when you write the Vectorflow library? [1]http://blog.mir.dlang.io/glas/benchmark/openblas/2016/09/23/glas-gemm-benchmark.html http://blog.mir.dlang.io/glas/benchmark/openblas/2016/09/23/...
- chillee 5y agoIn "small" setups, you're actually more likely to be overhead-bound if anything (especially on CPU).
- mochomocha 5y agoWould you mind defining "overhead"?
- spullara 5y agoLoading the model, preprocessing, post processing, i.e. things not even done in the NN.
- mochomocha 5y agoIf you're loading the model for every prediction, you're doing something really wrong. Pre-processing and post-processing are indeed usually trumping the core computation itself, but as you mentioned it's not done in the NN.
- blindseer 5y agoFor a scientific / academic focused language, I can’t help but feel like blog posts about Julia always seem to evangelize performance in the language in a biased or subjective manner. Anyone else get the same sense?
- rich_sasha 5y ago+1. And aggressively dismiss anyone who doesn’t love it. I’d quite like to like Julia, but that’s one of my sticking points.
- jstx1 5y agoThey’re taking what they can. The language isn’t getting the adoption that it would have if it was as good as the small community around it claims that it is. It just comes across as a cult.
- kolbe 5y agoI really like Julia. I wouldn't say I come off like a cultist, but I do think they spend an annoying amount of time on PR for themselves. I think Julia is stuck in a spot where they are better than say R/Python for scientific computing on a fundamental language structure level and package manager level, but that it's not good enough to make up for the fact that those languages have a more robust ecosystem.
- short_sells_poo 5y agoI think many people underestimate the inertia present in a large community, particularly one like scientific computing. I think Julia is much-much better than R/Python at basically everything important (as a language and platform). I certainly wouldn't call the python ecosystem robust. It's decades upon decades of ugly hacks and hacks upon hacks. The pandas internals are a Kafkaesque nightmare. But R/Python is where millions of people work daily and produce a huge amount of mindshare and libraries. You'd need an incredible amount of resources to replicate that in any other space. I'd argue though that this mindshare is not there due to Python/R as a language (and platform) being better, but simply because there were no better alternatives 10-15 years ago and by now the sheer inertia makes it impossible to stop. You'd need to convince a sufficient portion of people to move to Julia at more or less the same time. Few people want to be first movers. These tend to be the ones who actually care about the qualities of the platform, not just "get the work done and clock out". If you are a data scientist, you won't be paid for moving to Julia. You'll be paid for coming up with working models. And you take a serious risk by moving to a new platform with little adoption. The platform might die, taking your tools and processes with it. You won't be able to rely on your colleagues advice about technical issues. You can run into bugs more frequently simply because fewer eyes have looked at the ecosystem. The end result is that few people take the plunge.
- brrrrrm 5y agoI guess JAX would be a more fair comparison? It has a JIT compiler
- adgjlsfhk1 5y agoI'd be really interested in a Jax implementation. Any chance you can PR one?
- brrrrrm 5y agohere's an impl: https://jott.live/code/train_matrix_exp_estim_jax.py https://jott.live/code/train_matrix_exp_estim_jax.py it works but surprisingly slow on my devices
- adgjlsfhk1 5y agoThanks!
- ChrisRackauckas 5y agoInitial Train Loss: 5.9414 Initial Test Loss: 5.9414 Took: 175.90 seconds Train Loss: 0.0108 Test Loss: 0.0108 Took: 183.89 seconds Train Loss: 0.0016 Test Loss: 0.0016 Took: 188.95 seconds Train Loss: 0.0010 Test Loss: 0.0010 SimpleChains is about 0.5 seconds on this computer, about 5 seconds to compile. So you're right, it is slow, too slow, like 300x faster than Jax, and I suspect something is wrong with the Jax script. It should probably be closer to like 2.5x or so because of memory handling and vectorization stuff, but not any slower than PyTorch.
- patrickkidger 5y agoI've just run this myself. Tidying up the JAX script, I find that it runs in ~15 seconds on my laptop CPU. Likewise, fixing the crash error in the blog post (UndefVarError: alloc_threaded_grad not defined) then I find that the Julia implementation also runs in ~15 seconds on my laptop CPU. So, 1x faster than JAX? Clearly we're getting very different results here, so (a) I'd like to get to the bottom of this, but (b) perhaps we should be more cautious about making performance claims. (EDIT: updated script in this GitHub gist: https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20eaa38f2ad3d38 https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...)
- kolbe 5y agoI've been doing optimization work in C#, and I think findings like this one are unremarkable. I will readily admit C# is slower and less flexible than say C++, but for small things, I can always beat calling out to external code by writing optimized C#. This is usually because even a slower language like C# is still good enough to make the cost of marshalling objects or an IAT lookup become far more costly than the difference between the optimized C++ and C#. I don't know if thats the reason here, but I have my suspicions.
- sydthrowaway 5y agoPython will be done soon.
- mpfundstein 5y agoone might hope
- typon 5y agoDoes anyone care about small neural networks? What can a small NN do that XGBoost can't?
- adgjlsfhk1 5y agoBetter accuracy vs time (especially with optimized frameworks).
- mrtranscendence 5y agoIt's interesting, because the common wisdom that I've encountered is that neural networks are better when you can feed them large amounts of data. But some of the use cases here (e.g. approximating the solution to a set of partial differential equations) are so far outside the kind of work I do that I have a hard time conceptualizing them or how they work.
- deleted 5y ago[deleted]
- tryptophan 5y agoJulia can be as good as it wants, but until I can download JuliaStudio(r studio) or Juliconda(anaconda) and get to Just Work™, forget it. I tried to use Julia, and all I got was visual studio code complaints about this and that and how the julia plugin cant detect something...ugh The website doesn't even explain how to get started. I figured out that VScode + plugin was the way to go based of reddit.
- bobbylarrybobby 5y agoAll I had to do to "install" Julia on my Mac was download the .app and symlink its executable into /usr/local/bin (or wherever). Iirc using the vscode extension requires a handful of packages but they tell you the ones you need to install.
- Buttons840 5y agohttps://github.com/JuliaLang/juliaup https://github.com/JuliaLang/juliaup Rest your weary installers friend. Never install a new version of Julia by hand again.
- ChrisRackauckas 5y agoEase of installation is definitely a fair criticism until juliaup is made the default way that people are told to install and update on the website. We should change this sooner rather than later.
- machinekob 5y agoYou don't need juliaconda with julia cause default package manager just works good enough. But i can agree community is way 2small and compared to big languages tooling are extremely poor.
- lmeyerov 5y agoIf I'm understanding right, this CPU vs. GPU bit seems apples-to-oranges for how GPU/parallel alg research would advocate optimizing this: - if you have one tiny network, no one cares, use pure python - if you have a bunch of tiny networks, yes, doing a ton of tiny GPU kernel calls will kill you with overhead <-- strawman because ... - ... you can likely combine them into bulk sparse kernel calls and spread them out. If more of an iterative simulation, same thing, just over time. This goes back to the 80's with segmented scans, cray, & connection machines. Put all the data together, figure out regular bulk operators, even if individual parts are unbalanced or otherwise irregular, and voila. Since then, a lot of basic tricks for regularizing (sparsity, speculation, ...), and has gone into overdrive w/ GPU era. If this is as worthy the investments these specialists are doing, as it seems to be many years of R&D by smart groups, I'm surprised they're comparing to an unoptimized GPU alg. Presumably, at some point, they did a more reasonable one? It used to be that honest CPU<>GPU implementations were generally within 2X of one another here once you considered normalizers like performance-per-watt (= per $), though new CPU HW is closer to GPU HW and new GPU HW is fancier, so one of the funnier results is, algorithmically, you have to use the same fundamental GPU-era tricks either way.
- adgjlsfhk1 5y agoThis approach sometimes works, but not always. For a simple example, if you are solving control problems/ or Reinforcement Learning, you need to keep your batch sizes small to keep latency low.
- gnufx 5y agoThe article asks "Which Micro-optimizations matter for BLAS3?", implying small dimensions, but doesn't actually tell me. The problem is well-studied, depending on what you consider "small". The most important thing is to avoid the packing step below an appropriate threshold. Implementations include libxsmm, blasfeo, and the "sup" version in blis (with papers on libxsmm and blasfeo). Eigen might also be relevant. https://libxsmm.readthedocs.io/ https://libxsmm.readthedocs.io/ https://blasfeo.syscop.de/ https://blasfeo.syscop.de/ https://github.com/flame/blis https://github.com/flame/blis
- celrod 5y agoBlis is very bad at small sizes, last I checked. Their approach was based on having a single microkernel, and setting up calls to it. Last I checked, blasfeo did not support AVX512, and this performed poorly on CPUs supporting it. I'm not really familiar with xsmm, but someone showed me this: https://haampie.github.io/smm-bench/cascadelake/ https://haampie.github.io/smm-bench/cascadelake/ https://haampie.github.io/smm-bench/skylake-avx512/ https://haampie.github.io/smm-bench/skylake-avx512/ LoopVectorization.jl performs best with AVX512, but not badly without it: https://haampie.github.io/smm-bench/znver2/ https://haampie.github.io/smm-bench/znver2/ Versus Eigen, LoopVectorization.jl does better at small sizes (less than a couple hundred, up until packing matters) on my computer.
- adgjlsfhk1 5y agoOne thing to note is that LoopVectorization is much more general than traditional BLAS interfaces since it can automatically generate complicated efficient kernels eg tanh(A*v+y). The traditional model of calling out to hand optimized BLAS is severely lacking because that means that your code has to make use of existing hand designed kernels which ends up leading to extra passes over memory.
- gnufx 5y agoWhat does that have to do with what I wrote? At least compare maths function apples with apples. Most of the point of optimized L3 BLAS is optimizing for memory so it's compute-bound asymptotically but, for instance, header-only libxsmm isn't normal BLAS. (I'm used to just using -Ofast to get vectorization, not having to include a special library.)