#!/usr/bin/env python3 """Weight transfer from Qwen3.5-2B donor to Spider-FLEXITOKENS architecture. Implements the weight transfer pipeline per D-09 and D-10: - Loads Qwen3.5-2B via HF transformers - Filters to full_attention layers only (discards linear_attention) - SVD decomposition converts standard GQA attention to MLA format - Direct copies where shapes match (o_proj, layer norms) - Reinitializes incompatible weights (embeddings, boundary predictor, FFN) - Reports transfer coverage as percentage Usage: python scripts/transfer_weights.py --donor Qwen/Qwen3.5-2B --output models/Spider-FLEXITOKENS-init/ --config spider_flexitokens_997m """ import argparse import hashlib import json import math import os import sys from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # Import canonical Spider architecture components from spider.py # (replaces previously duplicated code — per VERIFICATION gap #2, #4, #5) from spider import ( SENTINEL_TOKENS, is_sentinel_token, create_modality_mask, BoundaryPredictor, downsample, upsample, SpiderConfig as _CanonicalSpiderConfig, spider_flexitokens_997m as _canonical_config_fn, ) # Reverse mapping for sentinel token IDs to names (IN-01 fix: computed once) _TOKEN_NAMES_BY_ID = {v: k for k, v in SENTINEL_TOKENS.items()} # ============================================================================ # Sentinel Token Vocabulary — imported from spider.py (D-06, D-11) # ============================================================================ # SENTINEL_TOKENS, is_sentinel_token, create_modality_mask are now imported # from spider.py. _SENTINEL_PAIRS and _MODALITY_SENTINEL_IDS are used # locally for transfer logic. _SENTINEL_PAIRS = [ (SENTINEL_TOKENS['IMG_START'], SENTINEL_TOKENS['IMG_END']), # (259, 260) (SENTINEL_TOKENS['AUD_START'], SENTINEL_TOKENS['AUD_END']), # (261, 262) (SENTINEL_TOKENS['VID_START'], SENTINEL_TOKENS['VID_END']), # (263, 264) ] _MODALITY_SENTINEL_IDS = {259, 260, 261, 262, 263, 264} # ============================================================================ # BoundaryPredictor — imported from spider.py (D-04, D-11) # ============================================================================ # BoundaryPredictor is now imported from spider.py. # ============================================================================ # Downsample / Upsample — imported from spider.py (D-05, D-08, D-11) # ============================================================================ # downsample, upsample, _downsample_common, _downsample_final are now # imported from spider.py. # ============================================================================ # Spider Configuration # ============================================================================ @dataclass class SpiderConfig: """Spider-FLEXITOKENS model configuration (hidden_size=2048). Based on mythos-fineweb-moe.py SpiderPortalConfig with byte-level tokenization and MLA attention. Mirrors canonical spider.py config. """ # Core architecture vocab_size: int = 272 # 256 bytes + 16 specials (D-06) hidden_size: int = 2048 num_hidden_layers: int = 6 # recurrent layers num_attention_heads: int = 16 num_key_value_heads: int = 4 # not used directly in MLA but kept for compat intermediate_size: int = 1024 hidden_act: str = "silu" # MoE configuration (D-20, D-21: shared-projection MoE) num_experts: int = 32 num_experts_per_tok: int = 2 num_shared_experts: int = 1 router_aux_loss_coef: float = 0.05 shared_intermediate_size: int = 6144 expert_core_rank: int = 256 shared_expert_intermediate_size: int = 7424 prelude_coda_intermediate_size: int = 4096 # RDT configuration max_loop_iters: int = 16 act_threshold: float = 0.5 prelude_layers: int = 2 coda_layers: int = 2 lora_rank: int = 128 loop_embed_dim: int = 128 # MLA parameters (DeepSeek-V2 style) kv_lora_rank: int = 128 q_lora_rank: int = 256 qk_rope_head_dim: int = 64 qk_nope_head_dim: int = 64 v_head_dim: int = 64 # Attention / RoPE max_position_embeddings: int = 262144 # 256k context rope_theta: float = 10000000.0 rope_scaling: Optional[Dict] = field(default_factory=lambda: { "type": "yarn", "factor": 8.0, "original_max_position_embeddings": 32768, }) sliding_window: int = 8192 # local attention window attention_dropout: float = 0.0 rms_norm_eps: float = 1e-6 initializer_range: float = 0.02 # Embeddings / head tie_word_embeddings: bool = True # Tied per D-06 (byte-level vocab) # Metadata model_type: str = "spider" torch_dtype: str = "bfloat16" # BoundaryPredictor bp_d_inner: int = 8192 # Engram (N-gram memory, D-20 revision) engram_layers: list = None # set in __post_init__ engram_table_size: int = 8191 engram_heads: int = 4 engram_dim: int = 128 engram_offload: bool = True # Multimodal vision_hidden_size: int = 2048 audio_hidden_size: int = 512 vision_num_frames: int = 60 vision_tokens_per_frame: int = 256 vision_temporal_tokens: int = 64 vision_temporal_layers: int = 2 @property def head_dim(self): return self.qk_nope_head_dim + self.qk_rope_head_dim # 128 def __post_init__(self): if self.engram_layers is None: self.engram_layers = [1, 4] def spider_flexitokens_997m() -> SpiderConfig: """Spider-FLEXITOKENS 997M config.""" return SpiderConfig() # ============================================================================ # Dummy Donor (for testing without downloading 6GB model) # ============================================================================ def create_dummy_donor(num_layers: int = 4, full_attention_layers: Optional[List[int]] = None, mini: bool = False): """Create a dummy Qwen3.5-2B-like donor state dict and config. Mimics the structure of Qwen3.5-2B with: - hidden_size=2048, num_heads=8, num_kv_heads=2, head_dim=256 - full_attention and linear_attention layer identification - intermediate_size=6144 - vocab_size=248320 Args: num_layers: Number of layers to create full_attention_layers: Indices of full_attention layers (default: all) mini: If True, use smaller tensors for fast testing Returns: Dict with "state_dict", "config" keys """ hidden_size = 2048 num_heads = 8 num_kv_heads = 2 head_dim = 256 # Qwen3.5-2B: 2048 / 8 = 256 intermediate_size = 6144 vocab_size = 248320 if full_attention_layers is None: # Default: all layers are full_attention for testing full_attention_layers = list(range(num_layers)) # Scale factor for mini mode (reduces tensor sizes for fast testing) scale = 8 if mini else 1 hs = hidden_size // scale n_h = max(num_heads // scale, 1) n_kv_h = max(num_kv_heads // scale, 1) hd = head_dim # Keep head_dim the same for shape correctness inter = intermediate_size // scale vs = min(vocab_size, 1024) if mini else vocab_size state_dict = {} # Embeddings state_dict["model.embed_tokens.weight"] = torch.randn(vs, hs) * 0.02 # Per-layer weights for i in range(num_layers): prefix = f"model.layers.{i}" # Attention projections (Qwen3.5-2B layout) state_dict[f"{prefix}.self_attn.q_proj.weight"] = torch.randn(n_h * hd, hs) * 0.02 state_dict[f"{prefix}.self_attn.k_proj.weight"] = torch.randn(n_kv_h * hd, hs) * 0.02 state_dict[f"{prefix}.self_attn.v_proj.weight"] = torch.randn(n_kv_h * hd, hs) * 0.02 state_dict[f"{prefix}.self_attn.o_proj.weight"] = torch.randn(hs, hs) * 0.02 # Layer norms state_dict[f"{prefix}.input_layernorm.weight"] = torch.ones(hs, dtype=torch.float32) state_dict[f"{prefix}.post_attention_layernorm.weight"] = torch.ones(hs, dtype=torch.float32) # FFN (SwiGLU: gate + up + down) state_dict[f"{prefix}.mlp.gate_proj.weight"] = torch.randn(inter, hs) * 0.02 state_dict[f"{prefix}.mlp.up_proj.weight"] = torch.randn(inter, hs) * 0.02 state_dict[f"{prefix}.mlp.down_proj.weight"] = torch.randn(hs, inter) * 0.02 # Final norm state_dict["model.norm.weight"] = torch.ones(hs, dtype=torch.float32) # LM head state_dict["lm_head.weight"] = torch.randn(vs, hs) * 0.02 config = { "hidden_size": hs, "num_attention_heads": n_h, "num_key_value_heads": n_kv_h, "head_dim": hd, "intermediate_size": inter, "vocab_size": vs, "num_hidden_layers": num_layers, "full_attention_layers": full_attention_layers, "model_type": "qwen3", "mini": mini, } return {"state_dict": state_dict, "config": config} # ============================================================================ # SVD Decomposition for MLA Conversion # ============================================================================ def decompose_attention_svd( weight: torch.Tensor, lora_rank: int, ) -> Tuple[torch.Tensor, torch.Tensor]: """SVD decompose a weight matrix into low-rank a_proj and b_proj. Per D-10: Decompression (b_proj) matrices initialized from SVD; compression (a_proj) matrices are reinitialized with Kaiming init. Args: weight: Weight matrix of shape [in_features, out_features] or [out_features, in_features]. For Linear(in, out, bias=False), PyTorch stores weight as [out_features, in_features]. lora_rank: Target rank for the low-rank decomposition. Returns: Tuple of (a_proj, b_proj) where: - a_proj: [in_features, lora_rank] — compression (REINITIALIZED by caller) - b_proj: [lora_rank, out_features] — decompression (from SVD) """ # Ensure weight is 2D (IN-03 fix: proper ValueError instead of assert) if weight.dim() != 2: raise ValueError(f"Expected 2D weight, got {weight.dim()}D") # Determine the orientation: we want to decompose W ≈ a @ b # where a: [in_features, rank] and b: [rank, out_features] # PyTorch Linear stores as [out_features, in_features] # We decompose W.T so: W.T = U @ diag(S) @ Vh # a = U[:, :rank] @ diag(S[:rank]) shape [in_features, rank] # b = Vh[:rank, :] shape [rank, out_features] # Work in float32 for SVD stability weight_f32 = weight.float() # SVD decomposition U, S, Vh = torch.linalg.svd(weight_f32, full_matrices=False) # Truncate to target rank a_proj = U[:, :lora_rank] @ torch.diag(S[:lora_rank]) # [in_features, rank] b_proj = Vh[:lora_rank, :] # [rank, out_features] return a_proj, b_proj # ============================================================================ # MoE Expert Splitting # ============================================================================ def split_dense_to_moe( spider_state_dict: Dict[str, torch.Tensor], config: SpiderConfig, noise_scale: float = 0.02, ) -> Dict[str, torch.Tensor]: """Initialize SharedProjectionMoE expert cores and router per D-20/D-21. Per D-21: W_gate and W_transform are randomly initialized with small normal noise (std=0.02) to break symmetry. shared_up, shared_down, and shared_expert are already populated by transfer_qwen_to_spider. Args: spider_state_dict: Spider model state dict (mutated in-place) config: Spider model config noise_scale: Noise std for expert core initialization Returns: Updated state dict with SharedProjectionMoE weights """ for layer_idx in range(config.num_hidden_layers): rec_prefix = f"model.recurrent_layers.{layer_idx}.moe" # W_gate: [num_experts, hidden_size, expert_core_rank] w_gate_key = f"{rec_prefix}.W_gate" if w_gate_key not in spider_state_dict: spider_state_dict[w_gate_key] = ( torch.randn(config.num_experts, config.hidden_size, config.expert_core_rank) * noise_scale ) # W_transform: [num_experts, expert_core_rank, shared_intermediate_size] w_transform_key = f"{rec_prefix}.W_transform" if w_transform_key not in spider_state_dict: spider_state_dict[w_transform_key] = ( torch.randn(config.num_experts, config.expert_core_rank, config.shared_intermediate_size) * noise_scale ) # Router weight: [num_experts, hidden_size] router_key = f"{rec_prefix}.router.weight" if router_key not in spider_state_dict: spider_state_dict[router_key] = ( torch.randn(config.num_experts, config.hidden_size) * config.initializer_range ) # Router bias: [num_experts] router_bias_key = f"{rec_prefix}.router.bias" if router_bias_key not in spider_state_dict: spider_state_dict[router_bias_key] = torch.zeros(config.num_experts, dtype=torch.float32) return spider_state_dict # ============================================================================ # Get Spider Parameter Shapes # ============================================================================ def get_spider_param_shapes(config: SpiderConfig) -> Dict[str, Tuple[int, ...]]: """Return expected parameter shapes for the Spider model. Used for validation that all shapes match after weight transfer. """ shapes = {} # Embeddings shapes["embed_tokens.weight"] = (config.vocab_size, config.hidden_size) shapes["lm_head.weight"] = (config.vocab_size, config.hidden_size) # BoundaryPredictor: nn.Sequential(Linear(2048, 8192), GELU(), Linear(8192, 1)) shapes["boundary_predictor.0.weight"] = (config.bp_d_inner, config.hidden_size) shapes["boundary_predictor.0.bias"] = (config.bp_d_inner,) shapes["boundary_predictor.2.weight"] = (1, config.bp_d_inner) shapes["boundary_predictor.2.bias"] = (1,) # null_group for downsample shapes["null_group.weight"] = (config.hidden_size,) # down_ln for downsample shapes["down_ln.weight"] = (config.hidden_size,) shapes["down_ln.bias"] = (config.hidden_size,) head_dim = config.head_dim # 128 for section, num_layers in [ ("prelude_layers", config.prelude_layers), ("coda_layers", config.coda_layers), ]: for i in range(num_layers): prefix = f"model.{section}.{i}" # MLA attention projections shapes[f"{prefix}.self_attn.q_a_proj.weight"] = (config.q_lora_rank, config.hidden_size) shapes[f"{prefix}.self_attn.q_a_layernorm.weight"] = (config.q_lora_rank,) shapes[f"{prefix}.self_attn.q_b_proj.weight"] = (config.num_attention_heads * head_dim, config.q_lora_rank) shapes[f"{prefix}.self_attn.kv_a_proj_with_mqa.weight"] = (config.kv_lora_rank + config.qk_rope_head_dim, config.hidden_size) shapes[f"{prefix}.self_attn.kv_a_layernorm.weight"] = (config.kv_lora_rank,) shapes[f"{prefix}.self_attn.kv_b_proj.weight"] = (config.num_attention_heads * (config.qk_nope_head_dim + config.v_head_dim), config.kv_lora_rank) shapes[f"{prefix}.self_attn.o_proj.weight"] = (config.hidden_size, config.num_attention_heads * config.v_head_dim) # Layer norms shapes[f"{prefix}.input_layernorm.weight"] = (config.hidden_size,) shapes[f"{prefix}.post_attention_layernorm.weight"] = (config.hidden_size,) # FFN (dense for prelude/coda, uses SpiderExpert SwiGLU with prelude_coda_intermediate_size) dense_inter = config.prelude_coda_intermediate_size shapes[f"{prefix}.ffn.gate_proj.weight"] = (dense_inter, config.hidden_size) shapes[f"{prefix}.ffn.up_proj.weight"] = (dense_inter, config.hidden_size) shapes[f"{prefix}.ffn.down_proj.weight"] = (config.hidden_size, dense_inter) # Recurrent (MoE) layers for i in range(config.num_hidden_layers): prefix = f"model.recurrent_layers.{i}" # MLA attention shapes[f"{prefix}.self_attn.q_a_proj.weight"] = (config.q_lora_rank, config.hidden_size) shapes[f"{prefix}.self_attn.q_a_layernorm.weight"] = (config.q_lora_rank,) shapes[f"{prefix}.self_attn.q_b_proj.weight"] = (config.num_attention_heads * head_dim, config.q_lora_rank) shapes[f"{prefix}.self_attn.kv_a_proj_with_mqa.weight"] = (config.kv_lora_rank + config.qk_rope_head_dim, config.hidden_size) shapes[f"{prefix}.self_attn.kv_a_layernorm.weight"] = (config.kv_lora_rank,) shapes[f"{prefix}.self_attn.kv_b_proj.weight"] = (config.num_attention_heads * (config.qk_nope_head_dim + config.v_head_dim), config.kv_lora_rank) shapes[f"{prefix}.self_attn.o_proj.weight"] = (config.hidden_size, config.num_attention_heads * config.v_head_dim) # Layer norms shapes[f"{prefix}.input_layernorm.weight"] = (config.hidden_size,) shapes[f"{prefix}.post_attention_layernorm.weight"] = (config.hidden_size,) # MoE: SharedProjectionMoE (D-20, D-21) # shared_up: Linear(hidden, shared_inter=6144) shapes[f"{prefix}.moe.shared_up.weight"] = (config.shared_intermediate_size, config.hidden_size) # shared_down: Linear(shared_inter=6144, hidden) shapes[f"{prefix}.moe.shared_down.weight"] = (config.hidden_size, config.shared_intermediate_size) # W_gate: Parameter [num_experts, hidden, expert_core_rank] shapes[f"{prefix}.moe.W_gate"] = (config.num_experts, config.hidden_size, config.expert_core_rank) # W_transform: Parameter [num_experts, expert_core_rank, shared_inter] shapes[f"{prefix}.moe.W_transform"] = (config.num_experts, config.expert_core_rank, config.shared_intermediate_size) # shared_expert: SpiderExpert with inter=shared_expert_intermediate_size shapes[f"{prefix}.moe.shared_expert.gate_proj.weight"] = (config.shared_expert_intermediate_size, config.hidden_size) shapes[f"{prefix}.moe.shared_expert.up_proj.weight"] = (config.shared_expert_intermediate_size, config.hidden_size) shapes[f"{prefix}.moe.shared_expert.down_proj.weight"] = (config.hidden_size, config.shared_expert_intermediate_size) # Router shapes[f"{prefix}.moe.router.weight"] = (config.num_experts, config.hidden_size) shapes[f"{prefix}.moe.router.bias"] = (config.num_experts,) # LoRA adapter shapes[f"{prefix}.lora_adapter.down.weight"] = (config.lora_rank, config.hidden_size) shapes[f"{prefix}.lora_adapter.B"] = (config.lora_rank, config.hidden_size) shapes[f"{prefix}.lora_adapter.scale.weight"] = (config.max_loop_iters, config.lora_rank) # ACT halting shapes[f"{prefix}.act_halting.halt_predictor.weight"] = (1, config.hidden_size) shapes[f"{prefix}.act_halting.halt_predictor.bias"] = (1,) # Engram (layers 1 and 4 only) if i in config.engram_layers: engram_mem_dim = config.engram_heads * config.engram_dim shapes[f"{prefix}.engram.W_k.weight"] = (config.hidden_size, engram_mem_dim * 2) shapes[f"{prefix}.engram.W_v.weight"] = (config.hidden_size, engram_mem_dim * 2) shapes[f"{prefix}.engram.conv.weight"] = (config.hidden_size, 1, 4) shapes[f"{prefix}.engram.conv.bias"] = (config.hidden_size,) shapes[f"{prefix}.engram.q_norm.weight"] = (config.hidden_size,) shapes[f"{prefix}.engram.k_norm.weight"] = (config.hidden_size,) shapes[f"{prefix}.engram.embed"] = (2, config.engram_heads, config.engram_table_size, config.engram_dim) shapes[f"{prefix}.engram.hash_seeds"] = (config.engram_heads * 2,) shapes[f"{prefix}.post_engram_layernorm.weight"] = (config.hidden_size,) # LTI injection shapes["model.injection.log_A"] = (config.hidden_size,) shapes["model.injection.delta_t"] = () shapes["model.injection.B.weight"] = (config.hidden_size, config.hidden_size) # Final norm shapes["model.norm.weight"] = (config.hidden_size,) # Loop embedding dimension (config attribute, not a parameter) # shapes["model.loop_embed_dim"] = () # ACT halting for model level shapes["model.act_halting.halt_predictor.weight"] = (1, config.hidden_size) shapes["model.act_halting.halt_predictor.bias"] = (1,) return shapes # ============================================================================ # Weight Adaptation Helper # ============================================================================ def _adapt_weight(weight, target_out, target_in): """Adapt a donor weight matrix to Spider dimensions via padding/cropping. When donor hidden_size differs from Spider's (e.g., in mini test mode), we pad or crop the weight matrix to match target dimensions. Args: weight: [out_features, in_features] weight tensor from donor target_out: Target output dimension target_in: Target input dimension Returns: Adapted weight tensor of shape [target_out, target_in] """ out_dim, in_dim = weight.shape # Create target-sized tensor with Kaiming init adapted = torch.empty(target_out, target_in) nn.init.kaiming_uniform_(adapted, a=math.sqrt(5)) # Copy what fits from donor copy_out = min(out_dim, target_out) copy_in = min(in_dim, target_in) adapted[:copy_out, :copy_in] = weight[:copy_out, :copy_in] return adapted # ============================================================================ # Main Transfer Function # ============================================================================ def transfer_qwen_to_spider( donor_state_dict: Dict[str, torch.Tensor], donor_config: Dict, spider_config: SpiderConfig, noise_scale: float = 0.02, ) -> Dict: """Transfer weights from Qwen3.5-2B donor to Spider-FLEXITOKENS architecture. Per D-09: Qwen3.5-2B is the weight donor. Per D-10: SVD decomposition converts standard GQA attention to MLA format. Transfer rules: - o_proj [2048, 2048]: direct copy from donor - q_proj → SVD → q_b_proj (q_a_proj reinitialized with Kaiming) - k_proj + v_proj → SVD → kv_b_proj (kv_a_proj reinitialized with Kaiming) - Layer norms [2048]: direct copy - Embeddings: REINIT [272, 2048] (byte-level) - BoundaryPredictor: REINIT (no pre-trained source) - FFN: REINIT (intermediate_size mismatch 6144 vs 1024) - LoRA, ACT, LTI: REINIT (Spider-specific modules) Args: donor_state_dict: Qwen3.5-2B state dict donor_config: Donor model config dict spider_config: Spider model config noise_scale: Noise scale for MoE expert perturbation Returns: Dict with "spider_state_dict", "transfer_coverage", "layer_mapping" """ hidden_size = spider_config.hidden_size q_lora_rank = spider_config.q_lora_rank kv_lora_rank = spider_config.kv_lora_rank num_heads = spider_config.num_attention_heads head_dim = spider_config.head_dim qk_nope_head_dim = spider_config.qk_nope_head_dim qk_rope_head_dim = spider_config.qk_rope_head_dim v_head_dim = spider_config.v_head_dim # Donor dimensions (may differ from Spider in mini/test mode) donor_hidden_size = donor_config.get("hidden_size", hidden_size) donor_num_heads = donor_config.get("num_attention_heads", 8) donor_num_kv_heads = donor_config.get("num_key_value_heads", 2) donor_head_dim = donor_config.get("head_dim", 256) donor_intermediate_size = donor_config.get("intermediate_size", 6144) # Track parameter counts for coverage report donor_param_count = 0 reinit_param_count = 0 donor_params = set() # keys that came from donor reinit_params = set() # keys that were reinitialized spider_sd = {} # Determine layer mapping from donor to Spider full_attention_layers = donor_config.get("full_attention_layers", []) num_donor_layers = donor_config.get("num_hidden_layers", 24) # Map donor layers to Spider sections: # prelude: 2 layers, recurrent: 6 layers, coda: 2 layers = 10 total # Use full_attention layers preferentially available_fa = list(full_attention_layers) # Build layer mapping: spider_layer_idx → donor_layer_idx layer_mapping = {} required_layers = ( spider_config.prelude_layers + spider_config.num_hidden_layers + spider_config.coda_layers ) # Fill from full_attention layers first, then fallback to any layer donor_pool = list(available_fa) if len(donor_pool) < required_layers: # Add remaining layers (including linear_attention) for norms all_layers = list(range(num_donor_layers)) for l in all_layers: if l not in donor_pool: donor_pool.append(l) for i in range(required_layers): if i < len(donor_pool): layer_mapping[i] = donor_pool[i] else: layer_mapping[i] = None # No donor layer available def _kaiming_init(shape): """Kaiming uniform initialization for new parameters.""" tensor = torch.empty(shape) nn.init.kaiming_uniform_(tensor, a=math.sqrt(5)) return tensor def _zeros_init(shape): """Zero initialization.""" return torch.zeros(shape, dtype=torch.float32) # IN-02: explicit dtype def _ones_init(shape): """Ones initialization for layer norm weights.""" return torch.ones(shape, dtype=torch.float32) # ---- 1. Embeddings: REINIT for byte-level vocab ---- embed_weight = _kaiming_init((spider_config.vocab_size, hidden_size)) spider_sd["embed_tokens.weight"] = embed_weight reinit_param_count += embed_weight.numel() reinit_params.add("embed_tokens.weight") lm_head_weight = _kaiming_init((spider_config.vocab_size, hidden_size)) spider_sd["lm_head.weight"] = lm_head_weight reinit_param_count += lm_head_weight.numel() reinit_params.add("lm_head.weight") # ---- 2. BoundaryPredictor: REINIT (no pre-trained source) ---- bp_0_weight = _kaiming_init((spider_config.bp_d_inner, hidden_size)) bp_0_bias = _zeros_init((spider_config.bp_d_inner,)) bp_2_weight = _kaiming_init((1, spider_config.bp_d_inner)) bp_2_bias = _zeros_init((1,)) spider_sd["boundary_predictor.0.weight"] = bp_0_weight spider_sd["boundary_predictor.0.bias"] = bp_0_bias spider_sd["boundary_predictor.2.weight"] = bp_2_weight spider_sd["boundary_predictor.2.bias"] = bp_2_bias reinit_param_count += bp_0_weight.numel() + bp_0_bias.numel() reinit_param_count += bp_2_weight.numel() + bp_2_bias.numel() reinit_params.add("boundary_predictor.0.weight") reinit_params.add("boundary_predictor.2.weight") # ---- 3. null_group and down_ln for downsample/upsample ---- null_group = _zeros_init((hidden_size,)) spider_sd["null_group.weight"] = null_group reinit_param_count += null_group.numel() reinit_params.add("null_group.weight") down_ln_w = torch.ones(hidden_size, dtype=torch.float32) down_ln_b = _zeros_init((hidden_size,)) spider_sd["down_ln.weight"] = down_ln_w spider_sd["down_ln.bias"] = down_ln_b reinit_param_count += down_ln_w.numel() + down_ln_b.numel() reinit_params.add("down_ln.weight") # ---- 4. Layer-by-layer weight transfer ---- for section_name, num_layers in [ ("prelude_layers", spider_config.prelude_layers), ("recurrent_layers", spider_config.num_hidden_layers), ("coda_layers", spider_config.coda_layers), ]: is_recurrent = section_name == "recurrent_layers" for layer_idx in range(num_layers): # WR-02 fix: accumulate spider_layer_idx across sections so # coda layers map to distinct donor layers instead of reusing # prelude donor layers spider_layer_idx = ({ "prelude_layers": 0, "recurrent_layers": spider_config.prelude_layers, "coda_layers": spider_config.prelude_layers + spider_config.num_hidden_layers, }[section_name] + layer_idx) donor_layer_idx = layer_mapping.get(spider_layer_idx) prefix = f"model.{section_name}.{layer_idx}" if donor_layer_idx is not None: donor_prefix = f"model.layers.{donor_layer_idx}" else: donor_prefix = None # ---- Attention: MLA via SVD ---- # q_proj: [num_heads_donor * head_dim_donor, hidden_size_donor] if donor_prefix is not None: donor_q_key = f"{donor_prefix}.self_attn.q_proj.weight" donor_q = donor_state_dict.get(donor_q_key) else: donor_q = None if donor_q is not None and donor_q.shape[0] == donor_num_heads * donor_head_dim and donor_q.shape[1] == donor_hidden_size: # SVD decompose q_proj → q_b_proj # donor_q shape: [out, in] — PyTorch Linear stores [out, in] # We want: q_a_proj weight [q_lora_rank, hidden_size] and # q_b_proj weight [num_heads * head_dim, q_lora_rank] # SVD decompose: donor_q.T = [in, out] → a=[in,rank], b=[rank,out] # a_svd: [in, rank], b_svd: [rank, out] # q_a_proj.weight = a_svd.T = [rank, in] → matches nn.Linear(hidden, q_lora_rank) # q_b_proj.weight = b_svd.T = [out, rank] → matches nn.Linear(q_lora_rank, num_heads*head_dim) # When donor_hidden_size != hidden_size, we adapt the SVD decomposition effective_q = donor_q if donor_hidden_size != hidden_size: # Pad/crop donor_q to match Spider dimensions effective_q = _adapt_weight(donor_q, donor_num_heads * donor_head_dim, hidden_size) q_a_svd, q_b_svd = decompose_attention_svd(effective_q, q_lora_rank) # Per D-10: q_a_proj is REINITIALIZED, q_b_proj from SVD q_a_proj = _kaiming_init((q_lora_rank, hidden_size)) q_b_proj = q_b_svd.T # [out, rank] — transposed for PyTorch Linear [out, in] donor_param_count += q_b_proj.numel() reinit_param_count += q_a_proj.numel() donor_params.add(f"{prefix}.self_attn.q_b_proj.weight") reinit_params.add(f"{prefix}.self_attn.q_a_proj.weight") else: q_a_proj = _kaiming_init((q_lora_rank, hidden_size)) q_b_proj = _kaiming_init((num_heads * head_dim, q_lora_rank)) reinit_param_count += q_a_proj.numel() + q_b_proj.numel() reinit_params.add(f"{prefix}.self_attn.q_a_proj.weight") reinit_params.add(f"{prefix}.self_attn.q_b_proj.weight") spider_sd[f"{prefix}.self_attn.q_a_proj.weight"] = q_a_proj spider_sd[f"{prefix}.self_attn.q_b_proj.weight"] = q_b_proj # q_a_layernorm q_a_ln = torch.ones(q_lora_rank, dtype=torch.float32) spider_sd[f"{prefix}.self_attn.q_a_layernorm.weight"] = q_a_ln reinit_param_count += q_a_ln.numel() reinit_params.add(f"{prefix}.self_attn.q_a_layernorm.weight") # k_proj + v_proj → SVD → kv_a_proj_with_mqa, kv_b_proj if donor_prefix is not None: donor_k_key = f"{donor_prefix}.self_attn.k_proj.weight" donor_v_key = f"{donor_prefix}.self_attn.v_proj.weight" donor_k = donor_state_dict.get(donor_k_key) donor_v = donor_state_dict.get(donor_v_key) else: donor_k = None donor_v = None if donor_k is not None and donor_v is not None: # Concatenate k_proj and v_proj along output dim # donor_k: [num_kv_heads * head_dim_donor, hidden_size_donor] # donor_v: [num_kv_heads * head_dim_donor, hidden_size_donor] # Combined: [num_kv_heads * head_dim_donor * 2, hidden_size_donor] combined_kv = torch.cat([donor_k, donor_v], dim=0) # Adapt dimensions if donor hidden_size differs from Spider's if donor_hidden_size != hidden_size: combined_kv_out = donor_num_kv_heads * donor_head_dim * 2 combined_kv = _adapt_weight(combined_kv, combined_kv_out, hidden_size) # Transpose for SVD: we want [hidden_size, combined_kv_out] kv_a_svd, kv_b_svd = decompose_attention_svd(combined_kv.T, kv_lora_rank) # kv_a_svd: [hidden_size, rank], kv_b_svd: [rank, combined_kv_out] # Per D-10: kv_a_proj (compression) REINITIALIZED # kv_b_proj (decompression) from SVD # kv_a_proj_with_mqa.weight: [kv_lora_rank + qk_rope_head_dim, hidden_size] # = [128 + 64, 2048] = [192, 2048] kv_a_with_mqa = _kaiming_init( (kv_lora_rank + qk_rope_head_dim, hidden_size) ) # kv_b_proj.weight: [num_heads*(qk_nope+v_head), kv_lora_rank] # = [16*(64+64), 128] = [2048, 128] # SVD gives kv_b_svd: [128, 1024] → transpose: [1024, 128] # This is smaller than [2048, 128], so pad with Kaiming init kv_b_proj_weight = _kaiming_init( (num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank) ) # [2048, 128] svd_contribution = kv_b_svd.T # [1024, 128] # Copy SVD result into the beginning of kv_b_proj_weight rows_to_copy = min(svd_contribution.shape[0], kv_b_proj_weight.shape[0]) kv_b_proj_weight[:rows_to_copy, :] = svd_contribution[:rows_to_copy] # Count: SVD-initialized rows count as donor, padding as reinit donor_rows = rows_to_copy reinit_rows = kv_b_proj_weight.shape[0] - donor_rows donor_param_count += donor_rows * kv_b_proj_weight.shape[1] reinit_param_count += reinit_rows * kv_b_proj_weight.shape[1] reinit_param_count += kv_a_with_mqa.numel() donor_params.add(f"{prefix}.self_attn.kv_b_proj.weight") reinit_params.add(f"{prefix}.self_attn.kv_a_proj_with_mqa.weight") else: kv_a_with_mqa = _kaiming_init( (kv_lora_rank + qk_rope_head_dim, hidden_size) ) kv_b_proj_weight = _kaiming_init( (num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank) ) reinit_param_count += kv_a_with_mqa.numel() + kv_b_proj_weight.numel() reinit_params.add(f"{prefix}.self_attn.kv_a_proj_with_mqa.weight") reinit_params.add(f"{prefix}.self_attn.kv_b_proj.weight") spider_sd[f"{prefix}.self_attn.kv_a_proj_with_mqa.weight"] = kv_a_with_mqa spider_sd[f"{prefix}.self_attn.kv_b_proj.weight"] = kv_b_proj_weight # kv_a_layernorm kv_a_ln = torch.ones(kv_lora_rank, dtype=torch.float32) spider_sd[f"{prefix}.self_attn.kv_a_layernorm.weight"] = kv_a_ln reinit_param_count += kv_a_ln.numel() reinit_params.add(f"{prefix}.self_attn.kv_a_layernorm.weight") # o_proj: copy from donor where possible # Donor o_proj: [donor_hidden_size, donor_hidden_size] # Spider o_proj: [hidden_size, num_heads * v_head_dim] if donor_prefix is not None: donor_o_key = f"{donor_prefix}.self_attn.o_proj.weight" donor_o = donor_state_dict.get(donor_o_key) else: donor_o = None o_proj_shape = (hidden_size, num_heads * v_head_dim) # [2048, 1024] o_proj = _kaiming_init(o_proj_shape) if donor_o is not None: # Copy what fits from donor's o_proj rows_to_copy = min(donor_o.shape[0], o_proj.shape[0]) cols_to_copy = min(donor_o.shape[1], o_proj.shape[1]) o_proj[:rows_to_copy, :cols_to_copy] = donor_o[:rows_to_copy, :cols_to_copy] donor_param_count += rows_to_copy * cols_to_copy remaining = o_proj.numel() - rows_to_copy * cols_to_copy if remaining > 0: reinit_param_count += remaining donor_params.add(f"{prefix}.self_attn.o_proj.weight") else: reinit_param_count += o_proj.numel() reinit_params.add(f"{prefix}.self_attn.o_proj.weight") spider_sd[f"{prefix}.self_attn.o_proj.weight"] = o_proj # Layer norms: direct copy where shapes match, adapt otherwise for norm_name in ["input_layernorm.weight", "post_attention_layernorm.weight"]: if donor_prefix is not None: donor_norm_key = f"{donor_prefix}.{norm_name}" donor_norm = donor_state_dict.get(donor_norm_key) else: donor_norm = None if donor_norm is not None and donor_norm.shape == (hidden_size,): spider_sd[f"{prefix}.{norm_name}"] = donor_norm.clone() donor_param_count += donor_norm.numel() donor_params.add(f"{prefix}.{norm_name}") elif donor_norm is not None and donor_norm.shape[0] != hidden_size: # Adapt: pad/crop layer norm to match Spider hidden_size adapted_norm = torch.ones(hidden_size, dtype=torch.float32) copy_size = min(donor_norm.shape[0], hidden_size) adapted_norm[:copy_size] = donor_norm[:copy_size] spider_sd[f"{prefix}.{norm_name}"] = adapted_norm donor_param_count += copy_size reinit_param_count += hidden_size - copy_size donor_params.add(f"{prefix}.{norm_name}") else: ln = torch.ones(hidden_size, dtype=torch.float32) spider_sd[f"{prefix}.{norm_name}"] = ln reinit_param_count += ln.numel() reinit_params.add(f"{prefix}.{norm_name}") # ---- FFN / MoE ---- if is_recurrent: # SharedProjectionMoE (D-20, D-21): # shared_up: Linear(hidden, shared_inter=6144) — DIRECT copy from donor up_proj # shared_down: Linear(shared_inter=6144, hidden) — DIRECT copy from donor down_proj # shared_expert: SpiderExpert with inter=7424 — partial copy from donor FFN # W_gate, W_transform: random init — created by split_dense_to_moe # Stride mapping: Spider layer i → Qwen layer i*4 (layers 0,4,8,12,16,20) # for 6 recurrent layers out of 24 Qwen layers if donor_layer_idx is not None: qwen_layer_for_ffn = donor_layer_idx else: qwen_layer_for_ffn = None # ---- shared_up: direct copy from donor up_proj ---- # Spider shared_up.weight: [shared_inter=6144, hidden=2048] # Qwen up_proj.weight: [inter=6144, hidden=2048] — EXACT MATCH (D-32) shared_up_key = f"{prefix}.moe.shared_up.weight" shared_up_shape = (spider_config.shared_intermediate_size, hidden_size) if qwen_layer_for_ffn is not None: donor_up_key = f"model.layers.{qwen_layer_for_ffn}.mlp.up_proj.weight" donor_up = donor_state_dict.get(donor_up_key) else: donor_up = None if donor_up is not None and donor_up.shape == shared_up_shape: spider_sd[shared_up_key] = donor_up.clone().float() donor_param_count += donor_up.numel() donor_params.add(shared_up_key) elif donor_up is not None: shared_up_w = _kaiming_init(shared_up_shape) rows_copy = min(donor_up.shape[0], shared_up_shape[0]) cols_copy = min(donor_up.shape[1], shared_up_shape[1]) shared_up_w[:rows_copy, :cols_copy] = donor_up[:rows_copy, :cols_copy].float() spider_sd[shared_up_key] = shared_up_w donor_param_count += rows_copy * cols_copy reinit_param_count += shared_up_w.numel() - rows_copy * cols_copy donor_params.add(shared_up_key) else: spider_sd[shared_up_key] = _kaiming_init(shared_up_shape) reinit_param_count += shared_up_shape[0] * shared_up_shape[1] reinit_params.add(shared_up_key) # ---- shared_down: direct copy from donor down_proj ---- # Spider shared_down.weight: [hidden=2048, shared_inter=6144] # Qwen down_proj.weight: [hidden=2048, inter=6144] — EXACT MATCH (D-32) shared_down_key = f"{prefix}.moe.shared_down.weight" shared_down_shape = (hidden_size, spider_config.shared_intermediate_size) if qwen_layer_for_ffn is not None: donor_down_key = f"model.layers.{qwen_layer_for_ffn}.mlp.down_proj.weight" donor_down = donor_state_dict.get(donor_down_key) else: donor_down = None if donor_down is not None and donor_down.shape == shared_down_shape: spider_sd[shared_down_key] = donor_down.clone().float() donor_param_count += donor_down.numel() donor_params.add(shared_down_key) elif donor_down is not None: shared_down_w = _kaiming_init(shared_down_shape) rows_copy = min(donor_down.shape[0], shared_down_shape[0]) cols_copy = min(donor_down.shape[1], shared_down_shape[1]) shared_down_w[:rows_copy, :cols_copy] = donor_down[:rows_copy, :cols_copy].float() spider_sd[shared_down_key] = shared_down_w donor_param_count += rows_copy * cols_copy reinit_param_count += shared_down_w.numel() - rows_copy * cols_copy donor_params.add(shared_down_key) else: spider_sd[shared_down_key] = _kaiming_init(shared_down_shape) reinit_param_count += shared_down_shape[0] * shared_down_shape[1] reinit_params.add(shared_down_key) # ---- shared_expert: partial transfer from donor FFN (6144→7424) ---- # Spider shared_expert has inter=7424 (D-21: larger than donor's 6144) # First 6144 rows/cols from donor, remaining 1280 randomly initialized shared_expert_inter = spider_config.shared_expert_intermediate_size if qwen_layer_for_ffn is not None: donor_gate_key = f"model.layers.{qwen_layer_for_ffn}.mlp.gate_proj.weight" donor_se_up_key = f"model.layers.{qwen_layer_for_ffn}.mlp.up_proj.weight" donor_se_down_key = f"model.layers.{qwen_layer_for_ffn}.mlp.down_proj.weight" donor_se_gate = donor_state_dict.get(donor_gate_key) donor_se_up = donor_state_dict.get(donor_se_up_key) donor_se_down = donor_state_dict.get(donor_se_down_key) else: donor_se_gate = donor_se_up = donor_se_down = None for proj_name, spider_shape in [ ("gate_proj", (shared_expert_inter, hidden_size)), ("up_proj", (shared_expert_inter, hidden_size)), ("down_proj", (hidden_size, shared_expert_inter)), ]: key = f"{prefix}.moe.shared_expert.{proj_name}.weight" w = _kaiming_init(spider_shape) if proj_name in ("gate_proj", "up_proj"): donor_src = donor_se_gate if proj_name == "gate_proj" else donor_se_up if donor_src is not None: rows_copy = min(donor_src.shape[0], spider_shape[0]) cols_copy = min(donor_src.shape[1], spider_shape[1]) w[:rows_copy, :cols_copy] = donor_src[:rows_copy, :cols_copy].float() donor_param_count += rows_copy * cols_copy reinit_param_count += w.numel() - rows_copy * cols_copy donor_params.add(key) else: reinit_param_count += w.numel() reinit_params.add(key) else: # down_proj: [hidden, shared_expert_inter] if donor_se_down is not None: rows_copy = min(donor_se_down.shape[0], spider_shape[0]) cols_copy = min(donor_se_down.shape[1], spider_shape[1]) w[:rows_copy, :cols_copy] = donor_se_down[:rows_copy, :cols_copy].float() donor_param_count += rows_copy * cols_copy reinit_param_count += w.numel() - rows_copy * cols_copy donor_params.add(key) else: reinit_param_count += w.numel() reinit_params.add(key) spider_sd[key] = w # W_gate and W_transform will be created by split_dense_to_moe # LoRA adapter lora_down = _kaiming_init((spider_config.lora_rank, hidden_size)) lora_B = torch.zeros(spider_config.lora_rank, hidden_size, dtype=torch.float32) lora_scale = torch.zeros(spider_config.max_loop_iters, spider_config.lora_rank, dtype=torch.float32) spider_sd[f"{prefix}.lora_adapter.down.weight"] = lora_down spider_sd[f"{prefix}.lora_adapter.B"] = lora_B spider_sd[f"{prefix}.lora_adapter.scale.weight"] = lora_scale reinit_param_count += lora_down.numel() + lora_B.numel() + lora_scale.numel() reinit_params.add(f"{prefix}.lora_adapter.down.weight") # ACT halting halt_w = _kaiming_init((1, hidden_size)) halt_b = _zeros_init((1,)) spider_sd[f"{prefix}.act_halting.halt_predictor.weight"] = halt_w spider_sd[f"{prefix}.act_halting.halt_predictor.bias"] = halt_b reinit_param_count += halt_w.numel() + halt_b.numel() reinit_params.add(f"{prefix}.act_halting.halt_predictor.weight") # Engram (layers 1 and 4 only — D-20 revision) if layer_idx in spider_config.engram_layers: engram_mem_dim = spider_config.engram_heads * spider_config.engram_dim engram_W_k = _kaiming_init((hidden_size, engram_mem_dim * 2)) engram_W_v = _kaiming_init((hidden_size, engram_mem_dim * 2)) engram_conv_w = _kaiming_init((hidden_size, 1, 4)) engram_conv_b = _zeros_init((hidden_size,)) engram_q_norm = _ones_init((hidden_size,)) engram_k_norm = _ones_init((hidden_size,)) engram_embed = torch.zeros( 2, spider_config.engram_heads, spider_config.engram_table_size, spider_config.engram_dim ) engram_hash = torch.arange(spider_config.engram_heads * 2, dtype=torch.float32) post_engram_norm = _ones_init((hidden_size,)) spider_sd[f"{prefix}.engram.W_k.weight"] = engram_W_k spider_sd[f"{prefix}.engram.W_v.weight"] = engram_W_v spider_sd[f"{prefix}.engram.conv.weight"] = engram_conv_w spider_sd[f"{prefix}.engram.conv.bias"] = engram_conv_b spider_sd[f"{prefix}.engram.q_norm.weight"] = engram_q_norm spider_sd[f"{prefix}.engram.k_norm.weight"] = engram_k_norm spider_sd[f"{prefix}.engram.embed"] = engram_embed spider_sd[f"{prefix}.engram.hash_seeds"] = engram_hash spider_sd[f"{prefix}.post_engram_layernorm.weight"] = post_engram_norm engram_params = (engram_W_k.numel() + engram_W_v.numel() + engram_conv_w.numel() + engram_conv_b.numel() + engram_q_norm.numel() + engram_k_norm.numel() + engram_embed.numel() + engram_hash.numel() + post_engram_norm.numel()) reinit_param_count += engram_params else: # Dense FFN for prelude/coda: partial transfer from donor FFN # Spider uses prelude_coda_intermediate_size=4096 (D-21) # Donor has intermediate_size=6144 → copy first 4096 rows/cols dense_inter = spider_config.prelude_coda_intermediate_size if donor_layer_idx is not None: donor_gate_key = f"model.layers.{donor_layer_idx}.mlp.gate_proj.weight" donor_up_key = f"model.layers.{donor_layer_idx}.mlp.up_proj.weight" donor_down_key = f"model.layers.{donor_layer_idx}.mlp.down_proj.weight" donor_d_gate = donor_state_dict.get(donor_gate_key) donor_d_up = donor_state_dict.get(donor_up_key) donor_d_down = donor_state_dict.get(donor_down_key) else: donor_d_gate = donor_d_up = donor_d_down = None for proj_name, shape, donor_src in [ ("gate_proj", (dense_inter, hidden_size), donor_d_gate), ("up_proj", (dense_inter, hidden_size), donor_d_up), ("down_proj", (hidden_size, dense_inter), donor_d_down), ]: w = _kaiming_init(shape) key = f"{prefix}.ffn.{proj_name}.weight" if donor_src is not None: if proj_name in ("gate_proj", "up_proj"): rows_copy = min(donor_src.shape[0], shape[0]) cols_copy = min(donor_src.shape[1], shape[1]) w[:rows_copy, :cols_copy] = donor_src[:rows_copy, :cols_copy].float() else: rows_copy = min(donor_src.shape[0], shape[0]) cols_copy = min(donor_src.shape[1], shape[1]) w[:rows_copy, :cols_copy] = donor_src[:rows_copy, :cols_copy].float() donor_param_count += rows_copy * cols_copy reinit_param_count += w.numel() - rows_copy * cols_copy donor_params.add(key) else: reinit_param_count += w.numel() reinit_params.add(key) spider_sd[key] = w # ---- 5. LTI Injection: REINIT (Spider-specific) ---- log_A = torch.full((hidden_size,), -2.0) delta_t = torch.tensor(1.0) B_weight = torch.randn(hidden_size, hidden_size) * 0.01 spider_sd["model.injection.log_A"] = log_A spider_sd["model.injection.delta_t"] = delta_t spider_sd["model.injection.B.weight"] = B_weight reinit_param_count += log_A.numel() + delta_t.numel() + B_weight.numel() reinit_params.add("model.injection.B.weight") # ---- 6. Final norm: try to copy from donor, adapt dimensions ---- donor_final_norm = donor_state_dict.get("model.norm.weight") if donor_final_norm is not None and donor_final_norm.shape == (hidden_size,): spider_sd["model.norm.weight"] = donor_final_norm.clone() donor_param_count += donor_final_norm.numel() donor_params.add("model.norm.weight") elif donor_final_norm is not None: # Adapt: pad/crop to match Spider hidden_size adapted_norm = torch.ones(hidden_size, dtype=torch.float32) copy_size = min(donor_final_norm.shape[0], hidden_size) adapted_norm[:copy_size] = donor_final_norm[:copy_size] spider_sd["model.norm.weight"] = adapted_norm donor_param_count += copy_size reinit_param_count += hidden_size - copy_size donor_params.add("model.norm.weight") else: spider_sd["model.norm.weight"] = torch.ones(hidden_size, dtype=torch.float32) reinit_param_count += hidden_size reinit_params.add("model.norm.weight") # ---- 7. Model-level ACT halting: REINIT ---- halt_w = _kaiming_init((1, hidden_size)) halt_b = _zeros_init((1,)) spider_sd["model.act_halting.halt_predictor.weight"] = halt_w spider_sd["model.act_halting.halt_predictor.bias"] = halt_b reinit_param_count += halt_w.numel() + halt_b.numel() # ---- 8. Apply MoE expert splitting ---- spider_sd = split_dense_to_moe(spider_sd, spider_config, noise_scale=noise_scale) # Count SharedProjectionMoE params created by split_dense_to_moe for layer_idx in range(spider_config.num_hidden_layers): rec_prefix = f"model.recurrent_layers.{layer_idx}.moe" # W_gate and W_transform are random init for core_key in [f"{rec_prefix}.W_gate", f"{rec_prefix}.W_transform"]: if core_key in spider_sd and core_key not in reinit_params and core_key not in donor_params: reinit_param_count += spider_sd[core_key].numel() reinit_params.add(core_key) # Router for router_key in [f"{rec_prefix}.router.weight", f"{rec_prefix}.router.bias"]: if router_key in spider_sd and router_key not in reinit_params and router_key not in donor_params: reinit_param_count += spider_sd[router_key].numel() reinit_params.add(router_key) # ---- 9. Compute transfer coverage ---- total_params = donor_param_count + reinit_param_count if total_params > 0: donor_pct = (donor_param_count / total_params) * 100.0 reinit_pct = (reinit_param_count / total_params) * 100.0 else: donor_pct = 0.0 reinit_pct = 0.0 transfer_coverage = { "donor_params": donor_param_count, "reinit_params": reinit_param_count, "total_params": total_params, "donor_pct": round(donor_pct, 2), "reinit_pct": round(reinit_pct, 2), "donor_keys": sorted(donor_params), "reinit_keys": sorted(reinit_params), } # Print report print("=" * 60) print("Weight Transfer Report") print("=" * 60) print(f" Donor: Qwen3.5-2B ({donor_config.get('num_hidden_layers', '?')} layers)") print(f" Target: Spider-FLEXITOKENS ({spider_config.prelude_layers}+{spider_config.num_hidden_layers}+{spider_config.coda_layers} layers)") print(f" Full attention layers used: {len(full_attention_layers)}") print(f" Layer mapping: {layer_mapping}") print() print(f" Total params: {total_params:>12,} ({total_params/1e6:.1f}M)") print(f" Donor-originated: {donor_param_count:>12,} ({donor_param_count/1e6:.1f}M) = {donor_pct:.1f}%") print(f" Reinitialized: {reinit_param_count:>12,} ({reinit_param_count/1e6:.1f}M) = {reinit_pct:.1f}%") print() print(f" Transfer coverage: {donor_pct:.1f}% from donor, {reinit_pct:.1f}% reinitialized") print("=" * 60) return { "spider_state_dict": spider_sd, "transfer_coverage": transfer_coverage, "layer_mapping": layer_mapping, } # ============================================================================ # SpiderMoEModel — Multimodal Forward Pass (D-11, 02-03) # ============================================================================ class SpiderMoEModel(nn.Module): """Spider-FLEXITOKENS model with multimodal forward pass. Implements the full forward pass wiring per D-11: - Text bytes → embed → prelude layers → BoundaryPredictor → downsample → recurrent core → upsample → coda layers → lm_head → logits - Modality tokens (vision/audio/video) are injected at sentinel-marked positions and bypass the BoundaryPredictor entirely. - Sentinel-gated passthrough: modality_mask forces boundary=1.0 at sentinel+modality positions, preventing cross-modality merges. This is a simplified model that implements the forward pass logic without the full SpiderPortalMLA attention (which requires position IDs, KV cache, etc.). It uses simple linear projections to demonstrate the multimodal wiring and parameter budget. """ def __init__(self, config: SpiderConfig): super().__init__() self.config = config # Embeddings: 272 vocab (256 bytes + 16 specials) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) # LM head (not tied per D-06) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # BoundaryPredictor self.boundary_predictor = BoundaryPredictor(config) # null_group for downsample self.null_group = nn.Parameter(torch.zeros(config.hidden_size, dtype=torch.float32)) # IN-02 # Downsample layer norm self.down_ln = nn.LayerNorm(config.hidden_size) # Prelude layers (2 dense layers with simplified attention + FFN) self.prelude_layers = nn.ModuleList([ self._make_dense_layer(config) for _ in range(config.prelude_layers) ]) # Recurrent layers (6 MoE layers with simplified attention + MoE) self.recurrent_layers = nn.ModuleList([ self._make_moe_layer(config, i) for i in range(config.num_hidden_layers) ]) # Coda layers (2 dense layers with simplified attention + FFN) self.coda_layers = nn.ModuleList([ self._make_dense_layer(config) for _ in range(config.coda_layers) ]) # Final norm self.norm = nn.LayerNorm(config.hidden_size) # LTI injection self.injection = _LTIInjection(config) # ACT halting self.act_halting = _ACTHalting(config) # LoRA adapter (per recurrent layer) self.lora_adapters = nn.ModuleList([ _LoRAAdapter(config) for _ in range(config.num_hidden_layers) ]) self.loop_embed_dim = config.loop_embed_dim self.max_loop_iters = config.max_loop_iters def _make_dense_layer(self, config): """Create a simplified dense layer (prelude/coda).""" return _DenseLayer(config) def _make_moe_layer(self, config, layer_idx): """Create a simplified MoE layer (recurrent).""" return _MoELayer(config, layer_idx) def _inject_modality_features( self, hidden_states: torch.Tensor, input_ids: torch.Tensor, features: list, modality: str = 'IMG', ) -> torch.Tensor: """Replace placeholder embeddings with actual encoder features at modality regions. Per D-11: Modality tokens (vision, audio, video) are injected at sentinel-marked positions in the hidden_states sequence. The caller constructs input_ids with sentinel tokens (e.g., IMG_START, IMG_END) marking modality regions. Between sentinel pairs, the initial embeddings are placeholders — this method replaces them with the actual encoder features. T-02-06 mitigation: Validates feature shape and sentinel pair count. Args: hidden_states: [B, L, D] hidden states after embedding. input_ids: [B, L] token IDs with sentinel markers. features: List of tensors, one per sentinel pair per batch item. Each tensor has shape [num_tokens, hidden_size]. modality: Modality type prefix ('IMG', 'AUD', 'VID'). Returns: hidden_states with modality features injected at sentinel regions. Raises: ValueError: If feature shape doesn't match [num_tokens, hidden_size] or sentinel pair count doesn't match feature count. """ start_token = SENTINEL_TOKENS[f'{modality}_START'] end_token = SENTINEL_TOKENS[f'{modality}_END'] for b in range(hidden_states.shape[0]): starts = (input_ids[b] == start_token).nonzero(as_tuple=True)[0] ends = (input_ids[b] == end_token).nonzero(as_tuple=True)[0] if len(starts) != len(ends): raise ValueError( f"Batch {b}: mismatched {modality} sentinel pairs — " f"{len(starts)} {_TOKEN_NAMES_BY_ID[start_token]}(s) vs " f"{len(ends)} {_TOKEN_NAMES_BY_ID[end_token]}(s)." ) if len(starts) != len(features): raise ValueError( f"Batch {b}: {modality} sentinel pair count ({len(starts)}) " f"doesn't match feature count ({len(features)})." ) for s, e, feat in zip(starts, ends, features): # T-02-06: Validate feature shape num_tokens = e - s - 1 # tokens between sentinels if feat.shape[0] != num_tokens: raise ValueError( f"Batch {b}: {modality} feature has {feat.shape[0]} tokens " f"but sentinel region has {num_tokens} positions " f"(from pos {s+1} to {e-1})." ) if feat.shape[1] != hidden_states.shape[-1]: raise ValueError( f"Batch {b}: {modality} feature hidden_size {feat.shape[1]} " f"doesn't match model hidden_size {hidden_states.shape[-1]}." ) # Replace placeholder embeddings with actual features hidden_states[b, s + 1:e] = feat.to(hidden_states.dtype) return hidden_states def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_values: Optional[list] = None, inputs_embeds: Optional[torch.Tensor] = None, vision_features: Optional[list] = None, audio_features: Optional[list] = None, video_features: Optional[list] = None, **kwargs, ) -> torch.Tensor: """Forward pass with multimodal sentinel-gated passthrough. Per D-11: - All positions go through embed_tokens (bytes get byte embeddings, sentinels get special embeddings, modality tokens get placeholder embeddings) - External encoder features are injected at sentinel-marked positions - BoundaryPredictor operates on the embedded sequence with modality_mask - Text bytes go through BP → downsample → recurrent → upsample → coda → logits - Modality tokens bypass BP and enter downsampled sequence at sentinel positions Args: input_ids: [B, L] token IDs with optional sentinel markers. attention_mask: Optional attention mask (not used in simplified model). position_ids: Optional position IDs (not used in simplified model). past_key_values: Optional KV cache (not used in simplified model). inputs_embeds: Optional pre-computed embeddings. vision_features: Optional list of tensors, each [num_tokens, hidden_size]. audio_features: Optional list of tensors, each [num_tokens, hidden_size]. video_features: Optional list of tensors, each [num_tokens, hidden_size]. Returns: logits: [B, L, vocab_size] output logits. """ B, L = input_ids.shape # 1. Embed all tokens (bytes, sentinels, modality placeholders) if inputs_embeds is not None: hidden_states = inputs_embeds else: hidden_states = self.embed_tokens(input_ids) # [B, L, D] # 2. Inject external modality features at sentinel positions if vision_features is not None: hidden_states = self._inject_modality_features( hidden_states, input_ids, vision_features, 'IMG' ) if audio_features is not None: hidden_states = self._inject_modality_features( hidden_states, input_ids, audio_features, 'AUD' ) if video_features is not None: hidden_states = self._inject_modality_features( hidden_states, input_ids, video_features, 'VID' ) # 3. Prelude layers for layer in self.prelude_layers: hidden_states = layer(hidden_states) # 4. Boundary prediction with modality mask modality_mask = create_modality_mask(input_ids) # [B, L] soft_boundaries, hard_boundaries = self.boundary_predictor( hidden_states, modality_mask=modality_mask ) # 5. Downsample with boundaries # Apply layer norm before downsample hidden_states_normed = self.down_ln(hidden_states) null_group = self.null_group.unsqueeze(0).unsqueeze(0).expand(1, B, -1) shortened = downsample(hard_boundaries, hidden_states_normed, null_group) # shortened: [S, B, D] # 6. Recurrent core with RDT looping # Convert shortened from SBD to BLD for recurrent layers hidden_states = shortened.permute(1, 0, 2) # [B, S, D] n_loops = self.max_loop_iters input_embedding = hidden_states.clone() for t in range(n_loops): # Loop index embedding loop_emb = _loop_index_embedding(hidden_states, t, self.loop_embed_dim) if t > 0: # LTI injection injection = self.injection(hidden_states, input_embedding) hidden_states = hidden_states + injection # Recurrent layers for i, layer in enumerate(self.recurrent_layers): # LoRA adaptation for this loop iteration lora_out = self.lora_adapters[i](hidden_states, t) hidden_states = layer(hidden_states + lora_out * 0.01) # 7. Upsample back to original sequence length # Convert back to SBD for upsample hidden_states_sbd = hidden_states.permute(1, 0, 2) # [S, B, D] hidden_states = upsample(hard_boundaries, hidden_states_sbd) # [B, L, D] # 8. Coda layers for layer in self.coda_layers: hidden_states = layer(hidden_states) # 9. Final norm + LM head hidden_states = self.norm(hidden_states) logits = self.lm_head(hidden_states) # [B, L, vocab_size] return logits # ============================================================================ # Simplified sub-modules for SpiderMoEModel # ============================================================================ class _DenseLayer(nn.Module): """Simplified dense layer for prelude/coda (attention + FFN).""" def __init__(self, config: SpiderConfig): super().__init__() self.input_layernorm = nn.LayerNorm(config.hidden_size) self.post_attention_layernorm = nn.LayerNorm(config.hidden_size) # Simplified self-attention (single-head for parameter efficiency demo) self.self_attn = nn.MultiheadAttention( config.hidden_size, num_heads=4, batch_first=True ) # FFN with SwiGLU-like structure self.ffn = _SwiGLUFFN(config.hidden_size, config.prelude_coda_intermediate_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Self-attention with residual residual = hidden_states hidden_states = self.input_layernorm(hidden_states) attn_out, _ = self.self_attn( hidden_states, hidden_states, hidden_states ) hidden_states = residual + attn_out # FFN with residual residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) ffn_out = self.ffn(hidden_states) hidden_states = residual + ffn_out return hidden_states class _MoELayer(nn.Module): """Simplified MoE layer for recurrent core.""" def __init__(self, config: SpiderConfig, layer_idx: int = 0): super().__init__() self.input_layernorm = nn.LayerNorm(config.hidden_size) self.post_attention_layernorm = nn.LayerNorm(config.hidden_size) # Simplified self-attention self.self_attn = nn.MultiheadAttention( config.hidden_size, num_heads=4, batch_first=True ) # MoE FFN (SharedProjectionMoE per D-20, D-21) self.moe = _SharedProjectionMoE(config) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Self-attention with residual residual = hidden_states hidden_states = self.input_layernorm(hidden_states) attn_out, _ = self.self_attn( hidden_states, hidden_states, hidden_states ) hidden_states = residual + attn_out # MoE FFN with residual residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) moe_out, _z_loss = self.moe(hidden_states) hidden_states = residual + moe_out return hidden_states class _SwiGLUFFN(nn.Module): """SwiGLU FFN: gate_proj, up_proj, down_proj.""" def __init__(self, hidden_size: int, intermediate_size: int): 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: torch.Tensor) -> torch.Tensor: return self.down_proj(nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) class _SharedProjectionMoE(nn.Module): """SharedProjectionMoE matching spider.py architecture (D-20, D-21). Shared up/down projections computed once per token, rank-256 expert cores specialize on the shared representation. """ def __init__(self, config: SpiderConfig): super().__init__() self.num_experts = config.num_experts self.num_experts_per_tok = config.num_experts_per_tok self.shared_inter = config.shared_intermediate_size self.expert_core_rank = config.expert_core_rank self.hidden_size = config.hidden_size self.shared_up = nn.Linear(config.hidden_size, config.shared_intermediate_size, bias=False) self.shared_down = nn.Linear(config.shared_intermediate_size, config.hidden_size, bias=False) self.W_gate = nn.Parameter( torch.randn(config.num_experts, config.hidden_size, config.expert_core_rank) * 0.02 ) self.W_transform = nn.Parameter( torch.randn(config.num_experts, config.expert_core_rank, config.shared_intermediate_size) * 0.02 ) self.shared_expert = _SwiGLUFFN(config.hidden_size, config.shared_expert_intermediate_size) self.router = nn.Linear(config.hidden_size, config.num_experts, bias=True) self.router.bias = nn.Parameter(torch.zeros(config.num_experts, dtype=torch.float32)) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: B, L, D = x.shape shared_hidden = nn.functional.silu(self.shared_up(x)) shared_out = self.shared_expert(x) router_logits = self.router(x) router_probs = nn.functional.softmax(router_logits, dim=-1) top2_probs, top2_indices = router_probs.topk(self.num_experts_per_tok, dim=-1) top2_probs = top2_probs / top2_probs.sum(dim=-1, keepdim=True) x_flat = x.reshape(B * L, D) shared_hidden_flat = shared_hidden.reshape(B * L, self.shared_inter) routed_out = torch.zeros(B * L, D, device=x.device, dtype=x.dtype) for k in range(self.num_experts_per_tok): expert_indices = top2_indices[:, :, k].reshape(B * L) expert_weights = top2_probs[:, :, k].reshape(B * L) for e in range(self.num_experts): mask = (expert_indices == e) if not mask.any(): continue expert_input = x_flat[mask] expert_sh = shared_hidden_flat[mask] gate = expert_input @ self.W_gate[e] core = gate @ self.W_transform[e] expert_output = self.shared_down(core * expert_sh) routed_out[mask] += expert_weights[mask].unsqueeze(-1) * expert_output routed_out = routed_out.reshape(B, L, D) z_loss = (router_logits.logsumexp(dim=-1) ** 2).mean() return shared_out + routed_out, z_loss class _LTIInjection(nn.Module): """Linear Time-Invariant injection module.""" def __init__(self, config: SpiderConfig): super().__init__() self.log_A = nn.Parameter(torch.full((config.hidden_size,), -2.0)) self.delta_t = nn.Parameter(torch.tensor(1.0)) self.B_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False) def forward(self, h_t: torch.Tensor, e: torch.Tensor) -> torch.Tensor: A = torch.exp(self.log_A) decay = A * self.delta_t B_e = self.B_proj(e) return decay.unsqueeze(0).unsqueeze(0) * B_e class _ACTHalting(nn.Module): """Adaptive Computation Time halting module.""" def __init__(self, config: SpiderConfig): super().__init__() self.halt_predictor = nn.Linear(config.hidden_size, 1, bias=True) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return torch.sigmoid(self.halt_predictor(hidden_states)).squeeze(-1) class _LoRAAdapter(nn.Module): """LoRA adapter for per-loop adaptation in recurrent layers.""" def __init__(self, config: SpiderConfig): super().__init__() self.down = nn.Linear(config.hidden_size, config.lora_rank, bias=False) self.up = nn.Linear(config.lora_rank, config.hidden_size, bias=False) # CR-01 fix: zero init the up-projection per LoRA convention nn.init.zeros_(self.up.weight) self.scale_embeddings = nn.Embedding(config.max_loop_iters, config.lora_rank) def forward(self, x: torch.Tensor, loop_iter: int) -> torch.Tensor: down = self.down(x) scale = self.scale_embeddings(torch.tensor([loop_iter], device=x.device)) scaled = down * scale.squeeze(0) return self.up(scaled) def _loop_index_embedding( hidden_states: torch.Tensor, loop_iter: int, embed_dim: int, ) -> torch.Tensor: """Generate sinusoidal loop index embedding. Provides positional-like encoding for the loop iteration index, allowing the model to differentiate between iterations of the recurrent depth loop. """ B, L, D = hidden_states.shape device = hidden_states.device # Sinusoidal embedding for loop iteration pos = torch.tensor([loop_iter], device=device, dtype=hidden_states.dtype) dim = torch.arange(embed_dim, device=device, dtype=hidden_states.dtype) freq = pos / (10000 ** (2 * dim / embed_dim)) # Interleave sin and cos emb = torch.zeros(embed_dim, device=device, dtype=hidden_states.dtype) emb[0::2] = torch.sin(freq[::2][:emb[0::2].shape[0]]) emb[1::2] = torch.cos(freq[1::2][:emb[1::2].shape[0]]) # Broadcast to [B, L, embed_dim] and pad to D if needed emb = emb.unsqueeze(0).unsqueeze(0).expand(B, L, -1) if embed_dim < D: padding = torch.zeros(B, L, D - embed_dim, device=device, dtype=hidden_states.dtype) emb = torch.cat([emb, padding], dim=-1) elif embed_dim > D: emb = emb[:, :, :D] return emb # ============================================================================ # Save & Config Export # ============================================================================ def save_spider_model( spider_state_dict: Dict[str, torch.Tensor], config: SpiderConfig, output_dir: Path, ): """Save Spider model state dict and config to output directory. Handles weight tying per safetensors pattern from init_spiderportal.py. """ output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) # Handle weight tying: safetensors refuses shared tensors save_sd = {} for name, param in spider_state_dict.items(): # Ensure tensor is contiguous (required by safetensors) # Transposes (.T) and slices can produce non-contiguous tensors save_sd[name] = param.contiguous() if config.tie_word_embeddings and "lm_head.weight" in save_sd: del save_sd["lm_head.weight"] print(" Note: lm_head.weight tied to embed_tokens.weight (saved once)") # Save as safetensors try: from safetensors.torch import save_file save_file(save_sd, output_dir / "model.safetensors") except ImportError: # Fallback to PyTorch save torch.save(save_sd, output_dir / "model.pt") print(" Warning: safetensors not available, saved as model.pt") # Save config cfg_dict = { "architectures": ["SpiderForConditionalGeneration"], "model_type": config.model_type, "vocab_size": config.vocab_size, "hidden_size": config.hidden_size, "num_hidden_layers": config.num_hidden_layers, "num_attention_heads": config.num_attention_heads, "num_key_value_heads": config.num_key_value_heads, "intermediate_size": config.intermediate_size, "hidden_act": config.hidden_act, "max_position_embeddings": config.max_position_embeddings, "rope_theta": config.rope_theta, "rope_scaling": config.rope_scaling, "sliding_window": config.sliding_window, "rms_norm_eps": config.rms_norm_eps, "initializer_range": config.initializer_range, "tie_word_embeddings": config.tie_word_embeddings, "torch_dtype": config.torch_dtype, # MoE "num_experts": config.num_experts, "num_experts_per_tok": config.num_experts_per_tok, "num_shared_experts": config.num_shared_experts, "router_aux_loss_coef": config.router_aux_loss_coef, "shared_intermediate_size": config.shared_intermediate_size, "expert_core_rank": config.expert_core_rank, "shared_expert_intermediate_size": config.shared_expert_intermediate_size, "prelude_coda_intermediate_size": config.prelude_coda_intermediate_size, # MLA "kv_lora_rank": config.kv_lora_rank, "q_lora_rank": config.q_lora_rank, "qk_rope_head_dim": config.qk_rope_head_dim, "qk_nope_head_dim": config.qk_nope_head_dim, "v_head_dim": config.v_head_dim, # RDT "max_loop_iters": config.max_loop_iters, "act_threshold": config.act_threshold, "prelude_layers": config.prelude_layers, "coda_layers": config.coda_layers, "lora_rank": config.lora_rank, # BoundaryPredictor "bp_d_inner": config.bp_d_inner, # Multimodal "vision_hidden_size": config.vision_hidden_size, "audio_hidden_size": config.audio_hidden_size, "vision_num_frames": config.vision_num_frames, "vision_tokens_per_frame": config.vision_tokens_per_frame, "vision_temporal_tokens": config.vision_temporal_tokens, "vision_temporal_layers": config.vision_temporal_layers, } with open(output_dir / "config.json", "w") as f: json.dump(cfg_dict, f, indent=2) # Compute SHA256 of model file for integrity check (T-02-03 mitigation) model_file = output_dir / "model.safetensors" if not model_file.exists(): model_file = output_dir / "model.pt" if model_file.exists(): sha256 = hashlib.sha256() with open(model_file, "rb") as f: for chunk in iter(lambda: f.read(8192), b""): sha256.update(chunk) print(f" Model SHA256: {sha256.hexdigest()[:16]}...") with open(output_dir / "model.sha256", "w") as f: f.write(sha256.hexdigest()) print(f" Saved to {output_dir}") if model_file.exists(): print(f" Model file size: {model_file.stat().st_size / 1e6:.1f} MB") # ============================================================================ # CLI Entry Point # ============================================================================ def main(): parser = argparse.ArgumentParser( description="Transfer weights from Qwen3.5-2B to Spider-FLEXITOKENS" ) parser.add_argument( "--donor", type=str, default="Qwen/Qwen3.5-2B", help="HuggingFace model ID or local path for donor model" ) parser.add_argument( "--output", type=str, default="models/Spider-FLEXITOKENS-init/", help="Output directory for Spider model" ) parser.add_argument( "--config", type=str, default="spider_flexitokens_997m", help="Spider model configuration name" ) parser.add_argument( "--noise-scale", type=float, default=0.02, help="Noise scale for MoE expert perturbation" ) parser.add_argument( "--dry-run", action="store_true", help="Run with dummy donor (no download required)" ) args = parser.parse_args() # Select config config_map = { "spider_flexitokens_997m": spider_flexitokens_997m(), } spider_config = config_map.get(args.config, spider_flexitokens_997m()) if args.dry_run: print("DRY RUN: Using dummy donor (no download)") donor = create_dummy_donor(num_layers=10, full_attention_layers=list(range(10))) donor_sd = donor["state_dict"] donor_cfg = donor["config"] else: # Load actual Qwen3.5-2B from HuggingFace print(f"Loading donor model: {args.donor}") try: from transformers import AutoModelForCausalLM, AutoConfig donor_model = AutoModelForCausalLM.from_pretrained( args.donor, torch_dtype=torch.bfloat16, device_map="cpu" ) donor_cfg_obj = AutoConfig.from_pretrained(args.donor) # Extract full_attention layers from Qwen3.5-2B config # Qwen3.5-2B has hybrid attention: some full, some linear full_attention_layers = getattr( donor_cfg_obj, "full_attention_layers", None ) if full_attention_layers is None: # Fallback: assume layers with attention_type == "full" # Qwen3.5-2B: 18 linear + 6 full attention in 24 layers full_attention_layers = [] for i in range(donor_cfg_obj.num_hidden_layers): layer_cfg = getattr(donor_cfg_obj, f"layer_{i}", None) if layer_cfg and getattr(layer_cfg, "attention_type", "full") == "full": full_attention_layers.append(i) if not full_attention_layers: # If no layer-level info, use known pattern for Qwen3.5-2B full_attention_layers = [3, 7, 11, 15, 19, 23] donor_sd = donor_model.state_dict() donor_cfg = { "hidden_size": donor_cfg_obj.hidden_size, "num_attention_heads": donor_cfg_obj.num_attention_heads, "num_key_value_heads": getattr(donor_cfg_obj, "num_key_value_heads", 2), "head_dim": getattr(donor_cfg_obj, "head_dim", donor_cfg_obj.hidden_size // donor_cfg_obj.num_attention_heads), "intermediate_size": donor_cfg_obj.intermediate_size, "vocab_size": donor_cfg_obj.vocab_size, "num_hidden_layers": donor_cfg_obj.num_hidden_layers, "full_attention_layers": full_attention_layers, "model_type": getattr(donor_cfg_obj, "model_type", "qwen3"), } except ImportError: print("Error: transformers library required for loading donor model.") print("Install with: pip install transformers") sys.exit(1) except Exception as e: print(f"Error loading donor model: {e}") print("Use --dry-run for testing without download.") sys.exit(1) # Run transfer result = transfer_qwen_to_spider( donor_state_dict=donor_sd, donor_config=donor_cfg, spider_config=spider_config, noise_scale=args.noise_scale, ) # Save save_spider_model( spider_state_dict=result["spider_state_dict"], config=spider_config, output_dir=Path(args.output), ) print("\nWeight transfer complete!") if __name__ == "__main__": main()