# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2023, KBLab at the National Library of Sweden. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Script to convert NeMo Megatron T5/UL2 model to Huggingface T5 model. Based off of NVIDIA's conversion script at: https://github.com/NVIDIA/NeMo/blob/main/scripts/nlp_language_modeling/hf_t5-v1_1_to_nemo.py . We reverse their conversion process. NOTE: You may want to double check the conversion if you are using a custom config with shared_decoder_tokens_head_embeddings=False. """ import argparse import os import collections import sys import torch from nemo.collections.nlp.models.language_modeling.megatron_t5_model import MegatronT5Model from omegaconf.omegaconf import OmegaConf from pytorch_lightning.trainer.trainer import Trainer from transformers import AutoTokenizer, T5Config, T5ForConditionalGeneration # Make hidden_size, num_heads, kv_dim configurable as args with argparse def load_nemo_megatron_model(checkpoint_path, devices=1, num_nodes=1, accelerator="gpu"): trainer = Trainer(devices=devices, num_nodes=num_nodes, accelerator=accelerator) model = MegatronT5Model.load_from_checkpoint(checkpoint_path, trainer=trainer) return model def load_huggingface_t5_model(model_config_path): """ # You need to configure config yourself based on your hparams during training # See examples of UL2 hugginface configs: # https://huggingface.co/google/flan-ul2/blob/main/config.json # https://huggingface.co/Finnish-NLP/ul2-base-nl36-finnish/blob/main/config.json """ t5_config = T5Config.from_pretrained(model_config_path) t5_model = T5ForConditionalGeneration(t5_config) return t5_model def _get_model_type_block_layer_hf(k): """ Get info from Huggingface model block and layer names Returns model_type, block number, layer number. """ if k.startswith("encoder"): model_type = "encoder" elif k.startswith("decoder"): model_type = "decoder" else: raise ValueError(f"Unknown model type for {k}") return model_type, int(k.split(".")[2]), int(k.split(".")[4]) def _get_model_type_layer_nemo(k): """ Get info from NeMo layer names. Returns model_type, layer number. 5th element in the split is the layer number. """ print(k) if "encoder" in k: model_type = "encoder" elif "decoder" in k: model_type = "decoder" else: raise ValueError(f"Unknown model type for {k}") return model_type, int(k.split(".")[5]) def fix_query_key_value_ordering(param, checkpoint_version, num_splits, num_heads, hidden_size): # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :] # for compatibility with later versions of NVIDIA Megatron-LM. # The inverse operation is performed inside Megatron-LM to read checkpoints: # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209 # If param is the weight tensor of the self-attention block, the returned tensor # will have to be transposed one more time to be read by HuggingFace BERT. input_shape = param.size() if checkpoint_version == 1.0: # version 1.0 stores [num_heads * hidden_size * num_splits, :] saved_shape = (num_heads, hidden_size, num_splits) + input_shape[1:] param = param.view(*saved_shape) param = param.transpose(0, 2) param = param.transpose(1, 2).contiguous() elif checkpoint_version >= 2.0: # other versions store [num_heads * num_splits * hidden_size, :] saved_shape = (num_heads, num_splits, hidden_size) + input_shape[1:] param = param.view(*saved_shape) param = param.transpose(0, 1).contiguous() param = param.view(*input_shape) return param def convert_nemo_to_hf( nemo_weights, fix_qkv_ordering=False, hidden_size=768, num_heads=12, kv_dim=64, checkpoint_version=2.0 ): """ Convert NeMo Megatron T5/UL2 model to Huggingface T5 model. Args: nemo_weights (dict): NeMo model weights (state dict). fix_qkv_ordering (bool): Whether to fix the query, key, value ordering in the self-attention blocks. hidden_size (int): Hidden size of the model. num_heads (int): Number of attention heads. kv_dim (int): Projection weights dimension in multi-head attention. Generally: hidden_size // num_heads. checkpoint_version (float): Megatron checkpoint version (No idea how to get this from the checkpoint itself). Returns: hf_weights (dict): Huggingface model weights (state dict). """ print(f"Found {len(nemo_weights.keys())} keys in the NeMo checkpoint") hf_weights = collections.OrderedDict() for k, v in nemo_weights.items(): ################################################# ###### Enc-Dec Embeddings and Output Layer ###### ################################################# # Tied decoder embedding and decoder output layer. if k == "enc_dec_model.decoder_embedding.word_embeddings.weight": # shared.weight, lm_head.weight, decoder.embed_tokens.weight and encoder.embed_tokens.weight # are the same in HF when tied_word_embeddings=True in T5Config. # Corresponding setting in NeMo config: share_decoder_tokens_head_embeddings=True (share decoder vocab embeddings and decoder LM Head) # and share_token_embeddings=True (share encoder/decoder vocab embeddings). # Shared decoder embeddings and LM head yield best result according to: https://aclanthology.org/2021.emnlp-main.465.pdf#page=7 . # Check if encoder and decoder token embeddings are the same. is_shared_encdec = torch.allclose( v, nemo_weights["enc_dec_model.encoder_embedding.word_embeddings.weight"] ) if is_shared_encdec: print("Found shared encoder and decoder embeddings") hf_weights["shared.weight"] = v else: ValueError( ( f"Found separate encoder and decoder embeddings in NeMo checkpoint. \n" f"Not supported in T5 HF implementation. \n" f"You should probably set 'share_token_embeddings' to True in your NeMo config. \n" ) ) if k == "enc_dec_model.tokens_head.weight": # This weight doesn't seem to exist in Nemo when share_decoder_tokens_head_embeddings=True. # Don't worry though. If you set tie_word_embeddings=True in HF, this weight will be # created automatically when loading the model in HF and tied to # shared.weight / decoder.embed_tokens.weight. hf_weights["lm_head.weight"] = v print(f"Mapped {k} to lm_head.weight") elif k == "enc_dec_model.tokens_head.bias": # HF doesn't have a bias for lm_head.weight ValueError( ( f"Found bias for lm_head.weight in NeMo checkpoint. This is not supported in HF T5 implementation. \n" f"You should probably set 'tokens_head_bias' to False in your NeMo config. \n" f"If your checkpoint is from older version of Megatron, you may also need to set 'share_decoder_tokens_head_embeddings' to False in NeMo config. \n" f"See: https://github.com/NVIDIA/NeMo/blob/557c4b7ae766faf050374e6b9a862e2e67385b10/nemo/collections/nlp/models/language_modeling/megatron_lm_encoder_decoder_model.py#L231-L236" ) ) # hf_weights["lm_head.bias"] = v # print(f"Mapped {k} to lm_head.bias") # Decoder embeddings elif k == "enc_dec_model.decoder_embedding.word_embeddings.weight": hf_weights["decoder.embed_tokens.weight"] = v elif k == "enc_dec_model.encoder_embedding.word_embeddings.weight": hf_weights["encoder.embed_tokens.weight"] = v print(f"Mapped {k} to encoder.embed_tokens.weight") ################################################# ################# RPE Weights ################### ################################################# elif k == "enc_dec_model.encoder_relative_position_embedding.relative_position_embedding.weight": hf_weights["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = v print(f"Mapped {k} to encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight") elif k == "enc_dec_model.decoder_relative_position_embedding.relative_position_embedding.weight": hf_weights["decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = v print(f"Mapped {k} to decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight") ################################################# #################$ LayerNorm #################### ################################################# # Block in HF corresponds to layer in NeMo. # Layer in HF does not correspond to anything in NeMo. # In Huggingface: Layer 0 is input layer norm, layer 1 is layer norm on self attn output, # layer 2 is layer norm for cross attn output in decoder. # In NeMo, some layernorm layers (final layernorms) don't have layer number in the name. # We take care of these early so _get_model_type_layer_nemo function doesn't fail. elif "layernorm" in k: if "final" in k: model_type = "encoder" if "encoder" in k else "decoder" # Layer 2 in HF is always FFN + LayerNorm hf_weights[f"{model_type}.final_layer_norm.weight"] = v print(f"Mapped {k} to {model_type}.final_layer_norm.weight") # if "bias" in k: # hf_weights[f"{model_type}.block.final_layer_norm.bias"] = v # print(f"Mapped {k} to {model_type}.block.final_layer_norm.bias") else: model_type, layer_number = _get_model_type_layer_nemo(k) if "input_layernorm" in k and model_type == "encoder": # Input layer norm is always layer 0 in HF hf_weights[f"encoder.block.{layer_number}.layer.0.layer_norm.weight"] = v print(f"Mapped {k} to encoder.block.{layer_number}.layer.0.layer_norm.weight") # if "bias" in k: # hf_weights[f"encoder.block.{layer_number}.layer.0.layer_norm.bias"] = v # print(f"Mapped {k} to encoder.block.{layer_number}.layer.0.layer_norm.bias") elif "post_attention_layernorm" in k and model_type == "encoder": # Layer 1 in HF is layer norm for self attn output hf_weights[f"{model_type}.block.{layer_number}.layer.1.layer_norm.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.layer_norm.weight") # if "bias" in k: # hf_weights[f"{model_type}.block.{layer_number}.layer.1.layer_norm.bias"] = v # print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.layer_norm.bias") elif "input_layernorm" in k and model_type == "decoder": # Input layer norm is always layer 0 in HF hf_weights[f"decoder.block.{layer_number}.layer.0.layer_norm.weight"] = v print(f"Mapped {k} to decoder.block.{layer_number}.layer.0.layer_norm.weight") # if "bias" in k: # hf_weights[f"decoder.block.{layer_number}.layer.0.layer_norm.bias"] = v # print(f"Mapped {k} to decoder.block.{layer_number}.layer.0.layer_norm.bias") elif "post_attention_layernorm" in k and model_type == "decoder": # Layer 1 in HF is layer norm for self attn output hf_weights[f"{model_type}.block.{layer_number}.layer.1.layer_norm.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.layer_norm.weight") # if "bias" in k: # hf_weights[f"{model_type}.block.{layer_number}.layer.1.layer_norm.bias"] = v # print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.layer_norm.bias") elif "post_inter_attention_layernorm" in k and model_type == "decoder": # Layer 2 in HF is layer norm for cross attn output hf_weights[f"{model_type}.block.{layer_number}.layer.2.layer_norm.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.2.layer_norm.weight") # if "bias" in k: # hf_weights[f"{model_type}.block.{layer_number}.layer.2.layer_norm.bias"] = v # print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.2.layer_norm.bias") else: raise ValueError("Unknown layer_norm key: {}".format(k)) ################################################# ############### Attention Layers ################ ################################################# # Self-Attention # Q, k, V in NeMo-Megatron is bundled into a single matrix. elif "self_attention.query_key_value.weight" in k: # Example naming in HF: # encoder.block.0.layer.0.SelfAttention.q.weight # decoder.block.0.layer.0.SelfAttention.q.weight # Model type is either "encoder" or "decoder" model_type, layer_number = _get_model_type_layer_nemo(k) if fix_qkv_ordering: out_val = fix_query_key_value_ordering( v, checkpoint_version=checkpoint_version, num_splits=3, num_heads=num_heads, hidden_size=kv_dim ) else: out_val = v q_weights = out_val[0 * hidden_size : 1 * hidden_size, :] k_weights = out_val[1 * hidden_size : 2 * hidden_size, :] v_weights = out_val[2 * hidden_size : 3 * hidden_size, :] # Layer 0 in HF is always self attn hf_weights[f"{model_type}.block.{layer_number}.layer.0.SelfAttention.q.weight"] = q_weights hf_weights[f"{model_type}.block.{layer_number}.layer.0.SelfAttention.k.weight"] = k_weights hf_weights[f"{model_type}.block.{layer_number}.layer.0.SelfAttention.v.weight"] = v_weights print( ( f"Mapped {k} to: \n", f"{model_type}.block.{layer_number}.layer.0.SelfAttention.q.weight \n", f"{model_type}.block.{layer_number}.layer.0.SelfAttention.k.weight \n", f"{model_type}.block.{layer_number}.layer.0.SelfAttention.v.weight \n", ) ) # If we trained with bias=True in NeMo we will have bias terms for all weight matrices. # Huggingface doesn't support optional bias terms in their T5 implementation. elif "self_attention.query_key_value.bias" in k: ValueError( "Bias terms for most weights are not supported in Huggingface T5. Train with bias=False in NeMo config." ) # Output self-attn matrix. elif "self_attention.dense.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) # Layer 0 in HF still always self attn hf_weights[f"{model_type}.block.{layer_number}.layer.0.SelfAttention.o.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.0.SelfAttention.o.weight") # Cross-Attention projection matrices are merged into K, V matrices in NeMo-Megatron. # Need to split them into K, V matrices in HF. elif "inter_attention.key_value.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) if fix_qkv_ordering: out_val = fix_query_key_value_ordering( v, checkpoint_version=checkpoint_version, num_splits=2, num_heads=num_heads, hidden_size=kv_dim ) else: out_val = v # Layer 1 in HF is always cross attn k_weights = out_val[0 * hidden_size : 1 * hidden_size, :] v_weights = out_val[1 * hidden_size : 2 * hidden_size, :] hf_weights[f"decoder.block.{layer_number}.layer.1.EncDecAttention.k.weight"] = k_weights hf_weights[f"decoder.block.{layer_number}.layer.1.EncDecAttention.v.weight"] = v_weights print( ( f"Mapped {k} to: \n", f"decoder.block.{layer_number}.layer.1.EncDecAttention.k.weight \n", f"decoder.block.{layer_number}.layer.1.EncDecAttention.v.weight \n", ) ) # Cross-Attention Q matrix is separate in NeMo-Megatron and HF. elif "inter_attention.query.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) # Layer 1 in HF is always cross attn hf_weights[f"decoder.block.{layer_number}.layer.1.EncDecAttention.q.weight"] = v print(f"Mapped {k} to decoder.block.{layer_number}.layer.1.EncDecAttention.q.weight") # Output cross-attention matrix. elif "inter_attention.dense.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) # Layer 1 in HF is always cross attn hf_weights[f"decoder.block.{layer_number}.layer.1.EncDecAttention.o.weight"] = v print(f"Mapped {k} to decoder.block.{layer_number}.layer.1.EncDecAttention.o.weight") ################################################# #################$ FFN Layers ################### ################################################# elif "mlp.dense_h_to_4h.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) if model_type == "encoder": # FFN + LayerNorm is always layer 1 in HF encoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.1.DenseReluDense.wi_0.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.DenseReluDense.wi_0.weight") elif model_type == "decoder": # FFN + LayerNorm is always layer 2 in HF decoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.2.DenseReluDense.wi_0.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.2.DenseReluDense.wi_0.weight") elif "mlp.dense_h_to_4h_2.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) if model_type == "encoder": # FFN + LayerNorm is always layer 1 in HF encoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.1.DenseReluDense.wi_1.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.DenseReluDense.wi_1.weight") elif model_type == "decoder": # FFN + LayerNorm is always layer 2 in HF decoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.2.DenseReluDense.wi_1.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.2.DenseReluDense.wi_1.weight") elif "mlp.dense_4h_to_h.weight" in k: model_type, layer_number = _get_model_type_layer_nemo(k) # Layer 2 in HF is always FFN + LayerNorm if model_type == "encoder": # FFN + LayerNorm is always layer 1 in HF encoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.1.DenseReluDense.wo.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.1.DenseReluDense.wo.weight") elif model_type == "decoder": # FFN + LayerNorm is always layer 2 in HF decoder attention blocks. hf_weights[f"{model_type}.block.{layer_number}.layer.2.DenseReluDense.wo.weight"] = v print(f"Mapped {k} to {model_type}.block.{layer_number}.layer.2.DenseReluDense.wo.weight") else: raise ValueError(f"Unknown key: {k}") print("Done mapping weights. \n") print(f"Total keys in converted Huggingface weight mapping: {len(hf_weights.keys())} \n") return hf_weights # singularity shell --nv data/nemo2302 def compare_weights_hf_nemo(model, hf_weights, hf_config_path, hf_model_path=None): """ Compares the weights of a Huggingface initialized model against Nemo model converted to HF. Prints if there are any missing keys that were expected but not mapped. Also compares parameter count of HF initialized model against original unconverted Nemo model. Args: model: NeMo model hf_weights: Dictionary of Huggingface weights hf_config_path: Path to Huggingface config file to initialize model from. hf_model_path: Path to Huggingface Hub or local HF model folder, if you alternatively want to load/initialize from an existing model on HF Hub or disk (optional) """ if args.hf_model_path: # If user supplies a HF hub model path, or local converted model, we load the model from there. hf_model = T5ForConditionalGeneration.from_pretrained(hf_model_path) else: # Otherwise, we load the model from the config. hf_model = load_huggingface_t5_model(hf_config_path) print(f"Total keys in converted Huggingface weight mapping: {len(hf_weights.keys())} \n") print(f"Total keys in Huggingface model initialized from config or HF Hub: {len(hf_model.state_dict().keys())} \n") # Count the number of parameters in the model print( f"Number of parameters in HF model initialized from config or HF hub: {sum(p.numel() for p in hf_model.parameters() if p.requires_grad)}" ) # Number of parameters in Nemo model print(f"Number of parameters in Nemo model: {sum(p.numel() for p in model.parameters() if p.requires_grad)} \n") # Check the set difference between the two sets of model keys (model loaded from config and converted model) print( ( f"Keys in converted HF weight mapping but missing in HF model initialized from config.json: \n" f"{set(hf_weights.keys()) - set(hf_model.state_dict().keys())} \n" ) ) print( ( f"Keys in HF model initialized from config.json but missing in converted HF weight mapping: \n" f"{set(hf_model.state_dict().keys()) - set(hf_weights.keys())} \n" ) ) print( ( f"It is expected that lm_head.weight is missing from converted HF weight mapping \n" f"if you have set share_decoder_tokens_head_embeddings=True in your Nemo config. \n" f"This weight doesn't exist in Nemo, as it is shared with the decoder token embeddings. \n \n" f"In Huggingface, weights for lm_head.weight and decoder token embeddings are generally duplicated \n" f"in the state_dict. When missing, the lm_head.weight is automatically initialized from shared decoder \n" f"token embeddings weights if your HF config.json has tie_word_embeddings=True." ) ) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Convert Nemo T5/UL2 model to Huggingface T5/UL2 model") parser.add_argument( "--nemo_model_path", type=str, required=True, help="Path to Nemo T5/UL2 model .ckpt file", ) parser.add_argument( "--hf_config_path", type=str, required=True, help="Path to Huggingface T5 config.json", ) parser.add_argument( "--hf_model_path", type=str, required=False, help="Path to Huggingface T5 model, local folder or HF hub model", ) parser.add_argument( "--output_path", type=str, required=True, help="Folder to save converted Huggingface T5/UL2 model in", ) parser.add_argument("--hidden_size", type=int, default=768, help="Hidden size of Nemo model") parser.add_argument("--num_heads", type=int, default=12, help="Number of attention heads in Nemo model") # Default False if --fix_qkv not specified parser.add_argument("--fix_qkv", action="store_true", help="Fix QKV weights in converted HF model") parser.add_argument("--checkpoint_version", type=float, default=2.0, help="Checkpoint version of Nemo model") parser.add_argument( "--kv_dim", type=int, default=64, help="Key/Value dimension of Nemo model. Typically hidden_size // num_heads" ) args = parser.parse_args() #### Convert Nemo T5/UL2 model to Huggingface T5/UL2 model model = load_nemo_megatron_model(checkpoint_path=args.nemo_model_path) nemo_weights = model.state_dict() hf_weights = convert_nemo_to_hf( nemo_weights=nemo_weights, fix_qkv_ordering=args.fix_qkv, hidden_size=args.hidden_size, num_heads=args.num_heads, kv_dim=args.kv_dim, checkpoint_version=args.checkpoint_version, ) # We trained with a HF tokenizer, we grab it from the Nemo model. tokenizer = model.tokenizer.__dict__["tokenizer"] # We manually create HF config.json that matches architecture of the nemo model # (or grab one from existing model on HF Hub and modify where necessary). # See example config.json config = T5Config.from_json_file(args.hf_config_path) # Save config config.save_pretrained(args.output_path) print(f"Saved config to {os.path.join(args.output_path, 'config.json')}") # Save tokenizer tokenizer.save_pretrained(args.output_path) print(f"Saved tokenizer to {os.path.join(args.output_path, 'tokenizer.json')}") # Save the converted weights to a file torch.save(hf_weights, os.path.join(args.output_path, "pytorch_model.bin")) print(f"Saved converted weights to {os.path.join(args.output_path, 'pytorch_model.bin')}") # Sanity check compare_weights_hf_nemo(model, hf_weights, hf_config_path=args.hf_config_path)