Accelerating Growth Through Optimizing GPU Usage // Sahil Khanna // AI in Production 2025

Key Takeaways

we have our last talk on this track coming up with sahil s how's it going good how are you I'm doing well I'm really excited this is our last but certainly not least talk I feel like every talk today has been really really awesome i' love having to pay attention and listen to all of them it's just been very informative conversations we're gonna dig in today for our last talk around optimizing GPU usage which is a very cool topic um yeah from this speaker s is a senior machine learning engineer at Adobe and my initial when I think about Adobe I wouldn't immediately think about GPU optimization I might think about like modal or like Nvidia but Adobe puts out such fantastic software and such cool engineering um and obviously they're must be doing a lot of things with generative AI today so it makes a lot of sense they would focus on this um why don't you give a little background on on where you're coming from what you work on and then you can jump into your talk yeah for sure first of all like I'm Tred to be here especially big thanks to all the organizers and speakers like so far all the talks have been fantastic I've learned so much in few hours so thank you for that um like for me I'm currently uh working as a machine learning engineer at Adobe specifically within a firy organization so I work with a team which focuses on training infrastructure and Frameworks um before that I worked like few years in Industry with Etsy and insta card building um machine Learning System for inference and training um so that's about me uh when we are ready we can start talking about the talk that we yeah fantastic um if you go ahead and share your screen I'll uh bring it in and we can take it off okay let me just figure that out okay sorry about that second no no worries okay I think I share the cor correct screen yeah I see your slides okay I'm going to jump off uh take it away when you're ready okay let's do that uh yeah so I think today what I'm going to talk about is some of the challenges um we faced optimized GPU when we try to scale generative AI model training and in the foll slides we're going to discuss some of the challenges as well as solution and our journey to solve those Solutions so let's get into it um before we discuss the technical aspects I'm not sure who for those like who are not familiar with firefly so Firefly you can I definitely encourage everyone to go to this link you can play with it uh Firefly actually currently is a diving force behind a lot of AI capabilities within a Doby applications um most of the features which it support some of them I also like added some Snippets of the video like you can currently generate videos uh using text you can generate images using text you can like uh give a image and ask Firefly to generate a video uh and there are many more functionality which you can do um in the interest of time I'm going to skip to the technical part on how in the back end we are actually powering a firefly website and how what kind of platform we have developed to empower the generative a models to train and then scale at this um scale at obious level so let's just get into that um okay so first of all uh to support the capability of firi application uh my team has developed a compute platform with the goal of improving the developer productivity by simplifying the access to GPU instances however it's not an easy task we will touch on Technical and management challenges to make GPU available for all these jobs and why it makes it hard to optimize the utilization even just before going there I just wanted to provide a little bit reference on the scale we are currently working on so we currently Le hundreds of distributed jobs in the platform and we manage thousands of GPU nodes so let's get into some of the challenges uh we faced so so I've categorized them into four category uh four parts here uh but these are not by any means comprehensive list um they have I've combined them I've only mentioned four for this discussion but I'm sure there are other challenges which we can also discuss if you have more time for this but like first of all I want you to uh specify why it it is little bit important for us to worry about this problem so the first challenge we have is a high demand uh we train a lot of gen models which help us generate videos audios images and it is currently very hard to make any progress or do experimentation or development if you don't have access to this big GPU machines like 800 h200 h00 things like that so there's very high demand from all the research Engineers uh in the organization uh unfortunately with this high demand we don't we have very limited Supply uh so we have to reserve capacity in advance for many years so that we can fulfill all the demands we have um in the in the company uh but with the reservations you can imagine that it comes there's a more management challenge around optimizing the cost um making sure we can fairly allocate all the available capacity we have among different projects and then efficiently use all the allocated capacity for these projects uh and on top of that uh the world is not very idle like we also have very high failure rate in this machines nod when we run it at scale so we frequently have to recycle the hardware bad Hardware with the new hardware and so it also introduces a a lot of Performance challenges which we have to tackle within the platform so let's first talk about the first two categories of high demand and resource management and see how it actually leads to some of the technical requirements from the system so the first of all uh in order to support this scale we have to support scheduling jobs on multiple clusters because we don't want to botland on a single cluster in addition to that we have to ensure that we can provide dedicated capacity to all the critical projects uh and also reuse all the available capacity for the noncritical projects when it's available now in order to support these requirements we end up developing our own inhouse scheduler uh we couldn't find any open source Solutions which can support these majorly three features one is to schedule jobs across multiple cluster and have state of wall where you can share notes between the jobs among clusters second is to support quota management so that you can enable some granted capacity for critical workloads and in addition to that we wanted to support preemption so that whenever the Avail resources are available and not used we can use it for noncritical jobs but always uh we have this uh ability to take it back and use it for critical job when they need it and all of these features work great for all the production training jobs but it created more challenges when we started uh supporting development jobs so we going to talk little bit about what kind of challenges we encountered when we started uh supporting the development uh development workflow so in order to support collaboration uh on the platform we started supporting interactive sessions uh so the whole process is like people can just create an interactive session run the jupyter notebook get access to these GPO nodes and they can do their development and experimentation of um the main challenges with this kind of workflow is that these sessions are mostly idle when people are not working on it or if they're not running anything and also they consist of a lot of custom State because if someone is interacting to a machine they set up a lot of environment in that machine uh and just doing a preemption on this machine is not ideal because it's very disruptive to use a workflow we can lose the whole the configuration they have spent like ours configuring so we had to solve this problem uh so in order to solve this problem we then worked on two main features one feature which we had to build was snapshotting so the whole idea is that we will take a snapshot of the EXP the environment state of your job before we actually stop the job uh so this the advantage of this is like the next time someone wants to to start a new job they can resume their job within one few seconds from the last environment stat and they will don't have to like spend another hours to set up the environment and this actually allowed us to support preemption effectively without provide without giving a bad user experience throughout the cluster and actually optimize the usage of those notes uh uh for all the jobs we have in addition to that we also started implementing a lot of policies on our system so one of the policy which we implemented was reclaim policy where we reclaim resources when the jobs are idle and use those resources for other applications in need and this helped us uh optimize the ideal GPU system uh and actually get to the better performance um overall of all the resources we have and this all of these features so far which in this journey are showed has helped us support both the production and development for workflow and everything was great uh if the world is Idol unfortunately that not that is not the case everything uh uh doesn't work there's a lot of challenges around reliability when it comes to using a shared infrastructure on cloud and that brings up into the last CH step uh last set of challenges which we encountered uh which also affected the experience as well as our performance of the training jobs so most most of this machines which we use on cloud they often encounter a lot of Hardware failures uh they have configuration issues where you can't schedule jobs on it the overheat and there other connectivity issue and many other issues um and all of these issues leads to longer training time sometimes disruption in training and low GP utilization because now we are like spending most of the time recovering and being ID and uh all of these challenges uh actually made us build our inhouse auto recovery system and the whole idea behind this system is there's a centralized brain which monitors the progress of a job it also gather the data from the nodes about its Health it also gathers data from uh other part of the systems around processes the failures which are happening in order to make a decision how to recover from different failures inside the job and autor resume training this has really helped us uh ensure that we successfully completes our training because uh generative models trainings are very very long it it takes days or weeks to complete so we won so this the system helped us ensure that we complete successfully and also uh complete it faster optimize all the resources we have recover faster uh from the issues happening because of bad Hardware or connectivity so all of these challenges um have led to this architecture where at the on the top on the extreme left we have all the user experience interfaces we support a UI a python SDK a CLI which lets user interact with our apis uh where we and then we have this layer where we store all the API all the data uh gathered from the user requirements and then generate events um all these events take actions by interacting with different components in the system uh the the first component for example if you take an example of a job how it get schedule the first component which uh takes an action is global scheduler it reads the requirement of a job it has a state of the world it knows how many jobs are running what are the notes available uh how many uh quota we can provide to this project and things like that and based on that it takes a de decision where to schedule this job and whether to schedule this job and then that action goes to a cluster manager cluster manager so we have a manager per cluster of kubernetes and each cluster manager responsibility is to talk to the kubernetes uh cluster master and schedule the actual ports which schedules a job and then this cluster managers takes the action from scheduler it knows the spe and it schedules a job with all the required um PS and infrastructure around it and we also run a few more agents on each nde and for each Port so that we can keep collecting data about the node as well as the port and also these different agents help us run these different actions which I mentioned pre previously for instance if you want to take a snapshot we talk to the node manager the node manager take a snapshot of the local data on that node and put it to S3 same similarly like we talked to forward agent if you want to like capture some metrics if you want to run tracing and all those things and and this is like a very um uh this is like very this is how like all these components interact with each other on a very high level um so I didn't um have so much time to include more details into it so uh I only put high level hopefully I can answer some questions if you have about this architecture uh if you have more time uh before ending this Pres presentation I wanted to obviously there are like so many people uh who have contributed to this platform uh I I only mention few of them who actively develop however there are many more who had invaluable contributions and because of the constraint space constraint I could list all of them uh so thank you for everyone listening to me and hopefully you have some questions I can answer wow s thank you so much that was such an awesome way to end the uh at least our track of the conference um I'm curious before asking any questions or pulling from the audience do you have any like you shared from other people do you have any like personal links or LinkedIn or GitHub or any place where people can follow you and learn more yeah I didn't add it here but let's see I uh just want to make sure yeah so you can see my LinkedIn profile so you can definitely reach out to me on L connect them can sweet very cool um yeah okay we brought we brought that back on yeah that that's awesome I'm curious I asked the same question um in a in another talk but like what are some of the most interesting applications that you personally have been leveraging as a user like dog fooding the platform uh that's a very good question like because our problems are very complex at scale most of the Adobe Systems are inhouse um previously there are few systems which we have used um and experimented with in Adobe as well as in my previous U job one of the common system is Ray for running distributed jobs we heavily use torch um for these um system uh for the training so torch adoption has been like very recently increased significantly and Adobe is also like heavily invested in it even like we contribute towards fixing bugs as well oh that's super cool um anything on the other side like the most interesting models like diffusion models or Vision models that have come out through this platform that you've you've gotten to watch happen yeah uh so application wise like I think this thank you not also I have used Firefly to create it from the image even I've used it to create a lot of symbols um for my projects like mcods and things like that I I am very I don't have the Innovative Edge in me so like I use firi very heavily to like uh get something interesting out there for my uh creative Pursuits but that's very cool um so is fire the platform that um I remember I saw Adobe release a fe I think I saw Adobe release a feature where you could take like svgs and and Orient them using AI like you could reposition them um in in somehow in three dimensions was that through Firefly yeah so Firefly currently so they are like this basic things where you can generate new videos audios but then it also integrates with Adobe other uh Solutions like Photoshop and things like that where you can actually with prompting can make edits and do a lot of uh these Transformations automatically and all of these models are powered and hosted by Firefly but those are integrated with all the doe Enterprise wishes that makes sense that's very cool um we got a question from rul it's actually a little bit similar to another question I had but he said when you mentioned cross cluster training um like how how hard is back propagation like how difficult does that become and how did you solve it so right now we try to like bin pack all the trainings within a cluster but because we have so many nodes and cluster um running a single cluster doesn't allow us to uh run these many nodes so we had to like should do have to create multiple cluster in order to support this scale but we try to right now like per job we don't have that much scale so we try to fit a single job within a cluster so that we can especially within a zone so that we can use the fastest connectivity possible and have the maximum through yeah that makes sense and it seems like he was saying network network bandwidth wise but it sounds like bin packing is like the way you're getting around that yeah bin packing is the way me um I had a question in terms of you had a whole slide around like obviously being fault tolerant and I think a lot about gpus like from from modal because they have a lot of like very cool serverless functionality I'm sure you guys have built a lot of that inh housee how do you um like how how frequently are your gpus failing such that you had to have that resiliency like in a typical training job how many are you using what percent are failing like what are some of those numbers if you can share them um I don't have stats but normally like every training run there are at least every day few incidents where either the job GPU will overheat and won't be available or the node doesn't have GPU anymore or there's a connectivity issues between nodes so we had to recycle them automatically so that training can keep moving uh we we are currently on AWS um but yeah right now I don't have like a stats around it we just deploy this auto recovery system where we're collecting more stats that's very cool um are you and you guys aren't using any like open source platforms for this training you you rolled it yourself we use torch and torch elastic but the the platform which Provisions the infrastructure is inhouse um we did recently encounter some issue with torch elastic and how it interacts with platform so we are also trying to see how we can either closely integrate that into within the platform or create something similar for Raj's case that's awesome um very cool well s thank you so much for coming on uh it's a really great last talk um yeah we appreciate it thank you very much [Music]