Upload vLLM-compatible full-width draft conversion from 1ed51485

README.md

@@ -0,0 +1,50 @@
+---
+license: cc-by-nc-4.0
+language:
+- en
+datasets:
+- nvidia/Nemotron-ClimbMix
+tags:
+- stairformer
+- asymmetric
+- vllm
+- causal-lm
+- pretraining
+- climbmix
+- custom_code
+---
+
+# atrost/climbmix-stairformer-353m-extracted-nested-94m-1p2b-h100-vllm
+
+vLLM compatibility conversion of `atrost/climbmix-stairformer-353m-extracted-nested-94m-1p2b-h100`.
+
+The source checkpoint is a compact asymmetric draft model: its entry layer uses
+the large 1280-wide embedding space, then later layers and the LM head operate in
+the 160-wide nested space. vLLM's generic Transformers causal wrapper assumes one
+hidden size for the base model and LM head, so this repo expands the checkpoint
+into a full-width StairFormer-shaped model with masked/zeroed suffix weights.
+
+The converted model preserves the source logits up to floating-point differences,
+but it is a compatibility artifact rather than an optimized 94M draft runtime.
+
+- Source checkpoint: `atrost/climbmix-stairformer-353m-extracted-nested-94m-1p2b-h100`
+- Source revision: `1ed51485d6ed35267425e46c77bd9daf86b781aa`
+- Local parity max logits diff: `4.292e-06`
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo_id = "atrost/climbmix-stairformer-353m-extracted-nested-94m-1p2b-h100-vllm"
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    trust_remote_code=True,
+    torch_dtype="auto",
+)
+```
+
+```python
+from vllm import LLM
+
+llm = LLM(model="atrost/climbmix-stairformer-353m-extracted-nested-94m-1p2b-h100-vllm", trust_remote_code=True, model_impl="transformers")
+```

config.json

@@ -0,0 +1,40 @@
+{
+  "architectures": [
+    "StairFormerForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_stairformer.StairFormerModel",
+    "AutoModelForCausalLM": "modeling_stairformer.StairFormerForCausalLM"
+  },
+  "bos_token_id": 50256,
+  "dtype": "float32",
+  "eos_token_id": 50256,
+  "head_dim": 80,
+  "hidden_act": "silu",
+  "hidden_size": 1280,
+  "initializer_range": 0.02,
+  "intermediate_size": 3584,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 8,
+  "pad_token_id": 50256,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "stairformer_dense_entry_layers": 1,
+  "stairformer_nested_loss_alpha": 0.1,
+  "stairformer_prefix_hidden_size": 160,
+  "stairformer_prefix_intermediate_size": 448,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.7.0",
+  "use_cache": true,
+  "vocab_size": 50304
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 50256,
+  "transformers_version": "5.7.0",
+  "use_cache": true
+}

model.safetensors

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

modeling_stairformer.py

@@ -0,0 +1,579 @@
+#!/usr/bin/env python3
+"""Remote-code definitions for StairFormer and asymmetric nested Llama checkpoints."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+
+from transformers import LlamaConfig, LlamaForCausalLM, LlamaPreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer,
+    LlamaModel,
+    LlamaRMSNorm,
+    LlamaRotaryEmbedding,
+)
+
+
+@dataclass
+class StairFormerCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    loss_full: Optional[torch.FloatTensor] = None
+    loss_nested: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    nested_logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class AsymmetricNestedLlamaOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+
+
+def configure_stairformer_config(config: LlamaConfig) -> LlamaConfig:
+    if getattr(config, "_attn_implementation", None) is None:
+        config._attn_implementation = "eager"
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    query_heads_per_kv = config.num_attention_heads // config.num_key_value_heads
+    config.stairformer_prefix_hidden_size = getattr(
+        config,
+        "stairformer_prefix_hidden_size",
+        head_dim * query_heads_per_kv,
+    )
+    config.stairformer_prefix_intermediate_size = getattr(
+        config,
+        "stairformer_prefix_intermediate_size",
+        config.intermediate_size * config.stairformer_prefix_hidden_size // config.hidden_size,
+    )
+    config.stairformer_dense_entry_layers = getattr(config, "stairformer_dense_entry_layers", 1)
+    config.stairformer_nested_loss_alpha = getattr(
+        config,
+        "stairformer_nested_loss_alpha",
+        getattr(config, "stairformer_nested_loss_weight", 0.1),
+    )
+    return config
+
+
+class BlockLowerTriangularLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        prefix_in_features: int,
+        prefix_out_features: int,
+        bias: bool = False,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.in_features = in_features
+        self.out_features = out_features
+        self.prefix_in_features = prefix_in_features
+        self.prefix_out_features = prefix_out_features
+        self.weight = nn.Parameter(torch.empty(out_features, in_features, **factory_kwargs))
+        self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs)) if bias else None
+        mask = torch.ones(out_features, in_features, device=device, dtype=torch.bool)
+        mask[:prefix_out_features, prefix_in_features:] = False
+        self.register_buffer("weight_mask", mask)
+        self.reset_parameters()
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear: nn.Linear,
+        prefix_in_features: int,
+        prefix_out_features: int,
+    ) -> "BlockLowerTriangularLinear":
+        masked = cls(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            prefix_in_features=prefix_in_features,
+            prefix_out_features=prefix_out_features,
+            bias=linear.bias is not None,
+            device=linear.weight.device,
+            dtype=linear.weight.dtype,
+        )
+        with torch.no_grad():
+            masked.weight.copy_(linear.weight)
+            masked.weight.mul_(masked.weight_mask)
+            if linear.bias is not None:
+                masked.bias.copy_(linear.bias)
+        return masked
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / fan_in**0.5 if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+        with torch.no_grad():
+            self.weight.mul_(self.weight_mask)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return F.linear(input, self.weight * self.weight_mask, self.bias)
+
+
+class BlockPrefixNorm(nn.Module):
+    def __init__(self, hidden_size: int, prefix_hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.prefix_hidden_size = prefix_hidden_size
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    @classmethod
+    def from_rms_norm(cls, norm: LlamaRMSNorm, prefix_hidden_size: int) -> "BlockPrefixNorm":
+        block_norm = cls(norm.weight.numel(), prefix_hidden_size, norm.variance_epsilon)
+        with torch.no_grad():
+            block_norm.weight.copy_(norm.weight)
+        return block_norm
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        prefix = hidden_states[..., : self.prefix_hidden_size]
+        suffix = hidden_states[..., self.prefix_hidden_size :]
+        prefix = prefix * torch.rsqrt(prefix.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
+        suffix = suffix * torch.rsqrt(hidden_states.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
+        return self.weight * torch.cat((prefix, suffix), dim=-1).to(input_dtype)
+
+
+class BlockPrefixRMSNorm(BlockPrefixNorm):
+    """Backward-compatible alias; new models instantiate BlockPrefixNorm.
+
+    vLLM's Transformers backend replaces classes whose names end with
+    "RMSNorm". The prefix norm is behaviorally different from a standard RMSNorm,
+    so StairFormerModel uses BlockPrefixNorm to keep it intact under vLLM.
+    """
+
+
+def apply_stairformer_modules(model: LlamaModel, config: LlamaConfig) -> None:
+    prefix_hidden_size = config.stairformer_prefix_hidden_size
+    prefix_intermediate_size = config.stairformer_prefix_intermediate_size
+    dense_entry_layers = config.stairformer_dense_entry_layers
+    prefix_key_value_size = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+
+    for layer_idx, layer in enumerate(model.layers):
+        if layer_idx < dense_entry_layers:
+            continue
+        layer.input_layernorm = BlockPrefixNorm.from_rms_norm(
+            layer.input_layernorm,
+            prefix_hidden_size,
+        )
+        layer.post_attention_layernorm = BlockPrefixNorm.from_rms_norm(
+            layer.post_attention_layernorm,
+            prefix_hidden_size,
+        )
+        attn = layer.self_attn
+        attn.q_proj = BlockLowerTriangularLinear.from_linear(attn.q_proj, prefix_hidden_size, prefix_hidden_size)
+        attn.k_proj = BlockLowerTriangularLinear.from_linear(attn.k_proj, prefix_hidden_size, prefix_key_value_size)
+        attn.v_proj = BlockLowerTriangularLinear.from_linear(attn.v_proj, prefix_hidden_size, prefix_key_value_size)
+        attn.o_proj = BlockLowerTriangularLinear.from_linear(attn.o_proj, prefix_hidden_size, prefix_hidden_size)
+        mlp = layer.mlp
+        mlp.gate_proj = BlockLowerTriangularLinear.from_linear(
+            mlp.gate_proj,
+            prefix_hidden_size,
+            prefix_intermediate_size,
+        )
+        mlp.up_proj = BlockLowerTriangularLinear.from_linear(
+            mlp.up_proj,
+            prefix_hidden_size,
+            prefix_intermediate_size,
+        )
+        mlp.down_proj = BlockLowerTriangularLinear.from_linear(
+            mlp.down_proj,
+            prefix_intermediate_size,
+            prefix_hidden_size,
+        )
+    model.norm = BlockPrefixNorm.from_rms_norm(model.norm, prefix_hidden_size)
+
+
+class StairFormerModel(LlamaModel):
+    config_class = LlamaConfig
+    _supports_attention_backend = True
+    _keys_to_ignore_on_save = [r".*weight_mask"]
+    _keys_to_ignore_on_load_missing = [r".*weight_mask"]
+
+    def __init__(self, config: LlamaConfig):
+        configure_stairformer_config(config)
+        super().__init__(config)
+        apply_stairformer_modules(self, config)
+
+
+class StairFormerForCausalLM(LlamaForCausalLM):
+    config_class = LlamaConfig
+    _supports_attention_backend = True
+    _keys_to_ignore_on_save = [r".*weight_mask"]
+    _keys_to_ignore_on_load_missing = [r".*weight_mask"]
+
+    def __init__(self, config: LlamaConfig):
+        configure_stairformer_config(config)
+        super().__init__(config)
+        self._apply_stairformer_modules()
+
+    @property
+    def prefix_hidden_size(self) -> int:
+        return self.config.stairformer_prefix_hidden_size
+
+    @property
+    def prefix_intermediate_size(self) -> int:
+        return self.config.stairformer_prefix_intermediate_size
+
+    @property
+    def dense_entry_layers(self) -> int:
+        return self.config.stairformer_dense_entry_layers
+
+    @property
+    def prefix_key_value_size(self) -> int:
+        return getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads)
+
+    @property
+    def nested_loss_alpha(self) -> float:
+        return float(self.config.stairformer_nested_loss_alpha)
+
+    def _apply_stairformer_modules(self) -> None:
+        apply_stairformer_modules(self.model, self.config)
+
+    def nested_lm_head_weight(self) -> torch.Tensor:
+        return self.lm_head.weight[:, : self.prefix_hidden_size]
+
+    def nested_logits_from_hidden(self, hidden_states: torch.Tensor, num_logits_to_keep: int = 0) -> torch.Tensor:
+        prefix_hidden = hidden_states[:, -num_logits_to_keep:, : self.prefix_hidden_size]
+        return F.linear(prefix_hidden, self.nested_lm_head_weight()).float()
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[List[torch.FloatTensor], Tuple[Tuple[torch.FloatTensor]]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = 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,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+        output_nested_logits: bool = False,
+    ):
+        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,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+        )
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float()
+        nested_logits = self.nested_logits_from_hidden(hidden_states, num_logits_to_keep) if (output_nested_logits or labels is not None) else None
+        loss = loss_full = loss_nested = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            shift_labels = labels[..., 1:].contiguous().view(-1)
+            loss_full = loss_fct(logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size), shift_labels.to(logits.device))
+            loss_nested = loss_fct(nested_logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size), shift_labels.to(nested_logits.device))
+            alpha = self.nested_loss_alpha
+            loss = (1.0 - alpha) * loss_full + alpha * loss_nested
+        return StairFormerCausalLMOutputWithPast(
+            loss=loss,
+            loss_full=loss_full,
+            loss_nested=loss_nested,
+            logits=logits,
+            nested_logits=nested_logits if output_nested_logits else None,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class AsymmetricNestedLlamaModel(LlamaPreTrainedModel):
+    config_class = LlamaConfig
+    _supports_attention_backend = True
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        configure_stairformer_config(config)
+        self.prefix_hidden_size = config.stairformer_prefix_hidden_size
+        self.prefix_intermediate_size = config.stairformer_prefix_intermediate_size
+        self.dense_entry_layers = config.stairformer_dense_entry_layers
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.entry_rotary_emb = LlamaRotaryEmbedding(config)
+        self.entry_layers = nn.ModuleList(
+            [LlamaDecoderLayer(config, layer_idx=i) for i in range(self.dense_entry_layers)]
+        )
+        self.prefix_config = self._make_prefix_config(config)
+        self.prefix_rotary_emb = LlamaRotaryEmbedding(self.prefix_config)
+        self.layers = nn.ModuleList(
+            [
+                LlamaDecoderLayer(self.prefix_config, layer_idx=self.dense_entry_layers + i)
+                for i in range(config.num_hidden_layers - self.dense_entry_layers)
+            ]
+        )
+        self.norm = LlamaRMSNorm(self.prefix_hidden_size, eps=config.rms_norm_eps)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def get_decoder(self):
+        return self
+
+    def _make_prefix_config(self, config: LlamaConfig) -> LlamaConfig:
+        prefix_config = LlamaConfig(**config.to_dict())
+        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        prefix_config.hidden_size = self.prefix_hidden_size
+        prefix_config.intermediate_size = self.prefix_intermediate_size
+        prefix_config.num_attention_heads = self.prefix_hidden_size // head_dim
+        prefix_config.num_key_value_heads = 1
+        prefix_config.tie_word_embeddings = False
+        if getattr(prefix_config, "_attn_implementation", None) is None:
+            prefix_config._attn_implementation = "eager"
+        return prefix_config
+
+    @staticmethod
+    def _make_causal_mask(
+        batch_size: int,
+        seq_len: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        min_value = torch.finfo(dtype).min
+        mask = torch.full((seq_len, seq_len), min_value, dtype=dtype, device=device)
+        mask = torch.triu(mask, diagonal=1)
+        mask = mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len)
+        if attention_mask is not None:
+            padding_mask = (1.0 - attention_mask[:, None, None, :].to(dtype)) * min_value
+            mask = mask + padding_mask
+        return mask
+
+    @staticmethod
+    def _run_decoder_layer(layer: nn.Module, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
+        try:
+            layer_outputs = layer(hidden_states, **kwargs)
+        except TypeError:
+            kwargs.pop("position_embeddings", None)
+            layer_outputs = layer(hidden_states, **kwargs)
+        return layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[List[torch.FloatTensor], Tuple[Tuple[torch.FloatTensor]]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: bool = True,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+        **kwargs,
+    ):
+        del output_attentions, output_hidden_states, cache_position, num_logits_to_keep
+        if past_key_values is not None:
+            raise NotImplementedError("AsymmetricNestedLlamaModel does not support KV cache reuse yet.")
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        if inputs_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = inputs_embeds
+        batch_size, seq_len, _ = hidden_states.shape
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)
+        causal_mask = self._make_causal_mask(
+            batch_size,
+            seq_len,
+            hidden_states.dtype,
+            hidden_states.device,
+            attention_mask,
+        )
+        entry_position_embeddings = self.entry_rotary_emb(hidden_states, position_ids)
+        for layer in self.entry_layers:
+            hidden_states = self._run_decoder_layer(
+                layer,
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                use_cache=False,
+                position_embeddings=entry_position_embeddings,
+                **kwargs,
+            )
+        hidden_states = hidden_states[..., : self.prefix_hidden_size]
+        prefix_position_embeddings = self.prefix_rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = self._run_decoder_layer(
+                layer,
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                use_cache=False,
+                position_embeddings=prefix_position_embeddings,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        if not return_dict:
+            return (hidden_states,)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=None if use_cache else None,
+        )
+
+
+class AsymmetricNestedLlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
+    config_class = LlamaConfig
+    _supports_attention_backend = True
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__(config)
+        configure_stairformer_config(config)
+        self.prefix_hidden_size = config.stairformer_prefix_hidden_size
+        self.prefix_intermediate_size = config.stairformer_prefix_intermediate_size
+        self.dense_entry_layers = config.stairformer_dense_entry_layers
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.entry_rotary_emb = LlamaRotaryEmbedding(config)
+        self.entry_layers = nn.ModuleList(
+            [LlamaDecoderLayer(config, layer_idx=i) for i in range(self.dense_entry_layers)]
+        )
+        self.prefix_config = self._make_prefix_config(config)
+        self.prefix_rotary_emb = LlamaRotaryEmbedding(self.prefix_config)
+        self.layers = nn.ModuleList(
+            [
+                LlamaDecoderLayer(self.prefix_config, layer_idx=self.dense_entry_layers + i)
+                for i in range(config.num_hidden_layers - self.dense_entry_layers)
+            ]
+        )
+        self.norm = LlamaRMSNorm(self.prefix_hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(self.prefix_hidden_size, self.vocab_size, bias=False)
+        self.post_init()
+
+    get_input_embeddings = AsymmetricNestedLlamaModel.get_input_embeddings
+    set_input_embeddings = AsymmetricNestedLlamaModel.set_input_embeddings
+    _make_prefix_config = AsymmetricNestedLlamaModel._make_prefix_config
+    _make_causal_mask = staticmethod(AsymmetricNestedLlamaModel._make_causal_mask)
+    _run_decoder_layer = staticmethod(AsymmetricNestedLlamaModel._run_decoder_layer)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[List[torch.FloatTensor], Tuple[Tuple[torch.FloatTensor]]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: bool = True,
+        cache_position: Optional[torch.LongTensor] = None,
+        num_logits_to_keep: int = 0,
+        **kwargs,
+    ):
+        del output_attentions, output_hidden_states, cache_position
+        if past_key_values is not None:
+            raise NotImplementedError("AsymmetricNestedLlamaForCausalLM does not support KV cache reuse yet.")
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        if inputs_embeds is None:
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = inputs_embeds
+        batch_size, seq_len, _ = hidden_states.shape
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)
+        causal_mask = self._make_causal_mask(
+            batch_size,
+            seq_len,
+            hidden_states.dtype,
+            hidden_states.device,
+            attention_mask,
+        )
+        entry_position_embeddings = self.entry_rotary_emb(hidden_states, position_ids)
+        for layer in self.entry_layers:
+            hidden_states = self._run_decoder_layer(
+                layer,
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                use_cache=False,
+                position_embeddings=entry_position_embeddings,
+                **kwargs,
+            )
+        hidden_states = hidden_states[..., : self.prefix_hidden_size]
+        prefix_position_embeddings = self.prefix_rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = self._run_decoder_layer(
+                layer,
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                use_cache=False,
+                position_embeddings=prefix_position_embeddings,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float()
+        loss = None
+        if labels is not None:
+            loss = CrossEntropyLoss()(
+                logits[..., :-1, :].contiguous().view(-1, self.vocab_size),
+                labels[..., 1:].contiguous().view(-1).to(logits.device),
+            )
+        if not return_dict:
+            output = (logits, hidden_states)
+            return (loss,) + output if loss is not None else output
+        return AsymmetricNestedLlamaOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=hidden_states,
+            past_key_values=None if use_cache else None,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        **kwargs,
+    ):
+        del past_key_values
+        model_inputs = {"inputs_embeds": inputs_embeds} if inputs_embeds is not None else {"input_ids": input_ids}
+        model_inputs.update(
+            attention_mask=attention_mask,
+            use_cache=False,
+        )
+        return model_inputs

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 2048,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}