rdiehlmartinez commited on
Commit
132c5b7
·
1 Parent(s): 81ed588

Saving HF Model -- Step 0

Browse files
Files changed (3) hide show
  1. config.json +22 -0
  2. model.safetensors +3 -0
  3. pico_decoder.py +615 -0
config.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "activation_hidden_dim": 6144,
3
+ "architectures": [
4
+ "PicoHF"
5
+ ],
6
+ "attention_n_heads": 12,
7
+ "attention_n_kv_heads": 4,
8
+ "auto_map": {
9
+ "AutoConfig": "pico.PicoHFConfig",
10
+ "AutoModelForCausalLM": "pico.PicoHF"
11
+ },
12
+ "batch_size": 1024,
13
+ "d_model": 1536,
14
+ "max_seq_len": 2048,
15
+ "model_type": "pico",
16
+ "n_layers": 12,
17
+ "norm_eps": 1e-06,
18
+ "position_emb_theta": 10000.0,
19
+ "torch_dtype": "float32",
20
+ "transformers_version": "4.48.1",
21
+ "vocab_size": 50304
22
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5533b7fc9ddfb2f170cff03059e363880a255b613e1d379c2cca8659e10c512b
3
+ size 2279246680
pico_decoder.py ADDED
@@ -0,0 +1,615 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Pico Decoder: A Lightweight Causal Transformer Language Model
3
+
4
+ Pico Decoder uses a simple LLAMA-style transformer architecture, written for clarity and educational purposes.
5
+
6
+ Everything is written with a modular design for easy modification and experimentation.
7
+
8
+ Key features:
9
+ - RMSNorm for layer normalization
10
+ - Rotary Positional Embeddings (RoPE)
11
+ - Multi-head attention with KV-cache support
12
+ - SwiGLU activation function
13
+ - Residual connections throughout
14
+
15
+ - KV-cache for faster autoregressive generation
16
+
17
+ References:
18
+ - RoPE: https://arxiv.org/abs/2104.09864
19
+ - SwiGLU: https://arxiv.org/abs/2002.05202
20
+ - LLAMA: https://arxiv.org/abs/2302.13971
21
+
22
+ Adapted from:
23
+ - OLMO: https://github.com/allenai/OLMo
24
+ - LLAMA: https://github.com/meta/llama
25
+ """
26
+
27
+ import torch
28
+ import torch.nn as nn
29
+ import torch.nn.functional as F
30
+ from torch.nn.attention import sdpa_kernel, SDPBackend
31
+
32
+ from dataclasses import asdict
33
+
34
+ from transformers import PretrainedConfig, PreTrainedModel
35
+ from transformers.modeling_outputs import CausalLMOutputWithPast, CausalLMOutput
36
+
37
+ # typing imports
38
+ from typing import Union, Tuple, Optional, TYPE_CHECKING, Dict, Any
39
+
40
+ try:
41
+ if TYPE_CHECKING:
42
+ # We need to do this to avoid importing these when creating the HF-compatible models
43
+ from src.config import ModelConfig
44
+ except ImportError:
45
+ pass
46
+
47
+ ########################################################
48
+ #
49
+ # Layer Normalization
50
+ #
51
+ ########################################################
52
+
53
+
54
+ class RMSNorm(torch.nn.Module):
55
+ """Root Mean Square Layer Normalization.
56
+
57
+ A variant of Layer Normalization that uses RMS statistics instead of mean/variance,
58
+ resulting in improved stability and performance.
59
+
60
+ Args:
61
+ config (Union[ModelConfig, PicoHFConfig]): Configuration object containing normalization parameters
62
+ - config.norm_eps: Small constant for numerical stability
63
+ - config.d_model: Model dimension for the weight parameter
64
+
65
+ References:
66
+ https://arxiv.org/abs/1910.07467
67
+ """
68
+
69
+ def __init__(self, config: Union["ModelConfig", "PicoDecoderHFConfig"]):
70
+ super().__init__()
71
+ self.eps = config.norm_eps
72
+ self.weight = nn.Parameter(torch.ones(config.d_model))
73
+
74
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
75
+ """
76
+ Normalizes the input tensor by its RMS value.
77
+ """
78
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
79
+
80
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
81
+ """
82
+ Applies RMS normalization to the input tensor and scales it by the weight parameter.
83
+ """
84
+ output = self._norm(x.float()).type_as(x)
85
+ return output * self.weight
86
+
87
+
88
+ ########################################################
89
+ #
90
+ # Positional Embedding
91
+ #
92
+ ########################################################
93
+
94
+
95
+ class RoPE(nn.Module):
96
+ """Rotary Positional Embeddings (RoPE).
97
+
98
+ Implements position-dependent rotation of keys and queries in attention mechanism,
99
+ allowing better modeling of relative positions in sequences. Uses complex number
100
+ operations for efficient rotation.
101
+
102
+ Args:
103
+ config (Union[ModelConfig, PicoHFConfig]): Model configuration containing:
104
+ - config.position_emb_theta: Base for frequency computation
105
+ - config.d_model: Model dimension
106
+ - config.attention_n_heads: Number of attention heads
107
+ - config.max_seq_len: Maximum sequence length
108
+
109
+ References:
110
+ https://arxiv.org/abs/2104.09864
111
+ """
112
+
113
+ _freqs_cis_tensor: torch.Tensor | None = None
114
+
115
+ def __init__(self, config: Union["ModelConfig", "PicoDecoderHFConfig"]):
116
+ super().__init__()
117
+
118
+ self.theta = config.position_emb_theta
119
+ self.dim = config.d_model // config.attention_n_heads
120
+
121
+ max_seq_len = config.max_seq_len
122
+
123
+ # only gets set once, and then reused for all RoPE instances
124
+ if RoPE._freqs_cis_tensor is None:
125
+ RoPE._freqs_cis_tensor = self._setup_freqs_cis(
126
+ max_seq_len, self.theta, self.dim
127
+ )
128
+
129
+ # register _freqs_cis buffer
130
+ # can be easily recomputed so persistent=False
131
+ self.register_buffer("_freqs_cis", self._freqs_cis_tensor, persistent=False)
132
+
133
+ @classmethod
134
+ def _setup_freqs_cis(cls, seq_len: int, theta: float, dim: int) -> torch.Tensor:
135
+ """Setup Frequency Tensor for RoPE Embeddings
136
+
137
+ Initializes the complex frequency tensor that is used to compute the RoPE embeddings.
138
+
139
+ Note other implementations will use cos and sin directly, but using the complex
140
+ number representation is (probably?) more efficient:
141
+
142
+ e^(theta * i * t) = cos(theta * t) + i * sin(theta * t) [Euler's formula]
143
+ """
144
+ _freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
145
+ positions = torch.arange(seq_len)
146
+ freqs = torch.outer(positions, _freqs)
147
+ return torch.polar(torch.ones_like(freqs), freqs) # complex64
148
+
149
+ def get_freqs_cis(
150
+ self, input_shape: torch.Size, start_pos: int, end_pos: int
151
+ ) -> torch.Tensor:
152
+ """Reshape Frequency Tensor for RoPE Embeddings
153
+
154
+ Makes the frequency tensor broadcastable with the input tensor.
155
+ """
156
+ _freqs_cis = self._freqs_cis[start_pos:end_pos]
157
+ ndim = len(input_shape)
158
+ assert 0 <= 1 < ndim
159
+ assert _freqs_cis.shape == (input_shape[1], input_shape[-1])
160
+
161
+ # TODO: Check whether this is correct (might be able to remove this)
162
+ shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(input_shape)]
163
+ return _freqs_cis.view(*shape)
164
+
165
+ def forward(
166
+ self,
167
+ queries: torch.Tensor,
168
+ keys: torch.Tensor,
169
+ start_pos: int = 0,
170
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
171
+ """Apply RoPE Embeddings to Queries and Keys
172
+
173
+ Applies the rotary positional embeddings to the input tensors via complex num multiplication
174
+
175
+ NOTE: The start_pos is used if we want to use the kv_cache in the attention mechanism.
176
+ """
177
+ queries_ = torch.view_as_complex(
178
+ queries.float().reshape(*queries.shape[:-1], -1, 2)
179
+ )
180
+ keys_ = torch.view_as_complex(keys.float().reshape(*keys.shape[:-1], -1, 2))
181
+
182
+ input_shape = (
183
+ queries_.shape
184
+ ) # same as keys: (batch_size, seq_len, n_heads, head_dim/2)
185
+ freqs_start_pos = start_pos
186
+ freqs_end_pos = freqs_start_pos + queries_.shape[1]
187
+
188
+ freqs_cis = self.get_freqs_cis(input_shape, freqs_start_pos, freqs_end_pos)
189
+
190
+ queries_rotated = torch.view_as_real(queries_ * freqs_cis).flatten(3)
191
+ keys_rotated = torch.view_as_real(keys_ * freqs_cis).flatten(3)
192
+ return queries_rotated.type_as(queries), keys_rotated.type_as(keys)
193
+
194
+
195
+ ########################################################
196
+ #
197
+ # Attention
198
+ #
199
+ ########################################################
200
+
201
+
202
+ class Attention(nn.Module):
203
+ """Multi-head Attention with Group Query Attention support.
204
+
205
+ Implements scaled dot-product attention and supports:
206
+ - Grouped Query Attention (GQA)
207
+ - Key-Value caching for efficient inference
208
+ - RoPE integration
209
+
210
+ Args:
211
+ config (Union[ModelConfig, PretrainedConfig]): Configuration containing:
212
+ - config.attention_n_heads: Number of attention heads
213
+ - config.attention_n_kv_heads: Number of key/value heads
214
+ - config.d_model: Model dimension
215
+ - config.batch_size: Maximum batch size
216
+ - config.max_seq_len: Maximum sequence length
217
+
218
+ Shape:
219
+ - Input: (batch_size, seq_len, d_model)
220
+ - Output: (batch_size, seq_len, d_model)
221
+ """
222
+
223
+ def __init__(
224
+ self,
225
+ config: Union["ModelConfig", "PicoDecoderHFConfig"],
226
+ ):
227
+ super().__init__()
228
+
229
+ self.n_heads = config.attention_n_heads
230
+ self.n_kv_heads = config.attention_n_kv_heads
231
+
232
+ self.batch_size = config.batch_size
233
+ self.max_seq_len = config.max_seq_len
234
+
235
+ d_model = config.d_model
236
+ self.head_dim = d_model // self.n_heads
237
+
238
+ self.n_rep = self.n_heads // self.n_kv_heads
239
+
240
+ self.q_proj = nn.Linear(d_model, self.n_heads * self.head_dim, bias=False)
241
+ self.k_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=False)
242
+ self.v_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=False)
243
+ self.o_proj = nn.Linear(self.n_heads * self.head_dim, d_model, bias=False)
244
+
245
+ self.rope = RoPE(config)
246
+
247
+ def forward(
248
+ self,
249
+ input: torch.Tensor,
250
+ mask: Optional[torch.Tensor] = None,
251
+ past_key_values: Optional[Tuple[torch.Tensor, ...]] = None,
252
+ use_cache: bool = False,
253
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
254
+ """Forward pass for the attention mechanism.
255
+
256
+ Computes queries, keys, and values for the attention mechanism. Applies rotary positional
257
+ embeddings to the queries and keys, and then computes attention scores and outputs.
258
+
259
+ For an introduction to the attention mechanism, see:
260
+ https://arxiv.org/abs/1706.03762
261
+
262
+ A few things to note:
263
+ - The past_key_values is used to implement the KV cache, which is used to speed up
264
+ generation by caching the KV pairs from previous forward passes. This is useful when doing
265
+ tasks that require generating multiple tokens conditioned on previous tokens (e.g. language
266
+ modeling, text generation, etc.). The way the KV cache is implemented is that each layer has
267
+ its own KV cache - this KV cache is implemented as a tuple.
268
+ """
269
+ bsz, seq_len, _ = input.shape
270
+ _queries, _keys, _values = (
271
+ self.q_proj(input),
272
+ self.k_proj(input),
273
+ self.v_proj(input),
274
+ )
275
+
276
+ # Reshaping for multi-head attention
277
+ queries = _queries.view(bsz, seq_len, self.n_heads, self.head_dim)
278
+ keys = _keys.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
279
+ values = _values.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
280
+
281
+ # The start position is used to apply the RoPE embeddings to only the new tokens
282
+ # when using the kv_cache in the attention mechanism.
283
+ # We want to start from the last position in the cache.
284
+ start_pos = past_key_values[0].shape[1] if past_key_values is not None else 0
285
+
286
+ # apply rotary positional embeddings
287
+ queries, keys = self.rope(queries, keys, start_pos)
288
+
289
+ if past_key_values is not None:
290
+ keys = torch.cat([past_key_values[0], keys], dim=1)
291
+ values = torch.cat([past_key_values[1], values], dim=1)
292
+
293
+ if use_cache:
294
+ cached_keys = keys
295
+ cached_values = values
296
+ else:
297
+ cached_keys = None
298
+ cached_values = None
299
+
300
+ queries = queries.transpose(1, 2)
301
+ keys = keys.transpose(1, 2)
302
+ values = values.transpose(1, 2)
303
+
304
+ apply_gqa = self.n_rep > 1
305
+ if apply_gqa and queries.device.type == "mps":
306
+ # NOTE: MPS does not support GQA in the SDPA kernel, but we can repeat the keys and values
307
+ # outside of the kernel to get the same effect.
308
+ # See: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
309
+ keys = keys.repeat_interleave(self.n_rep, dim=-3)
310
+ values = values.repeat_interleave(self.n_rep, dim=-3)
311
+ apply_gqa = False
312
+
313
+ backends = [SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]
314
+
315
+ with sdpa_kernel(backends=backends):
316
+ attn_output = F.scaled_dot_product_attention(
317
+ queries.contiguous(),
318
+ keys.contiguous(),
319
+ values.contiguous(),
320
+ attn_mask=mask.to(queries.dtype),
321
+ enable_gqa=apply_gqa,
322
+ )
323
+
324
+ attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, seq_len, -1)
325
+ output = self.o_proj(attn_output)
326
+
327
+ return output, (cached_keys, cached_values)
328
+
329
+
330
+ ########################################################
331
+ #
332
+ # SwiGLU (Combines MLP and Activation)
333
+ #
334
+ ########################################################
335
+
336
+
337
+ class SwiGLU(nn.Module):
338
+ """SwiGLU Activation Function with Linear Projections.
339
+
340
+ Implements the SwiGLU activation function combined with linear transformations,
341
+ serving as the feed-forward network in transformer blocks.
342
+
343
+ Args:
344
+ config (Union[ModelConfig, PicoDecoderHFConfig]): Configuration containing:
345
+ - config.d_model: Model dimension
346
+ - config.activation_hidden_dim: Hidden dimension (typically 4 * d_model)
347
+
348
+ References:
349
+ https://arxiv.org/abs/2002.05202
350
+ """
351
+
352
+ def __init__(self, config: Union["ModelConfig", "PicoDecoderHFConfig"]):
353
+ super().__init__()
354
+
355
+ model_dim = config.d_model
356
+ act_hidden_dim = config.activation_hidden_dim # usually 4 * d_model
357
+
358
+ self.w_0 = nn.Linear(model_dim, act_hidden_dim, bias=False)
359
+ self.w_1 = nn.Linear(model_dim, act_hidden_dim, bias=False)
360
+ self.w_2 = nn.Linear(act_hidden_dim, model_dim, bias=False)
361
+
362
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
363
+ return self.w_2(F.silu(self.w_0(x)) * self.w_1(x))
364
+
365
+
366
+ ########################################################
367
+ #
368
+ # PicoDecoderBlock
369
+ #
370
+ ########################################################
371
+
372
+
373
+ class PicoDecoderBlock(nn.Module):
374
+ """Single Transformer Block with Attention and Feed-forward layers.
375
+
376
+ Implements a standard transformer block with:
377
+ - Multi-head attention with normalization and residual connection
378
+ - SwiGLU feed-forward network with normalization and residual connection
379
+
380
+ Args:
381
+ config (Union[ModelConfig, PicoDecoderHFConfig]): Model configuration; either a dataclass or
382
+ a HuggingFace PicoDecoderHFConfig
383
+ """
384
+
385
+ def __init__(
386
+ self,
387
+ config: Union["ModelConfig", "PicoDecoderHFConfig"],
388
+ ):
389
+ super().__init__()
390
+
391
+ self.attention = Attention(config)
392
+ self.swiglu = SwiGLU(config)
393
+ self.attention_norm = RMSNorm(config)
394
+ self.swiglu_norm = RMSNorm(config)
395
+
396
+ def forward(
397
+ self,
398
+ input: torch.Tensor,
399
+ mask: Optional[torch.Tensor] = None,
400
+ past_key_values: Optional[Tuple[torch.Tensor]] = None,
401
+ use_cache: bool = False,
402
+ ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
403
+ attention_output, cached_key_values = self.attention(
404
+ self.attention_norm(input),
405
+ mask=mask,
406
+ past_key_values=past_key_values,
407
+ use_cache=use_cache,
408
+ )
409
+ # NOTE: cached_key_values is None if use_cache is False
410
+
411
+ h = input + attention_output
412
+ out = h + self.swiglu(self.swiglu_norm(h))
413
+ return out, cached_key_values
414
+
415
+
416
+ ########################################################
417
+ #
418
+ # Pico Decoder (Causal Transformer Model)
419
+ #
420
+ ########################################################
421
+
422
+
423
+ class PicoDecoder(nn.Module):
424
+ """
425
+ Pico Decoder: combines the embedding, causal decoder blocks, and output projection into a
426
+ single autoregressive model.
427
+
428
+ For more information on the model, see the classes for the modules that make up the model.
429
+ """
430
+
431
+ def __init__(
432
+ self,
433
+ model_config: Union["ModelConfig", "PicoDecoderHFConfig"],
434
+ ):
435
+ super().__init__()
436
+ self.config = model_config
437
+
438
+ self.embedding_proj = nn.Embedding(self.config.vocab_size, self.config.d_model)
439
+ self.layers = nn.ModuleList(
440
+ [PicoDecoderBlock(self.config) for _ in range(self.config.n_layers)]
441
+ )
442
+ self.output_norm = RMSNorm(self.config)
443
+ self.de_embedding_proj = nn.Linear(
444
+ self.config.d_model, self.config.vocab_size, bias=False
445
+ )
446
+
447
+ def convert_to_hf_model(self) -> "PicoDecoderHF":
448
+ """Convert the Lightning model to a HuggingFace model."""
449
+ # Create HF config without fabric-specific settings
450
+ hf_config = PicoDecoderHFConfig.from_dataclass(self.config)
451
+
452
+ # Create new HF model
453
+ hf_model = PicoDecoderHF(hf_config)
454
+
455
+ # Copy state dict, excluding fabric-specific keys
456
+ hf_model.load_state_dict(self.state_dict(prefix="pico_decoder."))
457
+
458
+ return hf_model
459
+
460
+ def forward(
461
+ self,
462
+ input_ids: torch.Tensor,
463
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
464
+ use_cache: bool = False,
465
+ ) -> Tuple[torch.Tensor, Optional[Tuple[Tuple[torch.Tensor, torch.Tensor]]]]:
466
+ """
467
+ This is the forward pass for the entire Pico model. It boils down to:
468
+ - Embedding the input ids
469
+ - Creating a causal mask
470
+ - Processing through the pico layers
471
+ - Projecting the output to logits
472
+
473
+ NOTE: One feature that might be confusing is the KV cache. The KV cache is used to speed up
474
+ generation by caching the KV pairs from previous forward passes. This is useful when doing
475
+ tasks that require generating multiple tokens conditioned on previous tokens (e.g. language
476
+ modeling, text generation, etc.). The way the KV cache is implemented is that each layer has
477
+ its own KV cache which is stored as a tuple. The whole model then stores a tuple of these
478
+ KV caches (so a tuple of tuples).
479
+ """
480
+
481
+ seq_len = input_ids.shape[-1]
482
+ h = self.embedding_proj(input_ids)
483
+
484
+ # Calculate start position from past cached KV pairs. Remember that each layer has its
485
+ # own KV Cache. So when we index past_key_values, we need to index into the KV pairs for the
486
+ # correct layer and then for either the keys or values.
487
+ start_pos = 0 if past_key_values is None else past_key_values[0][0].shape[1]
488
+
489
+ # Create causal mask for current sequence
490
+ mask = None
491
+ if seq_len > 1:
492
+ mask = torch.full((seq_len, seq_len), float("-inf"))
493
+ mask = torch.triu(mask, diagonal=1)
494
+
495
+ # If using KV cache, extend mask to cover cached sequence length
496
+ if past_key_values is not None:
497
+ # Add zeros for cached tokens (we can attend to all of them)
498
+ mask = torch.hstack([torch.zeros((seq_len, start_pos)), mask])
499
+
500
+ mask = mask.to(h.device)
501
+
502
+ # NOTE: If we are using the cache, we need to store the cached KV pairs for each layer
503
+ # in a tuple. Each layer will have its own cached KV pair which we aggregate in a tuple.
504
+ cached_key_values = () if use_cache else None
505
+
506
+ # Process through transformer blocks
507
+ for idx, layer in enumerate(self.layers):
508
+ layer_past_key_values = (
509
+ past_key_values[idx] if past_key_values is not None else None
510
+ )
511
+
512
+ h, layer_cached_key_values = layer(
513
+ h, mask=mask, past_key_values=layer_past_key_values, use_cache=use_cache
514
+ )
515
+
516
+ if use_cache:
517
+ cached_key_values += (layer_cached_key_values,)
518
+
519
+ # Final norm and projection
520
+ h = self.output_norm(h)
521
+ logits = self.de_embedding_proj(h).float()
522
+
523
+ return logits, cached_key_values
524
+
525
+
526
+ ########################################################
527
+ #
528
+ # HuggingFace Wrapper
529
+ #
530
+ ########################################################
531
+
532
+ """
533
+ HuggingFace wrapper for the Pico model.
534
+
535
+ Many evaluation frameworks require a model be setup as a HuggingFace model, so we provide a simple
536
+ wrapper that does just that. When we save checkpoints of the Pico model, we save both the normal
537
+ Pico model as well as the model wrapped in this HuggingFace class.
538
+
539
+ This also lets you do cool things like:
540
+
541
+ `model = AutoModelForCausalLM.from_pretrained("path/to/checkpoint")`
542
+ """
543
+
544
+
545
+ class PicoDecoderHFConfig(PretrainedConfig):
546
+ """HuggingFace config for Pico model."""
547
+
548
+ model_type = "pico_decoder"
549
+
550
+ @classmethod
551
+ def from_dict(cls, config_dict: Dict[str, Any], **kwargs) -> "PicoDecoderHFConfig":
552
+ # NOTE The typical from_dict method doesn't actually set the attributes unless they are
553
+ # defined in the constructor.
554
+
555
+ pico_config = cls(**kwargs)
556
+
557
+ # Because this class is just a wrapper around the ModelConfig dataclass, we need to do
558
+ # a little extra work to ensure that the attributes are actually set.
559
+ for key, value in config_dict.items():
560
+ setattr(pico_config, key, value)
561
+
562
+ return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
563
+ unused_kwargs = {
564
+ key: value for key, value in kwargs.items() if not hasattr(pico_config, key)
565
+ }
566
+
567
+ if return_unused_kwargs:
568
+ return pico_config, unused_kwargs
569
+ return pico_config
570
+
571
+ @classmethod
572
+ def from_dataclass(cls, model_config: "ModelConfig"):
573
+ return cls.from_dict(asdict(model_config))
574
+
575
+
576
+ class PicoDecoderHF(PreTrainedModel):
577
+ """HuggingFace wrapper for Pico model."""
578
+
579
+ config_class = PicoDecoderHFConfig
580
+ _no_split_modules = ["PicoBlock", "Attention", "SwiGLU", "RMSNorm"]
581
+
582
+ def __init__(self, config: PicoDecoderHFConfig):
583
+ super().__init__(config)
584
+ self.pico_decoder = PicoDecoder(config)
585
+
586
+ def forward(
587
+ self,
588
+ input_ids: torch.Tensor,
589
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
590
+ use_cache: bool = False,
591
+ **kwargs,
592
+ ) -> Union[CausalLMOutput, CausalLMOutputWithPast]:
593
+ """HuggingFace forward pass wrapper.
594
+
595
+ Forwards pass for the HuggingFace version of the Pico Model. Basic wrapper around the
596
+ Pico model's forward pass, and returns the output as a HuggingFace CausalLMOutput.
597
+ """
598
+ logits, past_key_values = self.pico_decoder(
599
+ input_ids, past_key_values, use_cache
600
+ )
601
+ if use_cache:
602
+ return CausalLMOutputWithPast(
603
+ logits=logits,
604
+ past_key_values=past_key_values,
605
+ )
606
+ else:
607
+ return CausalLMOutput(
608
+ logits=logits,
609
+ )
610
+
611
+
612
+ # Register for auto classes
613
+ PicoDecoderHFConfig.register_for_auto_class()
614
+ PicoDecoderHF.register_for_auto_class("AutoModel")
615
+ PicoDecoderHF.register_for_auto_class("AutoModelForCausalLM")