De 15% a 90% de utilização da GPU: corrija o pipeline de dados, não o modelo

Autoria: AI Engineer · Assistir no YouTube

Vídeo de autoria do canal indicado, não da IANIA; todos os direitos pertencem ao criador. Transcrição, resumo e tradução gerados automaticamente; o conteúdo no idioma original é preservado na versão em idioma original.

Idioma original: inglês

Vídeo original: From 15% to 90% GPU Utilization: Fix the Data Pipeline, Not the Model · Canal: AI Engineer · 16:33 · Assistir no YouTube

Tópicos: Dados, bancos de dados e embeddings, Inferência e otimização de modelos, Infraestrutura e hardware de IA, Modelos multimodais

Resumo

A palestra mostra como gargalos de carregamento, processamento e transporte de dados podem deixar GPUs subutilizadas durante o treinamento multimodal. Partindo de um pipeline serial com apenas 15% de utilização, o palestrante apresenta concorrência, pré-busca e uso do armazenamento de objetos do Ray para manter os dados disponíveis e reduzir cópias. Ao escalar o sistema, surge um gargalo de rede, que é mitigado ao distribuir os workers entre nós e ajustar opções de recuperação sem cópia e agendamento. A principal recomendação é perfilar o pipeline e identificar o gargalo atual, pois ele pode mudar conforme a escala.

Pontos-chave

O desempenho do treinamento multimodal depende de manter as GPUs abastecidas com dados, não apenas de otimizar kernels.
O pipeline de dados inclui carregamento, transformação e transporte dos dados até as GPUs.
A concorrência permite carregar e processar amostras independentes em paralelo.
A pré-busca mantém uma fila de dados pronta antes que o trainer a solicite, mas exige atenção ao checkpointing e ao aquecimento inicial da fila.
O armazenamento de objetos do Ray evita cópias desnecessárias de tensores multimodais entre processos.
Ao escalar o sistema, a interface de rede de nós que concentram workers pode virar um gargalo.
Distribuir workers entre nós e ajustar a recuperação sem cópia e o agendamento pode melhorar o throughput em larga escala.
O perfilamento da linha de base ajuda a identificar gargalos que podem mudar com a escala.

Capítulos

Início	Capítulo	Resumo
00:01	O gargalo está no pipeline de dados	A palestra apresenta a meta de aumentar a utilização da GPU alimentando o treinamento multimodal com dados continuamente.
03:27	Pipeline de linha de base	O palestrante descreve o modelo, o formato dos dados e o pipeline serial que resulta em baixa utilização das GPUs.
06:12	Concorrência no carregamento e processamento	Workers paralelos e Ray actors reduzem o tempo de preparação dos dados, embora o trainer ainda possa esperar por eles.
08:56	Pré-busca para manter dados disponíveis	Uma fila preenchida em segundo plano antecipa o trabalho e reduz a espera do trainer por dados.
10:35	Redução de cópias entre processos	O uso do armazenamento de objetos do Ray diminui a transferência de grandes tensores multimodais entre processos.
12:42	Escala e gargalos de rede	Ao aumentar a escala, a rede dos nós com workers se torna um gargalo, levando à distribuição dos workers e ao ajuste de parâmetros.
15:26	Aprendizados e recomendações	O palestrante recomenda identificar e perfilar o gargalo atual, lembrando que as configurações ideais podem mudar com a escala.

Citações

A parte mais difícil de otimizar o throughput do treinamento multimodal não é necessariamente a otimização de kernels; é garantir que as GPUs recebam dados continuamente. — 00:35

O objetivo é garantir que os dados nunca sejam o gargalo. — 03:27

As GPUs são caras, mas às vezes as maiores batalhas acontecem fora delas. — 15:26

Perfilar a linha de base ajuda você a avançar. — 16:33

Conceitos e entidades

Ray actors (ferramenta): Unidades de execução do Ray usadas para distribuir o processamento de dados entre workers.
Ray object store (ferramenta): Armazenamento do Ray que permite compartilhar referências a dados e reduzir cópias entre processos.
Megatron (ferramenta): Framework usado no exemplo para treinar o modelo e distribuir o trabalho entre ranks de paralelismo de dados.
Qwen3-VL (modelo): Modelo multimodal mencionado como base para o treinamento descrito na palestra.
S3 (ferramenta): Armazenamento de objetos presumido como origem dos dados no exemplo.
Python GIL (conceito): Mecanismo do Python, também chamado de global interpreter lock, que afeta a execução concorrente de threads.
Pré-busca (conceito): Técnica que prepara dados com antecedência para que estejam disponíveis quando o treinamento os solicitar.
Recuperação sem cópia (conceito): Forma de acessar dados sem criar cópias desnecessárias durante a recuperação.

Tradução completa (pt-BR)

00:01 [música] Olá, pessoal. Obrigado por estarem aqui comigo. Hoje vou falar sobre ray actors, vision tokens e o GIL: como projetar um pipeline de dados multimodal de SFT que mantenha as GPUs ocupadas. Há muita conversa sobre como

00:35 otimizar paradigmas de treinamento. Mas quero oferecer uma perspectiva diferente sobre todo o ecossistema necessário para treinar um modelo. A pergunta que vou fazer aqui é: seu loop de treinamento realmente depende das GPUs neste momento? Ele está bloqueado pelas GPUs? Acho que vamos aprender, ao ver este experimento, que a parte mais difícil de otimizar a vazão do treinamento multimodal não é necessariamente a otimização de kernels. É garantir que as GPUs recebam dados constantemente.

01:09 Então, vamos revisar a visão geral do tema desta palestra. De todo modo, vamos estabelecer um pipeline de treinamento multimodal de referência, encontrar gargalos e, depois, resolvê-los um de cada vez. Em linhas gerais, eles são: um, concorrência — com que rapidez você consegue processar uma imagem antes de entregá-la ao modelo; dois, pré-busca — como garantir que haja dados disponíveis sob demanda; três, como minimizar as cópias de dados entre processos diferentes; e quatro, f

01:43 otimizações adicionais do framework — apenas ajustar os detalhes e os recursos extras. Então, vamos estabelecer nossas premissas. Todos sabemos que as GPUs têm custos elevados e, em geral, queremos maximizar sua utilização enquanto estão à nossa disposição. O processamento multimodal é diferente do processamento apenas de texto, porque há muitas operações que exigem bastante da CPU e precisam acontecer antes de os dados serem entregues ao modelo. Isso inclui, por exemplo, carregar, redimensionar ou normalizar

02:20 imagens usando uma biblioteca como Pillow; para áudio, há muito carregamento e recorte, e você pode usar arrays do NumPy; já os vídeos são uma combinação de ambos. Também é verdade que, em geral, seus clusters de treinamento têm muito mais CPUs do que GPUs disponíveis, mas, neste pipeline, quem fica sem dados é a GPU. Hoje vamos apresentar uma equação chamada proporção entre tempo de espera e tempo de processamento, que formulei, em geral, como o tempo necessário para preparar os dados e transferi-los para as GPUs. Depois, dividimos tudo isso

02:56 pelo tempo necessário para fazer essas duas coisas e também treinar. Basicamente, o denominador é o pipeline inteiro. Vamos ver que, durante nossa execução, a utilização vai passar de 15% para 90%. Na configuração de referência, 85% do tempo é gasto esperando pela obtenção dos dados. Depois dessas otimizações, esse valor fica abaixo de 10%. E, para

03:27 o escopo desta palestra, as GPUs não são o objetivo. O objetivo é garantir que os dados nunca sejam o gargalo. Aqui está um pouco do contexto sobre como será nossa stack. Vamos usar um modelo baseado no Qwen 3VL e tentar treiná-lo em paradigmas de chat, VQA e conversas multimodais com VLM. Recapitulando: antes de as GPUs fazerem qualquer cálculo, nós

03:59 precisamos carregar os dados, transformá-los e transportá-los até as GPUs. Certo, esta é uma visão geral de como são nossos dados. Vamos usar arquivos JSONL, e as imagens são referenciadas por caminhos relativos ao local onde o exemplo está armazenado. Vamos supor que nossos dados estejam armazenados apenas no S3. Neste exemplo, por exemplo, parece que a imagem 33471.jpeg

04:32 é a foto de um ônibus azul. Este é o layout do nosso sistema. À esquerda, temos um gerador de dados, responsável por carregar, transformar e transportar os dados. Isso está identificado como distribuição para os ranks de paralelismo de dados do Megatron. Os ranks paralelos recebem partições dos dados, conforme determinado pelo coordenador. Eles agrupam os dados e executam a passagem para a frente e a passagem para trás. O Megatron cuida de otimizar os

05:05 pesos do modelo entre os ranks de DP. Vamos supor que ele esteja funcionando como uma caixa-preta. Então, vamos começar com o pipeline de referência ingênuo. Para cada exemplo, buscamos as imagens no armazenamento de objetos. Depois, tokenizamos, redimensionamos as imagens, normalizamos e as entregamos ao modelo. Quando fazemos isso,

05:37 vemos que as operações em vermelho são seriais: buscar todas as imagens uma por uma e, depois, processá-las também uma por uma. Se o treinador solicitar um tamanho de batch de, por exemplo, 256, será preciso processar todas essas imagens em série. Podemos ver que o tempo de treinamento fica entre 15 e 20 segundos. No entanto, o tempo gasto para buscar os dados passa de 100 segundos. Por isso, nossas GPUs têm apenas 15% de utilização. Também é possível ver um gráfico de chama de perfilamento

06:12 que mostra que passamos a maior parte do tempo fazendo carregamentos do S3; depois, há uma pequena cauda no final, referente ao carregamento da imagem propriamente dito na biblioteca Pillow. Certo, solução um: vamos tentar adicionar concorrência. Em vez desses processos individuais e seriais em vermelho, vamos tentar paralelizá-los usando um pool de workers. Quando você solicitar 256 exemplos, talvez possamos processar todos de uma vez. Por exemplo, nós

06:47 temos dois custos. Um é buscar as imagens; o outro, processá-las. Isso inclui decodificação, pré-processamento de visão e tokenização do texto. As amostras são independentes umas das outras. Isso significa que são trivialmente paralelizáveis. Então, há decisões de projeto a tomar: que formas de concorrência você quer usar? Há a programação assíncrona, adequada para operações de E/S, como fazer consultas SQL ou solicitações de rede. Há também threads, que você

07:20 pode usar para cargas de trabalho intensivas em CPU que liberam o GIL do Python, ou seja, o bloqueio global do interpretador. Não é óbvio quais operações realmente liberam o bloqueio, então você precisa investigar um pouco mais. Mas os usos comuns incluem tokenizers da biblioteca Hugging Face Transformers, a API de transforms do Pillow ou algumas operações do pandas. Já os processos são mais pesados, mas também proporcionam isolamento total entre as cargas de trabalho, então você pode aproveitar vários núcleos da sua

07:52 máquina à medida que conclui o trabalho. Uma análise aprofundada da solução certa mereceria uma palestra própria, mas há muitos recursos online que podem ajudar você a decidir para seu caso de uso. O importante para a nossa solução é corresponder ao nosso gargalo, em vez de simplesmente escolher uma ferramenta que já tenha funcionado para nós antes. Vamos escolher async IO para buscar as imagens e usar ray actors como nossos processadores, com a biblioteca de multiprocessamento. ray actors são essencialmente iguais a processos

08:23 comuns, mas oferecem APIs para comunicação entre workers que nos permitem não ter que reinventar a roda. O desempenho do seu processo é tão lento quanto seu elo mais lento. Podemos ver que nossa proporção entre tempo de espera e tempo de processamento diminuiu. No segundo exemplo, em rosa, o tempo gasto buscando dados por batch caiu para cerca de 40 segundos. Isso nos dá uma aceleração de 25%. No entanto, como você

08:56 pode ver, há uma grande lacuna que começa a se formar, especialmente à medida que avançamos por muitas etapas. Então, será que podemos reduzir o gargalo de esperar o treinador solicitar os dados antes de preparar nosso pool de workers para buscá-los? Podemos buscar os dados de forma antecipada? A solução dois que desenvolvemos foi ter uma thread em segundo plano que preenche uma fila antecipadamente e se mantém sempre à frente do consumidor.

09:27 o que significa que seu produtor está produzindo exemplos em uma proporção maior que seu consumidor, e assim o treinador nunca fica bloqueado esperando para solicitar os dados. Quando solicita os dados, eles estão disponíveis sob demanda. Há uma ressalva no pré-carregamento: ele torna o processo de checkpoint e retomada do modelo e do carregador de dados um pouco mais complexo. Além disso, pode haver um período de inicialização a frio, durante o qual a fila está sendo preenchida para ficar pronta para o treinamento. Mas, se você trabalhar

10:01 com esses detalhes, acho que o pré-carregamento é uma ótima ideia. Então, acho que garantir que o produtor esteja sempre à frente do consumidor vai trazer grandes benefícios aqui. Sem surpresa, chegamos perto de zero segundos de busca e processamento de dados por lote, porque já tínhamos buscado e processado os dados muito antes de o treinador precisar deles. Com isso, resolvemos efetivamente os gargalos tanto no carregamento quanto na transformação dos dados. Depois dessas duas medidas, nós

10:35 achamos que tínhamos terminado: o throughput aumentou e as GPUs estavam sendo alimentadas. No entanto, observamos que os custos de transporte de dados ainda persistiam — na verdade, tinham aumentado. Cada amostra transfere grandes arrays de imagens através dos limites entre processos, várias vezes. Cada amostra é serializada no worker, copiada para o coordenador e depois serializada novamente no treinador. Para cada amostra e cada etapa, temos uma cópia de memória que ninguém realmente pediu. Então, como isso se manifesta é que seu

11:06 gerador de dados pode, muitas vezes… E os tensores multimodais são enormes. Eles têm megabytes por amostra. Por padrão, cada amostra é retornada por valor, o que significa que esses arrays custosos são copiados de um processo para outro. Uma solução é usar um armazenamento de objetos. O Ray oferece um método conveniente para isso: cada amostra carrega uma referência de objeto, que é apenas um ponteiro para os dados. E então o

11:36 driver mantém apenas o ponteiro. O tensor nunca precisa ir e voltar até o nó principal. Os ranks de paralelismo de dados do Megatron podem simplesmente chamar diretamente o armazenamento do Ray para buscar os dados conforme necessário. Portanto, ao escolher uma solução de paralelismo e concorrência, é preciso sempre considerar os custos adicionais de comunicação. Planeje isso com cuidado. Certo. E conseguimos eliminar uma chamada extra do tipo put object que nossa solução de pré-carregamento fazia e com a qual não precisávamos lidar ao usar o armazenamento do Ray.

12:12 Podemos ver que nossos custos de comunicação diminuem, assim como a proporção do tempo gasto esperando. Recapitulando: começamos com uma linha de base serial. Adicionamos concorrência para carregar e processar os dados. Adicionamos pré-carregamento para nos anteciparmos às solicitações de treinamento e usamos o armazenamento do Ray para minimizar a transferência de dados. Reduzimos a proporção de tempo de espera de 85% na linha de base para 20%. E dá para ver que o padrão em dente de serra

12:42 que aparece no gráfico superior foi atenuado em grande medida, o que significa que a utilização das GPUs aumentou. Ao escalar — agora que temos um bom pipeline de treinamento multimodal —, vamos tentar ampliá-lo para uma execução de produção. Ao fazer isso, observamos que a proporção do tempo de espera no mesmo pipeline aumentou, embora só tivéssemos escalado alguns números. Simplesmente multiplicamos por quatro os ranks de paralelismo de dados, adicionamos cerca de 200 fluxos de dados e aumentamos o tamanho do lote em quatro vezes.

13:14 Acontece que, como o gerador de dados e o conjunto de workers estão colocados no mesmo nó para acelerar a transferência de dados, a placa de rede desse nó se torna um gargalo. Isso significa que, ao concentrarmos todo esse trabalho de processamento em nós específicos, a taxa de transferência de rede deles não consegue acompanhar a demanda do pipeline de treinamento.

13:45 Nesse ponto, investigamos um pouco os frameworks que estamos usando e vimos que, por padrão, o Ray prefere agendar os workers nos mesmos nós e colocá-los juntos. Isso significa que uma máquina pode virar um gargalo antes de começarmos a usar outra. Então consideramos a possibilidade de desagregar todo esse processo, distribuindo os workers por vários nós. Assim, eles podem fazer o próprio processamento de imagens

14:18 e armazená-las para que, de forma preguiçosa e just-in-time, os ranks de paralelismo de dados as recuperem e usem. Então, esses são os dois parâmetros, ou flags, que ajustamos: um é a recuperação sem cópia e o outro é um mecanismo de agendamento distribuído. Nenhum deles ajudou no experimento em pequena escala, mas, juntos, tiveram um efeito sinérgico no experimento em grande escala. Então, uma configuração pode

14:51 ter um desempenho ruim em pequena escala, mas muito bom em escala maior. Portanto, é sempre preciso pensar em como fazer ablações ao longo dos diferentes níveis de escala. Essas foram as duas flags que acabaram fazendo uma enorme diferença na velocidade do nosso pipeline de processamento de dados. E este é o pipeline final que desenvolvemos: um sistema concorrente, eager e ciente do padrão de acesso aos dados, que compartilha o treinamento entre cargas e mantém a GPU ativa, proporcionando fluidez e estabilidade ao treinamento.

15:26 No geral, conseguimos um aumento adicional de 50% no throughput em escala usando apenas essas duas flags. É o mesmo parâmetro, mas com resultados opostos, dependendo do gargalo. Essas são as lições que tiramos desta palestra, se não houver mais nada. As GPUs são caras, mas, às vezes, as maiores batalhas acontecem fora da GPU. Então, pergunte primeiro: você está limitado pela CPU ou pela GPU? Está limitado por E/S ou por computação? Resolver um gargalo muitas vezes revela outro. No nosso caso, primeiro foi o throughput, depois

15:58 passou a ser o tempo de espera, depois a memória e, por fim, o custo de recuperação. As técnicas clássicas não são tudo. Paralelismo e pré-carregamento importam, e acho que há bastante conteúdo sobre isso. Mas também importam a semântica dos frameworks que você está usando, as restrições com que está lidando e o fato de que a escala sempre muda a resposta. Parâmetros como recuperação sem cópia e distribuição dos workers só se tornam importantes em escalas maiores e podem não produzir melhorias evidentes em escalas menores.

16:33 No geral, traçar o perfil da linha de base ajuda a fazer diferença. Obrigado a todos por assistirem à palestra. Se tiverem alguma pergunta, será um prazer responder. Obrigado.

Transcrição original

[00:01] [music] Hello everyone. Thank you for joining me here. Um yeah, today I’m going to talk about ray actors, vision tokens, and the gill. um engineering a multimodal SFT data pipeline that keeps the GPUs busy. So there’s a lot of talk about you know how

[00:35] to optimize training paradigms. Um but I want to give a different perspective on the whole ecosystem of what it takes to get a model trained. Um and so the question I’m going to be asking here is does your training loop really depend on the GPUs right now? Is it blocked on the GPUs? And I think we’re going to learn that we’re going to see in this experiment that the hardest part of optimizing multimodal train throughput is not necessarily kernel optimization. It is making sure that the G GPUs are constantly fed with data.

[01:09] So let’s go over the overview of what this topic what this talk will be about. Anyway, so we will establish a baseline multimodal training pipeline um and then we will encounter bottlenecks and then address them one at a time. So loosely they are one concurrency how fast can you process one image before feeding it to the model two pre-fetching how can we make sure that we have data on demand three um how can we minimize data copies across different processes and four f

[01:43] additional framework optimizations um so just tweaking your bells and whistles um yeah so let’s set up our premise We all know that GPUs incur expensive costs and are generally want to maximize their utilization while we have them at hand. Multimodal processing is different than text only um because there’s a lot of CPU inensive operations that have to happen before you actually feed the model um the data. So whether that is um you know loading resizing or normalizing

[02:20] images using like a library like pill for audio you do a lot of loading and cropping you can use like numpy arrays and then videos are some sort of combination of both um and then it’s also true that usually your training clusters have a lot more CPUs than GPUs at hand yet the GPU is what is starving in this pipeline. So today we’re going to introduce this equation called like weight time ratio which is generally I formulated as like the time it takes to prep your data and transfer it to the GPUs. Um and then we divide all of that

[02:56] by the time it takes to do both of those and the train as well. So basically the denominator is like your total pipeline. Um and we’re going to see that during our run we will go from a utilization of 15% to utilization of 90%. So our baseline has like 85% of its spends 85% of its time um backed on data acquisition whereas after we do these optimizations less than 10 and so for

[03:27] the scope of this talk GPUs are not the goal. It’s making sure that the data is never the bottleneck is the goal. And so here’s a little bit of background as to like what our stack will look like. um we kind of take like a Quen 3VL based model um kind of try to train it on the paradigms of chat VQA and like multi-term VLM uh conversations. So to recap before the GPUs do any math um we

[03:59] have to load the data we have to transform the data and then we have to we have to transport it to the GPUs. Okay, this is a this is a overview of what our data actually looks like. So, we will be using JSON L files. Um, and then images are linked relatively uh to where the example lives. And then we’re going to assume that our data is just stored in S3. Um, in this example, for example, it seems like that image 33471.jpeg

[04:32] seems to be a picture of like a blue bus. Um, yeah. So this is our system layout. On the left we have a data generator which is responsible for the data loading, data transforming and data transporting. Um and the that’s like labeled as like fan out to the mega megatron data parallel ranks. Um so the parallel ranks receive partitions of the data as decided by the coordinator. Um it collates it performs forward pass and backwards pass on this data. So, Megatron takes care of optimizing the

[05:05] model weights across the DP ranks. And we’re going to assume that it’s like a it’s kind of it is kind of working as a black box. Um, so let’s start with the naive baseline path. So for each example, we fetch the images from object store. Um, we kind of like we tokenize it, we resize it, we write resize the images, normalize them, and then feed them to the model. Um yeah and then when we do that

[05:37] we can see that there’s serial operations in the red which is namely getting all the images one by one and then processing them one by one. So if your if your trainer is asking for a batch size of like 256 then you have to serially process all these images and then we can see that our train time is somewhere between 15 to 20 seconds. uh however the data the time spent to actually fetch the data is more than a 100 seconds. So that’s why we result in only 15% utilization of our GPUs. Um and then you can even see a profiling flame

[06:12] graph that we spend most of the time doing S3 loads and then you can see a long a little tail at the end about like loading the actual image on the pill uh library. So okay, solution one, let’s try to add some concurrency to this. Instead of those red singular uh you know serial processes, uh let’s try to parallelize it by kind of having like a worker pool. So when you request 256 examples, maybe let’s just process all of those at once. Um yeah, so for example, for example, we

[06:47] have two costs. One is fetching the images, two is processing. So that’s like decoding, vision pre-processing, and then tokenizing your text. Um, samples are very are independent. So this just means it’s like embarrassingly parallelizable. Um, and so yeah, you have a design decisions here. What forms of concurrency do you want? So there’s asynchronous programming, uh, which is suitable for like IO work such as making SQL queries or like performing network requests. There’s threads uh which you

[07:20] can use for like CPU intensive uh workloads that relinquish the Python gill like the global interpreter lock. Um so this it’s not straightforward what operations actually release the lock. So you have to do a little more of your investigation but common uses are like using hugging face transformers tokenizers using the pill transforms API or like some pandas operations processes um are more heavyweight but also provide total isolation between your workload so you can leverage multiple cores on your

[07:52] machine as you finish your work. Um, yeah, doing a deep dive on the right solution would warrant its own talk, but like there’s plenty of online resources that can help you decide for your use case. What’s important for our solution is that we need to match our bottleneck. Not simply going with a lever that’s worked for us before. Um, so we’ll choose async IO for fetching images here and then use ray actors as our processors uh as our process multipprocessing library. Um, ray actors are essentially the same as like general

[08:23] processes, but they offer us APIs to communicate across workers that just make it so that we don’t have to reinvent the wheel. So, yeah, your process is only as slow as your slowest link. Um, and we can see that our weight time ratio has gone down. Um, on the second example, which is in the pink, uh, you can see that the time spent fetching data per batch has gone basically, yeah, around like around like 40 40 seconds. Um, so this is this gives us a 25% speed up. However, you yeah that’s the as you

[08:56] can see the there’s a huge gap that starts generating especially as we’re going through many many steps. Um, so can we somehow limit our bottleneck of waiting on the trainer to request before provisioning our worker pool to fetch the data? So can we just like preemptively fetch our data? Um so we came up with solution two which was that we have a background thread that fills up a queue ahead of time and it’s always stays ahead of the consumer

[09:27] which means that you your producer is producing examples at a multiplier of your tra of your consumer and so the trainer is never bottlenecked on actually requesting the data. When it requests data it’s it is available on demand. There’s a caveat to prefetching. It makes your model and data loader checkpointing and resumption like a little intricate of a process and it can also suffer from a cold start period in which your queue is actually getting filled up uh so it can be ready for training. Um but uh if you work

[10:01] around those those details I think pre-fetching is a is a great idea. Um yeah, so I think making sure your producer is always ahead of your consumer will be will pay great dividends for you here. Um unsurprisingly, we achieved close to 0 seconds of data fetching and processing per batch because we’ve already fetched data and processed it way before the trainer even needs it. Um and so we’ve effectively solved both the data load and the transform bottlenecks. So after these two mitigations, we

[10:35] thought we were done. throughput up GPUs fed however we observe that the data transport costs still linger in fact they’ve gone up so every sample transfers a large amount of like image arrays across process boundaries multiple times every sample is pickled on the worker copied to the coordinator and then repickled onto the trainer per sample per step we have a mem copy that nobody really asked for so what h how this manifest is that your

[11:06] data generator can often oo. Um, and so multimodal tensors are huge. They’re megabytes per sample. Um, and so by default, every sample is returned by value, which means this that these expensive arrays are copied over from process to process. So a fix for this, you can come up with an object store. Um, this is a nice convenience method of that array provides you. Each sample carries a object reference which is just a pointer to the data. Um and then the

[11:36] driver just holds the pointer. The tensor never round trips to the head node. Um and the Megatron data parallel ranks they can just directly call the race store to fetch the data as they need. So with every choice of your parallelism and concurrency solution, you always have to think about the commu added communication costs. So plan wisely around that. Okay. Yeah. And we were able to eliminate an extraneous like put object uh kind of call that our prefetcher solution had that we did not

[12:12] have to deal with by using the ray store. And we can see that our communication costs go down and the ratio of the time spent waiting also goes down. So to recap, we took a serial baseline. We added concurrency for data loading and processing. We added pre-fetching to stay ahead of training requests and we added the race store to minimize data transfer. We went from a baseline weight ratio of 85% to 20%. And you can see that the saw to uh pattern

[12:42] that you see in the top graph is alleviated uh to a large extent which means that our utilization of the GPUs has gone up. So when we’re scaling up and now that we have like a good multimodal training pipeline, let’s try to scale it up to like a production run. While doing that, we observed that our weight time ratio for the same pipeline went up when all we did was scale some numbers. So we just added four times the data parallel rings. We added like 200 data streams and then batch size is x

[13:14] by4. Um so it turns out that since the data generator and the worker pool are colllocated for fast data transfer what happens is that there’s a bottleneck on the network interface card of that of that node. Uh which means that while we’re packing all of this working and processing into specific nodes that their network throughput is not able to keep up with the demand of the training pipeline.

[13:45] Um so at this point we did a little bit of digging into the frameworks that we are using and we saw that by default Ray schedules workers on this it prefers to schedule workers on the same nodes and colllocate workers which means that it’s easy to bottleneck one machine before you move to another machine. So what we did was consider a possibility where we disagregate this whole process where workers are spread across multiple nodes. Um so they can do their own image

[14:18] processing and storing and in a lazy just in time manner are retrieved and used by our DP rings. So yeah to re and then so these are the two parameter two flags that we kind of tuned one is a zero copy retrieval and second is a spread scheduling mechanism. Um neither of these actually helped in a small scale experiment but together they had um a synergistic effect in a large scale experiment. Um yeah so config can lose a config can

[14:51] do poorly at a small scale but can do really good at a larger scale. So you should always think about how to ablate across your your scaling ladders. And these were the two flags that ended up making a huge difference uh in how how fast our data processing pipeline was. And so now this is the final pipeline that we’ve accomplished a concurrent eager um and a data access pattern aware system that shares the training across loads and enables that GAPU stay warm while providing smoothness and stability in training.

[15:26] Overall we achieved a 50% additional throughput at scale just from these two flags. It’s the same knob opposite verdict set by the bottleneck. So these are the learnings we get from this talk if nothing else. Um GPUs are expensive but sometimes the biggest battles happen outside of the GPU. So first ask yourself are you CPU bound or GPU bound? Are you IO bound or computebound? And then solving one bottleneck often exposes another. So first it was throughput for us then it

[15:58] became weight time then it became memory then it became retrieval cost. Um, classic moves are not the end. You know, parallelism and prefetching matter and I think a lot of content has like covered this. Uh, but so do the semantics of the frameworks that you’re using and the and the constraints that you’re dealing with and scale always changes the answer. Tuning parameters like zero copy waste um and spreading workers become important only at larger scales and they may not manifest as obvious improvements at smaller levels.

[16:33] Overall, profiling your baseline helps you move the needle. Um, yeah, thank you everyone for attending the talk. Uh, if you have any questions, I’m happy to answer them. Thanks.