Model save

README.md

@@ -0,0 +1,184 @@
+---
+library_name: transformers
+tags:
+- axolotl
+- generated_from_trainer
+datasets:
+- llm-jp/magpie-sft-v1.0
+model-index:
+- name: plamo-2-1b-gorilla-chat2
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+[<img src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/axolotl-ai-cloud/axolotl)
+<details><summary>See axolotl config</summary>
+
+axolotl version: `0.7.0`
+```yaml
+
+
+# モデルの設定
+base_model: /notebooks/plamo-2-1b-gorilla-chat2              # HuggingFace上のモデル名
+model_type: AutoModelForCausalLM         # モデルのロードに使用するクラス
+tokenizer_type: AutoTokenizer           # トークナイザのロードに使用するクラス
+trust_remote_code: true                 # リモートのカスタムコードを信頼してモデルをロード
+
+hub_model_id: zamagi/fft-1
+hub_strategy: "end"
+push_dataset_to_hub:
+hf_use_auth_token: true
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_cross_entropy: false
+liger_rope: true
+liger_rms_norm: true
+liger_swiglu: true
+liger_fused_linear_cross_entropy: true
+
+# 8bit/4bit設定（8bitモードでメモリ削減）
+load_in_8bit: false   #f                      # 8bit量子化されたモデルをロード
+load_in_4bit: false                     # 4bit量子化は使用しない
+strict: false                           # 重みの厳密な一致を要求しない（追加トークン等がある場合に許容）
+
+chat_template: tokenizer_default
+
+# データセットの設定
+datasets:
+  - path: llm-jp/magpie-sft-v1.0         # 使用するデータセット（Hugging Face上のデータセット名）
+    type: chat_template                 # 会話形式のデータセットを使用
+    field_messages: conversations       # 会話データが格納されたフィールド名
+    message_property_mappings:         # メッセージ内のプロパティ名のマッピング
+      role: role                       # 役割（ユーザー/システム/アシスタント）を示すフィールド
+      content: content                 # メッセージ内容を示すフィールド
+    roles_to_train: ["assistant"]       # 学習対象とする役割（アシスタントの発話のみ学習）
+    train_on_eos: last
+dataset_prepared_path: notebooks/last_run_prepared  # 前処理済みデータの保存先（キャッシュ用）
+
+
+shuffle_merged_datasets: true
+dataset_prepared_path: /notebooks/data/fft-data
+#val_set_size: 0.003
+output_dir: /notebooks/data/27b-fft-out-1
+dataset_keep_in_memory: false
+
+gpu_memory_limit: 48GiB
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+pad_to_sequence_len: true
+
+adapter:
+lora_model_dir:
+lora_r:
+lora_alpha:
+lora_dropout:
+lora_target_linear:
+lora_fan_in_fan_out:
+
+
+# トレーニングの設定
+gradient_accumulation_steps: 4
+micro_batch_size: 8
+num_epochs: 2
+optimizer: paged_adamw_8bit
+lr_scheduler:
+cosine_min_lr_ratio: 0.1
+learning_rate: 0.00001
+max_steps: 5000
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: false
+
+#wandb: false
+#wandb_project: 27b-fft
+#wandb_entity: aratako-lm
+#wandb_watch:
+#wandb_name: attempt-01
+#wandb_log_model:
+
+gradient_checkpointing: true
+early_stopping_patience:
+auto_resume_from_checkpoints: true
+local_rank:
+logging_steps: 1
+xformers_attention: 
+flash_attention: 
+
+save_strategy: steps
+save_steps: 100
+save_total_limit: 2
+
+warmup_steps: 50
+eval_steps: 100
+eval_batch_size: 1
+eval_table_size:
+eval_max_new_tokens:
+
+debug:
+deepspeed: /notebooks/axolotl/deepspeed_configs/zero3_bf16.json
+weight_decay: 0.01
+fsdp:
+fsdp_config:
+
+
+
+# 出力の保存設定
+output_dir: /notebooks/output/plamo-2-1b-gorilla-chat2   # チェックポイントや最終モデルの出力先ディレクトリ
+hub_model_id: zamagi/plamo-2-1b-gorilla-chat2   # (オプション) Hugging Face Hubにアップロードする場合のリポジトリ名
+
+
+```
+
+</details><br>
+
+# plamo-2-1b-gorilla-chat2
+
+This model was trained from scratch on the llm-jp/magpie-sft-v1.0 dataset.
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 1e-05
+- train_batch_size: 8
+- eval_batch_size: 1
+- seed: 42
+- distributed_type: multi-GPU
+- gradient_accumulation_steps: 4
+- total_train_batch_size: 32
+- optimizer: Use OptimizerNames.PAGED_ADAMW_8BIT with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_steps: 50
+- training_steps: 1112
+
+### Training results
+
+
+
+### Framework versions
+
+- Transformers 4.49.0
+- Pytorch 2.5.1+cu124
+- Datasets 3.2.0
+- Tokenizers 0.21.1

config.json

@@ -0,0 +1,51 @@
+{
+  "_name_or_path": "/notebooks/plamo-2-1b-gorilla-chat2",
+  "architectures": [
+    "PlamoForCausalLM"
+  ],
+  "attention_window_size": 2048,
+  "auto_map": {
+    "AutoConfig": "modeling_plamo.PlamoConfig",
+    "AutoModelForCausalLM": "pfnet/plamo-2-1b--modeling_plamo.PlamoForCausalLM"
+  },
+  "bos_token_id": 1,
+  "capacity_factor": 1.0,
+  "eos_token_id": 2,
+  "eval_attention_n_bit": null,
+  "eval_mlp_n_bit": null,
+  "expert_dropout": 0.0,
+  "fp8_accum_dtype": "bfloat16",
+  "full_attention_idx": [],
+  "group_size": 1024,
+  "hidden_size": 2048,
+  "hidden_size_per_head": 128,
+  "image_feature_size": null,
+  "image_proj_type": "linear",
+  "image_token_id": null,
+  "intermediate_size": 8192,
+  "k_expert": null,
+  "linear_type": "fp8",
+  "mamba_chunk_size": 256,
+  "mamba_d_conv": 4,
+  "mamba_d_state": 64,
+  "mamba_enabled": true,
+  "mamba_num_heads": 32,
+  "mamba_step": 2,
+  "max_position_embeddings": 10485760,
+  "model_type": "plamo",
+  "n_expert": null,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 1,
+  "rms_norm_eps": 1e-06,
+  "shared_intermediate_size": null,
+  "sliding_window": 2048,
+  "sparse_intermediate_size": null,
+  "sparse_step": null,
+  "tokenizer_class": "PlamoTokenizer",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.49.0",
+  "use_cache": false,
+  "use_predefined_initial_state": false,
+  "vocab_size": 100000
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 2,
+  "transformers_version": "4.49.0"
+}

model.safetensors

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

modeling_plamo.py

@@ -0,0 +1,1699 @@
+import enum
+import math
+import warnings
+from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Union
+
+try:
+    # It is difficult to install mamba_ssm in login node because
+    # it requires GPU for installation
+    import mamba_ssm
+except ModuleNotFoundError:
+    warnings.warn("mamba_ssm could not be imported", stacklevel=2)
+try:
+    # It is difficult to install causal_conv1d in login node because
+    # it requires GPU for installation
+    import causal_conv1d.causal_conv1d_interface as causal_conv1d
+except ModuleNotFoundError:
+    warnings.warn("causal_conv1d could not be imported", stacklevel=2)
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+
+def _is_first_token(mask: torch.Tensor) -> torch.Tensor:
+    assert mask.dtype == torch.bool
+    B, Nh, q_len, kv_len = mask.shape
+    mask = mask[:, :, :, -q_len:]
+    cont = q_len != kv_len
+    v = False if cont else True
+    out = torch.logical_not(torch.diagonal(mask, offset=-1, dim1=-2, dim2=-1).bool())
+    out = torch.cat(
+        [
+            torch.full(size=(B, Nh, 1), dtype=torch.bool, device=out.device, fill_value=v),
+            out,
+        ],
+        dim=-1,
+    )
+    return out
+
+
+def _swiglu(h: torch.Tensor) -> torch.Tensor:
+    h0, h1 = h.chunk(2, dim=-1)
+    return torch.nn.functional.silu(h0) * h1
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(
+        self, dim: int, max_position_embeddings: int = 2048, base: int = 10000, device: Optional[torch.device] = None
+    ) -> None:
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len: int, device: Any, dtype: Any) -> None:
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)  # type: ignore
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
+
+    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),  # type: ignore
+            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),  # type: ignore
+        )
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
+    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    x_embed = (x * cos) + (_rotate_half(x) * sin)
+    return x_embed
+
+
+class LinearType(str, enum.Enum):
+    Normal = "normal"
+    Fp8 = "fp8"
+    Fp8Retain = "fp8-retain"
+
+
+class PlamoConfig(PretrainedConfig):  # type: ignore
+    model_type: str = "plamo"
+
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        num_hidden_layers: int = 32,
+        rms_norm_eps: float = 1e-6,
+        tie_word_embeddings: bool = True,
+        # Attention
+        num_attention_heads: int = 32,
+        num_key_value_heads: int = 4,
+        hidden_size_per_head: int = 128,
+        max_position_embeddings: int = 2048,
+        attention_window_size: int = 2048,
+        full_attention_idx: list[int] | None = None,
+        # Mamba
+        mamba_d_state: int = 64,
+        mamba_d_conv: int = 4,
+        mamba_num_heads: int = 64,
+        mamba_step: int = 2,
+        mamba_chunk_size: int = 256,
+        mamba_enabled: bool = True,
+        # MLP
+        intermediate_size: int = 13312,
+        # Tokenizer
+        vocab_size: int = 32000,
+        tokenizer_class: str = "PlamoTokenizer",
+        pad_token_id: Optional[int] = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        # Multimodal
+        image_token_id: Optional[int] = None,
+        image_feature_size: Optional[int] = None,
+        image_proj_type: Literal["linear", "mlp"] = "linear",
+        # FP8
+        linear_type: LinearType = LinearType.Normal,
+        fp8_accum_dtype: Optional[str] = None,
+        # Evaluation
+        eval_attention_n_bit: Optional[int] = None,
+        eval_mlp_n_bit: Optional[int] = None,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        # max_position_embeddings is often used to determine the max length during inference,
+        # but samba should have extrapolation abilities
+        self.max_position_embeddings = max(10 * 1024 * 1024, max_position_embeddings)
+        self.hidden_size = hidden_size
+        self.rms_norm_eps = rms_norm_eps
+
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size_per_head = hidden_size_per_head
+        self.num_key_value_heads = num_key_value_heads
+        self.attention_window_size = attention_window_size
+        self.full_attention_idx = full_attention_idx if full_attention_idx is not None else []
+
+        self.mamba_d_state = mamba_d_state
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_num_heads = mamba_num_heads
+        self.mamba_step = mamba_step
+        self.mamba_chunk_size = mamba_chunk_size
+        self.mamba_enabled = mamba_enabled
+
+        self.intermediate_size = intermediate_size
+
+        self.vocab_size = vocab_size
+
+        self.image_token_id = image_token_id
+        self.image_feature_size = image_feature_size
+        self.image_proj_type = image_proj_type
+
+        self.linear_type = linear_type
+        self.fp8_accum_dtype = fp8_accum_dtype
+
+        self.eval_attention_n_bit = eval_attention_n_bit
+        self.eval_mlp_n_bit = eval_mlp_n_bit
+        self.use_cache = use_cache
+
+        # fields for vLLM
+        self.sliding_window = attention_window_size
+
+        super().__init__(
+            tokenizer_class=tokenizer_class,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+class PlamoAttentionCache(torch.nn.Module):
+    def __init__(self, key: torch.Tensor, value: torch.Tensor) -> None:
+        super().__init__()
+        B, nh, L, c = key.shape
+        assert len(value.shape) == 4
+        assert value.shape[0] == B
+        assert value.shape[2] == L
+        self.register_parameter("key", torch.nn.Parameter(key, requires_grad=False))
+        self.register_parameter("value", torch.nn.Parameter(value, requires_grad=False))
+
+
+class PlamoMambaCache(torch.nn.Module):
+    def __init__(self, conv_state: torch.Tensor, ssm_state: torch.Tensor) -> None:
+        super().__init__()
+        # conv_state: [B, C, d_conv]
+        # ssm_state: [B, nhead, nchanel_per_head, d_state]
+        assert len(conv_state.shape) == 3
+        assert len(ssm_state.shape) == 4
+        assert conv_state.shape[0] == ssm_state.shape[0]
+        self.register_parameter("conv_state", torch.nn.Parameter(conv_state, requires_grad=False))
+        self.register_parameter("ssm_state", torch.nn.Parameter(ssm_state, requires_grad=False))
+
+
+PlamoLayerCache = PlamoAttentionCache | PlamoMambaCache
+
+
+class PlamoCache(torch.nn.Module):
+    """
+    stores states of the model for fast decoding.
+    `transformers` uses `transformers.Cache` for this purpose, but the interface and variable names are
+    deeply dependent on Transformers architecture (e.g., `key_states`) and it is difficult to use
+    other architectures (e.g., Mamba).
+    This class provides a similar interface to `transformers.Cache`, but is designed to also handle
+    the state of Mamba properly.
+    """
+
+    def __init__(self, config: PlamoConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.cache = torch.nn.ModuleList([None for _ in range(config.num_hidden_layers)])  # type: ignore
+
+    def append_kv(self, key: torch.Tensor, value: torch.Tensor, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        c = self.cache[layer_idx]
+        if c is None:
+            return key, value
+        assert isinstance(c, PlamoAttentionCache)
+
+        def _validate(cache: torch.Tensor, new_tensor: torch.Tensor) -> None:
+            assert len(cache.shape) == 4
+            assert len(new_tensor.shape) == 4
+            assert cache.shape[0] == new_tensor.shape[0]
+            assert cache.shape[1] == new_tensor.shape[1]
+            assert cache.shape[3] == new_tensor.shape[3]
+
+        _validate(c.key, key)
+        _validate(c.value, value)
+        assert key.shape[2] == value.shape[2]
+        return torch.cat([c.key, key], dim=2), torch.cat([c.value, value], dim=2)
+
+    def update_attention(
+        self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int
+    ) -> PlamoAttentionCache:
+        full_attn = layer_idx in self.config.full_attention_idx
+        window_size = self.config.attention_window_size
+
+        if self.cache[layer_idx] is None:
+            if full_attn:
+                self.cache[layer_idx] = PlamoAttentionCache(key_states, value_states)
+            else:
+                self.cache[layer_idx] = PlamoAttentionCache(
+                    key_states[:, :, -window_size:, :], value_states[:, :, -window_size:, :]
+                )
+        else:
+            c = self.cache[layer_idx]
+            assert isinstance(c, PlamoAttentionCache)
+            k, v = self.append_kv(key_states, value_states, layer_idx)
+            if full_attn:
+                c.key.data = k
+                c.value.data = v
+            else:
+                c.key.data = k[:, :, -window_size:, :]
+                c.value.data = v[:, :, -window_size:, :]
+        return self.cache[layer_idx]  # type: ignore
+
+    def update_mamba(self, conv_state: torch.Tensor, ssm_state: torch.Tensor, layer_idx: int) -> PlamoMambaCache:
+        if self.cache[layer_idx] is None:
+            self.cache[layer_idx] = PlamoMambaCache(conv_state, ssm_state)
+        else:
+            c = self.cache[layer_idx]
+            assert isinstance(c, PlamoMambaCache)
+            assert c.conv_state.shape == conv_state.shape
+            assert c.ssm_state.shape == ssm_state.shape
+            c.conv_state.data = conv_state
+            c.ssm_state.data = ssm_state
+        return self.cache[layer_idx]  # type: ignore
+
+    def __getitem__(self, layer_idx: int) -> PlamoLayerCache | None:
+        assert layer_idx < len(self.cache)
+        layer_cache = self.cache[layer_idx]
+        return layer_cache  # type: ignore
+
+    def __len__(self) -> int:
+        return len(self.cache)
+
+    def get_seq_length(self, layer_idx: Optional[int] = None) -> int:
+        if layer_idx is not None:
+            c = self.cache[layer_idx]
+            assert isinstance(c, PlamoAttentionCache)
+            return c.key.shape[2]  # type: ignore
+
+        sequence_length: int | None = None
+        for layer_cache in self.cache:
+            if isinstance(layer_cache, PlamoAttentionCache):
+                sequence_length = (
+                    max(layer_cache.key.shape[2], sequence_length)
+                    if sequence_length is not None
+                    else layer_cache.key.shape[2]
+                )
+        assert sequence_length is not None
+        return sequence_length
+
+    def get_max_length(self) -> int | None:
+        return None
+
+    def get_usable_length(self, new_seq_length: int, layer_idx: Optional[int] = 0) -> int:
+        """Given the sequence length of the new inputs, returns the usable length of the cache."""
+        # Cache without size limit -> all cache is usable
+        # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache
+        #   length, we will need to evict part of the cache (and thus not all cache is usable)
+        max_length = self.get_max_length()
+        previous_seq_length = self.get_seq_length(layer_idx)
+        if max_length is not None and previous_seq_length + new_seq_length > max_length:
+            return max_length - new_seq_length
+        return previous_seq_length
+
+    def reorder_cache(self, beam_idx: torch.Tensor) -> None:
+        def _mamba(cache: PlamoMambaCache) -> PlamoMambaCache:
+            return PlamoMambaCache(
+                conv_state=cache.conv_state.index_select(0, beam_idx),
+                ssm_state=cache.ssm_state.index_select(0, beam_idx),
+            )
+
+        def _attention(cache: PlamoAttentionCache) -> PlamoAttentionCache:
+            return PlamoAttentionCache(
+                key=cache.key.index_select(0, beam_idx),
+                value=cache.value.index_select(0, beam_idx),
+            )
+
+        for i in range(len(self.cache)):
+            if self.cache[i] is None:
+                continue
+            layer_cache = self.cache[i]
+            if isinstance(layer_cache, PlamoMambaCache):
+                self.cache[i] = _mamba(layer_cache)
+            else:
+                assert isinstance(layer_cache, PlamoAttentionCache)
+                self.cache[i] = _attention(layer_cache)
+
+    @property
+    def seen_tokens(self) -> int | None:
+        return None
+
+
+class DecoderInput(NamedTuple):
+    hidden_states: torch.Tensor
+    attention_mask: Optional[torch.Tensor] = None
+    past_states: Optional[PlamoCache] = None
+    output_hidden_states: Optional[bool] = False
+    output_attentions: Optional[bool] = False
+    gradient_checkpointing: bool = False
+    input_ids: Optional[torch.Tensor] = None
+
+
+class DecoderOutput(NamedTuple):
+    hidden_states: torch.Tensor
+    all_hidden_states: Optional[Tuple[torch.Tensor, ...]]
+    all_self_attns: Optional[Tuple[torch.Tensor, ...]]
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: Tuple[int, int], dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+) -> torch.Tensor:
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), float("-inf"), device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None) -> torch.Tensor:
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), float("-inf"))  # type: ignore
+
+
+def _rms_norm(
+    hidden_states: torch.Tensor, weight: Optional[torch.Tensor], eps: float, offset: float = 1.0
+) -> torch.Tensor:
+    input_dtype = hidden_states.dtype
+    hidden_states = hidden_states.to(torch.float32)
+    variance = hidden_states.pow(2).mean(-1, keepdim=True)
+    hidden_states = hidden_states * torch.rsqrt(variance + eps)
+    hidden_states = hidden_states.to(input_dtype)
+    if weight is not None:
+        hidden_states = (offset + weight) * hidden_states
+    return hidden_states
+
+
+class RMSNorm(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+        offset: float = 1.0,
+        device: Optional[Union[torch.device, str]] = None,
+    ) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size, device=device))
+        self.variance_epsilon = eps
+        self.offset = offset
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return _rms_norm(hidden_states, self.weight, self.variance_epsilon, offset=self.offset)
+
+
+def get_initial_dt_bias(num_heads: int) -> torch.Tensor:
+    dt_min = 0.001
+    dt_max = 0.1
+    dt = torch.exp(torch.rand(num_heads) * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min))
+    dt = torch.clamp(dt, 1e-4)
+    inv_dt = dt + torch.log(-torch.expm1(-dt))
+    return inv_dt
+
+
+def get_initial_A(num_heads: int) -> torch.Tensor:
+    A = torch.arange(1, num_heads + 1, dtype=torch.float32)
+    return torch.log(A)
+
+
+def _bf16_supported_in_triton() -> bool:
+    # newer torch (2.2.0 and later?) supports bfloat16 even when using Voltas
+    # but triton cannot compile bf16 kernels for Volta
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 8
+
+
+def _get_trition_dtype(dtype: torch.dtype) -> torch.dtype:
+    if dtype != torch.bfloat16:
+        return dtype
+    if _bf16_supported_in_triton():
+        return dtype
+    return torch.float32
+
+
+def ssd_update_state(
+    ssm_state: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor,
+    z: torch.Tensor,
+    dt_bias: torch.Tensor,
+    dt_softplus: bool,
+) -> torch.Tensor:
+    assert ssm_state.dtype == torch.float32
+    if dt.is_cuda:
+        dtype = _get_trition_dtype(x.dtype)
+    else:
+        dtype = x.dtype
+    if dt.is_cuda:
+        f = mamba_ssm.ops.triton.selective_state_update.selective_state_update
+    else:
+        f = mamba_ssm.ops.triton.selective_state_update.selective_state_update_ref
+
+    hidden_size_per_head = x.shape[-1]
+    d_state = B.shape[-1]
+    A = A[:, None, None].expand(-1, hidden_size_per_head, d_state).float()
+    dt = dt[..., None].expand(-1, -1, hidden_size_per_head)
+    dt_bias = dt_bias[:, None].expand(-1, hidden_size_per_head)
+    D = D[:, None].expand(-1, hidden_size_per_head)
+    assert ssm_state.dtype == torch.float32
+    out = f(
+        ssm_state,
+        x.to(dtype),
+        dt.to(dtype),
+        A.float(),
+        B.to(dtype),
+        C.to(dtype),
+        D.float(),
+        z.to(dtype),
+        dt_bias.float(),
+        dt_softplus=dt_softplus,
+    )
+    return out[:, None]  # type: ignore
+
+
+def _ssd_chunk_scan_combined_naive(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor,
+    z: torch.Tensor,
+    dt_bias: torch.Tensor,
+    dt_softplus: bool,
+    seq_idx: torch.Tensor | None,
+    ssm_state: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert ssm_state.dtype == torch.float32
+    length = x.shape[1]
+    ys = []
+    for i in range(length):
+        if i != 0 and seq_idx is not None:
+            ssm_state = torch.where(
+                (seq_idx[:, i - 1] != seq_idx[:, i])[:, None, None, None],
+                torch.zeros_like(ssm_state),
+                ssm_state,
+            )
+        y = ssd_update_state(
+            ssm_state,
+            x[:, i],
+            dt[:, i],
+            A,
+            B[:, i],
+            C[:, i],
+            D,
+            z=z[:, i],
+            dt_bias=dt_bias,
+            dt_softplus=dt_softplus,
+        )
+        ys.append(y)
+    return torch.cat(ys, dim=1), ssm_state
+
+
+def _ssd_chunk_scan_combined_cpu(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor,
+    z: torch.Tensor,
+    dt_bias: torch.Tensor,
+    dt_softplus: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # (bsize, nhead, nchunk, chunk_size)
+    dt = dt.float()  # We want high precision for this before cumsum
+    dt = dt.permute(0, 2, 1).unflatten(2, (-1, chunk_size))  # type: ignore
+    if dt_bias is not None:
+        dt = dt + dt_bias[None, :, None, None]
+    if dt_softplus:
+        dt = F.softplus(dt)
+    dA = dt * A[None, :, None, None]
+    dA_cumsum = torch.cumsum(dA, dim=-1)
+
+    _, _, nheads, _ = x.shape
+    dstate = B.shape[-1]
+    _ = dt.shape[2]
+
+    with torch.profiler.record_function("ssd_chunk_scan_combined_cpu_chunk_state"):
+        # Following is equivalent to `mamba_ssm.ops.triton.ssd_combined.chunk_state_ref(B, x, dt, dA_cumsum)`
+        # But `einsum` in the above function is too slow in CPU.
+        x_ = torch.unflatten(x, 1, (-1, chunk_size))
+        assert B.shape[2] == nheads  # B should be already expanded
+        B_ = torch.unflatten(B, 1, (-1, chunk_size)).to(x.dtype)  # (bsize, nchunk, chunk_size, nheads, dstate)
+        decay_states = torch.exp((dA_cumsum[:, :, :, -1:] - dA_cumsum)).to(x.dtype)
+        dt_ = dt.to(x.dtype)
+
+        # einsum("bclhn,bhcl,bhcl,bclhp->bchpn", B_, decay_states, dt_, x_)
+        B_ = B_.permute(0, 1, 3, 4, 2)  # bchnl
+        tmp = dt_ * decay_states  # bhcl
+        tmp = tmp.permute(0, 2, 1, 3)[:, :, :, None]  # bch1l
+        tmp = B_ * tmp  # bchnl
+        x_ = x_.permute(0, 1, 3, 2, 4)  # bchlp
+        tmp = tmp @ x_  # bchnp
+        states = tmp.permute(0, 1, 2, 4, 3)  # bchpn
+
+    states_dtype = states.dtype
+    if states.dtype not in [torch.float32, torch.float64]:
+        states = states.to(torch.float32)
+    with torch.profiler.record_function("ssd_chunk_scan_combined_cpu_state_passing"):
+        out, last_state = mamba_ssm.ops.triton.ssd_combined.state_passing_ref(
+            states.flatten(start_dim=-2, end_dim=-1),
+            dA_cumsum[:, :, :, -1],
+        )
+    states = torch.unflatten(out, -1, (-1, dstate))
+    last_state = torch.unflatten(last_state, -1, (-1, dstate))
+    states = states.to(states_dtype)
+    with torch.profiler.record_function("ssd_chunk_scan_combined_cpu_chunk_scan"):
+        out = mamba_ssm.ops.triton.ssd_combined.chunk_scan_ref(B, C, x, dt, dA_cumsum, states, D=D, z=z)
+
+    return out, last_state
+
+
+@torch.profiler.record_function("ssd_chunk_scan_combined")
+def ssd_chunk_scan_combined(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    chunk_size: int,
+    D: torch.Tensor,
+    z: torch.Tensor,
+    dt_bias: torch.Tensor,
+    dt_softplus: bool,
+    return_final_states: bool,
+    seq_idx: torch.Tensor | None,
+    ssm_state: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor:
+    if seq_idx is not None:
+        assert seq_idx.dtype == torch.int32
+        assert ssm_state is None
+        assert not return_final_states
+    if ssm_state is not None:
+        assert ssm_state.dtype == torch.float32
+        assert seq_idx is None
+
+    length = x.shape[1]
+
+    """
+    state will be updates by following:
+    ```
+    dt = softplus(dt)
+    dA = exp(dt * A)
+    state_next = state * dA + dB * x
+    ```
+
+    To avoid updating state, we set dt to -inf and x to 0
+    because `softplus(-inf) = 0` and `exp(0) = 1`
+    """
+    pad = (chunk_size - length % chunk_size) % chunk_size
+    x = torch.nn.functional.pad(x, pad=[0, 0, 0, 0, pad, 0], value=0.0)
+    dt = torch.nn.functional.pad(dt, pad=[0, 0, pad, 0], value=float("-inf"))
+    B = torch.nn.functional.pad(B, pad=[0, 0, 0, 0, pad, 0], value=0.0)
+    C = torch.nn.functional.pad(C, pad=[0, 0, 0, 0, pad, 0], value=0.0)
+    z = torch.nn.functional.pad(z, pad=[0, 0, 0, 0, pad, 0], value=0.0)
+    if seq_idx is not None:
+        seq_idx = torch.nn.functional.pad(seq_idx, pad=[pad, 0], value=0)
+
+    length = x.shape[1]
+    assert length % chunk_size == 0, (length, chunk_size)
+
+    if dt.is_cuda:
+        dtype = _get_trition_dtype(x.dtype)
+        out = mamba_ssm.ops.triton.ssd_combined.mamba_chunk_scan_combined(  # type: ignore
+            x.to(dtype),
+            dt.to(dtype),
+            A.float(),
+            B.to(dtype),
+            C.to(dtype),
+            chunk_size,
+            D=D.float(),
+            z=z.to(dtype),
+            initial_states=ssm_state,
+            dt_bias=dt_bias.float(),
+            dt_softplus=dt_softplus,
+            seq_idx=seq_idx,
+            return_final_states=return_final_states,
+        )
+        if return_final_states:
+            return out[0][:, pad:], out[1]
+        else:
+            assert isinstance(out, torch.Tensor)
+            return out[:, pad:]
+    else:
+        if ssm_state is None and seq_idx is None:
+            tmp = _ssd_chunk_scan_combined_cpu(
+                x,
+                dt,
+                A,
+                B,
+                C,
+                chunk_size,
+                D=D,
+                z=z,
+                dt_bias=dt_bias.float(),
+                dt_softplus=dt_softplus,
+            )
+        else:
+            if ssm_state is None:
+                bsize, _, num_heads, channel = x.shape
+                state = B.shape[-1]
+                ssm_state = torch.zeros(bsize, num_heads, channel, state, dtype=torch.float32, device=x.device)
+            tmp = _ssd_chunk_scan_combined_naive(
+                x, dt, A, B, C, D, z=z, dt_bias=dt_bias, dt_softplus=dt_softplus, seq_idx=seq_idx, ssm_state=ssm_state
+            )
+        tmp = (tmp[0][:, pad:], tmp[1])
+        if return_final_states:
+            return tmp
+        else:
+            return tmp[0]
+
+
+def _causal_conv1d_update(
+    conv_state: torch.Tensor, weight: torch.Tensor, xBC: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    dtype = conv_state.dtype
+    xBC = xBC.to(dtype)
+    weight = weight.to(dtype)
+    if conv_state.is_cuda:
+        x = causal_conv1d.causal_conv1d_update(
+            x=xBC,
+            conv_state=conv_state,
+            weight=weight[:, 0, :],
+            activation="silu",
+        )
+        return x, conv_state
+    else:
+        x = causal_conv1d.causal_conv1d_update_ref(
+            x=xBC,
+            conv_state=conv_state,
+            weight=weight[:, 0, :],
+            activation="silu",
+        )
+        return x, conv_state
+
+
+def _causal_conv1d_naive(
+    conv_state: torch.Tensor, weight: torch.Tensor, x: torch.Tensor, seq_idx: torch.Tensor | None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    length = x.shape[-1]
+    out = torch.zeros_like(x)
+    for i in range(length):
+        if i != 0 and seq_idx is not None:
+            conv_state = torch.where(
+                (seq_idx[:, i - 1] != seq_idx[:, i])[:, None, None],
+                torch.zeros_like(conv_state),
+                conv_state,
+            )
+        out[:, :, i : i + 1], conv_state = _causal_conv1d_update(conv_state, weight, x[:, :, i : i + 1])
+    return out, conv_state
+
+
+@torch.profiler.record_function("causal_conv1d")
+def _causal_conv1d(
+    conv_state: torch.Tensor | None, weight: torch.Tensor, x: torch.Tensor, seq_idx: torch.Tensor | None
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    dtype = x.dtype
+    if conv_state is not None:
+        dtype = conv_state.dtype
+        assert seq_idx is None
+    if seq_idx is not None:
+        assert seq_idx.dtype == torch.int32
+        assert conv_state is None
+    weight = weight.to(dtype)
+    x = x.to(dtype)
+
+    return_final_states = conv_state is not None
+    if weight.is_cuda:
+        if x.stride(1) != 1:
+            # to channel-last format
+            x = x.transpose(-1, -2).contiguous().transpose(-1, -2)
+        if conv_state is not None:
+            if conv_state.stride(1) != 1:
+                # to channel-last format
+                conv_state = conv_state.transpose(-1, -2).contiguous().transpose(-1, -2)
+        tmp = causal_conv1d.causal_conv1d_fn(
+            x=x,
+            weight=weight[:, 0, :],
+            initial_states=conv_state,
+            return_final_states=conv_state is not None,
+            activation="silu",
+            seq_idx=seq_idx,
+        )
+        if conv_state is not None:
+            x, conv_state = tmp
+        else:
+            x = tmp
+    else:
+        if seq_idx is None:
+            x, conv_state = causal_conv1d.causal_conv1d_ref(
+                x=x,
+                initial_states=conv_state,
+                return_final_states=True,
+                weight=weight[:, 0, :],
+                activation="silu",
+            )
+        else:
+            if conv_state is None:
+                bsize = x.shape[0]
+                dim = weight.shape[0]
+                d_conv = weight.shape[-1]
+                conv_state = torch.zeros(bsize, dim, d_conv - 1, dtype=x.dtype, device=x.device)
+            x, conv_state = _causal_conv1d_naive(conv_state, weight, x, seq_idx)
+    if return_final_states:
+        return x, conv_state
+    else:
+        return x, None
+
+
+class Mamba(torch.nn.Module):
+    def __init__(self, config: PlamoConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.d_state = config.mamba_d_state
+        self.d_conv = config.mamba_d_conv
+        self.chunk_size = config.mamba_chunk_size
+        self.num_heads = config.mamba_num_heads
+        # TODO add mamba_hidden_size_per_head config (?)
+        self.hidden_size_per_head = config.hidden_size_per_head
+
+        self.intermediate_size = self.num_heads * self.hidden_size_per_head
+
+        self.in_proj = torch.nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False)
+        self.conv1d = torch.nn.Conv1d(
+            in_channels=self.intermediate_size,
+            out_channels=self.intermediate_size,
+            bias=False,  # TODO the original implementation uses bias
+            kernel_size=self.d_conv,
+            groups=self.intermediate_size,
+            padding=0,
+        )
+        self.dt_dim = max(64, self.hidden_size // 16)
+        # Notes:
+        # Mamba2 removes this linear projection for simplicity (Figure 6 in the paper),
+        # but it may degrade the ability of content-length extrapolation.
+        self.bcdt_proj = torch.nn.Linear(
+            self.intermediate_size,
+            self.dt_dim + 2 * self.d_state,
+            bias=False,
+        )
+        self.dt_proj = torch.nn.Linear(self.dt_dim, self.num_heads, bias=False)
+
+        self.dt_bias = torch.nn.Parameter(get_initial_dt_bias(self.num_heads))
+        self.A_log = torch.nn.Parameter(get_initial_A(self.num_heads))
+        self.D = torch.nn.Parameter(torch.ones(self.num_heads))
+
+        # TODO norm weight before gating like Mamba2
+        self.dt_norm_weight = torch.nn.Parameter(torch.ones(self.dt_dim))
+        self.B_norm_weight = torch.nn.Parameter(torch.ones(self.d_state))
+        self.C_norm_weight = torch.nn.Parameter(torch.ones(self.d_state))
+
+        self.out_proj = torch.nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def _no_weight_decay_param_names(self) -> set[str]:
+        return set(["D", "dt_bias", "A_log"])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_states: Optional[PlamoCache] = None,
+    ) -> Tuple[torch.Tensor, Optional[PlamoCache]]:
+        bsize, length, _ = hidden_states.shape
+        is_update = length == 1 and past_states is not None
+
+        bool_mask: torch.Tensor | None = None
+        seq_idx: torch.Tensor | None = None
+        if attention_mask is not None:
+            if len(attention_mask.shape) == 2:
+                attention_mask = attention_mask[None, None].expand(bsize, 1, -1, -1)
+            assert len(attention_mask.shape) == 4
+
+            if past_states is None:
+                # TODO: support seq_idx with cache
+                bool_mask_4d = attention_mask == 0
+                is_first_token = _is_first_token(bool_mask_4d)[:, 0, :]
+                seq_idx = torch.cumsum(is_first_token, dim=-1) - 1
+                seq_idx = seq_idx.to(torch.int32)
+
+            # `generate` function creates attention mask that contains past tokens,
+            # but mamba does not use them
+            attention_mask = attention_mask[:, 0, -length:, -length:]
+            bool_mask = torch.diagonal(attention_mask, dim1=-2, dim2=-1) == 0
+
+        conv_state: torch.Tensor | None
+        ssm_state: torch.Tensor | None
+        if past_states is None:
+            conv_state = None
+            ssm_state = None
+        elif past_states[self.layer_idx] is None:
+            conv_state = torch.zeros(
+                bsize, self.intermediate_size, self.d_conv - 1, dtype=hidden_states.dtype, device=hidden_states.device
+            )
+            ssm_state = torch.zeros(
+                bsize,
+                self.num_heads,
+                self.hidden_size_per_head,
+                self.d_state,
+                dtype=torch.float32,
+                device=hidden_states.device,
+            )
+        else:
+            c = past_states[self.layer_idx]
+            assert isinstance(c, PlamoMambaCache)
+            conv_state = c.conv_state
+            ssm_state = c.ssm_state
+
+        zx = self.in_proj(hidden_states)
+        zx = zx.reshape(bsize, length, self.num_heads, -1)
+        # z: (bsize, length, num_heads, hidden_size_per_head)
+        # x: (bsize, length, num_heads, hidden_size_per_head)
+        z, x = torch.split(zx, [self.hidden_size_per_head, self.hidden_size_per_head], dim=-1)
+
+        # conv
+        x = x.reshape(bsize, length, -1).transpose(1, 2)  # (bsize, intermediate_size, length)
+        if bool_mask is not None:
+            x = torch.where(bool_mask[:, None, :], x, 0.0)
+        if is_update:
+            assert conv_state is not None
+            x, conv_state = _causal_conv1d_update(conv_state, self.conv1d.weight, x)
+        else:
+            x, conv_state = _causal_conv1d(conv_state, self.conv1d.weight, x, seq_idx=seq_idx)
+        x = x.to(dtype=hidden_states.dtype)
+        x = x.transpose(1, 2)  # (bsize, length, intermediate_size)
+        x = x.reshape(bsize, length, -1)
+        # x: (bsize, length, num_heads, hidden_size_per_head)
+        # B: (bsize, length, 1, d_state)
+        # C: (bsize, length, 1, d_state)
+        # dt: (bsize, length, dt_dim)
+        BCdt = self.bcdt_proj(x)
+        x = x.reshape(bsize, length, self.num_heads, -1)
+        B, C, dt = torch.split(BCdt, [self.d_state, self.d_state, self.dt_dim], dim=-1)
+        B = B[:, :, None, :]
+        C = C[:, :, None, :]
+
+        A = -torch.exp(self.A_log.float())  # (num_heads,)
+        dt = _rms_norm(dt, None, self.config.rms_norm_eps) * self.dt_norm_weight[None, None, :]
+        B = _rms_norm(B, None, self.config.rms_norm_eps) * self.B_norm_weight[None, None, None, :]
+        C = _rms_norm(C, None, self.config.rms_norm_eps) * self.C_norm_weight[None, None, None, :]
+
+        # (bsize, length, num_heads, 1)
+        dt = self.dt_proj(dt)[..., None]
+
+        # TODO it may not be required
+        B = B.expand(-1, -1, self.num_heads, -1)
+        C = C.expand(-1, -1, self.num_heads, -1)
+
+        if bool_mask is not None:
+            """
+            state will be updates by following:
+            ```
+            dt = softplus(dt)
+            dA = exp(dt * A)
+            state_next = state * dA + dB * x
+            ```
+
+            To avoid updating state, we set dt to -inf and x to 0
+            because `softplus(-inf) = 0` and `exp(0) = 1`
+            """
+            dt = torch.where(bool_mask[:, :, None, None], dt, float("-inf"))
+            x = torch.where(bool_mask[:, :, None, None], x, 0.0)
+
+        # ssm
+        if is_update:
+            assert ssm_state is not None
+            out = ssd_update_state(
+                ssm_state,
+                x[:, 0],
+                dt[:, 0].reshape(bsize, -1),
+                A,
+                B[:, 0],
+                C[:, 0],
+                D=self.D,
+                z=z[:, 0],
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+            )
+        else:
+            tmp = ssd_chunk_scan_combined(
+                x,
+                dt.reshape(bsize, length, -1),
+                A,
+                B,
+                C,
+                self.chunk_size,
+                D=self.D,
+                z=z,
+                dt_bias=self.dt_bias,
+                dt_softplus=True,
+                return_final_states=past_states is not None,
+                seq_idx=seq_idx,
+                ssm_state=ssm_state,
+            )
+            if past_states is not None:
+                out, ssm_state = tmp
+            else:
+                assert isinstance(tmp, torch.Tensor)
+                out = tmp
+
+        y = self.out_proj(out.reshape(bsize, length, -1))
+
+        if past_states is not None:
+            assert ssm_state is not None
+            assert conv_state is not None
+            past_states.update_mamba(conv_state, ssm_state, self.layer_idx)
+
+        return y, past_states
+
+
+def swa_mask(q_len: int, kv_len: int, device: torch.device, window_size: int) -> torch.Tensor:
+    max_len = max(q_len, kv_len)
+    mask = (
+        torch.ones(max_len, max_len, dtype=torch.bool, device=device)
+        .triu(diagonal=-window_size)
+        .tril(diagonal=window_size)
+    )
+    return mask[-q_len:, -kv_len:]
+
+
+class Attention(torch.nn.Module):
+    def __init__(self, config: PlamoConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        head_dim = config.hidden_size_per_head
+        self.max_position_embeddings = config.max_position_embeddings
+
+        self.q_num_heads = config.num_attention_heads
+        self.qk_dim = self.v_dim = head_dim
+        self.k_num_heads = self.v_num_heads = config.num_key_value_heads
+        assert self.q_num_heads % self.k_num_heads == 0
+        self.n_group = self.q_num_heads // self.k_num_heads
+
+        self.q_proj_dim = self.q_num_heads * self.qk_dim
+        self.k_proj_dim = self.k_num_heads * self.qk_dim
+        self.v_proj_dim = self.k_num_heads * self.v_dim
+        self.qkv_proj = nn.Linear(self.hidden_size, self.q_proj_dim + self.k_proj_dim + self.v_proj_dim, bias=False)
+        self.o_proj = nn.Linear(self.q_num_heads * self.v_dim, self.hidden_size, bias=False)
+
+        self.q_weight = torch.nn.Parameter(torch.ones((self.q_num_heads, self.qk_dim)))
+        self.k_weight = torch.nn.Parameter(torch.ones((self.k_num_heads, self.qk_dim)))
+
+        self.rotary_emb = RotaryEmbedding(self.qk_dim, max_position_embeddings=self.config.attention_window_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_states: Optional[PlamoCache] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[PlamoCache]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv = self.qkv_proj(hidden_states)
+        query_states, key_states, value_states = torch.split(
+            qkv, [self.q_proj_dim, self.k_proj_dim, self.v_proj_dim], dim=-1
+        )
+        query_states = query_states.view(bsz, q_len, self.q_num_heads, self.qk_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.k_num_heads, self.qk_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.v_num_heads, self.v_dim).transpose(1, 2)
+
+        attn_dtype = query_states.dtype
+
+        query_states = _rms_norm(query_states, None, 1e-6) * self.q_weight[None, :, None]
+        key_states = _rms_norm(key_states, None, 1e-6) * self.k_weight[None, :, None]
+
+        if past_states is not None:
+            # reuse k, v, self_attention
+            key_states_new = key_states
+            value_states_new = value_states
+            key_states, value_states = past_states.append_kv(key_states, value_states, self.layer_idx)  # type: ignore
+            past_states.update_attention(key_states_new, value_states_new, self.layer_idx)
+
+        kv_seq_len = key_states.shape[-2]
+        device = hidden_states.device
+        position_ids = torch.arange(kv_seq_len, dtype=torch.long, device=device)[None]
+        q_position_ids = position_ids[:, -query_states.shape[2] :]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states = _rotary_pos_emb(query_states, cos, sin, q_position_ids)
+        key_states = _rotary_pos_emb(key_states, cos, sin, position_ids)
+        # [bsz, nh, t, hd]
+
+        def _expand_kv(t: torch.Tensor, repeat: int, target: int) -> torch.Tensor:
+            t = torch.repeat_interleave(t, repeat, dim=1)
+            return t[:, :target]
+
+        # expand shared kv
+        assert self.k_num_heads == self.v_num_heads
+        key_states = _expand_kv(key_states, self.n_group, self.q_num_heads)
+        value_states = _expand_kv(value_states, self.n_group, self.q_num_heads)
+
+        full_attn = self.layer_idx in self.config.full_attention_idx
+
+        query_states = query_states.to(attn_dtype)
+        key_states = key_states.to(attn_dtype)
+        value_states = value_states.to(attn_dtype)
+        if attention_mask is not None and attention_mask.dtype != torch.bool:
+            attention_mask = attention_mask.to(attn_dtype)
+        if attention_mask is None:
+            if not full_attn:
+                assert key_states.shape[2] <= self.config.attention_window_size + 1
+            attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, is_causal=True)
+        else:
+            if attention_mask.dtype == torch.bool:
+                attention_mask = torch.where(attention_mask, torch.tensor(0.0, dtype=torch.float), float("-inf"))
+            if len(attention_mask.shape) == 2:
+                attention_mask = attention_mask[None, None]
+            assert len(attention_mask.shape) == 4
+
+            if not full_attn:
+                m_swa = swa_mask(
+                    query_states.shape[2], key_states.shape[2], query_states.device, self.config.attention_window_size
+                )
+                # `generate` function creates attention mask that does not consider sliding window
+                m_swa = m_swa[None, None]
+                attention_mask = attention_mask[:, :, -query_states.shape[2] :, -key_states.shape[2] :]
+                attention_mask = torch.where(m_swa, attention_mask, float("-inf"))
+
+            # like AttentionMaskConverter._unmask_unattended in huggingface.transfoermers,
+            # we need to attend to all tokens in masked rows for `scaled_dot_product_attention`
+            bool_mask = torch.logical_not(torch.isneginf(attention_mask))
+            valid_tokens = torch.sum(bool_mask, dim=-1).bool()  # (..., q_len)
+            attention_mask = torch.where(valid_tokens[..., None], attention_mask, float(0.0))
+            attn_output = F.scaled_dot_product_attention(
+                query_states, key_states, value_states, attn_mask=attention_mask
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+
+        attn_output = attn_output.reshape(bsz, q_len, self.q_num_heads * self.v_dim)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_states
+
+
+class MLP(nn.Module):
+    def __init__(self, config: PlamoConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_up_proj = torch.nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=False)
+        self.down_proj = torch.nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.gate_up_proj(x)
+        h = _swiglu(h)
+        return self.down_proj(h)  # type: ignore
+
+
+class PlamoDecoderLayer(torch.nn.Module):
+    def __init__(self, config: PlamoConfig, is_mamba: bool, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.is_mamba = is_mamba
+        self.mixer: torch.nn.Module
+        if is_mamba:
+            self.mixer = Mamba(config, layer_idx)
+        else:
+            self.mixer = Attention(config, layer_idx)
+        self.mlp = MLP(config)
+        """
+        Notes: The model performance was degraded when setting all offsets to 1.
+        """
+        self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, offset=1.0)
+        self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / 5)
+        self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, offset=1.0)
+        self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / (5**1.5))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_state: Optional[PlamoCache] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[Any, ...]:
+        # from LlamaDecoder
+        residual = hidden_states
+        hidden_states = self.pre_mixer_norm(hidden_states)
+
+        # Self Attention
+        if self.is_mamba:
+            hidden_states_sa, present_key_value = self.mixer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                past_states=past_state,
+            )
+            self_attn_weights = None
+        else:
+            hidden_states_sa, self_attn_weights, present_key_value = self.mixer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                past_states=past_state,
+                output_attentions=output_attentions,
+            )
+
+        hidden_states_sa = self.post_mixer_norm(hidden_states_sa)
+        hidden_states = residual + hidden_states_sa
+
+        residual = hidden_states
+        hidden_states = self.pre_mlp_norm(hidden_states)
+
+        # Fully Connected
+        hidden_states_mlp = self.mlp(hidden_states)
+
+        # Residual
+        hidden_states_mlp = self.post_mlp_norm(hidden_states_mlp)
+        hidden_states = residual + hidden_states_mlp
+
+        outputs: Any = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        return outputs  # type: ignore
+
+
+def is_mamba(config: PlamoConfig, i: int) -> bool:
+    if not config.mamba_enabled:
+        return False
+    assert config.mamba_step > 1
+    assert i < config.num_hidden_layers
+
+    if config.num_hidden_layers <= (config.mamba_step // 2):
+        # use attention in last layer
+        return i != config.num_hidden_layers - 1
+    return (i % config.mamba_step) != (config.mamba_step // 2)
+
+
+class PlamoDecoder(torch.nn.Module):
+    def __init__(self, config: PlamoConfig) -> None:
+        super().__init__()
+
+        self.layers = torch.nn.ModuleList(
+            [
+                PlamoDecoderLayer(config, is_mamba=is_mamba(config, i), layer_idx=i)
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(self, x: DecoderInput) -> DecoderOutput:
+        all_hidden_states: Optional[Tuple[torch.Tensor, ...]] = () if x.output_hidden_states else None
+        all_self_attns: Optional[Tuple[torch.Tensor, ...]] = () if x.output_attentions else None
+        hidden_states = x.hidden_states
+
+        for decoder_layer in self.layers:
+            if x.output_hidden_states:
+                assert all_hidden_states is not None
+                all_hidden_states += (hidden_states,)
+
+            if self.training and x.gradient_checkpointing:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    x.attention_mask,
+                    x.past_states,
+                    x.output_attentions,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=x.attention_mask,
+                    past_state=x.past_states,
+                    output_attentions=x.output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if x.output_attentions:
+                assert layer_outputs[1] is not None
+                assert all_self_attns is not None
+                all_self_attns += (layer_outputs[1],)
+        return DecoderOutput(hidden_states, all_hidden_states, all_self_attns)
+
+
+class PlamoPreTrainedModel(PreTrainedModel):  # type: ignore
+    config_class = PlamoConfig
+    _no_split_modules: List[str]
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["PlamoDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
+
+    def _init_weights(self, module: torch.nn.Module) -> None:
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class PlamoModel(PlamoPreTrainedModel):
+    def __init__(self, config: PlamoConfig):
+        super().__init__(config)
+        assert config.eval_attention_n_bit is None
+        assert config.eval_mlp_n_bit is None
+
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        if config.image_feature_size is not None:
+            if config.image_proj_type == "mlp":
+                self.image_proj = MLPImageProjector(config)  # type: ignore
+            elif config.image_proj_type == "linear":
+                self.image_proj = nn.Linear(config.image_feature_size, config.hidden_size, bias=False)  # type: ignore
+            else:
+                raise ValueError(f"Unknown image_proj_type: {config.image_proj_type}")
+        self.layers = PlamoDecoder(config)  # type: ignore
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: torch.nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int, int],
+        inputs_embeds: Optional[torch.Tensor],
+        past_key_values_length: int,
+    ) -> Optional[torch.Tensor]:
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask: Optional[torch.Tensor] = None
+        if input_shape[-1] > 1:
+            assert inputs_embeds is not None
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+            input_shape = (input_shape[0], combined_attention_mask.shape[2])
+
+        if attention_mask is not None:
+            if attention_mask.dim() == 4:
+                # Custom 4D attention mask
+                expanded_attn_mask = attention_mask
+            else:
+                # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+                assert inputs_embeds is not None
+                expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+                    inputs_embeds.device
+                )
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[PlamoCache] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        image_features: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        assert input_ids is not None
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        else:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values.get_seq_length()
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if image_features is not None:
+            assert self.config.image_token_id is not None
+            image_embeds = self.image_proj(image_features)
+            assert image_embeds.shape == inputs_embeds.shape, (image_embeds.shape, inputs_embeds.shape)
+            mask = input_ids == self.config.image_token_id
+            inputs_embeds[mask] = image_embeds[mask]
+
+        # embed positions
+        require_attn_mask = False
+        if not self.training or past_key_values is not None:
+            require_attn_mask = True
+        if seq_length_with_past >= self.config.attention_window_size:
+            require_attn_mask = True
+        if require_attn_mask and attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+            )
+        if attention_mask is not None:
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = PlamoCache(self.config)
+
+        # decoder layers
+        out = self.layers(
+            DecoderInput(
+                hidden_states,
+                attention_mask,
+                past_key_values,
+                output_hidden_states,
+                output_attentions,
+                self.gradient_checkpointing,
+            )
+        )
+        assert isinstance(out, DecoderOutput)
+        hidden_states = out.hidden_states
+        all_hidden_states = out.all_hidden_states
+        all_self_attns = out.all_self_attns
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            assert all_hidden_states is not None
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns] if v is not None
+            )
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class PlamoForCausalLM(PlamoPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    # Without this, the model cannot be loaded into a meta device.
+    # Relevant code:
+    # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/modeling_utils.py#L4376-L4381
+    # https://github.com/huggingface/transformers/blob/v4.44.2/src/transformers/modeling_utils.py#L356
+    # https://github.com/pytorch/pytorch/blob/v2.4.1/torch/nn/modules/module.py#L2068
+    _supports_param_buffer_assignment = False
+
+    def __init__(self, config: PlamoConfig) -> None:
+        super().__init__(config)
+        self.model = PlamoModel(config)
+
+        self.vocab_size = config.vocab_size
+        vocab_size = ((self.vocab_size + 15) // 16) * 16
+        self.lm_head: torch.nn.Module = nn.Linear(config.hidden_size, vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: torch.nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> torch.nn.Module:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None:
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder: PlamoModel) -> None:
+        self.model = decoder
+
+    def get_decoder(self) -> PlamoModel:
+        return self.model
+
+    def forward(  # type: ignore
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[PlamoCache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        image_features: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, LlamaForCausalLM
+
+        >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you consciours? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you."
+        ```"""
+        assert input_ids is not None
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            image_features=image_features,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits[..., : self.vocab_size]
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.Tensor,
+        past_key_values: Optional[PlamoCache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        image_features: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+            if image_features is not None:
+                image_features = image_features[:, -1:, :]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs: Dict[str, Any] = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "image_features": image_features,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values: PlamoCache, beam_idx: torch.Tensor) -> PlamoCache:
+        past_key_values.reorder_cache(beam_idx)
+        return past_key_values
+
+
+class MLPImageProjector(nn.Module):
+    def __init__(self, config: PlamoConfig) -> None:
+        super().__init__()
+        self.config = config
+
+        assert config.image_feature_size is not None  # for typing
+
+        # nn.LayerNorm is not supported by PFVM, so use RMSNorm + Bias instead to approximate this.
+        self.norm0 = RMSNorm(config.image_feature_size, eps=config.rms_norm_eps)
+        self.bias0 = Bias(config.image_feature_size)
+
+        # PFVM doesn't support Linear with bias, so add bias manually afterwards.
+        self.linear1 = nn.Linear(config.image_feature_size, config.hidden_size, bias=False)
+        self.bias1 = Bias(config.hidden_size)
+        self.act1 = nn.GELU()
+
+        self.linear2 = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.bias2 = Bias(config.hidden_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.norm0(hidden_states)
+        hidden_states = self.bias0(hidden_states)
+
+        hidden_states = self.linear1(hidden_states)
+        hidden_states = self.bias1(hidden_states)
+        hidden_states = self.act1(hidden_states)
+
+        hidden_states = self.linear2(hidden_states)
+        hidden_states = self.bias2(hidden_states)
+
+        return hidden_states
+
+
+class Bias(nn.Module):
+    def __init__(self, num_features: int) -> None:
+        super().__init__()
+        self._bias = nn.Parameter(torch.zeros((num_features,)))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return x + self._bias

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|plamo:bos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|plamo:eos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|plamo:pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|plamo:unk|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenization_plamo.py

@@ -0,0 +1,392 @@
+import json
+import math
+import os
+from shutil import copyfile
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+# NOTE: numba does not support type hints for njit: https://github.com/python/mypy/issues/16149
+from numba import njit  # type: ignore[attr-defined]
+from numba.core import types
+from numba.typed import Dict, List
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.jsonl"}
+logger = logging.get_logger(__name__)
+
+INVALID_SCORE = -20000000
+UNKNOWN_SCORE = -10000000
+
+TABLE_PIECE_LENGTH = 0
+TABLE_TOKEN_ID = 1
+TABLE_SCORE = 2
+TABLE_PIECE_ID = 3
+
+PATH_TOKEN_LENGTH = 0
+PATH_TOKEN_ID = 1
+PATH_NUM_TOKENS = 2
+
+
+class AhoCorasick:
+    def __init__(self) -> None:
+        # List of tokens in the vocabulary.
+        self._tokens: list[str]
+
+        # A mapping from a byte code point to a token ID, used for byte fallback.
+        self._bytes: np.ndarray
+
+        # A mapping from a suffix's piece code to a suffix ID.
+        #
+        # Typically, the Aho-Corasick algorithm builds a Trie and adds suffix links between nodes
+        # of the Trie. In this implementation, a suffix ID corresponds to a node in the trie, and
+        # a piece code to an edge (in other words, a pair of a node and the next character).
+        #
+        # A piece code is a 64-bit integer:
+        # - The upper 32 bits store the Unicode code point of the first character.
+        # - The lower 32 bits store the suffix ID of the remaining suffix.
+        #
+        # A suffix ID is an integer indicating the starting position in the _table.
+        self._to_suffix_id: Dict[types.int64, types.int32]
+
+        # Flattened table representing the Trie structure for the Aho-Corasick algorithm.
+        # It stores information including scores for each piece (prefix) within each suffix.
+        # It is flattened for memory efficiency and performance. Suffixes are stored in
+        # lexicographical order of their reversed strings, which improves memory access locality
+        # when exploring new characters starting from the string's end. Pieces within a suffix are
+        # stored in the decreasing order of their lengths.
+        #
+        # Each piece (a prefix fo the suffix) contains four pieces of information:
+        # - TABLE_PIECE_LENGTH: Length of the piece.
+        # - TABLE_TOKEN_ID: Token ID (or -1 if the piece is not a valid token).
+        # - TABLE_SCORE: Score (or INVALID_SCORE if the piece is not a valid token).
+        # - TABLE_PIECE_ID: Piece ID of the suffix.
+        #
+        # Each suffix also includes a sentinel row with a length of 1, a score of UNKNOWN_SCORE,
+        # and a token ID of -1. Sentinel rows are identified by the score being UNKNOWN_SCORE.
+        self._table: np.ndarray
+
+    def build(self, vocab: list[Any]) -> None:
+        self._bytes = np.zeros(256, dtype=np.int32)
+        self._to_suffix_id = Dict.empty(key_type=types.int64, value_type=types.int32)
+
+        # Build suffix_to_score and token_to_token_id.
+        # The suffix_to_score dictionary maps a suffix to its score. It also includes all suffixes
+        # of the token for the Trie structure for the Aho-Corasick algorithm. If a suffix is not a
+        # valid token, its score is set to math.nan.
+        # The token_to_token_id dictionary maps a token to its token ID.
+        suffix_to_score: dict[str, float] = {}
+        token_to_token_id: dict[str, int] = {}
+        self._tokens = []
+        for token_id, row in enumerate(vocab):
+            assert isinstance(row[0], str), row
+            assert isinstance(row[1], (int, float)), row
+
+            token = str(row[0])
+            self._tokens.append(token)
+            token_to_token_id[token] = token_id
+
+            # Special handling for byte tokens.
+            if len(row) > 2 and row[2] == "BYTE":
+                assert len(token) == 6 and token.startswith("<0x") and token.endswith(">"), row[0]
+                self._bytes[int(row[0][3:5], 16)] = token_id
+                continue
+
+            suffix_to_score[token] = float(row[1])
+            # Ensure that all suffixes are included in suffix_to_score.
+            for i in range(1, len(token)):
+                suffix_to_score[token[i:]] = suffix_to_score.get(token[i:], math.nan)
+
+        # Ensure all byte tokens are set.
+        for i in range(256):
+            assert self._bytes[i] != 0, f"Byte token for <0x{i:02X}> is not set."
+
+        # List suffixes in lexicographical order of their reversed strings.
+        suffixes = list(suffix_to_score.keys())
+        suffixes.append("")
+        suffixes.sort(key=lambda x: x[::-1])
+
+        # Build suffix_to_id, which is a mapping from a suffix to a suffix ID, and _to_suffix_id,
+        # which is a mapping from a piece code to a suffix ID.
+        suffix_to_id: dict[str, int] = {}
+        num_pieces = 0
+        for s in suffixes:
+            suffix_to_id[s] = num_pieces
+            if s != "":
+                self._to_suffix_id[ord(s[0]) << 32 | suffix_to_id[s[1:]]] = np.int32(num_pieces)
+            num_pieces += 1 + sum(s[:i] in suffix_to_score for i in range(1, len(s) + 1))
+        assert suffix_to_id[""] == 0, suffix_to_id[""]
+
+        # Build _table, which is a flattened table representing the Trie structure for the Aho-Corasick.
+        self._table = np.zeros((num_pieces, 4), dtype=np.int32)
+        i = 0
+        for suffix in suffixes:
+            # Add all prefixes of the suffix to the table.
+            for piece_length in range(len(suffix), 0, -1):
+                piece = suffix[:piece_length]
+                score = suffix_to_score.get(piece, None)
+                if score is None:
+                    continue
+                self._table[i, TABLE_PIECE_LENGTH] = piece_length
+                self._table[i, TABLE_TOKEN_ID] = token_to_token_id.get(piece, -1)
+                self._table[i, TABLE_SCORE] = round(score * 1e4) if math.isfinite(score) else INVALID_SCORE
+                self._table[i, TABLE_PIECE_ID] = suffix_to_id[piece]
+                i += 1
+
+            # Add a sentinel row.
+            self._table[i, TABLE_PIECE_LENGTH] = 1
+            self._table[i, TABLE_TOKEN_ID] = -1
+            self._table[i, TABLE_SCORE] = UNKNOWN_SCORE
+            i += 1
+        assert i == num_pieces, (i, num_pieces)
+
+    @staticmethod
+    @njit
+    def _encode(
+        to_suffix_id: Dict[types.int64, types.int32],
+        table: np.ndarray,
+        bytes: np.ndarray,
+        data: np.ndarray,
+    ) -> np.ndarray:
+        # Initialize scores array with a high value and set the score at the end to 0.
+        # This array keeps track of the minimum cost (best score) to encode from each position to the end.
+        scores = np.full((len(data) + 1,), 2**60, dtype=np.int64)
+        scores[-1] = 0
+
+        # Path array to store the best path information.
+        # The path array keeps track of token length, token ID, and number of tokens needed to encode.
+        path = np.zeros((len(data) + 1, 3), dtype=np.int32)
+
+        # Initialize suffix_id to 0, which represents the root of the Trie.
+        suffix_id = 0
+
+        # Process the input data from the end to the beginning.
+        for i in range(len(data) - 1, -1, -1):
+            c = data[i]
+
+            # Find the next suffix ID by iterating the suffix IDs of prefixes of the current suffix.
+            # NOTE: If no suffix ID is found, suffix_id will be set to 0.
+            for p in range(suffix_id, len(table)):
+                suffix_id = to_suffix_id.get(c << 32 | table[p, TABLE_PIECE_ID], np.int32(0))
+                # If a next suffix ID is found or a sentinel row is reached, break the loop.
+                if suffix_id > 0 or table[p, TABLE_SCORE] == UNKNOWN_SCORE:
+                    break
+
+            # Update the best path to the current position. If multiple paths have the same score,
+            # this chooses the longest prefix as the best path (table is sorted in the decreasing
+            # order of piece length).
+            for p in range(suffix_id, len(table)):
+                score = table[p, TABLE_SCORE]
+                if score > INVALID_SCORE:
+                    piece_length = table[p, TABLE_PIECE_LENGTH]
+                    s = scores[i + piece_length] - score
+                    if s < scores[i]:
+                        scores[i] = s
+                        path[i, PATH_TOKEN_LENGTH] = piece_length
+                        path[i, PATH_TOKEN_ID] = table[p, TABLE_TOKEN_ID]
+                        path[i, PATH_NUM_TOKENS] = path[i + piece_length, PATH_NUM_TOKENS] + 1
+                        if score == UNKNOWN_SCORE:
+                            # Add number of bytes to represent `c` in UTF-8 (minus 1; 1 is already
+                            # added above).
+                            path[i, PATH_NUM_TOKENS] += (c >= 0x80) + (c >= 0x800) + (c >= 0x10000)
+
+                # If it reaches a sentinel row, break the loop.
+                if score == UNKNOWN_SCORE:
+                    break
+
+        # Decode the best path from the beginning to get the token IDs.
+        pos = 0
+        token_ids = np.zeros(path[0, PATH_NUM_TOKENS], dtype=np.int32)
+        token_pos = 0
+        while pos < len(data):
+            if path[pos, PATH_TOKEN_ID] >= 0:
+                token_ids[token_pos] = path[pos, PATH_TOKEN_ID]
+                token_pos += 1
+            else:
+                # Fall back to byte tokens.
+                c = data[pos]
+                s = 1 + (c >= 0x80) + (c >= 0x800) + (c >= 0x10000)
+                # Add byte tokens representing UTF-8 bytes.
+                for i in range(s):
+                    b = c if s == 1 else (0xF00 >> s) & 0xFF if i == 0 else 0x80
+                    token_ids[token_pos] = bytes[b | ((c >> (s - i - 1) * 6) & 0x3F)]
+                    token_pos += 1
+
+            # Ensure that pos should increase by at least 1.
+            assert path[pos, PATH_TOKEN_LENGTH] > 0, (pos, path[pos])
+            pos += path[pos, PATH_TOKEN_LENGTH]
+
+        return token_ids
+
+    def encode(self, data: str) -> np.ndarray:
+        """Encodes a string into a sequence of token IDs."""
+        return np.asarray(
+            self._encode(
+                self._to_suffix_id,
+                self._table,
+                self._bytes,
+                # Convert a string into a numpy array of Unicode code points.
+                # NOTE: This skips UTF-32 BOM.
+                np.frombuffer(data.encode("utf-32"), dtype=np.int32)[1:],
+            )
+        )
+
+    def encode_as_tokens(self, data: str) -> list[str]:
+        """Encodes a string into a sequence of tokens."""
+        return [self._tokens[token_id] for token_id in self.encode(data)]
+
+
+class PlamoTokenizer(PreTrainedTokenizer):  # type: ignore
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    _save_files = [
+        "special_tokens_map.json",
+        "tokenization_plamo.py",
+        "tokenizer.jsonl",
+        "tokenizer_config.json",
+    ]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        unk_token: str = "<|plamo:unk|>",
+        bos_token: str = "<|plamo:bos|>",
+        eos_token: str = "<|plamo:eos|>",
+        pad_token: str = "<|plamo:pad|>",
+        cls_token: Optional[str] = None,
+        sep_token: Optional[str] = None,
+        mask_token: Optional[str] = None,
+        clean_up_tokenization_spaces: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Tokenizer for PLaMo.
+
+        Args:
+            vocab_file (str): Vocabrary file path.
+            unk_token (str): Unknown token.
+            bos_token (str): Beginning of sentence token.
+            eos_token (str): End of sentence token.
+            pad_token (str): Padding token.
+            cls_token (str):
+                Classification token, to extract a summary of an input sequence leveraging self-attention along the
+                full depth of the model.
+            sep_token (str): Separation token, to separate context and query in an input sequence.
+            mask_token (str): Mask token, to use when training a model with masked-language modeling.
+            clean_up_tokenization_spaces (bool): Whether or not to clean up the tokenization spaces.
+            num_threads (int):
+                Number of threads. This value will be ignored if one of `PLAMO_TOKENIZER_NUM_THREADS` or
+                `RAYON_NUM_THREADS` is set as an environment variable.
+        """
+        if "add_bos_token" not in kwargs:
+            kwargs["add_bos_token"] = False
+        if "add_eos_token" not in kwargs:
+            kwargs["add_eos_token"] = False
+        self.data: list[Any] = [json.loads(line) for line in open(vocab_file, "r", encoding="utf-8")]
+        self.vocab: dict[str, int] = {v[0]: i for i, v in enumerate(self.data)}
+        self.aho_corasick = AhoCorasick()
+        self.aho_corasick.build(self.data)
+        self.vocab_file = vocab_file
+        self.add_bos_token = kwargs["add_bos_token"]
+        self.add_eos_token = kwargs["add_eos_token"]
+
+        super().__init__(
+            vocab_file=vocab_file,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            sep_token=sep_token,
+            mask_token=mask_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    # the functions below are copied from hf transformers LlamaTokenizer's implementation to fix the behaviour of the tokenizer
+    # https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/models/llama/tokenization_llama.py
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["aho_corasick"] = None
+        return state
+
+    def __setstate__(self, d: dict[str, Any]) -> None:
+        self.__dict__ = d
+        self.aho_corasick = AhoCorasick()
+        self.aho_corasick.build(self.data)
+
+    @property
+    def vocab_size(self) -> Any:
+        """Returns vocab size"""
+        return len(self.data)
+
+    def token_to_score(self, token: str) -> Optional[float]:
+        """Returns score of the token"""
+        token_id = self.vocab.get(token, None)
+        return None if token_id is None else self.data[token_id][1]
+
+    def get_vocab(self) -> dict[str, int]:
+        """Returns vocab as a dict"""
+        vocab = self.vocab.copy()
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Converts a sequence of tokens (string) in a single string."""
+        return b"".join(
+            [bytes([int(t[3:5], 16)]) if t.startswith("<0x") else t.encode("utf-8") for t in tokens]
+        ).decode("utf-8", errors="replace")
+
+    def _tokenize(self, text: str) -> Any:
+        """Returns a tokenized string."""
+        return self.aho_corasick.encode_as_tokens(text)
+
+    def _convert_token_to_id(self, token: str) -> Any:
+        """Converts a token (str) in an id using the vocab."""
+        return self.vocab.get(token, 0)
+
+    def _convert_id_to_token(self, index: int) -> Any:
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.data[index][0]
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return ("",)
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "w") as f:
+                for token in self.data:
+                    print(json.dumps(token, ensure_ascii=False), file=f)
+
+        return (out_vocab_file,)

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|plamo:unk|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|plamo:bos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|plamo:eos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|plamo:pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_plamo.PlamoTokenizer",
+      null
+    ]
+  },
+  "bos_token": "<|plamo:bos|>",
+  "chat_template": "{{bos_token}}{% for message in messages %}{% if message['role'] == 'user' %}{{ '\\n\\n### 指示:\\n' + message['content'] }}{% elif message['role'] == 'system' %}{{ '以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。' }}{% elif message['role'] == 'assistant' %}{{ '\\n\\n### 応答:\\n' + message['content'] + eos_token }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '\\n\\n### 応答:\\n' }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": null,
+  "eos_token": "<|plamo:eos|>",
+  "extra_special_tokens": {},
+  "local_file_only": true,
+  "mask_token": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|plamo:pad|>",
+  "sep_token": null,
+  "tokenizer_class": "PlamoTokenizer",
+  "unk_token": "<|plamo:unk|>"
+}

training_args.bin

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