| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers import PreTrainedModel, GenerationMixin |
| from transformers.modeling_outputs import MoeCausalLMOutputWithPast |
|
|
| from core import RMSNorm, precompute_freqs_cis, Block, MOEFeedForward |
| from models.lm.config import LMConfig |
|
|
|
|
| class LM(nn.Module): |
| def __init__(self, config: LMConfig): |
| super().__init__() |
| self.config = config |
| self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.dropout = nn.Dropout(config.dropout) |
| self.layers = nn.ModuleList([Block(l, config) for l in range(self.num_hidden_layers)]) |
| self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling) |
| self.register_buffer("freqs_cos", freqs_cos, persistent=False) |
| self.register_buffer("freqs_sin", freqs_sin, persistent=False) |
|
|
| def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs): |
| batch_size, seq_length = input_ids.shape |
| if hasattr(past_key_values, 'layers'): |
| past_key_values = None |
| past_key_values = past_key_values or [None] * len(self.layers) |
| start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0 |
| hidden_states = self.dropout(self.embed_tokens(input_ids)) |
| if self.freqs_cos[0, 0] != 1.0: |
| freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling) |
| self.freqs_cos, self.freqs_sin = freqs_cos.to(device=hidden_states.device, dtype=hidden_states.dtype), freqs_sin.to(device=hidden_states.device, dtype=hidden_states.dtype) |
| position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length]) |
| presents = [] |
| for layer, past_key_value in zip(self.layers, past_key_values): |
| hidden_states, present = layer( |
| hidden_states, |
| position_embeddings, |
| past_key_value=past_key_value, |
| use_cache=use_cache, |
| attention_mask=attention_mask |
| ) |
| presents.append(present) |
| hidden_states = self.norm(hidden_states) |
| aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze()) |
| return hidden_states, presents, aux_loss |
|
|
|
|
| class LMForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = LMConfig |
| model_type = "omni" |
| _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} |
|
|
| def __init__(self, config: LMConfig = None): |
| self.config = config or LMConfig() |
| super().__init__(self.config) |
| self.model = LM(self.config) |
| self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) |
| if self.config.tie_word_embeddings: |
| self.model.embed_tokens.weight = self.lm_head.weight |
| self.post_init() |
|
|
| def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs): |
| hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs) |
| slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep |
| logits = self.lm_head(hidden_states[:, slice_indices, :]) |
| loss = None |
| if labels is not None: |
| x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous() |
| loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100) |
| return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states) |
|
|
| @torch.inference_mode() |
| def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs): |
| input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1) |
| attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None |
| past_key_values = kwargs.pop("past_key_values", None) |
| finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device) |
| if streamer: |
| streamer.put(input_ids.cpu()) |
| for _ in range(max_new_tokens): |
| past_len = past_key_values[0][0].shape[1] if past_key_values else 0 |
| outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs) |
| attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None |
| logits = outputs.logits[:, -1, :] / temperature |
| if repetition_penalty != 1.0: |
| for i in range(input_ids.shape[0]): |
| seen = torch.unique(input_ids[i]) |
| score = logits[i, seen] |
| logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty) |
| if top_k > 0: |
| logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf') |
| if top_p < 1.0: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p |
| mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0 |
| logits[mask.scatter(1, sorted_indices, mask)] = -float('inf') |
| next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True) |
| if eos_token_id is not None: |
| next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token) |
| input_ids = torch.cat([input_ids, next_token], dim=-1) |
| past_key_values = outputs.past_key_values if use_cache else None |
| if streamer: |
| streamer.put(next_token.cpu()) |
| if eos_token_id is not None: |
| finished |= next_token.squeeze(-1).eq(eos_token_id) |
| if finished.all(): |
| break |
| if streamer: |
| streamer.end() |
| if kwargs.get("return_kv"): |
| return {'generated_ids': input_ids, 'past_kv': past_key_values} |
| return input_ids |
|
|
