| """ |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| HIDDEN_SIZE = 5120 |
| VOCAB_SIZE = 248320 |
| INTERMEDIATE_SIZE = 26624 |
| NUM_LAYERS = 128 |
| RMS_NORM_EPS = 1e-6 |
|
|
| |
| |
| LAYER_TYPES = ["linear_attention"] * 3 + ["full_attention"] |
| LAYER_TYPES = (LAYER_TYPES * (NUM_LAYERS // len(LAYER_TYPES) + 1))[:NUM_LAYERS] |
|
|
| |
| 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 |
| DELTA_VALUE_DIM = LINEAR_NUM_VALUE_HEADS * LINEAR_VALUE_HEAD_DIM |
| DELTA_CONV_DIM = DELTA_KEY_DIM * 2 + DELTA_VALUE_DIM |
| DELTA_CHUNK_SIZE = 64 |
|
|
| |
| 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) |
| """ |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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 |
| |
| |
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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_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): |
| |
| 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) |
| freqs = self._apply_interleaved_mrope(freqs) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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) |
| 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) |
|
|
| 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) |
|
|
| |
| 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__() |
| |
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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: |
| |
| |
| 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) |
| |
| |
| 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)) |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| |
| |
| 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 = {} |
| 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 |
|
|
| |
| |
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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)}") |
|
|
