""" JiRackDeltaNet_70b.py ============================= Собственная (hand-rolled) архитектура "JiRack" -- в стиле прежнего DS32B-проекта: один файл, свой TransformerBlock, никакого AutoConfig/AutoModelForCausalLM внутри класса архитектуры, никакого model_id в конструкторе. Все размерности зашиты константами ниже -- подтверждены живым инспектом реального checkpoint'а (inspect_qwen38_keys.py, 851 тензор, сверено). Математика DeltaNet-рекуррентности, свёртки, гейтинга и partial-rotary mRoPE-attention ПОРТИРОВАНА из реального modeling_qwen3_5.py (Apache 2.0, HuggingFace/Qwen), а не переизобретена -- чтобы не рисковать тихой числовой ошибкой в тонкой рекуррентной формуле. Портирован смысл, но код переписан под наши имена/стиль и урезан: без vision-башни, без KV-cache классов transformers (для тренировки полный форвард по последовательности не нужен; для генерации -- полный recompute на каждом шаге, без инкрементального кэша; это сознательное упрощение, при желании кэш можно добавить отдельно). Имена подмодулей (`linear_attn` / `self_attn`, `in_proj_qkv`, `out_proj`, `mlp.gate_proj` и т.д.) намеренно СОВПАДАЮТ с именами в реальном state_dict Qwen3_5ForCausalLM -- поэтому веса из HF-чекпоинта грузятся напрямую через load_state_dict(), без ручного key-mapping. Файл умеет: - собрать модель (JiRackQwen38ForCausalLM) с нуля, случайные веса; - принять state_dict в формате HF Qwen3_5ForCausalLM (после AutoModelForCausalLM.from_pretrained(...).state_dict() в конвертере) и загрузить его напрямую; - пропатчить нужные nn.Linear -> BitLinear (TERNARIZE_PATTERNS); - сохранить/загрузить свой JiRack-чекпоинт (save_checkpoint/load_checkpoint). Конвертер (convert_qwen38_to_jirack.py) по-прежнему знает про model_id и HF Hub -- но это знание живёт ТОЛЬКО в конвертере, не в архитектуре. """ import os import re import torch import torch.nn as nn import torch.nn.functional as F # ===================================================================== # 1. CONFIRMED АРХИТЕКТУРА -- жёстко зашитые константы (не AutoConfig) # ===================================================================== HIDDEN_SIZE = 5120 VOCAB_SIZE = 248320 INTERMEDIATE_SIZE = 26624 NUM_LAYERS = 128 RMS_NORM_EPS = 1e-6 # Чередование слоёв: full_attention_interval = 4 -> # 3x linear_attention, 1x full_attention, повторяется на все 64 слоя. LAYER_TYPES = ["linear_attention"] * 3 + ["full_attention"] LAYER_TYPES = (LAYER_TYPES * (NUM_LAYERS // len(LAYER_TYPES) + 1))[:NUM_LAYERS] # --- Gated DeltaNet (linear_attention) --- LINEAR_NUM_KEY_HEADS = 16 LINEAR_KEY_HEAD_DIM = 128 LINEAR_NUM_VALUE_HEADS = 48 LINEAR_VALUE_HEAD_DIM = 128 LINEAR_CONV_KERNEL_DIM = 4 DELTA_KEY_DIM = LINEAR_NUM_KEY_HEADS * LINEAR_KEY_HEAD_DIM # 2048 DELTA_VALUE_DIM = LINEAR_NUM_VALUE_HEADS * LINEAR_VALUE_HEAD_DIM # 6144 DELTA_CONV_DIM = DELTA_KEY_DIM * 2 + DELTA_VALUE_DIM # 10240 DELTA_CHUNK_SIZE = 64 # --- Gated Attention (full_attention) --- NUM_ATTENTION_HEADS = 24 NUM_KEY_VALUE_HEADS = 4 HEAD_DIM = 256 PARTIAL_ROTARY_FACTOR = 0.25 ROPE_THETA = 1e7 MROPE_SECTION = [11, 11, 10] ATTENTION_BIAS = False ARCH_SUMMARY = f""" Qwen layout (qwen3_5), hand-rolled JiRack TransformerBlock, text-only. hidden_size={HIDDEN_SIZE}, vocab_size={VOCAB_SIZE}, intermediate_size={INTERMEDIATE_SIZE}, num_hidden_layers={NUM_LAYERS} Чередование (interval=4): 3x linear_attention (DeltaNet), 1x full_attention """ KEPT_FULL_PRECISION = """ Всегда остаются full precision (не GEMM, либо слишком чувствительны): - dt_bias, A_log (сырые Parameters, decay DeltaNet) - linear_attn.in_proj_b/in_proj_a (по умолчанию, см. GATE_PROJ_PATTERNS) - conv1d (depthwise grouped conv, не GEMM) - все нормы: RMSNormGated, q_norm/k_norm, model.norm - embed_tokens, lm_head - mtp.* (вся MTP/NextN-голова) (draft-слой спекулятивного декода; тернаризация убивает acceptance rate, комьюнити-GGUF держат её в Q8_0) """ # ===================================================================== # 2. BitLinear -- наш тернарный слой # ===================================================================== class BitLinear(nn.Linear): """b1.58-style per-tensor absmean ternary weights, per-token int8 activations, непрерывный lambda-warmup (STE).""" def __init__(self, in_features, out_features, bias=False): super().__init__(in_features, out_features, bias=bias) self.eps = 1e-5 self.register_buffer("lambda_", torch.zeros(()), persistent=True) @classmethod def from_linear(cls, lin: nn.Linear) -> "BitLinear": bl = cls(lin.in_features, lin.out_features, bias=lin.bias is not None) bl = bl.to(dtype=lin.weight.dtype, device=lin.weight.device) with torch.no_grad(): bl.weight.copy_(lin.weight) if lin.bias is not None: bl.bias.copy_(lin.bias) return bl def forward(self, x: torch.Tensor) -> torch.Tensor: if not self.training and float(self.lambda_) < 1e-6: return F.linear(x, self.weight, self.bias) lam = self.lambda_.to(x.dtype) w = self.weight gamma = w.float().abs().mean().clamp(min=self.eps).to(w.dtype) w_quant = torch.clamp(torch.round(w / gamma), -1.0, 1.0) * gamma w_effective = w + lam * (w_quant - w).detach() x_scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp(min=self.eps) x_quant = torch.clamp(torch.round(x * x_scale), -128.0, 127.0) / x_scale x_effective = x + lam * (x_quant - x).detach() return F.linear(x_effective, w_effective, self.bias) def set_lambda(self, v: float): self.lambda_.fill_(v) # ===================================================================== # 3. Карта: что тернаризуем, что оставляем full precision # ===================================================================== TERNARIZE_PATTERNS = [ r"linear_attn\.in_proj_qkv$", r"linear_attn\.in_proj_z$", r"linear_attn\.out_proj$", r"self_attn\.q_proj$", r"self_attn\.k_proj$", r"self_attn\.v_proj$", r"self_attn\.o_proj$", r"mlp\.gate_proj$", r"mlp\.up_proj$", r"mlp\.down_proj$", ] GATE_PROJ_PATTERNS = [ r"linear_attn\.in_proj_b$", r"linear_attn\.in_proj_a$", ] def apply_bitlinear_patch(model: nn.Module, include_gate_proj: bool = False): """Точечная замена nn.Linear -> BitLinear по TERNARIZE_PATTERNS.""" patterns = list(TERNARIZE_PATTERNS) if include_gate_proj: patterns += GATE_PROJ_PATTERNS compiled = [re.compile(p) for p in patterns] patched = [] for name, module in list(model.named_modules()): if not isinstance(module, nn.Linear): continue # MTP/NextN-голова НЕ тернаризуется -- всегда full precision. # Суффиксные паттерны (self_attn.q_proj$ и т.п.) иначе зацепили бы # mtp.layers.0.self_attn.q_proj. if name.startswith("mtp"): continue if any(p.search(name) for p in compiled): parent_name, _, child_name = name.rpartition(".") parent = model.get_submodule(parent_name) if parent_name else model setattr(parent, child_name, BitLinear.from_linear(module)) patched.append(name) return patched # ===================================================================== # 4. Общие блоки (порт математики из реального modeling_qwen3_5.py) # ===================================================================== class RMSNorm(nn.Module): """(x * w).to(dtype), w хранится как (weight) с offset +1 -- как в оригинале Qwen3_5RMSNorm (инициализация нулями, эффективный вес 1+w).""" def __init__(self, dim: int, eps: float = RMS_NORM_EPS): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) x = x * (1.0 + self.weight.float()) return x.to(dtype) class RMSNormGated(nn.Module): """RMSNorm с гейтингом через SiLU(gate), применяется ПЕРЕД гейтингом (норма считается по hidden_states, гейт применяется после).""" def __init__(self, dim: int, eps: float = RMS_NORM_EPS): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: dtype = x.dtype x = x.float() var = x.pow(2).mean(-1, keepdim=True) x = x * torch.rsqrt(var + self.eps) x = self.weight * x.to(dtype) x = x * F.silu(gate.float()) return x.to(dtype) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Partial rotary: только первые cos.shape[-1] каналов вращаются, остальное (q_pass/k_pass) проходит без изменений.""" cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_embed = torch.cat([(q_rot * cos) + (rotate_half(q_rot) * sin), q_pass], dim=-1) k_embed = torch.cat([(k_rot * cos) + (rotate_half(k_rot) * sin), k_pass], dim=-1) return q_embed, k_embed class TextRotaryEmbedding(nn.Module): """mRoPE с partial_rotary_factor. Для чистого текста (без vision) 3 позиционные "оси" (T,H,W) совпадают -> interleaved mRoPE вырождается в обычный RoPE, но мы всё равно считаем честно тем же путём, что и оригинал, для полной числовой идентичности.""" def __init__(self, head_dim=HEAD_DIM, partial_rotary_factor=PARTIAL_ROTARY_FACTOR, theta=ROPE_THETA, mrope_section=MROPE_SECTION): super().__init__() self.rotary_dim = int(head_dim * partial_rotary_factor) inv_freq = 1.0 / (theta ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) / self.rotary_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.mrope_section = mrope_section def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: # freqs: (3, bs, seq, rotary_dim // 2) freqs_t = freqs[0].clone() for dim, offset in enumerate((1, 2), start=1): length = self.mrope_section[dim] * 3 idx = slice(offset, length, 3) freqs_t[..., idx] = freqs[dim, ..., idx] return freqs_t @torch.no_grad() def forward(self, x: torch.Tensor, position_ids: torch.LongTensor): # position_ids: (bs, seq) -> реплицируем в 3 идентичные оси (text-only) if position_ids.ndim == 2: position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) inv_freq_exp = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) pos_exp = position_ids[:, :, None, :].float() freqs = (inv_freq_exp @ pos_exp).transpose(2, 3) # (3, bs, seq, rotary_dim//2) freqs = self._apply_interleaved_mrope(freqs) # (bs, seq, rotary_dim//2) emb = torch.cat((freqs, freqs), dim=-1) # (bs, seq, rotary_dim) return emb.cos().to(x.dtype), emb.sin().to(x.dtype) def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6): return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) def causal_conv1d(hidden_states: torch.Tensor, weight: torch.Tensor, activation: str = "silu"): """hidden_states: (batch, channels, seq_len). Depthwise causal conv, left-padded by kernel_size - 1, затем активация.""" channels, seq_len = hidden_states.shape[1], hidden_states.shape[2] padding = weight.shape[-1] - 1 out = F.conv1d(hidden_states, weight=weight.unsqueeze(1), bias=None, padding=padding, groups=channels) out = out[:, :, :seq_len] return F.silu(out) if activation == "silu" else out def chunked_gated_delta_rule(query, key, value, g, beta, chunk_size=DELTA_CHUNK_SIZE): """Порт torch_chunk_gated_delta_rule из modeling_qwen3_5.py -- без initial_state/output_final_state (тренировка/полный форвард без инкрементального кэша). query/key/value: (b, h, t, d), g/beta: (b, h, t).""" initial_dtype = query.dtype query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) query, key, value, beta, g = [x.contiguous().to(torch.float32) for x in (query, key, value, beta, g)] b, h, t, k_dim = key.shape v_dim = value.shape[-1] pad = (chunk_size - t % chunk_size) % chunk_size query = F.pad(query, (0, 0, 0, pad)) key = F.pad(key, (0, 0, 0, pad)) value = F.pad(value, (0, 0, 0, pad)) beta = F.pad(beta, (0, pad)) g = F.pad(g, (0, pad)) total_t = t + pad scale = 1 / (query.shape[-1] ** 0.5) query = query * scale v_beta = value * beta.unsqueeze(-1) k_beta = key * beta.unsqueeze(-1) query, key, value, k_beta, v_beta = [ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta) ] g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) g = g.cumsum(dim=-1) decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) for i in range(1, chunk_size): row = attn[..., i, :i].clone() sub = attn[..., :i, :i].clone() attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) value = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) last_state = torch.zeros(b, h, k_dim, v_dim, dtype=value.dtype, device=value.device) core_out = torch.zeros_like(value) for i in range(total_t // chunk_size): q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] attn_i = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i] v_prime = k_cumdecay[:, :, i] @ last_state v_new = v_i - v_prime attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state core_out[:, :, i] = attn_inter + attn_i @ v_new last_state = ( last_state * g[:, :, i, -1, None, None].exp() + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new ) core_out = core_out.reshape(core_out.shape[0], core_out.shape[1], -1, core_out.shape[-1]) core_out = core_out[:, :, :t] core_out = core_out.transpose(1, 2).contiguous().to(initial_dtype) return core_out # ===================================================================== # 5. Наш собственный TransformerBlock (и подмодули linear_attn/self_attn) # ===================================================================== class GatedDeltaNetBlock(nn.Module): """linear_attention слой (Gated DeltaNet). Имена подмодулей совпадают с HF state_dict, чтобы веса грузились напрямую.""" def __init__(self): super().__init__() self.conv1d = nn.Conv1d( in_channels=DELTA_CONV_DIM, out_channels=DELTA_CONV_DIM, bias=False, kernel_size=LINEAR_CONV_KERNEL_DIM, groups=DELTA_CONV_DIM, padding=LINEAR_CONV_KERNEL_DIM - 1, ) self.dt_bias = nn.Parameter(torch.ones(LINEAR_NUM_VALUE_HEADS)) A = torch.empty(LINEAR_NUM_VALUE_HEADS).uniform_(0.01, 16) self.A_log = nn.Parameter(torch.log(A)) self.norm = RMSNormGated(LINEAR_VALUE_HEAD_DIM, eps=RMS_NORM_EPS) self.out_proj = nn.Linear(DELTA_VALUE_DIM, HIDDEN_SIZE, bias=False) self.in_proj_qkv = nn.Linear(HIDDEN_SIZE, DELTA_KEY_DIM * 2 + DELTA_VALUE_DIM, bias=False) self.in_proj_z = nn.Linear(HIDDEN_SIZE, DELTA_VALUE_DIM, bias=False) self.in_proj_b = nn.Linear(HIDDEN_SIZE, LINEAR_NUM_VALUE_HEADS, bias=False) self.in_proj_a = nn.Linear(HIDDEN_SIZE, LINEAR_NUM_VALUE_HEADS, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: bsz, seq_len, _ = hidden_states.shape mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2) # (b, conv_dim, t) z = self.in_proj_z(hidden_states).reshape(bsz, seq_len, -1, LINEAR_VALUE_HEAD_DIM) b = self.in_proj_b(hidden_states) a = self.in_proj_a(hidden_states) mixed_qkv = causal_conv1d(mixed_qkv, self.conv1d.weight.squeeze(1), activation="silu") mixed_qkv = mixed_qkv.transpose(1, 2) # (b, t, conv_dim) query, key, value = torch.split(mixed_qkv, [DELTA_KEY_DIM, DELTA_KEY_DIM, DELTA_VALUE_DIM], dim=-1) query = query.reshape(bsz, seq_len, -1, LINEAR_KEY_HEAD_DIM) key = key.reshape(bsz, seq_len, -1, LINEAR_KEY_HEAD_DIM) value = value.reshape(bsz, seq_len, -1, LINEAR_VALUE_HEAD_DIM) beta = b.sigmoid() g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) rep = LINEAR_NUM_VALUE_HEADS // LINEAR_NUM_KEY_HEADS if rep > 1: query = query.repeat_interleave(rep, dim=2) key = key.repeat_interleave(rep, dim=2) # -> (b, h, t, d) for the chunked kernel query_h = query.transpose(1, 2) key_h = key.transpose(1, 2) value_h = value.transpose(1, 2) beta_h = beta.transpose(1, 2) g_h = g.transpose(1, 2) core_attn_out = chunked_gated_delta_rule(query_h, key_h, value_h, g_h, beta_h) core_attn_out = core_attn_out.reshape(-1, LINEAR_VALUE_HEAD_DIM) z_flat = z.reshape(-1, LINEAR_VALUE_HEAD_DIM) core_attn_out = self.norm(core_attn_out, z_flat) core_attn_out = core_attn_out.reshape(bsz, seq_len, -1) return self.out_proj(core_attn_out) class AttentionBlock(nn.Module): """full_attention слой (Gated Attention + partial rotary mRoPE).""" def __init__(self): super().__init__() self.num_key_value_groups = NUM_ATTENTION_HEADS // NUM_KEY_VALUE_HEADS self.scaling = HEAD_DIM ** -0.5 self.q_proj = nn.Linear(HIDDEN_SIZE, NUM_ATTENTION_HEADS * HEAD_DIM * 2, bias=ATTENTION_BIAS) self.k_proj = nn.Linear(HIDDEN_SIZE, NUM_KEY_VALUE_HEADS * HEAD_DIM, bias=ATTENTION_BIAS) self.v_proj = nn.Linear(HIDDEN_SIZE, NUM_KEY_VALUE_HEADS * HEAD_DIM, bias=ATTENTION_BIAS) self.o_proj = nn.Linear(NUM_ATTENTION_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=ATTENTION_BIAS) self.q_norm = RMSNorm(HEAD_DIM, eps=RMS_NORM_EPS) self.k_norm = RMSNorm(HEAD_DIM, eps=RMS_NORM_EPS) def forward(self, hidden_states, cos, sin, causal_mask): bsz, seq_len, _ = hidden_states.shape hidden_shape = (bsz, seq_len, -1, HEAD_DIM) query_states, gate = torch.chunk( self.q_proj(hidden_states).view(bsz, seq_len, -1, HEAD_DIM * 2), 2, dim=-1 ) gate = gate.reshape(bsz, seq_len, -1) query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) rep = self.num_key_value_groups key_states = key_states.repeat_interleave(rep, dim=1) value_states = value_states.repeat_interleave(rep, dim=1) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling attn_weights = attn_weights + causal_mask attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, seq_len, -1) attn_output = attn_output * torch.sigmoid(gate) return self.o_proj(attn_output) class MLPBlock(nn.Module): def __init__(self): super().__init__() self.gate_proj = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False) self.up_proj = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False) self.down_proj = nn.Linear(INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class TransformerBlock(nn.Module): """Наш decoder layer -- linear_attn ИЛИ self_attn, в зависимости от layer_type, + общий MLP. Имена атрибутов совпадают с HF state_dict.""" def __init__(self, layer_idx: int, layer_type: str = None): super().__init__() # layer_type задан явно -> используем его (нужно для MTP-слоя, # который всегда full_attention и лежит ВНЕ LAYER_TYPES). self.block_type = layer_type if layer_type is not None else LAYER_TYPES[layer_idx] if self.block_type == "linear_attention": self.linear_attn = GatedDeltaNetBlock() else: self.self_attn = AttentionBlock() self.mlp = MLPBlock() self.input_layernorm = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) self.post_attention_layernorm = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) def forward(self, hidden_states, cos, sin, causal_mask): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) if self.block_type == "linear_attention": hidden_states = self.linear_attn(hidden_states) else: hidden_states = self.self_attn(hidden_states, cos, sin, causal_mask) 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 # ===================================================================== # 5b. MTP / NextN-голова (draft-слой спекулятивного декода) # ===================================================================== class MTPBlock(nn.Module): """Multi-Token Prediction голова -- ровно как в реальном чекпоинте Имена подмодулей СОВПАДАЮТ с ключами HF state_dict (top-level `mtp.`): mtp.pre_fc_norm_embedding.weight mtp.pre_fc_norm_hidden.weight mtp.fc.weight [HIDDEN, 2*HIDDEN] mtp.layers.0.* (обычный full_attention decoder) mtp.norm.weight (final norm перед общим lm_head) embed_tokens и lm_head ОБЩИЕ со стволом (dedicated embeddings = false), здесь их нет. В GGUF это blk.{NUM_LAYERS}.nextn.eh_proj / shared_head_norm + тензоры draft-слоя. Механика (DeepSeek-style, порядок как в vLLM Qwen3NextMTP): h = fc(cat(pre_fc_norm_embedding(emb_next), pre_fc_norm_hidden(h_prev))) h = layers[0](h) # full attention logits2 = lm_head(norm(h)) """ def __init__(self): super().__init__() self.pre_fc_norm_embedding = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) self.pre_fc_norm_hidden = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) self.fc = nn.Linear(2 * HIDDEN_SIZE, HIDDEN_SIZE, bias=False) self.layers = nn.ModuleList([ TransformerBlock(layer_idx=NUM_LAYERS, layer_type="full_attention"), ]) self.norm = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) def forward(self, inputs_embeds, previous_hidden, cos, sin, causal_mask): """inputs_embeds -- embed_tokens(токены, сдвинутые на +1) previous_hidden -- hidden ствола ДО финального model.norm (см. контракт: MTP сам применяет pre_fc_norm_hidden).""" h = self.fc(torch.cat( (self.pre_fc_norm_embedding(inputs_embeds), self.pre_fc_norm_hidden(previous_hidden)), dim=-1, )) h = self.layers[0](h, cos, sin, causal_mask) return self.norm(h) # ===================================================================== # 6. Полная модель # ===================================================================== class JiRackQwen38Model(nn.Module): def __init__(self): super().__init__() self.embed_tokens = nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE) self.layers = nn.ModuleList([TransformerBlock(i) for i in range(NUM_LAYERS)]) self.norm = RMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) self.rotary_emb = TextRotaryEmbedding() def forward(self, input_ids: torch.LongTensor, return_prenorm: bool = False): bsz, seq_len = input_ids.shape hidden_states = self.embed_tokens(input_ids) position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(bsz, -1) cos, sin = self.rotary_emb(hidden_states, position_ids) causal_mask = torch.full((seq_len, seq_len), float("-inf"), device=input_ids.device, dtype=torch.float32) causal_mask = torch.triu(causal_mask, diagonal=1)[None, None, :, :].to(hidden_states.dtype) for layer in self.layers: hidden_states = layer(hidden_states, cos, sin, causal_mask) if return_prenorm: # Для MTP: голова сама применяет pre_fc_norm_hidden, финальный # model.norm сюда НЕ входит (иначе двойная нормализация). return self.norm(hidden_states), hidden_states, (cos, sin, causal_mask) return self.norm(hidden_states) class JiRackQwen38ForCausalLM(nn.Module): """Верхнеуровневая модель. Имена атрибутов (`model`, `lm_head`) совпадают с HF Qwen3_5ForCausalLM -- state_dict грузится напрямую.""" def __init__(self): super().__init__() self.model = JiRackQwen38Model() self.lm_head = nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False) # MTP/NextN-голова: top-level атрибут `mtp` -- ключи state_dict # получаются ровно mtp.*, как в реальном HF-чекпоинте. self.mtp = MTPBlock() def forward(self, input_ids: torch.LongTensor) -> torch.Tensor: hidden_states = self.model(input_ids) return self.lm_head(hidden_states) def mtp_forward(self, input_ids: torch.LongTensor): """Возвращает (logits_t1, logits_t2): обычные логиты ствола и draft-логиты MTP-головы (предсказание токена t+2 из позиции t). Для инференса/валидации MTP; в GGUF-конвертации не участвует.""" _, prenorm_hidden, (cos, sin, causal_mask) = self.model( input_ids, return_prenorm=True) logits_t1 = self.lm_head(self.model.norm(prenorm_hidden)) # Вход MTP -- эмбеддинги токенов, сдвинутых на +1 (следующий токен # уже известен на шаге верификации). Последнюю позицию дублируем, # чтобы сохранить форму (draft для неё не используется). shifted = torch.cat([input_ids[:, 1:], input_ids[:, -1:]], dim=1) emb_next = self.model.embed_tokens(shifted) h2 = self.mtp(emb_next, prenorm_hidden, cos, sin, causal_mask) logits_t2 = self.lm_head(h2) return logits_t1, logits_t2 # ----------------------------------------------------------------- # Загрузка реальных весов (state_dict уже скачан конвертером через # HF Hub -- сюда просто передаётся словарь тензоров) # ----------------------------------------------------------------- def load_hf_state_dict(self, state_dict: dict, strict: bool = False): """Грузит state_dict в формате оригинального Qwen3_5ForCausalLM. mtp.* ключи теперь ГРУЗЯТСЯ в self.mtp (MTPBlock) -- фильтра больше нет. strict=False оставлен из-за служебных lambda_-буферов BitLinear, которых нет в реальном чекпоинте.""" missing, unexpected = self.load_state_dict(state_dict, strict=strict) real_missing = [k for k in missing if not k.endswith("lambda_")] if real_missing: print(f"⚠️ {len(real_missing)} missing keys (первые 10): {real_missing[:10]}") if unexpected: print(f"⚠️ {len(unexpected)} unexpected keys (первые 10): {unexpected[:10]}") return missing, unexpected # ----------------------------------------------------------------- # Подсадка РЕАЛЬНЫХ mtp.* весов из оригинального HF-чекпоинта. # Нужна, потому что старые JiRack .pt сохранены ДО появления MTPBlock # и mtp-весов не содержат (там была бы случайная инициализация). # ----------------------------------------------------------------- def load_mtp_from_hf_dir(self, hf_dir: str): """Читает из папки оригинального HF-чекпоинта ТОЛЬКО mtp.* тензоры (по index.json или перебором *.safetensors) и грузит их в self.mtp. Возвращает число загруженных тензоров.""" import glob import json as _json from safetensors import safe_open index_path = os.path.join(hf_dir, "model.safetensors.index.json") mtp_map = {} # key -> shard file if os.path.exists(index_path): with open(index_path) as f: weight_map = _json.load(f)["weight_map"] for k, shard in weight_map.items(): if k.startswith("mtp."): mtp_map[k] = os.path.join(hf_dir, shard) else: for shard in glob.glob(os.path.join(hf_dir, "*.safetensors")): with safe_open(shard, framework="pt", device="cpu") as f: for k in f.keys(): if k.startswith("mtp."): mtp_map[k] = shard if not mtp_map: raise FileNotFoundError( f"В {hf_dir} не найдено ни одного mtp.* тензора. Нужен " f"ОРИГИНАЛЬНЫЙ чекпоинт Qwen (не JiRack-экспорт!). " f"Если оригинальные шарды удалены -- перекачай с HF Hub " f"только шарды с mtp.* по model.safetensors.index.json.") own_dtype = next(self.mtp.parameters()).dtype loaded = {} by_shard = {} for k, shard in mtp_map.items(): by_shard.setdefault(shard, []).append(k) for shard, keys in by_shard.items(): with safe_open(shard, framework="pt", device="cpu") as f: for k in keys: loaded[k] = f.get_tensor(k).to(own_dtype) missing, unexpected = self.load_state_dict(loaded, strict=False) n_ok = len(loaded) print(f"✅ MTP: загружено {n_ok} тензоров из {len(by_shard)} шардов ({hf_dir})") if unexpected: print(f"⚠️ MTP: unexpected keys: {unexpected[:10]}") return n_ok # ----------------------------------------------------------------- # JiRack-чекпоинт (свой формат, с lambda/patched_linear_names) # ----------------------------------------------------------------- def save_checkpoint(self, path: str, lam: float, global_step: int = 0, patched_linear_names=None): torch.save({ "model": self.state_dict(), "lambda": lam, "global_step": global_step, "patched_linear_names": patched_linear_names or [], }, path) @classmethod def load_checkpoint(cls, path: str, dtype=torch.bfloat16, include_gate_proj: bool = False): ckpt = torch.load(path, map_location="cpu", weights_only=False) model = cls() patched_names = ckpt.get("patched_linear_names", []) if patched_names: for name in patched_names: *parent_path, leaf = name.split(".") parent = model for p in parent_path: parent = getattr(parent, p) old = getattr(parent, leaf) if not isinstance(old, BitLinear): setattr(parent, leaf, BitLinear.from_linear(old)) else: apply_bitlinear_patch(model, include_gate_proj=include_gate_proj) has_mtp_in_ckpt = any(k.startswith("mtp.") for k in ckpt["model"]) model.load_state_dict(ckpt["model"], strict=False) model = model.to(dtype) if not has_mtp_in_ckpt: print( "⚠️ В этом .pt НЕТ mtp.* весов (чекпоинт сохранён до " "появления MTPBlock). Голова сейчас со СЛУЧАЙНОЙ " "инициализацией! Перед экспортом обязательно вызови " "model.load_mtp_from_hf_dir('<папка оригинального HF-чекпоинта>')." ) lam = ckpt.get("lambda", 0.0) for m in model.modules(): if isinstance(m, BitLinear): m.set_lambda(lam) print(f"✅ Загружено. global_step={ckpt.get('global_step', 0)}, lambda={lam}") return model # ===================================================================== # 7. Демонстрация (без скачивания весов -- только структура) # ===================================================================== if __name__ == "__main__": print(ARCH_SUMMARY) print(KEPT_FULL_PRECISION) model = JiRackQwen38ForCausalLM() n_params = sum(p.numel() for p in model.parameters()) print(f"✅ Модель собрана (случайные веса). Параметров: {n_params / 1e9:.2f}B") patched = apply_bitlinear_patch(model) print(f"✅ Заменено подмодулей на BitLinear: {len(patched)}") print("Примеры первых 10:") for n in patched[:10]: print(f" {n}") ids = torch.randint(0, VOCAB_SIZE, (1, 8)) with torch.no_grad(): logits = model(ids) print(f"✅ Пробный forward (случайные веса) -> logits shape: {tuple(logits.shape)}")