Upload TinyLlamaV2

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "TinyLlamaV2"
+  ],
+  "auto_map": {
+    "AutoConfig": "tiny_llama.TinyLlamaConfig",
+    "AutoModel": "tiny_llama.TinyLlamaV2"
+  },
+  "dropout": 0.1,
+  "dtype": "float32",
+  "hidden_dim": 768,
+  "intermediate_dim": 3072,
+  "max_seq_len": 1024,
+  "model_type": "tiny_llama_v2",
+  "n_heads": 12,
+  "n_kv_heads": 4,
+  "n_layers": 12,
+  "norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "vocab_size": 50257
+}

model.safetensors

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

tiny_llama.py

@@ -0,0 +1,315 @@
+"""TinyTransformer v2 as a Transformers custom model, servable by vLLM's Transformers modeling backend.
+
+A copy of `MLExamples/TinyTransformer/version2_pytorch_fused/tiny_llama_v2.py`'s model from the https://github.com/amd/HPCTrainingExamples — except where v2's code
+can't work when served. Every such change is marked `# SERVING CHANGE N:` below; there are three:
+  1. Attention goes through `ALL_ATTENTION_FUNCTIONS` (keyed by `layer_idx`) instead of calling SDPA
+     directly, so vLLM can plug in its own attention kernel and KV cache —
+     https://docs.vllm.ai/en/latest/models/supported_models/#writing-custom-models. GQA's K/V head
+     repeat moves into that interface too (transformers' `repeat_kv` gives the same head order as
+     v2's `repeat_interleave`): vLLM sizes its KV cache by `n_kv_heads`.
+  2. RoPE angles come from explicit `position_ids` instead of `start_pos + arange(seq_len)`, which
+     assumes every call starts at position `start_pos` — wrong for vLLM's decoding steps and packed
+     batches. vLLM may also hand in precomputed `inputs_embeds` instead of `input_ids`.
+  3. transformers plumbing, so transformers/vLLM can load it: `PretrainedConfig`/`PreTrainedModel`
+     base classes, and `TinyLlamaV2.forward` stops after the final norm and returns the hidden
+     states. vLLM builds this class via `auto_map.AutoModel` and puts its own lm_head on top, tied to
+     `get_input_embeddings()` (`tie_word_embeddings`) — which is exactly v2's `output_proj`, itself
+     tied to `token_embedding`.
+
+Also left out, without affecting outputs:
+  - `record_function` profiler ranges and dropout (inactive at inference).
+  - The unfused Q/K/V and gate/up projection branches, and the manual-attention/explicit-mask branch:
+    the workshop checkpoint uses the fused layout and the SDPA path (`convert_and_push.py` refuses
+    any other checkpoint), so `FusionConfig` is gone too.
+  - `output_proj` and the loss (see 3), `get_fusion_statistics`, and `_init_weights`: random
+    initialization, which loading the checkpoint overwrites.
+"""
+
+from typing import ClassVar
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.integrations.sdpa_attention import sdpa_attention_forward
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+
+class TinyLlamaConfig(PretrainedConfig):
+    """v2's TinyLlamaConfig: same name, fields and defaults. SERVING CHANGE 3: a PretrainedConfig
+    subclass instead of a dataclass, so transformers/vLLM can read it from config.json."""
+
+    model_type = "tiny_llama_v2"
+    # SERVING CHANGE 3: transformers/vLLM read the standard names; alias them to v2's fields.
+    attribute_map: ClassVar[dict[str, str]] = {
+        "hidden_size": "hidden_dim",
+        "num_hidden_layers": "n_layers",
+        "num_attention_heads": "n_heads",
+        "num_key_value_heads": "n_kv_heads",
+        "intermediate_size": "intermediate_dim",
+        "rms_norm_eps": "norm_eps",
+        "max_position_embeddings": "max_seq_len",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 1000,
+        hidden_dim: int = 256,
+        n_layers: int = 4,
+        n_heads: int = 8,
+        n_kv_heads: int = 4,
+        intermediate_dim: int = 512,
+        max_seq_len: int = 128,
+        rope_theta: float = 10000.0,
+        norm_eps: float = 1e-6,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.intermediate_dim = intermediate_dim
+        self.max_seq_len = max_seq_len
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        self.dropout = (
+            dropout  # kept for fidelity with v2's config; unused at inference
+        )
+        # v2 always ties output_proj to token_embedding; vLLM reads this flag to tie its lm_head.
+        kwargs["tie_word_embeddings"] = True
+        super().__init__(**kwargs)
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization - optimized for fusion."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+        x = x / rms
+        return x * self.weight
+
+
+class RotaryPositionEmbedding(nn.Module):
+    """Rotary Position Embeddings, v2's interleaved-pair formula.
+
+    SERVING CHANGE 2: angles are computed for the given `position_ids` rather than cached for
+    `start_pos .. start_pos + seq_len`. v2 builds positions and frequencies in the activations'
+    dtype (float32 for this checkpoint); so does this.
+    """
+
+    def __init__(self, dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+
+    def cos_sin(self, position_ids: torch.Tensor, dtype: torch.dtype):
+        # position_ids: [batch, seq_len] -> cos/sin: [batch, seq_len, 1, dim // 2]
+        t = position_ids.to(dtype)
+        inv_freq = 1.0 / (
+            self.theta
+            ** (torch.arange(0, self.dim, 2, device=t.device, dtype=dtype) / self.dim)
+        )
+        freqs = t[..., None] * inv_freq
+        return torch.cos(freqs).to(dtype).unsqueeze(2), torch.sin(freqs).to(
+            dtype
+        ).unsqueeze(2)
+
+    @staticmethod
+    def apply_rotary_embedding(
+        q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ):
+        batch_size, seq_len, n_heads_q, head_dim = q.shape
+        _, _, n_heads_k, _ = k.shape
+
+        # Reshape for rotary embedding
+        q = q.reshape(batch_size, seq_len, n_heads_q, head_dim // 2, 2)
+        k = k.reshape(batch_size, seq_len, n_heads_k, head_dim // 2, 2)
+
+        q_rot = torch.stack(
+            [q[..., 0] * cos - q[..., 1] * sin, q[..., 0] * sin + q[..., 1] * cos],
+            dim=-1,
+        )
+
+        k_rot = torch.stack(
+            [k[..., 0] * cos - k[..., 1] * sin, k[..., 0] * sin + k[..., 1] * cos],
+            dim=-1,
+        )
+
+        # Reshape back
+        q_rot = q_rot.reshape(batch_size, seq_len, n_heads_q, head_dim)
+        k_rot = k_rot.reshape(batch_size, seq_len, n_heads_k, head_dim)
+
+        return q_rot, k_rot
+
+
+class FusedAttention(nn.Module):
+    """Optimized attention with QKV fusion."""
+
+    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.head_dim = config.hidden_dim // config.n_heads
+        self.scale = self.head_dim**-0.5
+        # SERVING CHANGE 1: what the attention interface reads off the module.
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_key_value_groups = self.n_heads // self.n_kv_heads
+        self.is_causal = True
+
+        # Fused QKV projection - 3 operations combined into 1
+        self.qkv_proj = nn.Linear(
+            config.hidden_dim,
+            (config.n_heads + 2 * config.n_kv_heads) * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            config.n_heads * self.head_dim, config.hidden_dim, bias=False
+        )
+
+    def forward(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+
+        # Fused QKV projection, split into Q, K, V
+        qkv = self.qkv_proj(x)
+        q_size = self.n_heads * self.head_dim
+        kv_size = self.n_kv_heads * self.head_dim
+
+        q = qkv[:, :, :q_size].view(batch_size, seq_len, self.n_heads, self.head_dim)
+        k = qkv[:, :, q_size : q_size + kv_size].view(
+            batch_size, seq_len, self.n_kv_heads, self.head_dim
+        )
+        v = qkv[:, :, q_size + kv_size :].view(
+            batch_size, seq_len, self.n_kv_heads, self.head_dim
+        )
+
+        # Apply rotary position embeddings
+        q, k = RotaryPositionEmbedding.apply_rotary_embedding(q, k, cos, sin)
+
+        # Transpose for attention computation
+        q = q.transpose(1, 2)  # (batch, n_heads, seq_len, head_dim)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # SERVING CHANGE 1: v2 repeats the K/V heads here and calls
+        #     F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
+        # directly. transformers' "sdpa" interface is that same call, GQA repeat included, and lets
+        # vLLM swap in its own kernel and KV cache. It returns [batch, seq_len, n_heads, head_dim],
+        # already transposed back.
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, sdpa_attention_forward
+        )
+        attn_output, _ = attention_interface(
+            self, q, k, v, None, dropout=0.0, scaling=self.scale
+        )
+
+        attn_output = attn_output.contiguous().view(batch_size, seq_len, -1)
+        return self.o_proj(attn_output)
+
+
+class FusedSwiGLU(nn.Module):
+    """Optimized SwiGLU with fused gate/up projections."""
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__()
+        # Fused gate and up projection - 2 operations combined into 1
+        self.gate_up_proj = nn.Linear(
+            config.hidden_dim, 2 * config.intermediate_dim, bias=False
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_dim, config.hidden_dim, bias=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        intermediate = F.silu(gate) * up
+        return self.down_proj(intermediate)
+
+
+class FusedTransformerBlock(nn.Module):
+    """Optimized transformer block with fusion capabilities."""
+
+    def __init__(self, config: TinyLlamaConfig, layer_idx: int):
+        super().__init__()
+        self.attention = FusedAttention(config, layer_idx)
+        self.feed_forward = FusedSwiGLU(config)
+        self.norm1 = RMSNorm(config.hidden_dim, config.norm_eps)
+        self.norm2 = RMSNorm(config.hidden_dim, config.norm_eps)
+
+    def forward(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        # Pre-norm attention with residual connection
+        x = x + self.attention(self.norm1(x), cos, sin)
+        # Pre-norm feed-forward with residual connection
+        x = x + self.feed_forward(self.norm2(x))
+        return x
+
+
+class TinyLlamaV2(PreTrainedModel):
+    """Tiny LLaMA V2 with comprehensive fusion optimizations.
+
+    SERVING CHANGE 3: a PreTrainedModel ending at the final norm (see module docstring).
+    """
+
+    config_class = TinyLlamaConfig
+    _supports_sdpa = True
+    # vLLM's Transformers backend requires this (or source it can inspect) to use its attention.
+    _supports_attention_backend = True
+
+    def __init__(self, config: TinyLlamaConfig):
+        super().__init__(config)
+
+        # Token embeddings
+        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
+
+        # Transformer blocks with fusion
+        self.blocks = nn.ModuleList(
+            [
+                FusedTransformerBlock(config, layer_idx)
+                for layer_idx in range(config.n_layers)
+            ]
+        )
+
+        # Final norm
+        self.norm = RMSNorm(config.hidden_dim, config.norm_eps)
+
+        self.rope = RotaryPositionEmbedding(
+            config.hidden_dim // config.n_heads, config.rope_theta
+        )
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.token_embedding
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        # SERVING CHANGE 2: vLLM may pass embeddings, and always passes positions; default to v2's
+        # 0..seq_len-1. vLLM also passes use_cache/return_dict etc., unused here (**kwargs).
+        x = self.token_embedding(input_ids) if inputs_embeds is None else inputs_embeds
+        if position_ids is None:
+            position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
+        cos, sin = self.rope.cos_sin(position_ids, x.dtype)
+
+        for block in self.blocks:
+            x = block(x, cos, sin)
+
+        # SERVING CHANGE 3: stop at the final norm; vLLM applies its tied lm_head. A tuple — vLLM
+        # reads the hidden states as outputs[0].
+        return (self.norm(x),)