Can we decouple the data preprocssing/tokenization step from the fine-tuning phase?

[Electronic-Waste]

Hi @joecummings, after we discussed in kubeflow/trainer#2410 and Kubeflow WG Training Call, we found that torchtune is an amazing tool for fine-tuning LLMs and decided to adopt it as our low-level runtime for Kubeflow LLM Trainer. For now, we've started the implementation based on torchtune. Thanks for your engagement in the discussion.

However, we want to decouple the data preprocessing and tokenization step from the main fine-tuning phase, so as to:

• Reduce the time for using GPUs: we will wrap torchtune into a container and request for GPU resource for it (GPU is expensive and paid according to usage time)
• Integrate the data preprocessing / tokenization step with our data initializer: Do these steps ahead of fine-tuning and offload them to CPU

We wonder if torchtune have best practice to achieve these goals. And we'll appreciate it if you could offer some precious suggestions. Thanks!

Also /cc @andreyvelich @tenzen-y @astefanutti @deepanker13 @saileshd1402 @seanlaii

[andreyvelich]

Thank you for starting this issue @Electronic-Waste!

For the context, here is the recording when we discussed how to optimize GPU cost by decoupling data processing (e.g. tokenization) step from the actual fine-tuning: https://youtu.be/ALMJTr582rQ?t=1000

If we agree on that, we can leverage the power of Kubernetes to orchestrate heterogeneous workloads and reduce un-necessary load on GPUs (at least for tabular data which doesn't require GPU for processing).

For example we can run :

# Pre-process data on ReplicatedJob with CPUs
tune initialize <dataset>

#  Fine-tune LLM on ReplicatedJob with GPUs
tune run --nnodes 10 --nproc_per_node 4 full_finetune_distributed --config llama3_1/8B_full

cc @shravan-achar @akshaychitneni @bigsur0 @franciscojavierarceo

[ebsmothers]

Hi @Electronic-Waste and @andreyvelich, @joecummings is out for this week so I can take a look in the meantime. A bit of a noob question: does this mean that the torchtune training code needs to execute purely on GPU and the relationship with CPU is a read-only one? For something like tokenization, we do it on CPU anyways so the only gain here is that we do it once and cache the results? I think dataloading is not really a major bottleneck and suspect that fully preprocessing the dataset would not really save much GPU time. (The story may be a bit different for multimodal -- if I remember correctly we spent about 5% of time on data preprocessing for text + image data, but in that case there may actually be some advantage to doing preproc on GPU.) Admittedly I'm not as familiar with some of the orchestration considerations of kubeflow so do let me know if there's something simple I'm missing here.

[andreyvelich]

Does this mean that the torchtune training code needs to execute purely on GPU and the relationship with CPU is a read-only one? For something like tokenization, we do it on CPU anyways so the only gain here is that we do it once and cache the results?

That is correct @ebsmothers.
As @Electronic-Waste showed in the diagram above, TrainJob can orchestrate Kubernetes workloads in a way that Kubernetes CPU-nodes can run data processing step, and Kubernetes GPU-nodes can run training step.
That will allow us to reduce GPU idle time when tokenizations is performed.
Additional diagrams can be viewed here: https://github.com/kubeflow/trainer/tree/master/docs/proposals/2170-kubeflow-trainer-v2#llm-fine-tuning-diagram

After tokenization is complete step we can do one of the two things:

1. Save data to disk and use ReadOnlyMany PVC to mount this disk for each GPU node.
2. Transfer the data over the network ( cc @akshaychitneni @shravan-achar).

The story may be a bit different for multimodal -- if I remember correctly we spent about 5%

I think for multi-node it is even more important, since when we perform distributed training across 100 VMs, we need to pre-process data in every VM before start fine-tuning with tune run.

but in that case there may actually be some advantage to doing preproc on GPU.)

Please can you share more about it ? When do you see advantages to do preproc on GPU over CPU ?