add pretrained model

.gitignore

@@ -0,0 +1 @@
+__pycache__

config_mistral.py

@@ -0,0 +1,70 @@
+""" Example python script to generate a YAML config file which can be used to run a training with nanotron. Refer to "examples" section in the `/README.md` for more information.
+
+Usage:
+```
+python config_tiny_mistral.py
+```
+"""
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+from nanotron.config import (
+    CheckpointsArgs,
+    Config,
+    DataArgs,
+    GeneralArgs,
+    LoggingArgs,
+    LRSchedulerArgs,
+    ModelArgs,
+    OptimizerArgs,
+    ParallelismArgs,
+    PretrainDatasetsArgs,
+    RandomInit,
+    TokenizerArgs,
+    TokensArgs,
+)
+from nanotron.logging import human_format
+
+
+@dataclass
+class MistralConfig:
+    """Configuration for a MISTRAL model
+
+    Be careful on having a coherent typing as we use it to reconstruct the model from yaml
+    """
+
+    attn_pdrop: float = 0.0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    hidden_act: str = "silu"
+    hidden_size: int = 4096
+    initializer_range: float = 0.02
+    intermediate_size: int = 14336
+    is_mistral_config: bool = True  # We use this help differentiate models in yaml/python conversion
+    max_position_embeddings: int = 32768
+    num_attention_heads: int = 32
+    num_hidden_layers: int = 32
+    num_key_value_heads: Optional[int] = 8
+    pad_token_id: Optional[int] = None
+    pretraining_tp: int = 1
+    rms_norm_eps: float = 1e-05
+    rope_theta: float = 10000.0
+    sliding_window_size: int = 4096
+    tie_word_embeddings: bool = False
+    use_cache: bool = True
+    vocab_size: int = 32000
+
+    def __post_init__(self):
+        # for backward compatibility
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+def get_num_params(model_config: MistralConfig) -> int:
+    num_params = model_config.vocab_size * model_config.hidden_size * 2 + \
+        model_config.num_hidden_layers * (
+            3 * model_config.hidden_size * model_config.intermediate_size
+            + 2 * model_config.hidden_size * model_config.hidden_size
+            + 2 * model_config.hidden_size * (model_config.hidden_size / (model_config.num_attention_heads / model_config.num_key_value_heads))
+        )
+    return num_params

config_mistral_7b.py

@@ -0,0 +1,88 @@
+""" Example python script to generate a YAML config file which can be used to run a training with nanotron. Refer to "examples" section in the `/README.md` for more information.
+
+Usage:
+```
+python config_tiny_mistral.py
+```
+"""
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+from nanotron.config import (
+    CheckpointsArgs,
+    Config,
+    DataArgs,
+    GeneralArgs,
+    LoggingArgs,
+    LRSchedulerArgs,
+    ModelArgs,
+    OptimizerArgs,
+    ParallelismArgs,
+    PretrainDatasetsArgs,
+    RandomInit,
+    TokenizerArgs,
+    TokensArgs,
+)
+from nanotron.logging import human_format
+
+from config_mistral import MistralConfig, get_num_params
+
+
+MODEL_CONFIG = MistralConfig(
+    # Config for Mistral 7B
+    attn_pdrop=0.0,
+    bos_token_id=1,
+    eos_token_id=2,
+    hidden_act="silu",
+    hidden_size=4096,
+    initializer_range=0.02,
+    intermediate_size=14336,
+    max_position_embeddings=32768,
+    num_attention_heads=32,
+    num_hidden_layers=32,
+    num_key_value_heads=8,
+    pretraining_tp=1,
+    rms_norm_eps=1e-05,
+    rope_theta=10000.0,
+    sliding_window_size=4096,
+    tie_word_embeddings=False,
+    use_cache=True,
+    vocab_size=32000,
+)
+
+num_params = human_format(get_num_params(MODEL_CONFIG)).replace(".", "p")
+
+print(f"Model has {num_params} parameters")
+
+PARALLELISM = ParallelismArgs(
+    dp=2,
+    pp=2,
+    tp=2,
+    pp_engine="1f1b",
+    tp_mode="REDUCE_SCATTER",
+    tp_linear_async_communication=True,
+    recompute_granularity="selective",
+)
+
+CONFIG = Config(
+    general=GeneralArgs(project="mistralai", run="Mistral-7B-v0.1", seed=42),
+    checkpoints=None,
+    parallelism=PARALLELISM,
+    model=ModelArgs(init_method=RandomInit(std=0.025), model_config=MODEL_CONFIG),
+    tokenizer=TokenizerArgs("mistralai/Mistral-7B-v0.1"),
+    optimizer=None,
+    logging=None,
+    tokens=None,
+    data=None,
+    profiler=None,
+)
+
+if __name__ == "__main__":
+    file_path = os.path.abspath(__file__)
+
+    file_path = file_path.replace(".py", ".yaml")
+    # Save config as YAML file
+    config.save_as_yaml(file_path)
+
+    # You can now train a model with this config using `/run_train.py`

config_mistral_7b.yaml

@@ -0,0 +1,53 @@
+checkpoints: null
+data: null
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: false
+  project: mistralai
+  run: Mistral-7B-v0.1
+  seed: 42
+  step: null
+logging: null
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: bfloat16
+  init_method:
+    std: 0.025
+  make_vocab_size_divisible_by: 1
+  model_config:
+    attn_pdrop: 0.0
+    bos_token_id: 1
+    eos_token_id: 2
+    hidden_act: silu
+    hidden_size: 4096
+    initializer_range: 0.02
+    intermediate_size: 14336
+    is_mistral_config: true
+    max_position_embeddings: 32768
+    num_attention_heads: 32
+    num_hidden_layers: 32
+    num_key_value_heads: 8
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_theta: 10000.0
+    sliding_window_size: 4096
+    tie_word_embeddings: false
+    use_cache: true
+    vocab_size: 32000
+optimizer: null
+parallelism:
+  dp: 2
+  pp: 2
+  pp_engine: 1f1b
+  recompute_granularity: SELECTIVE
+  tp: 2
+  tp_linear_async_communication: true
+  tp_mode: REDUCE_SCATTER
+profiler: null
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: mistralai/Mistral-7B-v0.1
+  tokenizer_revision: null
+tokens: null

config_tiny_mistral.py

@@ -26,41 +26,12 @@ from nanotron.config import (
 )
 from nanotron.logging import human_format
 
-
-@dataclass
-class MistralConfig:
-    """Configuration for a MISTRAL model
-
-    Be careful on having a coherent typing as we use it to reconstruct the model from yaml
-    """
-
-    bos_token_id: int = 1
-    eos_token_id: int = 2
-    hidden_act: str = "silu"
-    hidden_size: int = 4096
-    initializer_range: float = 0.02
-    intermediate_size: int = 11008
-    is_mistral_config: bool = True  # We use this help differentiate models in yaml/python conversion
-    max_position_embeddings: int = 2048
-    num_attention_heads: int = 32
-    num_hidden_layers: int = 32
-    num_key_value_heads: Optional[int] = None
-    pad_token_id: Optional[int] = None
-    pretraining_tp: int = 1
-    rms_norm_eps: float = 1e-6
-    rope_scaling: Optional[dict] = None
-    tie_word_embeddings: bool = False
-    use_cache: bool = True
-    vocab_size: int = 32000
-
-    def __post_init__(self):
-        # for backward compatibility
-        if self.num_key_value_heads is None:
-            self.num_key_value_heads = self.num_attention_heads
+from config_mistral import MistralConfig, get_num_params
 
 
 model_config = MistralConfig(
     # Config for a tiny model model with 1.62M parameters
+    attn_pdrop=0.0,
     bos_token_id=1,
     eos_token_id=2,
     hidden_act="silu",
@@ -73,20 +44,13 @@ model_config = MistralConfig(
     num_key_value_heads=4,
     pretraining_tp=1,
     rms_norm_eps=1e-05,
-    rope_scaling=None,
+    rope_theta=10000.0,
     tie_word_embeddings=True,
     use_cache=True,
     vocab_size=256,
 )
 
-num_params = human_format(
-    model_config.vocab_size * model_config.hidden_size * 2
-    + model_config.num_hidden_layers
-    * (
-        3 * model_config.hidden_size * model_config.intermediate_size
-        + 4 * model_config.hidden_size * model_config.hidden_size
-    )
-).replace(".", "p")
+num_params = human_format(get_num_params(model_config)).replace(".", "p")
 
 print(f"Model has {num_params} parameters")
 
@@ -141,9 +105,10 @@ config = Config(
 )
 
 if __name__ == "__main__":
-    dir = os.path.dirname(__file__)
+    file_path = os.path.abspath(__file__)
 
+    file_path = file_path.replace(".py", ".yaml")
     # Save config as YAML file
-    config.save_as_yaml(f"{dir}/config_tiny_mistral.yaml")
+    config.save_as_yaml(file_path)
 
     # You can now train a model with this config using `/run_train.py`

config_tiny_mistral.yaml

@@ -0,0 +1,92 @@
+checkpoints:
+  checkpoint_interval: 10
+  checkpoints_path: /fsx/thomwolf/github/textbooks-proj/brrr/models/checkpoints
+  checkpoints_path_is_shared_file_system: false
+  resume_checkpoint_path: null
+  save_initial_state: false
+data:
+  dataset:
+    dataset_overwrite_cache: false
+    dataset_processing_num_proc_per_process: 1
+    hf_dataset_config_name: null
+    hf_dataset_or_datasets: HuggingFaceH4/testing_alpaca_small
+    hf_dataset_splits: train
+    text_column_name: completion
+  num_loading_workers: 1
+  seed: 42
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: false
+  project: debug
+  run: tiny_mistral
+  seed: 42
+  step: null
+logging:
+  iteration_step_info_interval: 1
+  log_level: info
+  log_level_replica: info
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: bfloat16
+  init_method:
+    std: 0.025
+  make_vocab_size_divisible_by: 1
+  model_config:
+    attn_pdrop: 0.0
+    bos_token_id: 1
+    eos_token_id: 2
+    hidden_act: silu
+    hidden_size: 16
+    initializer_range: 0.02
+    intermediate_size: 64
+    is_mistral_config: true
+    max_position_embeddings: 256
+    num_attention_heads: 4
+    num_hidden_layers: 2
+    num_key_value_heads: 4
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_theta: 10000.0
+    sliding_window_size: 4096
+    tie_word_embeddings: true
+    use_cache: true
+    vocab_size: 256
+optimizer:
+  accumulate_grad_in_fp32: true
+  adam_beta1: 0.9
+  adam_beta2: 0.95
+  adam_eps: 1.0e-08
+  clip_grad: 1.0
+  learning_rate_scheduler:
+    learning_rate: 0.0003
+    lr_decay_steps: 8
+    lr_decay_style: cosine
+    lr_warmup_steps: 2
+    lr_warmup_style: linear
+    min_decay_lr: 1.0e-05
+  torch_adam_is_fused: true
+  weight_decay: 0.01
+  zero_stage: 0
+parallelism:
+  dp: 2
+  pp: 2
+  pp_engine: 1f1b
+  recompute_granularity: SELECTIVE
+  tp: 2
+  tp_linear_async_communication: true
+  tp_mode: REDUCE_SCATTER
+profiler: null
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: gpt2
+  tokenizer_revision: null
+tokens:
+  batch_accumulation_per_replica: 1
+  limit_test_batches: 0
+  limit_val_batches: 0
+  micro_batch_size: 2
+  sequence_length: 32
+  train_steps: 10
+  val_check_interval: -1

convert_trfrs_to_brrr.py

@@ -0,0 +1,262 @@
+# ruff: noqa: E402
+"""
+This module converts a transformers LlamaForCausalLM to a brrr model
+
+Command:
+torchrun  --nproc_per_node=1 convert_trfrs_to_brrr.py \
+    --model_name  mistralai/Mistral-7B-v0.1 \
+        --save_path ./pretrained/Mistral-7B-v0.1
+"""
+import argparse
+import sys
+from dataclasses import asdict
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+
+from brrr.trainer import DistributedTrainer
+
+sys.path.append(Path(__file__).parent.parent.as_posix())
+import os
+
+from nanotron.parallel.parameters import NanotronParameter, sanity_check
+from nanotron.parallel.pipeline_parallel.engine import (
+    AllForwardAllBackwardPipelineEngine,
+)
+from nanotron.parallel.tensor_parallel.nn import TensorParallelLinearMode
+from transformers import MistralConfig as MistralConfig_trfs, MistralForCausalLM
+
+import nanotron.distributed as dist
+from nanotron.config import ParallelismArgs, RecomputeGranularity
+from nanotron.parallel.context import ParallelContext
+from nanotron.models import build_model
+from nanotron.trainer import mark_tied_parameters
+from nanotron.serialize import save_meta, save_weights, save
+
+from modeling_mistral import MistralForTraining
+from config_mistral_7b import PARALLELISM as PARALLELISM_BRRR, CONFIG as CONFIG_BRRR
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Convert transformers weights to brrr weights")
+    parser.add_argument("--model_name", type=str, default="mistralai/Mistral-7B-v0.1")
+    parser.add_argument("--save_path", type=str, default="pretrained/Mistral-7B-v0.1")
+    parser.add_argument("--dp", type=int, default=1)
+    parser.add_argument("--pp", type=int, default=1)
+    parser.add_argument("--tp", type=int, default=1)
+    return parser.parse_args()
+
+
+def permute_for_rotary(tensor, num_heads, per_head_hidden_size, hidden_size):
+    return (
+        tensor.view(num_heads, 2, per_head_hidden_size // 2, hidden_size)
+        .transpose(1, 2)
+        .contiguous()
+        .view(num_heads * per_head_hidden_size, hidden_size)
+    )
+
+
+def get_transformers_weight(
+    name: str, ref_module_state_dict: Dict[str, torch.Tensor], ref_module: MistralForCausalLM, get_grad: bool = False
+) -> torch.Tensor:
+    """From our brrr implementation, we get the equivalent tensor in transformers implementation"""
+    config = ref_module.config
+    brrr_prefix = "model."
+    assert name.startswith(brrr_prefix)
+    name = name[len(brrr_prefix) :]
+
+    path = name.split(".")
+    path.remove("pp_block")
+    name = ".".join(path)
+
+    if get_grad is False:
+
+        def get_tensor(path: str):
+            return ref_module_state_dict[path]
+
+        def get_tensors(path: List[str]):
+            return [get_tensor(p) for p in path]
+
+    else:
+
+        def get_tensor(path: str):
+            weight = ref_module.get_parameter(path)
+            return weight.grad
+
+        def get_tensors(path: List[str]):
+            return [get_tensor(p) for p in path]
+
+    if name == "token_position_embeddings.token_embedding.weight":
+        return get_tensor("model.embed_tokens.weight")
+
+    elif name == "lm_head.weight":
+        # This only used when weights are not shared
+        return get_tensor("lm_head.weight")
+
+    elif name == "final_layer_norm.weight":
+        return get_tensor("model.norm.weight")
+
+    if path[0] == "decoder":
+        transformer_path = ["model"] + ["layers"] + [path[1]]
+
+        if path[2] == "attn":
+            path[2] = "self_attn"
+
+        if path[2] == "ff":
+            path[2] = "mlp"
+
+        if path[3] == "qkv_proj":
+            proj_names = ["q_proj", "k_proj", "v_proj"]
+            tensor_list = get_tensors(
+                [".".join(transformer_path + path[2:3] + [proj_name] + path[4:]) for proj_name in proj_names]
+            )
+            # Permute q/k
+            per_head_hidden_size = config.hidden_size // config.num_attention_heads
+            # Permute q
+            print(f"Permuting q {tensor_list[0].shape}")
+            tensor_list[0] = permute_for_rotary(
+                tensor=tensor_list[0],
+                num_heads=config.num_attention_heads,
+                per_head_hidden_size=per_head_hidden_size,
+                hidden_size=config.hidden_size,
+            )
+            # Permute k
+            print(f"Permuting k {tensor_list[1].shape}")
+            tensor_list[1] = permute_for_rotary(
+                tensor=tensor_list[1],
+                num_heads=config.num_key_value_heads,
+                per_head_hidden_size=per_head_hidden_size,
+                hidden_size=config.hidden_size,
+            )
+            return torch.cat(tensor_list, dim=0)
+
+        if path[3] == "gate_up_proj":
+            tensor_list = get_tensors(
+                [
+                    ".".join(transformer_path + path[2:3] + [proj_name] + path[4:])
+                    for proj_name in ["gate_proj", "up_proj"]
+                ]
+            )
+            return torch.cat(tensor_list, dim=0)
+
+        return get_tensor(".".join(transformer_path + path[2:]))
+
+    else:
+        raise ValueError(f"Couldn't find transformer equivalent of {name}")
+
+
+def convert_trfrs_to_brrr(dp, pp, tp, model_name="huggyllama/llama-7b", save_path="pretrained/llama-7b"):
+    # check save_path doesnt exist or is empty
+    save_path = Path(save_path)
+    # assert not save_path.exists() or len(list(save_path.iterdir())) == 0, f"save_path {save_path} is not empty"
+
+    parallel_config = PARALLELISM_BRRR
+    
+    parallel_config.dp = dp
+    parallel_config.pp = pp
+    parallel_config.tp = tp
+
+    # Initialise all process groups
+    parallel_context = ParallelContext(
+        data_parallel_size=parallel_config.dp,
+        pipeline_parallel_size=parallel_config.pp,
+        tensor_parallel_size=parallel_config.tp,
+    )
+    # params
+    dtype = torch.bfloat16  # Flash attention doesn't support fp32
+
+    # Initialise brrr model
+    model_config_brrr = CONFIG_BRRR.model.model_config
+
+    model = build_model(
+        model_builder=lambda: MistralForTraining(
+            config=model_config_brrr,
+            parallel_context=parallel_context,
+            parallel_config=parallel_config,
+            random_states=None,
+        ),
+        dtype=dtype,
+        parallel_context=parallel_context,
+        device=torch.device("cpu"),
+    )
+
+    # Initialise transformers model
+    device_map = {}
+    current_pp_rank = dist.get_rank(group=parallel_context.pp_pg)
+    device_map["model.embed_tokens"] = (
+        model.model.token_position_embeddings.rank
+        if current_pp_rank == model.model.token_position_embeddings.rank
+        else "meta"
+    )
+    for i in range(model_config_brrr.num_hidden_layers):
+        device_map[f"model.layers.{i}"] = (
+            model.model.decoder[i].rank if current_pp_rank == model.model.decoder[i].rank else "meta"
+        )
+    device_map["model.norm"] = (
+        model.model.final_layer_norm.rank if current_pp_rank == model.model.final_layer_norm.rank else "meta"
+    )
+    device_map["lm_head"] = model.model.lm_head.rank if current_pp_rank == model.model.lm_head.rank else "meta"
+    model_ref = MistralForCausalLM.from_pretrained(model_name, torch_dtype=dtype, device_map=device_map)
+
+    # Copy weights from trfrs to brrr
+    ref_state_dict = model_ref.state_dict()
+    for name, param in model.named_parameters():
+        print(f"Syncing {name}")
+        ref_param = get_transformers_weight(name=name, ref_module_state_dict=ref_state_dict, ref_module=model_ref)
+
+        param_is_tp_sharded = (
+            isinstance(param, NanotronParameter)
+            and param.is_sharded
+            and parallel_context.world_ranks_to_pg[param.get_sharded_info().global_ranks] == parallel_context.tp_pg
+        )
+
+        if param_is_tp_sharded:
+            sharded_info = param.get_sharded_info()
+            # copy param data (not just the reference)
+            with torch.no_grad():
+                for local_global_slices_pair in sharded_info.local_global_slices_pairs:
+                    local_slices = local_global_slices_pair.local_slices
+                    global_slices = local_global_slices_pair.global_slices
+                    param[local_slices].copy_(ref_param[global_slices])
+        else:
+            assert (
+                ref_param.shape == param.shape
+            ), f"Parameter shape don't match for {name}\n{ref_param.shape} != {param.shape}"
+            # copy param data (not just the reference)
+            with torch.no_grad():
+                param.copy_(ref_param)
+                ref_param = None
+                # torch.cuda.empty_cache()
+
+    # TODO @nouamanetazi: assert weights are the same
+    # Marks parameters as NanotronParameters
+    mark_tied_parameters(model=model, parallel_context=parallel_context, parallel_config=parallel_config)
+
+    sanity_check(root_module=model)
+
+    checkpoint_metadata = {
+        "last_train_step": 0,
+        "consumed_train_samples": 0,
+    }
+    save(config=CONFIG_BRRR, model=model, optimizer=None, lr_scheduler=None, parallel_context=parallel_context, root_folder=save_path,
+         should_save_optimizer=False, should_save_lr_scheduler=False, checkpoint_metadata=checkpoint_metadata,
+         sanity_checks=False)
+    # save_weights(model=model, parallel_context=parallel_context, root_folder=save_path)
+    # save_meta(root_folder=save_path, parallel_context=parallel_context, checkpoint_metadata=checkpoint_metadata)
+
+    if dist.get_rank(parallel_context.world_pg) == 0:
+        print(save_path)
+        import json
+
+        with open(save_path / "model_config.json", mode="w") as fo:
+            fo.write(json.dumps(asdict(CONFIG_BRRR.model.model_config), indent=4))
+
+
+def main():
+    args = get_args()
+    convert_trfrs_to_brrr(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

modeling_mistral.py

@@ -15,6 +15,7 @@
 """ PyTorch Mistral model.
 """
 from typing import Dict, Optional, Union
+import inspect
 
 import torch
 from flash_attn import bert_padding
@@ -46,12 +47,15 @@ from nanotron.parallel.tensor_parallel.nn import (
 )
 from nanotron.random import RandomStates
 from nanotron.utils import checkpoint_method
+from nanotron.nn.activations import ACT2FN
 from torch import nn
-from transformers import MistralConfig
-from transformers.activations import ACT2FN
+
+from config_mistral_7b import MistralConfig
 
 logger = logging.get_logger(__name__)
 
+_flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_varlen_func).parameters)
+
 
 class RotaryEmbedding(nn.Module):
     def __init__(self, dim: int, end: int, theta: float = 10000.0):
@@ -189,15 +193,22 @@ class CoreAttention(nn.Module):
         ), f"Hidden size {config.hidden_size} must be divisible by number of attention heads {config.num_attention_heads}."
         self.d_qk = config.hidden_size // config.num_attention_heads
         self.d_v = config.hidden_size // config.num_attention_heads
+        self.dropout = config.attn_pdrop
 
         self.checkpoint_attention = False  # Because flash_attn already does checkpointing
 
+        if config.sliding_window_size is not None:
+            assert (
+                _flash_supports_window_size
+            ), "Current version of flash-attn doesn't support sliding window: `pip install flash-attn>=2.3`"
+        self.sliding_window_size = config.sliding_window_size  # if layer_idx not in config.global_attn_layers else None
+
     @checkpoint_method(attr_name="checkpoint_attention")
     def forward(
         self,
-        query_states: torch.Tensor,  # [batch_size * q_length, n_local_q_heads, inner_dim]
-        key_states: torch.Tensor,  # [batch_size * kv_length, n_local_kv_heads, inner_dim]
-        value_states: torch.Tensor,  # [batch_size * kv_length, n_local_kv_heads, inner_dim]
+        query_states: torch.Tensor,  # [batch_size * q_length, num_heads, inner_dim]
+        key_states: torch.Tensor,  # [batch_size * kv_length, 1, inner_dim]
+        value_states: torch.Tensor,  # [batch_size * kv_length, 1, inner_dim]
         q_sequence_mask: torch.Tensor,  # torch.BoolTensor [batch_size, q_length] (can be broadcasted to that size)
         kv_sequence_mask: torch.Tensor,  # torch.BoolTensor [batch_size, kv_length] (can be broadcasted to that size)
     ):
@@ -218,9 +229,10 @@ class CoreAttention(nn.Module):
             cu_seqlens_k=cu_seqlens_k,
             max_seqlen_q=q_sequence_mask.shape[1],
             max_seqlen_k=kv_sequence_mask.shape[1],
-            dropout_p=0.0,
-            softmax_scale=None,  # This already defaults to the scale I'm interested in
+            dropout_p=self.dropout if self.training else 0.0,
+            softmax_scale=None,  # defaults to 1/sqrt(d_qk)
             causal=causal,
+            window_size=(self.sliding_window_size - 1, 0) if self.sliding_window_size is not None else (-1, -1),
             return_attn_probs=False,
         )
 
@@ -318,10 +330,11 @@ class CausalSelfAttention(nn.Module, AttachableStore):
         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,
@@ -852,7 +865,6 @@ class MistralForTraining(NanotronModel):
         super().__init__()
         import warnings
 
-        warnings.warn("This is just a Llama Model, not a Mistral one for demo purpose. Please fix implementation")
         self.model = MistralModel(config=config, parallel_context=parallel_context, parallel_config=parallel_config)
         self.loss = PipelineBlock(
             p2p=self.model.p2p,
@@ -1044,12 +1056,13 @@ def get_flops(
     num_layers,
     hidden_size,
     num_heads,
-    num_key_value_heads,
     vocab_size,
     seq_len,
-    ffn_hidden_size,
+    kv_channels=None,
+    ffn_hidden_size=None,
     batch_size=1,
     recompute_granularity=None,
+    glu_activation=False,
 ):
     """Counts flops in an decoder-only model
     Args:
@@ -1066,33 +1079,43 @@ def get_flops(
         model_flops: flops in the model (should be independent of the hardware and model implementation)
         hardware_flops: flops in the hardware (actual flops performed on the hardware). Check 6.3 in https://arxiv.org/pdf/2205.05198.pdf
     """
-    if num_key_value_heads is None:
-        num_key_value_heads = num_heads
-    hidden_size_per_head = hidden_size // num_heads
+    if kv_channels is None:
+        assert hidden_size % num_heads == 0
+        kv_channels = hidden_size // num_heads
+    if ffn_hidden_size is None:
+        ffn_hidden_size = 4 * hidden_size
+
     # In the following we mark the reduced dimension with parentheses
     # decoder
-    # self attention
-    ## qkv projection
-    decoder_qkv_proj_flops_fwd = (
-        2 * num_layers * batch_size * seq_len * (hidden_size) * num_heads * hidden_size_per_head
-        + 2 * num_layers * batch_size * seq_len * (hidden_size) * 2 * num_key_value_heads * hidden_size_per_head
-    )
+    # self attention (MQA)
+    ## q projection
+    decoder_q_proj_flops_fwd = 2 * num_layers * batch_size * seq_len * (hidden_size) * num_heads * kv_channels
+    ## kv projection, shared across heads
+    decoder_kv_proj_flops_fwd = 2 * num_layers * batch_size * seq_len * (hidden_size) * 2 * kv_channels
     ## qk logits
-    decoder_qk_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (hidden_size_per_head) * seq_len
+    decoder_qk_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (kv_channels) * seq_len
+    ### SWA (sliding window attention / local attention)
+    # window_size = 4096
+    # decoder_qk_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (kv_channels) * window_size
     ## v logits
-    decoder_v_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (seq_len) * hidden_size_per_head
+    decoder_v_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (seq_len) * kv_channels
+    # decoder_v_logits_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (window_size) * kv_channels
     ## attn out
-    decoder_attn_out_flops_fwd = (
-        2 * num_layers * batch_size * num_heads * seq_len * (hidden_size_per_head) * hidden_size
-    )
+    decoder_attn_out_flops_fwd = 2 * num_layers * batch_size * num_heads * seq_len * (kv_channels) * hidden_size
     # FF
     ## 1st layer
-    decoder_ffn_1_flops_fwd = 4 * num_layers * batch_size * seq_len * (hidden_size) * ffn_hidden_size
+    decoder_ffn_1_flops_fwd = 2 * num_layers * batch_size * seq_len * (hidden_size) * ffn_hidden_size
+    if glu_activation:
+        # 3 matmuls instead of 2 in FFN
+        # ref. https://arxiv.org/pdf/2002.05202.pdf
+        # Used for example in T5 v1.1
+        decoder_ffn_1_flops_fwd = 4 * num_layers * batch_size * seq_len * (hidden_size) * ffn_hidden_size
     ## 2nd layer
     decoder_ffn_2_flops_fwd = 2 * num_layers * batch_size * seq_len * (ffn_hidden_size) * hidden_size
 
     decoder_flops_fwd = (
-        decoder_qkv_proj_flops_fwd
+        decoder_q_proj_flops_fwd
+        + decoder_kv_proj_flops_fwd
         + decoder_qk_logits_flops_fwd
         + decoder_v_logits_flops_fwd
         + decoder_attn_out_flops_fwd

pretrained/Mistral-7B-v0.1/checkpoint_metadata.json

@@ -0,0 +1,9 @@
+{
+  "dp": 1,
+  "metas": {
+    "consumed_train_samples": 0,
+    "last_train_step": 0
+  },
+  "tp": 1,
+  "version": "1.2"
+}

pretrained/Mistral-7B-v0.1/config.yaml

@@ -0,0 +1,53 @@
+checkpoints: null
+data: null
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: false
+  project: mistralai
+  run: Mistral-7B-v0.1
+  seed: 42
+  step: null
+logging: null
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: bfloat16
+  init_method:
+    std: 0.025
+  make_vocab_size_divisible_by: 1
+  model_config:
+    attn_pdrop: 0.0
+    bos_token_id: 1
+    eos_token_id: 2
+    hidden_act: silu
+    hidden_size: 4096
+    initializer_range: 0.02
+    intermediate_size: 14336
+    is_mistral_config: true
+    max_position_embeddings: 32768
+    num_attention_heads: 32
+    num_hidden_layers: 32
+    num_key_value_heads: 8
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_theta: 10000.0
+    sliding_window_size: 4096
+    tie_word_embeddings: false
+    use_cache: true
+    vocab_size: 32000
+optimizer: null
+parallelism:
+  dp: 1
+  pp: 1
+  pp_engine: 1f1b
+  recompute_granularity: SELECTIVE
+  tp: 1
+  tp_linear_async_communication: true
+  tp_mode: REDUCE_SCATTER
+profiler: null
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: mistralai/Mistral-7B-v0.1
+  tokenizer_revision: null
+tokens: null

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e6ef1be3d2daa611724f02567159bf507c9a9ea276d9771387a01f4942cafb6
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3d4484e1f9505c97b6ac37472d9526e95470e6aef462fec6ae461b63e4ff77a
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c286c58dfce1f3f030c13b90c41c831d05c4323da3d50e23fe434f38f81535b
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d69d412c683fc926a26a71536dd2530877994cfa6e4e9ae3f3a31f6861596b0
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6148707c761016f313ee2ada71e15b9eb977878148fa382806eea4ef30a145e6
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/0/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05588e50e8fafc16c332e2f7a3d3830c9e59d29c35858d439a98ba4e418eba78
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cffeb63cbb1f271fd7ab983b618dfe4a4fc2b6b3763b9332fc324d378207210d
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a1263d75614c4a9710ebc5a57fdec732b9348c1f57ace1887ce296e1805b529
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71cd7738481e6bcbacbc76ce206545fb2fe6d995f7e1a733b408c3fe92f7356c
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31de05f9c50d9e94fe67936a973c86840f82ed2aad1494806baa81df8bbf9bf8
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70761ee840fbdc950501814ff397c72e9d8bbc7be2030329f391c12eb5b73a0f
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/1/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1846ddc1c4ca9d8e03184f2fa34911398202f0edc310df5ea408a323a5f23ee8
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59d82cc1e5ec1f364e7401e17e58d53f62a39658799aeb4902060236ebb0cb60
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:117c7832cefd9a767282b27c852f00ed4ce7888a8abb7e2f9257a0b2fed60608
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c93dc35263c0efa22d22795777c009e4f9365cf1ef413b69880d14433d1069e8
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad3fefd66e98df8fee62bd0fe451b18ca1a14545b72e570d499dce0464368b81
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6346f7c23987a4c529ac1b63b5f6f56b4392981ffcaaf2cb84cf2bf5b2bc36a7
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/10/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6cd70662e84b3d81b4f4512929d00d9377515c2dfe75d78109edce27c57d834
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c17ee146b384be81a4b9cb06960728dd540d6650a5798abcc95315bb0daf2ca
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:497bbea1882c42d134dc4458194d71cd3d7d609b06e91e715093e0c243962116
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7fbc23c909758daf76a1d647d1beefb4c3cc62a4aa04f98679e22d79cc6813e
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef4f8de0f891e6d79255af98deda246f04c0775694835f696a1a8b0738f492da
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:844adcbf23618ae38fbffaf92f7d38ce7d853be5643485bb25f8f839c0f2819c
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/11/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1823cbba03a4ec4683cc6a41eab51e34cec90e92cea7af388d0c675abe451284
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9da7a13da9a196108d2efd875884aa8629533e8143255eef5915417ac592d9c0
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31c8fb0c121f6977e10c7277544259157152d28de9559c8aae8236925398329f
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90fd4f646b1f5ca201781cc77b713093ab9a67d4ee8de11c669a486a2896d773
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:368a3142cb9f085a2da1db74de226b13c509467cbea81da25f27db8842347443
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48e00b1107d1e575c2425fa8368e92eb714b59825153206ae4cccc36eb4e8e45
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/12/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c3300f1d0af46ca69fc14397728055e302b2955b8b9adfd9705b68a683377b1
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0d2409179997ff51079156414cb112c82b964976a8023f5088b1dd7ab28f50c
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c51efd448a50b2c75046bfb12f2703ce19e56b07f4f9e94f7a36f673c70517b8
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe220fd602b0f41f30f7bca607c400adacadb7b5e31f81a28d7e103fd5c0b0a8
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa8d2abd973723ddb4ea4cb2188fa767ea55168bc2519e170e728619fde864c4
+size 117440752

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/mlp/gate_up_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeeb6ae12119f5ffd390add932122f819d81d50886389567eb41103451b36d24
+size 234881328

pretrained/Mistral-7B-v0.1/model/model/decoder/13/pp_block/post_attention_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55a66303f00f120db3fba8ae4d13725fc2c22cd6fd3babc1d66dc0fdee7eb45f
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/14/pp_block/attn/o_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9fe17a044a248163b4c45783386f7d414e6217ae9657c8983d54a84e85aae8e
+size 33554672

pretrained/Mistral-7B-v0.1/model/model/decoder/14/pp_block/attn/qkv_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b727021f1c6625dae8fb240904bf838985311afe7b5e19f62839c3563072e75a
+size 50332000

pretrained/Mistral-7B-v0.1/model/model/decoder/14/pp_block/input_layernorm/model_weight.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5f069023e761eacf5747a540516a9215c49e3efc8614ffaa7fa4ca016c67075
+size 8288

pretrained/Mistral-7B-v0.1/model/model/decoder/14/pp_block/mlp/down_proj/model_weight_pp-rank-0-of-1_tp-rank-0-of-1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:993eebcd3451163de4981bd6f7cd82a2bc0bfcc29a67cbbedceb502a5036466c
+size 117440752