| """ |
| modeling_smartcoder_moe.py |
| Custom model class for SmartCoderMoE. |
| |
| Architecture (from tensor inspection): |
| - vocab_size: 65536, hidden: 2048, layers: 40 |
| - Attention: q[2048,2048], k/v[512,2048] — 16 heads, 4 KV heads, head_dim=128 |
| - MLP (hybrid dense + MoE): |
| dense_fc: [8192, 2048] up |
| dense_proj: [2048, 8192] down |
| experts_fc: [32, 512, 2048] expert up (batched) |
| experts_proj: [32, 2048, 512] expert down (batched) |
| router: [32, 2048] router logits |
| - LayerNorm: weight+bias (input_layernorm, post_attention_layernorm) |
| - Final norm: model.norm.weight/bias |
| """ |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel, PretrainedConfig |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| from typing import Optional, Tuple, List |
|
|
|
|
| |
| class SmartCoderMoEConfig(PretrainedConfig): |
| model_type = "smartcoder_moe" |
|
|
| def __init__( |
| self, |
| vocab_size=65536, |
| hidden_size=2048, |
| num_hidden_layers=40, |
| num_attention_heads=16, |
| num_key_value_heads=4, |
| dense_intermediate_size=8192, |
| num_experts=32, |
| expert_intermediate_size=512, |
| num_experts_per_tok=2, |
| max_position_embeddings=16384, |
| rope_theta=10000.0, |
| rms_norm_eps=1e-5, |
| pad_token_id=0, |
| bos_token_id=1, |
| eos_token_id=0, |
| tie_word_embeddings=False, |
| **kwargs, |
| ): |
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_size |
| self.num_hidden_layers = num_hidden_layers |
| self.num_attention_heads = num_attention_heads |
| self.num_key_value_heads = num_key_value_heads |
| self.head_dim = hidden_size // num_attention_heads |
| self.dense_intermediate_size = dense_intermediate_size |
| self.num_experts = num_experts |
| self.expert_intermediate_size = expert_intermediate_size |
| self.num_experts_per_tok = num_experts_per_tok |
| self.max_position_embeddings = max_position_embeddings |
| self.rope_theta = rope_theta |
| self.rms_norm_eps = rms_norm_eps |
| super().__init__( |
| pad_token_id=pad_token_id, |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| tie_word_embeddings=tie_word_embeddings, |
| **kwargs, |
| ) |
|
|
|
|
| |
| def rotate_half(x): |
| x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:] |
| return torch.cat([-x2, x1], dim=-1) |
|
|
| def apply_rotary_emb(q, k, cos, sin): |
| return (q * cos) + (rotate_half(q) * sin), \ |
| (k * cos) + (rotate_half(k) * sin) |
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim, max_pos=16384, base=10000.0): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq) |
| self.max_pos = max_pos |
| self._build_cache(max_pos) |
|
|
| def _build_cache(self, seq_len): |
| t = torch.arange(seq_len, device=self.inv_freq.device).float() |
| freqs = torch.outer(t, self.inv_freq) |
| emb = torch.cat([freqs, freqs], dim=-1) |
| self.register_buffer("cos_cached", emb.cos()[None, None, :, :]) |
| self.register_buffer("sin_cached", emb.sin()[None, None, :, :]) |
|
|
| def forward(self, seq_len): |
| if seq_len > self.max_pos: |
| self._build_cache(seq_len) |
| return self.cos_cached[:, :, :seq_len, :], \ |
| self.sin_cached[:, :, :seq_len, :] |
|
|
|
|
| |
| class LayerNormWithBias(nn.Module): |
| def __init__(self, hidden_size, eps=1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.bias = nn.Parameter(torch.zeros(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return F.layer_norm(x, x.shape[-1:], self.weight, self.bias, self.eps) |
|
|
|
|
| |
| class SmartCoderAttention(nn.Module): |
| def __init__(self, config: SmartCoderMoEConfig): |
| super().__init__() |
| self.hidden_size = config.hidden_size |
| self.num_heads = config.num_attention_heads |
| self.num_kv_heads = config.num_key_value_heads |
| self.head_dim = config.head_dim |
| self.num_kv_groups = self.num_heads // self.num_kv_heads |
|
|
| self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * config.head_dim, bias=True) |
| self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=True) |
| self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * config.head_dim, bias=True) |
| self.o_proj = nn.Linear(config.num_attention_heads * config.head_dim, config.hidden_size, bias=True) |
|
|
| self.rotary_emb = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta) |
|
|
| def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False): |
| B, T, _ = hidden_states.shape |
|
|
| q = self.q_proj(hidden_states).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(hidden_states).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) |
|
|
| cos, sin = self.rotary_emb(T) |
| cos = cos[:, :, :T, :self.head_dim] |
| sin = sin[:, :, :T, :self.head_dim] |
| q, k = apply_rotary_emb(q, k, cos, sin) |
|
|
| if past_key_value is not None: |
| k = torch.cat([past_key_value[0], k], dim=2) |
| v = torch.cat([past_key_value[1], v], dim=2) |
| present = (k, v) if use_cache else None |
|
|
| |
| k = k.repeat_interleave(self.num_kv_groups, dim=1) |
| v = v.repeat_interleave(self.num_kv_groups, dim=1) |
|
|
| scale = math.sqrt(self.head_dim) |
| attn = torch.matmul(q, k.transpose(-2, -1)) / scale |
|
|
| kv_len = k.shape[2] |
| causal_mask = torch.triu( |
| torch.full((T, kv_len), float("-inf"), device=q.device, dtype=q.dtype), |
| diagonal=1 + kv_len - T |
| ) |
| attn = attn + causal_mask.unsqueeze(0).unsqueeze(0) |
|
|
| if attention_mask is not None: |
| attn = attn + attention_mask |
|
|
| attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype) |
| out = torch.matmul(attn, v) |
| out = out.transpose(1, 2).contiguous().view(B, T, -1) |
| return self.o_proj(out), present |
|
|
|
|
| |
| class SmartCoderMoEMLP(nn.Module): |
| """ |
| Hybrid Dense + MoE MLP. |
| dense path: hidden -> dense_fc (8192) -> gelu -> dense_proj (2048) |
| expert path: router picks top-k experts from experts_fc/experts_proj |
| output = dense_out + expert_out |
| """ |
| def __init__(self, config: SmartCoderMoEConfig): |
| super().__init__() |
| H = config.hidden_size |
| DI = config.dense_intermediate_size |
| NE = config.num_experts |
| EI = config.expert_intermediate_size |
| K = config.num_experts_per_tok |
|
|
| self.num_experts = NE |
| self.top_k = K |
|
|
| |
| self.dense_fc = nn.Linear(H, DI, bias=True) |
| self.dense_proj = nn.Linear(DI, H, bias=True) |
|
|
| |
| |
| |
| self.experts_fc = nn.Parameter(torch.empty(NE, EI, H)) |
| self.experts_proj = nn.Parameter(torch.empty(NE, H, EI)) |
| self.router = nn.Linear(H, NE, bias=False) |
|
|
| def forward(self, x): |
| B, T, H = x.shape |
|
|
| |
| dense_out = self.dense_proj(F.gelu(self.dense_fc(x))) |
|
|
| |
| router_logits = self.router(x) |
| router_weights = F.softmax(router_logits, dim=-1) |
| top_weights, top_indices = router_weights.topk(self.top_k, dim=-1) |
| top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True) |
|
|
| |
| expert_out = torch.zeros_like(x) |
| x_flat = x.view(B * T, H) |
|
|
| for k in range(self.top_k): |
| expert_ids = top_indices[:, :, k].reshape(B * T) |
| weights = top_weights[:, :, k].reshape(B * T, 1) |
|
|
| |
| |
| |
| fc_w = self.experts_fc[expert_ids] |
| proj_w = self.experts_proj[expert_ids] |
|
|
| |
| hidden = F.gelu(torch.bmm(fc_w, x_flat.unsqueeze(-1)).squeeze(-1)) |
| |
| out = torch.bmm(proj_w, hidden.unsqueeze(-1)).squeeze(-1) |
|
|
| expert_out = expert_out + (out * weights).view(B, T, H) |
|
|
| return dense_out + expert_out |
|
|
|
|
| |
| class SmartCoderDecoderLayer(nn.Module): |
| def __init__(self, config: SmartCoderMoEConfig): |
| super().__init__() |
| self.input_layernorm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps) |
| self.self_attn = SmartCoderAttention(config) |
| self.post_attention_layernorm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps) |
| self.mlp = SmartCoderMoEMLP(config) |
|
|
| def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False): |
| |
| residual = hidden_states |
| hidden_states = self.input_layernorm(hidden_states) |
| hidden_states, present = self.self_attn( |
| hidden_states, attention_mask=attention_mask, |
| past_key_value=past_key_value, use_cache=use_cache |
| ) |
| hidden_states = residual + hidden_states |
|
|
| |
| residual = hidden_states |
| hidden_states = self.post_attention_layernorm(hidden_states) |
| hidden_states = self.mlp(hidden_states) |
| hidden_states = residual + hidden_states |
|
|
| return hidden_states, present |
|
|
|
|
| |
| class SmartCoderMoEModel(nn.Module): |
| def __init__(self, config: SmartCoderMoEConfig): |
| super().__init__() |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList([ |
| SmartCoderDecoderLayer(config) for _ in range(config.num_hidden_layers) |
| ]) |
| self.norm = LayerNormWithBias(config.hidden_size, config.rms_norm_eps) |
|
|
| def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False): |
| hidden_states = self.embed_tokens(input_ids) |
| presents = [] if use_cache else None |
|
|
| for i, layer in enumerate(self.layers): |
| pkv = past_key_values[i] if past_key_values else None |
| hidden_states, present = layer( |
| hidden_states, attention_mask=attention_mask, |
| past_key_value=pkv, use_cache=use_cache |
| ) |
| if use_cache: |
| presents.append(present) |
|
|
| hidden_states = self.norm(hidden_states) |
| return hidden_states, presents |
|
|
|
|
| |
| class SmartCoderMoEForCausalLM(PreTrainedModel): |
| config_class = SmartCoderMoEConfig |
| base_model_prefix = "model" |
| supports_gradient_checkpointing = False |
|
|
| def __init__(self, config: SmartCoderMoEConfig): |
| super().__init__(config) |
| self.model = SmartCoderMoEModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.model.embed_tokens |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| past_key_values=None, |
| inputs_embeds=None, |
| labels=None, |
| use_cache=True, |
| **kwargs, |
| ): |
| hidden_states, presents = self.model( |
| input_ids, attention_mask=attention_mask, |
| past_key_values=past_key_values, use_cache=use_cache |
| ) |
| logits = self.lm_head(hidden_states) |
|
|
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss = F.cross_entropy( |
| shift_logits.view(-1, shift_logits.size(-1)), |
| shift_labels.view(-1), |
| ignore_index=-100, |
| ) |
|
|
| return CausalLMOutputWithPast( |
| loss=loss, |
| logits=logits, |
| past_key_values=presents, |
| ) |
|
|
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs): |
| if past_key_values: |
| input_ids = input_ids[:, -1:] |
| return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True} |
|
|
|
|
| |
| def load_smartcoder_moe(model_id="Johnblick187/SmartCoderMoE", dtype=torch.bfloat16): |
| """Load SmartCoderMoE with correct custom architecture.""" |
| import os |
| from huggingface_hub import snapshot_download |
| from safetensors.torch import load_file |
|
|
| os.environ["HF_HUB_DISABLE_XET"] = "1" |
|
|
| print(f"Downloading {model_id}...") |
| model_dir = snapshot_download(model_id) |
|
|
| config = SmartCoderMoEConfig() |
| print("Initializing model...") |
| model = SmartCoderMoEForCausalLM(config) |
|
|
| print("Loading weights...") |
| from pathlib import Path |
| sf_files = sorted(Path(model_dir).glob("*.safetensors")) |
| state_dict = {} |
| for f in sf_files: |
| state_dict.update(load_file(str(f))) |
|
|
| |
| |
| |
|
|
| missing, unexpected = model.load_state_dict(state_dict, strict=False) |
| if missing: |
| print(f"Missing keys: {missing[:5]}{'...' if len(missing)>5 else ''}") |
| if unexpected: |
| print(f"Unexpected keys: {unexpected[:5]}{'...' if len(unexpected)>5 else ''}") |
|
|
| model = model.to(dtype) |
| print(f"Loaded! Params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B") |
| return model, config |
|
|
|
|
| if __name__ == "__main__": |
| from transformers import AutoTokenizer |
| import torch |
|
|
| model, config = load_smartcoder_moe() |
| model.eval() |
| model = model.cuda() |
|
|
| tokenizer = AutoTokenizer.from_pretrained("Johnblick187/SmartCoderMoE", trust_remote_code=True) |
|
|
| prompt = "def fibonacci(n):" |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) |
| input_len = inputs["input_ids"].shape[-1] |
|
|
| with torch.no_grad(): |
| out = model.generate( |
| **inputs, |
| max_new_tokens=150, |
| do_sample=True, |
| temperature=0.7, |
| top_p=0.95, |
| repetition_penalty=1.3, |
| pad_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| print(tokenizer.decode(out[0][input_len:], skip_special_tokens=True)) |
|
|
