[BigScience176B] Model conversion from Megatron-LM to transformers
Author: younesbelkada
Pull request: no

### Feature request

Creating here a thread for the model conversion of the BigScience-176B model from Megatron-LM to the transformers library. I will summarize here what I have done so far, the current status of the conversion procedure, as well as a summary of the 'discoveries' that I have made and the small details that we have to care about when dealing with the conversion procedure. I did my work by forking @thomwolf 's fork. The tests has been done on the DGX machine (4 NVIDIA A100).

## :cherry_blossom: Big picture

- Generate some samples with a recent checkpoint
- Testing the exactness of the logits / hidden states values that we obtain using the same input, between the model from Megatron-LM and the converted model. We use a small GPT2 trained on a dummy dataset (2 sentences). This model has been pushed on the hub and being used for integration tests.
- Apply these tests on a recent checkpoint of the 176B model to make sure about the robustness of the tests

## :paperclip: Main links:

- First PR: thomwolf/transformers#1
- WIP PR: thomwolf/transformers#2
- Final PR: #16514
- [The Small debug-GPT2 model](https://huggingface.co/bigscience/bigscience-small-testing)

## :hammer: Current status

- For now, all tests pass on the DGX's GPU (using different conda environments between the Megatron-LM model & the transformers model) with ```assertEqual```. 
- The tests does not pass with ```assertEqual``` when running them on the CPU, but they pass with ```assertAlmostEqual```, with a tolerance of (0.05) for the logits after the ```LayerNorm``` on the Embedding layer and a tolerance ```1e-06``` on the final logits. Check the tests [here](https://github.com/thomwolf/transformers/blob/bigscience176b/tests/bigscience176b/test_embeddings_bigscience176b.py). This behavior of *non-exactness* seem to be expected and we can not do much about that according to pytorch/pytorch#76052
- Added simple reconstruction and encoding tests on the BigScience tokenizer 

## :pushpin: Tips for conversion

+ Explicitly specify the dtype of your modules when initializing them seem to be helpful to ensure exact reproducibility - added an argument `dtype` on the config file
+ Concatenating the weights from Row-parallelized weights seem to return unsimilar results, I made a reproducible script and raised an issue  pytorch/pytorch#76232 the solution for now is to [manually aggregate the results across each TP rank](https://github.com/younesbelkada/transformers/blob/c83999f991137bd475d409a4b70f4903b256e608/src/transformers/models/bigscience176b/modeling_bigscience176b.py#L251). Needs further investigation for possible improvement of the conversion. 

## :white_check_mark: Next steps

- [x] Fix integration tests on the PR thomwolf/transformers#2
- [x] Define which checkpoint to use for the next tests
- [x] Convert the model with the selected checkpoints and compare the hidden states values between the 2 models.  -> fixed some issues in this new commit a4fa70c1a5042fdca7d0fbf26b0aad6ca99fdadc
- [x] `MixedFusedLayerNorm` and `FusedScaledSoftmax` seem to be replaceable respectively by `LayerNorm` and `Softmax` from `torch.nn`. Verify this assumption on the new checkpoints. 
- [ ] Convert a sharded version of the large model and try the tests on that



cc @thomwolf @suzana-ilic @thomasw21 @stas00 

### Motivation

The feature request is related to the [BigScience workshop](https://bigscience.huggingface.co/), where a large Language Model is currently being trained using [Megatron-LM](https://github.com/NVIDIA/Megatron-LM).

### Your contribution

Ultimately submitting a PR to add the BigScience-176B model to the transformers library - by ensuring the exactness of the operations between the converted model and the original trained model on Megatron-LM.

Comments:
I will start by linking this thread to this PR. A possible next step and improvement could be to get some ideas of the equivalency tests that has been made by this PR bigscience-workshop/Megatron-DeepSpeed#121

Also I think the big picture is we want to generate ASAP. (perhaps even before checking exactitude of conversion :S)


As a side note, I'm working on a solution to do model parallelism/offload while maximizing the GPU(s) memory/RAM available which should be useful to run this model on all kinds of setups (albeit more slowly). Should land in Accelerate in the coming weeks :-) 

After running several tests on JZ using a small model (but way larger than the debug model):
+ I fixed issues regarding some operations that were not considered in the previous version of the model a4fa70c1a5042fdca7d0fbf26b0aad6ca99fdadc 
+ There are still small discrepancies most likely due to  `MixedFusedLayerNorm` used in Meg-DS. A workaround is to use the `apex.normalization.FusedLayerNorm` but needs apex to be installed. 

But overall the model outputs the same logits with little discrepancies (`torch.test.assert_all_close` pass for eg with `rtol=1e-7, atol=1e-02`), exact same argmax for the same input as well as some metrics (min max mean) which are the same. I still need to dig more to understand what causes this very small differences.

I will move soon to trying these tests with the large model and see how it will impact everything (maybe we will be able to figure out what causes this little changes)

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Closing this issue since the whole discussion around the model conversion has been done in the WIP PR: #17202