Llama training with FP8

[pbelevich]

No description provided.

[KeitaW]

Overall, this looks a great feature addition to 10.FSDP. Do you think TE support can be added in the test case instead of creating new one?

[KeitaW]

slurm-*.out should already be excluded by https://github.com/aws-samples/awsome-distributed-training/blob/main/.gitignore

[KeitaW]

You may want to refer original code https://github.com/NVIDIA/TransformerEngine/blob/16a469df6bbc77e1c32e48e8e5fd3082dbc2d18e/docs/examples/te_llama/te_llama.py

[pbelevich]

@KeitaW thanks for the review! I was thinking about adding FP8 support to FSDP example, but there are two aspects why I decided to create a separate example for this:

1. Transformer Engine requires Nvidia's container to run(or as alternative relatively complicated process of building from source with CUDA headers, CUDNN etc). And I don't want to complicate FSDP example with it.
2. This example is bound to Llama model(taken from TE examples), but FSDP example supports multiple models that I don't want to rewrite with FP8.

So, in terms of importance this example is about LLama with FP8. FSDP training here is just kind of scaffolding.

[sbhavani]

@pbelevich FYI the AWS DLC for PyTorch also includes TE