Draft PR Adding mistral 0.1

README.md

@@ -507,13 +507,14 @@ Though NeoX supports a number of different architectural configurations, includi
 
 NeoX supports export of compatible models into the following architectures:
 - GPTNeoXForCausalLM
-- LlamaForCausalLM (GQA Support Coming Soon -- all Llama 1 models and Llama 2 / Codellama up to size 13B supported)
+- LlamaForCausalLM
+- MistralForCausalLM
 
 Training a model which does not fit into one of these Hugging Face Transformers architectures cleanly will require writing custom modeling code for the exported model.
 
 To convert a GPT-NeoX library checkpoint to Hugging Face-loadable format, run:
 ```bash
-python ./tools/ckpts/convert_neox_to_hf.py --input_dir /path/to/model/global_stepXXX --config_file your_config.yml --output_dir hf_model/save/location --precision {auto,fp16,bf16,fp32} --architecture {neox,llama}
+python ./tools/ckpts/convert_neox_to_hf.py --input_dir /path/to/model/global_stepXXX --config_file your_config.yml --output_dir hf_model/save/location --precision {auto,fp16,bf16,fp32} --architecture {neox,mistral,llama}
 ```
 
 Then to upload a model to [the Hugging Face Hub](https://huggingface.co/), run:
@@ -527,11 +528,11 @@ and input the requested information, including HF hub user token.
 
 NeoX supplies several utilities for converting a pretrained model checkpoint into a format that can be trained within the library.
 
-The following models can be loaded in GPT-NeoX:
+The following models or model families can be loaded in GPT-NeoX:
 - Llama 1
-- Llama 2 (Up to size 13B)
-- CodeLlama (Up to size 13B)
-- Mistral-7b-v0.1 (Coming Soon!)
+- Llama 2
+- CodeLlama
+- Mistral-7b-v0.1
 
 We provide two utilities for converting from two different checkpoint formats into a format compatible with GPT-NeoX.
 

configs/llama2/13B.yml

@@ -0,0 +1,26 @@
+{
+  "pipe_parallel_size": 1,
+  "model_parallel_size": 2,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 40,
+  "hidden_size": 5120,
+  "num_attention_heads": 40,
+  "seq_length": 4096,
+  "max_position_embeddings": 4096,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+}

configs/llama2/70B.yml

@@ -0,0 +1,31 @@
+{
+  "pipe_parallel_size": 1,
+  "model_parallel_size": 1,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 80,
+  "hidden_size": 8192,
+  "intermediate_size": 28672,
+  "num_attention_heads": 64,
+  "num_kv_heads": 8,
+  "seq_length": 4096,
+  "max_position_embeddings": 4096,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "rotary_emb_base": 1000000,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "attention_config": [[["flash"], 48]],
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+}

configs/llama2/7B.yml

@@ -0,0 +1,26 @@
+{
+  "pipe_parallel_size": 1,
+  "model_parallel_size": 1,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 32,
+  "hidden_size": 4096,
+  "num_attention_heads": 32,
+  "seq_length": 4096,
+  "max_position_embeddings": 4096,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+}

configs/llama2/codellama_34B.yml

@@ -0,0 +1,32 @@
+{
+  "pipe_parallel_size": 1,
+  "model_parallel_size": 1,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 48,
+  "hidden_size": 8192,
+  "num_attention_heads": 64,
+  "num_kv_heads": 8,
+  # Codellama was uptrained on 16k token sequence lengths
+  # with rotary_emb_base adjusted to 1_000_000.
+  "seq_length": 16384,
+  "max_position_embeddings": 16384,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "rotary_emb_base": 1000000,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "attention_config": [[["flash"], 48]],
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+}

configs/llama2/codellama_7B.yml

@@ -0,0 +1,31 @@
+{
+  "pipe_parallel_size": 1,
+  "model_parallel_size": 1,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 32,
+  "hidden_size": 4096,
+  "num_attention_heads": 32,
+  # Codellama was uptrained on 16k token sequence lengths
+  # with rotary_emb_base adjusted to 1_000_000.
+  "seq_length": 16384,
+  "max_position_embeddings": 16384,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "rotary_emb_base": 1000000,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "attention_config": [[["flash"], 32]],
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+}

configs/llemma/34B.yml

@@ -0,0 +1,101 @@
+{
+  "pipe_parallel_size": 0,
+  "model_parallel_size": 8,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 48,
+  "hidden_size": 8192,
+  "num_attention_heads": 64,
+  "attention_type": "groupedquery",
+  "num_kv_heads": 8,
+  # NB: These rotary embedding and sequence length parameters
+  # May differ from CodeLlama configs. They match what we used for
+  # Llemma continued pretraining. See https://arxiv.org/abs/2310.10631
+  # For detailed discussion
+  "seq_length": 4096,
+  "max_position_embeddings": 4096,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "rotary_emb_base": 1000000,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "attention_config": [[["flash"], 48]],
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+
+   "optimizer": {
+     "type": "Adam",
+     "params": {
+       "lr": 0.00005,
+       "betas": [0.9, 0.95],
+       "eps": 1.0e-8
+     }
+   },
+
+   "zero_optimization": {
+    "stage": 1,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 1260000000,
+    "overlap_comm": true,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 1260000000,
+    "contiguous_gradients": true,
+    "cpu_offload": false
+  },
+
+  # trained on 256 gpus
+  "train_micro_batch_size_per_gpu": 2,
+  "gradient_accumulation_steps": 16,
+  "data_impl": "mmap",
+
+  "checkpoint_activations": true,
+  "checkpoint_num_layers": 1,
+  "partition_activations": true,
+  "synchronize_each_layer": true,
+
+  "gradient_clipping": 1.0,
+  "weight_decay": 0.1,
+  "hidden_dropout": 0,
+  "attention_dropout": 0,
+
+  "precision": "bfloat16",
+  "fp32_allreduce": true,
+  "bf16": {
+    "enabled": true
+  },
+  "data_types": {
+    "grad_accum_dtype": "fp32"
+  },
+
+  "train_iters": 12000,
+  "lr_decay_iters": 12000,
+  "distributed_backend": "nccl",
+  "lr_decay_style": "cosine",
+  "min_lr": 1.65e-6,
+  "warmup": 0.042, # warmup for ~500 iters
+  "checkpoint_factor": 250,
+  "eval_interval": 250,
+  "eval_iters": 25,
+
+  "log_interval": 1,
+  "steps_per_print": 1,
+  "wall_clock_breakdown": true,
+
+  "tokenizer_type": "SPMTokenizer",
+  #"vocab-file": # use 'tokenizer.model' from Meta CodeLlama download
+
+  # "load": "" # set to same as "save" to resume from intermediate finetuning step
+  #"load": MP=8 CodeLlama-34B checkpoint, converted from Meta CodeLlama download.
+  # When resuming from mid-finetuning run, change "load" to the same as save location.
+  "finetune": true, # set to false once resuming from intermediate finetuning step
+}

configs/llemma/7B.yml

@@ -0,0 +1,100 @@
+{
+  "pipe_parallel_size": 0,
+  "model_parallel_size": 2,
+  "make_vocab_size_divisible_by": 1,
+
+  # model settings
+  "num_layers": 32,
+  "hidden_size": 4096,
+  "num_attention_heads": 32,
+  # NB: These rotary embedding and sequence length parameters
+  # May differ from CodeLlama configs. They match what we used for
+  # Llemma continued pretraining. See https://arxiv.org/abs/2310.10631
+  # For detailed discussion
+  "seq_length": 4096,
+  "max_position_embeddings": 4096,
+  "pos_emb": "rotary",
+  "rotary_pct": 1,
+  "rotary_emb_base": 10000,
+  "no_weight_tying": true,
+  "gpt_j_residual": false,
+  "output_layer_parallelism": "column",
+  "norm": "rmsnorm",
+  "rms_norm_epsilon": 1.0e-5,
+
+  "attention_config": [[["flash"], 32]],
+
+  "scaled_upper_triang_masked_softmax_fusion": true,
+  "bias_gelu_fusion": false,
+  "use_bias_in_norms": false,
+  "use_bias_in_attn_linear": false,
+  "mlp_type": "llama",
+  "activation": "silu",
+
+   "optimizer": {
+     "type": "Adam",
+     "params": {
+       "lr": 0.0001,
+       "betas": [0.9, 0.95],
+       "eps": 1.0e-8
+     }
+   },
+
+   "zero_optimization": {
+    "stage": 1,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 1260000000,
+    "overlap_comm": true,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 1260000000,
+    "contiguous_gradients": true,
+    "cpu_offload": false
+  },
+
+  # trained on 256 gpus
+  "train_micro_batch_size_per_gpu": 4,
+  "gradient_accumulation_steps": 2,
+  "data_impl": "mmap",
+
+  "checkpoint_activations": true,
+  "checkpoint_num_layers": 1,
+  "partition_activations": true,
+  "synchronize_each_layer": true,
+
+  "gradient_clipping": 1.0,
+  "weight_decay": 0.1,
+  "hidden_dropout": 0,
+  "attention_dropout": 0,
+
+  "precision": "bfloat16",
+  "fp32_allreduce": true,
+  "bf16": {
+    "enabled": true
+  },
+  "data_types": {
+    "grad_accum_dtype": "fp32"
+  },
+
+  "train_iters": 48000,
+  "lr_decay_iters": 48000,
+  "distributed_backend": "nccl",
+  "lr_decay_style": "cosine",
+  "decay_lr_to": 0.033,
+  "warmup_iters": 500,
+  "checkpoint_factor": 500,
+  "eval_interval": 250,
+  "eval_iters": 50,
+
+  "log_interval": 1,
+  "steps_per_print": 1,
+  "wall_clock_breakdown": true,
+
+  "tokenizer_type": "SPMTokenizer",
+  "vocab-file": "codellama/tokenizer.model", # use tokenizer.model from Meta CodeLlama download
+
+  "save": "/path/to/save/llema-replication",
+  #"load": "", # once run is started, to restart from intermediate ckpt use "load" = "save"
+  "load": "/path/to/converted/codellama_7b_weights_with_mp2",
+
+  "finetune": true, # set to false once resuming from intermediate finetuning step
+}