Update modeling_omega.py

modeling_omega.py

@@ -8,22 +8,22 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 from .configuration_omega import OmegaConfig
 
 # ============================================================
-# 1. QUANTIZATION UTILITIES
+# 1. QUANTIZATION UTILITIES (Bit-Unpacking)
 # ============================================================
 
 class OmegaQuantLinear(nn.Module):
     """
-    Decompresses 4-bit packed Int32 weights into BFloat16 on-the-fly.
-    Layout: 8 weights per Int32, packed via (w << (i*4)).
+    On-the-fly dequantization kernel.
+    Unpacks 8x4-bit weights from Int32 containers into BFloat16 registers.
     """
     def __init__(self, in_features, out_features, bias=False):
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
         
-        # Binary containers
-        self.register_buffer('qweight', torch.zeros((in_features, out_features // 8), dtype=torch.int32))
-        self.register_buffer('scales', torch.zeros((in_features, out_features // 128), dtype=torch.bfloat16))
+        # Binary containers as defined in the Forge process
+        self.register_buffer('qweight', torch.zeros((out_features, in_features // 8), dtype=torch.int32))
+        self.register_buffer('scales', torch.zeros((out_features, in_features // 128), dtype=torch.bfloat16))
         
         if bias:
             self.register_buffer('bias', torch.zeros((out_features), dtype=torch.bfloat16))
@@ -32,33 +32,30 @@ class OmegaQuantLinear(nn.Module):
 
     def _unpack(self):
         """
-        Efficient GPU-side bit-unpacking logic.
-        Maps [0, 15] back to [-8, 7] and applies scales.
+        JIT Unpacking: Int32 -> 8x 4-bit -> BFloat16.
+        Memory Boundary: Only materializes weights for the current layer execution.
         """
         device = self.qweight.device
-        # 1. Expand Int32 to 8 interleaved 4-bit slots
-        # shape: [In, Out/8, 8]
-        unpacked = torch.zeros((self.in_features, self.out_features // 8, 8), 
+        # 1. De-interleave bits
+        # Resulting shape: [Out, In/8, 8]
+        unpacked = torch.zeros((self.out_features, self.in_features // 8, 8), 
                                dtype=torch.int32, device=device)
         
         for i in range(8):
             unpacked[..., i] = (self.qweight >> (i * 4)) & 0xF
             
-        # 2. Reshape to original weight dimensions
-        # q_weight: [In, Out]
-        q_weight = unpacked.view(self.in_features, self.out_features)
+        # 2. Reshape to [Out, In] and apply offset (Forge used +8)
+        q_weight = unpacked.view(self.out_features, self.in_features)
+        fp_weight = (q_weight.to(torch.bfloat16) - 8.0)
         
-        # 3. Apply Offset (+8 was added during packing) and Cast
-        # Apply group-wise scaling (group_size=128)
-        weight = (q_weight.to(torch.bfloat16) - 8.0)
-        
-        # scales shape: [In, Out/128]. Upsample to [In, Out]
+        # 3. Apply scales (Block size 128)
         s = self.scales.repeat_interleave(128, dim=1)
-        return (weight * s).T # Return [Out, In] for F.linear
+        return fp_weight * s
 
     def forward(self, x):
-        weight = self._unpack()
-        return F.linear(x, weight, self.bias)
+        # We materialise the BF16 weight locally within the function scope
+        # It is garbage collected immediately after the linear operation
+        return F.linear(x, self._unpack(), self.bias)
 
 # ============================================================
 # 2. CORE UTILITIES
@@ -93,14 +90,13 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
     return q_embed, k_embed
 
 # ============================================================
-# 3. COMPONENTS (UPGRADED TO QUANTLINEAR)
+# 3. COMPONENTS
 # ============================================================
 
 class OmegaAttention(nn.Module):
     def __init__(self, config, layer_idx):
         super().__init__()
         self.config = config
-        # All Projections now use QuantLinear
         self.q_a_proj = OmegaQuantLinear(config.hidden_size, config.q_lora_rank, bias=False)
         self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
         self.q_b_proj = OmegaQuantLinear(config.q_lora_rank, config.num_attention_heads * (config.qk_nope_head_dim + config.qk_rope_head_dim), bias=False)
@@ -110,13 +106,13 @@ class OmegaAttention(nn.Module):
         self.o_proj = OmegaQuantLinear(config.num_attention_heads * config.v_head_dim, config.hidden_size, bias=False)
         self.rotary_emb = OmegaRotaryEmbedding(config.qk_rope_head_dim)
 
-    def forward(self, x, attention_mask=None, position_ids=None, **kwargs):
-        bsz, q_len, _ = x.size()
-        q_a = self.q_a_layernorm(self.q_a_proj(x))
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, **kwargs):
+        bsz, q_len, _ = hidden_states.size()
+        q_a = self.q_a_layernorm(self.q_a_proj(hidden_states))
         q = self.q_b_proj(q_a).view(bsz, q_len, self.config.num_attention_heads, -1).transpose(1, 2)
         q_nope, q_pe = q.split([self.config.qk_nope_head_dim, self.config.qk_rope_head_dim], dim=-1)
         
-        kv_a_raw = self.kv_a_proj_with_mqa(x)
+        kv_a_raw = self.kv_a_proj_with_mqa(hidden_states)
         kv_a, k_pe = kv_a_raw.split([self.config.kv_lora_rank, self.config.qk_rope_head_dim], dim=-1)
         k_pe = k_pe.view(bsz, q_len, 1, -1).transpose(1, 2)
         
@@ -131,7 +127,7 @@ class OmegaAttention(nn.Module):
         
         attn = (q @ k.transpose(-1, -2)) * (q.shape[-1]**-0.5)
         if attention_mask is not None: attn += attention_mask
-        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(x.dtype)
+        attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(hidden_states.dtype)
         out = (attn @ v).transpose(1, 2).reshape(bsz, q_len, -1)
         return self.o_proj(out), None, None
 
@@ -141,9 +137,7 @@ class OmegaConsensusMLP(nn.Module):
         self.config = config
         self.layer_idx = layer_idx
         
-        # Dynamic handling for Layer 0 (Dense) vs Layer 1+ (Omega)
-        # Note: If Layer 0 was patched to Dense during Forge, 
-        # it will be caught by the standard gate_proj naming.
+        # Layer 0 is Dense (18432 wide). Subsequent layers are Council-style.
         if layer_idx == 0:
             self.gate_proj = OmegaQuantLinear(config.hidden_size, 18432, bias=False)
             self.up_proj = OmegaQuantLinear(config.hidden_size, 18432, bias=False)
@@ -155,17 +149,18 @@ class OmegaConsensusMLP(nn.Module):
                 "down_proj": OmegaQuantLinear(config.intermediate_size, config.hidden_size, bias=False)
             })
             self.wide_gate_up = OmegaQuantLinear(config.hidden_size, config.expert_intermediate * config.num_experts_per_layer * 2, bias=False)
-            
-            # 3D parameter dequantization (packed as 2D for safetensors compatibility)
-            # Expert down weights are treated as a specialized 2D block
             self.expert_down_weights_packed = OmegaQuantLinear(config.num_experts_per_layer * config.expert_intermediate, config.hidden_size, bias=False)
             self.register_buffer("importance", torch.ones(config.num_experts_per_layer))
 
     def forward(self, x):
-        if hasattr(self, 'gate_proj'): # Standard Dense Path
-            return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-        
-        # Omega Consensus Path
+        # Path 1: Dense SwiGLU (Layer 0)
+        if self.layer_idx == 0:
+            # SwiGLU: down(silu(gate(x)) * up(x))
+            gate = self.gate_proj(x)
+            up = self.up_proj(x)
+            return self.down_proj(F.silu(gate) * up)
+        
+        # Path 2: Omega Consensus (Layer 1-60)
         c_gate = self.chairman.gate_proj(x)
         c_up = self.chairman.up_proj(x)
         shared_out = self.chairman.down_proj(F.silu(c_gate) * c_up)
@@ -174,12 +169,11 @@ class OmegaConsensusMLP(nn.Module):
         gate, up = wide_out.chunk(2, dim=-1)
         lane_acts = F.silu(gate) * up
         
-        # Dequantize 3D block
-        # We unpack the weights and reshape to [E, I, H]
+        # Unpack 3D slab: [E*I, H] -> [E, I, H]
         e_weights = self.expert_down_weights_packed._unpack().view(self.config.num_experts_per_layer, self.config.expert_intermediate, -1)
         expert_outs = torch.einsum('bsei,eih->bseh', lane_acts, e_weights)
         
-        # Agreement Governance (FP32 precision for stability)
+        # FP32 Agreement Metric for governance stability
         c_norm = torch.norm(shared_out, dim=-1, keepdim=True) + 1e-6
         agreement = torch.einsum('bseh,bsh->bse', expert_outs.float(), shared_out.float()) / c_norm.float()
         mask = torch.maximum(F.relu(agreement / self.config.tau).to(x.dtype), 
@@ -189,7 +183,7 @@ class OmegaConsensusMLP(nn.Module):
         return shared_out + council_out
 
 # ============================================================
-# 4. WRAPPER LAYERS
+# 4. TOP LEVEL
 # ============================================================
 
 class OmegaDecoderLayer(nn.Module):
@@ -207,7 +201,6 @@ class OmegaForCausalLM(PreTrainedModel):
     config_class = OmegaConfig
     def __init__(self, config):
         super().__init__(config)
-        # Keep sensitive Bookends in BF16
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList([OmegaDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
         self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)