MoR-TC-v2.1-2-ties

MoR-TC-v2.1-2-ties This model is a merge of "liminerity/MoR-TC-v2.1" and "liminerity/MoR-TC-v2.1-2" \n "liminerity/MoR-TC-v2.1" was trained on the first half of "cognitivecomputations/dolphin" and 2.1-2 was trained on the second half.\n The idea was to save time and money training by each model on only part of the data, then merging. \n The following code can be used to inference this model:\n

import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from transformers import GPT2Tokenizer
from safetensors.torch import load_file
from huggingface_hub import snapshot_download
import sys

class Config:
    def __init__(self, **kwargs):
        self.vocab_size = 50257
        self.d_model = 1024
        self.n_head = 16
        self.d_k = self.d_model // self.n_head
        self.d_ff = 4096
        self.max_depth = 4
        self.num_recursive_layers = 6
        self.balancing_weight = 0.01
        self.temperature = 1.0
        self.seq_len = 512
        self.batch_size = 16
        self.window_size = 2048
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        for key, value in kwargs.items():
            setattr(self, key, value)
        if hasattr(self, 'd_model') and hasattr(self, 'n_head'):
            self.d_k = self.d_model // self.n_head

class RecursiveLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.w_q = nn.Linear(config.d_model, config.d_model)
        self.w_k = nn.Linear(config.d_model, config.d_model)
        self.w_v = nn.Linear(config.d_model, config.d_model)
        self.attn_out = nn.Linear(config.d_model, config.d_model)
        self.ffn = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.GELU(),
            nn.Linear(config.d_ff, config.d_model)
        )
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)

    def forward(self, h, active_mask):
        batch_size, seq_len, _ = h.shape

        # Project current hidden state for Q, K, V
        q = self.w_q(h).view(batch_size, seq_len, self.config.n_head, self.config.d_k)
        k = self.w_k(h).view(batch_size, seq_len, self.config.n_head, self.config.d_k)
        v = self.w_v(h).view(batch_size, seq_len, self.config.n_head, self.config.d_k)

        q = q.permute(0, 2, 1, 3)  # [batch, head, seq, d_k]
        k = k.permute(0, 2, 1, 3)  # [batch, head, seq, d_k]
        v = v.permute(0, 2, 1, 3)  # [batch, head, seq, d_k]

        # Create causal mask with windowing
        attn_mask = torch.ones(seq_len, seq_len, device=h.device, dtype=torch.bool)
        attn_mask = torch.tril(attn_mask, diagonal=0)  # Causal lower triangle
        attn_mask = torch.triu(attn_mask, diagonal=-self.config.window_size)  # Windowing

        # Expand mask for batch and heads
        attn_mask = attn_mask.view(1, 1, seq_len, seq_len)

        # Compute attention scores
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.config.d_k)
        attn_scores = attn_scores.masked_fill(~attn_mask, float('-inf'))
        attn_probs = F.softmax(attn_scores, dim=-1)

        # Apply attention
        attn_out = torch.matmul(attn_probs, v)
        attn_out = attn_out.permute(0, 2, 1, 3).contiguous()
        attn_out = attn_out.view(batch_size, seq_len, self.config.d_model)
        attn_out = self.attn_out(attn_out)

        # Apply active mask
        active_mask_expanded = active_mask.unsqueeze(-1)
        attn_out = attn_out * active_mask_expanded

        # Residual connection and norm
        h = h + attn_out
        h = self.norm1(h)

        # FFN
        ffn_out = self.ffn(h) * active_mask_expanded
        h = h + ffn_out
        h = self.norm2(h)

        return h

class Router(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(config.d_model, config.d_model // 2),
            nn.GELU(),
            nn.Linear(config.d_model // 2, config.max_depth)
        )
        self.temperature = config.temperature

    def forward(self, h, train=True):
        logits = self.linear(h)
        if train:
            probs = F.gumbel_softmax(logits, tau=self.temperature, dim=-1)
            return probs, F.softmax(logits, dim=-1)
        else:
            probs = F.softmax(logits, dim=-1)
            return probs, probs

class MixtureRecursions(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_embed = nn.Embedding(config.seq_len, config.d_model)
        self.first_layer = nn.Sequential(
            nn.Linear(config.d_model, config.d_model),
            nn.GELU(),
            nn.LayerNorm(config.d_model)
        )
        self.recursive_layers = nn.ModuleList([
            RecursiveLayer(config) for _ in range(config.num_recursive_layers)
        ])
        self.router = Router(config)
        self.final_norm = nn.LayerNorm(config.d_model)
        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, x, targets=None):
        device = x.device
        batch_size, seq_len = x.shape
        pos_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
        pos_emb = self.pos_embed(pos_ids)
        tok_emb = self.embed(x)
        h = tok_emb + pos_emb
        h = self.first_layer(h)

        # Get router assignments
        router_probs, router_soft = self.router(h)
        assigned_depths = router_probs.argmax(dim=-1) + 1

        # Process through recursive layers
        for depth in range(1, self.config.max_depth + 1):
            active_mask = (assigned_depths >= depth)
            layer_idx = (depth - 1) % self.config.num_recursive_layers
            h = self.recursive_layers[layer_idx](h, active_mask)

        h = self.final_norm(h)
        logits = self.head(h)
        loss = None
        balancing_loss = None

        if targets is not None:
            logits = logits[:, :-1, :].contiguous()
            targets = targets[:, 1:].contiguous()
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

            # Balancing loss
            router_decision = router_probs.sum(dim=[0, 1])
            router_decision = router_decision / (batch_size * seq_len)
            balancing_loss = torch.var(router_decision) * self.config.balancing_weight

            return logits, loss, balancing_loss

        return logits, loss, balancing_loss

# --- Download and load everything as before ---
repo_id = "liminerity/MoR-TC-v2"
model_dir = snapshot_download(repo_id=repo_id)
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
with open(f"{model_dir}/config.json", 'r') as f:
    hf_config = json.load(f)

config_map = {
    'vocab_size': 'vocab_size',
    'dim': 'd_model',
    'num_layers': 'num_recursive_layers',
    'num_heads': 'n_head',
    'max_recursion': 'max_depth',
    'max_position_embeddings': 'seq_len',
    'balancing_weight': 'balancing_weight',
    'temperature': 'temperature',
    'window_size': 'window_size'
}
mapped_config = {config_map[k]: v for k, v in hf_config.items() if k in config_map}
mapped_config['d_ff'] = hf_config['ffn_expansion'] * mapped_config['d_model']
config = Config(**mapped_config)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MixtureRecursions(config).to(device)

weights = load_file(f"{model_dir}/model.safetensors", device=str(device))
model.load_state_dict(weights)
model.eval()

# --- Autoregressive Generation Loop without KV Cache ---
def autoregressive_generate(
    model, tokenizer, input_text, max_new_tokens=500, temperature=0.3, line_width=71):

    model.eval()
    device = next(model.parameters()).device
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
    current_ids = input_ids
    generated_text = input_text

    # Print initial text without newline
    print(input_text, end="", flush=True)

    for _ in range(max_new_tokens):
        # Truncate if sequence gets too long
        if current_ids.shape[1] >= config.seq_len:
            current_ids = current_ids[:, -config.seq_len:]

        with torch.no_grad():
            # Run model on current sequence
            logits = model(current_ids)[0]

        # Get next token
        next_token_logits = logits[0, -1, :] / temperature
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token_id = torch.multinomial(probs, num_samples=1).item()

        # Append new token
        current_ids = torch.cat(
            [current_ids, torch.tensor([[next_token_id]], device=device)], dim=1
        )

        # Decode and print token by token
        new_token = tokenizer.decode([next_token_id])
        generated_text += new_token
        print(new_token, end="", flush=True)

    print()  # Final newline

# Test streaming generation of 500 tokens
input_text = "The future of AI is"
autoregressive_generate(model, tokenizer, input_text, max_new_tokens=500, temperature=config.temperature)

The following code was used to merge the two models using the ties method:\n\n

# Install required libraries
#!pip install transformers huggingface-hub safetensors

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file, save_file
from transformers import GPT2Tokenizer
import os
import shutil
import json
def push_to_hub(save_folder):
    REPO_NAME = "liminerity/MoR-TC-v2.1-2-ties"  # Replace with your Hugging Face username and desired model name
    MODEL_DIR = save_folder  # Directory where we saved the model

    # Create repository and push files
    api = HfApi()
    api.create_repo(
        repo_id=REPO_NAME,
        repo_type="model",
        exist_ok=True  # Will not error if repo already exists
    )

    api.upload_folder(
        folder_path=MODEL_DIR,
        repo_id=REPO_NAME,
        repo_type="model"
    )

    print(f"Model successfully pushed to: https://huggingface.co/{REPO_NAME}")
# Configuration
SPARSITY = 0.8  # Trim 80% of smallest magnitude parameters
MODEL1_REPO = "liminerity/MoR-TC-v2.1-2"
MODEL2_REPO = "liminerity/MoR-TC-v2.1"
MERGED_MODEL_DIR = "MoR-TC-merged"
save_folder = MERGED_MODEL_DIR
# Download models
model1_dir = snapshot_download(repo_id=MODEL1_REPO)
model2_dir = snapshot_download(repo_id=MODEL2_REPO)

# Load state_dicts
state_dict1 = load_file(os.path.join(model1_dir, "model.safetensors"))
state_dict2 = load_file(os.path.join(model2_dir, "model.safetensors"))

# Create base state_dict (average of both models)
base_state_dict = {}
for name in state_dict1:
    base_state_dict[name] = (state_dict1[name] + state_dict2[name]) / 2

# Prepare merged state_dict
merged_state_dict = {}

# TIES-Merging: Trim, Elect Sign, Disjoint Merge
for name in base_state_dict:
    base_param = base_state_dict[name]
    param1 = state_dict1[name]
    param2 = state_dict2[name]

    # Compute deltas
    delta1 = param1 - base_param
    delta2 = param2 - base_param

    # Trim: Set smallest magnitude parameters to zero
    k1 = int(delta1.numel() * SPARSITY)
    k2 = int(delta2.numel() * SPARSITY)

    if k1 > 0:
        flat_d1 = delta1.view(-1)
        _, indices = torch.topk(flat_d1.abs(), k1, largest=False)
        flat_d1[indices] = 0

    if k2 > 0:
        flat_d2 = delta2.view(-1)
        _, indices = torch.topk(flat_d2.abs(), k2, largest=False)
        flat_d2[indices] = 0

    # Elect Sign: Determine dominant direction
    total_delta = delta1 + delta2
    elected_sign = torch.sign(total_delta)

    # Nullify conflicting updates
    mask1 = (delta1 != 0) & (torch.sign(delta1) != elected_sign)
    delta1[mask1] = 0

    mask2 = (delta2 != 0) & (torch.sign(delta2) != elected_sign)
    delta2[mask2] = 0

    # Disjoint Merge: Average aligned updates
    count = (delta1 != 0).float() + (delta2 != 0).float()
    merged_delta = (delta1 + delta2) / torch.clamp(count, min=1.0)

    # Combine with base
    merged_state_dict[name] = base_param + merged_delta

# Save merged model
os.makedirs(MERGED_MODEL_DIR, exist_ok=True)
save_file(merged_state_dict, os.path.join(MERGED_MODEL_DIR, "model.safetensors"))

# Copy config from model1
shutil.copy(os.path.join(model1_dir, "config.json"),
            os.path.join(MERGED_MODEL_DIR, "config.json"))

# Save tokenizer from model1
tokenizer = GPT2Tokenizer.from_pretrained(model1_dir)
tokenizer.save_pretrained(MERGED_MODEL_DIR)

print(f"Merged model saved to: {MERGED_MODEL_DIR}")
push_to_hub(save_folder)