Add distopt support for FP8 params and BF16 optimizer state
State: closed | Pull request: yes
Labels: name: core

name: NLP

name: CI

# What does this PR do ?

Adds support in the Apex distributed Adam optimizer to support FP8 parameters (using experimental FP8 tensors from Transformer Engine) and BF16 optimizer state.

**Collection**: NLP

# Changelog 
- Adds distributed optimizer support for FP8 parameters
- Adds the option to initialize GPT with FP8 parameters
- Add support for non-FP32 distributed optimizer state

# Usage

Run GPT, e.g. with the config at https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/language_modeling/conf/megatron_gpt_config.yaml.

Enable FP8 support with `model.fp8=True`, FP8 parameters with `model.fp8_params=True`, the distributed optimizer with `model.optim.name=distributed_fused_adam`, and BF16 optimizer state with `model.optim.dtype=bf16`.

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [x] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Equivalent to https://github.com/NVIDIA/NeMo/pull/7885 and https://github.com/NVIDIA/NeMo/pull/7895.
* Depends on https://github.com/NVIDIA/apex/pull/1749
* Depends on ~https://github.com/NVIDIA/TransformerEngine/pull/524~ https://github.com/NVIDIA/TransformerEngine/pull/529
* Uses Transformer Engine `Float8Tensor` added in https://github.com/NVIDIA/TransformerEngine/pull/452
* Closes https://github.com/NVIDIA/NeMo/pull/7469, closes https://github.com/NVIDIA/NeMo/pull/7565, closes https://github.com/NVIDIA/NeMo/pull/7885, closes https://github.com/NVIDIA/NeMo/pull/7895.


I trained the llama model with 2 nodes using model.fp8=True, model.fp8_params=True, model.optim.name=distributed_fused_adam, model.optim.dtype=bf16.
I got this error when saving checkpoint

<img width="1421" alt="Screenshot 2023-11-20 at 17 27 23" src="https://github.com/NVIDIA/NeMo/assets/25007718/5ed193c7-44a4-4387-a2b5-f2cf40aa87af">


How can i solve this problem

Thank you!

@chiendb97 Thanks for the report! I reproduce the error and have a fix at https://github.com/NVIDIA/TransformerEngine/pull/529. I haven't tested thoroughly, but I was able to save and load a checkpoint for LLaMa with FP8 params.

@timmoon10 Is there any blocker for the review? With this PR, the memory allocation per param is: 1(fp8)+1(fp8 transpose)+2(BF16 gradient)+[2(bf16 weight)+2(bf16 momentum)+2(bf16 variance)]/dp=4+6/dp. Is my understanding correct?

@lhb8125 The only remaining blocker is fixing some convergence issues when running with SFT. I don't fully understand that issue yet, but I don't expect it would require major changes. Starting a review would be great.

I've found a bug when using this PR with LLaMa SFT. Bugfix: https://github.com/NVIDIA/TransformerEngine/pull/567

This does not affect GPT pretraining though. I think this is ready to review and merge.

jenkins

The Jenkins failure is because it is using an old version of Apex. The Dockerfile and README have been updated with the required Apex.

What Base PyTorch version is needed? We can update it in the Jenkinsfile

jenkins

jenkins

jenkins

This PR requires [Apex commit a2f6683](https://github.com/NVIDIA/apex/commit/a2f6683b10fb4c29ab57c9e3d16957db76a8a5ba) (11/19), which is included in the PyTorch 23.11 container. However, [Apex commit 6ff4548](https://github.com/NVIDIA/apex/commit/6ff45486f432f91eb86937a0def5eb5f2cf792ae) (9/27) breaks the API used by Megatron-LM and causes the Jenkins pipeline to fail.

Probably the easiest fix is to tweak Apex so that [6ff4548](https://github.com/NVIDIA/apex/commit/6ff45486f432f91eb86937a0def5eb5f2cf792ae) is backward compatible. I'll try to get a PR opened in Apex, although it'll probably be 2024 before we can get it in. Once that issue is debugged, I propose that we build a custom container to use in Jenkins instead of waiting for it to get included in the official PyTorch container.

jenkins

jenkins

jenkins

@ericharper 
Can you do the final check and help merge?

jenkins

@athitten 
Can you help review this PR? I think @ericharper is away.

jenkins

jenkins

@timmoon10 Can you check why this fails the CI?

The error message is "No space left on device", so I suspect it's related to the recent file system issues on DLCluster. I find I often need to run a couple times to get past these errors, as well as segfaults coming from ASR.

jenkins

jenkins

jenkins

jenkins

jenkins

jenkins