rio model analysis

gistfile1.txt

"""
Minimal verification: is Rio-3.5-Open-397B ≈ α·Nex-N2-Pro + (1-α)·Qwen3.5-397B-A17B?

Strategy: download the safetensors header via HTTP range request, then fetch
only the specific tensor bytes we need. No multi-GB shard downloads.
"""

import json
import struct
import numpy as np
from huggingface_hub import hf_hub_url, get_token
import urllib.request

MODELS = {
    "rio":  "prefeitura-rio/Rio-3.5-Open-397B",
    "nex":  "nex-agi/Nex-N2-Pro",
    "qwen": "Qwen/Qwen3.5-397B-A17B",
}

PROBE_TENSORS = [
    # MoE router weights — 2D, [num_experts, hidden_dim]
    "model.language_model.layers.0.mlp.gate.weight",
    "model.language_model.layers.15.mlp.gate.weight",
    "model.language_model.layers.30.mlp.gate.weight",
    "model.language_model.layers.45.mlp.gate.weight",
    "model.language_model.layers.59.mlp.gate.weight",
    # Shared expert gate
    "model.language_model.layers.0.mlp.shared_expert_gate.weight",
    "model.language_model.layers.30.mlp.shared_expert_gate.weight",
    # Norms for comparison
    "model.language_model.layers.0.input_layernorm.weight",
    "model.language_model.layers.30.input_layernorm.weight",
    "model.language_model.norm.weight",
]

DTYPE_MAP = {
    "BF16": (2, ">u2"),  # we'll decode bf16 manually
    "F16":  (2, "<e"),
    "F32":  (4, "<f"),
}


def bf16_to_f32(raw_bytes):
    u16 = np.frombuffer(raw_bytes, dtype=np.uint16)
    u32 = u16.astype(np.uint32) << 16
    return u32.view(np.float32)


def fetch_range(url, start, end, token=None):
    req = urllib.request.Request(url)
    req.add_header("Range", f"bytes={start}-{end-1}")
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    with urllib.request.urlopen(req) as resp:
        return resp.read()


def get_safetensors_header(url, token=None):
    first_8 = fetch_range(url, 0, 8, token)
    header_len = struct.unpack("<Q", first_8)[0]
    header_bytes = fetch_range(url, 8, 8 + header_len, token)
    header = json.loads(header_bytes)
    return header, 8 + header_len


def fetch_tensor(url, header, data_offset, tensor_name, token=None):
    meta = header[tensor_name]
    dtype_str = meta["dtype"]
    shape = meta["shape"]
    offsets = meta["data_offsets"]
    abs_start = data_offset + offsets[0]
    abs_end = data_offset + offsets[1]

    raw = fetch_range(url, abs_start, abs_end, token)

    if dtype_str == "BF16":
        return bf16_to_f32(raw)
    elif dtype_str == "F16":
        return np.frombuffer(raw, dtype=np.float16).astype(np.float32)
    elif dtype_str == "F32":
        return np.frombuffer(raw, dtype=np.float32)
    else:
        raise ValueError(f"Unknown dtype: {dtype_str}")


def main():
    from huggingface_hub import hf_hub_download
    token = get_token()

    # Step 1: load index files to find which shard holds each tensor
    shard_map = {}
    for key, repo in MODELS.items():
        path = hf_hub_download(repo, "model.safetensors.index.json")
        with open(path) as f:
            idx = json.load(f)
        shard_map[key] = idx["weight_map"]

    # Step 2: for each probe tensor, identify the shard and cache headers
    header_cache = {}  # (repo, shard_filename) -> (header, data_offset, url)

    def get_header_for(model_key, tensor_name):
        repo = MODELS[model_key]
        shard_fn = shard_map[model_key][tensor_name]
        cache_key = (model_key, shard_fn)
        if cache_key not in header_cache:
            url = hf_hub_url(repo, shard_fn)
            print(f"  Fetching header: {model_key}/{shard_fn}")
            header, data_offset = get_safetensors_header(url, token)
            header_cache[cache_key] = (header, data_offset, url)
        return header_cache[cache_key]

    # Step 3: fetch tensors and compute collinearity
    print("=" * 82)
    print(f"{'Tensor':<48} {'#params':>10} {'α':>6} {'cos_fit':>8}")
    print("=" * 82)

    alphas = []
    cos_fits = []

    for tname in PROBE_TENSORS:
        tensors = {}
        for mkey in ["rio", "nex", "qwen"]:
            header, data_offset, url = get_header_for(mkey, tname)
            tensors[mkey] = fetch_tensor(url, header, data_offset, tname, token)

        rio = tensors["rio"]
        nex = tensors["nex"]
        qwen = tensors["qwen"]

        delta_rio = (rio - qwen).astype(np.float64).flatten()
        delta_nex = (nex - qwen).astype(np.float64).flatten()

        # cosine similarity
        dot = np.dot(delta_rio, delta_nex)
        norm_rio = np.linalg.norm(delta_rio)
        norm_nex = np.linalg.norm(delta_nex)
        cos = dot / (norm_rio * norm_nex + 1e-30)

        # alpha = projection of delta_rio onto delta_nex direction
        alpha = dot / (np.dot(delta_nex, delta_nex) + 1e-30)

        short = tname.replace("model.language_model.", "")
        n_params = rio.size
        print(f"  {short:<48} {n_params:>10,} {alpha:>6.4f} {cos:>8.6f}")
        alphas.append(alpha)
        cos_fits.append(cos)

    print("=" * 82)
    print(f"  {'MEAN':<48} {'':>10} {np.mean(alphas):>6.4f} {np.mean(cos_fits):>8.6f}")
    print(f"  {'STD':<48} {'':>10} {np.std(alphas):>6.4f} {np.std(cos_fits):>8.6f}")
    print()

    mean_cos = np.mean(cos_fits)
    mean_alpha = np.mean(alphas)
    if mean_cos > 0.99 and 0.5 < mean_alpha < 0.7:
        print(f"RESULT: Consistent with a linear merge at α≈{mean_alpha:.3f}")
        print(f"        cos_fit={mean_cos:.6f} — this is NOT explainable by chance.")
    elif mean_cos > 0.9:
        print(f"RESULT: High similarity (cos={mean_cos:.4f}) but not a clean merge.")
    else:
        print(f"RESULT: No evidence of a linear merge (cos={mean_cos:.4f}).")


if __name__ == "__main__":
    main()