初始化项目,由ModelHub XC社区提供模型
Model: squ11z1/Gravity-2 Source: Original Platform
This commit is contained in:
38
.gitattributes
vendored
Normal file
38
.gitattributes
vendored
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.model filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||||
|
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||||
|
*.xz 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
|
||||||
|
gravity-2-f16.gguf filter=lfs diff=lfs merge=lfs -text
|
||||||
|
gravity-2-Q4_K_M.gguf filter=lfs diff=lfs merge=lfs -text
|
||||||
85
README.md
Normal file
85
README.md
Normal file
@@ -0,0 +1,85 @@
|
|||||||
|
---
|
||||||
|
license: mit
|
||||||
|
pipeline_tag: text-generation
|
||||||
|
tags: [research, experimental, gravity-attention, qwen2]
|
||||||
|
---
|
||||||
|
|
||||||
|
# Gravity-2
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
**Experimental research model by squ11z1.**
|
||||||
|
|
||||||
|
A 3B reasoning model in which the standard
|
||||||
|
scaled-dot-product attention is replaced by a physically-motivated **gravity attention**,
|
||||||
|
then adapted with LoRA. This card documents a **stage-1 proof-of-mechanism**
|
||||||
|
|
||||||
|
## The experiment
|
||||||
|
|
||||||
|
Transformer attention scores tokens by **alignment** — the dot product `q·k`. Gravity-2
|
||||||
|
asks a different question: *what if tokens attended by **proximity** instead?* We replace
|
||||||
|
the score with an inverse-square law borrowed from gravitation — each token is pulled
|
||||||
|
toward others that are close in query/key space, weighted by a learnable per-head "mass":
|
||||||
|
|
||||||
|
```
|
||||||
|
M_h²
|
||||||
|
score(i, j) = ───────────────────── → softmax_j( score )
|
||||||
|
‖q_i − k_j‖² + ε
|
||||||
|
```
|
||||||
|
|
||||||
|
- **M_h = softplus(gravity_mass_log[h])** — one learnable mass per **query head** (16 / layer),
|
||||||
|
initialised at 0.5; `softplus` keeps it strictly positive.
|
||||||
|
- **‖q_i − k_j‖²** — squared L2 distance, computed stably as `‖q‖² + ‖k‖² − 2·q·k`.
|
||||||
|
- **ε = 0.1** — softening length; prevents the `q → k` singularity.
|
||||||
|
- The raw gravity scores are then passed through the **usual softmax** (see Limitations).
|
||||||
|
|
||||||
|
### Why it's interesting
|
||||||
|
- **Different inductive bias.** Dot-product attention rewards directional alignment;
|
||||||
|
inverse-distance rewards *locality* in the learned embedding geometry — a metric prior
|
||||||
|
rather than an inner-product one.
|
||||||
|
- **Interpretable per-head masses.** Each head learns a scalar "mass" controlling how
|
||||||
|
sharply it concentrates — a compact, inspectable knob (see `figures/04_mass_heatmap.png`).
|
||||||
|
- **A bridge to physics-style sparsity.** An inverse-square field is naturally local, which
|
||||||
|
later stages (pruning / QUBO, "Gravity-6") aim to exploit for structured sparsity.
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
Qwen2-3B class: 36 layers, hidden 2048, **16 query heads / 2 KV heads (GQA, group size 8)**,
|
||||||
|
head_dim 128. The 2 KV heads are `repeat_kv`-expanded to 16 before the distance, so each
|
||||||
|
query head gets its own mass. Integrated via the transformers-5.x `AttentionInterface`
|
||||||
|
(a registered `"gravity"` op + eager causal-mask reuse) — RoPE / KV-cache / masking are
|
||||||
|
left to the framework; only the score function changes.
|
||||||
|
|
||||||
|
## Results
|
||||||
|
|
||||||
|
| | |
|
||||||
|
|---|---|
|
||||||
|
|  |  |
|
||||||
|
|  |  |
|
||||||
|
|  |  |
|
||||||
|
|
||||||
|
## Honest limitations
|
||||||
|
- **Not "pure" gravity.** The inverse-square scores are renormalised by a **softmax on top**
|
||||||
|
(`softmax_j(M²/(d²+ε))`). Without it training was unstable, but it means this is a
|
||||||
|
*distance-biased softmax attention*, not a literal gravitational field — the normalisation
|
||||||
|
reintroduces global competition between keys.
|
||||||
|
- **MHA → GQA transfer is an open question.** The mechanism was first prototyped on MHA
|
||||||
|
(1 KV head per query head). Here it runs on GQA by `repeat_kv`-expanding 2 KV heads to 16
|
||||||
|
and giving each query head its own mass; whether this is the right granularity (vs. one
|
||||||
|
mass per KV group) is **unresolved** and may matter for convergence.
|
||||||
|
- **Loading requires the patch** (below). **GGUF builds run standard attention, not gravity**
|
||||||
|
(llama.cpp has no kernel for `M²/(‖q−k‖²+ε)`) — the `*.gguf` files are format placeholders
|
||||||
|
and produce incorrect output.
|
||||||
|
|
||||||
|
## Loading (requires the gravity patch)
|
||||||
|
```bash
|
||||||
|
python load_gravity2.py # from_pretrained -> patch_qwen_with_gravity -> load gravity_mass_log.pt
|
||||||
|
```
|
||||||
|
Weights are LoRA-merged into the base but were trained under gravity scoring; loading them
|
||||||
|
under vanilla attention gives garbage. `config.json` ships `_attn_implementation="eager"`
|
||||||
|
only so the checkpoint loads — the patch switches it to gravity.
|
||||||
|
|
||||||
|
## License & attribution
|
||||||
|
Released under the **MIT License**. This is a **derivative work of
|
||||||
|
[`WeiboAI/VibeThinker-3B`](https://huggingface.co/WeiboAI/VibeThinker-3B)** (the base model
|
||||||
|
for the experiment), which is distributed under the **MIT License**; that license is
|
||||||
|
inherited here and the original authors are credited accordingly.
|
||||||
54
chat_template.jinja
Normal file
54
chat_template.jinja
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
{%- if tools %}
|
||||||
|
{{- '<|im_start|>system\n' }}
|
||||||
|
{%- if messages[0]['role'] == 'system' %}
|
||||||
|
{{- messages[0]['content'] }}
|
||||||
|
{%- else %}
|
||||||
|
{{- 'You are a helpful assistant.' }}
|
||||||
|
{%- endif %}
|
||||||
|
{{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||||
|
{%- for tool in tools %}
|
||||||
|
{{- "\n" }}
|
||||||
|
{{- tool | tojson }}
|
||||||
|
{%- endfor %}
|
||||||
|
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||||
|
{%- else %}
|
||||||
|
{%- if messages[0]['role'] == 'system' %}
|
||||||
|
{{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
|
||||||
|
{%- else %}
|
||||||
|
{{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endif %}
|
||||||
|
{%- for message in messages %}
|
||||||
|
{%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
|
||||||
|
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
|
||||||
|
{%- elif message.role == "assistant" %}
|
||||||
|
{{- '<|im_start|>' + message.role }}
|
||||||
|
{%- if message.content %}
|
||||||
|
{{- '\n' + message.content }}
|
||||||
|
{%- endif %}
|
||||||
|
{%- for tool_call in message.tool_calls %}
|
||||||
|
{%- if tool_call.function is defined %}
|
||||||
|
{%- set tool_call = tool_call.function %}
|
||||||
|
{%- endif %}
|
||||||
|
{{- '\n<tool_call>\n{"name": "' }}
|
||||||
|
{{- tool_call.name }}
|
||||||
|
{{- '", "arguments": ' }}
|
||||||
|
{{- tool_call.arguments | tojson }}
|
||||||
|
{{- '}\n</tool_call>' }}
|
||||||
|
{%- endfor %}
|
||||||
|
{{- '<|im_end|>\n' }}
|
||||||
|
{%- elif message.role == "tool" %}
|
||||||
|
{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
|
||||||
|
{{- '<|im_start|>user' }}
|
||||||
|
{%- endif %}
|
||||||
|
{{- '\n<tool_response>\n' }}
|
||||||
|
{{- message.content }}
|
||||||
|
{{- '\n</tool_response>' }}
|
||||||
|
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||||
|
{{- '<|im_end|>\n' }}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endif %}
|
||||||
|
{%- endfor %}
|
||||||
|
{%- if add_generation_prompt %}
|
||||||
|
{{- '<|im_start|>assistant\n' }}
|
||||||
|
{%- endif %}
|
||||||
69
config.json
Normal file
69
config.json
Normal file
@@ -0,0 +1,69 @@
|
|||||||
|
{
|
||||||
|
"architectures": [
|
||||||
|
"Qwen2ForCausalLM"
|
||||||
|
],
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"bos_token_id": 151643,
|
||||||
|
"dtype": "bfloat16",
|
||||||
|
"eos_token_id": 151643,
|
||||||
|
"hidden_act": "silu",
|
||||||
|
"hidden_size": 2048,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 11008,
|
||||||
|
"layer_types": [
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention",
|
||||||
|
"full_attention"
|
||||||
|
],
|
||||||
|
"max_position_embeddings": 131072,
|
||||||
|
"max_window_layers": 36,
|
||||||
|
"model_type": "qwen2",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_hidden_layers": 36,
|
||||||
|
"num_key_value_heads": 2,
|
||||||
|
"pad_token_id": null,
|
||||||
|
"rms_norm_eps": 1e-06,
|
||||||
|
"rope_parameters": {
|
||||||
|
"rope_theta": 1000000.0,
|
||||||
|
"rope_type": "default"
|
||||||
|
},
|
||||||
|
"sliding_window": null,
|
||||||
|
"tie_word_embeddings": true,
|
||||||
|
"transformers_version": "5.12.1",
|
||||||
|
"use_cache": false,
|
||||||
|
"use_sliding_window": false,
|
||||||
|
"vocab_size": 151936
|
||||||
|
}
|
||||||
BIN
figures/01_loss.png
Normal file
BIN
figures/01_loss.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 41 KiB |
BIN
figures/02_mass_band.png
Normal file
BIN
figures/02_mass_band.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 58 KiB |
BIN
figures/03_gradnorm.png
Normal file
BIN
figures/03_gradnorm.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
BIN
figures/04_mass_heatmap.png
Normal file
BIN
figures/04_mass_heatmap.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 57 KiB |
BIN
figures/05_aer_entropy.png
Normal file
BIN
figures/05_aer_entropy.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 56 KiB |
BIN
figures/06_concept.png
Normal file
BIN
figures/06_concept.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 77 KiB |
6
generation_config.json
Normal file
6
generation_config.json
Normal file
@@ -0,0 +1,6 @@
|
|||||||
|
{
|
||||||
|
"bos_token_id": 151643,
|
||||||
|
"eos_token_id": 151643,
|
||||||
|
"max_new_tokens": 2048,
|
||||||
|
"transformers_version": "5.12.1"
|
||||||
|
}
|
||||||
3
gravity-2-Q4_K_M.gguf
Normal file
3
gravity-2-Q4_K_M.gguf
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:68492d7cc67d8796ff80a06347a8f7c7754d164a711f531bce4dfedfc4b646f7
|
||||||
|
size 1929902304
|
||||||
3
gravity-2-f16.gguf
Normal file
3
gravity-2-f16.gguf
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:b1cab036db87b080b897cfc4637219bece09c9bee7634c9bce5009abcf5d0b53
|
||||||
|
size 6178316512
|
||||||
110
gravity_attention_qwen.py
Normal file
110
gravity_attention_qwen.py
Normal file
@@ -0,0 +1,110 @@
|
|||||||
|
"""
|
||||||
|
Gravity-2 attention for Qwen2 / VibeThinker-3B (transformers 5.x interface).
|
||||||
|
|
||||||
|
Replaces softmax(QKᵀ·scaling) with a physically-motivated score:
|
||||||
|
|
||||||
|
score(i,j) = M_h² / (||q_i − k_j||² + eps) # then standard softmax over j
|
||||||
|
|
||||||
|
• M_h = softplus(gravity_mass_log[h]) — one learnable mass per QUERY head (16/layer)
|
||||||
|
• ||q_i − k_j||² = ||q||² + ||k||² − 2·q·k # GQA: K repeated 2→16 first
|
||||||
|
• eps guards the singularity at q==k
|
||||||
|
|
||||||
|
Integration uses the transformers-5.x AttentionInterface dispatch (NOT a forward
|
||||||
|
monkeypatch): we register a "gravity" attention fn + alias its mask to "eager" so
|
||||||
|
the framework keeps building the additive causal mask, handling RoPE/cache itself.
|
||||||
|
"""
|
||||||
|
import math
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from transformers.models.qwen2.modeling_qwen2 import repeat_kv
|
||||||
|
from transformers.modeling_utils import AttentionInterface
|
||||||
|
from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
|
||||||
|
|
||||||
|
ATTN_NAME = "gravity"
|
||||||
|
|
||||||
|
|
||||||
|
def gravity_attention_forward(module, query, key, value, attention_mask,
|
||||||
|
scaling=None, dropout=0.0, **kwargs):
|
||||||
|
"""AttentionInterface contract.
|
||||||
|
|
||||||
|
query: (B, Hq, Tq, D) key/value: (B, Hkv, Tk, D)
|
||||||
|
returns: (attn_output (B, Tq, Hq, D), attn_weights (B, Hq, Tq, Tk))
|
||||||
|
`scaling` is intentionally ignored — gravity replaces the 1/√d scale.
|
||||||
|
"""
|
||||||
|
# GQA: expand 2 KV heads up to 16 so distances live in per-query-head space
|
||||||
|
key = repeat_kv(key, module.num_key_value_groups)
|
||||||
|
value = repeat_kv(value, module.num_key_value_groups)
|
||||||
|
|
||||||
|
# ||q_i - k_j||^2 in fp32 for numerical stability
|
||||||
|
q = query.float()
|
||||||
|
k = key.float()
|
||||||
|
q_sq = (q * q).sum(-1, keepdim=True) # (B,Hq,Tq,1)
|
||||||
|
k_sq = (k * k).sum(-1, keepdim=True).transpose(-2, -1) # (B,Hq,1,Tk)
|
||||||
|
qk = torch.matmul(q, k.transpose(-2, -1)) # (B,Hq,Tq,Tk)
|
||||||
|
d_sq = (q_sq + k_sq - 2.0 * qk).clamp_min(0.0)
|
||||||
|
|
||||||
|
mass = F.softplus(module.gravity_mass_log).float().view(1, -1, 1, 1) # (1,Hq,1,1)
|
||||||
|
scores = (mass * mass) / (d_sq + module.gravity_eps) # (B,Hq,Tq,Tk), fp32
|
||||||
|
|
||||||
|
if attention_mask is not None:
|
||||||
|
# additive causal mask (eager-style), already correct length
|
||||||
|
scores = scores + attention_mask[..., : key.shape[-2]].float()
|
||||||
|
|
||||||
|
attn = F.softmax(scores, dim=-1, dtype=torch.float32)
|
||||||
|
|
||||||
|
# AER: optionally stash mean per-row attention entropy (flag-gated, ~free when off)
|
||||||
|
if getattr(module, "_capture_entropy", False):
|
||||||
|
ent = -(attn.clamp_min(1e-12) * attn.clamp_min(1e-12).log()).sum(-1)
|
||||||
|
module._last_entropy = ent.mean().detach()
|
||||||
|
|
||||||
|
attn = F.dropout(attn, p=dropout, training=module.training)
|
||||||
|
attn = attn.to(value.dtype)
|
||||||
|
|
||||||
|
out = torch.matmul(attn, value) # (B,Hq,Tq,D)
|
||||||
|
out = out.transpose(1, 2).contiguous() # (B,Tq,Hq,D)
|
||||||
|
return out, attn
|
||||||
|
|
||||||
|
|
||||||
|
_REGISTERED = False
|
||||||
|
|
||||||
|
|
||||||
|
def _register():
|
||||||
|
global _REGISTERED
|
||||||
|
if _REGISTERED:
|
||||||
|
return
|
||||||
|
AttentionInterface.register(ATTN_NAME, gravity_attention_forward)
|
||||||
|
# reuse the eager additive-mask builder for our custom impl
|
||||||
|
ALL_MASK_ATTENTION_FUNCTIONS.register(ATTN_NAME, ALL_MASK_ATTENTION_FUNCTIONS["eager"])
|
||||||
|
_REGISTERED = True
|
||||||
|
|
||||||
|
|
||||||
|
def patch_qwen_with_gravity(model, eps: float = 0.1, init_mass: float = 0.5):
|
||||||
|
"""Add per-head gravity_mass_log to every Qwen2 self-attn and switch dispatch.
|
||||||
|
|
||||||
|
Leaves q/k/v/o_proj weights untouched. gravity_mass_log kept in fp32.
|
||||||
|
"""
|
||||||
|
_register()
|
||||||
|
init_log = math.log(math.exp(init_mass) - 1.0) # softplus^{-1}(init_mass)
|
||||||
|
H = model.config.num_attention_heads
|
||||||
|
n = 0
|
||||||
|
for layer in model.model.layers:
|
||||||
|
attn = layer.self_attn
|
||||||
|
dev = attn.q_proj.weight.device
|
||||||
|
attn.gravity_mass_log = nn.Parameter(
|
||||||
|
torch.full((H,), init_log, device=dev, dtype=torch.float32)
|
||||||
|
)
|
||||||
|
attn.gravity_eps = float(eps)
|
||||||
|
# config object is shared, but set defensively
|
||||||
|
attn.config._attn_implementation = ATTN_NAME
|
||||||
|
n += 1
|
||||||
|
model.config._attn_implementation = ATTN_NAME
|
||||||
|
print(f"[gravity] patched {n} Qwen2 layers (heads={H}, eps={eps}, init_mass={init_mass})")
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
|
def gravity_mass_state_dict(model):
|
||||||
|
"""Extract only the gravity_mass_log params (for saving separately from base)."""
|
||||||
|
return {f"model.layers.{i}.self_attn.gravity_mass_log":
|
||||||
|
layer.self_attn.gravity_mass_log.detach().cpu()
|
||||||
|
for i, layer in enumerate(model.model.layers)}
|
||||||
3
gravity_mass_log.pt
Normal file
3
gravity_mass_log.pt
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:b37e7594676837dbf0b980970007938a992c059effcd5989c46e98ca1ced4c84
|
||||||
|
size 14513
|
||||||
16
load_gravity2.py
Normal file
16
load_gravity2.py
Normal file
@@ -0,0 +1,16 @@
|
|||||||
|
import torch
|
||||||
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||||
|
from gravity_attention_qwen import patch_qwen_with_gravity
|
||||||
|
|
||||||
|
REPO = "." # or "squ11z1/Gravity-2"
|
||||||
|
tok = AutoTokenizer.from_pretrained(REPO)
|
||||||
|
model = AutoModelForCausalLM.from_pretrained(REPO, dtype=torch.bfloat16,
|
||||||
|
device_map="cuda", attn_implementation="eager")
|
||||||
|
patch_qwen_with_gravity(model) # re-enable gravity attention
|
||||||
|
masses = torch.load(f"{REPO}/gravity_mass_log.pt", map_location="cuda")
|
||||||
|
for i, layer in enumerate(model.model.layers):
|
||||||
|
layer.self_attn.gravity_mass_log.data.copy_(masses[f"model.layers.{i}.self_attn.gravity_mass_log"].cuda())
|
||||||
|
model.eval()
|
||||||
|
ids = tok.apply_chat_template([{"role":"user","content":"What is 24*17?"}],
|
||||||
|
add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].cuda()
|
||||||
|
print(tok.decode(model.generate(ids, max_new_tokens=200)[0, ids.shape[1]:], skip_special_tokens=True))
|
||||||
3
model.safetensors
Normal file
3
model.safetensors
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:e645da04da9ff8321fa0f2f9894b7db9975d416ecd55e70e03af4dde88ccdfac
|
||||||
|
size 6171933008
|
||||||
3
tokenizer.json
Normal file
3
tokenizer.json
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:b17e16899b7fab7e695509f84bac5f10ed12804f0a590e935941e5af7f092f7f
|
||||||
|
size 11422263
|
||||||
16
tokenizer_config.json
Normal file
16
tokenizer_config.json
Normal file
@@ -0,0 +1,16 @@
|
|||||||
|
{
|
||||||
|
"add_prefix_space": false,
|
||||||
|
"backend": "tokenizers",
|
||||||
|
"bos_token": null,
|
||||||
|
"clean_up_tokenization_spaces": false,
|
||||||
|
"eos_token": "<|endoftext|>",
|
||||||
|
"errors": "replace",
|
||||||
|
"is_local": false,
|
||||||
|
"local_files_only": false,
|
||||||
|
"model_max_length": 131072,
|
||||||
|
"pad_token": "<|endoftext|>",
|
||||||
|
"padding_side": "right",
|
||||||
|
"split_special_tokens": false,
|
||||||
|
"tokenizer_class": "Qwen2Tokenizer",
|
||||||
|
"unk_token": null
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user