#!/usr/bin/env python3 """ Distill GPT-2 (124M) into an FEM-LLM with piecewise linear mesh approximation. Benchmarks original vs. FEM-LLM speed and memory. """ import torch import torch.nn as nn import numpy as np from transformers import GPT2Model, GPT2Config, GPT2LMHeadModel from datasets import load_dataset from sklearn.cluster import MiniBatchKMeans from tqdm import tqdm import time from collections import defaultdict # --------------------------- # 1. Load GPT-2 (124M) and a small corpus # --------------------------- print("Loading GPT-2 124M...") device = "cuda" if torch.cuda.is_available() else "cpu" model_orig = GPT2LMHeadModel.from_pretrained("gpt2").to(device) model_orig.eval() config = model_orig.config hidden_size = config.n_embd # 768 num_layers = config.n_layer # 12 # Use a tiny slice of WikiText-2 for distillation (first 5000 tokens) print("Loading dataset (WikiText-2, 5000 tokens)...") dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train") text = " ".join(dataset["text"][:200]) # ~5000 tokens inputs = model_orig.tokenizer(text, return_tensors="pt", truncation=True, max_length=1024) input_ids = inputs["input_ids"].to(device) seq_len = input_ids.shape[1] # --------------------------- # 2. Collect hidden states and FFN outputs for each layer # --------------------------- # We'll register hooks to capture: # - Input to each FFN (i.e., hidden state after attention + residual) # - Output of each FFN (before residual) ffn_inputs = defaultdict(list) # layer_idx -> list of [batch, seq, d] (flattened) ffn_outputs = defaultdict(list) def make_hook(layer_idx, is_input): def hook(module, inp, out): # inp is tuple; first element is the input to the module if is_input: # For FFN input: the tensor before the FFN (after attention + residual) x = inp[0].detach().cpu().numpy() ffn_inputs[layer_idx].append(x.reshape(-1, hidden_size)) else: # FFN output (after activation and second linear) y = out.detach().cpu().numpy() ffn_outputs[layer_idx].append(y.reshape(-1, hidden_size)) return hook # Register hooks on each GPT-2 block's MLP (which is the FFN) for i, block in enumerate(model_orig.transformer.h): # Input to the MLP: the hidden state right before the MLP module block.mlp.register_forward_pre_hook(make_hook(i, is_input=True)) # Output of the MLP: the tensor after the MLP (before residual addition) block.mlp.register_forward_hook(make_hook(i, is_input=False)) # Forward pass to collect data print("Collecting hidden states and FFN outputs...") with torch.no_grad(): _ = model_orig(input_ids) # Concatenate all collected data per layer for layer in range(num_layers): if ffn_inputs[layer]: ffn_inputs[layer] = np.concatenate(ffn_inputs[layer], axis=0) ffn_outputs[layer] = np.concatenate(ffn_outputs[layer], axis=0) else: raise ValueError(f"No data collected for layer {layer}") print(f"Collected {ffn_inputs[0].shape[0]} samples per layer (seq * batch).") # --------------------------- # 3. Build mesh (cluster hidden states) for each layer independently # --------------------------- num_nodes = 512 # K = number of mesh nodes (coarse but effective) print(f"Building mesh with K={num_nodes} nodes per layer using MiniBatchKMeans...") kmeans_models = [] node_centers = [] # list of [K, d] per layer for layer in tqdm(range(num_layers), desc="Clustering layers"): kmeans = MiniBatchKMeans(n_clusters=num_nodes, batch_size=10000, random_state=42, n_init=3) kmeans.fit(ffn_inputs[layer]) kmeans_models.append(kmeans) node_centers.append(kmeans.cluster_centers_) # --------------------------- # 4. For each node, fit a linear regression: a_node * x + b_node ≈ FFN_output(x) # --------------------------- print("Fitting linear regressors per node...") node_coeffs = [] # list of dict: node_idx -> (a, b) where a is [d], b scalar for layer in range(num_layers): X = ffn_inputs[layer] # (N, d) Y = ffn_outputs[layer] # (N, d) # we actually need to predict full d-dim vector? # Wait: FFN output is a vector of dimension d (768). We can fit a separate linear map per output dim, # but a full matrix per node would be d*d parameters -> too many. Instead we treat each output dimension # independently (like a multi‑output linear regression with shared input). # Or we can use a single linear layer: y = A x + b, where A is d×d. But that's dense again. # To stay true to "O(d) per token", we approximate FFN(x) component‑wise: each output dimension i # has its own a_i (vector) and b_i scalar. Then total per node is d*(d+1) parameters, which is large. # But in the theory, the element function is a_e * x + b_e where a_e is a *row vector* -> output is scalar. # That would force us to predict a single scalar, not a vector. Hmm. # Correction: The FEM-LLM replaces *each* linear layer with an element‑wise linear map that outputs a # vector of the same dimension. That is a_e is a matrix (d x d) and b_e is a vector (d). Then inference # becomes a matrix‑vector product per element: O(d^2) again, not good. # # The "O(d)" claim in the theory comes from assuming the output is a single value (like a scalar field). # For LLMs, the hidden dimension is large. To keep O(d) we must use a **diagonal** approximation: # y_i = a_i * x_i + b_i (i per dimension) # i.e., element‑wise scaling and shift. This outputs a vector and costs O(d) per node. # # Let's implement diagonal linear maps per node. # Assign each sample to a node labels = kmeans_models[layer].labels_ # (N,) # For each node, compute diagonal regression: for each dimension d_idx, predict y_dim from x_dim # We'll store a (num_nodes, d) and b (num_nodes, d) a_node = np.zeros((num_nodes, hidden_size)) b_node = np.zeros((num_nodes, hidden_size)) for node_id in range(num_nodes): idxs = np.where(labels == node_id)[0] if len(idxs) < 100: # not enough samples: use global mean # fallback: use identity a_node[node_id] = 1.0 b_node[node_id] = 0.0 continue X_node = X[idxs] # (n, d) Y_node = Y[idxs] # (n, d) # For each dimension independently: solve y = a*x + b # Using least squares: stack [x, 1] for dim in range(hidden_size): x_dim = X_node[:, dim].reshape(-1, 1) y_dim = Y_node[:, dim].reshape(-1, 1) # Build design matrix [x, 1] A = np.hstack([x_dim, np.ones_like(x_dim)]) coeff, _, _, _ = np.linalg.lstsq(A, y_dim, rcond=None) a_node[node_id, dim] = coeff[0, 0] b_node[node_id, dim] = coeff[1, 0] node_coeffs.append((torch.tensor(a_node, dtype=torch.float32), torch.tensor(b_node, dtype=torch.float32))) print("Linear regressors fitted.") # --------------------------- # 5. Build FEM-LLM model that uses these piecewise linear FFNs # --------------------------- class FEMLinearFFN(nn.Module): """Element‑wise diagonal linear map: output = a * input + b, where a,b are per‑node and per‑dimension.""" def __init__(self, node_centers, a, b, kmeans): super().__init__() self.register_buffer("node_centers", torch.tensor(node_centers, dtype=torch.float32)) self.register_buffer("a", a) # (K, d) self.register_buffer("b", b) # (K, d) self.kmeans = kmeans # scikit model for assignment def forward(self, x): # x: (batch, seq, d) or (batch*d, d) orig_shape = x.shape x_flat = x.view(-1, hidden_size).cpu().numpy() # Assign each vector to nearest node node_ids = self.kmeans.predict(x_flat) # (N,) node_ids = torch.tensor(node_ids, device=x.device) # Lookup a and b a_e = self.a[node_ids] # (N, d) b_e = self.b[node_ids] # Compute output = a_e * x_flat + b_e (elementwise) y_flat = a_e * x.view(-1, hidden_size) + b_e return y_flat.view(orig_shape) class FEMGPT2Block(nn.Module): """A single transformer block where the MLP is replaced by FEMLinearFFN.""" def __init__(self, orig_block, node_centers, a, b, kmeans): super().__init__() # Copy the attention and layernorms from original block (frozen) self.ln_1 = orig_block.ln_1 self.attn = orig_block.attn self.ln_2 = orig_block.ln_2 self.mlp = FEMLinearFFN(node_centers, a, b, kmeans) def forward(self, x, *args, **kwargs): # Same as original block but using FEM MLP residual = x x = self.ln_1(x) x = self.attn(x, *args, **kwargs)[0] # GPT2Attention returns (attn_out, present) x = residual + x residual = x x = self.ln_2(x) x = self.mlp(x) x = residual + x return (x,) class FEMGPT2(nn.Module): """Wrapper that replaces all MLPs in GPT-2 with FEM piecewise linear approximations.""" def __init__(self, orig_model, per_layer_data): super().__init__() self.transformer = orig_model.transformer self.lm_head = orig_model.lm_head # Replace each block's mlp for i, block in enumerate(self.transformer.h): node_centers, a, b, kmeans = per_layer_data[i] new_block = FEMGPT2Block(block, node_centers, a, b, kmeans) self.transformer.h[i] = new_block def forward(self, input_ids, **kwargs): return self.lm_head(self.transformer(input_ids, **kwargs).last_hidden_state) # Build FEM model from original per_layer_data = [] for layer in range(num_layers): node_centers = node_centers[layer] # (K, d) a, b = node_coeffs[layer] kmeans = kmeans_models[layer] per_layer_data.append((node_centers, a, b, kmeans)) fem_model = FEMGPT2(model_orig, per_layer_data).to(device) fem_model.eval() # --------------------------- # 6. Benchmark: original GPT-2 vs FEM-LLM # --------------------------- test_input_ids = input_ids[:, :200] # first 200 tokens for test batch_size = 1 def benchmark(model, input_ids, num_runs=20, description=""): model.eval() # Warmup for _ in range(3): with torch.no_grad(): _ = model(input_ids) torch.cuda.synchronize() # Measure time start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() for _ in range(num_runs): with torch.no_grad(): _ = model(input_ids) end_event.record() torch.cuda.synchronize() elapsed_ms = start_event.elapsed_time(end_event) # milliseconds total_tokens = input_ids.shape[0] * input_ids.shape[1] * num_runs tokens_per_sec = total_tokens / (elapsed_ms / 1000.0) # Memory mem_alloc = torch.cuda.max_memory_allocated() / 1e6 # MB print(f"{description}: {tokens_per_sec:.1f} tokens/sec, peak memory {mem_alloc:.1f} MB") return tokens_per_sec, mem_alloc print("\n--- Benchmark ---") with torch.no_grad(): orig_tps, orig_mem = benchmark(model_orig, test_input_ids, num_runs=10, description="Original GPT-2") fem_tps, fem_mem = benchmark(fem_model, test_input_ids, num_runs=10, description="FEM-LLM (distilled)") print("\n--- Results ---") print(f"Speedup: {fem_tps / orig_tps:.2f}x") print(f"Memory reduction: {orig_mem / fem_mem:.2f}x (peak allocation)")