Converting pretrained Nemo T5/UL2 to Huggingface (request for help)
State: closed | Pull request: no
Labels: name: stale

**Is your feature request related to a problem? Please describe.**

An important consideration for us when deciding which library, framework or model implementation to use when pretraining models is whether the model will be accessible and easy to use for end users. In practice this translates to whether it is possible and straightforward to convert the model to a Huggingface compatible format. 

For Nemo, there exists a conversion script to [convert a model from Huggingface to Nemo format](https://github.com/NVIDIA/NeMo/blob/main/scripts/nlp_language_modeling/hf_t5-v1_1_to_nemo.py). However, there doesn't exist a resource for a conversion in the opposite direction. 

Since an `HF -> Nemo` script existed, we thought converting from `Nemo -> HF` would simply be a matter of reversing that script. We therefore [carefully configured our Nemo configs](https://huggingface.co/Lauler/UL2-nemo-conversion/tree/main/nemo_config/ul2-base-nl36) to return a model that weight wise seemed compatible with Huggingface's T5 implementation (e.g. no bias terms, tied encoder/decoder word embeddings). Then we wrote our [own conversion script](https://huggingface.co/Lauler/UL2-nemo-conversion/blob/main/convert_nemo_ul2_checkpoint.py) to reverse Nemo's existing one. 

Our script appears to work fine reversing the weight conversion performed in Nemo's `HF -> Nemo` script. Unfortunately, it does not work as well when attempting to convert a model initialized and pretrained with Nemo Megatron. 

We are aware that [Megatron-LM uses different ordering of QKV](https://github.com/NVIDIA/Megatron-LM/blob/42c1cf4279acea5a554500dcb552211f44cbec45/megatron/checkpointing.py#L209-L237) in the attention layers depending on the version of Megatron-LM used. We are also aware of an existing conversion script that Huggingface have created for converting Megatron-BERT to Huggingface, where they adapt the ordering of QKV in Megatron to [match the ordering used in Huggingface](https://github.com/NVIDIA/Megatron-LM/blob/42c1cf4279acea5a554500dcb552211f44cbec45/megatron/checkpointing.py#L209-L237). We copied this reordering, and added a `--fix_qkv` argument in our conversion script that applies the same reordering of QKV as Huggingface does. 

Unfortunately, none of the above solves the issue we have with our conversion script.

**Describe the solution you'd like**

We've uploaded Nemo model checkpoints, conversion scripts, and detailed instructions in the README on how to reproduce our attempted conversion in [this Huggingface repo](https://huggingface.co/Lauler/UL2-nemo-conversion
). We'd very much appreciate some help figuring out how T5 and UL2 models can be converted from Nemo Megatron to HF. Specifically:

1. Is our Nemo config appropriate to be able to perform a weight conversion? 
2. Does the issue lie with the QKV reordering in Megatron? What is the correct way of reordering QKV in self-attention and KV in cross attention for encoder-decoders? How does one find out which Megatron "checkpoint version" one's saved checkpoints are using?

To clone the HF repo:

```bash
git clone https://huggingface.co/Lauler/UL2-nemo-conversion
```

**Describe alternatives you've considered**

We searched for [UL2](https://huggingface.co/models?search=ul2) and T5 models on Huggingface to see whether we could find any that were successfully converted from Megatron-LM or Nemo. All the HF compatible ones we could find were either FLAX/JAX or models with unclear origin. Nvidia's own uploaded T5 models consist of [`.nemo` checkpoints with instructions to install Apex and Nemo](https://huggingface.co/nvidia/nemo-megatron-t5-3B).

Ideally, we'd like to use Nemo to pretrain language models. But realistically we are going to go for the option that means our users are able to use the final model. Nemo is a great library, but the reality is only a minority of our users would be able to use the models without the kind of well documented convenience functions that exist in HF.

So in other words: our main alternative at this point is switching to JAX/FLAX.  

**Additional context**

Would it be more appropriate to open an issue about this in the Megatron-LM repo rather than here?