{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Architecture Analyzer for Transformer Models\n",
    "\n",
    "This tool analyzes transformer model architectures for GPU hardware efficiency based on the paper \"The Case for Co-Designing Model Architectures with Hardware\" (Anthony et al., 2024): https://arxiv.org/abs/2401.14489\n",
    "\n",
    "The `ArchitectureAnalyzer` class evaluates transformer architectures against known hardware constraints to identify inefficient model dimensions and provide optimization suggestions. It focuses on making model dimensions align with GPU hardware requirements, particularly for optimal Tensor Core usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "GPU EFFICIENCY ANALYSIS: My Model\n",
      "================================================================================\n",
      "Model Configuration:\n",
      "  Model Type:          encoder_only\n",
      "  Hidden Size:         2560\n",
      "  Attention Heads:     32\n",
      "  Head Dimension:      80\n",
      "  Hidden Layers:       32\n",
      "  Vocabulary Size:     50257\n",
      "  Intermediate Size:   10240\n",
      "  Tensor Parallel:     1\n",
      "\n",
      "Efficiency Score: \u001b[93m78.8/100\u001b[0m (Good)\n",
      "\n",
      "Identified Issues:\n",
      "  1. Head dimension (80) is not divisible by 64\n",
      "  2. Vocabulary size (50257) is not divisible by 64\n",
      "\n",
      "Optimization Suggestions:\n",
      "  1. Change attention heads from 32 to 1 to get head dimension of 2560\n",
      "  2. Change attention heads from 32 to 2 to get head dimension of 1280\n",
      "  3. Change attention heads from 32 to 4 to get head dimension of 640\n",
      "  4. Change attention heads from 32 to 5 to get head dimension of 512\n",
      "  5. Change attention heads from 32 to 8 to get head dimension of 320\n",
      "  6. Change attention heads from 32 to 10 to get head dimension of 256\n",
      "  7. Change attention heads from 32 to 20 to get head dimension of 128\n",
      "  8. Pad vocabulary size from 50257 to 50304 (+47 tokens)\n",
      "\n",
      "Note: These recommendations aim to optimize GPU compute efficiency\n",
      "      without consideration for model quality or convergence.\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "from lobster.model.utils import Architecture, ArchitectureAnalyzer, GPUType, ModelType\n",
    "\n",
    "architecture = Architecture(\n",
    "    hidden_size=2560,\n",
    "    num_attention_heads=32, \n",
    "    num_hidden_layers=32,\n",
    "    vocab_size=50257,\n",
    "    model_type=ModelType.ENCODER_ONLY,\n",
    "    name=\"My Model\",\n",
    ")\n",
    "analyzer = ArchitectureAnalyzer(architecture, gpu_type = GPUType.A100)\n",
    "\n",
    "result = analyzer.analyze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explanation\n",
    "\n",
    "### 1. Head Dimension Alignment (Most Critical)\n",
    "- **Rule**: The head dimension (`hidden_size / num_attention_heads`) should be divisible by 64 (for A100 GPUs)\n",
    "- **Reason**: This ensures optimal Tensor Core usage, as Tensor Cores process matrix multiplications in blocks and work best with dimensions that are multiples of 64 (for A100)\n",
    "- **Penalty**: 25% score reduction when violated\n",
    "- **Detection**: `_check_head_dimension()` method\n",
    "- **Suggestions**: \n",
    "  - Decrease number of heads to get aligned head dimension\n",
    "  - Alternatively, increase the hidden size to get aligned head dimension\n",
    "\n",
    "### 2. Hidden Dimension Alignment\n",
    "- **Rule**: The hidden size should be divisible by 64\n",
    "- **Reason**: Optimal Tensor Core usage for matrix operations\n",
    "- **Penalty**: 20% score reduction when violated\n",
    "- **Detection**: `_check_hidden_dimension()` method\n",
    "- **Suggestion**: Round up to the nearest multiple of 64\n",
    "\n",
    "### 3. Vocabulary Size Alignment\n",
    "- **Rule**: The vocabulary size should be divisible by 64\n",
    "- **Reason**: Improves efficiency of embedding and classification layers\n",
    "- **Penalty**: 15% score reduction when violated\n",
    "- **Detection**: `_check_vocab_size()` method\n",
    "- **Suggestion**: Pad vocabulary size to the nearest multiple of 64\n",
    "\n",
    "### 4. Intermediate Size Alignment\n",
    "- **Rule**: The intermediate size (in MLP layers) should be divisible by 64\n",
    "- **Rule Extension**: For SwiGLU activations (which often use 8/3 coefficient), adjust the coefficient to get an aligned intermediate size\n",
    "- **Reason**: Optimize MLP layer matrix operations\n",
    "- **Detection**: `_check_intermediate_size()` method\n",
    "- **Suggestion**: \n",
    "  - For standard MLPs: Adjust to nearest multiple of 64\n",
    "  - For SwiGLU: Use a better coefficient that results in an aligned intermediate size\n",
    "\n",
    "### 5. Tensor Parallelism Rules\n",
    "When using tensor parallelism (splitting across multiple GPUs):\n",
    "- **Rule 1**: Hidden size should be divisible by tensor parallel size\n",
    "- **Penalty**: 15% score reduction when violated\n",
    "- **Rule 2**: `(batch_size * num_attention_heads)` should be divisible by tensor parallel size\n",
    "- **Penalty**: 10% score reduction when violated\n",
    "- **Detection**: `_check_tensor_parallelism()` method\n",
    "- **Suggestions**:\n",
    "  - Change tensor parallel size to a divisor of the hidden size\n",
    "  - Adjust attention heads for better tensor parallelism efficiency\n",
    "\n",
    "## General Recommendations\n",
    "\n",
    "1. **Batch size**: Make batch size as large as possible\n",
    "2. **Powers of two**: `batch_size * sequence_length`, `hidden_size / num_attention_heads`, and `hidden_size / tensor_parallel_size` should ideally be divisible by a power of two (ideally 64 or higher)\n",
    "3. **Tensor parallelism**: Keep tensor parallel size as small as possible while satisfying memory constraints\n",
    "4. **Attention heads**: Fewer, larger heads are often more efficient than many small heads\n",
    "5. **Layer count**: The number of layers should be divisible by the number of pipeline parallel stages (if using pipeline parallelism)\n",
    "\n",
    "## Efficiency Score Calculation\n",
    "\n",
    "The efficiency score starts at 100 and gets reduced based on violations:\n",
    "- Head dimension not aligned with 64: Up to -25\n",
    "- Hidden size not divisible by 64: -20\n",
    "- Vocabulary size not divisible by 64: -15\n",
    "- Tensor parallelism issues: Up to -25 combined"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
