#This script is used to download the model and pretrained tokenizer from huggingface then initiating it with the defined architecture. from re import M import torch import torch.nn as nn import torch.nn.functional as F import math from torch.utils.data import Dataset, DataLoader from torch.utils.checkpoint import checkpoint from tqdm import tqdm import matplotlib.pyplot as plt from torch.cuda.amp import autocast, GradScaler import numpy as np import os from safetensors.torch import save_file, load_file import json from transformers import PreTrainedTokenizerFast from huggingface_hub import hf_hub_download # Add this at the top to help with debugging os.environ['CUDA_LAUNCH_BLOCKING'] = '1' MODEL = "liminerity/MoR-deep" def save_huggingface_model(model, tokenizer, folder_path="MoR-v1"): # Create directory structure os.makedirs(folder_path, exist_ok=True) # 1. Save model weights in safetensors format weights = model.state_dict() save_file(weights, os.path.join(folder_path, "model.safetensors")) # 2. Create and save config.json config = { "vocab_size": VOCAB_SIZE, "dim": DIM, "num_layers": NUM_LAYERS, "num_heads": HEADS, "max_recursion": MAX_RECURSIONS, "num_experts": model.num_experts, "ffn_expansion": 4, "max_position_embeddings": 2048, "model_type": "MoR", "architecture": "MixtureOfRecursions", "hidden_act": "gelu" } with open(os.path.join(folder_path, "config.json"), "w") as f: json.dump(config, f, indent=2) # 3. Save tokenizer files hf_tokenizer = PreTrainedTokenizerFast( tokenizer_object=tokenizer, unk_token="[UNK]", pad_token="[PAD]", bos_token="[BOS]", eos_token="[EOS]", ) hf_tokenizer.save_pretrained(folder_path) # 4. Create safetensors index file index = { "metadata": {"total_size": sum(p.numel() * p.element_size() for p in model.parameters())}, "weight_map": {name: "model.safetensors" for name in weights.keys()} } with open(os.path.join(folder_path, "model.safetensors.index.json"), "w") as f: json.dump(index, f, indent=2) print(f"Model saved in Hugging Face format to {folder_path}/") def load_model_from_hub(repo_id=MODEL): # Download model files local_dir = f"./models/{repo_id}" if not os.path.exists(local_dir): print(f"Downloading model from {repo_id}...") os.makedirs(local_dir, exist_ok=True) # Download config config_path = hf_hub_download(repo_id, "config.json", cache_dir=local_dir) # Download safetensors safetensors_path = hf_hub_download(repo_id, "model.safetensors", cache_dir=local_dir) else: print(f"Using cached model from {local_dir}") config_path = os.path.join(local_dir, "config.json") safetensors_path = os.path.join(local_dir, "model.safetensors") # Load config with open(config_path, 'r') as f: config = json.load(f) # Load weights to inspect expert count weights = load_file(safetensors_path) # Infer number of experts from checkpoint weights NUM_EXPERTS = weights['expert_routers.0.gate.weight'].shape[0] print(f"Inferred number of experts from checkpoint: {NUM_EXPERTS}") # Update config with inferred value config['num_experts'] = NUM_EXPERTS # Use config values (with updated num_experts) to initialize model global VOCAB_SIZE, DIM, NUM_LAYERS, HEADS, MAX_RECURSIONS VOCAB_SIZE = config['vocab_size'] DIM = config['dim'] NUM_LAYERS = config['num_layers'] HEADS = config['num_heads'] MAX_RECURSIONS = config['max_recursion'] # Create model with CORRECTED expert count model = QuantizedMoRModel( vocab_size=VOCAB_SIZE, dim=DIM, num_layers=NUM_LAYERS, num_heads=HEADS, max_recursion=MAX_RECURSIONS, num_experts=NUM_EXPERTS # Now matches checkpoint ) model.load_state_dict(weights) return model # Initialize with default values (will be overridden by config) VOCAB_SIZE = 10000 DIM = 1536 NUM_LAYERS = 6 HEADS = 8 BATCH_SIZE = 32 SEQ_LEN = 512 MAX_RECURSIONS = 4 learn_rate = 5e-5 EPOCHS = 3 NUM_EXPERTS = 12 GRAD_ACCUM_STEPS = 4 # Gradient accumulation steps # ---------------------- # Dataset Preparation # ---------------------- def prepare_datasets(file_path, tokenizer, seq_len=SEQ_LEN, val_split=0.05): print("Preparing datasets with tokenizer...") # Read text file with open(file_path, 'r', encoding='utf-8') as f: text = f.read() # Tokenize in chunks to avoid memory issues chunk_size = 500000 # characters per chunk chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)] encoded_chunks = [] for chunk in tqdm(chunks, desc="Tokenizing text chunks"): # Tokenize without special tokens encoding = tokenizer.encode(chunk, add_special_tokens=False) input_ids = torch.tensor(encoding) encoded_chunks.append(input_ids) # Concatenate all tokenized chunks encoded = torch.cat(encoded_chunks) total_tokens = len(encoded) split_idx = int(total_tokens * (1 - val_split)) # Create datasets train_dataset = TextDataset(encoded[:split_idx], seq_len) val_dataset = TextDataset(encoded[split_idx:], seq_len) print(f"Training samples: {len(train_dataset)}") print(f"Validation samples: {len(val_dataset)}") print(f"Total tokens: {total_tokens}") return train_dataset, val_dataset class TextDataset(Dataset): def __init__(self, encoded_data, seq_len=SEQ_LEN): self.encoded = encoded_data self.seq_len = seq_len def __len__(self): return len(self.encoded) // self.seq_len def __getitem__(self, idx): start = idx * self.seq_len end = start + self.seq_len + 1 segment = self.encoded[start:end] return segment[:-1].clone(), segment[1:].clone() # ---------------------- # MoR Model Components # ---------------------- class ExpertChoiceRouter(nn.Module): """Expert Choice Routing: Experts select top-k tokens""" def __init__(self, dim, num_experts, k=2): super().__init__() self.num_experts = num_experts self.k = k self.gate = nn.Linear(dim, num_experts, bias=False) def forward(self, x): scores = self.gate(x) expert_weights, expert_indices = torch.topk(scores, self.k, dim=-1) return expert_weights.softmax(dim=-1), expert_indices # ---------------------- # 4-bit Quantization Utilities # ---------------------- class Quantizer4Bit(nn.Module): def __init__(self): super().__init__() @staticmethod def quantize(tensor): max_val = tensor.abs().max() scale = max_val / 7.5 if max_val > 1e-8 else 1.0 quantized = torch.clamp(torch.round(tensor / scale), -8, 7) return quantized.to(torch.int8), scale @staticmethod def dequantize(quantized, scale): return quantized.float() * scale def init_weights(module): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) 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) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) # ---------------------- # MoR Model Components with Quantization # ---------------------- class QuantizedRecursiveTransformerBlock(nn.Module): def __init__(self, dim, num_heads, ffn_expansion=4): super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.q_proj = nn.Linear(dim, dim) self.k_proj = nn.Linear(dim, dim) self.v_proj = nn.Linear(dim, dim) self.attn_out = nn.Linear(dim, dim) self.ffn = nn.Sequential( nn.Linear(dim, ffn_expansion * dim), nn.GELU(), nn.Linear(ffn_expansion * dim, dim) ) self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) def forward(self, x): return checkpoint(self._forward, x, use_reentrant=False) def _forward(self, x): residual = x x = self.norm1(x) q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) k_quant, k_scale = Quantizer4Bit.quantize(k) v_quant, v_scale = Quantizer4Bit.quantize(v) k = Quantizer4Bit.dequantize(k_quant, k_scale) v = Quantizer4Bit.dequantize(v_quant, v_scale) B, T, _ = q.shape q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) attn = attn.softmax(dim=-1) attn_out = (attn @ v).transpose(1, 2).contiguous().view(B, T, self.dim) attn_out = self.attn_out(attn_out) x = residual + attn_out x = x + self.ffn(self.norm2(x)) return x class RecursionDepthRouter(nn.Module): def __init__(self, dim, max_depth=4): super().__init__() self.max_depth = max_depth self.router = nn.Sequential( nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, max_depth) ) for layer in self.router: if isinstance(layer, nn.Linear): nn.init.xavier_uniform_(layer.weight) nn.init.zeros_(layer.bias) def forward(self, x): x_pooled = x.mean(dim=(0, 1)) router_logits = self.router(x_pooled) return router_logits.softmax(dim=-1) # ---------------------- # Main MoR Architecture # ---------------------- class QuantizedMoRModel(nn.Module): def __init__(self, vocab_size, dim=DIM, num_layers=NUM_LAYERS, num_heads=HEADS, max_recursion=MAX_RECURSIONS, num_experts=NUM_EXPERTS): super().__init__() self.dim = dim self.max_recursion = max_recursion self.num_experts = num_experts self.embedding = nn.Embedding(vocab_size, dim) self.pos_embed = nn.Embedding(2048, dim) self.init_layers = nn.ModuleList([ QuantizedRecursiveTransformerBlock(dim, num_heads) for _ in range(2) ]) self.cycle_depth = 3 self.recursive_blocks = nn.ModuleList([ QuantizedRecursiveTransformerBlock(dim, num_heads) for _ in range(self.cycle_depth) ]) self.recursion_routers = nn.ModuleList([ RecursionDepthRouter(dim, max_depth=max_recursion) for _ in range(num_layers - 4) ]) self.expert_routers = nn.ModuleList([ ExpertChoiceRouter(dim, num_experts) for _ in range(max_recursion) ]) self.final_layers = nn.ModuleList([ QuantizedRecursiveTransformerBlock(dim, num_heads) for _ in range(2) ]) self.ln_f = nn.LayerNorm(dim) self.head = nn.Linear(dim, vocab_size, bias=False) def forward(self, x): pos = torch.arange(0, x.shape[1], device=x.device) x = self.embedding(x) * 0.02 x = x + self.pos_embed(pos) for layer in self.init_layers: x = layer(x) * 0.8 batch_size, seq_len, _ = x.shape recursion_outputs = [] for router in self.recursion_routers: depth_probs = router(x) depth = torch.multinomial(depth_probs, 1).item() expert_weights, expert_indices = self.expert_routers[depth](x) full_weights = torch.zeros((batch_size, seq_len, self.num_experts), device=x.device) full_weights.scatter_(2, expert_indices, expert_weights) expert_outputs = [] for expert_idx in range(self.num_experts): expert_x = x * full_weights[:, :, expert_idx].unsqueeze(-1) out = self.recursive_blocks[depth % self.cycle_depth](expert_x) expert_outputs.append(out) x = sum(expert_outputs) recursion_outputs.append(x) if recursion_outputs: x = torch.stack(recursion_outputs).mean(dim=0) for layer in self.final_layers: x = layer(x) x = self.ln_f(x) logits = self.head(x) return logits # ---------------------- # Learning Rate Scheduler # ---------------------- def get_lr(current_step, total_steps, warmup_steps, max_lr): if current_step < warmup_steps: return max_lr * (current_step / warmup_steps) else: decay_ratio = (current_step - warmup_steps) / (total_steps - warmup_steps) return max_lr * 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # ---------------------- # Training Loop with Validation # ---------------------- def train_model(): # Load pre-trained tokenizer tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL) global VOCAB_SIZE VOCAB_SIZE = tokenizer.vocab_size # Update from actual tokenizer # Prepare datasets train_dataset, val_dataset = prepare_datasets("input.txt", tokenizer, SEQ_LEN, val_split=0.05) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True) # Load pre-trained model model = load_model_from_hub(MODEL) # Fixed to use MoR-v1 # Parameter counting total_params = sum(p.numel() for p in model.parameters()) print(f"Model Parameters: {total_params/1e6:.2f}M") # Optimizer optimizer = torch.optim.AdamW(model.parameters(), lr=learn_rate, weight_decay=0.01) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # Mixed precision training scaler = GradScaler() # Training setup total_steps = EPOCHS * len(train_loader) warmup_steps = int(0.1 * total_steps) print(f"Total training steps: {total_steps}, Warmup steps: {warmup_steps}") # Training loop train_losses = [] val_losses = [] best_val_loss = float('inf') for epoch in range(EPOCHS): model.train() epoch_train_loss = 0 accumulated_loss = 0 optimizer.zero_grad() for step, (inputs, targets) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1} Training")): global_step = epoch * len(train_loader) + step current_lr = get_lr(global_step, total_steps, warmup_steps, learn_rate) for param_group in optimizer.param_groups: param_group['lr'] = current_lr inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True) with autocast(): logits = model(inputs) loss = F.cross_entropy( logits.view(-1, VOCAB_SIZE), targets.view(-1), ignore_index=0 ) / GRAD_ACCUM_STEPS scaler.scale(loss).backward() accumulated_loss += loss.item() * GRAD_ACCUM_STEPS if step % 100 == 0: print(f"Step {global_step}: Batch Loss={accumulated_loss:.4f}, LR={current_lr:.2e}") if (step + 1) % GRAD_ACCUM_STEPS == 0 or step == len(train_loader) - 1: scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad() epoch_train_loss += accumulated_loss accumulated_loss = 0 avg_train_loss = epoch_train_loss / len(train_loader) train_losses.append(avg_train_loss) # Validation model.eval() epoch_val_loss = 0 with torch.no_grad(): for inputs, targets in tqdm(val_loader, desc=f"Epoch {epoch+1} Validation"): inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True) with autocast(): logits = model(inputs) loss = F.cross_entropy( logits.view(-1, VOCAB_SIZE), targets.view(-1), ignore_index=0 ) epoch_val_loss += loss.item() avg_val_loss = epoch_val_loss / len(val_loader) val_losses.append(avg_val_loss) if avg_val_loss < best_val_loss: best_val_loss = avg_val_loss save_huggingface_model(model, tokenizer, "MoR-v1-continued") print(f"Saved new best model with val loss: {best_val_loss:.4f}") print(f"Epoch {epoch+1} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | LR: {current_lr:.2e}") # Plot training and validation plt.figure(figsize=(10, 5)) plt.plot(train_losses, label='Training Loss') plt.plot(val_losses, label='Validation Loss') plt.title("Training and Validation Loss") plt.xlabel("Epoch") plt.ylabel("Loss") plt.legend() plt.savefig("training_validation_loss_continued.png") # Save final model save_huggingface_model(model, tokenizer, "MoR-v1-continued") print("Training complete. Models saved.") # ---------------------- # Execution # ---------------------- if __name__ == "__main__": train_model()