stas00 / ml-engineering

Machine Learning Engineering Open Book

Machine Learning Engineering Open Book

This is an open collection of methodologies, tools and step by step instructions to help with successful training of large language models and multi-modal models.

This is a technical material suitable for LLM/VLM training engineers and operators. That is the content here contains lots of scripts and copy-n-paste commands to enable you to quickly address your needs.

This repo is an ongoing brain dump of my experiences training Large Language Models (LLM) (and VLMs); a lot of the know-how I acquired while training the open-source BLOOM-176B model in 2022 and IDEFICS-80B multi-modal model in 2023. Currently, I'm working on developing/training open-source Retrieval Augmented Generation (RAG) models at Contextual.AI.

I've been compiling this information mostly for myself so that I could quickly find solutions I have already researched in the past and which have worked, but as usual I'm happy to share these with the wider ML community.

Table of Contents

My apologies if the layout is a bit unstable while I'm writing new chapters and gradually re-organizing the content to be more intuitive.

Part 1. Insights

1. The AI Battlefield Engineering - what you need to know in order to succeed

Part 2. Hardware

1. Compute - accelerators, CPUs, CPU memory.

2. Storage - local, distributed and shared file systems.

3. Network - intra- and inter-node networking.

Part 3. Orchestration

1. SLURM - the main orchestration environment

Part 4. Training

1. Training - model training related guides

Part 5. Development

1. Debugging and Troubleshooting - how to debug easy and difficult issues

2. And more debugging

3. Testing - numerous tips and tools to make test writing enjoyable

Part 6. Miscellaneous

1. Resources - LLM/VLM chronicles

Updates

I announce any significant updates on my twitter channel https://twitter.com/StasBekman.

PDF version

Download the PDF version of the book.

I will try to rebuild it once a week or so, but if you want the latest, the instructions for building are here.

Thanks to HuggingFace for giving me permission to host my book's PDF at the HF hub.

Discussions

If you want to discuss something related to ML engineering this repo has the community discussions available - so please don't hesitate to share your experience or start a new discussion about something you're passionate about.

Key comparison tables

High end accelerators:

• Theoretical accelerator TFLOPS
• Accelerator memory size and speed

Networks:

• Theoretical inter-node speed
• Theoretical intra-node speed

Shortcuts

Things that you are likely to need to find quickly and often.

Tools:

• all_reduce_bench.py - a much easier way to benchmark network throughput than nccl-tests.
• torch-distributed-gpu-test.py - a tool to quickly test your inter-node connectivity

Guides:

• debugging pytorch applications - quick copy-n-paste solutions to resolve hanging or breaking pytorch applications
• slurm for users - a slurm cheatsheet and tricks
• make tiny models/datasets/tokenizers
• LLM/VLM chronicles collection

Gratitude

None of this would have been possible without me being entrusted with doing the specific LLM/VLM trainings I have learned the initial know-how from. This is a privilege that only a few enjoy due to the prohibitively expensive cost of renting huge ML compute clusters. So hopefully the rest of the ML community will vicariously learn from these notes.

Special thanks go to Thom Wolf who proposed that I lead the BLOOM-176B training back when I didn't know anything about large scale training. This was the project that catapulted me into the intense learning process. And, of course, HuggingFace for giving me the opportunity to work full time on BLOOM-176B and later on IDEFICS-80B trainings.

Currently, I'm continue expanding my knowledge and experience while training models and building systems at Contextual.AI and I'm grateful to that opportunity.

I'd also like to say thanks to the numerous contributors who have been making this text awesome and error-free.

Contributing

If you found a bug, typo or would like to propose an improvement please don't hesitate to open an Issue or contribute a PR.

License

The content of this site is distributed under Attribution-ShareAlike 4.0 International.

My repositories map

✔ Machine Learning: ML Engineering Open Book | ML ways | Porting

✔ Guides: The Art of Debugging

✔ Applications: ipyexperiments

✔ Tools and Cheatsheets: bash | conda | git | jupyter-notebook | make | python | tensorboard | unix

ml-engineering's Issues

ML

GPU requirements and cost estimation.

Hey, I just came towards this repo and I highly appreciate the content that you put up here. I also have started to do fine-tuning and moving towards large scale fine-tuning. One of the key things before getting started is estimating the number of GPUs required and the estimated cost for doing the fine-tuning. Now there are so many variables like:

1. What GPU you use.
2. How many GPUs we are using.
3. Is this full-finetuning or using LoRA or QLoRA, etc. etc.

The worst part is when we take some amount of GPU but that is not fully utilized. So, I want to understand whether there is any empirical way of doing things or not. Yes, it might not be precise. However, it would be really useful.

If you want, I can dump some of my findings here.

convert markdown to pdf

Hi, I had a good read of this book! Wondering if we can convert the markdown files to PDF so that we can print it out to read. I would like to submit a pr for that! Let me know if you are interested!

Improve folder structure

Would it be possible to improve the folder structure so it somewhat matches your table of contents?

eg:
├── Part 1
│ ├── Topic 1
│ ├── Topic 2
├── Part 2
│ ├── Topic 3
│ ├── Topic 4

Convert to bfloat16 failing

I tried your shell script torch-checkpoint-convert-to-bf16 on a quantized model with no .bin files, but it didn't quite work. Here are my results:

(gptq) ~/gptq/models/TheBloke_StableBeluga2-70B-GPTQ$ ll
total 34506488
-rw-rw-r-- 1 matt matt        7020 Jul 28 12:44 LICENSE.txt
-rw-rw-r-- 1 matt matt       15560 Jul 28 12:44 README.md
-rw-rw-r-- 1 matt matt        4766 Jul 28 12:44 USE_POLICY.md
-rw-rw-r-- 1 matt matt         679 Jul 28 12:44 config.json
-rw-rw-r-- 1 matt matt         137 Jul 28 12:44 generation_config.json
-rw-rw-r-- 1 matt matt 35332232456 Jul 28 12:49 gptq_model-4bit--1g.safetensors
-rw-rw-r-- 1 matt matt         301 Jul 28 12:44 huggingface-metadata.txt
-rw-rw-r-- 1 matt matt         183 Jul 28 12:49 quantize_config.json
-rw-rw-r-- 1 matt matt         411 Jul 28 12:49 special_tokens_map.json
-rw-rw-r-- 1 matt matt     1842764 Jul 28 12:49 tokenizer.json
-rw-rw-r-- 1 matt matt      499723 Jul 28 12:49 tokenizer.model
-rw-rw-r-- 1 matt matt         649 Jul 28 12:49 tokenizer_config.json

(gptq) ~/gptq/models/TheBloke_StableBeluga2-70B-GPTQ$ bash ~/mr/utilz/torch-checkpoint-convert-to-bf16 
creating a new checkpoint under dir bf16
converting *bin torch files
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "<string>", line 1, in <listcomp>
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 791, in load
    with _open_file_like(f, 'rb') as opened_file:
         ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 271, in _open_file_like
    return _open_file(name_or_buffer, mode)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 252, in __init__
    super().__init__(open(name, mode))
                     ^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: '*bin'
converting *safetensors files
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "<string>", line 1, in <listcomp>
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 791, in load
    with _open_file_like(f, 'rb') as opened_file:
         ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 271, in _open_file_like
    return _open_file(name_or_buffer, mode)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/matt/miniconda3/envs/gptq/lib/python3.11/site-packages/torch/serialization.py", line 252, in __init__
    super().__init__(open(name, mode))
                     ^^^^^^^^^^^^^^^^
FileNotFoundError: [Errno 2] No such file or directory: '*bin'
the dir bf16 now contains a copy of the original checkpoint with bf16 weights

(gptq) ~/gptq/models/TheBloke_StableBeluga2-70B-GPTQ$ ll bf16/
total 2312
-rw-rw-r-- 1 matt matt     679 Aug 15 18:27 config.json
-rw-rw-r-- 1 matt matt     137 Aug 15 18:27 generation_config.json
-rw-rw-r-- 1 matt matt     183 Aug 15 18:27 quantize_config.json
-rw-rw-r-- 1 matt matt     411 Aug 15 18:27 special_tokens_map.json
-rw-rw-r-- 1 matt matt 1842764 Aug 15 18:27 tokenizer.json
-rw-rw-r-- 1 matt matt  499723 Aug 15 18:27 tokenizer.model
-rw-rw-r-- 1 matt matt     649 Aug 15 18:27 tokenizer_config.json

What are your thoughts? Thank you!

Daisy chain batch jobs

The job array works well to queue up multiple jobs:

https://github.com/stas00/ml-engineering/tree/master/fault-tolerance#queue-up-multiple-training-jobs

Another common approach is to "daisy chain" jobs by having the job script submit another job that is dependent on itself. For example, in train.slurm you'd have a line like:

# when train.slurm executes, have it submit another job dependent on itself
sbatch --dependency=$SLURM_JOBID train.slurm

This is usually done near the top of the script, before the command that actually launches the run.

One might also pair that with some logic to stop the chaining when the job is done. For example, the application or the user might touch a "run.done" file when it completes. Then the script can check for that file.

# exit right away if "run.done" file is detected
if [ -f run.done ] ; then
  exit 0
fi

# otherwise chain up another job
sbatch --dependency=$SLURM_JOBID train.slurm

# then launch the run
<<launch run>>>

Additionally, one could check for the "run.done" file after the run and attempt to cancel any already daisy-chained job.

I don't have a list of pros/cons vs the job array, but it's one more method I see in practice.

Missing `hparams` section

Hi Stas, thank you for making these notes public! They are an invaluable resource.

I noticed that the hparams folder, linked here in the readme, seems to be missing from the repository. Was this intentional?

discuss the solutions to Not fully recovering spikes

Not fully recovering spikes are quite common in model training at scale. Like what I have:

It's worth discussing the causes and potential solutions to that.
My general guesses are:

• some of the toxic data points at certain iterations
• The learning rate.

Welcome everyone to talk about your experiences dealing with the spikes...

Reference:

https://github.com/stas00/ml-engineering/blob/master/training/instabilities/training-loss-patterns.md

Parallel training hangs

Hi, I saw your toolbox link in a Huggingface issue and gave it a try. My four new GPUs hang when trying to fine tune a transformer, and they appear to do the same thing when running your torch-distributed-gpu-test.py tool, too. However, I'm not sure what the expected outcome is here. I should point out that I can fine tune a transformer with just a single GPU. I'm using Python 3.9.7, Transformers 4.17.0, PyTorch 1.11.0+cu113, NCCL 2.12.7 for CUDA 11.6, and four Nvidia A6000 GPUs.

$ NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 4 --nnodes 1 torch-distributed-gpu-test.py
WARNING:__main__:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
DeepWhite:21288:21288 [0] NCCL INFO Bootstrap : Using enp67s0:192.168.50.21<0>
DeepWhite:21288:21288 [0] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
DeepWhite:21288:21288 [0] NCCL INFO NET/IB : No device found.
DeepWhite:21288:21288 [0] NCCL INFO NET/Socket : Using [0]enp67s0:192.168.50.21<0>
DeepWhite:21288:21288 [0] NCCL INFO Using network Socket
NCCL version 2.10.3+cuda11.3
DeepWhite:21290:21290 [2] NCCL INFO Bootstrap : Using enp67s0:192.168.50.21<0>
DeepWhite:21289:21289 [1] NCCL INFO Bootstrap : Using enp67s0:192.168.50.21<0>
DeepWhite:21290:21290 [2] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
DeepWhite:21289:21289 [1] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
DeepWhite:21289:21289 [1] NCCL INFO NET/IB : No device found.
DeepWhite:21290:21290 [2] NCCL INFO NET/IB : No device found.
DeepWhite:21290:21290 [2] NCCL INFO NET/Socket : Using [0]enp67s0:192.168.50.21<0>
DeepWhite:21289:21289 [1] NCCL INFO NET/Socket : Using [0]enp67s0:192.168.50.21<0>
DeepWhite:21290:21290 [2] NCCL INFO Using network Socket
DeepWhite:21289:21289 [1] NCCL INFO Using network Socket
DeepWhite:21291:21291 [3] NCCL INFO Bootstrap : Using enp67s0:192.168.50.21<0>
DeepWhite:21291:21291 [3] NCCL INFO NET/Plugin : No plugin found (libnccl-net.so), using internal implementation
DeepWhite:21291:21291 [3] NCCL INFO NET/IB : No device found.
DeepWhite:21291:21291 [3] NCCL INFO NET/Socket : Using [0]enp67s0:192.168.50.21<0>
DeepWhite:21291:21291 [3] NCCL INFO Using network Socket
DeepWhite:21289:21327 [1] NCCL INFO Trees [0] 2/-1/-1->1->0 [1] 2/-1/-1->1->0
DeepWhite:21291:21329 [3] NCCL INFO Trees [0] -1/-1/-1->3->2 [1] -1/-1/-1->3->2
DeepWhite:21290:21328 [2] NCCL INFO Trees [0] 3/-1/-1->2->1 [1] 3/-1/-1->2->1
DeepWhite:21288:21326 [0] NCCL INFO Channel 00/02 :    0   1   2   3
DeepWhite:21288:21326 [0] NCCL INFO Channel 01/02 :    0   1   2   3
DeepWhite:21288:21326 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] 1/-1/-1->0->-1
DeepWhite:21288:21326 [0] NCCL INFO Setting affinity for GPU 0 to ffffffff,ffffffff
DeepWhite:21291:21329 [3] NCCL INFO Channel 00 : 3[4a000] -> 0[3000] via P2P/IPC
DeepWhite:21289:21327 [1] NCCL INFO Channel 00 : 1[21000] -> 2[49000] via P2P/IPC
DeepWhite:21288:21326 [0] NCCL INFO Channel 00 : 0[3000] -> 1[21000] via P2P/IPC
DeepWhite:21290:21328 [2] NCCL INFO Channel 00 : 2[49000] -> 3[4a000] via P2P/IPC
DeepWhite:21291:21329 [3] NCCL INFO Channel 01 : 3[4a000] -> 0[3000] via P2P/IPC
DeepWhite:21289:21327 [1] NCCL INFO Channel 01 : 1[21000] -> 2[49000] via P2P/IPC
DeepWhite:21288:21326 [0] NCCL INFO Channel 01 : 0[3000] -> 1[21000] via P2P/IPC
DeepWhite:21290:21328 [2] NCCL INFO Channel 01 : 2[49000] -> 3[4a000] via P2P/IPC
DeepWhite:21288:21326 [0] NCCL INFO Connected all rings
DeepWhite:21291:21329 [3] NCCL INFO Connected all rings
DeepWhite:21290:21328 [2] NCCL INFO Connected all rings
DeepWhite:21291:21329 [3] NCCL INFO Channel 00 : 3[4a000] -> 2[49000] via P2P/IPC
DeepWhite:21289:21327 [1] NCCL INFO Connected all rings
DeepWhite:21291:21329 [3] NCCL INFO Channel 01 : 3[4a000] -> 2[49000] via P2P/IPC
DeepWhite:21290:21328 [2] NCCL INFO Channel 00 : 2[49000] -> 1[21000] via P2P/IPC
DeepWhite:21289:21327 [1] NCCL INFO Channel 00 : 1[21000] -> 0[3000] via P2P/IPC
DeepWhite:21290:21328 [2] NCCL INFO Channel 01 : 2[49000] -> 1[21000] via P2P/IPC
DeepWhite:21289:21327 [1] NCCL INFO Channel 01 : 1[21000] -> 0[3000] via P2P/IPC
DeepWhite:21291:21329 [3] NCCL INFO Connected all trees
DeepWhite:21291:21329 [3] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 8/8/512
DeepWhite:21291:21329 [3] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
DeepWhite:21288:21326 [0] NCCL INFO Connected all trees
DeepWhite:21288:21326 [0] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 8/8/512
DeepWhite:21288:21326 [0] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
DeepWhite:21290:21328 [2] NCCL INFO Connected all trees
DeepWhite:21290:21328 [2] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 8/8/512
DeepWhite:21290:21328 [2] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
DeepWhite:21289:21327 [1] NCCL INFO Connected all trees
DeepWhite:21289:21327 [1] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 8/8/512
DeepWhite:21289:21327 [1] NCCL INFO 2 coll channels, 2 p2p channels, 2 p2p channels per peer
DeepWhite:21291:21329 [3] NCCL INFO comm 0x7f8894002fb0 rank 3 nranks 4 cudaDev 3 busId 4a000 - Init COMPLETE
DeepWhite:21289:21327 [1] NCCL INFO comm 0x7fd2c8002fb0 rank 1 nranks 4 cudaDev 1 busId 21000 - Init COMPLETE
DeepWhite:21290:21328 [2] NCCL INFO comm 0x7f7aa0002fb0 rank 2 nranks 4 cudaDev 2 busId 49000 - Init COMPLETE
DeepWhite:21288:21326 [0] NCCL INFO comm 0x7fb314002fb0 rank 0 nranks 4 cudaDev 0 busId 3000 - Init COMPLETE
DeepWhite:21288:21288 [0] NCCL INFO Launch mode Parallel

pip install -r build/requirements.txt fails due to github_md_utils

ERROR: Could not find a version that satisfies the requirement github_md_utils (from versions: none) ERROR: No matching distribution found for github_md_utils

Any ideas? Googling github_md_utils gives 0 hits: https://www.google.com/search?client=safari&rls=en&q=github_md_utils&ie=UTF-8&oe=UTF-8

Thanks!

Minor Typo in emulate multi node

Please add import torch.distributed as dist in the section
4. Create a test script to check both GPUs are used.

I tried creating a branch and raise a PR but could not do so due to

$:~/ml-engineering-stats00$ git push origin asaha/emulate-multi-node
remote: Permission to stas00/ml-engineering.git denied to anindya-saha.
fatal: unable to access 'https://github.com/stas00/ml-engineering.git/': The requested URL returned error: 403

Quarto Site

Hey @stas00 ,

I added a comment on my previous PR:

@stas00 if you're curious, I've done some re-org and tried rendering everything with Quarto in my fork

You can see it online at:

https://saforem2.github.io/ml-engineering/

This is mostly just a first (rough) pass (and copied + modified from existing Quarto sites I've made),
but I made an effort to fix / change things where appropriate (e.g. links, etc.)

disclaimer: I mostly just did this out of personal curiosity / testing, but thought I'd share incase you're interested 🤷🏻‍♂️

I know that I, personally have a hard time keeping track of GitHub comments ( especially on closed PRs 😂 ),
so I figured it probably made more sense to move the discussion to a separate issue (I hope that's okay)

Clarification for gradient memory in mixed precision training

Hi @stas00 ,

I was going over your excellent notes on model memory usage and I noticed the following:

Gradients

4 bytes * number of parameters for either fp32 or mixed precision training (gradients are almost always kept in fp32).
2 bytes * number of parameters for more recent works where half-precision is used

I was trying to see how you got the 4 bytes number for mixed precision! From what I've seen in the mixed precision paper and the ZeRO paper, the gradients are kept in half-precision and should take 2 bytes / parameter. However, during the optimization step, the gradients are converted to FP32 - but I assumed this happens on the fly and thus you don't have to write the full FP32 gradient tensor into HBM.

When I was benchmarking memory usage using some of your older BERT-based benchmarking code, the memory usage with FP16 training works out to be 11GB on my machine. If I use the 4 bytes number for gradients instead of 2 bytes, the estimated memory will end up coming closer to the 11GB amount! I was initially assuming this discrepancy is some overhead with temporary buffers (instead of number of bytes for gradient elements to be 4). I'm just curious to know more about this 4 byte number, and even better if you're aware of the implementation details for this in Apex or DeepSpeed. Thank you!

Conflicting opinions about streaming data from cloud storage?

(1) and (2) seem to express different opinions:

1. In the "3 Machine Learning IO needs" section, one of the bullet points under "Incoming suggestions from Ross Wightman to integrate" is "Note that once your datasets are optimally friendly for a large, distributed network filesystem, they can usually just be streamed from bucket storage in cloud systems that have that option. So better to move them off the network
   filesystem in that case."

2. The section "Local storage beats cloud storage" starts with "While cloud storage is cheaper the whole idea of fetching and processing your training data stream dynamically at training time is very problematic with a huge number of issues around it...It’s so much better to have enough disk space locally for data loading."

What am I missing?

TPU v4 has 1,200GB/s of mem bandwidth and not 2,400, right?

In the README.md:

On Google's documentation: https://cloud.google.com/tpu/docs/v4