Add lmul-attention version of Qwen/Qwen2-7B-Instruct

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,5 @@
+{
+  "<|endoftext|>": 151643,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644
+}

chat_template.jinja

@@ -0,0 +1,6 @@
+{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system
+You are a helpful assistant.<|im_end|>
+' }}{% endif %}{{'<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>' + '
+'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 3584,
+  "initializer_range": 0.02,
+  "intermediate_size": 18944,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 28,
+  "model_type": "qwen2",
+  "num_attention_heads": 28,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 4,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": 131072,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.52.4",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 152064
+}

lmul.py

@@ -0,0 +1,88 @@
+"""
+PyTorch-native implementation of the L-Mul algorithm.
+"""
+from __future__ import annotations
+
+import torch
+
+__all__ = [
+    "l_mul_tensor",
+    "l_mul_attention",
+]
+
+def l_mul_tensor(
+    x: torch.Tensor, y: torch.Tensor, *, offset: float = 1.0 / 16.0
+) -> torch.Tensor:
+    """
+    Approximates `x * y` element-wise using the L-Mul algorithm.
+    """
+    sign = torch.sign(x) * torch.sign(y)
+    
+    x_abs = torch.abs(x)
+    y_abs = torch.abs(y)
+    
+    # Decompose tensors into mantissa and exponent.
+    # torch.frexp gives mantissa in [0.5, 1.0)
+    mx, ex = torch.frexp(x_abs)
+    my, ey = torch.frexp(y_abs)
+
+    # The paper's logic implies a mantissa in [1.0, 2.0).
+    # We reconstruct this by multiplying the mantissa by 2 and adjusting the exponent.
+    mant_a = mx * 2.0
+    mant_b = my * 2.0
+    exp_a = ex - 1
+    exp_b = ey - 1
+
+    # Approximate multiplication using the L-Mul formula
+    result_mant = (mant_a - 1.0) + (mant_b - 1.0) + 1.0 + offset
+    result_exp = exp_a + exp_b
+    
+    # Reconstruct the final number
+    result = torch.ldexp(result_mant, result_exp)
+    
+    # Apply the correct sign and handle zero inputs
+    final_result = sign * result
+    final_result[x == 0] = 0
+    final_result[y == 0] = 0
+    
+    return final_result
+
+def l_mul_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    mask: torch.Tensor | None = None,
+    dropout: torch.nn.Module | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Scaled dot-product attention where matrix multiplications are replaced
+    by the L-Mul approximation for performance.
+    """
+    d_k = query.size(-1)
+
+    # Approximate Q @ K.T using l_mul_tensor.
+    # This requires broadcasting and summing to perform the matrix multiplication.
+    scores = l_mul_tensor(
+        query.unsqueeze(-1),
+        key.transpose(-2, -1).unsqueeze(-3)
+    ).sum(-2)
+    
+    scores = scores / (d_k ** 0.5)
+
+    if mask is not None:
+        scores = scores.masked_fill(mask == 0, -1e9)
+
+    attn_probs = torch.nn.functional.softmax(scores, dim=-1)
+    if dropout is not None:
+        attn_probs = dropout(attn_probs)
+
+    # Approximate Attn @ V using l_mul_tensor
+    output = l_mul_tensor(
+        attn_probs.unsqueeze(-1),
+        value.unsqueeze(-3)
+    ).sum(-2)
+    
+    return output, attn_probs
+
+
+

model.safetensors

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

special_tokens_map.json

@@ -0,0 +1,20 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,43 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}