update

.gitignore

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

README.md

@@ -0,0 +1,37 @@
+---
+library_name: nanotron
+---
+
+# ⚙️ Nano-Mistral
+
+Modeling code for Mistral to use with [Nanotron](https://github.com/huggingface/nanotron/)
+
+Also contains converted pretrained weights for Mistral-7B-0.1: https://huggingface.co/mistralai/Mistral-7B-v0.1
+
+## 🚀 Quickstart
+
+```bash
+# Generate a config file
+python config_tiny_mistral.py
+
+# Run training
+export CUDA_DEVICE_MAX_CONNECTIONS=1 # important for some distributed operations
+torchrun --nproc_per_node=8 run_train.py --config-file config_tiny_mistral.yaml
+```
+
+## 🚀 Run generation with pretrained Mistral-7B-0.1
+
+```bash
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+torchrun --nproc_per_node=1 run_generate.py --ckpt-path ./pretrained/Mistral-7B-v0.1
+```
+
+## 🚀 Use your custom model
+
+- Update the `MistralConfig` class in `config_tiny_mistral.py` to match your model's configuration
+- Update the `MistralForTraining` class in `modeling_mistral.py` to match your model's architecture
+- Pass the previous to the `DistributedTrainer` class in `run_train.py`:
+```python
+trainer = DistributedTrainer(config_file, model_class=MistralForTraining, model_config_class=MistralConfig)
+```
+- Run training as usual

config_minicpm.py

@@ -0,0 +1,119 @@
+""" 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 MiniCPMConfig:
+    """Configuration for a MiniCPM 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
+    pad_token_id: Optional[int] = None
+    hidden_act: str ="silu"
+    hidden_size: int =2304
+    initializer_range: float =0.1
+    intermediate_size: int =5760
+    max_position_embeddings: int =2048
+    num_attention_heads: int =36
+    num_hidden_layers: int =40
+    num_key_value_heads: int =36
+    pretraining_tp: int=1
+    rms_norm_eps: float=1e-05
+    rope_theta: float = 10000.0
+    tie_word_embeddings: bool =True
+    use_cache: bool =True
+    vocab_size: int = 122753
+    scale_emb: float = 12
+    dim_model_base: int= 256
+    scale_depth: float = 1.4
+
+    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: MiniCPMConfig) -> 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
+
+def get_num_params_no_embed(model_config: MiniCPMConfig) -> int:
+    num_params = 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
+
+MODEL_CONFIG = MiniCPMConfig()
+
+num_params = human_format(get_num_params(MODEL_CONFIG)).replace(".", "p")
+num_params_no_embed = human_format(get_num_params_no_embed(MODEL_CONFIG)).replace(".", "p")
+
+print(f"Model has {num_params} parameters or {num_params_no_embed} without embeddings")
+
+PARALLELISM = ParallelismArgs(
+    dp=1,
+    pp=1,
+    tp=1,
+    pp_engine="1f1b",
+    tp_mode="REDUCE_SCATTER",
+    tp_linear_async_communication=True,
+    recompute_granularity="selective",
+)
+
+CONFIG = Config(
+    general=GeneralArgs(project="openbmb", run="MiniCPM-2B-dpo-bf16", seed=42, step=0),
+    checkpoints=None,
+    parallelism=PARALLELISM,
+    model=ModelArgs(init_method=RandomInit(std=0.025), model_config=MODEL_CONFIG),
+    tokenizer=TokenizerArgs("openbmb/MiniCPM-2B-dpo-bf16"),
+    optimizer=None,
+    logging=None,
+    tokens=None,
+    data=None,
+    profiler=None,
+    lighteval=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_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  openbmb/MiniCPM-2B-dpo-bf16 \
+        --save_path ./pretrained/MiniCPM-2B-dpo-bf16
+"""
+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_minicpm import MiniCPMForTraining
+from config_minicpm 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="openbmb/MiniCPM-2B-dpo-bf16")
+    parser.add_argument("--save_path", type=str, default="pretrained/MiniCPM-2B-dpo-bf16")
+    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: MiniCPMForTraining(
+            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()

dataloader.py

@@ -0,0 +1,107 @@
+from nanotron import logging
+from nanotron.config import (
+    PretrainDatasetsArgs,
+)
+from nanotron.dataloader import (
+    clm_process,
+    dummy_infinite_data_generator,
+    get_datasets,
+    get_train_dataloader,
+)
+from nanotron.logging import log_rank
+from nanotron.parallel.pipeline_parallel.utils import get_input_output_pp_ranks
+from nanotron.trainer import DistributedTrainer
+from nanotron.utils import (
+    main_rank_first,
+)
+
+try:
+    from huggingface_hub import __version__ as hf_hub_version
+    from transformers import AutoTokenizer
+    from transformers import __version__ as tf_version
+except ImportError:
+    hf_hub_version = None
+    tf_version = None
+
+logger = logging.get_logger(__name__)
+
+
+def get_dataloader(trainer: DistributedTrainer):
+    """Returns a dataloader for training."""
+
+    # First, we need to know which ranks to feed the dataloader to
+    input_pp_rank, output_pp_rank = get_input_output_pp_ranks(model=trainer.model)
+
+    # Case 1: Dummy data generator
+    if trainer.config.data.dataset is None:
+        log_rank("Using dummy data generator", logger=logger, level=logging.INFO, rank=0)
+        dataloader = dummy_infinite_data_generator(
+            micro_batch_size=trainer.micro_batch_size,
+            sequence_length=trainer.sequence_length,
+            input_pp_rank=input_pp_rank,
+            output_pp_rank=output_pp_rank,
+            vocab_size=trainer.model_config.vocab_size,
+            seed=trainer.config.data.seed,
+            parallel_context=trainer.parallel_context,
+        )()
+
+    # Case 2: HuggingFace datasets
+    elif isinstance(trainer.config.data.dataset, PretrainDatasetsArgs):
+        log_rank("Using `datasets` library", logger=logger, level=logging.INFO, rank=0)
+        tokenizer_path = trainer.config.tokenizer.tokenizer_name_or_path
+        log_rank(
+            f"Loading tokenizer from {tokenizer_path} and transformers/hf_hub versions {tf_version, hf_hub_version}",
+            logger=logger,
+            level=logging.INFO,
+            rank=0,
+        )
+
+        # We need to the 1st device to process dataset and cache it, then other devices load from cache
+        with main_rank_first(trainer.parallel_context.world_pg):
+            # TODO @nouamanetazi: this may timeout before 1st device finishes processing dataset. Can we have a ctxmanager to modify timeout?
+            # TODO: generalise to include  for validation/test splits
+
+            # We load the raw dataset
+            raw_dataset = get_datasets(
+                hf_dataset_or_datasets=trainer.config.data.dataset.hf_dataset_or_datasets,
+                splits=trainer.config.data.dataset.hf_dataset_splits,
+            )["train"]
+
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.padding_side = "left"
+
+            # We apply the Causal Language Modeling preprocessing
+            train_dataset = clm_process(
+                raw_dataset=raw_dataset,
+                tokenizer=tokenizer,
+                text_column_name=trainer.config.data.dataset.text_column_name,
+                dataset_processing_num_proc_per_process=trainer.config.data.dataset.dataset_processing_num_proc_per_process,
+                dataset_overwrite_cache=trainer.config.data.dataset.dataset_overwrite_cache,
+                sequence_length=trainer.sequence_length,
+            )
+
+            # We load the processed dataset on the ranks requiring it
+            dataloader = get_train_dataloader(
+                train_dataset=train_dataset,
+                sequence_length=trainer.sequence_length,
+                parallel_context=trainer.parallel_context,
+                input_pp_rank=input_pp_rank,
+                output_pp_rank=output_pp_rank,
+                micro_batch_size=trainer.micro_batch_size,
+                consumed_train_samples=trainer.consumed_train_samples,
+                dataloader_num_workers=trainer.config.data.num_loading_workers,
+                seed_worker=trainer.config.data.seed,
+                dataloader_drop_last=True,
+            )
+            # Check if we have enough samples for train_steps
+            assert (
+                trainer.config.tokens.train_steps - trainer.start_iteration_step
+            ) * trainer.global_batch_size // trainer.parallel_context.dp_pg.size() < len(dataloader), (
+                f"Dataset is too small for steps ({len(dataloader)} < {(trainer.config.tokens.train_steps - trainer.start_iteration_step) * trainer.global_batch_size // trainer.parallel_context.dp_pg.size()}), "
+                f"Try train_steps<={len(dataloader) * trainer.parallel_context.dp_pg.size() // trainer.global_batch_size + trainer.start_iteration_step}"
+            )
+    else:
+        raise ValueError(f"Unhandled case of `self.config.data.dataset`. Got: {trainer.config.data.dataset}")
+
+    return dataloader

modeling_minicpm.py

@@ -0,0 +1,1147 @@
+# coding=utf-8
+# Copyright 2018 HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch MiniCPM model.
+"""
+from typing import Dict, Optional, Union
+import inspect
+import math
+
+import torch
+from flash_attn import bert_padding
+from flash_attn.flash_attn_interface import (
+    flash_attn_varlen_func,
+    flash_attn_with_kvcache,
+)
+from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding
+from nanotron import distributed as dist
+from nanotron import logging
+from nanotron.config import ParallelismArgs, RecomputeGranularity
+from nanotron.generation.generate_store import AttachableStore
+from nanotron.logging import log_rank
+from nanotron.models import NanotronModel
+from nanotron.nn.layer_norm import TritonRMSNorm
+from nanotron.parallel import ParallelContext
+from nanotron.parallel.parameters import NanotronParameter
+from nanotron.parallel.pipeline_parallel.block import (
+    PipelineBlock,
+    TensorPointer,
+)
+from nanotron.parallel.pipeline_parallel.p2p import P2P
+from nanotron.parallel.tensor_parallel.functional import sharded_cross_entropy
+from nanotron.parallel.tensor_parallel.nn import (
+    TensorParallelColumnLinear,
+    TensorParallelEmbedding,
+    TensorParallelLinearMode,
+    TensorParallelRowLinear,
+)
+from nanotron.random import RandomStates
+from nanotron.utils import checkpoint_method
+from nanotron.nn.activations import ACT2FN
+from torch import nn
+
+from config_minicpm import MiniCPMConfig
+
+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):
+        super().__init__()
+        assert dim % 2 == 0
+        self.dim = dim
+        self.end = end
+        self.theta = theta
+        # TODO @nouamane: Figure out why we can't set `DTypeInvariantTensor` ...
+        # TODO @thomasw21: Complex buffers break DDP, instead we store float and view them as complex
+        self.freqs_cis: torch.Tensor
+        self._initialized_buffer = False
+
+    def init_rotary_embeddings(self):
+        if self._initialized_buffer is True:
+            # Buffer if already initialized
+            return
+        self.register_buffer(
+            "freqs_cis",
+            torch.empty(self.end, self.dim // 2, 2, dtype=torch.float, device="cuda"),
+            persistent=False,
+        )
+        assert self.freqs_cis.device.type == "cuda"
+        # TODO @nouamane: One we figure out how to do the DTypeInvariantTensor, this can be removed and changed to an assert
+        if self.freqs_cis.dtype != torch.float:
+            self.freqs_cis = self.freqs_cis.to(torch.float)
+        assert self.freqs_cis.dtype == torch.float
+        freqs = 1.0 / (
+            self.theta
+            ** (torch.arange(0, self.dim, 2, dtype=torch.float, device="cuda")[: (self.dim // 2)] / self.dim)
+        )
+        t = torch.arange(self.end, device="cuda")
+        freqs = torch.outer(t, freqs).float()
+        complex_freqs = torch.polar(torch.ones_like(freqs), freqs)
+        freqs = torch.view_as_real(complex_freqs)
+        self.freqs_cis.copy_(freqs)
+        self._initialized_buffer = True
+
+    def forward(
+        self,
+        x: torch.Tensor,  # [batch_size, seq_length, num_heads, d_qk]
+        position_ids: Optional[torch.LongTensor],  # [batch_size, seq_length]
+    ):
+        batch_size, seq_length, num_heads, inner_dim = x.shape
+        while (
+            position_ids is not None and position_ids[-1, -1] >= self.end
+        ) or seq_length >= self.end:  # TODO @nouamane: check if this causes cpu-gpu sync
+            self.end *= 2
+            self._initialized_buffer = False
+        if self._initialized_buffer is False:
+            self.init_rotary_embeddings()
+        dtype = x.dtype
+        assert inner_dim % 2 == 0
+        x = x.view(
+            batch_size, seq_length, num_heads, inner_dim // 2, 2
+        )  # [batch_size, q_length, num_heads, inner_dim]
+        if x.dtype == torch.bfloat16:
+            x = x.float()
+        complex_x = torch.view_as_complex(x)  # [batch_size, q_length, num_heads, inner_dim // 2]
+        if position_ids is None:
+            freqs_cis = self.freqs_cis[None, :seq_length, None, :]
+        else:
+            # TODO(kunhao): Should None follow the num_heads dimension?
+            if position_ids[-1, -1] < 0 or position_ids[-1, -1] >= self.end:  # Quick test hopefully
+                raise ValueError(f"Position ids must be in the range [0, {self.end}), but got {position_ids}")
+            freqs_cis = self.freqs_cis[position_ids][:, :, None, :]
+        complex_freqs = torch.view_as_complex(freqs_cis)
+        x_out = torch.view_as_real(complex_x * complex_freqs).view(batch_size, seq_length, num_heads, inner_dim)
+        return x_out.type(dtype)
+
+
+class GLUActivation(nn.Module):
+    def __init__(self, act_fn_name: str):
+        super().__init__()
+        self.act = ACT2FN[act_fn_name]
+
+    def forward(self, merged_states: torch.Tensor):
+        gate_states, up_states = torch.split(merged_states, merged_states.shape[-1] // 2, dim=-1)
+        return self.act(gate_states) * up_states
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        config: MiniCPMConfig,
+        parallel_config: Optional[ParallelismArgs],
+        tp_pg: dist.ProcessGroup,
+    ):
+        super().__init__()
+
+        # TODO @thomasw21: refactor so that we store that default in a single place.
+        tp_mode = parallel_config.tp_mode if parallel_config is not None else TensorParallelLinearMode.ALL_REDUCE
+        tp_linear_async_communication = (
+            parallel_config.tp_linear_async_communication if parallel_config is not None else False
+        )
+
+        gate_up_contiguous_chunks = (
+            config.intermediate_size,  # shape of gate_linear
+            config.intermediate_size,  # shape of up_linear
+        )
+        self.gate_up_proj = TensorParallelColumnLinear(
+            config.hidden_size,
+            2 * config.intermediate_size,
+            pg=tp_pg,
+            mode=tp_mode,
+            bias=False,
+            async_communication=tp_linear_async_communication,
+            contiguous_chunks=gate_up_contiguous_chunks,
+        )
+
+        self.down_proj = TensorParallelRowLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            pg=tp_pg,
+            mode=tp_mode,
+            bias=False,
+            async_communication=tp_linear_async_communication and tp_mode is TensorParallelLinearMode.REDUCE_SCATTER,
+        )
+        # TODO @nouamane: why can't we torch.jit.script GLUActivation?
+        self.split_silu_mul = GLUActivation(config.hidden_act)
+
+    def forward(self, hidden_states):  # [seq_length, batch_size, hidden_dim]
+        merged_states = self.gate_up_proj(hidden_states)
+        hidden_states = self.down_proj(self.split_silu_mul(merged_states))
+        return {"hidden_states": hidden_states}
+
+
+class CoreAttention(nn.Module):
+    def __init__(self, config: MiniCPMConfig, parallel_config: Optional[ParallelismArgs], layer_idx: int):
+        super().__init__()
+        # TODO @thomasw21: GPT has a weird `d_kv` config which I'm guessing is essentically a `d_qkv`
+        assert (
+            config.hidden_size % config.num_attention_heads == 0
+        ), 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, 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)
+    ):
+        # TODO @thomasw21: Compute once, instead of computing for each layers.
+        cu_seqlens_q = torch.zeros((q_sequence_mask.shape[0] + 1), dtype=torch.int32, device=query_states.device)
+        cu_seqlens_k = torch.zeros((kv_sequence_mask.shape[0] + 1), dtype=torch.int32, device=query_states.device)
+        torch.cumsum(q_sequence_mask.sum(-1, dtype=torch.int32), dim=0, dtype=torch.int32, out=cu_seqlens_q[1:])
+        torch.cumsum(kv_sequence_mask.sum(-1, dtype=torch.int32), dim=0, dtype=torch.int32, out=cu_seqlens_k[1:])
+
+        # TODO(kunhao): flash attn's causal means that the query can only attend to the keys before it. This is not
+        # what we want if we are using kv cache. This is a hack as we always have q_length == 1 when using kv cache.
+        causal = False if q_sequence_mask.shape[1] == 1 else True
+        attn_output = flash_attn_varlen_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=q_sequence_mask.shape[1],
+            max_seqlen_k=kv_sequence_mask.shape[1],
+            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,
+        )
+
+        return attn_output
+
+
+def pad_to_right(tensor, mask, new_tensor=None):
+    """Transform a left-padded tensor into a right-padded tensor. (Useful for prefilling key/value states)
+    Args:
+        tensor: (batch_size, seqlen, d1, d2)
+        mask: (batch_size, seqlen)
+        new_tensor: (batch_size, new_tensor_seqlen, d1, d2)
+    Returns:
+        new_tensor: (batch_size, new_tensor_seqlen, d1, d2)
+        right_padded_mask: (batch_size, seqlen)
+    """
+    # First, we need to find the number of padding for each row
+    unpad_seqlens = mask.sum(1)
+    # Then, we need to find the maximum length of the tensor
+    max_seqlen = mask.shape[1]
+    # We can then create the indices to select the padded values
+    # The indices are the same for each row
+    indices = torch.arange(max_seqlen, device=mask.device)
+    # We can then create the mask for the padded values
+    right_padded_mask = indices < unpad_seqlens[:, None]
+    # We select the useful values
+    useful_values = tensor[mask]
+    # We create the new tensor (if not provided)
+    new_tensor = torch.zeros_like(tensor) if new_tensor is None else new_tensor
+    # We fill the new tensor with the useful values
+    new_tensor[:, : right_padded_mask.shape[1], :, :][right_padded_mask] = useful_values
+    return new_tensor, right_padded_mask
+
+
+class CausalSelfAttention(nn.Module, AttachableStore):
+    def __init__(
+        self,
+        config: MiniCPMConfig,
+        parallel_config: Optional[ParallelismArgs],
+        tp_pg: dist.ProcessGroup,
+        layer_idx: int,
+    ):
+        super().__init__()
+        # Tensor parallel considerations: We split tensors along head dimension
+        assert (
+            config.num_attention_heads % tp_pg.size() == 0
+        ), f"Number of attention heads ({config.num_attention_heads}) must be divisible by TP size ({tp_pg.size()})."
+        try:
+            assert (
+                config.num_key_value_heads % tp_pg.size() == 0
+            ), f"Number of key/value heads ({config.num_key_value_heads}) must be divisible by TP size ({tp_pg.size()})."
+        except AttributeError:
+            log_rank(
+                "WARNING: num_key_value_heads not defined, assuming it is equal to num_attention_heads",
+                logger=logger,
+                level=logging.WARNING,
+                rank=0,
+            )
+            # If num_key_value_heads is not defined, we assume that it is equal to num_attention_heads
+            config.num_key_value_heads = config.num_attention_heads
+        assert (
+            config.num_attention_heads % config.num_key_value_heads == 0
+        ), f"Number of attention heads ({config.num_attention_heads}) must be divisible by number of key/value heads ({config.num_key_value_heads})."
+        self.n_local_q_heads = config.num_attention_heads // tp_pg.size()
+        self.n_local_kv_heads = config.num_key_value_heads // tp_pg.size()
+        self.n_repeats = config.num_attention_heads // config.num_key_value_heads
+        self.is_gqa = config.num_attention_heads != config.num_key_value_heads  # Whether we are using GQA or not
+        self.d_qk = config.hidden_size // config.num_attention_heads
+        self.d_v = config.hidden_size // config.num_attention_heads
+        self.d_model = config.hidden_size
+
+        # TODO @thomasw21: refactor so that we store that default in a single place.
+        tp_mode = parallel_config.tp_mode if parallel_config is not None else TensorParallelLinearMode.ALL_REDUCE
+        tp_linear_async_communication = (
+            parallel_config.tp_linear_async_communication if parallel_config is not None else False
+        )
+
+        # build the slice config for self.qkv for save/load
+        # shard are done within the contiguous chunk
+        qkv_contiguous_chunks = (
+            config.num_attention_heads * self.d_qk,  # shape of q
+            config.num_key_value_heads * self.d_qk,  # shape of k
+            config.num_key_value_heads * self.d_qk,  # shape of v
+        )
+        self.qkv_proj = TensorParallelColumnLinear(
+            self.d_model,
+            config.num_attention_heads * self.d_qk + 2 * config.num_key_value_heads * self.d_qk,
+            pg=tp_pg,
+            mode=tp_mode,
+            bias=False,
+            async_communication=tp_linear_async_communication,
+            contiguous_chunks=qkv_contiguous_chunks,
+        )
+        # TODO(kunhao): We want to have only one version per device and not one version per layer.
+        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, base=config.rope_theta, interleaved=True)
+
+        self.o_proj = TensorParallelRowLinear(
+            config.num_attention_heads * self.d_qk,
+            self.d_model,
+            pg=tp_pg,
+            mode=tp_mode,
+            bias=False,
+            async_communication=tp_linear_async_communication,
+        )
+
+        self.attention = CoreAttention(
+            config,
+            parallel_config=parallel_config,
+            layer_idx=layer_idx,
+        )
+
+        self.prefill_kv_len = (
+            config.max_position_embeddings
+        )  # TODO @nouamane: compute based on free memory, because in rope we can surpass max_position_embeddings
+
+    def forward(
+        self,
+        hidden_states,  # [seq_length, batch_size, hidden_size]
+        sequence_mask,  # [batch_size, seq_length]
+    ):
+        qkv_states = self.qkv_proj(
+            hidden_states
+        )  # [seq_length, batch_size, n_local_q_heads * d_qk + 2 * n_local_kv_heads * d_qk]
+        q_length, batch_size, _ = qkv_states.shape
+
+        if self.is_gqa:
+            query_states, key_states, value_states = torch.split(
+                qkv_states,
+                [
+                    self.n_local_q_heads * self.d_qk,
+                    self.n_local_kv_heads * self.d_qk,
+                    self.n_local_kv_heads * self.d_qk,
+                ],
+                dim=-1,
+            )
+
+            query_states = (
+                query_states.transpose(0, 1).contiguous().view(batch_size, q_length, self.n_local_q_heads, self.d_qk)
+            )
+            key_states = (
+                key_states.transpose(0, 1).contiguous().view(batch_size, q_length, self.n_local_kv_heads, self.d_qk)
+            )
+            value_states = (
+                value_states.transpose(0, 1).contiguous().view(batch_size, q_length, self.n_local_kv_heads, self.d_qk)
+            )
+        else:
+            query_states, key_states, value_states = (
+                qkv_states.view(q_length, batch_size, 3, self.n_local_q_heads, self.d_qk)
+                .permute(2, 1, 0, 3, 4)
+                .contiguous()
+            )  # [3, batch_size, seq_length, n_local_q_heads, d_qk]
+
+        store = self.get_local_store()
+        if store is not None:  # Inference case
+            # Double check that we use store only at inference time
+            assert key_states.requires_grad is False
+            assert value_states.requires_grad is False
+            if "position_offsets" in store:
+                old_position_offsets = store["position_offsets"]
+                position_ids = old_position_offsets[:, None] + sequence_mask
+            else:
+                position_ids = torch.cumsum(sequence_mask, dim=-1, dtype=torch.int32) - 1
+            position_offsets = position_ids[:, -1]
+
+            # Compute rotary embeddings
+            # Note: keep track of old rotary embedding end to check if we need to enlarge k_cache and v_cache
+            old_rotary_embed_end = self.rotary_embedding.end
+            query_states = self.rotary_embedding(query_states, position_ids=position_ids)
+            key_states = self.rotary_embedding(key_states, position_ids=position_ids)
+
+            if "key" not in store:
+                # First inference iteration (Prefill)
+                # TODO @nouamane: support custom masking
+                # assert that [ False, False, False, False,  True,  True,  True,  True,  True,  True] is accepted
+                # but [ False, False, False, False,  True,  True,  False,  False,  True,  True] is not (can't mask in the middle of sequence)
+                assert ~(
+                    sequence_mask[:, :-1] & (~sequence_mask[:, 1:])  # True is never followed by False
+                ).any(), "Can't mask in the middle of sequence, please make sure that pads are at the left of the sequence if existing"
+
+                # preallocate k_cache, v_cache to self.prefill_kv_len
+                k_cache = torch.zeros(
+                    (
+                        batch_size,
+                        self.prefill_kv_len,
+                        self.n_local_kv_heads,
+                        self.d_qk,
+                    ),
+                    dtype=query_states.dtype,
+                    device=query_states.device,
+                )
+                v_cache = torch.zeros(
+                    (batch_size, self.prefill_kv_len, self.n_local_kv_heads, self.d_v),
+                    dtype=query_states.dtype,
+                    device=query_states.device,
+                )
+                # Remove pad tokens from key_states and concatenate samples in key_unpad
+                # cu_seqlens_k is the cumulative sequence lengths of key_states
+                (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(
+                    query_states,
+                    sequence_mask,
+                )
+                (key_unpad, indices_k, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(
+                    key_states, sequence_mask
+                )
+                (value_unpad, _, _, _) = bert_padding.unpad_input(value_states, sequence_mask)
+
+                output_unpad = flash_attn_varlen_func(
+                    q=query_unpad,  # (total_q, n_local_q_heads, d_qk)
+                    k=key_unpad,  # (total_kv, n_local_kv_heads, d_qk)
+                    v=value_unpad,  # (total_kv, n_local_kv_heads, d_v)
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    dropout_p=0.0,
+                    softmax_scale=None,
+                    causal=True,  # True in prefill phase, False in subsequent phases
+                    return_attn_probs=False,
+                )  # (total_unpadded, n_local_q_heads, d_v)
+
+                attention_output = bert_padding.pad_input(
+                    output_unpad, indices_q, batch_size, q_length
+                )  # (batch_size, q_length, n_local_q_heads, d_v)
+
+                pad_to_right(key_states, sequence_mask, new_tensor=k_cache)
+                pad_to_right(value_states, sequence_mask, new_tensor=v_cache)
+
+            else:
+                # Pull pre-computed key/value states
+                # Subsequent inference iterations (q_length=1)
+                k_cache = store["key"]
+                v_cache = store["value"]
+
+                # NOTE(fmom): According to flash_attn_with_kvcache, "If you pass in k / v, you must make sure that the cache is large enough to hold the new values"
+                # Since rotary embedding has changed (to enable larger context), we need to enlarge k_cache and v_cache
+                if self.rotary_embedding.end > old_rotary_embed_end:
+                    k_cache = torch.cat(
+                        [
+                            k_cache,
+                            torch.zeros(
+                                (
+                                    batch_size,
+                                    self.rotary_embedding.end - old_rotary_embed_end,
+                                    self.n_local_kv_heads,
+                                    self.d_qk,
+                                ),
+                                dtype=query_states.dtype,
+                                device=query_states.device,
+                            ),
+                        ],
+                        dim=1,
+                    )
+
+                    v_cache = torch.cat(
+                        [
+                            v_cache,
+                            torch.zeros(
+                                (
+                                    batch_size,
+                                    self.rotary_embedding.end - old_rotary_embed_end,
+                                    self.n_local_kv_heads,
+                                    self.d_v,
+                                ),
+                                dtype=query_states.dtype,
+                                device=query_states.device,
+                            ),
+                        ],
+                        dim=1,
+                    )
+
+                assert (
+                    k_cache.shape[1] == self.rotary_embedding.end
+                ), f"Cache size {k_cache.shape[1]} is smaller than rotary embedding end {self.rotary_embedding.end}"
+                assert (
+                    v_cache.shape[1] == self.rotary_embedding.end
+                ), f"Cache size {v_cache.shape[1]} is smaller than rotary embedding end {self.rotary_embedding.end}"
+
+                # [batch_size, seq_length, num_heads, d_qk]
+                query_states = query_states.view(
+                    batch_size, q_length, self.n_local_q_heads, self.d_qk
+                )  # [batch_size, q_length, self.n_heads, d_qk]
+                kv_length = key_states.shape[1]
+                key_states = key_states.view(
+                    batch_size, kv_length, self.n_local_kv_heads, self.d_qk
+                )  # [batch_size, kv_length, self.n_heads, d_qk]
+                value_states = value_states.view(
+                    batch_size, kv_length, self.n_local_kv_heads, self.d_v
+                )  # [batch_size, kv_length, self.n_heads, d_v]
+
+                attention_output = flash_attn_with_kvcache(
+                    query_states,
+                    k_cache,
+                    v_cache,
+                    key_states,
+                    value_states,
+                    rotary_cos=None,
+                    rotary_sin=None,
+                    # TODO @nouamane: seems like this doesnt help to indicate padding in (for first iteration it's just 0)
+                    cache_seqlens=position_offsets.contiguous(),
+                    softmax_scale=None,
+                    causal=True,
+                    rotary_interleaved=False,  # GPT-NeoX style
+                )
+
+            store.update(
+                {
+                    "key": k_cache,  # flash-attn has updated with new key_states using cache_seqlens
+                    "value": v_cache,
+                    "position_offsets": position_offsets,
+                }
+            )
+
+        else:  # Training case
+            # Apply rotary embeddings to query/key states
+            # NOTE: The layout is different from models/MiniCPM.py which is [batch_size, num_heads, seq_length, d_qk]
+            # Here it is, [batch_size, seq_length, num_heads, d_qk]
+            # [2, batch_size, seq_length, num_heads, d_qk]
+            key_value_states = torch.cat([key_states.unsqueeze(0), value_states.unsqueeze(0)], dim=0)
+            # [batch_size, seq_length, 2, num_heads, d_qk]
+            key_value_states = key_value_states.permute(1, 2, 0, 3, 4).contiguous()
+            query_states, key_value_states = self.flash_rotary_embedding(query_states, kv=key_value_states)
+            # [batch_size, seq_length, num_heads, d_qk]
+            key_states, value_states = torch.split(key_value_states, 1, dim=2)
+
+            q_sequence_mask = sequence_mask
+            kv_sequence_mask = sequence_mask
+
+            kv_length = key_states.shape[1]
+            # [batch_size, seq_length, num_heads, d_qk]
+            # Shaping for use in `flash-attn` version of flash-attn: `flash_attn_unpadded_func`
+            query_states = query_states.view(
+                batch_size * q_length, self.n_local_q_heads, self.d_qk
+            )  # [batch_size * q_length, self.n_heads, d_qk]
+
+            key_states = key_states.view(
+                batch_size * kv_length, self.n_local_kv_heads, self.d_qk
+            )  # [batch_size * kv_length, self.n_heads, d_qk]
+            value_states = value_states.view(
+                batch_size * kv_length, self.n_local_kv_heads, self.d_v
+            )  # [batch_size * kv_length, self.n_heads, d_v]
+
+            attention_output = self.attention(
+                query_states=query_states,
+                key_states=key_states,
+                value_states=value_states,
+                q_sequence_mask=q_sequence_mask,
+                kv_sequence_mask=kv_sequence_mask,
+            )
+
+        attention_output = (
+            attention_output.contiguous().view(batch_size, q_length, self.n_local_q_heads * self.d_v).transpose(0, 1)
+        )
+        output = self.o_proj(attention_output)
+
+        return {"hidden_states": output, "sequence_mask": sequence_mask}
+
+
+class MiniCPMDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: MiniCPMConfig,
+        parallel_config: Optional[ParallelismArgs],
+        tp_pg: dist.ProcessGroup,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.scale_depth = config.scale_depth
+        self.num_hidden_layers = config.num_hidden_layers
+        self.input_layernorm = TritonRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = CausalSelfAttention(
+            config=config,
+            parallel_config=parallel_config,
+            tp_pg=tp_pg,
+            layer_idx=layer_idx,
+        )
+
+        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,
+        hidden_states: Union[torch.Tensor, TensorPointer],
+        sequence_mask: Union[torch.Tensor, TensorPointer],
+    ) -> Dict[str, Union[torch.Tensor, TensorPointer]]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        output = self.attn(hidden_states=hidden_states, sequence_mask=sequence_mask)
+        hidden_states = output["hidden_states"]
+        hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states=hidden_states)["hidden_states"]
+        hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
+
+        return {
+            "hidden_states": hidden_states,
+            "sequence_mask": output["sequence_mask"],
+        }
+
+
+class Embedding(nn.Module, AttachableStore):
+    def __init__(self, tp_pg: dist.ProcessGroup, config: MiniCPMConfig, parallel_config: Optional[ParallelismArgs]):
+        super().__init__()
+        self.token_embedding = TensorParallelEmbedding(
+            num_embeddings=config.vocab_size,
+            embedding_dim=config.hidden_size,
+            padding_idx=config.pad_token_id,
+            pg=tp_pg,
+            mode=parallel_config.tp_mode if parallel_config is not None else TensorParallelLinearMode.ALL_REDUCE,
+        )
+        self.pg = tp_pg
+
+    def forward(self, input_ids: torch.Tensor, input_mask: torch.Tensor):  # [batch_size, seq_length]
+        store = self.get_local_store()
+        if store is not None:
+            if "past_length" in store:
+                past_length = store["past_length"]
+            else:
+                past_length = torch.zeros(1, dtype=torch.long, device=input_ids.device).expand(input_ids.shape[0])
+
+            cumsum_mask = input_mask.cumsum(-1, dtype=torch.long)
+            # Store new past_length in store
+            store["past_length"] = past_length + cumsum_mask[:, -1]
+
+        # Format input in `[seq_length, batch_size]` to support high TP with low batch_size
+        input_ids = input_ids.transpose(0, 1)
+        input_embeds = self.token_embedding(input_ids)
+        return {"input_embeds": input_embeds}
+
+
+class MiniCPMModel(nn.Module):
+    """Build pipeline graph"""
+
+    def __init__(
+        self,
+        config: MiniCPMConfig,
+        parallel_context: ParallelContext,
+        parallel_config: Optional[ParallelismArgs],
+    ):
+        super().__init__()
+
+        # Declare all the nodes
+        self.p2p = P2P(parallel_context.pp_pg, device=torch.device("cuda"))
+        self.config = config
+        self.parallel_config = parallel_config
+        self.parallel_context = parallel_context
+        self.tp_mode = parallel_config.tp_mode if parallel_config is not None else TensorParallelLinearMode.ALL_REDUCE
+        tp_linear_async_communication = (
+            parallel_config.tp_linear_async_communication if parallel_config is not None else False
+        )
+
+        self.token_position_embeddings = PipelineBlock(
+            p2p=self.p2p,
+            module_builder=Embedding,
+            module_kwargs={
+                "tp_pg": parallel_context.tp_pg,
+                "config": config,
+                "parallel_config": parallel_config,
+            },
+            module_input_keys={"input_ids", "input_mask"},
+            module_output_keys={"input_embeds"},
+        )
+
+        self.decoder = nn.ModuleList(
+            [
+                PipelineBlock(
+                    p2p=self.p2p,
+                    module_builder=MiniCPMDecoderLayer,
+                    module_kwargs={
+                        "config": config,
+                        "parallel_config": parallel_config,
+                        "tp_pg": parallel_context.tp_pg,
+                        "layer_idx": layer_idx,
+                    },
+                    module_input_keys={"hidden_states", "sequence_mask"},
+                    module_output_keys={"hidden_states", "sequence_mask"},
+                )
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.final_layer_norm = PipelineBlock(
+            p2p=self.p2p,
+            module_builder=TritonRMSNorm,
+            module_kwargs={"hidden_size": config.hidden_size, "eps": config.rms_norm_eps},
+            module_input_keys={"input"},
+            module_output_keys={"hidden_states"},
+        )  # TODO
+
+        self.lm_head = PipelineBlock(
+            p2p=self.p2p,
+            # Understand that this means that we return sharded logits that are going to need to be gathered
+            module_builder=TensorParallelColumnLinear,
+            module_kwargs={
+                "in_features": config.hidden_size,
+                "out_features": config.vocab_size,
+                "pg": parallel_context.tp_pg,
+                "bias": False,
+                # TODO @thomasw21: refactor so that we store that default in a single place.
+                "mode": self.tp_mode,
+                "async_communication": tp_linear_async_communication,
+            },
+            module_input_keys={"x"},
+            module_output_keys={"logits"},
+        )
+
+        self.cast_to_fp32 = PipelineBlock(
+            p2p=self.p2p,
+            module_builder=lambda: lambda x: x.float(),
+            module_kwargs={},
+            module_input_keys={"x"},
+            module_output_keys={"output"},
+        )
+
+    def forward(
+        self,
+        input_ids: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
+        input_mask: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
+    ):
+        return self.forward_with_hidden_states(input_ids=input_ids, input_mask=input_mask)[0]
+
+    def forward_with_hidden_states(
+        self,
+        input_ids: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
+        input_mask: Union[torch.Tensor, TensorPointer],  # [batch_size, seq_length]
+    ):
+        # all tensors are optional as most ranks don't need anything from the dataloader.
+
+        output = self.token_position_embeddings(input_ids=input_ids, input_mask=input_mask)
+
+        hidden_encoder_states = {
+            "hidden_states": output["input_embeds"] * self.config.scale_emb,
+            "sequence_mask": input_mask,
+        }
+        for encoder_block in self.decoder:
+            hidden_encoder_states = encoder_block(**hidden_encoder_states)
+
+        hidden_states = self.final_layer_norm(input=hidden_encoder_states["hidden_states"])["hidden_states"]
+
+        sharded_logits = self.lm_head(x=hidden_states / (self.config.hidden_size / self.config.dim_model_base))["logits"]
+
+        fp32_sharded_logits = self.cast_to_fp32(x=sharded_logits)["output"]
+
+        return fp32_sharded_logits, hidden_states
+
+    def get_block_compute_costs(self):
+        """Computes the compute cost of each block in the model so that we can do a better job of load balancing."""
+        model_config = self.config
+        d_ff = model_config.intermediate_size
+        d_qkv = model_config.hidden_size // model_config.num_attention_heads
+        block_compute_costs = {
+            # CausalSelfAttention (qkv proj + attn out) + MLP
+            MiniCPMDecoderLayer: 4 * model_config.num_attention_heads * d_qkv * model_config.hidden_size
+            + 3 * d_ff * model_config.hidden_size,
+            # This is the last lm_head
+            TensorParallelColumnLinear: model_config.vocab_size * model_config.hidden_size,
+        }
+        return block_compute_costs
+
+    def get_flops_per_sec(self, iteration_time_in_sec, sequence_length, global_batch_size):
+        """Get flops per second for a given model"""
+        world_size = self.parallel_context.world_pg.size()
+        try:
+            num_key_values_heads = self.config.num_key_value_heads
+        except AttributeError:
+            num_key_values_heads = self.config.num_attention_heads
+
+        model_flops, hardware_flops = get_flops(
+            num_layers=self.config.num_hidden_layers,
+            hidden_size=self.config.hidden_size,
+            num_heads=self.config.num_attention_heads,
+            num_key_value_heads=num_key_values_heads,
+            vocab_size=self.config.vocab_size,
+            ffn_hidden_size=self.config.intermediate_size,
+            seq_len=sequence_length,
+            batch_size=global_batch_size,
+            recompute_granularity=self.parallel_config.recompute_granularity,
+        )
+
+        model_flops_per_s = model_flops / (iteration_time_in_sec * world_size * 1e12)
+        hardware_flops_per_s = hardware_flops / (iteration_time_in_sec * world_size * 1e12)
+        return model_flops_per_s, hardware_flops_per_s
+
+
+@torch.jit.script
+def masked_mean(loss, label_mask, dtype):
+    # type: (Tensor, Tensor, torch.dtype) -> Tensor
+    return (loss * label_mask).sum(dtype=dtype) / label_mask.sum()
+
+
+class Loss(nn.Module):
+    def __init__(self, tp_pg: dist.ProcessGroup):
+        super().__init__()
+        self.tp_pg = tp_pg
+
+    def forward(
+        self,
+        sharded_logits: torch.Tensor,  # [seq_length, batch_size, logits]
+        label_ids: torch.Tensor,  # [batch_size, seq_length]
+        label_mask: torch.Tensor,  # [batch_size, seq_length]
+    ) -> Dict[str, torch.Tensor]:
+        # Megatron by defaults cast everything in fp32. `--f16-lm-cross-entropy` is an option you can use to keep current precision.
+        # https://github.com/NVIDIA/Megatron-LM/blob/f267e6186eae1d6e2055b412b00e2e545a8e896a/megatron/model/gpt_model.py#L38
+        loss = sharded_cross_entropy(
+            sharded_logits, label_ids.transpose(0, 1).contiguous(), group=self.tp_pg, dtype=torch.float
+        ).transpose(0, 1)
+        # TODO @thomasw21: It's unclear what kind of normalization we want to do.
+        loss = masked_mean(loss, label_mask, dtype=torch.float)
+        # I think indexing causes a sync we don't actually want
+        # loss = loss[label_mask].sum()
+        return {"loss": loss}
+
+
+class MiniCPMForTraining(NanotronModel):
+    def __init__(
+        self,
+        config: MiniCPMConfig,
+        parallel_context: ParallelContext,
+        parallel_config: Optional[ParallelismArgs],
+        random_states: Optional[RandomStates] = None,
+    ):
+        super().__init__()
+        import warnings
+
+        self.model = MiniCPMModel(config=config, parallel_context=parallel_context, parallel_config=parallel_config)
+        self.loss = PipelineBlock(
+            p2p=self.model.p2p,
+            module_builder=Loss,
+            module_kwargs={"tp_pg": parallel_context.tp_pg},
+            module_input_keys={
+                "sharded_logits",
+                "label_ids",
+                "label_mask",
+            },
+            module_output_keys={"loss"},
+        )
+        self.parallel_context = parallel_context
+        self.config = config
+        self.parallel_config = parallel_config
+
+    def forward(
+        self,
+        input_ids: Union[torch.Tensor, TensorPointer],
+        input_mask: Union[torch.Tensor, TensorPointer],
+        label_ids: Union[torch.Tensor, TensorPointer],
+        label_mask: Union[torch.Tensor, TensorPointer],
+    ) -> Dict[str, Union[torch.Tensor, TensorPointer]]:
+        sharded_logits = self.model(
+            input_ids=input_ids,
+            input_mask=input_mask,
+        )
+        loss = self.loss(
+            sharded_logits=sharded_logits,
+            label_ids=label_ids,
+            label_mask=label_mask,
+        )["loss"]
+        return {"loss": loss}
+
+    @torch.no_grad()
+    def init_model_randomly(self, init_method, scaled_init_method):
+        """Initialize model parameters randomly.
+        Args:
+            init_method (callable): Used for embedding/position/qkv weight in attention/first layer weight of mlp/ /lm_head/
+            scaled_init_method (callable): Used for o weight in attention/second layer weight of mlp/
+
+        Note:
+            Layernorm weight all 0 or 1 depending on `apply_layernorm_1p`
+        """
+        model = self
+        initialized_parameters = set()
+        # Handle tensor parallelism
+        module_id_to_prefix = {id(module): f"{module_name}." for module_name, module in model.named_modules()}
+        # Fix the root_model
+        module_id_to_prefix[id(model)] = ""
+
+        for module_name, module in model.named_modules():
+            if isinstance(module, TensorParallelColumnLinear):
+                # Somehow Megatron-LM does something super complicated, https://github.com/NVIDIA/Megatron-LM/blob/2360d732a399dd818d40cbe32828f65b260dee11/megatron/core/tensor_parallel/layers.py#L96
+                # What it does:
+                #  - instantiate a buffer of the `full size` in fp32
+                #  - run init method on it
+                #  - shard result to get only a specific shard
+                # Instead I'm lazy and just going to run init_method, since they are scalar independent
+                assert {"weight"} == {name for name, _ in module.named_parameters()} or {"weight"} == {
+                    name for name, _ in module.named_parameters()
+                }
+                for param_name, param in module.named_parameters():
+                    assert isinstance(param, NanotronParameter)
+                    if param.is_tied:
+                        tied_info = param.get_tied_info()
+                        full_param_name = tied_info.get_full_name_from_module_id_to_prefix(
+                            module_id_to_prefix=module_id_to_prefix
+                        )
+                    else:
+                        full_param_name = f"{module_name}.{param_name}"
+
+                    if full_param_name in initialized_parameters:
+                        # Already initialized
+                        continue
+
+                    if "weight" == param_name:
+                        init_method(param)
+                    elif "bias" == param_name:
+                        param.zero_()
+                    else:
+                        raise ValueError(f"Who the fuck is {param_name}?")
+
+                    assert full_param_name not in initialized_parameters
+                    initialized_parameters.add(full_param_name)
+            elif isinstance(module, TensorParallelRowLinear):
+                # Somehow Megatron-LM does something super complicated, https://github.com/NVIDIA/Megatron-LM/blob/2360d732a399dd818d40cbe32828f65b260dee11/megatron/core/tensor_parallel/layers.py#L96
+                # What it does:
+                #  - instantiate a buffer of the `full size` in fp32
+                #  - run init method on it
+                #  - shard result to get only a specific shard
+                # Instead I'm lazy and just going to run init_method, since they are scalar independent
+                assert {"weight"} == {name for name, _ in module.named_parameters()} or {"weight"} == {
+                    name for name, _ in module.named_parameters()
+                }
+                for param_name, param in module.named_parameters():
+                    assert isinstance(param, NanotronParameter)
+                    if param.is_tied:
+                        tied_info = param.get_tied_info()
+                        full_param_name = tied_info.get_full_name_from_module_id_to_prefix(
+                            module_id_to_prefix=module_id_to_prefix
+                        )
+                    else:
+                        full_param_name = f"{module_name}.{param_name}"
+
+                    if full_param_name in initialized_parameters:
+                        # Already initialized
+                        continue
+
+                    if "weight" == param_name:
+                        scaled_init_method(param)
+                    elif "bias" == param_name:
+                        param.zero_()
+                    else:
+                        raise ValueError(f"Who the fuck is {param_name}?")
+
+                    assert full_param_name not in initialized_parameters
+                    initialized_parameters.add(full_param_name)
+            elif isinstance(module, TritonRMSNorm):
+                assert {"weight"} == {name for name, _ in module.named_parameters()}
+                for param_name, param in module.named_parameters():
+                    assert isinstance(param, NanotronParameter)
+                    if param.is_tied:
+                        tied_info = param.get_tied_info()
+                        full_param_name = tied_info.get_full_name_from_module_id_to_prefix(
+                            module_id_to_prefix=module_id_to_prefix
+                        )
+                    else:
+                        full_param_name = f"{module_name}.{param_name}"
+
+                    if full_param_name in initialized_parameters:
+                        # Already initialized
+                        continue
+
+                    if "weight" == param_name:
+                        # TODO @thomasw21: Sometimes we actually want 0
+                        param.fill_(1)
+                    elif "bias" == param_name:
+                        param.zero_()
+                    else:
+                        raise ValueError(f"Who the fuck is {param_name}?")
+
+                    assert full_param_name not in initialized_parameters
+                    initialized_parameters.add(full_param_name)
+            elif isinstance(module, TensorParallelEmbedding):
+                # TODO @thomasw21: Handle tied embeddings
+                # Somehow Megatron-LM does something super complicated, https://github.com/NVIDIA/Megatron-LM/blob/2360d732a399dd818d40cbe32828f65b260dee11/megatron/core/tensor_parallel/layers.py#L96
+                # What it does:
+                #  - instantiate a buffer of the `full size` in fp32
+                #  - run init method on it
+                #  - shard result to get only a specific shard
+                # Instead I'm lazy and just going to run init_method, since they are scalar independent
+                assert {"weight"} == {name for name, _ in module.named_parameters()}
+
+                assert isinstance(module.weight, NanotronParameter)
+                if module.weight.is_tied:
+                    tied_info = module.weight.get_tied_info()
+                    full_param_name = tied_info.get_full_name_from_module_id_to_prefix(
+                        module_id_to_prefix=module_id_to_prefix
+                    )
+                else:
+                    full_param_name = f"{module_name}.weight"
+
+                if full_param_name in initialized_parameters:
+                    # Already initialized
+                    continue
+
+                init_method(module.weight)
+                assert full_param_name not in initialized_parameters
+                initialized_parameters.add(full_param_name)
+
+        assert initialized_parameters == {
+            param.get_tied_info().get_full_name_from_module_id_to_prefix(module_id_to_prefix=module_id_to_prefix)
+            if param.is_tied
+            else name
+            for name, param in model.named_parameters()
+        }, f"Somehow the initialized set of parameters don't match:\n - Expected: { {name for name, _ in model.named_parameters()} }\n - Got: {initialized_parameters}"
+
+    def get_block_compute_costs(self):
+        """Computes the compute cost of each block in the model so that we can do a better job of load balancing."""
+        return self.model.get_block_compute_costs()
+
+    def get_flops_per_sec(self, iteration_time_in_sec, sequence_length, global_batch_size):
+        """Get flops per second for a given model"""
+        return self.model.get_flops_per_sec(iteration_time_in_sec, sequence_length, global_batch_size)
+
+
+def get_flops(
+    num_layers,
+    hidden_size,
+    num_heads,
+    vocab_size,
+    seq_len,
+    kv_channels=None,
+    ffn_hidden_size=None,
+    batch_size=1,
+    recompute_granularity=None,
+    glu_activation=False,
+):
+    """Counts flops in an decoder-only model
+    Args:
+        num_layers: number of decoder layers
+        hidden_size: hidden size of the model
+        num_heads: number of heads in the model
+        num_key_value_heads: number of key/value heads in the model
+        ffn_hidden_size: hidden size of the FFN
+        vocab_size: size of the vocabulary
+        seq_len: sequence length of the decoder
+        batch_size: batch size
+        recompute_granularity: Activation recomputation method. Either None, FULL or SELECTIVE. Check Megatron-LM docs for more info.
+    Returns:
+        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 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 (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 * (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) * 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 * (kv_channels) * hidden_size
+    # FF
+    ## 1st layer
+    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_q_proj_flops_fwd
+        + decoder_kv_proj_flops_fwd
+        + decoder_qk_logits_flops_fwd
+        + decoder_v_logits_flops_fwd
+        + decoder_attn_out_flops_fwd
+        + decoder_ffn_1_flops_fwd
+        + decoder_ffn_2_flops_fwd
+    )
+
+    # lm head
+    lm_head_flops_fwd = 2 * batch_size * seq_len * (hidden_size) * vocab_size
+
+    # the bwd pass requires double the flops in case of matmuls to calculate the gradients with respect to
+    # both input and weight tensors
+    model_flops = 3 * (decoder_flops_fwd + lm_head_flops_fwd)  # 1 for fwd + 2 for bwd
+
+    if recompute_granularity is None:
+        hardware_flops = model_flops
+    elif recompute_granularity is RecomputeGranularity.FULL:
+        # Note: we don't recompute lm head activs
+        hardware_flops = model_flops + decoder_flops_fwd  # + activ recomputation
+    elif recompute_granularity is RecomputeGranularity.SELECTIVE:
+        # all terms with s^2 are flops that are recomputed
+        # ref. appendix A: https://arxiv.org/pdf/2205.05198.pdf
+        recomputed_decoder_flops = decoder_qk_logits_flops_fwd + decoder_v_logits_flops_fwd
+        hardware_flops = model_flops + recomputed_decoder_flops
+    else:
+        raise ValueError("recompute_granularity must be one of 'full' or 'selective'")
+
+    return model_flops, hardware_flops

pretrained/MiniCPM-2B-dpo-bf16/checkpoint_metadata.json

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

pretrained/MiniCPM-2B-dpo-bf16/config.yaml

@@ -0,0 +1,55 @@
+checkpoints: null
+data: null
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: false
+  project: openbmb
+  run: MiniCPM-2B-dpo-bf16
+  seed: 42
+  step: 0
+lighteval: 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
+    dim_model_base: 256
+    eos_token_id: 2
+    hidden_act: silu
+    hidden_size: 2304
+    initializer_range: 0.1
+    intermediate_size: 5760
+    max_position_embeddings: 2048
+    num_attention_heads: 36
+    num_hidden_layers: 40
+    num_key_value_heads: 36
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_theta: 10000.0
+    scale_depth: 1.4
+    scale_emb: 12
+    tie_word_embeddings: true
+    use_cache: true
+    vocab_size: 122753
+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: openbmb/MiniCPM-2B-dpo-bf16
+  tokenizer_revision: null
+tokens: null

pretrained/MiniCPM-2B-dpo-bf16/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:36d4080827449bba1682daf6cab4b546e2fc02bb6a5c62efff2470a7e83202f5
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:e8843d169ca3025d42ddfdb590fc51d01122c1f9212f9581d4dd319430f0e45b
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:beb60636c7e802f851754225517f7f375a0faff02721fd37a2d460a2ca61458d
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:b5fec63ea6756d045071b457036fe09b783b34570efa1705d1d5120611db651f
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:b895ee74603d6a88a024041a16266e351714207a61f966ceeb5e638b4ba1a4ed
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:a319fc293f19f00db8aa2811b1a730a38a5966459493151020240c7adcb5dfd6
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:495d95a04eea9c36ec15b8c1add6e51a52052ee0fd90b446683d7a412d01b526
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:908e028114e8dc887ca2e80245678ea7ab08662b8462bdf9111d5e342f9de729
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:a6a6bc147e3a4d94a72c4a48c86b7b4517662141923380b80e216897e0c647bf
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:88805e4d760f1e78e2d42e7ff18c1ffa61051571bf8b9eefc6fdd9f3843770e4
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:16ff202e99dfe71bb7619cc3580a9f59c788868f8f6fc58c6c694f5eae562c89
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:f240f0046868ee13ee354feb943d5c382358774d8e1533dd53e16040d9618124
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:56ab2210f9e6e6018e84b6e96ad6bdc525e46396049ff66caa0285fb7a170120
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:8736369fe6db6ca62bdd30014a6c920c6d6bb570c85b79d5ce43bff75366a499
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:e37863cb49b75e41eff480d6fb94a3240cbc884dc2e76215d36f8d6207fba5ee
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:6b7d3522e6a49963803f5e4248ff76e2a9d51a044c76cfac156bc581b794f492
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:d1fdd9830af039908d5a402c6edeccb2527d5cdddfbe020dd3e5adc4d33e0371
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:639ee973e90fb837781436fbf443facd7ed1f73f070041ffd71c14196d06e2ce
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:70f9c571fb7c57aa4a57d51fe4c5d05fe21a34e32a2f4896049c1d9bed6ea32b
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:3f4a20b46bd3098c0f1efdd665bba8f99adc2363257db31a8b7a0e4e73a540bd
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:c572d06613b5e508d138b5af0180b5b3a09c78abc7b06a07039fec736487d1e9
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:7d021e4f9a713730b67f2e8df642639628ebd5674eb960c39575c62e3d92212f
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:a290b7bf65372eb56ae8036fbe048ee0d94a39e01debf2c07600d2228bd6cb0c
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:03a6490c81420efe127f1f8dc68c7f8734a9b936ca901e2f5e2485d90d178469
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:f1d3c77097833aa4ea96466ae247220e1f2a4ef6957c65f3257ad8b1da84ef70
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:0ca582e7a768fa2766bf655a6acb54ee6d8c3ec1fe63a7b8bc5ebb53846828f2
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:0b66f9fc608826efaa57dff502e1f0ba1bea00996208f1e8e875f58459d82c65
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:36035e96236f5c5f39b19d5f53af38da02d05317849852bc5f2dd0eca315148a
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:d1c6047d8d338b808b94f73626c1a0ef9ac33a4c8cf964041ba8c9db3d8f9629
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:f3c926d53d3eb31a3c6aea0390794e16fd9ac99d4a99038a48de9ce750895ed4
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:5f1b53cb9c0ea5fe7080e68a0ad6e988aaf958315d3b58276d9ec7e2bd239e67
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:08eca1997fce9a5815f81299a83848066b309a2d031f88e5012878236ae4650c
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:ab3f62ae9b040d8aa3e1b9d3a8b81d89ad731c5d9bb32757a43b116fcfd79ca0
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:eb415ebe5cbfc55a84d7657c74cf1634a9f534a5980242e59484b9c2a041d896
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/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:4baf79283bdd79707aabf64311ee95bb4149435bb4469aa0e43aed65caa8c667
+size 53084456

pretrained/MiniCPM-2B-dpo-bf16/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:5543638dcaccb01fa706781c9dc67b820df41a7209c25494dc82b6db8e63b1f3
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:0883579e29c4b28ee7f4af72be5d77c2f045fca9757ca06f6d0a8c75de6abd0f
+size 10617072

pretrained/MiniCPM-2B-dpo-bf16/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:c436a6adf43608478642ab6efbf750101b4b3fc4d19425b649443309bde8417a
+size 31850848

pretrained/MiniCPM-2B-dpo-bf16/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:43b8fe2b4b03bbe4ddbd0dc572b4cfac0286624eb2865d96f9afdb75c727643f
+size 4704

pretrained/MiniCPM-2B-dpo-bf16/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:e6194cb03e3b7c3260e8a34e632f7736177a5ee15a46d184d90e6e94d5a6802e
+size 26542320

pretrained/MiniCPM-2B-dpo-bf16/model/model/decoder/14/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:ec68dd15fe1f1643fec13bb3cb19d8139beed19fdf48075c097baf16c5f23ba3
+size 53084456