GODELEV commited on
Commit
205cd74
Β·
verified Β·
1 Parent(s): 2ea9624

Upload 8 files

Browse files
config.json ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "RoseX1ForCausalLM"
4
+ ],
5
+ "model_type": "rose_x1",
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_rose_x1.RoseX1Config",
8
+ "AutoModelForCausalLM": "modeling_rose_x1.RoseX1ForCausalLM"
9
+ },
10
+ "vocab_size": 32768,
11
+ "hidden_size": 512,
12
+ "intermediate_size": 1728,
13
+ "num_hidden_layers": 24,
14
+ "num_attention_heads": 8,
15
+ "num_key_value_heads": 2,
16
+ "head_dim": 64,
17
+ "max_position_embeddings": 2048,
18
+ "hidden_act": "silu",
19
+ "rms_norm_eps": 1e-05,
20
+ "attention_bias": false,
21
+ "mlp_bias": false,
22
+ "attention_dropout": 0.0,
23
+ "tie_word_embeddings": true,
24
+ "rope_theta": 100000.0,
25
+ "rope_scaling": null,
26
+ "use_qk_norm": true,
27
+ "refresh_gate_enabled": true,
28
+ "refresh_gate_inject_layers": [
29
+ 9,
30
+ 18
31
+ ],
32
+ "refresh_gate_kernel_size": 9,
33
+ "bos_token_id": 2,
34
+ "eos_token_id": 3,
35
+ "pad_token_id": 0,
36
+ "torch_dtype": "float32",
37
+ "transformers_version": "4.40.0",
38
+ "use_cache": true,
39
+ "pretraining_tp": 1,
40
+ "initializer_range": 0.02
41
+ }
configuration_rose_x1.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """HuggingFace configuration for the Rose X1 architecture."""
2
+ from transformers import PretrainedConfig
3
+
4
+
5
+ class RoseX1Config(PretrainedConfig):
6
+ model_type = "rose_x1"
7
+
8
+ def __init__(
9
+ self,
10
+ vocab_size=32768,
11
+ hidden_size=512,
12
+ intermediate_size=1720,
13
+ num_hidden_layers=24,
14
+ num_attention_heads=8,
15
+ num_key_value_heads=2,
16
+ head_dim=None,
17
+ max_position_embeddings=2048,
18
+ hidden_act="silu",
19
+ rms_norm_eps=1e-5,
20
+ attention_bias=False,
21
+ mlp_bias=False,
22
+ attention_dropout=0.0,
23
+ tie_word_embeddings=True,
24
+ rope_theta=100000.0,
25
+ rope_scaling=None,
26
+ initializer_range=0.02,
27
+ use_cache=True,
28
+ # ── Rose X1 specifics ──────────────────────────────────────────────
29
+ use_qk_norm=True,
30
+ refresh_gate_enabled=True,
31
+ refresh_gate_inject_layers=None,
32
+ refresh_gate_kernel_size=9,
33
+ **kwargs,
34
+ ):
35
+ self.vocab_size = vocab_size
36
+ self.hidden_size = hidden_size
37
+ self.intermediate_size = intermediate_size
38
+ self.num_hidden_layers = num_hidden_layers
39
+ self.num_attention_heads = num_attention_heads
40
+ self.num_key_value_heads = num_key_value_heads
41
+ self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
42
+ self.max_position_embeddings = max_position_embeddings
43
+ self.hidden_act = hidden_act
44
+ self.rms_norm_eps = rms_norm_eps
45
+ self.attention_bias = attention_bias
46
+ self.mlp_bias = mlp_bias
47
+ self.attention_dropout = attention_dropout
48
+ self.tie_word_embeddings = tie_word_embeddings
49
+ self.rope_theta = rope_theta
50
+ self.rope_scaling = rope_scaling
51
+ self.initializer_range = initializer_range
52
+ self.use_cache = use_cache
53
+ self.use_qk_norm = use_qk_norm
54
+ self.refresh_gate_enabled = refresh_gate_enabled
55
+ self.refresh_gate_inject_layers = (
56
+ list(refresh_gate_inject_layers) if refresh_gate_inject_layers else []
57
+ )
58
+ self.refresh_gate_kernel_size = refresh_gate_kernel_size
59
+ super().__init__(**kwargs)
generation_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "bos_token_id": 2,
3
+ "eos_token_id": 3
4
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:03e67edf9946dcc34761280cbd779a2915ac2d291f18474924ea6898cd0567b2
3
+ size 391311896
modeling_rose_x1.py ADDED
@@ -0,0 +1,433 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Rose X1 model implementation for Hugging Face transformers.
2
+ Architecture (T-X4 family with XSA refresh gate):
3
+ RoPE (half-split) + RMSNorm + SwiGLU + grouped-query attention with
4
+ per-head QK-norm, plus an XSA *refresh gate* that re-injects the original
5
+ token embedding through a gated depthwise-causal-conv path on a subset of
6
+ layers (``config.refresh_gate_inject_layers``).
7
+ Cache design (after the T-X4 reference implementation):
8
+ KV cache uses HF's ``DynamicCache``. The refresh gate's conv history
9
+ (last ``kernel-1`` timesteps of the normalised attention output) is stored
10
+ in a plain dict monkey-patched onto the same ``DynamicCache`` object as
11
+ ``_refresh_conv_state``, so both share one lifetime and no custom Cache
12
+ subclass is needed.
13
+ """
14
+ from typing import Optional
15
+
16
+ import torch
17
+ import torch.nn as nn
18
+ from torch.nn import functional as F
19
+ from transformers import PreTrainedModel
20
+ from transformers.cache_utils import DynamicCache
21
+ from transformers.generation.utils import GenerationMixin
22
+ from transformers.modeling_outputs import CausalLMOutputWithPast
23
+
24
+ try:
25
+ from .configuration_rose_x1 import RoseX1Config
26
+ except ImportError:
27
+ from configuration_rose_x1 import RoseX1Config
28
+
29
+
30
+ # ═══════════════════════════════════════════════════════════════════════════
31
+ # Primitives
32
+ # ═══════════════════════════════════════════════════════════════════════════
33
+
34
+ class RMSNorm(nn.Module):
35
+ """RMSNorm with fp32 internal computation, cast back to input dtype."""
36
+ def __init__(self, dim: int, eps: float = 1e-5):
37
+ super().__init__()
38
+ self.eps = eps
39
+ self.weight = nn.Parameter(torch.ones(dim))
40
+
41
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
42
+ in_dtype = x.dtype
43
+ xf = x.float()
44
+ out = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
45
+ return (out * self.weight.float()).to(in_dtype)
46
+
47
+
48
+ def precompute_rope_cos_sin(head_dim: int, seq_len: int, theta: float = 100000.0):
49
+ """Precompute RoPE cos/sin tables. Returns (cos, sin) each (seq_len, head_dim//2)."""
50
+ freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
51
+ t = torch.arange(seq_len, dtype=torch.float32)
52
+ angles = torch.outer(t, freqs) # (seq_len, head_dim//2)
53
+ return angles.cos(), angles.sin()
54
+
55
+
56
+ def apply_rotary_emb(q: torch.Tensor, k: torch.Tensor,
57
+ cos: torch.Tensor, sin: torch.Tensor):
58
+ """Half-split RoPE (matches the trainer's ``apply_rope``).
59
+ cos / sin: (T, head_dim//2) β€” already sliced to the right positions.
60
+ q, k: (B, H, T, head_dim)
61
+ """
62
+ cos = cos.unsqueeze(0).unsqueeze(0).to(q.dtype) # (1,1,T,d//2)
63
+ sin = sin.unsqueeze(0).unsqueeze(0).to(q.dtype)
64
+ d2 = q.shape[-1] // 2
65
+
66
+ q1, q2 = q[..., :d2], q[..., d2:]
67
+ k1, k2 = k[..., :d2], k[..., d2:]
68
+
69
+ q_out = torch.cat([q1 * cos - q2 * sin, q2 * cos + q1 * sin], dim=-1)
70
+ k_out = torch.cat([k1 * cos - k2 * sin, k2 * cos + k1 * sin], dim=-1)
71
+ return q_out, k_out
72
+
73
+
74
+ # ═══════════════════════════════════════════════════════════════════════════
75
+ # Attention (GQA + QK-norm + RoPE)
76
+ # ═══════════════════════════════════════════════════════════════════════════
77
+
78
+ class RoseX1Attention(nn.Module):
79
+ def __init__(self, config: RoseX1Config, layer_idx: int):
80
+ super().__init__()
81
+ self.layer_idx = layer_idx
82
+ self.n_head = config.num_attention_heads
83
+ self.n_kv_heads = config.num_key_value_heads
84
+ self.head_dim = config.head_dim
85
+ self.n_rep = self.n_head // self.n_kv_heads
86
+
87
+ self.q_proj = nn.Linear(config.hidden_size, self.n_head * self.head_dim, bias=False)
88
+ self.k_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
89
+ self.v_proj = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
90
+ self.o_proj = nn.Linear(self.n_head * self.head_dim, config.hidden_size, bias=False)
91
+
92
+ # QK-norm: per-head RMSNorm on Q & K, applied BEFORE RoPE
93
+ self.use_qk_norm = bool(getattr(config, "use_qk_norm", False))
94
+ if self.use_qk_norm:
95
+ self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
96
+ self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
97
+
98
+ def forward(self, x, rope_cos, rope_sin,
99
+ past_key_value: Optional[DynamicCache] = None,
100
+ use_cache: bool = False,
101
+ attention_mask: Optional[torch.Tensor] = None):
102
+ B, T, _ = x.size()
103
+
104
+ q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
105
+ k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
106
+ v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
107
+
108
+ # ── QK-norm BEFORE RoPE (== trainer) ──────────────────────────────
109
+ if self.use_qk_norm:
110
+ q = self.q_norm(q)
111
+ k = self.k_norm(k)
112
+
113
+ # ── RoPE (half-split, cos/sin already sliced to current positions) ─
114
+ q, k = apply_rotary_emb(q, k, rope_cos, rope_sin)
115
+
116
+ # ── KV cache (DynamicCache.update handles concat internally) ──────
117
+ if past_key_value is not None:
118
+ k, v = past_key_value.update(k, v, self.layer_idx)
119
+
120
+ S = k.size(2)
121
+
122
+ # ── GQA expansion ─────────────────────────────────────────────────
123
+ k = k.unsqueeze(2).expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim) \
124
+ .reshape(B, self.n_head, S, self.head_dim)
125
+ v = v.unsqueeze(2).expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim) \
126
+ .reshape(B, self.n_head, S, self.head_dim)
127
+
128
+ # ── Attention mask ────────────────────────────────────────────────
129
+ # is_causal=True only for prefill (no cache) with T>1 and no padding
130
+ # mask. For decode (T=1) causal is trivially satisfied.
131
+ is_causal = (past_key_value is None
132
+ or past_key_value.get_seq_length(self.layer_idx) == T)
133
+ attn_mask = None
134
+ if attention_mask is not None:
135
+ key_pad = attention_mask.to(torch.bool)[:, None, None, :] # (B,1,1,S)
136
+ if is_causal and T > 1:
137
+ causal = torch.ones(T, S, dtype=torch.bool, device=x.device) \
138
+ .tril(diagonal=S - T)
139
+ attn_mask = key_pad & causal[None, None, :, :]
140
+ else:
141
+ attn_mask = key_pad.expand(B, 1, T, S)
142
+ is_causal = False
143
+
144
+ y = F.scaled_dot_product_attention(q, k, v,
145
+ attn_mask=attn_mask,
146
+ is_causal=is_causal)
147
+ y = y.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_dim)
148
+ return self.o_proj(y)
149
+
150
+
151
+ # ═══════════════════════════════════════════════════════════════════════════
152
+ # MLP (SwiGLU)
153
+ # ═══════════════════════════════════════════════════════════════════════════
154
+
155
+ class RoseX1MLP(nn.Module):
156
+ def __init__(self, config: RoseX1Config):
157
+ super().__init__()
158
+ self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
159
+ self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
160
+ self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
161
+
162
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
163
+ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
164
+
165
+
166
+ # ═══════════════════════════════════════════════════════════════════════════
167
+ # XSA Refresh Gate
168
+ # ═══════════════════════════════════════════════════════════════════════════
169
+
170
+ class RoseX1RefreshGate(nn.Module):
171
+ """Re-injects the original token embedding (e0) into the residual stream,
172
+ gated by a causal depthwise conv over the (detached) attention output.
173
+ Conv history for cached generation is read/written via ``conv_state``
174
+ (a plain dict living on the DynamicCache object).
175
+ """
176
+ def __init__(self, config: RoseX1Config):
177
+ super().__init__()
178
+ H = config.hidden_size
179
+ self.kernel_size = int(getattr(config, "refresh_gate_kernel_size", 9))
180
+
181
+ self.attn_norm = RMSNorm(H, eps=config.rms_norm_eps)
182
+ self.emb_norm = RMSNorm(H, eps=config.rms_norm_eps)
183
+ self.gate_proj = nn.Linear(H, H, bias=False)
184
+ self.value_proj = nn.Linear(H, H, bias=False)
185
+ self.out_proj = nn.Linear(H, H, bias=False)
186
+ self.out_norm = RMSNorm(H, eps=config.rms_norm_eps)
187
+
188
+ # padding attribute documents intent; forward uses F.conv1d(padding=0)
189
+ # with manual left-pad so the same weight works for cached & non-cached.
190
+ self.causal_conv = nn.Conv1d(H, H, self.kernel_size,
191
+ groups=H, bias=False,
192
+ padding=self.kernel_size - 1)
193
+ self.alpha = nn.Parameter(torch.tensor(0.1))
194
+
195
+ def forward(self, h, attn_out, e0, conv_state=None, layer_idx=None):
196
+ a = self.attn_norm(attn_out.detach())
197
+ e = self.emb_norm(e0)
198
+
199
+ k = self.kernel_size
200
+ B, T, D = a.shape
201
+
202
+ if conv_state is not None:
203
+ # ── cached generation: prepend stored history ─────────────────
204
+ prev = conv_state.get(layer_idx)
205
+ if prev is None or prev.size(0) != B:
206
+ prev = a.new_zeros(B, k - 1, D)
207
+ a_ext = torch.cat([prev, a], dim=1) # (B, k-1+T, D)
208
+ conv_state[layer_idx] = a_ext[:, -(k - 1):, :].detach()
209
+ else:
210
+ # ── no cache (training / full recompute): left-pad zeros ──────
211
+ a_ext = F.pad(a, (0, 0, k - 1, 0)) # (B, k-1+T, D)
212
+
213
+ # Manual left-pad + padding=0 conv (== trainer's CausalDepthwiseConv1d)
214
+ c = F.conv1d(a_ext.transpose(1, 2),
215
+ self.causal_conv.weight,
216
+ bias=None, padding=0, groups=D)
217
+ c = c.transpose(1, 2) # (B, T, D)
218
+
219
+ gate = self.gate_proj(a) + c
220
+ value = self.value_proj(e)
221
+ z = self.out_norm(self.out_proj(F.silu(gate) * value))
222
+ return h + self.alpha * z
223
+
224
+
225
+ # ═══════════════════════════════════════════════════════════════════════════
226
+ # Decoder layer
227
+ # ═══════════════════════════════════════════════════════════════════════════
228
+
229
+ class RoseX1DecoderLayer(nn.Module):
230
+ def __init__(self, config: RoseX1Config, layer_idx: int):
231
+ super().__init__()
232
+ self.layer_idx = layer_idx
233
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
234
+ self.self_attn = RoseX1Attention(config, layer_idx)
235
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
236
+ self.mlp = RoseX1MLP(config)
237
+
238
+ inject = list(getattr(config, "refresh_gate_inject_layers", []) or [])
239
+ self.has_refresh = (bool(getattr(config, "refresh_gate_enabled", False))
240
+ and layer_idx in inject)
241
+ if self.has_refresh:
242
+ self.refresh_gate = RoseX1RefreshGate(config)
243
+
244
+ def forward(self, x, e0, rope_cos, rope_sin,
245
+ past_key_value=None, use_cache=False,
246
+ attention_mask=None, conv_state=None):
247
+ attn_out = self.self_attn(self.input_layernorm(x), rope_cos, rope_sin,
248
+ past_key_value, use_cache, attention_mask)
249
+ x = x + attn_out
250
+ # Refresh gate fires AFTER attention residual, BEFORE FFN (== trainer)
251
+ if self.has_refresh:
252
+ x = self.refresh_gate(x, attn_out, e0,
253
+ conv_state=conv_state,
254
+ layer_idx=self.layer_idx)
255
+ x = x + self.mlp(self.post_attention_layernorm(x))
256
+ return x
257
+
258
+
259
+ # ═══════════════════════════════════════════════════════════════════════════
260
+ # Base / backbone / head
261
+ # ═══════════════════════════════════════════════════════════════════════════
262
+
263
+ class RoseX1PreTrainedModel(PreTrainedModel):
264
+ config_class = RoseX1Config
265
+ base_model_prefix = "model"
266
+ supports_gradient_checkpointing = False
267
+ _no_split_modules = ["RoseX1DecoderLayer"]
268
+ _supports_sdpa = True
269
+
270
+ def _init_weights(self, module):
271
+ std = self.config.initializer_range
272
+ if isinstance(module, nn.Linear):
273
+ nn.init.normal_(module.weight, mean=0.0, std=std)
274
+ if module.bias is not None:
275
+ nn.init.zeros_(module.bias)
276
+ elif isinstance(module, nn.Embedding):
277
+ nn.init.normal_(module.weight, mean=0.0, std=std)
278
+ elif isinstance(module, nn.Conv1d):
279
+ nn.init.normal_(module.weight, mean=0.0, std=std)
280
+ elif isinstance(module, RMSNorm):
281
+ nn.init.ones_(module.weight)
282
+
283
+
284
+ class RoseX1Model(nn.Module):
285
+ """Backbone: embed β†’ N Γ— decoder layer β†’ final norm."""
286
+ def __init__(self, config: RoseX1Config):
287
+ super().__init__()
288
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
289
+ self.dropout = nn.Dropout(config.attention_dropout)
290
+ self.layers = nn.ModuleList(
291
+ [RoseX1DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
292
+ )
293
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
294
+
295
+ def forward(self, input_ids, e0, rope_cos, rope_sin,
296
+ past_key_value=None, use_cache=False,
297
+ attention_mask=None, conv_state=None):
298
+ x = self.dropout(self.embed_tokens(input_ids))
299
+ for layer in self.layers:
300
+ x = layer(x, e0, rope_cos, rope_sin,
301
+ past_key_value, use_cache, attention_mask, conv_state)
302
+ return self.norm(x)
303
+
304
+
305
+ class RoseX1ForCausalLM(RoseX1PreTrainedModel, GenerationMixin):
306
+ # Dict format required by modern transformers' get_expanded_tied_weights_keys.
307
+ # Tells HF: "lm_head.weight is tied to model.embed_tokens.weight β€” if it's
308
+ # missing from the checkpoint, fill it from the embedding, don't warn."
309
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
310
+
311
+ def __init__(self, config: RoseX1Config):
312
+ super().__init__(config)
313
+ self.model = RoseX1Model(config)
314
+
315
+ # Always create lm_head; tie it when configured (standard HF pattern).
316
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
317
+ if config.tie_word_embeddings:
318
+ self.lm_head.weight = self.model.embed_tokens.weight
319
+
320
+ self._rope_cache = None # (cos, sin) cached on device
321
+ self.post_init()
322
+
323
+ # ── Embedding accessors (used by tie_weights / resize) ────────────────
324
+ def get_input_embeddings(self):
325
+ return self.model.embed_tokens
326
+
327
+ def set_input_embeddings(self, value):
328
+ self.model.embed_tokens = value
329
+
330
+ def get_output_embeddings(self):
331
+ return self.lm_head
332
+
333
+ def set_output_embeddings(self, new_embeddings):
334
+ self.lm_head = new_embeddings
335
+
336
+ # ── RoPE cache ────────────────────────────────────────────────────────
337
+ def _get_rope(self, seq_len: int, device: torch.device):
338
+ cache = self._rope_cache
339
+ if (cache is None
340
+ or cache[0].device != device
341
+ or cache[0].size(0) < seq_len):
342
+ cos, sin = precompute_rope_cos_sin(
343
+ self.config.head_dim, seq_len, self.config.rope_theta)
344
+ cache = (cos.to(device), sin.to(device))
345
+ self._rope_cache = cache
346
+ return cache[0][:seq_len], cache[1][:seq_len]
347
+
348
+ # ── Generation plumbing ───────────────────────────────────────────────
349
+ def prepare_inputs_for_generation(self, input_ids,
350
+ past_key_values=None,
351
+ attention_mask=None, **kwargs):
352
+ # When a cache with content exists, feed only the newest token.
353
+ if past_key_values is not None and past_key_values.get_seq_length() > 0:
354
+ input_ids = input_ids[:, -1:]
355
+ return {
356
+ "input_ids": input_ids,
357
+ "attention_mask": attention_mask,
358
+ "past_key_values": past_key_values,
359
+ "use_cache": True,
360
+ }
361
+
362
+ # ── Forward ───────────────────────────────────────────────────────────
363
+ def forward(
364
+ self,
365
+ input_ids: torch.Tensor,
366
+ attention_mask: Optional[torch.Tensor] = None,
367
+ labels: Optional[torch.Tensor] = None,
368
+ past_key_values: Optional[DynamicCache] = None,
369
+ use_cache: bool = False,
370
+ **kwargs,
371
+ ) -> CausalLMOutputWithPast:
372
+ B, T = input_ids.size()
373
+
374
+ # ── Conv-state cache for the refresh gate ─────────────────────────
375
+ # Monkey-patched onto the DynamicCache so it shares the cache's
376
+ # lifetime. No custom Cache subclass needed.
377
+ conv_state = None
378
+ if use_cache:
379
+ if past_key_values is None:
380
+ past_key_values = DynamicCache()
381
+ if not hasattr(past_key_values, "_refresh_conv_state"):
382
+ past_key_values._refresh_conv_state = {}
383
+ conv_state = past_key_values._refresh_conv_state
384
+
385
+ # ── Position from cache length (no explicit position_ids needed) ──
386
+ past_len = (past_key_values.get_seq_length()
387
+ if past_key_values is not None else 0)
388
+
389
+ # ── Embeddings ────────────────────────────────────────────────────
390
+ e0 = self.model.embed_tokens(input_ids) # original embedding for refresh gate
391
+
392
+ # ── RoPE: precompute up to past_len+T, slice to current positions ─
393
+ cos, sin = self._get_rope(past_len + T, input_ids.device)
394
+ cos, sin = cos[past_len:], sin[past_len:] # (T, head_dim//2)
395
+
396
+ # ── Backbone ──────────────────────────────────────────────────────
397
+ hidden = self.model(
398
+ input_ids, e0, cos, sin,
399
+ past_key_values if use_cache else None,
400
+ use_cache, attention_mask, conv_state,
401
+ )
402
+
403
+ # ── Head ──────────────────────────────────────────────────────────
404
+ logits = self.lm_head(hidden).float() # fp32 for stable logprobs
405
+
406
+ loss = None
407
+ if labels is not None:
408
+ loss = F.cross_entropy(
409
+ logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size),
410
+ labels[..., 1:].contiguous().view(-1),
411
+ ignore_index=-100,
412
+ )
413
+
414
+ return CausalLMOutputWithPast(
415
+ loss=loss,
416
+ logits=logits,
417
+ past_key_values=past_key_values if use_cache else None,
418
+ )
419
+
420
+
421
+ # ── Optional registration (lets model_type="rose_x1" resolve without auto_map) ──
422
+ try:
423
+ from transformers import AutoConfig, AutoModelForCausalLM
424
+ try:
425
+ AutoConfig.register("rose_x1", RoseX1Config)
426
+ except Exception:
427
+ pass
428
+ try:
429
+ AutoModelForCausalLM.register(RoseX1Config, RoseX1ForCausalLM)
430
+ except Exception:
431
+ pass
432
+ except Exception:
433
+ pass
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "bos_token": "<|bos|>",
4
+ "clean_up_tokenization_spaces": false,
5
+ "eos_token": "<|eos|>",
6
+ "extra_special_tokens": [
7
+ "<|unk_text|>",
8
+ "<|sot|>",
9
+ "<|eot|>",
10
+ "<|system|>",
11
+ "<|user|>",
12
+ "<|assistant|>",
13
+ "<|math|>",
14
+ "<|expr|>",
15
+ "<|answer|>",
16
+ "<|code|>",
17
+ "<|think|>",
18
+ "<|end_of_think|>"
19
+ ],
20
+ "is_local": false,
21
+ "local_files_only": false,
22
+ "mask_token": "<|mask|>",
23
+ "model_max_length": 32768,
24
+ "pad_token": "<|pad|>",
25
+ "tokenizer_class": "TokenizersBackend",
26
+ "unk_token": "<|unk|>"
27
+ }
training_meta.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "step": 38196,
3
+ "val_loss": 1.68247950930655,
4
+ "val_ppl": 5.378876424356594,
5
+ "params_M": 97.82,
6
+ "pushed_at": "2026-08-06T15:28:30.269550",
7
+ "tokens_seen": 80063705088,
8
+ "architecture": "Rose X1",
9
+ "features": [
10
+ "GQA",
11
+ "QK-Norm",
12
+ "RoPE",
13
+ "SwiGLU",
14
+ "RMSNorm",
15
+ "XSA-RefreshGate"
16
+ ],
17
+ "optimizer": "muon_adamw"
18
+ }