Felipe97/llama-cpp-compiled

1#!/usr/bin/env python32 3import os4import sys5import torch6import transformers7import json8import textwrap9import numpy as np10from pathlib import Path11 12 13def get_model_name_from_env_path(env_path_name):14    model_path = os.getenv(env_path_name)15    if not model_path:16        print(f"Error: {env_path_name} environment variable not set")17        sys.exit(1)18 19    if not os.path.exists(model_path):20        print(f"Error: Model file not found: {model_path}")21        sys.exit(1)22 23    name = os.path.basename(os.path.normpath(model_path))24    if name.endswith(".gguf"):25        name = name[:-5]26 27    return name28 29 30def summarize(tensor: torch.Tensor, name: str, max_seq: int = 3, max_vals: int = 3):31    """32    Print a tensor in llama.cpp debug style.33 34    Supports:35    - 2D tensors (seq, hidden)36    - 3D tensors (batch, seq, hidden)37    - 4D tensors (batch, seq, heads, dim_per_head) via flattening heads × dim_per_head38 39    Shows first and last max_vals of each vector per sequence position.40    """41    t = tensor.detach().to(torch.float32).cpu()42 43    # Determine dimensions44    if t.ndim == 3:45        _, s, _ = t.shape46    elif t.ndim == 2:47        _, s = 1, t.shape[0]48        t = t.unsqueeze(0)49    elif t.ndim == 4:50        _, s, _, _ = t.shape51    else:52        print(f"Skipping tensor due to unsupported dimensions: {t.ndim}")53        return54 55    ten_shape = t.shape56 57    print(f"ggml_debug: {name} = (f32)  ... = {{{ten_shape}}}")58    print("                                     [")59    print("                                      [")60 61    # Determine indices for first and last sequences62    first_indices = list(range(min(s, max_seq)))63    last_indices = list(range(max(0, s - max_seq), s))64 65    # Check if there's an overlap between first and last indices or if we're at the edge case of s = 2 * max_seq66    has_overlap = bool(set(first_indices) & set(last_indices)) or (max_seq * 2 == s)67 68    # Combine indices69    if has_overlap:70        # If there's overlap, just use the combined unique indices71        indices = sorted(list(set(first_indices + last_indices)))72        separator_index = None73    else:74        # If no overlap, we'll add a separator between first and last sequences75        indices = first_indices + last_indices76        separator_index = len(first_indices)77 78    for i, si in enumerate(indices):79        # Add separator if needed80        if separator_index is not None and i == separator_index:81            print("                                       ...")82 83        # Extract appropriate slice84        vec = t[0, si]85        if vec.ndim == 2:  # 4D case: flatten heads × dim_per_head86            flat = vec.flatten().tolist()87        else:  # 2D or 3D case88            flat = vec.tolist()89 90        # First and last slices91        first = flat[:max_vals]92        last = flat[-max_vals:] if len(flat) >= max_vals else flat93        first_str = ", ".join(f"{v:12.4f}" for v in first)94        last_str = ", ".join(f"{v:12.4f}" for v in last)95 96        print(f"                                       [{first_str}, ..., {last_str}]")97 98    print("                                      ],")99    print("                                     ]")100    print(f"                                     sum = {t.sum().item():.6f}\n")101 102 103def debug_hook(name):104    def fn(_m, input, output):105        if isinstance(input, torch.Tensor):106            summarize(input, name + "_in")107        elif isinstance(input, (tuple, list)) and len(input) > 0 and isinstance(input[0], torch.Tensor):108            summarize(input[0], name + "_in")109        if isinstance(output, torch.Tensor):110            summarize(output, name + "_out")111        elif isinstance(output, (tuple, list)) and len(output) > 0 and isinstance(output[0], torch.Tensor):112            summarize(output[0], name + "_out")113 114    return fn115 116 117def setup_rope_debug(model_module_path: str, function_name: str = "apply_rotary_pos_emb"):118    """119    Apply monkey patch to dump RoPE activations for debugging.120 121    Args:122        model_module_path: Path to the model module (e.g., "transformers.models.apertus.modeling_apertus")123        function_name: Name of the RoPE function to patch (default: "apply_rotary_pos_emb")124 125    Example:126        from utils.common import setup_rope_debug127        setup_rope_debug("transformers.models.apertus.modeling_apertus")128    """129    import importlib130 131    # Import the module and get the original function132    module = importlib.import_module(model_module_path)133    orig_rope = getattr(module, function_name)134 135    # Set torch print options for better debugging136    torch.set_printoptions(threshold=float('inf'))137    torch.set_printoptions(precision=6, sci_mode=False)138 139    def debug_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):140        # log inputs141        summarize(q, "RoPE.q_in")142        summarize(k, "RoPE.k_in")143 144        # call original145        q_out, k_out = orig_rope(q, k, cos, sin, position_ids, unsqueeze_dim)146 147        # log outputs148        summarize(q_out, "RoPE.q_out")149        summarize(k_out, "RoPE.k_out")150 151        return q_out, k_out152 153    # Patch it154    setattr(module, function_name, debug_rope)155    print(f"RoPE debug patching applied to {model_module_path}.{function_name}")156 157 158def save_output_data(data, tokens, prompt, model_name, type_suffix="", output_dir="data"):159    """160    Save output data (logits/embeddings), tokens, and prompt to files.161 162    Args:163        data:        numpy array of floats (logits or embeddings)164        tokens:      list or array of token IDs165        prompt:      string containing the input prompt166        model_name:  name of the model167        type_suffix: optional suffix like "-embeddings" (default: "")168        output_dir:  directory to save files (default: "data")169 170    Creates the following files in output_dir:171        - pytorch-{model_name}{type_suffix}.bin172        - pytorch-{model_name}{type_suffix}.txt173        - pytorch-{model_name}{type_suffix}-prompt.txt174        - pytorch-{model_name}{type_suffix}-tokens.bin175    """176    data_dir = Path(output_dir)177    data_dir.mkdir(exist_ok=True)178    base_path = data_dir / f"pytorch-{model_name}{type_suffix}"179 180    # Convert and flatten logits/embeddings181    data = data.cpu().numpy() if isinstance(data, torch.Tensor) else np.asarray(data)182    data = data.flatten() if data.ndim > 1 else data183 184    # Save logits/embedding files185    data.astype(np.float32).tofile(f"{base_path}.bin")186    print(f"Data saved to {base_path}.bin")187 188    with open(f"{base_path}.txt", "w") as f:189        f.writelines(f"{i}: {value:.6f}\n" for i, value in enumerate(data))190    print(f"Data saved to {base_path}.txt")191 192    # Convert and flatten tokens193    tokens = tokens.cpu().numpy() if isinstance(tokens, torch.Tensor) else np.asarray(tokens)194    tokens = tokens.flatten() if tokens.ndim > 1 else tokens195 196    # Save token binary file197    tokens.astype(np.int32).tofile(f"{base_path}-tokens.bin")198    print(f"Tokens saved to {base_path}-tokens.bin")199 200    # Save prompt file201    with open(f"{base_path}-prompt.txt", "w") as f:202        f.write(f"prompt: {prompt}\n")203        f.write(f"n_tokens: {len(tokens)}\n")204        f.write(f"token ids: {', '.join(str(int(tid)) for tid in tokens)}\n")205    print(f"Prompt saved to {base_path}-prompt.txt")206 207 208def compare_tokens(original, converted, type_suffix="", output_dir="data"):209    data_dir = Path(output_dir)210 211    # Read tokens from both models212    tokens1_file = data_dir / f"{original}{type_suffix}-tokens.bin"213    tokens2_file = data_dir / f"{converted}{type_suffix}-tokens.bin"214 215    if not tokens1_file.exists():216        print(f"Error: Token file not found: {tokens1_file}")217        return False218 219    if not tokens2_file.exists():220        print(f"Error: Token file not found: {tokens2_file}")221        return False222 223    tokens1 = np.fromfile(tokens1_file, dtype=np.int32)224    tokens2 = np.fromfile(tokens2_file, dtype=np.int32)225 226    print(f"\nComparing tokens between:")227    print(f"  Original : {original} ({len(tokens1)} tokens)")228    print(f"  Converted: {converted} ({len(tokens2)} tokens)")229 230    if len(tokens1) != len(tokens2):231        print(f"\n❌ Token count mismatch: {len(tokens1)} vs {len(tokens2)}")232        return False233 234    if np.array_equal(tokens1, tokens2):235        print(f"\n✅ All {len(tokens1)} tokens match!")236        return True237 238    mismatches = np.where(tokens1 != tokens2)[0]239    print(f"\n❌ Found {len(mismatches)} mismatched tokens:")240 241    num_to_show = min(len(mismatches), 10)242    for idx in mismatches[:num_to_show]:243        print(f"  Position {idx}: {tokens1[idx]} vs {tokens2[idx]}")244 245    if len(mismatches) > num_to_show:246        print(f"  ... and {len(mismatches) - num_to_show} more mismatches")247 248    return False249 250 251def show_version_warning(current_version, model_version):252    if not model_version:253        return False254 255    try:256        from packaging.version import parse, InvalidVersion257        try:258            return parse(current_version) < parse(model_version)259        except InvalidVersion:260            return current_version != model_version261    except ImportError:262        return current_version != model_version263 264def get_model_transformers_version(model_path):265    if not model_path:266        return None267 268    config_path = Path(model_path) / "config.json"269    if not config_path.is_file():270        return None271 272    try:273        with open(config_path, "r", encoding="utf-8") as f:274            config = json.load(f)275        return config.get("transformers_version")276    except (IOError, json.JSONDecodeError) as e:277        print(f"Warning: Could not read or parse {config_path}: {e}", file=sys.stderr)278        return None279 280def exit_with_warning(message, model_path):281    print(message)282 283    if model_path and transformers is not None:284        model_transformers_version = get_model_transformers_version(model_path)285        transformers_version       = transformers.__version__286        if show_version_warning(transformers_version, model_transformers_version):287            warning_message = f"""288                =====================================================================289                Verification failure might be due to a transformers version mismatch:290 291                Current transformers version: {transformers_version}292                Model's required version    : {model_transformers_version}293 294                Consider installing the version specified by the model's config:295                pip install transformers=={model_transformers_version}296                =====================================================================297            """298            print(textwrap.dedent(warning_message))299    sys.exit(1)300