phase1.py

import torch
import os
import argparse
from tqdm import tqdm
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from accelerate import init_empty_weights # <-- IMPORT THE FIX

# --- Configuration ---
SRC_HIDDEN_SIZE = 5120
SRC_INTERMEDIATE_SIZE = 25600
TGT_HIDDEN_SIZE = 8192
TGT_INTERMEDIATE_SIZE = 29568

DELTA_HIDDEN = TGT_HIDDEN_SIZE - SRC_HIDDEN_SIZE
DELTA_INTERMEDIATE = TGT_INTERMEDIATE_SIZE - SRC_INTERMEDIATE_SIZE

# --- Interpolation Functions ---
def linear_interpolation(block1, block2, weight=0.5):
    return (1 - weight) * block1 + weight * block2

def upscale_tensor(tensor: torch.Tensor, name: str) -> torch.Tensor:
    # This function is correct from the previous version, simplified for brevity
    if tensor.ndim == 1:
        if tensor.shape[0] == SRC_HIDDEN_SIZE:
            block1, block2 = tensor[:DELTA_HIDDEN], tensor[-DELTA_HIDDEN:]
            return torch.cat([tensor, linear_interpolation(block1, block2)], dim=0)
    elif tensor.ndim == 2:
        if "embed_tokens" in name or "lm_head" in name:
            if tensor.shape[1] == SRC_HIDDEN_SIZE:
                block1, block2 = tensor[:, :DELTA_HIDDEN], tensor[:, -DELTA_HIDDEN:]
                return torch.cat([tensor, linear_interpolation(block1, block2)], dim=1)
        elif "self_attn" in name:
            if "q_proj.weight" in name or "k_proj.weight" in name or "v_proj.weight" in name:
                block1, block2 = tensor[:, :DELTA_HIDDEN], tensor[:, -DELTA_HIDDEN:]
                return torch.cat([tensor, linear_interpolation(block1, block2)], dim=1)
            elif "o_proj.weight" in name:
                block1, block2 = tensor[:DELTA_HIDDEN, :], tensor[-DELTA_HIDDEN:, :]
                return torch.cat([tensor, linear_interpolation(block1, block2)], dim=0)
        elif "mlp" in name:
            if "gate_proj.weight" in name or "up_proj.weight" in name:
                row_block1, row_block2 = tensor[:DELTA_INTERMEDIATE, :], tensor[-DELTA_INTERMEDIATE:, :]
                upscaled_rows = torch.cat([tensor, linear_interpolation(row_block1, row_block2)], dim=0)
                col_block1, col_block2 = upscaled_rows[:, :DELTA_HIDDEN], upscaled_rows[:, -DELTA_HIDDEN:]
                return torch.cat([upscaled_rows, linear_interpolation(col_block1, col_block2)], dim=1)
            elif "down_proj.weight" in name:
                row_block1, row_block2 = tensor[:DELTA_HIDDEN, :], tensor[-DELTA_HIDDEN:, :]
                upscaled_rows = torch.cat([tensor, linear_interpolation(row_block1, row_block2)], dim=0)
                col_block1, col_block2 = upscaled_rows[:, :DELTA_INTERMEDIATE], upscaled_rows[:, -DELTA_INTERMEDIATE:]
                return torch.cat([upscaled_rows, linear_interpolation(col_block1, col_block2)], dim=1)
    return tensor

def run_test_inference(model_path, prompt):
    print("\n" + "="*50)
    print("Running test inference...")
    print("="*50)
    
    # Load the newly saved model to ensure it works from disk
    print(f"Loading model from disk: {model_path}")
    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

    print(f"Prompt: \"{prompt}\"")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    model.to(device)

    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(**inputs, max_new_tokens=50, do_sample=False)
    
    result_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
    print("\n--- Generated Text ---")
    print(result_text)
    print("----------------------")
    print("\nTest inference complete.")


def main(source_model_id, output_path):
    print(f"Loading source model: {source_model_id}")
    # Load on CPU to save VRAM
    source_model = AutoModelForCausalLM.from_pretrained(
        source_model_id, torch_dtype=torch.bfloat16, device_map="cpu", trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(source_model_id, trust_remote_code=True)
    
    source_state_dict = source_model.state_dict()
    new_state_dict = {}

    print("Up-scaling tensors using self-interpolation...")
    for name, tensor in tqdm(source_state_dict.items(), desc="Processing Tensors"):
        new_state_dict[name] = upscale_tensor(tensor.clone(), name)
    
    # Free up memory from the source model before creating the new one
    del source_model
    del source_state_dict
    
    print("\nCreating new model with target architecture...")
    config = AutoConfig.from_pretrained(source_model_id, trust_remote_code=True)
    config.hidden_size = TGT_HIDDEN_SIZE
    config.intermediate_size = TGT_INTERMEDIATE_SIZE
    config.torch_dtype = torch.bfloat16 # Ensure config specifies the correct dtype
    
    # --- THIS IS THE FIX ---
    # Initialize a 'skeleton' model without allocating memory. This is instant.
    print("Step 1: Initializing empty 'skeleton' model on meta device (should be instantaneous)...")
    with init_empty_weights():
        new_model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
    print("Empty model created successfully.")
    # ---------------------
    
    # The model is now on the 'meta' device. We need to tie the weights to the CPU.
    new_model.tie_weights()
    
    print("\nStep 2: Loading up-scaled weights into the new model (this may take time and RAM)...")
    # This step will materialize the weights on the CPU, consuming memory.
    new_model.load_state_dict(new_state_dict, assign=True)
    print("State dict loaded successfully.")
    
    print(f"\nStep 3: Saving up-scaled model and tokenizer to: {output_path}")
    os.makedirs(output_path, exist_ok=True)
    new_model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)
    
    print("\nPhase 1 (Self-Interpolation) Complete!")
    print(f"The up-scaled model is ready at '{output_path}' for use with MergeKit.")

    # --- Run Test Inference ---
    # We pass the path instead of the model object to test loading from disk
    run_test_inference(output_path, "The rain in Maine ")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Upscale Qwen3-32B to hypothetical Qwen3-72B dimensions via self-interpolation.")
    parser.add_argument("--source_model", type=str, default="Qwen/Qwen3-32B", help="The Hugging Face model ID of the source model.")
    parser.add_argument(
        "--output_path", type=str, default="./Qwen3-32B-Upscaled",
        help="The local directory path to save the up-scaled model. (default: ./Qwen3-32B-Upscaled)"
    )
    args = parser.parse_args()
    main(args.source_model, args.output_path)