Create qwen2to3_diagnostic.py

qwen2to3_diagnostic.py

@@ -0,0 +1,222 @@
+# file: qwen2to3_diagnostic.py
+
+import torch
+import os
+import json
+import re
+from datetime import datetime
+from tqdm import tqdm
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
+from transformers import Qwen3Config, Qwen3ForCausalLM
+from collections import Counter
+
+# --- DIAGNOSTIC HELPERS ---
+
+def log_tensor_stats(tensor, name):
+    """Prints statistics for a given tensor."""
+    if tensor.numel() == 0:
+        print(f"  - DIAGNOSTIC STATS for '{name}': Tensor is empty.")
+        return
+    print(
+        f"  - DIAGNOSTIC STATS for '{name}':\n"
+        f"    - Shape: {tensor.shape}, Dtype: {tensor.dtype}\n"
+        f"    - Mean: {tensor.float().mean().item():.4f}, Std: {tensor.float().std().item():.4f}\n"
+        f"    - Min: {tensor.float().min().item():.4f}, Max: {tensor.float().max().item():.4f}\n"
+        f"    - Has NaN: {torch.isnan(tensor).any().item()}, Has Inf: {torch.isinf(tensor).any().item()}"
+    )
+
+def verify_embedding_transfer(s_embeds, t_embeds, mapping, t_tokenizer, num_samples=5):
+    """Verifies that some shared token embeddings were copied correctly."""
+    print("\n  - DIAGNOSTIC: Verifying embedding transfer for sample tokens...")
+    verified_count = 0
+    for t_id, s_id in mapping.items():
+        if verified_count >= num_samples:
+            break
+        if s_id != -1:
+            token = t_tokenizer.convert_ids_to_tokens(t_id)
+            source_vec = s_embeds[s_id]
+            target_vec = t_embeds[t_id]
+            diff = torch.sum(torch.abs(source_vec - target_vec)).item()
+            if diff < 1e-6:
+                print(f"    - ✓ Token '{token}' (ID {t_id}) transferred successfully (diff: {diff:.2e}).")
+            else:
+                print(f"    - ✗ FAILED: Token '{token}' (ID {t_id}) has a large difference after transfer (diff: {diff:.2e}).")
+            verified_count += 1
+            
+def verify_grafted_layer(target_state_dict, donor_state_dict, target_layer_idx, donor_layer_idx):
+    """Verifies that cyclical grafting for q_norm/k_norm worked."""
+    print(f"\n  - DIAGNOSTIC: Verifying cyclical graft for target layer {target_layer_idx} from donor layer {donor_layer_idx}...")
+    for norm_type in ['q_norm', 'k_norm']:
+        target_key = f'model.layers.{target_layer_idx}.self_attn.{norm_type}.weight'
+        donor_key = f'model.layers.{donor_layer_idx}.self_attn.{norm_type}.weight'
+        diff = torch.sum(torch.abs(target_state_dict[target_key] - donor_state_dict[donor_key])).item()
+        if diff < 1e-6:
+            print(f"    - ✓ {norm_type} weights match (diff: {diff:.2e}).")
+        else:
+            print(f"    - ✗ FAILED: {norm_type} weights DO NOT match (diff: {diff:.2e}).")
+
+def check_for_nan_inf(state_dict):
+    """Scans the entire state_dict for NaN or Inf values."""
+    print("\n  - DIAGNOSTIC: Scanning final state dictionary for NaN/Inf values...")
+    found_issue = False
+    for key, tensor in tqdm(state_dict.items(), desc="Scanning tensors"):
+        if torch.isnan(tensor).any() or torch.isinf(tensor).any():
+            print(f"    - ✗ FAILED: Found NaN or Inf in tensor '{key}'!")
+            found_issue = True
+    if not found_issue:
+        print("    - ✓ All tensors in the final state dictionary are clean.")
+    return not found_issue
+
+# --- STANDARD HELPERS ---
+
+def create_vocab_mapping(s_tok, t_tok):
+    # ... (code is unchanged)
+    s_vocab, t_vocab = s_tok.get_vocab(), t_tok.get_vocab()
+    s_tok_to_id = {t: i for t, i in s_vocab.items()}
+    mapping = {t_id: s_tok_to_id.get(t, -1) for t, t_id in t_vocab.items()}
+    matches = sum(1 for v in mapping.values() if v != -1)
+    print(f"Vocabulary overlap: {matches}/{len(t_vocab)} tokens ({matches/len(t_vocab)*100:.1f}%) will be transferred.")
+    return mapping
+
+def verify_special_tokens(s_tok, t_tok, mapping):
+    # ... (code is unchanged)
+    print("\nVerifying special token mappings...")
+    for name, token_value in t_tok.special_tokens_map.items():
+        def _process_token(token_str):
+            if token_str and token_str in t_tok.get_vocab():
+                t_id = t_tok.convert_tokens_to_ids(token_str)
+                s_id = mapping.get(t_id, -1)
+                status = f"Mapped (T: {t_id} -> S: {s_id})" if s_id != -1 else "NOT FOUND in source (initialized with mean)"
+                print(f"  ✓ ('{token_str}'): {status}")
+        if isinstance(token_value, str): _process_token(token_value)
+        elif isinstance(token_value, list):
+            for token_str_in_list in token_value: _process_token(token_str_in_list)
+
+def create_hybrid_matrix(s_matrix, mapping, shape):
+    # ... (code is unchanged, but we'll add logging inside the main function)
+    mean_embedding = s_matrix.mean(dim=0, keepdim=True)
+    hybrid = torch.zeros(shape, dtype=s_matrix.dtype, device='cpu')
+    for t_id, s_id in mapping.items():
+        hybrid[t_id] = s_matrix[s_id] if s_id != -1 else mean_embedding
+    return hybrid.to(s_matrix.device)
+
+def validate_model_diagnostic(path):
+    print("\n[Step 6/6] Running DIAGNOSTIC validation...")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.bfloat16)
+        model.eval()
+        prompt = "The theory of relativity states that"
+        print(f"\nValidation Prompt: '{prompt}'")
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=25, do_sample=False, pad_token_id=tokenizer.eos_token_id)
+        
+        print("\n--- DIAGNOSTIC: RAW TOKEN IDs ---")
+        print(outputs[0].tolist())
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print("\n--- DIAGNOSTIC: Decoded Response ---")
+        print(f"'{response}'")
+        
+        if '�' in response:
+            print("\n  - ✗ VALIDATION FAILED: Found replacement character '�' in output. This indicates a tokenization/decoding issue.")
+            return
+
+        # A more robust check for coherence
+        if "states that states that" in response or "the the the" in response or len(set(response.split())) < 5:
+             print("\n  - ✗ VALIDATION FAILED: Output appears repetitive or incoherent.")
+        else:
+             print("\n  - ✓ Validation check passed: Model generated non-repetitive text.")
+
+    except Exception as e:
+        print(f"\n  ✗ Validation FAILED with an exception: {e}")
+
+# --- Main Conversion Logic ---
+def convert_qwen2_to_qwen3_diagnostic():
+    source_model_id, donor_model_id = "Qwen/Qwen2.5-72B-Instruct", "Qwen/Qwen3-32B"
+    target_model_path = "./Qwen3-72B"
+    print("Starting DIAGNOSTIC conversion process (v6.0)...")
+
+    # --- Step 1 & 2: Load everything ---
+    s_config = AutoConfig.from_pretrained(source_model_id)
+    d_config = AutoConfig.from_pretrained(donor_model_id)
+    dtype = torch.bfloat16
+    s_model = AutoModelForCausalLM.from_pretrained(source_model_id, torch_dtype=dtype, device_map="auto")
+    d_model = AutoModelForCausalLM.from_pretrained(donor_model_id, torch_dtype=dtype, device_map="auto")
+    s_tokenizer = AutoTokenizer.from_pretrained(source_model_id)
+    t_tokenizer = AutoTokenizer.from_pretrained(donor_model_id)
+    
+    # --- Step 3: Create Target ---
+    t_config = Qwen3Config(hidden_size=s_config.hidden_size, intermediate_size=s_config.intermediate_size, num_hidden_layers=s_config.num_hidden_layers, num_attention_heads=s_config.num_attention_heads, num_key_value_heads=s_config.num_key_value_heads, max_position_embeddings=s_config.max_position_embeddings, max_window_layers=s_config.max_window_layers, sliding_window=s_config.sliding_window, attention_bias=d_config.attention_bias, hidden_act=d_config.hidden_act, initializer_range=d_config.initializer_range, rms_norm_eps=d_config.rms_norm_eps, rope_theta=d_config.rope_theta, vocab_size=d_config.vocab_size, tie_word_embeddings=True)
+    with torch.device("meta"): t_model = Qwen3ForCausalLM(t_config)
+
+    # --- Step 4: Convert and DIAGNOSE Weights ---
+    print("\n[Step 4/6] Converting weights (DIAGNOSTIC mode)...")
+    s_state_dict = {k: v.cpu() for k, v in tqdm(s_model.state_dict().items(), desc="Source state dict to CPU")}
+    d_state_dict = {k: v.cpu() for k, v in tqdm(d_model.state_dict().items(), desc="Donor state dict to CPU")}
+    vocab_mapping = create_vocab_mapping(s_tokenizer, t_tokenizer)
+    verify_special_tokens(s_tokenizer, t_tokenizer, vocab_mapping)
+    
+    new_state_dict = {}
+    num_donor_layers = d_config.num_hidden_layers
+    
+    # --- Create and Diagnose Hybrid Embeddings ---
+    print("\n--- Creating and Diagnosing Embedding and LM Head ---")
+    
+    # Embeddings
+    print("Processing model.embed_tokens.weight...")
+    s_embeds = s_state_dict['model.embed_tokens.weight']
+    mean_embedding = s_embeds.mean(dim=0, keepdim=True)
+    log_tensor_stats(mean_embedding, "Mean Initializer Vector")
+    new_embed_matrix = create_hybrid_matrix(s_embeds, vocab_mapping, (t_config.vocab_size, t_config.hidden_size))
+    log_tensor_stats(new_embed_matrix, "Final Hybrid Embedding Matrix")
+    new_state_dict['model.embed_tokens.weight'] = new_embed_matrix
+    verify_embedding_transfer(s_embeds, new_embed_matrix, vocab_mapping, t_tokenizer)
+
+    # LM Head
+    print("\nProcessing lm_head.weight...")
+    s_lm_head = s_state_dict['lm_head.weight']
+    new_lm_head_matrix = create_hybrid_matrix(s_lm_head, vocab_mapping, (t_config.vocab_size, t_config.hidden_size))
+    log_tensor_stats(new_lm_head_matrix, "Final Hybrid LM Head Matrix")
+    new_state_dict['lm_head.weight'] = new_lm_head_matrix
+    
+    # --- Process remaining layers ---
+    print("\n--- Processing Transformer Layers ---")
+    # Get all keys except the ones we already handled
+    remaining_keys = [k for k in t_model.state_dict().keys() if 'embed_tokens' not in k and 'lm_head' not in k]
+
+    for key in tqdm(remaining_keys, desc="Transferring layer weights"):
+        if "q_norm" in key or "k_norm" in key:
+            match = re.search(r'layers\.(\d+)\.', key)
+            if match:
+                target_layer_idx = int(match.group(1))
+                donor_layer_idx = target_layer_idx % num_donor_layers
+                donor_key = key.replace(f'layers.{target_layer_idx}.', f'layers.{donor_layer_idx}.')
+                new_state_dict[key] = d_state_dict[donor_key].clone()
+        elif key in s_state_dict:
+            new_state_dict[key] = s_state_dict[key].clone()
+        else:
+            print(f"  ⚠️ Unhandled key: {key} (not in source, skipping)")
+
+    # --- Final Diagnostic Checks ---
+    verify_grafted_layer(new_state_dict, d_state_dict, target_layer_idx=num_donor_layers, donor_layer_idx=0)
+    all_clean = check_for_nan_inf(new_state_dict)
+    if not all_clean:
+        print("\nCRITICAL ERROR: NaN/Inf values detected. Aborting before save.")
+        return # Stop the process
+
+    # --- Step 5: Save ---
+    print("\n[Step 5/6] Saving final model and supporting files...")
+    t_model.load_state_dict(new_state_dict) # Load into meta-model
+    t_model.save_pretrained(target_model_path, safe_serialization=True, state_dict=new_state_dict)
+    t_tokenizer.save_pretrained(target_model_path)
+    print(f"✅ Model saved to: {target_model_path}")
+
+    # --- Step 6: Validate ---
+    del s_model, d_model, s_state_dict, d_state_dict, new_state_dict, t_model
+    torch.cuda.empty_cache()
+    validate_model_diagnostic(path=target_model_path)
+
+if __name__ == "__main__":
+    convert_qwen2_to_qwen3_diagnostic()