Support H100 training with FP8 in Trainer and Deepspeed
Author: michaelroyzen
Pull request: no
Labels: name: trainer

name: Feature request

name: DeepSpeed

### Feature request

Support H100 training with FP8 in Trainer and Deepspeed

### Motivation

FP8 should be much faster than FP16 on supported Hopper hardware. Particularly with Deepspeed integration @stas00

### Your contribution

Happy to help in any way that I can.

Comments:
cc @pacman100 

Any updates @pacman100 @sgugger?

The speedup is only going to show when fully training model over 6B parameters, which is why we haven't prioritized the support in the Trainer. It is baked in Accelerate though.

Thanks for the update @sgugger. I'm training 6B+ models with Trainer + DeepSpeed using an MPI launcher, so Trainer support would be helpful.

Would love your input how FP8 can be used with Trainer + DeepSpeed @stas00 

Hello, you can directly use the Accelerate Launcher with Trainer to use FP8 support out of the box.

Just do:
```
accelerate launch --mixed_precision fp8 training_script_using_trainer.py --kwarg1 value ...
```

With respect to FP8 support with DeepSPeed, can you raise an issue with the DeepSpeed team?

@michaelroyzen FP8 has been tested with all DeepSpeed ZeRO stages and is compatible with [Transformer Engine](https://github.com/NVIDIA/TransformerEngine). There's a basic FP8 / DeepSpeed test [here](https://github.com/microsoft/DeepSpeed/blob/f060407829f87da32a267a60d26d13a68dc11c61/tests/unit/runtime/half_precision/test_fp8.py). Feel free to raise an issue in Transformer Engine github if it's not working.

Hi @pacman100, it seems like accelerate does not support fp8 with [neither DeepSpeed ZeRO nor FSDP](https://github.com/huggingface/accelerate/blob/e1247de01e0733c5d21075cb6f39b2605f4be123/src/accelerate/utils/dataclasses.py#L1386). I'd appreciate your guidance here.

@pacman100, is no longer with HF.

Tagging @muellerzr who is the new maintainer of deepspeed integration in accelerate

Thanks, Stas. @muellerzr would love to hear your thoughts. There doesn't seem to be a way to get HF Trainer/Accelerate FP8 training working with DeepSpeed ZeRO or FSDP despite @sbhavani's confirmation that it should work in ZeRO.

@michaelroyzen I'll be working on enabling them this week, I'll have a report back to you by Friday at the latest :)

Thank you @muellerzr! Do you have an update? Is there anything I can help with?

> @michaelroyzen I'll be working on enabling them this week, I'll have a report back to you by Friday at the latest :)

Any update on where to find resources to enable this?

> Any update on where to find resources to enable this?

I would also like to know if it has been achieved now?

Working my way through it. Open PR is here: https://github.com/huggingface/accelerate/pull/2983 facing a few issues where it seems params aren't being trained, but looking to resolve that part today

On the accelerate side zero1 and zero2 are supported. With Zero3 I'm facing some issues with weight adjustments since we do an inplace conversion of the model parameters when converting to Zero3. Will be looking into that tommorow but if you want FSDP or Zero1/Zero2, this should be doable now if you do `accelerate launch` w/ a config that specifies mixed_precision="fp8"

(and make sure you pip install from the branch in the PR)

Just curious, does this work with replacing the linear layers, or can I use a custom linear like BitLinear?

@Codys12 we check for nn.Linear and nn.LayerNorm. You can certainly manually create/convert a TE model and pass it in though and everything should still be fine. 

Good to close this now :) 

Please install accelerate via `pip install git+https://github.com/huggingface/accelerate` until our next release.
From there do `accelerate config` to configure FP8, followed by `accelerate launch` and you can then use FP8 w/ DeepSpeed (or FSDP) with the Trainer 🤗 

Thanks for the update. I tried the latest code and examples from `accelerate` and it doesn't seem to work natively with `Trainer` (may need to manually set up the training loop). Any further updates?

> @Codys12 we check for nn.Linear and nn.LayerNorm. You can certainly manually create/convert a TE model and pass it in though and everything should still be fine.

This check is done at `src/accelerate/utils/transformer_engine.py`, and nn.Linear and nn.LayerNorm, are converted to their te counterparts, but will it be able to autocast non-te modules to FP8? i.e BitLinear?

I have a great need for FP8 compute and will do whatever I need to for implementation, I just want to be sure I am looking in the right place. Thanks!

No, you would need to replace/create your own BitLinear layer yourself in te, then pass that model in. TE and MS-AMP only allow for converting of normal nn.Linear and nn.LN. (If you look on the TE repo, there are a number of layers you can use to build your model off of, and maybe your own custom layer)

Any update?