MLOps Community’s Post

Transcript

The next topic that I wanted to talk about, which is deepspeed and all your work that you've been doing on that, because it feels like DeepSpeed helps out with pre and post training, right? Yeah, yeah, yeah, yeah. Right now we also do quantization. Yeah. So there you go. It's like the trifecta. And can you explain to people what deepspeed is? I think we can call deep speed like a third party library based on Pytorch. What we did is that we support some new features that Python doesn't have previously. The most important one is called the zero optimizer. Basically it's like a data parallel training paradigm, but it is very memory efficient, different from like traditional data parallel training that each GPU need to maintain the whole model parameter in zero strategy, each GPU only maintain a portion of the whole model parameter and one specific layer need to do some computation. Every GPU do all gather or do like a recollect of the weight from other GPUs to finish this layer's compute and then just erase or dump this layer's weight. Once we need to call the computation of this layer again we do all gather again on the weights across the gpu. So it's memory efficient. That's a huge performance gain compared with vanilla Pytorch. Because in Pytorch previously it doesn't support such memory efficient data parallel training paradigm. And right now they just doing something similar called fsdp which is similar as zero, but it's a Pytorch native library. This is one big piece. And the other thing we did is we do some kind of data offloading to further release memory pressure on the GPU set. So if a chunk of data doesn't need to compute at the moment, we move it from GPU memory back to CPU memory such that we have more GPU memory to use for other computation. And after we need to compute some tensors on the CPU set, we just move it back to GPU memory and do the compute. Once we get the results, we move it back to save your memory, to save the memory.