From 5ad074227709522c232844b89320b8d51063f185 Mon Sep 17 00:00:00 2001 From: loubnabnl Date: Mon, 13 Nov 2023 15:34:10 +0000 Subject: [PATCH 1/5] add model --- src/transformers/models/gpt_bigcode/README.md | 18 ++ .../convert_megatron_checkpoint.py | 299 ++++++++++++++++++ 2 files changed, 317 insertions(+) create mode 100644 src/transformers/models/gpt_bigcode/README.md create mode 100644 src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py diff --git a/src/transformers/models/gpt_bigcode/README.md b/src/transformers/models/gpt_bigcode/README.md new file mode 100644 index 0000000000..947250ad86 --- /dev/null +++ b/src/transformers/models/gpt_bigcode/README.md @@ -0,0 +1,18 @@ +## Conversion to `transformers` + +To convert a model from Megatron-LM to transformers use: +```bash +source ~/.bashrc +export PYTHONPATH=Megatron-LM +export PYTHONPATH=transformers/src:$PYTHONPATH + +cd transformers/src/transformers/models + +python gpt_bigcode/convert_megatron_checkpoint.py \ + --path_to_checkpoint /fsx/bigcode/experiments/pretraining/starcoder2-1B/checkpoints/iter_0200000/mp_rank_00/model_optim_rng.pt \ + --save_dir /fsx/bigcode/experiments/pretraining/starcoder2-1B/checkpoints/conversions \ + --test_generation \ + --tokenizer_path /fsx/loubna/data/tokenizer/starcoder2-smol-internal-1 +``` + +For `fast-llm` use `convert_fast_llm_checkpoint.py`. For cloning and pushing models from existng iterations directly to HF hub check `push_checkpoints.py`. \ No newline at end of file diff --git a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py new file mode 100644 index 0000000000..079be1c5c1 --- /dev/null +++ b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py @@ -0,0 +1,299 @@ +#################################################################################################### + +# Copyright (c) 2021-, NVIDIA CORPORATION. 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. + +#################################################################################################### + +# +# Note: If when running this conversion script you're getting an exception: +# ModuleNotFoundError: No module named 'megatron.model.enums' +# you need to tell python where to find the clone of Megatron-LM, e.g.: +# +# cd /tmp +# git clone https://github.com/NVIDIA/Megatron-LM +# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ... +# +# if you already have it cloned elsewhere, simply adjust the path to the existing path +# +# If the training was done using a Megatron-LM fork, e.g., +# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one +# in your path, i.e., /path/to/Megatron-DeepSpeed/ +# + +import argparse +import os +import re + +import torch + +from transformers.models.gpt_bigcode import GPTBigCodeConfig, GPTBigCodeForCausalLM, GPTBigCodeModel + + +# The simple map of names for "automated" rules. +NAME_MAP = { + "attention.dense": ".attn.c_proj.", + "self_attention.dense": ".attn.c_proj.", + "mlp.dense_h_to_4h": ".mlp.c_fc.", + "mlp.dense_4h_to_h": ".mlp.c_proj.", + "self_attention.query_key_value": ".attn.c_attn.", + "self_attention.query": ".attn.q_attn.", + "self_attention.key_value": ".attn.kv_attn.", +} + + +def recursive_print(name, val, spaces=0): + # Format the message. + if name is None: + msg = None + else: + fmt = "." * max(0, spaces - 2) + "# {:" + str(50 - spaces) + "s}" + msg = fmt.format(name) + + # Print and recurse (if needed). + if isinstance(val, dict): + if msg is not None: + print(msg) + for k in val.keys(): + recursive_print(k, val[k], spaces + 2) + elif isinstance(val, torch.Tensor): + print(msg, ":", val.size()) + else: + print(msg, ":", val) + + +def convert_megatron_checkpoint(input_state_dict, merge_qkv): + # The converted output model. + output_state_dict = {} + ds_args = input_state_dict["args"] + + if ds_args is not None: + # @loubnabnl fastllm uses gelu? + if ds_args.bias_gelu_fusion: + activation_function = "gelu_pytorch_tanh" + elif ds_args.openai_gelu: + activation_function = "gelu_new" + else: + activation_function = "gelu" + else: + # in the very early days this used to be "gelu_new" + activation_function = "gelu_new" + + if ds_args.attention_head_type == "multihead": + multi_query = False + else: + assert ds_args.attention_head_type == "multiquery" + # @loubnabnl we don't use the no-merge-kv anymore? + # attention_type = 2 if merge_qkv else 3 + multi_query = True + + attention_softmax_in_fp32 = ds_args.attention_softmax_in_fp32 or ds_args.apply_query_key_layer_scaling + + # Spell out all parameters in case the defaults change. + config = GPTBigCodeConfig( + architectures=["GPTBigCodeLMHeadModel"], + vocab_size=ds_args.padded_vocab_size, + n_positions=ds_args.max_position_embeddings, + n_embd=ds_args.hidden_size, + n_layer=ds_args.num_layers, + n_head=ds_args.num_attention_heads, + n_inner=ds_args.ffn_hidden_size, + activation_function=activation_function, + multi_query=multi_query, + resid_pdrop=0.1, + embd_pdrop=0.1, + attn_pdrop=0.1, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + summary_type="cls_index", + summary_use_proj=True, + summary_activation=None, + summary_proj_to_labels=True, + summary_first_dropout=0.1, + scale_attn_weights=True, + use_cache=True, + bos_token_id=0, + eos_token_id=0, + attention_softmax_in_fp32=attention_softmax_in_fp32, + scale_attention_softmax_in_fp32=True, + use_rotary_embeddings=ds_args.use_rotary_position_embeddings, + rotary_embedding_scale=ds_args.rotary_theta, + use_position_embeddings=ds_args.add_position_embedding, + ) + + from pprint import pprint + pprint(vars(ds_args)) + pprint(config) + + # Megatron-LM checkpoint version + checkpoint_version = input_state_dict["checkpoint_version"] + if checkpoint_version < 2.0: + raise NotImplementedError(f"Checkpoint version {checkpoint_version} not supported.") + + # The model. + model = input_state_dict["model"]["language_model"] + + # The word embeddings, truncated to to vocab_size rows. + word_embeddings = model["embedding"]["word_embeddings"]["weight"][: config.vocab_size, :] + output_state_dict["transformer.wte.weight"] = word_embeddings + + # The position embeddings. + output_state_dict["transformer.wpe.weight"] = model["embedding"]["position_embeddings"]["weight"] + + # The transformer. + transformer = model["transformer"] if "transformer" in model else model["encoder"] + + # The regex to extract layer names. + layer_re = re.compile("layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)") + + # Extract the layers. + for key, val in transformer.items(): + # Match the name. + m = layer_re.match(key) + + # Stop if that's not a layer + if m is None: + break + + # The index of the layer. + layer_idx = int(m.group(1)) + # The name of the operation. + op_name = m.group(2) + # Is it a weight or a bias? + weight_or_bias = m.group(3) + + # The name of the layer. + layer_name = f"transformer.h.{layer_idx}" + + # For layernorm(s), simply store the layer norm. + if op_name.endswith("layernorm"): + + ln_name = "ln_1" if op_name.startswith("input") else "ln_2" + output_state_dict[layer_name + "." + ln_name + "." + weight_or_bias] = val + + # Concatenate QKV matrix. + elif merge_qkv and (op_name == "self_attention.key_value"): + # Query is before key_value in the dict. + query = output_state_dict.pop(layer_name + ".attn.q_attn." + weight_or_bias) + out_val = torch.cat([query, val], dim=0) + output_state_dict[layer_name + ".attn.c_attn." + weight_or_bias] = out_val + + # Copy the parameters. + else: + output_state_dict[layer_name + NAME_MAP[op_name] + weight_or_bias] = val + + # DEBUG. + assert config.n_layer == layer_idx + 1 + + # The final layernorm. + output_state_dict["transformer.ln_f.weight"] = transformer["final_layernorm.weight"] + output_state_dict["transformer.ln_f.bias"] = transformer["final_layernorm.bias"] + + # For LM head, transformers' wants the matrix to weight embeddings. + output_state_dict["lm_head.weight"] = word_embeddings + + # It should be done! + return config, output_state_dict + + +def test_conversion(checkpoint_path, tokenizer_path): + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + model = GPTBigCodeForCausalLM.from_pretrained(checkpoint_path, torch_dtype=torch.bfloat16, device_map="cuda") + text = "def fibonnaci(n" + inputs = tokenizer(text, return_tensors="pt").to("cuda") + print(f"Testing generation with prompt '{text}'") + print(f"Input ids: {inputs['input_ids']}") + output = model.generate(**inputs, max_new_tokens=128, do_sample=False) + print(tokenizer.decode(output[0])) + + +def main(argv=None): + # Create the argument parser. + parser = argparse.ArgumentParser() + parser.add_argument("--print-checkpoint-structure", action="store_true") + parser.add_argument( + "--path_to_checkpoint", + type=str, + help="Path to the checkpoint file (.zip archive or direct .pt file)", + ) + parser.add_argument( + "--no_merge_qkv", + dest="merge_qkv", + action="store_false", + help="Do not merge the query and key_value tensors (MQA).", + ) + parser.add_argument( + "--custom_model", + action="store_true", + help="Save as custom model so it can be used with huggingface transformers.", + ) + parser.add_argument( + "--save_dir", help="Path where the converted model is saved. Will use the checkpoint directory if not provided" + ) + parser.add_argument( + "--tokenizer_path", + type=str, + help="Path to the tokenizer or repo name on the HF hub for testing", + ) + parser.add_argument( + "--test_generation", + action="store_true", + help="Test generation with the converted model", + ) + args = parser.parse_args(argv) + + # Extract the basename. + basename = args.save_dir or os.path.dirname(args.path_to_checkpoint) + + # Load the model. + print(f"Extracting PyTorch state dictionary from {args.path_to_checkpoint}") + input_state_dict = torch.load(args.path_to_checkpoint, map_location="cpu") + + # Convert. + print("Converting") + config, output_state_dict = convert_megatron_checkpoint(input_state_dict, args.merge_qkv) + + # Print the structure of converted state dict. + if args.print_checkpoint_structure: + recursive_print(None, output_state_dict) + + if args.custom_model: + # Save custom model + GPTBigCodeConfig.register_for_auto_class() + GPTBigCodeModel.register_for_auto_class("AutoModelForCausalLM") + hf_model = GPTBigCodeForCausalLM(config) + hf_model.load_state_dict(output_state_dict) + hf_model.save_pretrained(basename) + + else: + # Store the config to file. + print("Saving config") + config.save_pretrained(basename) + + # Store the state_dict to file. + output_checkpoint_file = os.path.join(basename, "pytorch_model.bin") + print(f'Saving checkpoint to "{output_checkpoint_file}"') + torch.save(output_state_dict, output_checkpoint_file) + + # test model + if args.test_generation: + print(f"Testing converted model at {args.save_dir}") + if args.tokenizer_path is None: + raise ValueError("Please provide a tokenizer path for testing") + test_conversion(checkpoint_path=args.save_dir, tokenizer_path=args.tokenizer_path) + + +if __name__ == "__main__": + main() From b2d84f5ec8f9852258fcbb2dcb8ac519fa6697d5 Mon Sep 17 00:00:00 2001 From: Loubna Ben Allal <44069155+loubnabnl@users.noreply.github.com> Date: Mon, 13 Nov 2023 17:01:18 +0000 Subject: [PATCH 2/5] Update src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py Co-authored-by: RaymondLi0 --- .../models/gpt_bigcode/convert_megatron_checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py index 079be1c5c1..6fb4c672b0 100644 --- a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py +++ b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py @@ -128,7 +128,7 @@ def convert_megatron_checkpoint(input_state_dict, merge_qkv): attention_softmax_in_fp32=attention_softmax_in_fp32, scale_attention_softmax_in_fp32=True, use_rotary_embeddings=ds_args.use_rotary_position_embeddings, - rotary_embedding_scale=ds_args.rotary_theta, + rotary_embedding_scale=-math.log(ds_args.rotary_theta), use_position_embeddings=ds_args.add_position_embedding, ) From c144d9e97dc15c9476abfdde4c1337f93127e5b6 Mon Sep 17 00:00:00 2001 From: loubnabnl Date: Mon, 13 Nov 2023 17:02:29 +0000 Subject: [PATCH 3/5] import math --- .../models/gpt_bigcode/convert_megatron_checkpoint.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py index 6fb4c672b0..0bed939e1b 100644 --- a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py +++ b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py @@ -37,6 +37,7 @@ import re import torch +import math from transformers.models.gpt_bigcode import GPTBigCodeConfig, GPTBigCodeForCausalLM, GPTBigCodeModel From 1c243bad519f44fdee257bb75845e3345d81b3dd Mon Sep 17 00:00:00 2001 From: loubnabnl Date: Mon, 13 Nov 2023 18:53:43 +0000 Subject: [PATCH 4/5] apply raymond's suggestion/megatron's rope --- src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py index 387e78eb46..e7583893c5 100644 --- a/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py +++ b/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py @@ -89,7 +89,7 @@ def _apply_rotary_embeddings( * Convert back tho the input format. # TODO: Full precision only needed for bfloat16? (Doesn't support complex numbers) """ - complex_tensor = torch.view_as_complex(tensor.float().view(*tensor.shape[:-1], -1, rope_frequencies.size(-1), 2)) + complex_tensor = torch.view_as_complex(tensor.float().view(*tensor.shape[:-1], -1, 2, rope_frequencies.size(-1)).transpose(-2, -1).contiguous()) return torch.view_as_real(complex_tensor * rope_frequencies).view_as(tensor).type_as(tensor) From fd0fbef15e9e36fc90b6b4acbc273649c91d6f47 Mon Sep 17 00:00:00 2001 From: loubnabnl Date: Mon, 13 Nov 2023 18:54:32 +0000 Subject: [PATCH 5/5] add antother example to the test --- .../convert_megatron_checkpoint.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py index 0bed939e1b..ccb92e5c8f 100644 --- a/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py +++ b/src/transformers/models/gpt_bigcode/convert_megatron_checkpoint.py @@ -208,16 +208,20 @@ def convert_megatron_checkpoint(input_state_dict, merge_qkv): return config, output_state_dict -def test_conversion(checkpoint_path, tokenizer_path): +def test_conversion(checkpoint_path, tokenizer_path, device="cpu", prompt=None): from transformers import AutoTokenizer + from transformers.models.gpt_bigcode import GPTBigCodeForCausalLM tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - model = GPTBigCodeForCausalLM.from_pretrained(checkpoint_path, torch_dtype=torch.bfloat16, device_map="cuda") - text = "def fibonnaci(n" - inputs = tokenizer(text, return_tensors="pt").to("cuda") - print(f"Testing generation with prompt '{text}'") - print(f"Input ids: {inputs['input_ids']}") - output = model.generate(**inputs, max_new_tokens=128, do_sample=False) - print(tokenizer.decode(output[0])) + model = GPTBigCodeForCausalLM.from_pretrained(checkpoint_path, torch_dtype=torch.bfloat16, device_map=device) + prompt_1 = 'def separate_paren_groups(paren_string: str) -> List[str]:\n """ Input to this function is a string containing multiple groups of nested parentheses. Your goal is to\n separate those group into separate strings and return the list of those.\n Separate groups are balanced (each open brace is properly closed) and not nested within each other\n Ignore any spaces in the input string.\n >>> separate_paren_groups(\'( ) (( )) (( )( ))\')\n [\'()\', \'(())\', \'(()())\']\n """' + prompt_2 = 'def fibonnaci(n' + prompts = [prompt_1, prompt_2] + for text in prompts: + inputs = tokenizer(text, return_tensors="pt").to(device) + print(f"Testing generation with prompt '{text}'") + print(f"Input ids: {inputs['input_ids']}") + output = model.generate(**inputs, max_new_tokens=128, do_sample=False) + print(tokenizer.decode(output[0])) def main(argv=None):