# model.py import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass @dataclass class ModelConfig: embedding_dim: int hidden_dim: int num_attention_heads: int layer_count: int max_sequence_length: int rope_theta: float vocab_size: int # ---- 以下 TokenEmbedding / RotaryEmbedding / MHA / FFN / Block / GPT ---- # （あなたが提示したコードをそのまま貼る） # added top-p and top-k filtering in generate function # set vocab_size in config.py # MHA with KV cache + RoPE + PyTorch SDPA. # This traditional implementation is easier to understand, and still efficient in practice. # GQA and MLA is a great way for long-text inference with reduced KV cache size, # but both comes with slight loss increase and no efficiency merits during training phase. # KV cache does not help training speed. Codebase will be simpler without it. # KV cache supports multi-turn continuation by RoPE with position offset. # No Dropout. Dataset is large enough and regularization is not necessary. import torch import torch.nn as nn import torch.nn.functional as F class TokenEmbedding(nn.Module): def __init__(self, config): super().__init__() self.token_embedding_table = nn.Embedding(config.vocab_size, config.embedding_dim) # keep embedding in default dtype (autocast will handle bf16 when enabled) def forward(self, input_indices): return self.token_embedding_table(input_indices) class RotaryEmbedding(nn.Module): def __init__(self, dim, max_seq_len=2048, rope_theta=1e6): super().__init__() inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2) / dim)) position_index = torch.arange(max_seq_len) frequency_matrix = torch.einsum('i,j->ij', position_index, inv_freq) cosine = torch.cos(frequency_matrix)[None, None, :, :] sine = torch.sin(frequency_matrix)[None, None, :, :] self.register_buffer("cos_cached", cosine, persistent=False) self.register_buffer("sin_cached", sine, persistent=False) def apply_rotary_emb(self, x, position_offset=0): sequence_length = x.size(2) cosine = self.cos_cached[:, :, position_offset:position_offset + sequence_length, :] sine = self.sin_cached[:, :, position_offset:position_offset + sequence_length, :] x_even = x[..., 0::2] x_odd = x[..., 1::2] rotated_even = x_even * cosine - x_odd * sine rotated_odd = x_odd * cosine + x_even * sine rotated = torch.empty_like(x) rotated[..., 0::2] = rotated_even rotated[..., 1::2] = rotated_odd return rotated class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() self.config = config self.num_heads = config.num_attention_heads self.embed_dim = config.embedding_dim self.head_dim = self.embed_dim // self.num_heads # QKV projection self.query_fc = nn.Linear(self.embed_dim, self.embed_dim, bias=False) self.key_fc = nn.Linear(self.embed_dim, self.embed_dim, bias=False) self.value_fc = nn.Linear(self.embed_dim, self.embed_dim, bias=False) # Rotary Positional Embedding (RoPE) self.rotary_emb = RotaryEmbedding( dim=self.head_dim, max_seq_len=config.max_sequence_length, rope_theta=config.rope_theta ) self.output_projection = nn.Linear(self.embed_dim, self.embed_dim) self.register_buffer( "causal_mask", torch.tril(torch.ones( config.max_sequence_length, config.max_sequence_length, dtype=torch.bool )), persistent=False ) # KV cache self.register_buffer("cache_k", None, persistent=False) self.register_buffer("cache_v", None, persistent=False) self.current_pos = 0 # -------------------------------------------------- # router # -------------------------------------------------- def forward(self, x, use_cache=False): input_len = x.size(1) if use_cache is False: return self.forward_no_cache(x) elif use_cache is True and input_len > 1: return self.forward_prefill(x) elif use_cache is True and input_len == 1: # Hi scenario also starts with T==1 return self.forward_cached_decoding(x) else: raise RuntimeError("Unexpected condition in MultiHeadAttention forward") # -------------------------------------------------- # (1) no cache : training # -------------------------------------------------- def forward_no_cache(self, x): B, T, C = x.shape Q = self.query_fc(x) K = self.key_fc(x) V = self.value_fc(x) Q = Q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) K = K.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) V = V.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # RoPE : offset = 0 Q = self.rotary_emb.apply_rotary_emb(Q, position_offset=0) K = self.rotary_emb.apply_rotary_emb(K, position_offset=0) out = F.scaled_dot_product_attention( Q, K, V, attn_mask=None, is_causal=True ) out = out.transpose(1, 2).contiguous().view(B, T, C) out = self.output_projection(out) return out # -------------------------------------------------- # (2) prefill : initialize KV cache # -------------------------------------------------- def forward_prefill(self, x): B, T, C = x.shape Q = self.query_fc(x) K = self.key_fc(x) V = self.value_fc(x) Q = Q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) K = K.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) V = V.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # init cache if self.cache_k is None: self.cache_k = torch.zeros( B, self.num_heads, self.config.max_sequence_length, self.head_dim, device=x.device, dtype=K.dtype ) self.cache_v = torch.zeros( B, self.num_heads, self.config.max_sequence_length, self.head_dim, device=x.device, dtype=V.dtype ) self.current_pos = 0 # RoPE : offset = current_pos (supports multi-turn continuation) Q = self.rotary_emb.apply_rotary_emb(Q, position_offset=self.current_pos) K = self.rotary_emb.apply_rotary_emb(K, position_offset=self.current_pos) # prevent overflow if self.current_pos + T > self.config.max_sequence_length: raise RuntimeError("KV cache exceeded max_sequence_length") self.cache_k[:, :, self.current_pos:self.current_pos + T, :] = K self.cache_v[:, :, self.current_pos:self.current_pos + T, :] = V K = self.cache_k[:, :, :self.current_pos + T, :] V = self.cache_v[:, :, :self.current_pos + T, :] attn_mask = self.causal_mask[ self.current_pos : self.current_pos + T, : self.current_pos + T ] out = F.scaled_dot_product_attention( Q, K, V, attn_mask=attn_mask, is_causal=False ) self.current_pos += T out = out.transpose(1, 2).contiguous().view(B, T, C) out = self.output_projection(out) return out # -------------------------------------------------- # (3) decode : cached decoding (1 token) # -------------------------------------------------- def forward_cached_decoding(self, x): B, T, C = x.shape assert T == 1, "cached decoding expects T==1" Q = self.query_fc(x) K = self.key_fc(x) V = self.value_fc(x) Q = Q.view(B, 1, self.num_heads, self.head_dim).transpose(1, 2) K = K.view(B, 1, self.num_heads, self.head_dim).transpose(1, 2) V = V.view(B, 1, self.num_heads, self.head_dim).transpose(1, 2) # This is not usually needed since prefill should have initialized the cache. # Just in case for "Hi" scenario, which starts with single token input. if self.cache_k is None: self.cache_k = torch.zeros( B, self.num_heads, self.config.max_sequence_length, self.head_dim, device=x.device, dtype=K.dtype ) self.cache_v = torch.zeros( B, self.num_heads, self.config.max_sequence_length, self.head_dim, device=x.device, dtype=V.dtype ) self.current_pos = 0 if self.current_pos + 1 >= self.config.max_sequence_length: raise RuntimeError("KV cache exceeded max_sequence_length") # RoPE : offset = current_pos Q = self.rotary_emb.apply_rotary_emb(Q, position_offset=self.current_pos) K = self.rotary_emb.apply_rotary_emb(K, position_offset=self.current_pos) self.cache_k[:, :, self.current_pos:self.current_pos + 1, :] = K self.cache_v[:, :, self.current_pos:self.current_pos + 1, :] = V K = self.cache_k[:, :, :self.current_pos + 1, :] V = self.cache_v[:, :, :self.current_pos + 1, :] out = F.scaled_dot_product_attention( Q, K, V, attn_mask=None, is_causal=False ) self.current_pos += 1 out = out.transpose(1, 2).contiguous().view(B, T, C) out = self.output_projection(out) return out def reset_cache(self): self.cache_k = None self.cache_v = None self.current_pos = 0 class FeedForward(nn.Module): def __init__(self, config): super().__init__() self.net = nn.Sequential( nn.Linear(config.embedding_dim, config.hidden_dim, bias=False), nn.ReLU(), nn.Linear(config.hidden_dim, config.embedding_dim, bias=False), ) def forward(self, input_tensor): return self.net(input_tensor) class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.layer_norm1 = nn.LayerNorm(config.embedding_dim) self.layer_norm2 = nn.LayerNorm(config.embedding_dim) self.multihead_attention = MultiHeadAttention(config=config) self.feed_forward = FeedForward(config=config) def forward(self, input_tensor, use_cache=False): normed_input = self.layer_norm1(input_tensor) attention_output = self.multihead_attention(normed_input, use_cache=use_cache) residual_attention = attention_output + input_tensor normed_attention = self.layer_norm2(residual_attention) feedforward_output = self.feed_forward(normed_attention) final_output = feedforward_output + residual_attention return final_output class VocabularyLogits(nn.Module): def __init__(self, config): super().__init__() self.output_norm = nn.LayerNorm(config.embedding_dim) self.vocab_projection = nn.Linear(config.embedding_dim, config.vocab_size, bias=False) def forward(self, transformer_block_output): x = transformer_block_output normalized_output = self.output_norm(x) vocab_logits = self.vocab_projection(normalized_output) return vocab_logits class GPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.token_embedding_layer = TokenEmbedding(config=config) self.blocks = nn.ModuleList([TransformerBlock(config=config) for _ in range(config.layer_count)]) self.vocab_projection = VocabularyLogits(config=config) self.criterion = nn.CrossEntropyLoss() def forward(self, input_indices, target_indices, use_cache=False): token_embeddings = self.token_embedding_layer.forward(input_indices) x = token_embeddings for block in self.blocks: x = block(x, use_cache=use_cache) logits = self.vocab_projection(x) if target_indices is None: return logits, None batch_size, token_len, vocab_size = logits.shape logits_flat = logits.view(batch_size * token_len, vocab_size) targets_flat = target_indices.view(batch_size * token_len) loss = self.criterion(logits_flat, targets_flat) return logits, loss def generate(self, input_indices, max_new_tokens, temperature=1.0, use_cache=True, reset_cache=True, top_k=None, # ### NEW ### top_p=None, # ### NEW ### ): self.eval() if reset_cache: for block in self.blocks: block.multihead_attention.reset_cache() next_token = None for i in range(max_new_tokens): if use_cache: if i == 0: logits, _ = self.forward(input_indices, None, use_cache=True) else: logits, _ = self.forward(next_token, None, use_cache=True) else: logits, _ = self.forward(input_indices, None, use_cache=False) """ DELETE last_logits = logits[:, -1, :] / temperature probs = F.softmax(last_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) """ ### NEW ### last_logits = logits[:, -1, :] / temperature if top_k is not None: top_k = min(top_k, last_logits.size(-1)) values, _ = torch.topk(last_logits, top_k) min_value = values[:, -1].unsqueeze(-1) last_logits = torch.where( last_logits < min_value, torch.full_like(last_logits, float("-inf")), last_logits, ) if top_p is not None: sorted_logits, sorted_indices = torch.sort(last_logits, descending=True) sorted_probs = F.softmax(sorted_logits, dim=-1) cumulative_probs = torch.cumsum(sorted_probs, dim=-1) sorted_mask = cumulative_probs > top_p sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() sorted_mask[..., 0] = False sorted_logits = torch.where( sorted_mask, torch.full_like(sorted_logits, float("-inf")), sorted_logits, ) last_logits = torch.zeros_like(last_logits).scatter( -1, sorted_indices, sorted_logits ) probs = F.softmax(last_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) ### NEW ### yield int(next_token.item()) input_indices = torch.cat((input_indices, next_token), dim=1)