import gradio as gr import torch import os import sys import json import gc from pathlib import Path import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from typing import Dict, List, Optional, Tuple import numpy as np from PIL import Image import io # Add the ffg_experiment_suite to the path sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'ffg_experiment_suite')) # Import modules from the surgeon toolkit from src.models import load_assets from src.grafting import fast_fisher_graft, magnitude_graft, fish_mask_graft from src.analysis import (_create_sparsity_distribution_plot, _set_publication_fonts, load_masks_from_run) # Constants AVAILABLE_MODELS = [ { "name": "Math Reasoning", "base": "meta-llama/Meta-Llama-3.1-8B", "finetuned": "pmahdavi/Llama-3.1-8B-math-reasoning", "optimizer_states": "pmahdavi/Llama-3.1-8B-math-reasoning:export/exp_avg_sq.safetensors" }, { "name": "Coding", "base": "meta-llama/Meta-Llama-3.1-8B", "finetuned": "pmahdavi/Llama-3.1-8B-coding-tulu3-ebs128-lr5e6-wsdcr0p4", "optimizer_states": "pmahdavi/Llama-3.1-8B-coding-tulu3-ebs128-lr5e6-wsdcr0p4:export_full_state_checkpoint-1100/exp_avg_sq.safetensors" }, { "name": "Instruction Following", "base": "meta-llama/Meta-Llama-3.1-8B", "finetuned": "pmahdavi/Llama-3.1-8B-precise-if", "optimizer_states": "pmahdavi/Llama-3.1-8B-precise-if:export/exp_avg_sq.safetensors" }, { "name": "General", "base": "meta-llama/Meta-Llama-3.1-8B", "finetuned": "pmahdavi/Llama-3.1-8B-general", "optimizer_states": "pmahdavi/Llama-3.1-8B-general:export/exp_avg_sq.safetensors" }, { "name": "Knowledge Recall", "base": "meta-llama/Meta-Llama-3.1-8B", "finetuned": "pmahdavi/Llama-3.1-8B-knowledge-recall", "optimizer_states": "pmahdavi/Llama-3.1-8B-knowledge-recall:export/exp_avg_sq.safetensors" } ] class FFGMaskExplorer: def __init__(self): self.current_masks = None self.current_stats = None def generate_masks(self, model_selection: str, sparsity_ratio: float, grafting_method: str, device_type: str, progress=gr.Progress()): """Generate masks for a single model configuration.""" # Find the selected model config model_config = None for model in AVAILABLE_MODELS: if model["name"] == model_selection: model_config = model break if not model_config: return None, None, "Model not found!" progress(0.1, desc="Loading models...") try: # Prepare config for load_assets config = { "base_model_id": model_config["base"], "finetuned_model_id": model_config["finetuned"], "optimizer_states_file": model_config["optimizer_states"], "device": device_type.lower(), "dtype": "bfloat16" } # Load models and optimizer states pretrained_model, finetuned_model, optimizer_v_state, tokenizer = load_assets(config) progress(0.5, desc="Generating masks...") # Generate masks based on method if grafting_method == "Fast Fisher (FFG)": grafted_model, stats_dict, masks_dict = fast_fisher_graft( pretrained_model, finetuned_model, optimizer_v_state, sparsity_ratio ) elif grafting_method == "Magnitude": grafted_model, stats_dict, masks_dict = magnitude_graft( pretrained_model, finetuned_model, sparsity_ratio ) elif grafting_method == "Fish-Mask": grafted_model, stats_dict, masks_dict = fish_mask_graft( pretrained_model, finetuned_model, optimizer_v_state, sparsity_ratio ) else: raise ValueError(f"Unknown grafting method: {grafting_method}") # Store results self.current_masks = masks_dict self.current_stats = stats_dict progress(0.8, desc="Creating visualizations...") # Generate visualizations viz_images = self.create_basic_visualizations(masks_dict, stats_dict) # Clean up memory del pretrained_model, finetuned_model, grafted_model if optimizer_v_state is not None: del optimizer_v_state gc.collect() torch.cuda.empty_cache() progress(1.0, desc="Complete!") # Format statistics for display stats_text = self.format_statistics(stats_dict) return viz_images, stats_text, "Success!" except Exception as e: gc.collect() torch.cuda.empty_cache() return None, None, f"Error: {str(e)}" def create_basic_visualizations(self, masks_dict: Dict, stats_dict: Dict) -> List[Image.Image]: """Create basic visualizations from the masks.""" images = [] # Set publication fonts _set_publication_fonts(scale_factor=1.0) # 1. Overall statistics plot fig, ax = plt.subplots(figsize=(10, 6)) stats_data = { 'Total Parameters': stats_dict['total_params'], 'Kept Parameters': stats_dict['kept_params'], 'Pruned Parameters': stats_dict['total_params'] - stats_dict['kept_params'] } ax.bar(stats_data.keys(), stats_data.values(), color=['blue', 'green', 'red']) ax.set_ylabel('Number of Parameters') ax.set_title(f'Grafting Statistics (Sparsity: {stats_dict["final_sparsity"]:.2%})') # Add value labels on bars for i, (key, value) in enumerate(stats_data.items()): ax.text(i, value, f'{value:,}', ha='center', va='bottom') plt.tight_layout() images.append(self.fig_to_image(fig)) plt.close(fig) # 2. Layer-wise sparsity plot fig, ax = plt.subplots(figsize=(14, 8)) # Extract layer information layer_sparsities = [] layer_names = [] for name, mask in masks_dict.items(): if mask is not None and mask.numel() > 0: sparsity = 1.0 - mask.float().mean().item() layer_sparsities.append(sparsity) # Shorten layer name for display short_name = name.replace('model.layers.', 'L').replace('.weight', '') if len(short_name) > 20: short_name = short_name[:17] + '...' layer_names.append(short_name) # Limit to first 50 layers for clarity if len(layer_names) > 50: layer_names = layer_names[:50] layer_sparsities = layer_sparsities[:50] ax.barh(range(len(layer_names)), layer_sparsities, color='skyblue') ax.set_yticks(range(len(layer_names))) ax.set_yticklabels(layer_names, fontsize=8) ax.set_xlabel('Sparsity Ratio') ax.set_title('Layer-wise Sparsity Distribution') ax.set_xlim(0, 1) plt.tight_layout() images.append(self.fig_to_image(fig)) plt.close(fig) # 3. Sample mask heatmap (first few layers) num_samples = min(4, len(masks_dict)) fig, axes = plt.subplots(2, 2, figsize=(12, 10)) axes = axes.flatten() for idx, (name, mask) in enumerate(list(masks_dict.items())[:num_samples]): if mask is None or mask.numel() == 0: continue ax = axes[idx] # Convert mask to numpy and visualize mask_np = mask.cpu().float().numpy() # For 2D tensors, show directly if mask.ndim == 2: im = ax.imshow(mask_np, cmap='RdBu_r', aspect='auto', vmin=0, vmax=1) else: # For 1D or higher dim, reshape or show first 2D slice if mask.ndim == 1: # Reshape 1D to approximate square size = int(np.sqrt(mask.numel())) if size * size == mask.numel(): mask_np = mask_np.reshape(size, size) else: # Pad and reshape target_size = size + 1 padded = np.zeros(target_size * target_size) padded[:mask.numel()] = mask_np.flatten() mask_np = padded.reshape(target_size, target_size) im = ax.imshow(mask_np, cmap='RdBu_r', aspect='auto', vmin=0, vmax=1) else: # For higher dimensions, show a 2D slice mask_np = mask_np.reshape(mask_np.shape[0], -1)[:min(mask_np.shape[0], 512), :min(mask_np.shape[1], 512)] im = ax.imshow(mask_np, cmap='RdBu_r', aspect='auto', vmin=0, vmax=1) # Add colorbar plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) # Set title short_name = name.replace('model.layers.', 'L').replace('.weight', '') if len(short_name) > 30: short_name = short_name[:27] + '...' ax.set_title(short_name, fontsize=10) ax.set_xlabel('Dimension 1') ax.set_ylabel('Dimension 0') # Hide unused axes for idx in range(num_samples, len(axes)): axes[idx].axis('off') plt.suptitle('Sample Mask Visualizations (1=kept, 0=pruned)', fontsize=14) plt.tight_layout() images.append(self.fig_to_image(fig)) plt.close(fig) return images def fig_to_image(self, fig) -> Image.Image: """Convert matplotlib figure to PIL Image.""" buf = io.BytesIO() fig.savefig(buf, format='png', dpi=150, bbox_inches='tight') buf.seek(0) return Image.open(buf) def format_statistics(self, stats_dict: Dict) -> str: """Format statistics dictionary as readable text.""" lines = [ "### Grafting Statistics", f"- **Total Parameters**: {stats_dict['total_params']:,}", f"- **Kept Parameters**: {stats_dict['kept_params']:,}", f"- **Pruned Parameters**: {stats_dict['total_params'] - stats_dict['kept_params']:,}", f"- **Final Sparsity**: {stats_dict['final_sparsity']:.4f}", f"- **Threshold**: {stats_dict.get('threshold', 'N/A')}" ] return "\n".join(lines) # Initialize the explorer explorer = FFGMaskExplorer() # Create Gradio interface with gr.Blocks(title="FFG Mask Explorer", theme=gr.themes.Base()) as app: gr.Markdown(""" # 🔬 FFG Mask Explorer Interactive tool for generating and visualizing Fast Fisher Grafting (FFG) masks on fine-tuned language models. Based on the paper: [Harnessing Optimization Dynamics for Curvature-Informed Model Merging](https://arxiv.org/abs/2509.11167) ### How to use: 1. Select a pre-configured model or enter custom model IDs 2. Choose sparsity ratio (fraction of parameters to KEEP) 3. Select grafting method 4. Click Generate to create masks and visualizations """) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Model Configuration") model_dropdown = gr.Dropdown( choices=[m["name"] for m in AVAILABLE_MODELS], value=AVAILABLE_MODELS[0]["name"], label="Select Pre-configured Model", interactive=True ) with gr.Accordion("Custom Model (Advanced)", open=False): custom_base = gr.Textbox( label="Base Model ID", placeholder="e.g., meta-llama/Meta-Llama-3.1-8B" ) custom_finetuned = gr.Textbox( label="Fine-tuned Model ID", placeholder="e.g., username/model-name" ) custom_optimizer = gr.Textbox( label="Optimizer States Path", placeholder="e.g., username/model:export/exp_avg_sq.safetensors" ) sparsity_slider = gr.Slider( minimum=0.01, maximum=0.9, value=0.4, step=0.01, label="Sparsity Ratio (fraction to KEEP)", info="0.4 means keeping 40% of parameters" ) method_radio = gr.Radio( choices=["Fast Fisher (FFG)", "Magnitude", "Fish-Mask"], value="Fast Fisher (FFG)", label="Grafting Method", info="FFG uses optimizer second moments for importance" ) device_radio = gr.Radio( choices=["CUDA", "CPU"], value="CUDA", label="Device", info="CUDA recommended for faster processing" ) generate_btn = gr.Button("🚀 Generate Masks", variant="primary", size="lg") with gr.Column(scale=2): gr.Markdown("### Results") status_text = gr.Textbox(label="Status", interactive=False, value="Ready") with gr.Tabs(): with gr.TabItem("Visualizations"): gallery = gr.Gallery( label="Generated Visualizations", show_label=False, elem_id="gallery", columns=2, rows=2, object_fit="contain", height="auto" ) with gr.TabItem("Statistics"): stats_markdown = gr.Markdown("*Generate masks to see statistics*") with gr.Row(): download_masks_btn = gr.Button("💾 Download Masks", size="sm", interactive=False) download_viz_btn = gr.Button("📊 Download Visualizations", size="sm", interactive=False) # Event handlers def on_generate(model_selection, sparsity, method, device, progress=gr.Progress()): images, stats, status = explorer.generate_masks( model_selection, sparsity, method, device, progress ) # Enable download buttons if successful if images: return ( images, stats, status, gr.Button(interactive=True), gr.Button(interactive=True) ) else: return ( None, "*Generation failed*", status, gr.Button(interactive=False), gr.Button(interactive=False) ) generate_btn.click( fn=on_generate, inputs=[model_dropdown, sparsity_slider, method_radio, device_radio], outputs=[gallery, stats_markdown, status_text, download_masks_btn, download_viz_btn] ) gr.Markdown(""" --- ### About FFG Fast Fisher Grafting (FFG) uses the second moments from Adam optimizer to identify important parameters in fine-tuned models. This allows for more informed pruning compared to magnitude-based methods. ### Citation ```bibtex @misc{mahdavinia2025harnessingoptimizationdynamicscurvatureinformed, title={Harnessing Optimization Dynamics for Curvature-Informed Model Merging}, author={Pouria Mahdavinia and Hamed Mahdavi and Niloofar Mireshghallah and Mehrdad Mahdavi}, year={2025}, eprint={2509.11167}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2509.11167}, } ``` """) if __name__ == "__main__": app.launch()