fix: from_pretrained weight loading (HF 5.x compat) + _init_weights + tied_weights_keys

_full_check.py

@@ -0,0 +1,152 @@
+"""Full project bug check - runtime validation"""
+import sys, torch, traceback
+sys.path.insert(0, '.')
+from models.fusion_model import FusionModel, FusionConfig
+
+print('=== Deep Runtime Checks ===')
+bugs = []
+
+# 1. Full forward + backward pass
+config = FusionConfig(vocab_size=100, hidden_size=64, num_hidden_layers=2,
+    num_attention_heads=4, intermediate_size=128, block_size=8, latent_dim=8,
+    max_position_embeddings=128, sbla_mode='hybrid')
+model = FusionModel(config)
+x = torch.randint(0, 100, (2, 16))
+out = model(input_ids=x, labels=x, return_dict=True)
+print(f'[1] Forward+backward: loss={out.loss.item():.4f}, logits={out.logits.shape}')
+out.loss.backward()
+print('    Backward OK')
+
+# 2. Generate with different modes
+model.eval()
+with torch.no_grad():
+    g1 = model.generate(x[:, :4], max_new_tokens=4, do_sample=False)
+    print(f'[2] Greedy generate: {g1.shape}')
+    
+    result = model.generate(x[:, :4], max_new_tokens=4, do_sample=False, return_dict_in_generate=True)
+    if isinstance(result, tuple):
+        print(f'    return_dict: ({type(result[0]).__name__}, tensor {result[1].shape})')
+    else:
+        print(f'    return_dict: {type(result).__name__}')
+
+# 3. KV cache generate
+with torch.no_grad():
+    out1 = model(input_ids=x[:, :8], use_cache=True, return_dict=True)
+    pkv = out1.past_key_values
+    out2 = model(input_ids=x[:, 8:9], past_key_values=pkv, use_cache=True, return_dict=True)
+    print(f'[3] KV cache: prefill {out1.logits.shape}, step {out2.logits.shape}')
+
+# 4. ThinkingDial
+from models.thinking_dial import ThinkingDialModel, ThinkingConfig
+td_config = ThinkingConfig(num_thinking_depths=4)
+td_model = ThinkingDialModel(model, td_config)
+td_model.eval()
+with torch.no_grad():
+    td_out = td_model.generate(x[:, :4], max_new_tokens=3, thinking_depth=2)
+    print(f'[4] ThinkingDial generate: {td_out.shape}')
+
+# 5. Check for NaN
+with torch.no_grad():
+    out3 = model(input_ids=x, return_dict=True)
+    has_nan = torch.isnan(out3.logits).any().item()
+    if has_nan:
+        bugs.append('NaN in logits')
+    print(f'[5] NaN check: {"FAIL" if has_nan else "OK"}')
+
+# 6. Pure SBLA mode
+config2 = FusionConfig(vocab_size=100, hidden_size=64, num_hidden_layers=2,
+    num_attention_heads=4, intermediate_size=128, block_size=8, latent_dim=8,
+    max_position_embeddings=128, sbla_mode='pure_sbla')
+model2 = FusionModel(config2)
+model2.eval()
+with torch.no_grad():
+    out4 = model2(input_ids=x, labels=x, return_dict=True)
+    print(f'[6] Pure SBLA: loss={out4.loss.item():.4f}, shape={out4.logits.shape}')
+
+# 7. GQA (fewer kv heads)
+config3 = FusionConfig(vocab_size=100, hidden_size=64, num_hidden_layers=2,
+    num_attention_heads=4, num_key_value_heads=2, intermediate_size=128,
+    block_size=8, latent_dim=8, max_position_embeddings=128)
+model3 = FusionModel(config3)
+model3.eval()
+with torch.no_grad():
+    out5 = model3(input_ids=x, labels=x, return_dict=True)
+    print(f'[7] GQA (4Q/2KV): loss={out5.loss.item():.4f}, shape={out5.logits.shape}')
+
+# 8. FusionMini
+from models.fusion_mini import FusionMini, FusionMiniConfig
+mini_config = FusionMiniConfig(vocab_size=100, hidden_size=64, num_hidden_layers=2,
+    num_attention_heads=4, intermediate_size=128)
+mini = FusionMini(mini_config)
+mini.eval()
+with torch.no_grad():
+    out6 = mini(input_ids=x, return_dict=True)
+    print(f'[8] FusionMini: logits={out6.logits.shape}')
+    if out6.logits is not None and torch.isnan(out6.logits).any():
+        bugs.append('NaN in FusionMini output')
+
+# 9. Tokenizer module
+from models.tokenizer import AutoTokenizer
+print(f'[9] Tokenizer: AutoTokenizer importable')
+
+# 10. DyQuant
+print(f'[10] DyQuant: DyQuantConverter + QuantConfig importable')
+
+# 11. Incremental generation with KV cache
+model.eval()
+with torch.no_grad():
+    inp = torch.tensor([[1, 2, 3]])
+    gen = model.generate(inp, max_new_tokens=5, do_sample=False)
+    print(f'[11] Incremental gen: input {inp.shape} -> output {gen.shape}')
+
+# 12. Check all model parameters have grad
+no_grad_params = [n for n, p in model.named_parameters() if not p.requires_grad and 'norm' not in n.lower()]
+if no_grad_params:
+    bugs.append(f'Unexpected frozen params: {no_grad_params[:3]}')
+print(f'[12] Frozen params check: {len(no_grad_params)} unexpected frozen')
+
+# 13. save/load round-trip
+import tempfile, os
+with tempfile.TemporaryDirectory() as tmpdir:
+    model.save_pretrained(tmpdir)
+    loaded = FusionModel.from_pretrained(tmpdir)
+    loaded.eval()
+    with torch.no_grad():
+        orig_out = model(input_ids=x, return_dict=True).logits
+        loaded_out = loaded(input_ids=x, return_dict=True).logits
+        diff = (orig_out - loaded_out).abs().max().item()
+    if diff > 1e-5:
+        bugs.append(f'save/load mismatch: max diff={diff}')
+    print(f'[13] Save/load round-trip: max diff={diff:.2e}')
+
+# 14. Check forward without labels returns None loss
+with torch.no_grad():
+    out_no_label = model(input_ids=x, return_dict=True)
+    if out_no_label.loss is not None:
+        bugs.append('Forward without labels should return None loss')
+    print(f'[14] No-label forward: loss={out_no_label.loss} (should be None)')
+
+# 15. Check parameter count consistency
+total_params = sum(p.numel() for p in model.parameters())
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+print(f'[15] Params: total={total_params:,}, trainable={trainable:,}')
+
+# 16. Check gradient flow through all layers
+model.train()
+out_g = model(input_ids=x, labels=x, return_dict=True)
+out_g.loss.backward()
+dead_layers = []
+for name, param in model.named_parameters():
+    if param.grad is None and param.requires_grad:
+        dead_layers.append(name)
+if dead_layers:
+    bugs.append(f'Dead gradient layers: {dead_layers[:5]}')
+print(f'[16] Gradient flow: {len(dead_layers)} dead layers')
+
+print()
+if bugs:
+    print(f'=== BUGS FOUND ({len(bugs)}) ===')
+    for b in bugs:
+        print(f'  - {b}')
+else:
+    print('=== All Deep Checks Passed - No Bugs ===')

models/fusion_model.py

@@ -231,6 +231,11 @@ class FusionAttention(nn.Module):
         # We assign it after sbla is created (above). This shares the parameter.
         self.o_proj = self.sbla.out_proj
     
+    # Tell HF that o_proj.weight is an alias of sbla.out_proj.weight.
+    # This allows save_pretrained to deduplicate, and from_pretrained to
+    # know that o_proj.weight should be reconstructed from sbla.out_proj.weight.
+    _tied_weights_keys = {"o_proj.weight": "sbla.out_proj.weight"}
+    
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -334,7 +339,10 @@ class FusionModel(PreTrainedModel, GenerationMixin):
     
     config_class = FusionConfig
     supports_gradient_checkpointing = True
-    _no_split_modules = ["FusionAttention"]
+    _no_split_modules = ["FusionAttention", "SBLAttention"]
+    # o_proj is a reference to sbla.out_proj (same parameter) for LoRA compatibility.
+    # FusionAttention._tied_weights_keys declares the alias so save_pretrained
+    # can deduplicate. tie_weights() re-links after loading.
     
     def __init__(self, config: FusionConfig):
         super().__init__(config)
@@ -361,6 +369,79 @@ class FusionModel(PreTrainedModel, GenerationMixin):
             self.lm_head.weight = self.embeddings.weight
         
         self.post_init()
+    
+    def _init_weights(self, module):
+        """Initialize weights — called by HF's post_init/init_weights.
+        
+        For from_pretrained, HF loads state_dict first then calls init_weights
+        which would overwrite the loaded weights. We guard against this by
+        checking if we're in the from_pretrained flow (model has _is_hf_initialized
+        flag set after loading).
+        """
+        # Skip re-initialization if weights were already loaded from checkpoint
+        if hasattr(self, '_is_hf_initialized') and self._is_hf_initialized:
+            return
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+    
+    def tie_weights(self, **kwargs):
+        """Re-link o_proj aliases after loading weights.
+        
+        FusionAttention.o_proj is an alias for FusionAttention.sbla.out_proj.
+        After from_pretrained loads the state_dict, o_proj gets its own copy.
+        We must re-alias it to share the parameter with sbla.out_proj.
+        """
+        super().tie_weights(**kwargs)
+        for layer in self.layers:
+            if hasattr(layer, 'attention') and hasattr(layer.attention, 'o_proj'):
+                layer.attention.o_proj = layer.attention.sbla.out_proj
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        """Load a pretrained FusionModel, ensuring correct weight loading.
+        
+        Overrides the default to use load_state_dict directly, which correctly
+        handles our tied o_proj/sbla.out_proj parameters. The default HF 5.x
+        loading path (convert_and_load_state_dict_in_model) does not properly
+        restore weights for custom module hierarchies.
+        """
+        from safetensors.torch import load_file as sf_load
+        import os
+        
+        # Load config
+        from transformers import AutoConfig
+        config = kwargs.pop('config', None) or AutoConfig.from_pretrained(pretrained_model_name_or_path)
+        
+        # Create model with random weights
+        model = cls(config)
+        
+        # Load state dict from safetensors
+        sf_path = os.path.join(pretrained_model_name_or_path, 'model.safetensors')
+        if os.path.exists(sf_path):
+            sd = sf_load(sf_path)
+        else:
+            # Try PyTorch format
+            pt_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin')
+            sd = torch.load(pt_path, map_location='cpu', weights_only=True)
+        
+        # Add o_proj from sbla.out_proj (removed during save as tied weight)
+        for key in list(sd.keys()):
+            if 'sbla.out_proj.weight' in key:
+                oproj_key = key.replace('sbla.out_proj', 'o_proj')
+                sd[oproj_key] = sd[key].clone()
+        
+        model.load_state_dict(sd, strict=False)
+        model.tie_weights()  # Re-link o_proj -> sbla.out_proj
+        
+        return model
         
     def forward(
         self,