Lightning Talk: In-Transit Machine Learning Using PyTorch on Frontier Exascale System- Vineeth Gutta

Key Takeaways

hey everyone I'm venth I'm a PhD student at University of Delaware and I'm advised by Dr Sita chandan and I'm part of the computational research and programming lab there so today I'll be talking about a workflow um called intransit machine learning using pytorch on Frontier exoscale system um this is a use case of P torch for a scientific um application running at scale so this was a collaborative effort across um the us and Europe so this includes University of Delaware um oish National Lab and um hcdr which is a National Lab in Germany so to give a background about the science case um pick on GPU simulation is a particle in cell simulations um physics simulation um that tracks particles within a grid representing the space um I'm a computer scientist and the benefits of working with in interp project is that we get cool graphics so I'll show you cool graphics um so here you see a laser pulsing through a plasma Channel and the pulse displaces electrons in the channel creating a wake thus creating positive and negative charged regions which is accelerating the electrons to achieve higher and higher energies thus creating a high energy particle accelerator and the which is the goal of this project so we want fast and predictive simulations for the Next Generation laser plasma accelerators um so I want you to imagine that you're looking at the Golden Gate and there's a ship going underneath it um forming waves due to the water it displaces now imagine Surfers jumping into the wave and riding the wave they get accelerated don't they um in this analogy the San Francisco Bay is the plasma channel the ship is a laser pulse and the waves are the radiation and the Surfers are the individual particles um so what's the point of these simulations these simulations have various um applications like chemotherapy where concentrated energy needs to be maximized while reducing its effect on the surrounding regions and particles um so the challenge with all of these simulations is that they require solving complex non linear equations for trillions of particles creating terabytes of data and the use of machine learning allows us to learn complex Dynamics from this data so going Beyond just our application um the use of machine learning to learn from large volumes of sim simulation data is not a unique problem to pick on GPU um a couple more examples are that digital twins such as destination um Earth from the European commission uses climate and weather models such as um e esms and um ecmwf ifs both of which provide um generate large amounts of data and there's also the cerns large hron collider which is a particle accelerator that we're all familiar with so um where do we run all of this right um so the we run this on the frontier ex scale system this is currently the fastest supercomputer um it's located at the Oakridge National Lab and it contains approximately 37,000 total AMD Mi 250X gpus and pick on GPU was one of the eight research projects that was selected um in 2019 across the country um to to test um the system so to give a overview about the architecture of a single node uh on Frontier um so this has again the amdm 250X GPU so it technically has four gpus but what you see on the right is um it's logically presented as eight um that's because um the software sees amds Mi 250X gpus as AG a individual gpus um because of the way the architecture um shows up so each GPU contains two graphic compute dies so high level libraries like pytorch see two gpus um and pick on gpus pytorch code was single GPU for NVIDIA gpus uh my role was to test and validate the pytorch GDP scaling on the AMD Mi 250X gpus on Frontier when we first started it wasn't clear that to what extent Pi torch would scale um especially on AMD gpus um so the challenge is that as GPU acceleration of plasma simulation codes such as pick on GPU um Advanced and led to a speed up of the data generated um the volume of the data stream produced by the simulation um made offline and analysis of such simulations impossible due to file system and IO constraints um this is because the GP acceleration um led to the increase in data generated so there is a mismatch between throughput of the GPU memory and the file system throughput so the solution is to use intransit training of complex ml models which require processing of data produced at Source by the simulation um at extreme data rates through streaming um the data is streamed from the simulation using um two Pi torch using a custom implementation of open PMD which stands for particle mesh data um standard and the combination of open PMD API and audios 2 allow direct in-memory transport between one parallel producer code to an arbitrary number of parallel consumer codes which in our case was the pi torch code um and the data streams are established only between the processes that share data reducing um communication bottlenecks so why do we need streaming um this is mainly because of system constraints we face um so 20 to 30 terabits per second of parallel throughput is produced by the simulation ver versus the file system which only has about 10 terabits per second of bandwidth and just to give an example when scaling to just 25% % of Frontier system um so approximately 37,000 divid 4 is um about 9,000 gpus um one petabyte of data is produced for every time step and in our simulation there's approximately 1,000 time steps where each one lasts from 0.1 seconds to a second so we often see 1 to 10 pedabytes per second of particle data um requiring on the order of 10 exabytes of total dis space so the upshot here is that the storage can basically be exhausted in as little as 100 seconds so this is why the intransit ml um is necessary and to give you an overview of the idea of what's going on here we use in transit ml to bypass disk storage which is a bottleneck for applications like pick on GPU the simulation produces so much data that we must still do reduction before um the ml model can reconstruct and learn um correlations reconstructed data and learn correlations in our case we do in node streaming where the simulation um which is the producer and Pi torch which is the distributed consumer run on the same node but on different gpus so that inmemory transport of the data can be achieved and in this case the model learns correlations using unsupervised techniques um between the emitted radiation and the particle Dynamics so this is an overview of the um ml architecture um it uses an invertible neural network and it predicts particle Dynamics from radiation intensities so this is a general overview of it and um the goal is to again find plasma instabilities connecting radiation with the source of the plasma again this has various use cases in Material Science and cancer applications and due to the nature of the simulation and streaming um the model uses continuous learning but such continuous learning techniques run into challenges due to the rate of data that is produced and to give you an idea of the what the scale that was achieved um I want to preface this by saying that not all model architectures are as scalable as llms um llms are a recent phenomenon and there are still many ml models that do not scale as well and we ran into similar issues with our um model where we scaled it up to 100 nodes or 400 mi2 AMD Mii 250X gpus and we believe the model architecture was a primary constraint for scaling further because the communication bottleneck at that point um overwhelmed um any benefit we would gain from scaling further and this but the simulation part the pick on GPU the um the plasma physics part of it um Can scale up to 8,000 nodes on the system so we know that there's further room to um improve and the the chart you see on the right is showing efficiency as a scaling go as we scale on the ml model um the efficiency goes down primarily due to communication overhead um as pytorch GDP does a collective all reduce to average the gradients across each instance of the model that's distributed and efficiency here is defined as the lowest possible work workload benchmarked as 100% effici so we take uh what was the smallest run which was eight nodes and um use that as The Benchmark so the goal of this entire ml workflow was to extract physics insights from such application the model learns correlations from the radi uh from the radiation predictions and also correctly identifies the the instability regions so to evaluate the performance of the train model the radiation Spectra is inverted back to the original momentum um distribution but this remains an unsolved physics problem so um what we see is that it seems like the model learned a fundamental aspect of special relativity the which is the Doppler effect since it reproduces the correct cut off frequency which you see in the gray on in chart a on top um when predicting radiation Spectra but the momentum prediction was less accurate um despite that the model clearly identified the instability regions um so what was the point of again what was the point of this entire uh workflow um it's to prove that it's possible to learn correlations from physical from physics simulations in real time and this has an added benefit and the well the main benefit here is that in even in hetran systems um they can benefit from codes like this um by circumventing um disc capacity and bandwidth constraints that are posed um again this was a large team effort um so I want to thank everyone and I want to thank the pi torch community and AMD for all the support to make this um project happen um and I'll take any questions if you have any questions no time for questions thank you yeah