Add Llama 3.1 405B config

[awaelchli]

Adds the config for Llama 3.1 405B.

Did a basic test last night using 8xH100 and both TP and sequential generation.

Here are some of the outputs:

litgpt generate_tp checkpoints/meta-llama/Meta-Llama-3.1-405B --precision bf16-true --quantize bnb.nf4  --num_samples 5

What food do llamas eat? Llamas are herbivores, which means that they have a plant-based diet. They are found in the wild in South America and are a relative of the camel. Llamas are very hardy and are able to survive on a diet
Time for inference 1: 9.55 sec total, 5.24 tokens/sec                                                                                                        

What food do llamas eat? What do lizards eat in their natural environment? Lizards are a diverse group of reptiles that includes iguanas, geckos, chameleons, and many other species. They are found on every continent except Antarctica and in a
Time for inference 2: 8.12 sec total, 6.15 tokens/sec 

What food do llamas eat? Llamas are herbivores which means that they are primary consumers. They eat mostly plants. These animals will consume plants such as wheat, corn, grass, and oats, but will eat other plants that are available as well. How much does  
Time for inference 3: 8.57 sec total, 5.84 tokens/sec                                                                                                        

What food do llamas eat? Llamas are herbivores and therefore only eat plant-based foods, primarily native grasses, no grain or alfalfa. A llama’s stomach is very similar to a cow’s stomach. Both are ruminants and have a three-part stomach
Time for inference 4: 8.38 sec total, 5.97 tokens/sec                                                                                                        

What food do llamas eat? In nature, lamas eat grass and plant-like flora. In captivity, they should be fed low-protein grass, good-quality hay and llama feed. Treats for lamas include bite-sized apples, broccoli, and all-natural gummy bears  
Time for inference 5: 8.03 sec total, 6.23 tokens/sec   

litgpt generate_sequentially checkpoints/meta-llama/Meta-Llama-3.1-405B --precision bf16-true --quantize bnb.nf4  --num_samples 5

What food do llamas eat? Llamas are herbivores, which means they do not eat meat. Instead, they get to spend their days grazing on delicious plants and grasses. Most llamas are domesticated, so they are fed alfalfa and hay, usually.
Time for inference 1: 14.09 sec total, 3.55 tokens/sec                                                                                                       

What food do llamas eat? What do lizards eat in their natural environment? Lizards are popular among reptile lovers. They are often kept as pets and can be seen in some homes. If you're a big fan of lizards, you might be wondering what they eat
Time for inference 2: 13.15 sec total, 3.80 tokens/sec                                                                                                       

What food do llamas eat? Llamas are herbivores which means that they are primary consumers. They eat mostly plants. These animals will consume plants such as wheat, corn, grass, and oats, but also eat other plants that are native to the area that they live
Time for inference 3: 13.15 sec total, 3.80 tokens/sec                                                                                                       

What food do llamas eat? Llamas are herbivores and therefore only eat plant-based foods, primarily native grasses, no grain or alfalfa. A llama’s stomach is very similar to a cow’s stomach. What are the 3 things llamas eat? L
Time for inference 4: 13.15 sec total, 3.80 tokens/sec                                                                                                       

What food do llamas eat? In nature, lamas eat grass and ferns. In captivity, they should be fed grass and ferns, with hay. They can also be fed alfalfa and corn. What is a llama's home? A lama's home is the
Time for inference 5: 13.15 sec total, 3.80 tokens/sec 

[Andrei-Aksionov]

What about the main README and tutorials/download_model_weights.md?

[awaelchli]

I think updating the readme can be done last once we have all pieces in. What should I update in tutorials/download_model_weights.md?

[Andrei-Aksionov]

I think updating the readme can be done last once we have all pieces in.

Sebastian is working on 8B and 70B variants, so he adds only them to the table.
If he merges his PR after yours, he can forget to update the table accordingly.

What should I update in tutorials/download_model_weights.md?

There is a table with supported models and an example output from litgpt download list.

[rasbt]

Btw what do you think about the FP8 version: https://huggingface.co/meta-llama/Meta-Llama-3.1-405B-FP8

Could that be something we could support on H100 hardware?

[awaelchli]

We can't load these FP8 checkpoints directly (afaik). I think we have to do a bit of tinkering to load things fast.

[rasbt]

>       response_executor.shutdown(wait=False, cancel_futures=True)
E       TypeError: shutdown() got an unexpected keyword argument 'cancel_futures'

/opt/hostedtoolcache/Python/3.8.18/x64/lib/python3.8/site-packages/litserve/server.py:486: TypeError
----------------------------- Captured stderr call -----------------------------
Seed set to 123

Fetching 4 files:   0%|          | 0/4 [00:00<?, ?it/s]
Fetching 4 files:  75%|███████▌  | 3/4 [00:00<00:00, 11.47it/s]
Fetching 4 files: 100%|██████████| 4/4 [00:00<00:00, 15.28it/s]
=========================== short test summary info ============================
FAILED tests/test_serve.py::test_simple - TypeError: shutdown() got an unexpected keyword argument 'cancel_futures'

There was a new LitServe release 35 min ago. Could be related

[awaelchli]

@rasbt cancel_futures was added in Python 3.9 only. LitServe does not test with Python 3.8, and 3.8 will be deprecated soon so we should just remove the 3.8 workflow from LitGPT.

[rasbt]

Sounds reasonable, thanks