https://outerbounds.com/ logo
Join Slack
Powered by
Could anyone offer some guidance/resources on tool...
# ask-metaflow
l
Could anyone offer some guidance/resources on tooling for training foundation models? I just joined a lab that trains foundation LLMs (ttt, sts, tts, stt, etc.). As someone who has been doing 95% classical ML/DS this whole time, it's a dizzying learning curve. I'll add a bit about our setup, and a few questions in this thread...
Our setup. It's pretty vanilla. • GPUs: Lambda Labs, but we bias toward using our on-prem B200s • Frameworks: DeepSpeed, HF Trainer, FSDP, Megatron-LM, PyTorch • SaaS: Weights & Biases, Huggingface We're wondering if these things could speed us up: 1. Does Kubernetes/Metaflow add value to this use case? Or does containerization/framework overhead somehow outweigh just, SSH-ing and executing 
pixi run my_distributed_training_script.py
? 2. Is SLURM mutually exclusive or complementary with Metaflow/Kubernetes?
v
sounds fun! tl;dr - yeah, pretraining even a medium-scale LLM is a tad bit more complicated than ssh'ing to a box and running a Python script 😉 since your lab trains foundation models for living, I am sure you/they know much of this already, but I am dumping here a bunch of references if others are lurking in the thread • HF Ultrascale Playbook • NVIDIA NeMo toolkit • LLM360 (older but still useful as a reference) • Large-scale ML infra at Meta
👍 1
How Metaflow can help: • Data preprocessing workflows of various kinds - a huge part of pretraining is about data • Orchestrating distributed training with 
@deepspeed
, 
@metaflow_ray
, 
@torchrun
etc. • Checkpointing with 
@checkpoint
• Unified compute over 
@kubernetes
, 
@slurm
etc • Observability, logging, experiment tracking etc
re: Slurm vs. Kubernetes: Slurm is an HPC job scheduler whereas Kubernetes is a container orchestrator which you can turn into an HPC cluster with job scheduler with (a non-trivial amount of) effort Or you can combine the two Since the systems are quite different, serving slightly different use cases, there's no straightforward answer which one is better. Much depends on the technical preferences / capabilities of the team too
h
@victorious-lawyer-58417 Not to hijack this thread, but what do you think would need to change to make metaflow become a goto for training foundation models? I think metaflow’s unopinionated stance is somewhat hurting here. Do you think such functionality can live in an extension? Or is the core training loop aspect not what metaflow targets?
v
in my mind it's the good old "hidden technical debt of ML infra" picture again - the core training loop is a small part of a foundation model lab I think Metaflow can be helpful in the core training loop (especially with small/medium-size models - giant models are a very special beast in many ways), but I think Metaflow has even a broader scope if you think the full extent of everything that a lab training foundation models needs to do
but to answer your question directly: I am bullish on Metaflow orchestration distributed compute `@parallel`/`@torchun`/`@metaflow_ray` - much more can be done there. I'd love to test those extensions e.g. with NVIDIA's Kai scheduler, which comes in very handy with large GPU clusters (presumably)
another fun topic is post-training with GRPO and other RL techniques - "reward function as a service", like what

Eddie talked about here▾

- we have customers doing it with Metaflow today but the field is evolving quickly so tons of opportunities for improvments
👀 1
g
this is very interesting!
excited bird 1
13 Views
Open in Slack
PreviousNext
Powered by