Initial deployment: Add GPU Memory Calculator with Docker

- Add FastAPI web application
- Add all calculator modules (training, inference, multi-node)
- Configure Docker for Python 3.12
- Add requirements.txt with web dependencies
- Add Space README with metadata

.dockerignore

@@ -0,0 +1,92 @@
+# Git
+.git
+.gitignore
+.github
+
+# Docker
+Dockerfile
+.dockerignore
+
+# Python
+__pycache__
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+.venv/
+.env
+
+# Testing
+.pytest_cache/
+.coverage
+coverage.xml
+htmlcov/
+.tox/
+.mypy_cache/
+.ruff_cache/
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Claude
+.claude/
+.mcp.json
+
+# Documentation (source files included, but skip extras)
+docs/
+*.md
+!README.md
+
+# Project specific
+*.log
+.env
+.venv/
+
+# CI/CD
+CODE_OF_CONDUCT.md
+CONTRIBUTING.md
+MARKETING.md
+SECURITY.md
+CHANGELOG.md
+
+# Screenshots and images
+*.png
+*.jpg
+*.jpeg
+*.gif
+!screenshot.png
+
+# Test files
+tests/
+examples/
+configs/
+
+# MCP server config
+.mcp.json

Dockerfile

@@ -0,0 +1,40 @@
+# Dockerfile for Hugging Face Spaces
+# GPU Memory Calculator - FastAPI Web Application
+
+FROM python:3.12-slim
+
+# Set working directory
+WORKDIR /app
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PORT=7860
+
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        gcc \
+        && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better Docker layer caching
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy project files
+COPY . .
+
+# Install the package in editable mode
+RUN pip install --no-cache-dir -e .
+
+# Expose Hugging Face Spaces default port
+EXPOSE 7860
+
+# Health check endpoint
+HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
+    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:7860/').read()"
+
+# Run the FastAPI application with uvicorn
+CMD ["uvicorn", "web.app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,12 +1,62 @@
 ---
-title: Gpu Memory Calculator
-emoji: 😻
-colorFrom: gray
-colorTo: yellow
+title: GPU Memory Calculator
+emoji: 🎮
+colorFrom: blue
+colorTo: purple
 sdk: docker
 pinned: false
-license: apache-2.0
-short_description: Calculates GPU memory for training, inference, and more
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# GPU Memory Calculator
+
+Calculate GPU memory requirements for training and running Large Language Models (LLMs). Supports multiple training engines (PyTorch DDP, DeepSpeed ZeRO, Megatron-LM, FSDP), inference engines (HuggingFace, vLLM, TGI, TensorRT-LLM, SGLang), and multi-node training configurations.
+
+## Features
+
+- **Training Memory Calculation**: Calculate memory for PyTorch DDP, DeepSpeed ZeRO (0-3), Megatron-LM, FSDP, and hybrid approaches
+- **Inference Memory Calculation**: Estimate memory requirements for HuggingFace Transformers, vLLM, TGI, TensorRT-LLM, and SGLang
+- **Multi-Node Support**: Calculate network overhead for distributed training across multiple nodes
+- **Model Presets**: Pre-configured settings for popular models (LLaMA 2, GPT-3, Mixtral, GLM, Qwen, DeepSeek-MoE)
+- **Configuration Export**: Generate configs for Accelerate, Lightning, Axolotl, DeepSpeed, YAML, and JSON
+- **Batch Size Optimization**: Automatically find the maximum batch size that fits in GPU memory
+
+## Supported Training Engines
+
+- PyTorch DDP (Distributed Data Parallel)
+- DeepSpeed ZeRO (Stages 0-3) with CPU/NVMe offloading
+- Megatron-LM (Tensor + Pipeline Parallelism)
+- PyTorch FSDP (Fully Sharded Data Parallel)
+- Megatron-LM + DeepSpeed (Hybrid)
+
+## Supported Inference Engines
+
+- HuggingFace Transformers
+- vLLM (PagedAttention)
+- Text Generation Inference (TGI)
+- TensorRT-LLM
+- SGLang (RadixAttention)
+
+## How to Use
+
+1. **Select a preset model** or configure your own
+2. **Choose training/inference engine** and adjust parameters
+3. **Calculate** memory requirements instantly
+4. **Export** configurations to your preferred framework
+
+## Example Use Cases
+
+- Planning GPU requirements for LLM training
+- Optimizing batch sizes for your hardware
+- Comparing memory efficiency across engines
+- Estimating KV cache memory for inference
+- Calculating multi-node network overhead
+
+## Links
+
+- [GitHub Repository](https://github.com/George614/gpu-mem-calculator)
+- [Documentation](https://github.com/George614/gpu-mem-calculator/blob/main/README.md)
+
+## License
+
+MIT License - see [LICENSE](https://github.com/George614/gpu-mem-calculator/blob/main/LICENSE) for details.

cli/main.py

@@ -0,0 +1,399 @@
+"""CLI interface for GPU Memory Calculator."""
+
+import json
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import click
+
+if TYPE_CHECKING:
+    from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+    from gpu_mem_calculator.core.models import MemoryResult
+
+
+@click.group()
+@click.version_option(version="0.1.0")
+def main() -> None:
+    """GPU Memory Calculator for LLM Training.
+
+    Calculate GPU memory requirements for training Large Language Models
+    with various training engines (PyTorch DDP, DeepSpeed, Megatron-LM, FSDP).
+    """
+    pass
+
+
+@main.command()
+@click.option(
+    "--config",
+    "-c",
+    type=click.Path(exists=True),
+    help="Path to JSON configuration file",
+)
+@click.option(
+    "--preset",
+    "-p",
+    type=str,
+    help="Name of a preset model configuration",
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(),
+    help="Output file path (default: stdout)",
+)
+@click.option(
+    "--format",
+    "-f",
+    type=click.Choice(["json", "yaml", "table"]),
+    default="table",
+    help="Output format (default: table)",
+)
+def calculate(
+    config: str | None,
+    preset: str | None,
+    output: str | None,
+    format: Literal["json", "yaml", "table"],
+) -> None:
+    """Calculate GPU memory requirements from config file or preset.
+
+    Examples:
+        gpu-mem-calc calculate --config configs/llama2_7b.json
+        gpu-mem-calc calculate --preset llama2-7b
+        gpu-mem-calc calculate -p mixtral-8x7b --format json
+    """
+    if not config and not preset:
+        click.echo("Error: Either --config or --preset is required", err=True)
+        sys.exit(1)
+
+    if config and preset:
+        click.echo("Error: Cannot use both --config and --preset", err=True)
+        sys.exit(1)
+
+    try:
+        import tempfile
+
+        from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+
+        if preset:
+            # Load preset configuration
+            from gpu_mem_calculator.config.presets import get_preset_config
+
+            preset_config = get_preset_config(preset)
+            if preset_config is None:
+                click.echo(
+                    f"Error: Preset '{preset}' not found. "
+                    "Use 'gpu-mem-calc presets' to list available presets.",
+                    err=True,
+                )
+                sys.exit(1)
+
+            # Write preset to temp file for from_config_file
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+                json.dump(preset_config, f, indent=2)
+                temp_path = f.name
+
+            calculator = GPUMemoryCalculator.from_config_file(temp_path)
+            Path(temp_path).unlink()  # Clean up temp file
+        elif config:
+            calculator = GPUMemoryCalculator.from_config_file(config)
+        else:
+            # This should never happen due to the checks above
+            click.echo("Error: Either --config or --preset is required", err=True)
+            sys.exit(1)
+
+        result = calculator.calculate()
+
+        # Format output
+        if format == "json":
+            output_text = json.dumps(result.model_dump(mode="json"), indent=2)
+        elif format == "yaml":
+            try:
+                import yaml  # type: ignore[import-untyped]
+
+                output_text = yaml.dump(result.model_dump(mode="json"), default_flow_style=False)
+            except ImportError:
+                click.echo(
+                    "Error: YAML format requires PyYAML. Install with: pip install pyyaml",
+                    err=True,
+                )
+                sys.exit(1)
+        else:  # table
+            output_text = _format_result_as_table(result, calculator)
+
+        # Write output
+        if output:
+            Path(output).write_text(output_text)
+            click.echo(f"Results written to {output}")
+        else:
+            click.echo(output_text)
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.argument(
+    "params",
+    type=float,
+    required=True,
+)
+@click.option(
+    "--gpus",
+    "-g",
+    type=int,
+    default=1,
+    help="Number of GPUs (default: 1)",
+)
+@click.option(
+    "--gpu-mem",
+    "-m",
+    type=float,
+    default=80.0,
+    help="GPU memory in GB (default: 80.0)",
+)
+@click.option(
+    "--engine",
+    "-e",
+    type=click.Choice(["pytorch", "deepspeed", "megatron", "fsdp"]),
+    default="pytorch",
+    help="Training engine (default: pytorch)",
+)
+@click.option(
+    "--dtype",
+    "-d",
+    type=click.Choice(["fp32", "fp16", "bf16"]),
+    default="bf16",
+    help="Data type (default: bf16)",
+)
+def quick(
+    params: float,
+    gpus: int,
+    gpu_mem: float,
+    engine: str,
+    dtype: str,
+) -> None:
+    """Quick calculation from model size (in billions of parameters).
+
+    Example:
+        gpu-mem-calc quick 7 --gpus 8 --engine deepspeed
+    """
+    try:
+        from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+        from gpu_mem_calculator.core.models import (
+            DType,
+            EngineConfig,
+            EngineType,
+            GPUConfig,
+            ModelConfig,
+            ParallelismConfig,
+            TrainingConfig,
+        )
+
+        # Map engine string to EngineType
+        engine_map = {
+            "pytorch": EngineType.PYTORCH_DDP,
+            "deepspeed": EngineType.DEEPSPEED,
+            "megatron": EngineType.MEGATRON_LM,
+            "fsdp": EngineType.FSDP,
+        }
+
+        # Map dtype string to DType
+        dtype_map = {
+            "fp32": DType.FP32,
+            "fp16": DType.FP16,
+            "bf16": DType.BF16,
+        }
+
+        # Create a minimal config for quick calculation
+        # Estimate model architecture from parameter count
+        # Rough approximation based on typical transformer models
+        num_params = int(params * 1e9)
+
+        # Estimate hidden size and layers from param count
+        # These are rough approximations
+        if params <= 1:
+            hidden_size, num_layers = 768, 12
+        elif params <= 7:
+            hidden_size, num_layers = 4096, 32
+        elif params <= 13:
+            hidden_size, num_layers = 5120, 40
+        elif params <= 30:
+            hidden_size, num_layers = 6656, 60
+        elif params <= 65:
+            hidden_size, num_layers = 8192, 80
+        else:
+            hidden_size, num_layers = 12288, 96
+
+        model_config = ModelConfig(
+            name="quick-estimate",
+            num_parameters=num_params,
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            num_attention_heads=hidden_size // 128,
+            vocab_size=32000,
+            max_seq_len=2048,
+        )
+
+        training_config = TrainingConfig(
+            batch_size=1,
+            gradient_accumulation_steps=1,
+            dtype=dtype_map[dtype],
+        )
+
+        parallelism_config = ParallelismConfig(data_parallel_size=gpus)
+
+        engine_config = EngineConfig(
+            type=engine_map[engine],
+            zero_stage=2 if engine == "deepspeed" else None,
+        )
+
+        gpu_config = GPUConfig(num_gpus=gpus, gpu_memory_gb=gpu_mem)
+
+        calculator = GPUMemoryCalculator(
+            model_config=model_config,
+            training_config=training_config,
+            parallelism_config=parallelism_config,
+            engine_config=engine_config,
+            gpu_config=gpu_config,
+        )
+
+        result = calculator.calculate()
+
+        # Display results
+        click.echo(_format_result_as_table(result, calculator))
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.argument(
+    "config_path",
+    type=click.Path(exists=True),
+)
+def validate(config_path: str) -> None:
+    """Validate a configuration file.
+
+    Example:
+        gpu-mem-calc validate configs/my_config.json
+    """
+    try:
+        from gpu_mem_calculator.config import ConfigParser
+
+        ConfigParser.parse_full_config(config_path)
+        click.echo(f"✓ Configuration file '{config_path}' is valid")
+
+    except Exception as e:
+        click.echo(f"✗ Validation failed: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.option(
+    "--format",
+    "-f",
+    type=click.Choice(["list", "json", "table"]),
+    default="list",
+    help="Output format (default: list)",
+)
+def presets(format: str) -> None:
+    """List available model preset configurations.
+
+    Examples:
+        gpu-mem-calc presets
+        gpu-mem-calc presets --format table
+        gpu-mem-calc presets -f json
+    """
+    try:
+        from gpu_mem_calculator.config.presets import list_presets
+
+        all_presets = list_presets()
+
+        if not all_presets:
+            click.echo("No presets found.")
+            return
+
+        if format == "json":
+            click.echo(json.dumps(all_presets, indent=2))
+        elif format == "table":
+            from rich.console import Console
+            from rich.table import Table
+
+            console = Console()
+            table = Table(
+                title="Available Model Presets",
+                show_header=True,
+                header_style="bold magenta",
+            )
+            table.add_column("Preset Name", style="cyan", width=25)
+            table.add_column("Display Name", style="green", width=30)
+            table.add_column("Description", style="yellow")
+
+            for name, info in sorted(all_presets.items()):
+                table.add_row(name, info["display_name"], info["description"])
+
+            console.print(table)
+        else:  # list format
+            click.echo("Available model presets:\n")
+            for name, info in sorted(all_presets.items()):  # type: ignore[annotation-unchecked]
+                click.echo(f"  {name:25} - {info['display_name']}")
+                if info.get("description"):
+                    click.echo(f"{'':27}{info['description']}")
+                click.echo()
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+def _format_result_as_table(result: MemoryResult, calculator: "GPUMemoryCalculator") -> str:
+    """Format result as ASCII table."""
+    from rich.console import Console
+    from rich.table import Table
+
+    console = Console()
+
+    # Main results table
+    table = Table(
+        title="GPU Memory Calculation Results",
+        show_header=True,
+        header_style="bold magenta",
+    )
+    table.add_column("Metric", style="cyan", width=30)
+    table.add_column("Value", style="green")
+
+    # Memory results
+    table.add_row("Memory per GPU", f"{result.total_memory_per_gpu_gb:.2f} GB")
+    table.add_row("Total GPU Memory", f"{result.total_memory_all_gpus_gb:.2f} GB")
+    table.add_row("CPU Memory", f"{result.cpu_memory_gb:.2f} GB")
+    table.add_row("", "")  # Spacer
+
+    # Breakdown
+    table.add_row("Model Parameters", f"{result.breakdown.model_params_gb:.2f} GB")
+    table.add_row("Gradients", f"{result.breakdown.gradients_gb:.2f} GB")
+    table.add_row("Optimizer States", f"{result.breakdown.optimizer_states_gb:.2f} GB")
+    table.add_row("Activations", f"{result.breakdown.activations_gb:.2f} GB")
+    table.add_row("Overhead", f"{result.breakdown.overhead_gb:.2f} GB")
+    table.add_row("", "")  # Spacer
+
+    # Feasibility
+    status = "✓ Fits" if result.fits_on_gpu else "✗ OOM"
+    table.add_row("Status", status)
+    table.add_row("Memory Utilization", f"{result.memory_utilization_percent:.1f}%")
+    if result.recommended_batch_size:
+        table.add_row("Recommended Batch Size", str(result.recommended_batch_size))
+
+    # Capture table output
+    from io import StringIO
+
+    buffer = StringIO()
+    console.file = buffer
+    console.print(table)
+    return buffer.getvalue()
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,12 @@
+# GPU Memory Calculator - Requirements for Hugging Face Spaces
+
+# Core dependencies
+pydantic>=2.0.0
+click>=8.1.0
+pydantic-settings>=2.0.0
+rich>=13.0.0
+
+# Web dependencies
+fastapi>=0.100.0
+uvicorn[standard]>=0.23.0
+jinja2>=3.1.0

src/gpu_mem_calculator.egg-info/PKG-INFO

@@ -0,0 +1,720 @@
+Metadata-Version: 2.4
+Name: gpu-mem-calculator
+Version: 0.1.0
+Summary: GPU Memory Calculator for LLM Training
+Author: GPU Mem Calculator Team
+License: MIT
+Project-URL: Homepage, https://github.com/George614/gpu-mem-calculator
+Project-URL: Repository, https://github.com/George614/gpu-mem-calculator
+Project-URL: Issues, https://github.com/George614/gpu-mem-calculator/issues
+Keywords: gpu,memory,calculator,llm,large-language-model,training,deepspeed,megatron,pytorch,fsdp,transformer,machine-learning,deep-learning,distributed-training,zero-optimization
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: click>=8.1.0
+Requires-Dist: pydantic-settings>=2.0.0
+Requires-Dist: rich>=13.0.0
+Provides-Extra: web
+Requires-Dist: fastapi>=0.100.0; extra == "web"
+Requires-Dist: uvicorn[standard]>=0.23.0; extra == "web"
+Requires-Dist: jinja2>=3.1.0; extra == "web"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: black>=23.0.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: mypy>=1.5.0; extra == "dev"
+Dynamic: license-file
+
+# GPU Memory Calculator for LLM Training
+
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](CONTRIBUTING.md)
+
+A versatile Python application for calculating GPU memory requirements for training Large Language Models with support for multiple training engines including PyTorch DDP, DeepSpeed ZeRO, Megatron-LM, and FSDP.
+
+📖 **[Getting Started Guide](docs/GETTING_STARTED.md)** | 💬 **[FAQ](docs/FAQ.md)** | 🤝 **[Contributing](CONTRIBUTING.md)**
+
+<p align="center">
+  <img src="screenshot.png" alt="GPU Memory Calculator Screenshot" width="800">
+</p>
+
+## 🚀 Why Use This Tool?
+
+Training large language models requires careful memory planning. This calculator helps you:
+
+- **💰 Save costs** by determining the optimal GPU configuration before you start training
+- **⚡ Avoid OOM errors** by validating your training configuration fits in GPU memory
+- **📊 Compare strategies** across different training engines (DeepSpeed, Megatron, FSDP)
+- **🎯 Plan infrastructure** by knowing exactly how many GPUs you need
+- **📈 Scale efficiently** with detailed memory breakdowns for optimization
+
+Whether you're training a 7B parameter model on a single GPU or a 175B model across hundreds of GPUs, this tool provides accurate memory estimates based on proven formulas from DeepSpeed, Megatron-LM, and the latest research.
+
+## ✨ Features
+
+### Core Training Calculation
+- 🔧 **Multiple Training Engines**: Support for PyTorch DDP, DeepSpeed ZeRO (stages 1-3), Megatron-LM, Megatron+DeepSpeed, and PyTorch FSDP
+- 🖥️ **Dual Interface**: Both CLI and Web UI for flexible usage
+- 🎯 **Preset Models**: Quick-load configurations for popular models (LLaMA 2, GPT-3, etc.)
+- 📊 **Detailed Breakdown**: Memory breakdown by component (parameters, gradients, optimizer states, activations)
+- ✅ **Feasibility Analysis**: Check if your configuration fits on available GPU memory
+- ⚙️ **Easy Config**: JSON-based configuration files with human-readable parameter formats (e.g., "7B", "7000M")
+
+### 🆕 Inference Memory Calculation
+- 🚀 **Multi-Engine Support**: HuggingFace Transformers, vLLM, TGI, TensorRT-LLM
+- 💾 **KV Cache Optimization**: Quantization options (NONE, INT8, FP8, INT4)
+- 🔄 **Tensor Parallelism**: Automatic memory distribution across GPUs
+- 📈 **Throughput Estimation**: Tokens/second estimates for capacity planning
+- 🎯 **Batch Size Optimization**: Find maximum batch size for your hardware
+
+### 🆕 Multi-Node Training
+- 🌐 **Network Overhead Calculation**: AllReduce, AllGather, ReduceScatter, pipeline communication
+- 📡 **Interconnect Support**: InfiniBand, NVLink, Ethernet (10G/25G/100G/200G)
+- ⚡ **Hybrid Parallelism Optimization**: Automatic TP+PP+DP strategy optimization
+- 🔧 **ZeRO Stage Impact Analysis**: Compare communication overhead across ZeRO stages
+
+### 🆕 Framework Configuration Exporters
+- 📦 **Accelerate Export**: HuggingFace Accelerate config generation
+- ⚡ **Lightning Export**: PyTorch Lightning Trainer configuration
+- 🔥 **Axolotl Export**: YAML config for fine-tuning
+- 📄 **File Export**: Save to YAML/JSON formats
+- 🎛️ **Format Conversion**: Convert between different framework configs
+
+## 📦 Installation
+
+### Quick Start
+
+### Core Capabilities
+- **Multiple Training Engines**: Support for PyTorch DDP, DeepSpeed ZeRO (stages 0-3), Megatron-LM, Megatron+DeepSpeed, and PyTorch FSDP
+- **Dual Interface**: Both CLI and Web UI for flexible usage
+- **Preset Models**: Quick-load configurations for popular models (LLaMA 2, GPT-3, GLM, Mixtral, etc.)
+- **Detailed Breakdown**: Memory breakdown by component (parameters, gradients, optimizer states, activations)
+- **Feasibility Analysis**: Check if your configuration fits on available GPU memory
+- **Easy Config**: JSON-based configuration files with human-readable parameter formats (e.g., "7B", "7000M")
+
+### Web UI Enhancements
+- **Formula Explanations**: See exactly how memory is calculated with your values plugged in
+- **Real-time Validation**: Client-side validation prevents invalid configurations
+- **Smart Auto-calculation**: Optimized debouncing (1s) with minimum interval protection
+- **Export Capabilities**: Export to DeepSpeed config files, JSON, or copy to clipboard
+- **Batch Size Optimizer**: Automatically find maximum batch size that fits
+- **Comparison Mode**: Save and compare different configurations side-by-side
+- **Accessibility Features**: ARIA labels, keyboard navigation, colorblind-friendly charts
+
+### Advanced Features
+- **MoE Support**: Mixture of Experts models with configurable experts and top-k routing
+- **CPU/NVMe Offloading**: Offload optimizer states and parameters to CPU or NVMe storage
+- **Activation Checkpointing**: 5 levels from none to full checkpointing
+- **Sequence Parallelism**: Optimize memory for long sequences
+- **Result Caching**: Fast repeated calculations with built-in caching
+
+```bash
+pip install git+https://github.com/George614/gpu-mem-calculator.git
+```
+
+### From source
+
+```bash
+git clone https://github.com/George614/gpu-mem_calculator.git
+cd gpu_mem_calculator
+pip install -e .
+```
+
+### For Web UI support
+
+```bash
+pip install -e ".[web]"
+```
+
+### Development installation
+
+```bash
+pip install -e ".[dev]"
+```
+
+## 🎓 Use Cases
+
+### Research & Academia
+- Estimate GPU requirements for research projects before requesting compute resources
+- Plan multi-GPU training configurations for large-scale experiments
+- Compare memory efficiency of different training strategies
+
+### Industry & Production
+- Cost optimization: Choose the right GPU type and count for your training workload
+- Capacity planning: Forecast infrastructure needs for model development
+- Debugging: Diagnose OOM errors and optimize memory usage
+
+### Education & Learning
+- Understand how training configuration affects memory consumption
+- Learn about different distributed training strategies
+- Experiment with various optimization techniques safely
+
+## 🚀 Usage
+
+### Command Line Interface
+
+#### Using model presets (Recommended)
+
+The calculator includes pre-configured model presets for popular LLMs:
+
+```bash
+# List all available presets
+gpu-mem-calc presets
+
+# Calculate with a preset
+gpu-mem-calc calculate --preset llama2-7b
+gpu-mem-calc calculate --preset mixtral-8x7b --format json
+
+# List presets in table format
+gpu-mem-calc presets --format table
+```
+
+Available presets include:
+- **Dense Models**: LLaMA 2 (7B, 13B, 70B), GPT-3 (175B)
+- **MoE Models**: Mixtral 8x7B, GLM-4 (9B), GLM-4.7 (355B), GLM-4.5 Air (106B),
+  Qwen1.5-MoE-A2.7B, DeepSeek-MoE (16B)
+
+#### Calculate from config file
+
+```bash
+gpu-mem-calc calculate --config configs/llama2_7b_deepspeed.json
+```
+
+#### Quick calculation from model size
+
+```bash
+# Calculate memory for 7B model with 8x80GB GPUs using DeepSpeed
+gpu-mem-calc quick 7 --gpus 8 --engine deepspeed
+
+# With custom GPU memory
+gpu-mem-calc quick 70 --gpus 64 --gpu-mem 80 --engine megatron
+```
+
+#### Validate configuration
+
+```bash
+gpu-mem-calc validate configs/my_config.json
+```
+
+### Web Interface
+
+Start the web server:
+
+```bash
+python -m gpu_mem_calculator.web.app
+```
+
+Or using uvicorn directly:
+
+```bash
+uvicorn gpu_mem_calculator.web.app:app --reload
+```
+
+Then open your browser to `http://localhost:8000`
+
+### Python API
+
+#### Training Memory Calculation
+
+```python
+from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+from gpu_mem_calculator.core.models import (
+    ModelConfig,
+    TrainingConfig,
+    ParallelismConfig,
+    EngineConfig,
+    GPUConfig,
+)
+
+# Create configuration
+model_config = ModelConfig(
+    name="llama2-7b",
+    num_parameters=7_000_000_000,
+    num_layers=32,
+    hidden_size=4096,
+    num_attention_heads=32,
+    vocab_size=32000,
+    max_seq_len=4096,
+)
+
+training_config = TrainingConfig(
+    batch_size=4,
+    gradient_accumulation_steps=4,
+    dtype="bf16",
+    optimizer="adamw",
+)
+
+parallelism_config = ParallelismConfig(
+    data_parallel_size=8,
+)
+
+engine_config = EngineConfig(
+    type="deepspeed",
+    zero_stage=3,
+    offload_optimizer="cpu",
+)
+
+gpu_config = GPUConfig(
+    num_gpus=8,
+    gpu_memory_gb=80,
+)
+
+# Calculate memory
+calculator = GPUMemoryCalculator(
+    model_config=model_config,
+    training_config=training_config,
+    parallelism_config=parallelism_config,
+    engine_config=engine_config,
+    gpu_config=gpu_config,
+)
+
+result = calculator.calculate()
+
+print(f"Memory per GPU: {result.total_memory_per_gpu_gb:.2f} GB")
+print(f"Fits on GPU: {result.fits_on_gpu}")
+print(f"Utilization: {result.memory_utilization_percent:.1f}%")
+```
+
+#### 🆕 Inference Memory Calculation
+
+```python
+from gpu_mem_calculator.inference.calculator import InferenceMemoryCalculator
+from gpu_mem_calculator.core.models import (
+    ModelConfig,
+    InferenceConfig,
+    InferenceEngineType,
+    GPUConfig,
+)
+
+# Create configurations
+model_config = ModelConfig(
+    name="llama2-7b",
+    num_parameters=7_000_000_000,
+    num_layers=32,
+    hidden_size=4096,
+    num_attention_heads=32,
+    max_seq_len=4096,
+)
+
+inference_config = InferenceConfig(
+    batch_size=32,
+    kv_cache_quantization="int8",  # NONE, INT8, FP8, INT4
+    tensor_parallel_size=2,
+    gpu_memory_utilization=0.9,
+)
+
+gpu_config = GPUConfig(num_gpus=2, gpu_memory_gb=80)
+
+# Calculate for different inference engines
+calculator = InferenceMemoryCalculator(model_config, inference_config, gpu_config)
+
+# vLLM inference
+result_vllm = calculator.calculate(InferenceEngineType.VLLM)
+print(f"vLLM: {result_vllm.total_memory_per_gpu_gb:.2f} GB")
+print(f"Max batch size: {result_vllm.max_supported_batch_size}")
+print(f"Throughput: {result_vllm.estimated_throughput_tokens_per_sec:.0f} tokens/sec")
+
+# TensorRT-LLM inference
+result_trt = calculator.calculate(InferenceEngineType.TENSORRT_LLM)
+print(f"TensorRT-LLM: {result_trt.total_memory_per_gpu_gb:.2f} GB")
+```
+
+#### 🆕 Multi-Node Network Overhead
+
+```python
+from gpu_mem_calculator.core.multinode import MultiNodeCalculator
+from gpu_mem_calculator.core.models import (
+    NodeConfig,
+    InterconnectType,
+)
+
+# Configure multi-node setup
+node_config = NodeConfig(
+    num_nodes=4,
+    gpus_per_node=8,
+    interconnect_type=InterconnectType.INFINIBAND,
+)
+
+calculator = MultiNodeCalculator(
+    model_config=model_config,
+    training_config=training_config,
+    parallelism_config=parallelism_config,
+    node_config=node_config,
+    engine_config=engine_config,
+)
+
+# Calculate network overhead
+network_overhead = calculator.calculate_network_overhead()
+print(f"AllReduce: {network_overhead.allreduce_gb:.2f} GB")
+print(f"AllGather: {network_overhead.allgather_gb:.2f} GB")
+print(f"Time overhead: {network_overhead.estimated_overhead_ms_per_step:.2f} ms/step")
+
+# Optimize hybrid parallelism
+from gpu_mem_calculator.core.models import HybridParallelismConfig
+
+hybrid_config = HybridParallelismConfig(
+    auto_optimize=True,
+    prefer_pipeline_parallel=True,
+    enable_sequence_parallel=True,
+)
+
+optimized_parallelism = calculator.optimize_hybrid_parallelism(hybrid_config)
+print(f"Optimized TP: {optimized_parallelism.tensor_parallel_size}")
+print(f"Optimized PP: {optimized_parallelism.pipeline_parallel_size}")
+print(f"Optimized DP: {optimized_parallelism.data_parallel_size}")
+```
+
+#### 🆕 Export Framework Configurations
+
+```python
+from gpu_mem_calculator.exporters.manager import ExportManager, ExportFormat
+
+# Create export manager
+manager = ExportManager(
+    model_config=model_config,
+    training_config=training_config,
+    parallelism_config=parallelism_config,
+    engine_config=engine_config,
+    node_config=node_config,
+)
+
+# Export to different formats
+accelerate_config = manager.export(ExportFormat.ACCELERATE)
+lightning_config = manager.export(ExportFormat.LIGHTNING)
+axolotl_config = manager.export(ExportFormat.AXOLOTL)
+
+# Export to file
+manager.export_to_file(ExportFormat.ACCELERATE, "accelerate_config.yaml")
+manager.export_to_file(ExportFormat.JSON, "config.json")
+
+# Get DeepSpeed config
+deepspeed_config = manager.export(ExportFormat.DEEPSPEED)
+```
+
+## Configuration File Format
+
+```json
+{
+  "model": {
+    "name": "llama2-7b",
+    "num_parameters": "7B",
+    "num_layers": 32,
+    "hidden_size": 4096,
+    "num_attention_heads": 32,
+    "vocab_size": 32000,
+    "max_seq_len": 4096
+  },
+  "training": {
+    "batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "optimizer": "adamw",
+    "dtype": "bf16",
+    "activation_checkpointing": 1
+  },
+  "parallelism": {
+    "tensor_parallel_size": 1,
+    "pipeline_parallel_size": 1,
+    "data_parallel_size": 8,
+    "sequence_parallel": false
+  },
+  "engine": {
+    "type": "deepspeed",
+    "zero_stage": 3,
+    "offload_optimizer": "cpu",
+    "offload_param": "none"
+  },
+  "hardware": {
+    "num_gpus": 8,
+    "gpu_memory_gb": 80
+  }
+}
+```
+
+## Supported Training Engines
+
+### PyTorch DDP (Baseline)
+Standard Distributed Data Parallel training without memory optimizations.
+
+### DeepSpeed ZeRO
+- **ZeRO-1**: Shard optimizer states
+- **ZeRO-2**: Shard optimizer states + gradients
+- **ZeRO-3**: Shard everything (parameters, gradients, optimizer states)
+- Supports CPU/NVMe offloading
+
+### Megatron-LM
+Tensor and pipeline parallelism with activation checkpointing support.
+
+### Megatron + DeepSpeed
+Combines Megatron-LM's model parallelism with DeepSpeed ZeRO's optimizer sharding.
+
+### PyTorch FSDP
+Fully Sharded Data Parallel with multiple sharding strategies.
+
+## Memory Formulas
+
+The calculator uses formulas verified against authoritative sources:
+
+### Base Components
+
+**Model Parameters:**
+- FP16/BF16: `num_params × 2 bytes`
+- FP32: `num_params × 4 bytes`
+
+**Gradients:**
+- FP16/BF16: `num_params × 2 bytes`
+- FP32: `num_params × 4 bytes`
+
+**Optimizer States** (per optimizer type):
+- **Adam/AdamW**: `num_params × 12 bytes`
+  - 4 bytes: FP32 parameter copy
+  - 4 bytes: Momentum
+  - 4 bytes: Variance
+- **AdamW 8-bit**: `num_params × 2 bytes` (quantized)
+- **SGD**: `num_params × 4 bytes` (FP32 only, no momentum)
+
+**Activations:**
+- Approximation: `batch_size × seq_len × hidden_size × num_layers × ~16 bytes/token/layer`
+- Varies based on activation checkpointing level
+
+### DeepSpeed ZeRO Stages
+
+**ZeRO-0** (Baseline - same as PyTorch DDP):
+```
+total_per_gpu = 2×params + 2×params + 12×params + activations
+             = 16×params + activations
+```
+
+**ZeRO-1** (Shard optimizer states):
+```
+total_per_gpu = 2×params + 2×params + (12×params)/num_gpus + activations
+```
+
+**ZeRO-2** (Shard optimizer + gradients):
+```
+total_per_gpu = 2×params + (2×params)/num_gpus + (12×params)/num_gpus + activations
+```
+
+**ZeRO-3** (Shard everything):
+```
+total_per_gpu = largest_layer_memory + (16×params)/num_gpus + activations
+where largest_layer_memory ≈ 4×(num_params/10)
+```
+
+**CPU/NVMe Offloading:**
+- Optimizer states offloaded to CPU: 0 GB GPU memory
+- Parameters offloaded to CPU/NVMe: Dynamically gathered during compute
+
+### Verification
+
+All formulas have been verified against:
+- ✅ 18 comprehensive test scenarios (100% pass rate)
+- ✅ EleutherAI Transformer Math 101
+- ✅ Microsoft Research ZeRO Blog
+- ✅ DeepSpeed Official Documentation
+- ✅ PyTorch FSDP Documentation
+
+### References
+
+- [EleutherAI Transformer Math 101](https://blog.eleuther.ai/transformer-math/) - Comprehensive transformer memory breakdown
+- [Microsoft Research ZeRO Blog](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/) - ZeRO optimization techniques
+- [DeepSpeed Memory Documentation](https://deepspeed.readthedocs.io/en/latest/memory.html) - Official DeepSpeed memory formulas
+
+## Example Configurations
+
+### LLaMA 2 7B with DeepSpeed ZeRO-3
+```bash
+gpu-mem-calc calculate --config configs/llama2_7b_deepspeed.json
+```
+
+### GPT-3 175B with Megatron-LM
+```bash
+gpu-mem-calc calculate --config configs/gpt3_175b_megatron.json
+```
+
+### Custom 1B model with PyTorch DDP
+```bash
+gpu-mem-calc calculate --config configs/pytorch_ddp_example.json
+```
+
+## Web UI Features
+
+### Interactive Interface
+- **Real-time Calculations**: Auto-calculates as you adjust parameters (1s debounce)
+- **Client-side Validation**: Instant feedback on configuration errors before API calls
+- **Smart Presets**: Quick-load model configurations (LLaMA 2, GPT-3, GLM, Mixtral, Qwen, DeepSeek)
+- **Visual Breakdown**: Color-coded bar chart with patterns for colorblind accessibility
+- **Feasibility Status**: Clear indicators showing if configuration fits on GPU
+
+### Formula Explanations
+- **Detailed Breakdowns**: See exact formulas used with your values plugged in
+- **Component-by-Component**: Each memory component explained with formula and result
+- **Authoritative References**: Links to EleutherAI, Microsoft Research, DeepSpeed docs
+- **Engine-Specific Details**: Different formulas for PyTorch DDP, DeepSpeed ZeRO, FSDP, Megatron-LM
+
+### Advanced Tools
+- **Export to DeepSpeed**: Generate `deepspeed_config.json` files automatically
+- **Batch Size Optimizer**: Find maximum batch size that fits your GPU memory
+- **Config Persistence**: Save configurations to browser localStorage
+- **Comparison Mode**: Compare different configurations side-by-side
+
+### Accessibility
+- **ARIA Labels**: Full screen reader support throughout the interface
+- **Keyboard Navigation**: All features accessible via keyboard
+- **Colorblind-Friendly**: Patterns and textures supplement colors in charts
+- **High Contrast**: Clear visual indicators with multiple cues
+
+### API Endpoints
+- `POST /api/calculate` - Calculate GPU memory requirements
+- `POST /api/explain-formula` - Get detailed formula explanation
+- `POST /api/export/deepspeed` - Export DeepSpeed config file
+- `POST /api/optimize/batch-size` - Find maximum batch size
+- `GET /api/preset/{preset_name}` - Load model preset
+
+## Development
+
+### Running Tests
+
+```bash
+pytest tests/
+```
+
+### Test Coverage
+
+The calculator includes comprehensive testing:
+- **Unit Tests**: Core calculation logic for each engine type
+- **Integration Tests**: End-to-end configuration validation
+- **Formula Verification**: 18 scenarios verifying formula accuracy
+- **API Tests**: Web API endpoint testing
+- **Accessibility Tests**: Screen reader and keyboard navigation
+
+All formulas verified accurate against authoritative sources with 100% test pass rate.
+
+### Code Formatting
+
+```bash
+black src/ cli/ web/
+ruff check src/ cli/ web/
+```
+
+### Type Checking
+
+```bash
+mypy src/
+```
+
+## Recent Improvements
+
+### Latest Updates
+- ✨ Added formula explanation feature with detailed breakdowns
+- ✨ Added client-side validation for better UX
+- ✨ Added batch size optimizer API
+- ✨ Added DeepSpeed config export functionality
+- ✨ Added comprehensive input validation
+- ✨ Added result caching for performance
+- ♿ Added ARIA labels for full accessibility
+- ♿ Added colorblind patterns to charts
+- 🐛 Fixed optimizer formulas to be optimizer-specific
+- 🐛 Fixed Pydantic namespace warnings
+
+### Verification Status
+- ✅ All 18 test scenarios passing (100%)
+- ✅ Formulas verified against EleutherAI, Microsoft Research, DeepSpeed docs
+- ✅ Optimizer formulas corrected for AdamW, AdamW 8-bit, and SGD
+- ✅ ZeRO stage formulas validated (0, 1, 2, 3)
+- ✅ Engine type formulas validated (PyTorch DDP, DeepSpeed, FSDP, Megatron-LM)
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed guidelines.
+
+## 📚 References
+
+The memory calculations in this tool are based on authoritative sources:
+
+**Core Memory Formulas:**
+- [EleutherAI Transformer Math 101](https://blog.eleuther.ai/transformer-math/) - Comprehensive breakdown of transformer memory requirements
+- [Microsoft Research ZeRO Blog](https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/) - ZeRO optimization techniques
+- [Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/abs/2204.13323) - Activation checkpointing strategies
+
+**Engine Documentation:**
+- [DeepSpeed Memory Documentation](https://deepspeed.readthedocs.io/en/latest/memory.html) - Official DeepSpeed memory formulas
+- [NVIDIA Megatron-LM](https://github.com/NVIDIA/Megatron-LM) - Tensor and pipeline parallelism
+- [PyTorch FSDP Documentation](https://pytorch.org/docs/stable/fsdp.html) - Fully sharded data parallel
+- [PyTorch DDP Tutorial](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) - Distributed data parallel
+
+**Related Tools:**
+- [llm-analysis](https://github.com/cli99/llm-analysis) - LLM memory analysis
+- [vram-calculator](https://github.com/furiousteabag/vram-calculator) - VRAM calculation utilities
+
+## 🤝 Community & Support
+
+- 📖 [Documentation](README.md)
+- 🐛 [Issue Tracker](https://github.com/George614/gpu-mem-calculator/issues)
+- 💬 [Discussions](https://github.com/George614/gpu-mem-calculator/discussions)
+- 📧 Contact the maintainers via GitHub
+
+### Star History
+
+If you find this tool useful, please consider giving it a star! ⭐
+
+## 📋 Roadmap
+
+- [x] Inference memory calculation
+- [x] Multi-node training configurations
+- [x] Export to training framework configs (Accelerate, Lightning, Axolotl)
+- [ ] PyPI package distribution
+- [ ] Support for more model architectures (Vision Transformers, Diffusion models)
+- [ ] Real-time memory monitoring dashboard
+- [ ] CLI commands for inference and export features
+
+## 🙏 Acknowledgments
+
+This tool was inspired by and builds upon the excellent work of:
+- [DeepSpeed Memory Estimator](https://deepspeed.readthedocs.io/en/latest/memory.html) - ZeRO memory optimization formulas
+- [llm-analysis](https://github.com/cli99/llm-analysis) - LLM memory analysis methodology
+- [vram-calculator](https://github.com/furiousteabag/vram-calculator) - VRAM calculation approach
+
+Special thanks to the EleutherAI community for their comprehensive [Transformer Math 101](https://blog.eleuther.ai/transformer-math/) guide, which provides detailed formulas for transformer memory calculations.
+
+## 📄 License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## 📚 Citation
+
+If you use this tool in your research, please cite:
+
+```bibtex
+@software{gpu_mem_calculator,
+  title = {GPU Memory Calculator for LLM Training},
+  author = {GPU Mem Calculator Team},
+  year = {2024},
+  url = {https://github.com/George614/gpu-mem-calculator}
+}
+```
+
+---
+
+<p align="center">
+  Made with ❤️ for the ML community
+</p>
+
+<p align="center">
+  <a href="https://github.com/George614/gpu-mem-calculator/stargazers">⭐ Star us on GitHub</a> •
+  <a href="https://github.com/George614/gpu-mem-calculator/issues">🐛 Report a Bug</a> •
+  <a href="https://github.com/George614/gpu-mem-calculator/issues">💡 Request a Feature</a>
+</p>
+

src/gpu_mem_calculator.egg-info/SOURCES.txt

@@ -0,0 +1,46 @@
+LICENSE
+README.md
+pyproject.toml
+src/gpu_mem_calculator/__init__.py
+src/gpu_mem_calculator/py.typed
+src/gpu_mem_calculator.egg-info/PKG-INFO
+src/gpu_mem_calculator.egg-info/SOURCES.txt
+src/gpu_mem_calculator.egg-info/dependency_links.txt
+src/gpu_mem_calculator.egg-info/entry_points.txt
+src/gpu_mem_calculator.egg-info/requires.txt
+src/gpu_mem_calculator.egg-info/top_level.txt
+src/gpu_mem_calculator/cli/__init__.py
+src/gpu_mem_calculator/cli/main.py
+src/gpu_mem_calculator/config/__init__.py
+src/gpu_mem_calculator/config/parser.py
+src/gpu_mem_calculator/config/presets.py
+src/gpu_mem_calculator/core/__init__.py
+src/gpu_mem_calculator/core/calculator.py
+src/gpu_mem_calculator/core/formulas.py
+src/gpu_mem_calculator/core/models.py
+src/gpu_mem_calculator/core/multinode.py
+src/gpu_mem_calculator/engines/__init__.py
+src/gpu_mem_calculator/engines/base.py
+src/gpu_mem_calculator/engines/deepspeed.py
+src/gpu_mem_calculator/engines/fsdp.py
+src/gpu_mem_calculator/engines/megatron.py
+src/gpu_mem_calculator/engines/pytorch.py
+src/gpu_mem_calculator/exporters/__init__.py
+src/gpu_mem_calculator/exporters/accelerate.py
+src/gpu_mem_calculator/exporters/axolotl.py
+src/gpu_mem_calculator/exporters/lightning.py
+src/gpu_mem_calculator/exporters/manager.py
+src/gpu_mem_calculator/inference/__init__.py
+src/gpu_mem_calculator/inference/base.py
+src/gpu_mem_calculator/inference/calculator.py
+src/gpu_mem_calculator/inference/huggingface.py
+src/gpu_mem_calculator/inference/tensorrt_llm.py
+src/gpu_mem_calculator/inference/tgi.py
+src/gpu_mem_calculator/inference/vllm.py
+src/gpu_mem_calculator/utils/__init__.py
+src/gpu_mem_calculator/utils/precision.py
+tests/test_calculator.py
+tests/test_comprehensive.py
+tests/test_exporters.py
+tests/test_inference.py
+tests/test_multinode.py

src/gpu_mem_calculator.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/gpu_mem_calculator.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+gpu-mem-calc = gpu_mem_calculator.cli:main

src/gpu_mem_calculator.egg-info/requires.txt

@@ -0,0 +1,16 @@
+pydantic>=2.0.0
+click>=8.1.0
+pydantic-settings>=2.0.0
+rich>=13.0.0
+
+[dev]
+pytest>=7.0.0
+pytest-cov>=4.0.0
+black>=23.0.0
+ruff>=0.1.0
+mypy>=1.5.0
+
+[web]
+fastapi>=0.100.0
+uvicorn[standard]>=0.23.0
+jinja2>=3.1.0

src/gpu_mem_calculator.egg-info/top_level.txt

@@ -0,0 +1 @@
+gpu_mem_calculator

src/gpu_mem_calculator/__init__.py

@@ -0,0 +1,3 @@
+"""GPU Memory Calculator for LLM Training."""
+
+__version__ = "0.1.0"

src/gpu_mem_calculator/cli/__init__.py

@@ -0,0 +1,5 @@
+"""CLI interface for GPU Memory Calculator."""
+
+from gpu_mem_calculator.cli.main import main
+
+__all__ = ["main"]

src/gpu_mem_calculator/cli/main.py

@@ -0,0 +1,399 @@
+"""CLI interface for GPU Memory Calculator."""
+
+import json
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import click
+
+if TYPE_CHECKING:
+    from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+    from gpu_mem_calculator.core.models import MemoryResult
+
+
+@click.group()
+@click.version_option(version="0.1.0")
+def main() -> None:
+    """GPU Memory Calculator for LLM Training.
+
+    Calculate GPU memory requirements for training Large Language Models
+    with various training engines (PyTorch DDP, DeepSpeed, Megatron-LM, FSDP).
+    """
+    pass
+
+
+@main.command()
+@click.option(
+    "--config",
+    "-c",
+    type=click.Path(exists=True),
+    help="Path to JSON configuration file",
+)
+@click.option(
+    "--preset",
+    "-p",
+    type=str,
+    help="Name of a preset model configuration",
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(),
+    help="Output file path (default: stdout)",
+)
+@click.option(
+    "--format",
+    "-f",
+    type=click.Choice(["json", "yaml", "table"]),
+    default="table",
+    help="Output format (default: table)",
+)
+def calculate(
+    config: str | None,
+    preset: str | None,
+    output: str | None,
+    format: Literal["json", "yaml", "table"],
+) -> None:
+    """Calculate GPU memory requirements from config file or preset.
+
+    Examples:
+        gpu-mem-calc calculate --config configs/llama2_7b.json
+        gpu-mem-calc calculate --preset llama2-7b
+        gpu-mem-calc calculate -p mixtral-8x7b --format json
+    """
+    if not config and not preset:
+        click.echo("Error: Either --config or --preset is required", err=True)
+        sys.exit(1)
+
+    if config and preset:
+        click.echo("Error: Cannot use both --config and --preset", err=True)
+        sys.exit(1)
+
+    try:
+        import tempfile
+
+        from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+
+        if preset:
+            # Load preset configuration
+            from gpu_mem_calculator.config.presets import get_preset_config
+
+            preset_config = get_preset_config(preset)
+            if preset_config is None:
+                click.echo(
+                    f"Error: Preset '{preset}' not found. "
+                    "Use 'gpu-mem-calc presets' to list available presets.",
+                    err=True,
+                )
+                sys.exit(1)
+
+            # Write preset to temp file for from_config_file
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+                json.dump(preset_config, f, indent=2)
+                temp_path = f.name
+
+            calculator = GPUMemoryCalculator.from_config_file(temp_path)
+            Path(temp_path).unlink()  # Clean up temp file
+        elif config:
+            calculator = GPUMemoryCalculator.from_config_file(config)
+        else:
+            # This should never happen due to the checks above
+            click.echo("Error: Either --config or --preset is required", err=True)
+            sys.exit(1)
+
+        result = calculator.calculate()
+
+        # Format output
+        if format == "json":
+            output_text = json.dumps(result.model_dump(mode="json"), indent=2)
+        elif format == "yaml":
+            try:
+                import yaml  # type: ignore[import-untyped]
+
+                output_text = yaml.dump(result.model_dump(mode="json"), default_flow_style=False)
+            except ImportError:
+                click.echo(
+                    "Error: YAML format requires PyYAML. Install with: pip install pyyaml",
+                    err=True,
+                )
+                sys.exit(1)
+        else:  # table
+            output_text = _format_result_as_table(result, calculator)
+
+        # Write output
+        if output:
+            Path(output).write_text(output_text)
+            click.echo(f"Results written to {output}")
+        else:
+            click.echo(output_text)
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.argument(
+    "params",
+    type=float,
+    required=True,
+)
+@click.option(
+    "--gpus",
+    "-g",
+    type=int,
+    default=1,
+    help="Number of GPUs (default: 1)",
+)
+@click.option(
+    "--gpu-mem",
+    "-m",
+    type=float,
+    default=80.0,
+    help="GPU memory in GB (default: 80.0)",
+)
+@click.option(
+    "--engine",
+    "-e",
+    type=click.Choice(["pytorch", "deepspeed", "megatron", "fsdp"]),
+    default="pytorch",
+    help="Training engine (default: pytorch)",
+)
+@click.option(
+    "--dtype",
+    "-d",
+    type=click.Choice(["fp32", "fp16", "bf16"]),
+    default="bf16",
+    help="Data type (default: bf16)",
+)
+def quick(
+    params: float,
+    gpus: int,
+    gpu_mem: float,
+    engine: str,
+    dtype: str,
+) -> None:
+    """Quick calculation from model size (in billions of parameters).
+
+    Example:
+        gpu-mem-calc quick 7 --gpus 8 --engine deepspeed
+    """
+    try:
+        from gpu_mem_calculator.core.calculator import GPUMemoryCalculator
+        from gpu_mem_calculator.core.models import (
+            DType,
+            EngineConfig,
+            EngineType,
+            GPUConfig,
+            ModelConfig,
+            ParallelismConfig,
+            TrainingConfig,
+        )
+
+        # Map engine string to EngineType
+        engine_map = {
+            "pytorch": EngineType.PYTORCH_DDP,
+            "deepspeed": EngineType.DEEPSPEED,
+            "megatron": EngineType.MEGATRON_LM,
+            "fsdp": EngineType.FSDP,
+        }
+
+        # Map dtype string to DType
+        dtype_map = {
+            "fp32": DType.FP32,
+            "fp16": DType.FP16,
+            "bf16": DType.BF16,
+        }
+
+        # Create a minimal config for quick calculation
+        # Estimate model architecture from parameter count
+        # Rough approximation based on typical transformer models
+        num_params = int(params * 1e9)
+
+        # Estimate hidden size and layers from param count
+        # These are rough approximations
+        if params <= 1:
+            hidden_size, num_layers = 768, 12
+        elif params <= 7:
+            hidden_size, num_layers = 4096, 32
+        elif params <= 13:
+            hidden_size, num_layers = 5120, 40
+        elif params <= 30:
+            hidden_size, num_layers = 6656, 60
+        elif params <= 65:
+            hidden_size, num_layers = 8192, 80
+        else:
+            hidden_size, num_layers = 12288, 96
+
+        model_config = ModelConfig(
+            name="quick-estimate",
+            num_parameters=num_params,
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            num_attention_heads=hidden_size // 128,
+            vocab_size=32000,
+            max_seq_len=2048,
+        )
+
+        training_config = TrainingConfig(
+            batch_size=1,
+            gradient_accumulation_steps=1,
+            dtype=dtype_map[dtype],
+        )
+
+        parallelism_config = ParallelismConfig(data_parallel_size=gpus)
+
+        engine_config = EngineConfig(
+            type=engine_map[engine],
+            zero_stage=2 if engine == "deepspeed" else None,
+        )
+
+        gpu_config = GPUConfig(num_gpus=gpus, gpu_memory_gb=gpu_mem)
+
+        calculator = GPUMemoryCalculator(
+            model_config=model_config,
+            training_config=training_config,
+            parallelism_config=parallelism_config,
+            engine_config=engine_config,
+            gpu_config=gpu_config,
+        )
+
+        result = calculator.calculate()
+
+        # Display results
+        click.echo(_format_result_as_table(result, calculator))
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.argument(
+    "config_path",
+    type=click.Path(exists=True),
+)
+def validate(config_path: str) -> None:
+    """Validate a configuration file.
+
+    Example:
+        gpu-mem-calc validate configs/my_config.json
+    """
+    try:
+        from gpu_mem_calculator.config import ConfigParser
+
+        ConfigParser.parse_full_config(config_path)
+        click.echo(f"✓ Configuration file '{config_path}' is valid")
+
+    except Exception as e:
+        click.echo(f"✗ Validation failed: {e}", err=True)
+        sys.exit(1)
+
+
+@main.command()
+@click.option(
+    "--format",
+    "-f",
+    type=click.Choice(["list", "json", "table"]),
+    default="list",
+    help="Output format (default: list)",
+)
+def presets(format: str) -> None:
+    """List available model preset configurations.
+
+    Examples:
+        gpu-mem-calc presets
+        gpu-mem-calc presets --format table
+        gpu-mem-calc presets -f json
+    """
+    try:
+        from gpu_mem_calculator.config.presets import list_presets
+
+        all_presets = list_presets()
+
+        if not all_presets:
+            click.echo("No presets found.")
+            return
+
+        if format == "json":
+            click.echo(json.dumps(all_presets, indent=2))
+        elif format == "table":
+            from rich.console import Console
+            from rich.table import Table
+
+            console = Console()
+            table = Table(
+                title="Available Model Presets",
+                show_header=True,
+                header_style="bold magenta",
+            )
+            table.add_column("Preset Name", style="cyan", width=25)
+            table.add_column("Display Name", style="green", width=30)
+            table.add_column("Description", style="yellow")
+
+            for name, info in sorted(all_presets.items()):
+                table.add_row(name, info["display_name"], info["description"])
+
+            console.print(table)
+        else:  # list format
+            click.echo("Available model presets:\n")
+            for name, info in sorted(all_presets.items()):  # type: ignore[annotation-unchecked]
+                click.echo(f"  {name:25} - {info['display_name']}")
+                if info.get("description"):
+                    click.echo(f"{'':27}{info['description']}")
+                click.echo()
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+
+
+def _format_result_as_table(result: MemoryResult, calculator: "GPUMemoryCalculator") -> str:
+    """Format result as ASCII table."""
+    from rich.console import Console
+    from rich.table import Table
+
+    console = Console()
+
+    # Main results table
+    table = Table(
+        title="GPU Memory Calculation Results",
+        show_header=True,
+        header_style="bold magenta",
+    )
+    table.add_column("Metric", style="cyan", width=30)
+    table.add_column("Value", style="green")
+
+    # Memory results
+    table.add_row("Memory per GPU", f"{result.total_memory_per_gpu_gb:.2f} GB")
+    table.add_row("Total GPU Memory", f"{result.total_memory_all_gpus_gb:.2f} GB")
+    table.add_row("CPU Memory", f"{result.cpu_memory_gb:.2f} GB")
+    table.add_row("", "")  # Spacer
+
+    # Breakdown
+    table.add_row("Model Parameters", f"{result.breakdown.model_params_gb:.2f} GB")
+    table.add_row("Gradients", f"{result.breakdown.gradients_gb:.2f} GB")
+    table.add_row("Optimizer States", f"{result.breakdown.optimizer_states_gb:.2f} GB")
+    table.add_row("Activations", f"{result.breakdown.activations_gb:.2f} GB")
+    table.add_row("Overhead", f"{result.breakdown.overhead_gb:.2f} GB")
+    table.add_row("", "")  # Spacer
+
+    # Feasibility
+    status = "✓ Fits" if result.fits_on_gpu else "✗ OOM"
+    table.add_row("Status", status)
+    table.add_row("Memory Utilization", f"{result.memory_utilization_percent:.1f}%")
+    if result.recommended_batch_size:
+        table.add_row("Recommended Batch Size", str(result.recommended_batch_size))
+
+    # Capture table output
+    from io import StringIO
+
+    buffer = StringIO()
+    console.file = buffer
+    console.print(table)
+    return buffer.getvalue()
+
+
+if __name__ == "__main__":
+    main()

src/gpu_mem_calculator/config/__init__.py

@@ -0,0 +1,5 @@
+"""Configuration parsing and defaults."""
+
+from gpu_mem_calculator.config.parser import ConfigParser, load_config, save_config
+
+__all__ = ["ConfigParser", "load_config", "save_config"]

src/gpu_mem_calculator/config/parser.py

@@ -0,0 +1,323 @@
+"""Configuration file parser and utilities."""
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from pydantic import ValidationError
+
+from gpu_mem_calculator.core.models import (
+    DType,
+    EngineConfig,
+    EngineType,
+    GPUConfig,
+    ModelConfig,
+    OffloadDevice,
+    OptimizerType,
+    ParallelismConfig,
+    TrainingConfig,
+)
+
+
+class ConfigParseError(Exception):
+    """Error parsing configuration file."""
+
+    def __init__(self, message: str, errors: list[Any] | None = None):
+        super().__init__(message)
+        self.errors = errors or []
+
+
+class ConfigParser:
+    """Parse and validate configuration files."""
+
+    @staticmethod
+    def _convert_dtype(value: str) -> DType:
+        """Convert string dtype to DType enum."""
+        dtype_map = {
+            "float32": DType.FP32,
+            "fp32": DType.FP32,
+            "float16": DType.FP16,
+            "fp16": DType.FP16,
+            "bfloat16": DType.BF16,
+            "bf16": DType.BF16,
+            "int8": DType.INT8,
+            "int4": DType.INT4,
+        }
+        return dtype_map.get(value.lower(), DType.BF16)
+
+    @staticmethod
+    def _convert_optimizer(value: str) -> OptimizerType:
+        """Convert string optimizer to OptimizerType enum."""
+        opt_map = {
+            "adam": OptimizerType.ADAM,
+            "adamw": OptimizerType.ADAMW,
+            "sgd": OptimizerType.SGD,
+            "adamw_8bit": OptimizerType.ADAMW_8BIT,
+            "adamw-8bit": OptimizerType.ADAMW_8BIT,
+        }
+        return opt_map.get(value.lower(), OptimizerType.ADAMW)
+
+    @staticmethod
+    def _convert_engine(value: str) -> EngineType:
+        """Convert string engine to EngineType enum."""
+        engine_map = {
+            "pytorch": EngineType.PYTORCH_DDP,
+            "pytorch_ddp": EngineType.PYTORCH_DDP,
+            "ddp": EngineType.PYTORCH_DDP,
+            "deepspeed": EngineType.DEEPSPEED,
+            "megatron": EngineType.MEGATRON_LM,
+            "megatron_lm": EngineType.MEGATRON_LM,
+            "megatron-lm": EngineType.MEGATRON_LM,
+            "fsdp": EngineType.FSDP,
+            "megatron_deepspeed": EngineType.MEGATRON_DEEPSPEED,
+        }
+        return engine_map.get(value.lower(), EngineType.PYTORCH_DDP)
+
+    @staticmethod
+    def _convert_offload(value: str) -> OffloadDevice:
+        """Convert string offload to OffloadDevice enum."""
+        offload_map = {
+            "none": OffloadDevice.NONE,
+            "cpu": OffloadDevice.CPU,
+            "nvme": OffloadDevice.NVME,
+        }
+        return offload_map.get(value.lower(), OffloadDevice.NONE)
+
+    @staticmethod
+    def _parse_num_params(value: str | int | float) -> int:
+        """Parse number of parameters from various formats.
+
+        Supports:
+        - Raw integer: 7000000000
+        - Billions: "7B", "7b", "7e9"
+        - Millions: "7000M", "7000m", "7000e6"
+        """
+        if isinstance(value, int):
+            return value
+        if isinstance(value, float):
+            return int(value)
+
+        if isinstance(value, str):
+            value = value.strip().upper()
+
+            # Handle billions suffix
+            if value.endswith("B"):
+                return int(float(value[:-1]) * 1_000_000_000)
+
+            # Handle millions suffix
+            if value.endswith("M"):
+                return int(float(value[:-1]) * 1_000_000)
+
+            # Handle scientific notation
+            if "E" in value:
+                return int(float(value))
+
+            # Try direct conversion
+            return int(value)
+
+        raise ValueError(f"Cannot parse parameter count: {value}")
+
+    @classmethod
+    def parse_model_config(cls, data: dict[str, Any]) -> ModelConfig:
+        """Parse model configuration from dict.
+
+        Args:
+            data: Dictionary with model configuration
+
+        Returns:
+            ModelConfig object
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        try:
+            # Convert parameter count if it's a string
+            if "num_parameters" in data and isinstance(data["num_parameters"], str):
+                data["num_parameters"] = cls._parse_num_params(data["num_parameters"])
+
+            if "largest_layer_params" in data and isinstance(data["largest_layer_params"], str):
+                data["largest_layer_params"] = cls._parse_num_params(data["largest_layer_params"])
+
+            return ModelConfig(**data)
+        except ValidationError as e:
+            raise ConfigParseError("Invalid model configuration", e.errors()) from e
+
+    @classmethod
+    def parse_training_config(cls, data: dict[str, Any]) -> TrainingConfig:
+        """Parse training configuration from dict.
+
+        Args:
+            data: Dictionary with training configuration
+
+        Returns:
+            TrainingConfig object
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        try:
+            # Convert dtype
+            if "dtype" in data and isinstance(data["dtype"], str):
+                data["dtype"] = cls._convert_dtype(data["dtype"])
+
+            # Convert optimizer
+            if "optimizer" in data and isinstance(data["optimizer"], str):
+                data["optimizer"] = cls._convert_optimizer(data["optimizer"])
+
+            return TrainingConfig(**data)
+        except ValidationError as e:
+            raise ConfigParseError("Invalid training configuration", e.errors()) from e
+
+    @classmethod
+    def parse_parallelism_config(cls, data: dict[str, Any]) -> ParallelismConfig:
+        """Parse parallelism configuration from dict.
+
+        Args:
+            data: Dictionary with parallelism configuration
+
+        Returns:
+            ParallelismConfig object
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        try:
+            return ParallelismConfig(**data)
+        except ValidationError as e:
+            raise ConfigParseError("Invalid parallelism configuration", e.errors()) from e
+
+    @classmethod
+    def parse_engine_config(cls, data: dict[str, Any]) -> EngineConfig:
+        """Parse engine configuration from dict.
+
+        Args:
+            data: Dictionary with engine configuration
+
+        Returns:
+            EngineConfig object
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        try:
+            # Convert engine type
+            if "type" in data and isinstance(data["type"], str):
+                data["type"] = cls._convert_engine(data["type"])
+
+            # Convert offload options
+            if "offload_optimizer" in data and isinstance(data["offload_optimizer"], str):
+                data["offload_optimizer"] = cls._convert_offload(data["offload_optimizer"])
+
+            if "offload_param" in data and isinstance(data["offload_param"], str):
+                data["offload_param"] = cls._convert_offload(data["offload_param"])
+
+            return EngineConfig(**data)
+        except ValidationError as e:
+            raise ConfigParseError("Invalid engine configuration", e.errors()) from e
+
+    @classmethod
+    def parse_gpu_config(cls, data: dict[str, Any]) -> GPUConfig:
+        """Parse GPU configuration from dict.
+
+        Args:
+            data: Dictionary with GPU configuration
+
+        Returns:
+            GPUConfig object
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        try:
+            return GPUConfig(**data)
+        except ValidationError as e:
+            raise ConfigParseError("Invalid GPU configuration", e.errors()) from e
+
+    @classmethod
+    def parse_file(cls, config_path: str | Path) -> dict[str, Any]:
+        """Parse configuration from JSON file.
+
+        Args:
+            config_path: Path to configuration file
+
+        Returns:
+            Dictionary with parsed configuration
+
+        Raises:
+            ConfigParseError: If file cannot be read or parsed
+        """
+        path = Path(config_path)
+        if not path.exists():
+            raise ConfigParseError(f"Configuration file not found: {config_path}")
+
+        try:
+            with path.open("r") as f:
+                data = cast(dict[str, Any], json.load(f))
+            return data
+        except json.JSONDecodeError as e:
+            raise ConfigParseError(f"Invalid JSON in configuration file: {e}") from e
+        except Exception as e:
+            raise ConfigParseError(f"Error reading configuration file: {e}") from e
+
+    @classmethod
+    def parse_full_config(
+        cls,
+        config_path: str | Path,
+    ) -> tuple[ModelConfig, TrainingConfig, ParallelismConfig, EngineConfig, GPUConfig]:
+        """Parse complete configuration from file.
+
+        Args:
+            config_path: Path to configuration file
+
+        Returns:
+            Tuple of (ModelConfig, TrainingConfig, ParallelismConfig, EngineConfig, GPUConfig)
+
+        Raises:
+            ConfigParseError: If validation fails
+        """
+        data = cls.parse_file(config_path)
+
+        try:
+            model_config = cls.parse_model_config(data.get("model", {}))
+            training_config = cls.parse_training_config(data.get("training", {}))
+            parallelism_config = cls.parse_parallelism_config(data.get("parallelism", {}))
+            engine_config = cls.parse_engine_config(data.get("engine", {}))
+            gpu_config = cls.parse_gpu_config(data.get("hardware", {}))
+
+            return (
+                model_config,
+                training_config,
+                parallelism_config,
+                engine_config,
+                gpu_config,
+            )
+        except ConfigParseError:
+            raise
+        except Exception as e:
+            raise ConfigParseError(f"Unexpected error parsing configuration: {e}") from e
+
+
+def load_config(config_path: str | Path) -> dict[str, Any]:
+    """Load configuration from file.
+
+    Args:
+        config_path: Path to configuration file
+
+    Returns:
+        Dictionary with configuration data
+    """
+    return ConfigParser.parse_file(config_path)
+
+
+def save_config(data: dict[str, Any], output_path: str | Path) -> None:
+    """Save configuration to JSON file.
+
+    Args:
+        data: Configuration dictionary to save
+        output_path: Path to save configuration file
+    """
+    path = Path(output_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    with path.open("w") as f:
+        json.dump(data, f, indent=2)

src/gpu_mem_calculator/config/presets.py

@@ -0,0 +1,83 @@
+"""Preset model configurations loader.
+
+This module provides a centralized location for managing model preset
+configurations that can be used by both CLI and web interfaces.
+"""
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+# Base directory for the package
+BASE_DIR = Path(__file__).parent.parent.parent.parent
+
+
+def get_presets_file_path() -> Path:
+    """Get the path to the presets JSON file.
+
+    Returns:
+        Path to the presets JSON file
+    """
+    # Check for web/presets/models.json relative to project root
+    presets_path = BASE_DIR / "web" / "presets" / "models.json"
+    if presets_path.exists():
+        return presets_path
+
+    # Fallback to src directory for development installs
+    presets_path = BASE_DIR / "src" / "gpu_mem_calculator" / "presets" / "models.json"
+    return presets_path
+
+
+def load_presets() -> dict[str, dict[str, Any]]:
+    """Load all preset model configurations.
+
+    Returns:
+        Dictionary mapping preset names to their configurations.
+        Each preset has: display_name, description, config
+    """
+    presets_file = get_presets_file_path()
+
+    if not presets_file.exists():
+        return {}
+
+    try:
+        with presets_file.open("r") as f:
+            return cast(dict[str, dict[str, Any]], json.load(f))
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def get_preset_config(preset_name: str) -> dict[str, Any] | None:
+    """Get a specific preset configuration.
+
+    Args:
+        preset_name: Name of the preset to retrieve
+
+    Returns:
+        Preset configuration dict, or None if not found
+    """
+    presets = load_presets()
+    preset = presets.get(preset_name)
+
+    if preset is None:
+        return None
+
+    # Return just the config part (what the calculator needs)
+    return cast(dict[str, Any], preset.get("config", {}))
+
+
+def list_presets() -> dict[str, dict[str, str]]:
+    """List all available presets with metadata.
+
+    Returns:
+        Dictionary mapping preset names to their display metadata.
+        Each entry has: display_name, description
+    """
+    presets = load_presets()
+    return {
+        name: {
+            "display_name": preset.get("display_name", name),
+            "description": preset.get("description", ""),
+        }
+        for name, preset in presets.items()
+    }

src/gpu_mem_calculator/core/__init__.py

@@ -0,0 +1,24 @@
+"""Core memory calculation models and formulas."""
+
+from gpu_mem_calculator.core.formulas import Precision
+from gpu_mem_calculator.core.models import (
+    EngineConfig,
+    EngineType,
+    GPUConfig,
+    ModelConfig,
+    ParallelismConfig,
+    TrainingConfig,
+)
+
+__all__ = [
+    "ModelConfig",
+    "TrainingConfig",
+    "ParallelismConfig",
+    "EngineConfig",
+    "EngineType",
+    "GPUConfig",
+    "Precision",
+]
+
+# Import GPUMemoryCalculator separately to avoid circular import
+# Use: from gpu_mem_calculator.core.calculator import GPUMemoryCalculator

src/gpu_mem_calculator/core/calculator.py

@@ -0,0 +1,178 @@
+"""Main GPU memory calculator.
+
+Orchestrates the memory calculation by selecting the appropriate
+training engine and aggregating results.
+"""
+
+from gpu_mem_calculator.config.parser import ConfigParser
+from gpu_mem_calculator.core.models import (
+    EngineConfig,
+    EngineType,
+    GPUConfig,
+    MemoryResult,
+    ModelConfig,
+    NodeConfig,
+    ParallelismConfig,
+    TrainingConfig,
+)
+from gpu_mem_calculator.engines import (
+    DeepSpeedEngine,
+    FSDPEngine,
+    MegatronDeepSpeedEngine,
+    MegatronLMEngine,
+    PyTorchDDPEngine,
+)
+
+# Type alias for engine types
+EngineTypeAlias = (
+    PyTorchDDPEngine | DeepSpeedEngine | MegatronLMEngine | FSDPEngine | MegatronDeepSpeedEngine
+)
+
+
+class GPUMemoryCalculator:
+    """Main GPU memory calculator.
+
+    This class provides a high-level interface for calculating
+    GPU memory requirements for LLM training.
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        training_config: TrainingConfig,
+        parallelism_config: ParallelismConfig | None = None,
+        engine_config: EngineConfig | None = None,
+        gpu_config: GPUConfig | None = None,
+        node_config: NodeConfig | None = None,
+    ) -> None:
+        """Initialize the calculator.
+
+        Args:
+            model_config: Model architecture configuration
+            training_config: Training hyperparameters
+            parallelism_config: Parallelism settings (default: no parallelism)
+            engine_config: Training engine configuration (default: PyTorch DDP)
+            gpu_config: Hardware configuration (default: 1x 80GB GPU)
+            node_config: Multi-node configuration (default: single node)
+        """
+        self.model_config = model_config
+        self.training_config = training_config
+        self.parallelism_config = parallelism_config or ParallelismConfig()
+        self.engine_config = engine_config or EngineConfig()
+        self.gpu_config = gpu_config or GPUConfig()
+        self.node_config = node_config or NodeConfig()
+
+    def calculate(self) -> MemoryResult:
+        """Calculate GPU memory requirements.
+
+        Selects the appropriate training engine based on configuration
+        and returns the memory calculation result.
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        engine = self._get_engine()
+        return engine.calculate_memory()
+
+    def _get_engine(self) -> EngineTypeAlias:
+        """Get the appropriate training engine instance.
+
+        Returns:
+            Engine instance configured with current settings
+        """
+        match self.engine_config.type:
+            case EngineType.PYTORCH_DDP:
+                return PyTorchDDPEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+            case EngineType.DEEPSPEED:
+                return DeepSpeedEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+            case EngineType.MEGATRON_LM:
+                return MegatronLMEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+            case EngineType.FSDP:
+                return FSDPEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+            case EngineType.MEGATRON_DEEPSPEED:
+                return MegatronDeepSpeedEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+            case _:
+                # Default to PyTorch DDP
+                return PyTorchDDPEngine(
+                    model_config=self.model_config,
+                    training_config=self.training_config,
+                    parallelism_config=self.parallelism_config,
+                    engine_config=self.engine_config,
+                    gpu_config=self.gpu_config,
+                    node_config=self.node_config,
+                )
+
+    @classmethod
+    def from_config_file(
+        cls,
+        config_path: str,
+    ) -> "GPUMemoryCalculator":
+        """Create calculator from configuration file.
+
+        Args:
+            config_path: Path to JSON configuration file
+
+        Returns:
+            Configured GPUMemoryCalculator instance
+        """
+        model_config, training_config, parallelism_config, engine_config, gpu_config = (
+            ConfigParser.parse_full_config(config_path)
+        )
+
+        return cls(
+            model_config=model_config,
+            training_config=training_config,
+            parallelism_config=parallelism_config,
+            engine_config=engine_config,
+            gpu_config=gpu_config,
+        )
+
+    def to_dict(self) -> dict:
+        """Export calculator configuration to dictionary.
+
+        Returns:
+            Dictionary with all configuration
+        """
+        return {
+            "model": self.model_config.model_dump(),
+            "training": self.training_config.model_dump(),
+            "parallelism": self.parallelism_config.model_dump(),
+            "engine": self.engine_config.model_dump(),
+            "hardware": self.gpu_config.model_dump(),
+            "multinode": self.node_config.model_dump(),
+        }

src/gpu_mem_calculator/core/formulas.py

@@ -0,0 +1,268 @@
+"""Memory calculation formulas.
+
+This module contains the fundamental formulas for calculating GPU memory
+requirements for LLM training.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Precision:
+    """Precision information for a data type.
+
+    This is re-exported from utils.precision for convenience.
+    """
+
+    name: str
+    bits_per_param: int
+    bytes_per_param: float
+    is_integer: bool = False
+
+
+def calculate_parameter_memory(
+    num_params: int,
+    dtype: str,
+    num_gpus: int = 1,
+) -> float:
+    """Calculate memory in GB for model parameters.
+
+    Args:
+        num_params: Number of model parameters
+        dtype: Data type (e.g., "fp32", "fp16", "bf16", "int8", "int4")
+        num_gpus: Number of GPUs for distribution
+
+    Returns:
+        Memory in GB
+    """
+    from gpu_mem_calculator.utils.precision import gb_from_params
+
+    # Parameters are distributed across GPUs in data parallel training
+    # But for tensor/pipeline parallel, each GPU holds a portion
+    # We'll handle parallelism in the engine implementations
+    return gb_from_params(num_params, dtype)
+
+
+def calculate_gradient_memory(
+    num_params: int,
+    dtype: str,
+) -> float:
+    """Calculate memory in GB for gradients.
+
+    Gradients are typically stored in the same precision as parameters
+    for training (though updated in FP32).
+
+    Args:
+        num_params: Number of model parameters
+        dtype: Data type for gradients
+
+    Returns:
+        Memory in GB
+    """
+    from gpu_mem_calculator.utils.precision import gb_from_params
+
+    # Gradients are same size as parameters during training
+    return gb_from_params(num_params, dtype)
+
+
+def calculate_optimizer_memory(
+    num_params: int,
+    optimizer: str,
+) -> float:
+    """Calculate memory in GB for optimizer states.
+
+    Args:
+        num_params: Number of model parameters
+        optimizer: Optimizer type (adam, adamw, sgd, adamw_8bit)
+
+    Returns:
+        Memory in GB (for FP32 optimizer states)
+    """
+    from gpu_mem_calculator.utils.precision import gb_from_bytes
+
+    # Optimizer states are typically stored in FP32
+    # bytes_per_param = 4.0  # FP32
+
+    match optimizer.lower():
+        case "adam" | "adamw":
+            # Adam/AdamW optimizer states: 12 bytes per param
+            # - FP32 parameter copy: 4 bytes
+            # - Momentum (fp32): 4 bytes
+            # - Variance (fp32): 4 bytes
+            # Reference: https://blog.eleuther.ai/transformer-math/#optimizer-states
+            # Reference: https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
+            # Reference: https://deepspeed.readthedocs.io/en/latest/memory.html
+            optimizer_bytes_per_param = 12.0
+        case "adamw_8bit":
+            # 8-bit Adam: ~2 bytes per param (quantized states)
+            # Reference: bitsandbytes 8-bit optimizer
+            optimizer_bytes_per_param = 2.0
+        case "sgd":
+            # SGD: momentum (4 bytes) if using momentum, 0 if not
+            # Assuming momentum is used
+            optimizer_bytes_per_param = 4.0
+        case _:
+            # Default to Adam
+            optimizer_bytes_per_param = 12.0
+
+    total_bytes = num_params * optimizer_bytes_per_param
+    return gb_from_bytes(total_bytes)
+
+
+def calculate_activation_memory(
+    batch_size: int,
+    seq_len: int,
+    hidden_size: int,
+    num_layers: int,
+    num_attention_heads: int,
+    tensor_parallel_size: int = 1,
+    activation_checkpointing: int = 0,
+    moe_enabled: bool = False,
+    num_experts: int = 1,
+    top_k: int = 1,
+    expert_intermediate_size: int | None = None,
+) -> float:
+    """Calculate approximate memory in GB for activations.
+
+    This provides an estimate based on transformer architecture. Actual
+    activation memory depends on many factors including the specific
+    model implementation and framework.
+
+    Reference: https://blog.eleuther.ai/transformer-math/#activations
+    Reference: https://arxiv.org/abs/2204.13323 ("Reducing Activation Recomputation
+               in Large Transformer Models")
+
+    According to EleutherAI Transformer Math 101, for selective activation
+    checkpointing (the most common approach), the formula is:
+
+        sbhL(10 + 24/t) bytes
+
+    Where:
+    - s = sequence length (seq_len)
+    - b = batch size (batch_size)
+    - h = hidden size (hidden_size)
+    - L = number of layers (num_layers)
+    - t = tensor parallel size (tensor_parallel_size)
+
+    This implementation uses a simplified heuristic that approximates
+    this formula: hidden_size * 16 bytes per token per layer. This
+    provides a reasonable estimate for typical model configurations
+    while being simple to understand and modify.
+
+    For MoE models, activation memory is reduced because only top_k experts
+    are active per token, not all experts.
+
+    Args:
+        batch_size: Batch size per GPU
+        seq_len: Sequence length
+        hidden_size: Hidden dimension size
+        num_layers: Number of transformer layers
+        num_attention_heads: Number of attention heads
+        tensor_parallel_size: Tensor parallelism degree
+        activation_checkpointing: Checkpointing level (0-4)
+        moe_enabled: Whether model uses Mixture of Experts
+        num_experts: Total number of experts (for MoE)
+        top_k: Number of active experts per token (for MoE)
+        expert_intermediate_size: Expert intermediate layer size (for MoE)
+
+    Returns:
+        Memory in GB
+    """
+    from gpu_mem_calculator.utils.precision import gb_from_bytes
+
+    # Approximate activation memory per token per layer
+    # Based on EleutherAI formula: sbhL(10 + 24/t)
+    # For t=1: ~10-24 bytes per token per layer depending on architecture
+    # We use 16 as a middle-ground estimate
+    # This includes attention outputs, MLP activations, layer norms, etc.
+
+    bytes_per_token_per_layer = hidden_size * 16  # Heuristic estimate
+
+    # For MoE models, adjust activation memory based on active experts
+    moe_multiplier = 1.0
+    if moe_enabled and num_experts > 1:
+        # Only top_k experts are active per token
+        # Base ratio of active experts
+        expert_ratio = top_k / num_experts
+
+        # Add router overhead (gating network activations)
+        router_overhead = 0.1
+
+        moe_multiplier = min(1.0, expert_ratio + router_overhead)
+
+    # For MoE, experts typically have larger intermediate sizes
+    if moe_enabled and expert_intermediate_size:
+        # Scale up slightly for larger expert intermediate layers
+        # Typical expert intermediate size is 4x hidden_size (vs 2x for dense)
+        size_ratio = expert_intermediate_size / (hidden_size * 2)
+        moe_multiplier *= min(2.0, size_ratio)  # Cap at 2x increase
+
+    # Total activation memory
+    total_bytes = (
+        batch_size
+        * seq_len
+        * num_layers
+        * bytes_per_token_per_layer
+        * moe_multiplier
+        / tensor_parallel_size
+    )
+
+    # Adjust for activation checkpointing
+    # Level 0: No checkpointing (100% memory)
+    # Level 1: Checkpoint attention output (~80% memory)
+    # Level 2: Checkpoint attention input (~60% memory)
+    # Level 3: Checkpoint more (~40% memory)
+    # Level 4: Full checkpointing (~20% memory)
+    checkpoint_factors = [1.0, 0.8, 0.6, 0.4, 0.2]
+    checkpoint_factor = checkpoint_factors[min(activation_checkpointing, 4)]
+
+    total_bytes *= checkpoint_factor
+
+    return gb_from_bytes(total_bytes)
+
+
+def calculate_overhead(
+    total_memory: float,
+    overhead_factor: float = 0.2,
+) -> float:
+    """Calculate additional memory overhead.
+
+    This accounts for CUDA context, fragmentation, temporary buffers, etc.
+
+    Args:
+        total_memory: Total calculated memory in GB
+        overhead_factor: Fraction to add for overhead (default 20%)
+
+    Returns:
+        Overhead memory in GB
+    """
+    return total_memory * overhead_factor
+
+
+def estimate_largest_layer_params(
+    hidden_size: int,
+    num_attention_heads: int,
+    intermediate_size: int | None = None,
+) -> int:
+    """Estimate the largest layer parameters for ZeRO-3 calculations.
+
+    The largest layer is typically the MLP layer or attention projection.
+
+    Args:
+        hidden_size: Hidden dimension size
+        num_attention_heads: Number of attention heads
+        intermediate_size: MLP intermediate size (default 4 * hidden_size)
+
+    Returns:
+        Estimated number of parameters in the largest layer
+    """
+    if intermediate_size is None:
+        intermediate_size = 4 * hidden_size
+
+    # MLP layer: hidden_size * intermediate_size * 2 (for up and down projections)
+    mlp_params = hidden_size * intermediate_size * 2
+
+    # Attention output projection: hidden_size * hidden_size
+    attn_params = hidden_size * hidden_size
+
+    return max(mlp_params, attn_params)

src/gpu_mem_calculator/core/models.py

@@ -0,0 +1,568 @@
+"""Data models for GPU memory calculation."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal, cast
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic_core.core_schema import ValidationInfo as FieldValidationInfo
+
+
+class EngineType(str, Enum):
+    """Supported training engine types."""
+
+    PYTORCH_DDP = "pytorch_ddp"
+    DEEPSPEED = "deepspeed"
+    MEGATRON_LM = "megatron_lm"
+    FSDP = "fsdp"
+    MEGATRON_DEEPSPEED = "megatron_deepspeed"
+
+
+class InferenceEngineType(str, Enum):
+    """Supported inference engine types."""
+
+    HUGGINGFACE = "huggingface"
+    VLLM = "vllm"
+    TGI = "tgi"
+    TENSORRT_LLM = "tensorrt_llm"
+    TRTLLM = "trtllm"
+    SGLANG = "sglang"
+
+
+class OptimizerType(str, Enum):
+    """Supported optimizer types."""
+
+    ADAM = "adam"
+    ADAMW = "adamw"
+    SGD = "sgd"
+    ADAMW_8BIT = "adamw_8bit"
+
+
+class DType(str, Enum):
+    """Supported data types."""
+
+    FP32 = "fp32"
+    FP16 = "fp16"
+    BF16 = "bf16"
+    INT8 = "int8"
+    INT4 = "int4"
+
+
+class OffloadDevice(str, Enum):
+    """CPU offload options."""
+
+    NONE = "none"
+    CPU = "cpu"
+    NVME = "nvme"
+
+
+class ModelConfig(BaseModel):
+    """Model architecture configuration."""
+
+    name: str = Field(default="custom", description="Model name")
+    num_parameters: int = Field(gt=0, description="Total number of parameters")
+    num_layers: int = Field(gt=0, description="Number of transformer layers")
+    hidden_size: int = Field(gt=0, description="Hidden dimension size")
+    num_attention_heads: int = Field(gt=0, description="Number of attention heads")
+    vocab_size: int = Field(default=32000, gt=0, description="Vocabulary size")
+    max_seq_len: int = Field(default=2048, gt=0, description="Maximum sequence length")
+    largest_layer_params: int | None = Field(
+        default=None,
+        gt=0,
+        description="Largest layer parameters (auto-calculated if not provided)",
+    )
+
+    # MoE (Mixture of Experts) parameters
+    moe_enabled: bool = Field(default=False, description="Enable Mixture of Experts")
+    num_experts: int = Field(default=8, ge=1, description="Number of experts in MoE")
+    top_k: int = Field(default=2, ge=1, description="Number of experts activated per token (top-k)")
+    expert_intermediate_size: int | None = Field(
+        default=None,
+        gt=0,
+        description="Expert intermediate layer size (defaults to 4x hidden_size)",
+    )
+    shared_expert_intermediate_size: int | None = Field(
+        default=None,
+        gt=0,
+        description="Shared expert intermediate size (for models like GLM with shared experts)",
+    )
+
+    @model_validator(mode="after")
+    def calculate_largest_layer(self) -> ModelConfig:
+        """Calculate largest layer params if not provided."""
+        if self.largest_layer_params is not None:
+            return self
+        # Calculate it
+        hidden = self.hidden_size
+        moe_enabled = self.moe_enabled
+
+        if hidden and moe_enabled:
+            # For MoE: largest layer includes expert parameters
+            expert_intermediate = self.expert_intermediate_size or hidden * 4
+            self.largest_layer_params = int(hidden * expert_intermediate * 2)
+        elif hidden:
+            # Dense model: attention output + MLP
+            self.largest_layer_params = int(hidden * hidden * 4)
+        return self
+
+    @property
+    def effective_num_experts(self) -> int:
+        """Get effective number of experts (returns 1 if MoE disabled)."""
+        return self.num_experts if self.moe_enabled else 1
+
+    @property
+    def active_experts(self) -> int:
+        """Get number of active experts per token (top_k or 1 if dense)."""
+        return self.top_k if self.moe_enabled else 1
+
+
+class TrainingConfig(BaseModel):
+    """Training hyperparameters configuration."""
+
+    batch_size: int = Field(default=1, gt=0, description="Batch size per GPU")
+    gradient_accumulation_steps: int = Field(
+        default=1,
+        gt=0,
+        description="Gradient accumulation steps",
+    )
+    optimizer: OptimizerType = Field(default=OptimizerType.ADAMW, description="Optimizer type")
+    dtype: DType = Field(default=DType.BF16, description="Data type for training")
+    activation_checkpointing: int = Field(
+        default=0,
+        ge=0,
+        le=4,
+        description="Activation checkpointing level (0-4)",
+    )
+
+    @property
+    def effective_batch_size(self) -> int:
+        """Calculate effective batch size with gradient accumulation."""
+        return self.batch_size * self.gradient_accumulation_steps
+
+
+class ParallelismConfig(BaseModel):
+    """Parallelism configuration."""
+
+    tensor_parallel_size: int = Field(default=1, ge=1, description="Tensor parallelism degree")
+    pipeline_parallel_size: int = Field(default=1, ge=1, description="Pipeline parallelism degree")
+    data_parallel_size: int = Field(default=1, ge=1, description="Data parallelism degree")
+    sequence_parallel: bool = Field(default=False, description="Enable sequence parallelism")
+
+    @property
+    def total_parallel_size(self) -> int:
+        """Calculate total parallelism degree."""
+        return self.tensor_parallel_size * self.pipeline_parallel_size * self.data_parallel_size
+
+
+class EngineConfig(BaseModel):
+    """Training engine specific configuration."""
+
+    type: EngineType = Field(default=EngineType.PYTORCH_DDP, description="Training engine type")
+    zero_stage: int | None = Field(
+        default=None,
+        ge=0,
+        le=3,
+        description="DeepSpeed ZeRO stage (only for DeepSpeed engine)",
+    )
+    offload_optimizer: OffloadDevice = Field(
+        default=OffloadDevice.NONE,
+        description="CPU offload for optimizer states",
+    )
+    offload_param: OffloadDevice = Field(
+        default=OffloadDevice.NONE,
+        description="CPU offload for parameters",
+    )
+    zero_init: bool = Field(
+        default=True,
+        description="Use ZeRO initialization (only for DeepSpeed ZeRO-3)",
+    )
+    sharding_strategy: Literal["no_shard", "shard_grad_op", "full_shard"] = Field(
+        default="full_shard",
+        description="FSDP sharding strategy",
+    )
+
+
+class GPUConfig(BaseModel):
+    """Hardware configuration."""
+
+    num_gpus: int = Field(default=1, ge=1, description="Number of GPUs")
+    gpu_memory_gb: float = Field(default=80.0, gt=0, description="GPU memory in GB")
+    total_gpu_memory_gb: float | None = Field(
+        default=None,
+        description="Total GPU memory (calculated if not provided)",
+    )
+
+    @field_validator("total_gpu_memory_gb")
+    @classmethod
+    def calculate_total_memory(cls, v: float | None, info: FieldValidationInfo) -> float | None:
+        """Calculate total GPU memory if not provided."""
+        if v is None:
+            num_gpus = cast(int, info.data.get("num_gpus", 1))
+            gpu_mem = cast(float, info.data.get("gpu_memory_gb", 80.0))
+            return num_gpus * gpu_mem
+        return v
+
+
+class InterconnectType(str, Enum):
+    """Multi-node interconnect types."""
+
+    INFINIBAND = "infiniband"
+    NVLINK = "nvlink"
+    ETHERNET_10G = "ethernet_10g"
+    ETHERNET_25G = "ethernet_25g"
+    ETHERNET_100G = "ethernet_100g"
+    ETHERNET_200G = "ethernet_200g"
+
+
+class NodeConfig(BaseModel):
+    """Multi-node configuration."""
+
+    num_nodes: int = Field(default=1, ge=1, description="Number of nodes")
+    gpus_per_node: int | None = Field(
+        default=None,
+        ge=1,
+        description="GPUs per node (calculated from num_gpus if not provided)",
+    )
+    interconnect_type: InterconnectType = Field(
+        default=InterconnectType.INFINIBAND,
+        description="Interconnect type between nodes",
+    )
+    interconnect_bandwidth_gbps: float | None = Field(
+        default=None,
+        gt=0,
+        description="Interconnect bandwidth in Gbps (default: auto from type)",
+    )
+
+    @field_validator("gpus_per_node")
+    @classmethod
+    def calculate_gpus_per_node(cls, v: int | None, info: FieldValidationInfo) -> int | None:
+        """Calculate GPUs per node if not provided."""
+        if v is None:
+            num_nodes = cast(int, info.data.get("num_nodes", 1))
+            num_gpus = cast(int, info.data.get("num_gpus", 1))
+            return max(1, num_gpus // num_nodes)
+        return v
+
+    def get_interconnect_bandwidth_gbps(self) -> float:
+        """Get interconnect bandwidth in Gbps.
+
+        Returns bandwidth from config or default based on interconnect type.
+        """
+        if self.interconnect_bandwidth_gbps:
+            return self.interconnect_bandwidth_gbps
+
+        # Default bandwidth values for each interconnect type
+        bandwidth_defaults = {
+            InterconnectType.INFINIBAND: 200.0,  # HDR200 InfiniBand
+            InterconnectType.NVLINK: 300.0,  # NVLink/NVSwitch
+            InterconnectType.ETHERNET_10G: 10.0,
+            InterconnectType.ETHERNET_25G: 25.0,
+            InterconnectType.ETHERNET_100G: 100.0,
+            InterconnectType.ETHERNET_200G: 200.0,
+        }
+        return bandwidth_defaults.get(self.interconnect_type, 100.0)
+
+    @property
+    def is_multi_node(self) -> bool:
+        """Check if this is a multi-node configuration."""
+        return self.num_nodes > 1
+
+
+class NetworkOverhead(BaseModel):
+    """Network communication overhead for multi-node training."""
+
+    allreduce_gb: float = Field(default=0.0, ge=0, description="AllReduce communication in GB")
+    allgather_gb: float = Field(default=0.0, ge=0, description="AllGather communication in GB")
+    reducescatter_gb: float = Field(
+        default=0.0, ge=0, description="ReduceScatter communication in GB"
+    )
+    point_to_point_gb: float = Field(
+        default=0.0, ge=0, description="Point-to-point communication in GB"
+    )
+    total_overhead_gb: float = Field(default=0.0, ge=0, description="Total network overhead in GB")
+    estimated_overhead_ms_per_step: float | None = Field(
+        default=None,
+        description="Estimated communication overhead per training step in milliseconds",
+    )
+
+
+class HybridParallelismConfig(BaseModel):
+    """Hybrid parallelism configuration for optimal multi-node scaling."""
+
+    auto_optimize: bool = Field(
+        default=False,
+        description="Automatically optimize parallelism strategy for given hardware",
+    )
+    target_gpu_utilization: float = Field(
+        default=0.85,
+        gt=0.0,
+        le=1.0,
+        description="Target GPU memory utilization (0.0-1.0)",
+    )
+    prefer_pipeline_parallel: bool = Field(
+        default=False,
+        description="Prefer pipeline parallelism over data parallel for multi-node",
+    )
+    max_pipeline_chunks: int | None = Field(
+        default=None,
+        ge=1,
+        description="Maximum number of pipeline chunks (virtual stages)",
+    )
+    enable_sequence_parallel: bool = Field(
+        default=True,
+        description="Enable sequence parallelism for long sequences",
+    )
+    sequence_parallel_threshold: int = Field(
+        default=4096,
+        ge=1,
+        description="Sequence length threshold for enabling sequence parallel",
+    )
+
+
+class MemoryBreakdown(BaseModel):
+    """Memory calculation result breakdown."""
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    model_params_gb: float = Field(ge=0, description="Model parameters memory in GB")
+    gradients_gb: float = Field(ge=0, description="Gradients memory in GB")
+    optimizer_states_gb: float = Field(ge=0, description="Optimizer states memory in GB")
+    activations_gb: float = Field(ge=0, description="Activations memory in GB")
+    overhead_gb: float = Field(default=0.0, ge=0, description="Additional overhead in GB")
+
+    @property
+    def total_memory_gb(self) -> float:
+        """Total memory in GB."""
+        return (
+            self.model_params_gb
+            + self.gradients_gb
+            + self.optimizer_states_gb
+            + self.activations_gb
+            + self.overhead_gb
+        )
+
+
+class MemoryResult(BaseModel):
+    """Complete memory calculation result."""
+
+    total_memory_per_gpu_gb: float = Field(ge=0, description="Total memory per GPU in GB")
+    total_memory_all_gpus_gb: float = Field(ge=0, description="Total memory across all GPUs in GB")
+    cpu_memory_gb: float = Field(default=0.0, ge=0, description="CPU memory required in GB")
+    breakdown: MemoryBreakdown = Field(description="Memory breakdown by component")
+    network_overhead: NetworkOverhead | None = Field(
+        default=None,
+        description="Network communication overhead for multi-node training",
+    )
+    fits_on_gpu: bool = Field(description="Whether the config fits on available GPU")
+    memory_utilization_percent: float = Field(ge=0, description="Memory utilization percentage")
+    recommended_batch_size: int | None = Field(
+        default=None,
+        description="Recommended batch size if current doesn't fit",
+    )
+    multi_node_info: dict | None = Field(
+        default=None,
+        description="Additional multi-node configuration info",
+    )
+
+
+class KVCacheQuantization(str, Enum):
+    """KV cache quantization options."""
+
+    NONE = "none"
+    INT8 = "int8"
+    FP8 = "fp8"
+    INT4 = "int4"
+
+
+class InferenceMemoryBreakdown(BaseModel):
+    """Memory breakdown for inference workloads."""
+
+    model_config = ConfigDict(protected_namespaces=())
+
+    model_params_gb: float = Field(ge=0, description="Model parameters memory in GB")
+    kv_cache_gb: float = Field(ge=0, description="KV cache memory in GB")
+    activations_gb: float = Field(ge=0, description="Activation memory in GB")
+    overhead_gb: float = Field(default=0.0, ge=0, description="Additional overhead in GB")
+
+    @property
+    def total_memory_gb(self) -> float:
+        """Total memory in GB."""
+        return self.model_params_gb + self.kv_cache_gb + self.activations_gb + self.overhead_gb
+
+
+class InferenceConfig(BaseModel):
+    """Inference-specific configuration."""
+
+    batch_size: int = Field(default=1, gt=0, description="Batch size for inference")
+    max_seq_len: int | None = Field(
+        default=None,
+        gt=0,
+        description="Override max sequence length for inference (default: use model config)",
+    )
+    kv_cache_quantization: KVCacheQuantization = Field(
+        default=KVCacheQuantization.NONE,
+        description="KV cache quantization type",
+    )
+    use_kv_cache: bool = Field(default=True, description="Enable KV cache for generation")
+    tensor_parallel_size: int = Field(default=1, ge=1, description="Tensor parallelism degree")
+    enable_streaming: bool = Field(default=False, description="Enable streaming inference")
+
+    # Common inference options
+    gpu_memory_utilization: float = Field(
+        default=0.9,
+        gt=0.0,
+        le=1.0,
+        description="GPU memory utilization target (0.0-1.0)",
+    )
+
+    # TGI-specific options
+    max_total_tokens: int | None = Field(
+        default=None,
+        gt=0,
+        description="TGI: Maximum total tokens (input + output) - defines memory budget",
+    )
+    max_input_tokens: int | None = Field(
+        default=None,
+        gt=0,
+        description="TGI: Maximum input tokens",
+    )
+    max_batch_total_tokens: int | None = Field(
+        default=None,
+        gt=0,
+        description="TGI: Maximum total tokens across all batches",
+    )
+    tgi_quantize: Literal[
+        "none",
+        "awq",
+        "eetq",
+        "exl2",
+        "gptq",
+        "marlin",
+        "bitsandbytes",
+        "bitsandbytes-nf4",
+        "bitsandbytes-fp4",
+        "fp8",
+    ] = Field(
+        default="none",
+        description="TGI: Weight quantization method",
+    )
+    tgi_dtype: Literal["float16", "bfloat16"] = Field(
+        default="bfloat16",
+        description="TGI: Data type for inference",
+    )
+    sharded: bool = Field(default=False, description="TGI: Enable sharded inference")
+    num_shard: int | None = Field(
+        default=None,
+        ge=1,
+        description="TGI: Number of shards for sharded inference",
+    )
+
+    # vLLM-specific options
+    block_size: int | None = Field(
+        default=None,
+        ge=1,
+        description="vLLM: Block size for KV cache management (default: 16)",
+    )
+    swap_space_gb: float = Field(default=0.0, ge=0.0, description="vLLM: CPU swap space in GB")
+    enable_prefix_caching: bool = Field(default=False, description="vLLM: Enable prefix caching")
+    enforce_eager: bool = Field(
+        default=False,
+        description="vLLM: Enable eager mode (disable CUDA graph)",
+    )
+    max_num_batched_tokens: int | None = Field(
+        default=None,
+        gt=0,
+        description="vLLM: Maximum number of batched tokens",
+    )
+    max_num_seqs: int | None = Field(
+        default=None,
+        gt=0,
+        description="vLLM: Maximum number of sequences in a batch",
+    )
+    vllm_quantization: Literal["none", "awq", "gptq", "squeezellm", "fp8"] = Field(
+        default="none",
+        description="vLLM: Weight quantization method",
+    )
+
+    # TensorRT-LLM-specific options
+    trt_max_batch_size: int | None = Field(
+        default=None,
+        gt=0,
+        description="TensorRT-LLM: Maximum batch size",
+    )
+    trt_max_input_len: int | None = Field(
+        default=None,
+        gt=0,
+        description="TensorRT-LLM: Maximum input length",
+    )
+    trt_max_seq_len: int | None = Field(
+        default=None,
+        gt=0,
+        description="TensorRT-LLM: Maximum sequence length",
+    )
+    trt_max_beam_width: int | None = Field(
+        default=None,
+        ge=1,
+        description="TensorRT-LLM: Maximum beam width for beam search",
+    )
+
+    # SGLang-specific options
+    chunk_size: int | None = Field(
+        default=None,
+        ge=1,
+        description="SGLang: Prefill chunk size for long contexts (default: 8192)",
+    )
+    max_running_requests: int | None = Field(
+        default=None,
+        ge=1,
+        description="SGLang: Maximum number of concurrent requests",
+    )
+    disable_radix_cache: bool = Field(
+        default=False,
+        description="SGLang: Disable RadixAttention cache (for debugging)",
+    )
+    enable_p2p: bool = Field(
+        default=False,
+        description="SGLang: Enable P2P attention for multi-GPU",
+    )
+    disable_custom_all_reduce: bool = Field(
+        default=False,
+        description="SGLang: Disable custom all-reduce kernel",
+    )
+    attention_backend: Literal["flashinfer", "triton", "torch"] = Field(
+        default="flashinfer",
+        description="SGLang: Attention backend implementation",
+    )
+    enable_torch_compile: bool = Field(
+        default=False,
+        description="SGLang: Enable torch.compile for model optimization",
+    )
+    radix_cache_max_seq_len: int | None = Field(
+        default=None,
+        gt=0,
+        description="SGLang: Maximum sequence length for RadixCache",
+    )
+    speculative_algo: Literal["default", "medusa", "eagle"] = Field(
+        default="default",
+        description="SGLang: Speculative decoding algorithm",
+    )
+    multi_lora_enabled: bool = Field(default=False, description="SGLang: Enable multi-LoRA serving")
+
+
+class InferenceMemoryResult(BaseModel):
+    """Inference memory calculation result."""
+
+    total_memory_per_gpu_gb: float = Field(ge=0, description="Total memory per GPU in GB")
+    total_memory_all_gpus_gb: float = Field(ge=0, description="Total memory across all GPUs in GB")
+    breakdown: InferenceMemoryBreakdown = Field(description="Memory breakdown by component")
+    fits_on_gpu: bool = Field(description="Whether the config fits on available GPU")
+    memory_utilization_percent: float = Field(ge=0, description="Memory utilization percentage")
+    max_supported_batch_size: int | None = Field(
+        default=None,
+        description="Maximum batch size that fits in GPU memory",
+    )
+    estimated_throughput_tokens_per_sec: float | None = Field(
+        default=None,
+        description="Estimated throughput in tokens/second",
+    )

src/gpu_mem_calculator/core/multinode.py

@@ -0,0 +1,308 @@
+"""Multi-node training calculator.
+
+Handles network communication overhead calculation and hybrid
+parallelism optimization for multi-node training configurations.
+"""
+
+from gpu_mem_calculator.core.models import (
+    EngineConfig,
+    EngineType,
+    HybridParallelismConfig,
+    ModelConfig,
+    NetworkOverhead,
+    NodeConfig,
+    ParallelismConfig,
+    TrainingConfig,
+)
+
+
+class MultiNodeCalculator:
+    """Calculator for multi-node training overhead and optimization.
+
+    This class provides:
+    - Network communication overhead estimation
+    - Hybrid parallelism strategy optimization
+    - Multi-node performance modeling
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        training_config: TrainingConfig,
+        parallelism_config: ParallelismConfig,
+        node_config: NodeConfig,
+        engine_config: EngineConfig,
+    ) -> None:
+        """Initialize the multi-node calculator.
+
+        Args:
+            model_config: Model architecture configuration
+            training_config: Training hyperparameters
+            parallelism_config: Parallelism settings
+            node_config: Multi-node hardware configuration
+            engine_config: Training engine configuration
+        """
+        self.model_config = model_config
+        self.training_config = training_config
+        self.parallelism_config = parallelism_config
+        self.node_config = node_config
+        self.engine_config = engine_config
+
+    def calculate_network_overhead(self) -> NetworkOverhead:
+        """Calculate network communication overhead for multi-node training.
+
+        Estimates communication overhead for different collective operations
+        based on model size, parallelism strategy, and interconnect bandwidth.
+
+        Returns:
+            NetworkOverhead with detailed breakdown
+        """
+        if not self.node_config.is_multi_node:
+            return NetworkOverhead()
+
+        # Get model size in bytes
+        model_params = self.model_config.num_parameters
+        dtype_bytes = self._get_dtype_bytes()
+        model_size_bytes = int(model_params * dtype_bytes)
+
+        # Calculate communication for each collective operation
+        allreduce_gb = self._calculate_allreduce_overhead(model_size_bytes)
+        allgather_gb = self._calculate_allgather_overhead(model_size_bytes)
+        reducescatter_gb = self._calculate_reducescatter_overhead(model_size_bytes)
+        point_to_point_gb = self._calculate_pipeline_overhead(model_size_bytes)
+
+        total_overhead_gb = allreduce_gb + allgather_gb + reducescatter_gb + point_to_point_gb
+
+        # Estimate time overhead per step
+        overhead_ms = self._estimate_communication_time_ms(total_overhead_gb)
+
+        return NetworkOverhead(
+            allreduce_gb=allreduce_gb,
+            allgather_gb=allgather_gb,
+            reducescatter_gb=reducescatter_gb,
+            point_to_point_gb=point_to_point_gb,
+            total_overhead_gb=total_overhead_gb,
+            estimated_overhead_ms_per_step=overhead_ms,
+        )
+
+    def optimize_hybrid_parallelism(
+        self,
+        hybrid_config: HybridParallelismConfig,
+    ) -> ParallelismConfig:
+        """Optimize hybrid parallelism strategy for multi-node training.
+
+        Analyzes the hardware configuration and model characteristics
+        to recommend optimal parallelism degrees.
+
+        Args:
+            hybrid_config: Hybrid parallelism configuration and preferences
+
+        Returns:
+            Optimized ParallelismConfig
+        """
+        if not hybrid_config.auto_optimize:
+            return self.parallelism_config
+
+        num_nodes = self.node_config.num_nodes
+        gpus_per_node = self.node_config.gpus_per_node or 1
+        total_gpus = num_nodes * gpus_per_node
+
+        seq_len = self.model_config.max_seq_len
+
+        # Determine optimal parallelism strategy
+        if seq_len >= hybrid_config.sequence_parallel_threshold:
+            # Enable sequence parallel for long sequences
+            enable_sp = True
+        else:
+            enable_sp = hybrid_config.enable_sequence_parallel
+
+        # Calculate parallelism degrees
+        if hybrid_config.prefer_pipeline_parallel and num_nodes > 1:
+            # Prefer pipeline parallel across nodes
+            pp_size = int(min(num_nodes, 8))  # Limit pipeline stages
+            tp_size = int(min(gpus_per_node, 8))  # Tensor parallel within node
+            dp_size = int(total_gpus // (pp_size * tp_size))
+        else:
+            # Default: maximize data parallel
+            tp_size = 1
+            pp_size = 1
+            dp_size = int(total_gpus)
+
+        # Ensure all values are at least 1
+        tp_size = max(1, tp_size)
+        pp_size = max(1, pp_size)
+        dp_size = max(1, dp_size)
+
+        return ParallelismConfig(
+            tensor_parallel_size=tp_size,
+            pipeline_parallel_size=pp_size,
+            data_parallel_size=dp_size,
+            sequence_parallel=enable_sp,
+        )
+
+    def _calculate_allreduce_overhead(self, model_size_bytes: int) -> float:
+        """Calculate AllReduce communication overhead.
+
+        AllReduce is used for gradient averaging in data parallel training.
+        Algorithm: Ring AllReduce with O(2 * model_size) communication.
+
+        Args:
+            model_size_bytes: Model size in bytes
+
+        Returns:
+            Communication volume in GB
+        """
+        # Ring AllReduce: each GPU sends/receives 2 * model_size / num_gpus
+        # But we need the total across the network
+
+        # For gradient averaging: 2 * model_size (send + receive)
+        allreduce_bytes = 2 * model_size_bytes
+
+        # Adjust for collective operation efficiency
+        # In multi-node, cross-node traffic is the bottleneck
+        if self.node_config.is_multi_node:
+            # Only cross-node traffic matters
+            allreduce_bytes = int(allreduce_bytes / self.node_config.num_nodes)
+
+        return allreduce_bytes / (1024**3)
+
+    def _calculate_allgather_overhead(self, model_size_bytes: int) -> float:
+        """Calculate AllGather communication overhead.
+
+        AllGather is used in ZeRO-3 and tensor parallel for parameter gathering.
+
+        Args:
+            model_size_bytes: Model size in bytes
+
+        Returns:
+            Communication volume in GB
+        """
+        # AllGather: (num_gpus - 1) * model_size / num_gpus per GPU
+        # But for ZeRO-3, we gather all parameters
+        is_zero3 = (
+            self.engine_config.type == EngineType.DEEPSPEED and self.engine_config.zero_stage == 3
+        )
+
+        if is_zero3:
+            # ZeRO-3 gathers all parameters during forward pass
+            allgather_bytes = model_size_bytes
+        else:
+            # Standard allgather for tensor parallel
+            allgather_bytes = int(model_size_bytes / self.parallelism_config.tensor_parallel_size)
+
+        # Adjust for multi-node
+        if self.node_config.is_multi_node:
+            allgather_bytes = int(allgather_bytes / self.node_config.num_nodes)
+
+        return allgather_bytes / (1024**3)
+
+    def _calculate_reducescatter_overhead(self, model_size_bytes: int) -> float:
+        """Calculate ReduceScatter communication overhead.
+
+        ReduceScatter is used in ZeRO-2 and gradient sharding.
+
+        Args:
+            model_size_bytes: Model size in bytes
+
+        Returns:
+            Communication volume in GB
+        """
+        is_zero2 = (
+            self.engine_config.type == EngineType.DEEPSPEED and self.engine_config.zero_stage == 2
+        )
+
+        if is_zero2:
+            # ZeRO-2 scatters gradients
+            reducescatter_bytes = model_size_bytes
+        else:
+            # Standard reducescatter
+            reducescatter_bytes = int(model_size_bytes / self.parallelism_config.data_parallel_size)
+
+        # Adjust for multi-node
+        if self.node_config.is_multi_node:
+            reducescatter_bytes = int(reducescatter_bytes / self.node_config.num_nodes)
+
+        return reducescatter_bytes / (1024**3)
+
+    def _calculate_pipeline_overhead(self, model_size_bytes: int) -> float:
+        """Calculate pipeline parallel communication overhead.
+
+        Point-to-point communication between pipeline stages.
+
+        Args:
+            model_size_bytes: Model size in bytes
+
+        Returns:
+            Communication volume in GB
+        """
+        if self.parallelism_config.pipeline_parallel_size <= 1:
+            return 0.0
+
+        # Pipeline parallel sends activations between stages
+        # Approximate as layer activations
+        hidden_size = self.model_config.hidden_size
+        seq_len = self.model_config.max_seq_len
+        batch_size = self.training_config.batch_size
+        num_layers = self.model_config.num_layers
+
+        # Activation size per layer
+        activation_bytes = batch_size * seq_len * hidden_size * 2  # FP16/BF16
+
+        # Number of microbatches determines communication frequency
+        # For simplicity, assume num_stages communications per step
+        pp_size = self.parallelism_config.pipeline_parallel_size
+        pipeline_comm_bytes = activation_bytes * (num_layers // pp_size)
+
+        # Adjust for multi-node
+        if self.node_config.is_multi_node:
+            pipeline_comm_bytes = int(pipeline_comm_bytes / self.node_config.num_nodes)
+
+        return pipeline_comm_bytes / (1024**3)
+
+    def _estimate_communication_time_ms(self, total_gb: float) -> float:
+        """Estimate communication time per training step in milliseconds.
+
+        Args:
+            total_gb: Total communication volume in GB
+
+        Returns:
+            Estimated time in milliseconds
+        """
+        if total_gb == 0:
+            return 0.0
+
+        # Get bandwidth in GB/s
+        bandwidth_gbps = self.node_config.get_interconnect_bandwidth_gbps()
+        bandwidth_gbps_per_sec = bandwidth_gbps / 8  # Convert to GB/s
+
+        # Basic time = size / bandwidth
+        time_seconds = total_gb / bandwidth_gbps_per_sec
+
+        # Add latency overhead for collective operations
+        # Typical latency: 10-50 microseconds per hop
+        num_nodes = self.node_config.num_nodes
+        latency_overhead = num_nodes * 0.00005  # 50 microseconds per node
+
+        # Network efficiency factor (not 100% efficient)
+        efficiency = 0.85
+
+        total_time_seconds = (time_seconds / efficiency) + latency_overhead
+
+        return total_time_seconds * 1000  # Convert to ms
+
+    def _get_dtype_bytes(self) -> float:
+        """Get bytes per element based on dtype."""
+        dtype_map = {
+            "fp32": 4,
+            "fp16": 2,
+            "bf16": 2,
+            "int8": 1,
+            "int4": 0.5,
+        }
+        return dtype_map.get(self.training_config.dtype.value, 2)
+
+    def _calculate_model_size_gb(self) -> float:
+        """Calculate model size in GB."""
+        dtype_bytes = self._get_dtype_bytes()
+        model_size_bytes = self.model_config.num_parameters * dtype_bytes
+        return model_size_bytes / (1024**3)

src/gpu_mem_calculator/engines/__init__.py

@@ -0,0 +1,16 @@
+"""Training engine implementations."""
+
+from gpu_mem_calculator.engines.base import BaseEngine
+from gpu_mem_calculator.engines.deepspeed import DeepSpeedEngine
+from gpu_mem_calculator.engines.fsdp import FSDPEngine
+from gpu_mem_calculator.engines.megatron import MegatronDeepSpeedEngine, MegatronLMEngine
+from gpu_mem_calculator.engines.pytorch import PyTorchDDPEngine
+
+__all__ = [
+    "BaseEngine",
+    "PyTorchDDPEngine",
+    "DeepSpeedEngine",
+    "MegatronLMEngine",
+    "MegatronDeepSpeedEngine",
+    "FSDPEngine",
+]

src/gpu_mem_calculator/engines/base.py

@@ -0,0 +1,220 @@
+"""Base class for training engine implementations."""
+
+from abc import ABC, abstractmethod
+
+from gpu_mem_calculator.core.models import (
+    EngineConfig,
+    GPUConfig,
+    MemoryBreakdown,
+    MemoryResult,
+    ModelConfig,
+    NodeConfig,
+    ParallelismConfig,
+    TrainingConfig,
+)
+
+
+class BaseEngine(ABC):
+    """Abstract base class for training engine memory calculation.
+
+    Each training engine (PyTorch DDP, DeepSpeed, Megatron-LM, etc.)
+    should implement this interface to provide engine-specific
+    memory calculations.
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        training_config: TrainingConfig,
+        parallelism_config: ParallelismConfig,
+        engine_config: EngineConfig,
+        gpu_config: GPUConfig,
+        node_config: NodeConfig | None = None,
+    ) -> None:
+        """Initialize the engine with configuration.
+
+        Args:
+            model_config: Model architecture configuration
+            training_config: Training hyperparameters
+            parallelism_config: Parallelism settings
+            engine_config: Engine-specific configuration
+            gpu_config: Hardware configuration
+            node_config: Multi-node configuration (optional)
+        """
+        self.model_config = model_config
+        self.training_config = training_config
+        self.parallelism_config = parallelism_config
+        self.engine_config = engine_config
+        self.gpu_config = gpu_config
+        self.node_config = node_config or NodeConfig()
+
+    @abstractmethod
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory requirements for this engine.
+
+        This is the main method that should be implemented by each engine.
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        pass
+
+    def _check_feasibility(
+        self,
+        total_memory_per_gpu: float,
+    ) -> tuple[bool, float, int | None]:
+        """Check if the configuration fits on available GPU.
+
+        Args:
+            total_memory_per_gpu: Total memory required per GPU
+
+        Returns:
+            Tuple of (fits_on_gpu, utilization_percent, recommended_batch_size)
+        """
+        available_memory = self.gpu_config.gpu_memory_gb
+        utilization_percent = (total_memory_per_gpu / available_memory) * 100
+
+        fits_on_gpu = total_memory_per_gpu <= available_memory
+
+        # If doesn't fit, suggest a smaller batch size
+        recommended_batch_size = None
+        if not fits_on_gpu:
+            # Simple heuristic: scale batch size inversely with memory excess
+            excess_factor = total_memory_per_gpu / available_memory
+            recommended_batch_size = max(1, int(self.training_config.batch_size / excess_factor))
+
+        return fits_on_gpu, utilization_percent, recommended_batch_size
+
+    def _create_result(
+        self,
+        breakdown: MemoryBreakdown,
+        cpu_memory_gb: float = 0.0,
+    ) -> MemoryResult:
+        """Create a MemoryResult from breakdown.
+
+        Args:
+            breakdown: Memory breakdown by component
+            cpu_memory_gb: CPU memory required (default 0)
+
+        Returns:
+            Complete MemoryResult
+        """
+        total_memory_per_gpu = breakdown.total_memory_gb
+        total_memory_all_gpus = total_memory_per_gpu * self.gpu_config.num_gpus
+
+        fits_on_gpu, utilization_percent, recommended_batch_size = self._check_feasibility(
+            total_memory_per_gpu
+        )
+
+        # Calculate network overhead for multi-node configurations
+        network_overhead = None
+        multi_node_info = None
+        if self.node_config.is_multi_node:
+            from gpu_mem_calculator.core.multinode import MultiNodeCalculator
+
+            multinode_calc = MultiNodeCalculator(
+                model_config=self.model_config,
+                training_config=self.training_config,
+                parallelism_config=self.parallelism_config,
+                node_config=self.node_config,
+                engine_config=self.engine_config,
+            )
+            network_overhead = multinode_calc.calculate_network_overhead()
+
+            # Add multi-node info
+            multi_node_info = {
+                "num_nodes": self.node_config.num_nodes,
+                "gpus_per_node": self.node_config.gpus_per_node,
+                "interconnect_type": self.node_config.interconnect_type.value,
+                "interconnect_bandwidth_gbps": self.node_config.get_interconnect_bandwidth_gbps(),
+            }
+
+        return MemoryResult(
+            total_memory_per_gpu_gb=total_memory_per_gpu,
+            total_memory_all_gpus_gb=total_memory_all_gpus,
+            cpu_memory_gb=cpu_memory_gb,
+            breakdown=breakdown,
+            network_overhead=network_overhead,
+            fits_on_gpu=fits_on_gpu,
+            memory_utilization_percent=utilization_percent,
+            recommended_batch_size=recommended_batch_size,
+            multi_node_info=multi_node_info,
+        )
+
+    @property
+    def effective_batch_size(self) -> int:
+        """Calculate effective batch size with gradient accumulation."""
+        return (
+            self.training_config.batch_size
+            * self.training_config.gradient_accumulation_steps
+            * self.parallelism_config.data_parallel_size
+        )
+
+    @property
+    def total_num_gpus(self) -> int:
+        """Get total number of GPUs."""
+        return self.gpu_config.num_gpus
+
+    @property
+    def num_gpus_per_model(self) -> int:
+        """Get number of GPUs per model replica.
+
+        This is tensor_parallel * pipeline_parallel for distributed training.
+        """
+        return (
+            self.parallelism_config.tensor_parallel_size
+            * self.parallelism_config.pipeline_parallel_size
+        )
+
+    def calculate_moe_activation_multiplier(self) -> float:
+        """Calculate activation memory multiplier for MoE models.
+
+        For MoE models, activation memory depends on top_k (active experts per token)
+        rather than total number of experts. This is because only top_k experts
+        are activated per token during forward/backward pass.
+
+        Returns:
+            Multiplier for activation memory (1.0 for dense models, <1 for MoE)
+        """
+        if not self.model_config.moe_enabled:
+            return 1.0
+
+        # For MoE: only top_k experts are active per token
+        # Activation memory scales with active_experts / total_experts
+        # But we also have router overhead and gating network activations
+
+        num_experts = self.model_config.num_experts
+        top_k = self.model_config.top_k
+
+        # Base activation ratio: only top_k experts active
+        activation_ratio = top_k / num_experts
+
+        # Add router overhead (typically 5-15% extra for gating)
+        router_overhead = 0.1
+
+        # For models with shared experts (like GLM), adjust accordingly
+        if self.model_config.shared_expert_intermediate_size:
+            # Shared expert is always active, so add its contribution
+            # This is a simplified approximation
+            activation_ratio = activation_ratio + (1.0 / num_experts)
+
+        return min(1.0, activation_ratio + router_overhead)
+
+    def calculate_moe_parameter_ratio(self) -> float:
+        """Calculate effective parameter ratio for MoE models.
+
+        For MoE models, only top_k experts are used during forward pass,
+        but all expert parameters are stored in memory.
+
+        Returns:
+            Ratio of active parameters to total parameters (for memory estimation)
+        """
+        if not self.model_config.moe_enabled:
+            return 1.0
+
+        # All expert parameters are stored, but only top_k are used per token
+        # For gradient calculation, we need gradients for all experts
+        # So parameter storage = 1.0 (all params stored)
+        # But we can use this for inference-specific calculations
+
+        return 1.0  # All parameters stored in memory

src/gpu_mem_calculator/engines/deepspeed.py

@@ -0,0 +1,316 @@
+"""DeepSpeed ZeRO engine implementation.
+
+Implements memory calculations for DeepSpeed ZeRO stages 1, 2, and 3.
+Based on: https://deepspeed.readthedocs.io/en/latest/memory.html
+"""
+
+from gpu_mem_calculator.core.formulas import (
+    calculate_activation_memory,
+    calculate_overhead,
+    estimate_largest_layer_params,
+)
+from gpu_mem_calculator.core.models import (
+    MemoryBreakdown,
+    MemoryResult,
+    OffloadDevice,
+)
+from gpu_mem_calculator.engines.base import BaseEngine
+from gpu_mem_calculator.utils.precision import gb_from_bytes
+
+
+class DeepSpeedEngine(BaseEngine):
+    """DeepSpeed ZeRO memory calculation.
+
+    Implements ZeRO stages:
+    - ZeRO-1: Shard optimizer states
+    - ZeRO-2: Shard optimizer states + gradients
+    - ZeRO-3: Shard optimizer states + gradients + parameters
+    """
+
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory requirements for DeepSpeed ZeRO training.
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        zero_stage = self.engine_config.zero_stage or 0
+        offload_optimizer = self.engine_config.offload_optimizer
+        offload_param = self.engine_config.offload_param
+
+        # Get largest layer params for ZeRO-3
+        if self.model_config.largest_layer_params is None:
+            largest_layer_params = estimate_largest_layer_params(
+                hidden_size=self.model_config.hidden_size,
+                num_attention_heads=self.model_config.num_attention_heads,
+            )
+        else:
+            largest_layer_params = self.model_config.largest_layer_params
+
+        match zero_stage:
+            case 0:
+                return self._calculate_zero0()
+            case 1:
+                return self._calculate_zero1(offload_optimizer)
+            case 2:
+                return self._calculate_zero2(offload_optimizer)
+            case 3:
+                return self._calculate_zero3(
+                    offload_optimizer,
+                    offload_param,
+                    largest_layer_params,
+                )
+            case _:
+                # Default to ZeRO-2
+                return self._calculate_zero2(offload_optimizer)
+
+    def _calculate_zero0(self) -> MemoryResult:
+        """Calculate memory for ZeRO-0 (disabled, same as PyTorch DDP)."""
+        # Import here to avoid circular dependency
+        from gpu_mem_calculator.engines.pytorch import PyTorchDDPEngine
+
+        # ZeRO-0 is the same as PyTorch DDP
+        ddp_engine = PyTorchDDPEngine(
+            model_config=self.model_config,
+            training_config=self.training_config,
+            parallelism_config=self.parallelism_config,
+            engine_config=self.engine_config,
+            gpu_config=self.gpu_config,
+        )
+        return ddp_engine.calculate_memory()
+
+    def _calculate_zero1(
+        self,
+        offload_optimizer: OffloadDevice,
+    ) -> MemoryResult:
+        """Calculate memory for ZeRO-1 (shard optimizer states).
+
+        ZeRO-1 shards optimizer states across data parallel GPUs.
+
+        Reference: https://deepspeed.readthedocs.io/en/latest/memory.html
+        Reference: https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
+
+        Memory formula:
+        - offload_optimizer=cpu: 2 * params (fp16 params only on GPU)
+        - offload_optimizer=none: 4 * params (fp16 params + fp32 params) +
+          12 * params / num_gpus (sharded optimizer states)
+
+        Note: Optimizer states = 12 bytes per param for Adam/AdamW
+        - 4 bytes: FP32 parameter copy
+        - 4 bytes: Momentum (FP32)
+        - 4 bytes: Variance (FP32)
+        """
+        num_params = self.model_config.num_parameters
+        num_gpus = self.total_num_gpus
+
+        # Model parameters (fp16/bf16 on GPU)
+        model_params_gb = gb_from_bytes(num_params * 2)  # FP16/BF16 = 2 bytes
+
+        # Gradients (fp16 on GPU)
+        gradients_gb = gb_from_bytes(num_params * 2)
+
+        # Optimizer states (sharded across GPUs, possibly offloaded to CPU)
+        # 12 bytes per param for Adam/AdamW (FP32 params copy + momentum + variance)
+        if offload_optimizer == OffloadDevice.CPU:
+            # Offloaded to CPU, minimal GPU memory for optimizer
+            optimizer_gb = 0.0
+            cpu_memory_gb = gb_from_bytes(num_params * 12)  # Full optimizer on CPU
+        else:
+            # Sharded across GPUs: 12 bytes / num_gpus per GPU
+            optimizer_gb = gb_from_bytes((num_params * 12) / num_gpus)
+            cpu_memory_gb = 0.0
+
+        # Activations (same as baseline)
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = calculate_overhead(base_memory)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown, cpu_memory_gb)
+
+    def _calculate_zero2(
+        self,
+        offload_optimizer: OffloadDevice,
+    ) -> MemoryResult:
+        """Calculate memory for ZeRO-2 (shard optimizer + gradients).
+
+        ZeRO-2 shards optimizer states AND gradients across data parallel GPUs.
+
+        Reference: https://deepspeed.readthedocs.io/en/latest/memory.html
+        Reference: https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
+
+        Memory formula:
+        - offload_optimizer=cpu: 2 * params (fp16 params) +
+          (2 * params / num_gpus) (sharded fp16 grads)
+        - offload_optimizer=none: 2 * params (fp16 params) +
+          2 * params / num_gpus (sharded fp16 grads) +
+          12 * params / num_gpus (sharded optimizer states)
+
+        Note: Unlike ZeRO-1, ZeRO-2 shards gradients across GPUs
+        """
+        num_params = self.model_config.num_parameters
+        num_gpus = self.total_num_gpus
+
+        # Model parameters (fp16/bf16 on GPU) - NOT sharded in ZeRO-2
+        model_params_gb = gb_from_bytes(num_params * 2)  # FP16/BF16 = 2 bytes
+
+        # Gradients (fp16 on GPU) - SHARDED in ZeRO-2
+        gradients_gb = gb_from_bytes((num_params * 2) / num_gpus)
+
+        # Optimizer states (sharded across GPUs, possibly offloaded to CPU)
+        # 12 bytes per param for Adam/AdamW (FP32 params copy + momentum + variance)
+        if offload_optimizer == OffloadDevice.CPU:
+            # Offloaded to CPU, minimal GPU memory for optimizer
+            optimizer_gb = 0.0
+            cpu_memory_gb = gb_from_bytes(num_params * 12)  # Full optimizer on CPU
+        else:
+            # Sharded across GPUs: 12 bytes / num_gpus per GPU
+            optimizer_gb = gb_from_bytes((num_params * 12) / num_gpus)
+            cpu_memory_gb = 0.0
+
+        # Activations (same as baseline)
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = calculate_overhead(base_memory)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown, cpu_memory_gb)
+
+    def _calculate_zero3(
+        self,
+        offload_optimizer: OffloadDevice,
+        offload_param: OffloadDevice,
+        largest_layer_params: int,
+    ) -> MemoryResult:
+        """Calculate memory for ZeRO-3 (shard params + optimizer + gradients).
+
+        ZeRO-3 shards everything across GPUs.
+
+        Reference: https://deepspeed.readthedocs.io/en/latest/memory.html
+        Reference: https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
+
+        Memory formula:
+        - largest_layer_memory = 4 * largest_layer_params (fp16 params + fp16 grads)
+
+        Case 1 (no offload):
+          largest_layer_memory + 18 * params / num_gpus
+          (where 18 = 16 bytes optimizer states + 2 bytes fp16 params)
+
+        Case 2 (param + optimizer offload to CPU):
+          largest_layer_memory (main limit is CPU RAM)
+
+        Case 3 (optimizer offload to CPU only):
+          largest_layer_memory + 2 * params / num_gpus
+
+        Note: Optimizer states = 16 bytes per param for Adam/AdamW (FP32)
+        - 4 bytes: FP32 parameter copy
+        - 4 bytes: Momentum (FP32)
+        - 4 bytes: Variance (FP32)
+        - 4 bytes: Gradient (FP32 copy for optimizer update)
+        """
+        num_params = self.model_config.num_parameters
+        num_gpus = self.total_num_gpus
+
+        # Largest layer memory (fp16 params + fp16 grads gathered on one GPU)
+        largest_layer_memory_gb = gb_from_bytes(largest_layer_params * 4)
+
+        # Calculate memory based on offload configuration
+        if offload_param == OffloadDevice.CPU and offload_optimizer == OffloadDevice.CPU:
+            # Case 2: Both params and optimizer offloaded to CPU
+            # Only need largest layer on GPU at a time
+            params_per_gpu_gb = 0.0
+            gradients_per_gpu_gb = 0.0
+            optimizer_gb = 0.0
+            cpu_memory_gb = gb_from_bytes(num_params * 18)  # Full model on CPU
+        elif offload_optimizer == OffloadDevice.CPU:
+            # Case 3: Only optimizer offloaded to CPU
+            params_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+            gradients_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+            optimizer_gb = 0.0
+            cpu_memory_gb = gb_from_bytes(num_params * 16)  # Optimizer on CPU
+        else:
+            # Case 1: No offload
+            params_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+            gradients_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+            optimizer_gb = gb_from_bytes((num_params * 16) / num_gpus)  # FP32
+            cpu_memory_gb = 0.0
+
+        # Model params = largest layer for ZeRO-3
+        model_params_gb = largest_layer_memory_gb
+
+        # Activations
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Overhead
+        base_memory = (
+            model_params_gb
+            + params_per_gpu_gb
+            + gradients_per_gpu_gb
+            + optimizer_gb
+            + activations_gb
+        )
+        overhead_gb = calculate_overhead(base_memory)
+
+        # For ZeRO-3, we combine params/gradients/optimizer into model_params in breakdown
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb + params_per_gpu_gb,
+            gradients_gb=gradients_per_gpu_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown, cpu_memory_gb)

src/gpu_mem_calculator/engines/fsdp.py

@@ -0,0 +1,213 @@
+"""FSDP (Fully Sharded Data Parallel) engine implementation.
+
+Implements memory calculations for PyTorch FSDP.
+
+Reference: https://pytorch.org/docs/stable/fsdp.html
+Reference: https://blog.eleuther.ai/transformer-math/
+"""
+
+from gpu_mem_calculator.core.formulas import (
+    calculate_activation_memory,
+    calculate_overhead,
+    estimate_largest_layer_params,
+)
+from gpu_mem_calculator.core.models import (
+    MemoryBreakdown,
+    MemoryResult,
+)
+from gpu_mem_calculator.engines.base import BaseEngine
+from gpu_mem_calculator.utils.precision import gb_from_bytes
+
+
+class FSDPEngine(BaseEngine):
+    """PyTorch FSDP memory calculation.
+
+    FSDP shards model parameters, gradients, and optimizer states
+    across data parallel GPUs, similar to DeepSpeed ZeRO-3.
+
+    Sharding strategies:
+    - NO_SHARD: Equivalent to DDP (no sharding)
+    - SHARD_GRAD_OP: Shard gradients and optimizer states (like ZeRO-2)
+    - FULL_SHARD: Shard everything (like ZeRO-3)
+    """
+
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory requirements for FSDP training.
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        sharding_strategy = self.engine_config.sharding_strategy
+
+        # Get largest layer params for FULL_SHARD
+        if self.model_config.largest_layer_params is None:
+            largest_layer_params = estimate_largest_layer_params(
+                hidden_size=self.model_config.hidden_size,
+                num_attention_heads=self.model_config.num_attention_heads,
+            )
+        else:
+            largest_layer_params = self.model_config.largest_layer_params
+
+        match sharding_strategy:
+            case "no_shard":
+                return self._calculate_no_shard()
+            case "shard_grad_op":
+                return self._calculate_shard_grad_op()
+            case "full_shard":
+                return self._calculate_full_shard(largest_layer_params)
+            case _:
+                # Default to full shard
+                return self._calculate_full_shard(largest_layer_params)
+
+    def _calculate_no_shard(self) -> MemoryResult:
+        """Calculate memory for NO_SHARD (same as DDP).
+
+        No sharding - each GPU holds a full copy of the model.
+        """
+        # Import PyTorch DDP engine
+        from gpu_mem_calculator.engines.pytorch import PyTorchDDPEngine
+
+        ddp_engine = PyTorchDDPEngine(
+            model_config=self.model_config,
+            training_config=self.training_config,
+            parallelism_config=self.parallelism_config,
+            engine_config=self.engine_config,
+            gpu_config=self.gpu_config,
+        )
+        return ddp_engine.calculate_memory()
+
+    def _calculate_shard_grad_op(self) -> MemoryResult:
+        """Calculate memory for SHARD_GRAD_OP.
+
+        Shards gradients and optimizer states across GPUs.
+        Similar to DeepSpeed ZeRO-2.
+
+        Reference: https://pytorch.org/tutorials/intermediate/FSDP_advanced.html
+        Reference: https://blog.eleuther.ai/transformer-math/
+
+        Memory formula:
+        - Model parameters: Full model on each GPU (not sharded)
+        - Gradients: Sharded across GPUs
+        - Optimizer states: Sharded across GPUs (12 bytes per param for Adam/AdamW)
+
+        Note: Optimizer states = 12 bytes per param for Adam/AdamW
+        - 4 bytes: FP32 parameter copy
+        - 4 bytes: Momentum (FP32)
+        - 4 bytes: Variance (FP32)
+        """
+        num_params = self.model_config.num_parameters
+        num_gpus = self.total_num_gpus
+
+        # Model parameters (full model on each GPU)
+        model_params_gb = gb_from_bytes(num_params * 2)  # FP16/BF16
+
+        # Gradients (sharded)
+        gradients_gb = gb_from_bytes((num_params * 2) / num_gpus)
+
+        # Optimizer states (sharded) - 12 bytes per param for Adam/AdamW
+        optimizer_gb = gb_from_bytes((num_params * 12) / num_gpus)  # FP32
+
+        # Activations
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = calculate_overhead(base_memory)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown)
+
+    def _calculate_full_shard(self, largest_layer_params: int) -> MemoryResult:
+        """Calculate memory for FULL_SHARD.
+
+        Shards parameters, gradients, and optimizer states.
+        Similar to DeepSpeed ZeRO-3.
+
+        Reference: https://pytorch.org/tutorials/intermediate/FSDP_advanced.html
+        Reference: https://blog.eleuther.ai/transformer-math/
+
+        Memory formula:
+        - Largest layer: 4 * largest_layer_params (fp16 params + fp16 grads)
+        - Remaining parameters and gradients: Sharded across GPUs (2 bytes fp16 each)
+        - Optimizer states: Sharded across GPUs (12 bytes per param for Adam/AdamW in FP32)
+
+        Total per GPU: largest_layer_memory + 2 * params / num_gpus +
+                       2 * params / num_gpus + 12 * params / num_gpus
+                    = largest_layer_memory + 16 * params / num_gpus
+
+        Note: FSDP typically uses 12 bytes for optimizer states (not 16 like DeepSpeed ZeRO-3)
+        because FSDP doesn't keep an additional FP32 gradient copy in the optimizer states.
+        """
+        num_params = self.model_config.num_parameters
+        num_gpus = self.total_num_gpus
+
+        # Largest layer memory (fp16 params + fp16 grads gathered during compute)
+        largest_layer_memory_gb = gb_from_bytes(largest_layer_params * 4)
+
+        # Sharded parameters (fp16)
+        params_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+
+        # Sharded gradients (fp16)
+        gradients_per_gpu_gb = gb_from_bytes((num_params * 2) / num_gpus)
+
+        # Sharded optimizer states (FP32 for Adam/AdamW)
+        # 12 bytes per param: 4 bytes fp32 params copy + 4 bytes momentum + 4 bytes variance
+        optimizer_per_gpu_gb = gb_from_bytes((num_params * 12) / num_gpus)
+
+        # Model params in breakdown: largest layer (gathered) + sharded params
+        # This represents the total parameter memory on each GPU
+        model_params_gb = largest_layer_memory_gb + params_per_gpu_gb
+
+        # Activations
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Overhead
+        base_memory = (
+            largest_layer_memory_gb
+            + params_per_gpu_gb
+            + gradients_per_gpu_gb
+            + optimizer_per_gpu_gb
+            + activations_gb
+        )
+        overhead_gb = calculate_overhead(base_memory)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_per_gpu_gb,
+            optimizer_states_gb=optimizer_per_gpu_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown)

src/gpu_mem_calculator/engines/megatron.py

@@ -0,0 +1,257 @@
+"""Megatron-LM engine implementation.
+
+Implements memory calculations for Megatron-LM with tensor, pipeline,
+and sequence parallelism.
+
+Reference: https://github.com/NVIDIA/Megatron-LM
+Reference: https://arxiv.org/abs/1909.08053
+Reference: https://blog.eleuther.ai/transformer-math/
+"""
+
+from gpu_mem_calculator.core.formulas import (
+    calculate_activation_memory,
+    calculate_gradient_memory,
+    calculate_optimizer_memory,
+    calculate_overhead,
+    calculate_parameter_memory,
+)
+from gpu_mem_calculator.core.models import (
+    MemoryBreakdown,
+    MemoryResult,
+)
+from gpu_mem_calculator.engines.base import BaseEngine
+from gpu_mem_calculator.utils.precision import gb_from_bytes
+
+
+class MegatronLMEngine(BaseEngine):
+    """Megatron-LM memory calculation.
+
+    Megatron-LM uses tensor parallelism to split individual layers across GPUs,
+    and optionally pipeline parallelism to split layers across GPUs.
+    """
+
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory requirements for Megatron-LM training.
+
+        Megatron-LM memory characteristics:
+        - Parameters are sharded across tensor parallel GPUs
+        - Gradients are sharded across tensor parallel GPUs
+        - Optimizer states can be sharded or replicated
+        - Activations depend on tensor/pipeline/sequence parallelism
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        tp_size = self.parallelism_config.tensor_parallel_size
+        pp_size = self.parallelism_config.pipeline_parallel_size
+        seq_parallel = self.parallelism_config.sequence_parallel
+
+        # 1. Model parameters (sharded by tensor parallelism)
+        # Each TP GPU holds 1/tp of the parameters
+        params_per_gpu = self.model_config.num_parameters / tp_size
+        model_params_gb = calculate_parameter_memory(
+            num_params=int(params_per_gpu),
+            dtype=self.training_config.dtype.value,
+        )
+
+        # 2. Gradients (sharded by tensor parallelism)
+        gradients_gb = calculate_gradient_memory(
+            num_params=int(params_per_gpu),
+            dtype=self.training_config.dtype.value,
+        )
+
+        # 3. Optimizer states
+        # In Megatron-LM, optimizer states are typically sharded similarly to parameters
+        # for tensor parallelism, but this can vary based on configuration
+        optimizer_gb = calculate_optimizer_memory(
+            num_params=int(params_per_gpu),
+            optimizer=self.training_config.optimizer.value,
+        )
+
+        # 4. Activations
+        # Activations are affected by:
+        # - Tensor parallelism: splits activations across TP GPUs
+        # - Pipeline parallelism: only holds activations for current stage
+        # - Sequence parallelism: splits sequence dimension
+        activations_gb = self._calculate_megatron_activations(
+            tp_size=tp_size,
+            pp_size=pp_size,
+            seq_parallel=seq_parallel,
+        )
+
+        # 5. Overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = calculate_overhead(base_memory)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown)
+
+    def _calculate_megatron_activations(
+        self,
+        tp_size: int,
+        pp_size: int,
+        seq_parallel: bool,
+    ) -> float:
+        """Calculate activation memory for Megatron-LM.
+
+        Megatron-LM activations are affected by parallelism strategy:
+        - Tensor parallelism: splits hidden dimension
+        - Pipeline parallelism: only current stage's activations
+        - Sequence parallelism: splits sequence dimension
+
+        Args:
+            tp_size: Tensor parallelism size
+            pp_size: Pipeline parallelism size
+            seq_parallel: Whether sequence parallelism is enabled
+
+        Returns:
+            Activation memory in GB
+        """
+
+        # Base activation memory
+        base_activations = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=tp_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Adjust for pipeline parallelism
+        # Each PP stage only holds num_layers / pp_size layers
+        pp_factor = 1.0 / pp_size
+
+        # Adjust for sequence parallelism
+        # If enabled, splits sequence dimension across TP GPUs
+        if seq_parallel and tp_size > 1:
+            seq_factor = 1.0 / tp_size
+        else:
+            seq_factor = 1.0
+
+        return base_activations * pp_factor * seq_factor
+
+
+class MegatronDeepSpeedEngine(BaseEngine):
+    """Megatron-LM + DeepSpeed combined engine.
+
+    This combines Megatron-LM's tensor/pipeline parallelism with
+    DeepSpeed ZeRO's optimizer/gradient sharding.
+    """
+
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory for Megatron-LM + DeepSpeed.
+
+        This uses:
+        - Megatron-LM for tensor/pipeline parallelism and activation memory
+        - DeepSpeed ZeRO for optimizer/gradient sharding
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        # Import DeepSpeed engine
+
+        # First calculate activation memory using Megatron-LM approach
+        tp_size = self.parallelism_config.tensor_parallel_size
+        pp_size = self.parallelism_config.pipeline_parallel_size
+        seq_parallel = self.parallelism_config.sequence_parallel
+
+        activations_gb = self._calculate_megatron_activations(
+            tp_size=tp_size,
+            pp_size=pp_size,
+            seq_parallel=seq_parallel,
+        )
+
+        # For parameters, gradients, optimizer - use DeepSpeed ZeRO logic
+        # But account for tensor parallelism (parameters are already split by TP)
+        tp_size = self.parallelism_config.tensor_parallel_size
+        params_per_gpu = self.model_config.num_parameters / tp_size
+
+        zero_stage = self.engine_config.zero_stage or 2
+        offload_optimizer = self.engine_config.offload_optimizer
+
+        # Model parameters (sharded by TP, then possibly by ZeRO)
+        if zero_stage >= 3:
+            # ZeRO-3 shards further
+            dp_size = self.parallelism_config.data_parallel_size
+            model_params_gb = gb_from_bytes((params_per_gpu * 2) / dp_size)
+        else:
+            # ZeRO-0/1/2 keeps parameters on each TP GPU
+            model_params_gb = gb_from_bytes(params_per_gpu * 2)
+
+        # Gradients
+        if zero_stage >= 2:
+            dp_size = self.parallelism_config.data_parallel_size
+            gradients_gb = gb_from_bytes((params_per_gpu * 2) / dp_size)
+        else:
+            gradients_gb = gb_from_bytes(params_per_gpu * 2)
+
+        # Optimizer states (12 bytes per param for Adam/AdamW in FP32)
+        if offload_optimizer.value == "cpu":
+            optimizer_gb = 0.0
+        else:
+            if zero_stage >= 1:
+                dp_size = self.parallelism_config.data_parallel_size
+                optimizer_gb = gb_from_bytes((params_per_gpu * 12) / dp_size)
+            else:
+                optimizer_gb = gb_from_bytes(params_per_gpu * 12)
+
+        # Overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = gb_from_bytes(base_memory * 0.2)
+
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown)
+
+    def _calculate_megatron_activations(
+        self,
+        tp_size: int,
+        pp_size: int,
+        seq_parallel: bool,
+    ) -> float:
+        """Calculate activation memory for Megatron-LM."""
+
+        # Base activation memory
+        base_activations = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=tp_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # Adjust for pipeline parallelism
+        pp_factor = 1.0 / pp_size
+
+        # Adjust for sequence parallelism
+        if seq_parallel and tp_size > 1:
+            seq_factor = 1.0 / tp_size
+        else:
+            seq_factor = 1.0
+
+        return base_activations * pp_factor * seq_factor

src/gpu_mem_calculator/engines/pytorch.py

@@ -0,0 +1,88 @@
+"""PyTorch DDP (Distributed Data Parallel) engine implementation.
+
+This is the baseline implementation without any memory optimizations.
+
+Reference: https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
+Reference: https://blog.eleuther.ai/transformer-math/
+"""
+
+from gpu_mem_calculator.core.formulas import (
+    calculate_activation_memory,
+    calculate_gradient_memory,
+    calculate_optimizer_memory,
+    calculate_overhead,
+    calculate_parameter_memory,
+)
+from gpu_mem_calculator.core.models import (
+    MemoryBreakdown,
+    MemoryResult,
+)
+from gpu_mem_calculator.engines.base import BaseEngine
+
+
+class PyTorchDDPEngine(BaseEngine):
+    """PyTorch DDP memory calculation.
+
+    DDP replicates the model on each GPU, so memory is not sharded.
+    Each GPU holds a full copy of the model, gradients, and optimizer states.
+    """
+
+    def calculate_memory(self) -> MemoryResult:
+        """Calculate memory requirements for PyTorch DDP training.
+
+        For DDP:
+        - Model parameters: Full model on each GPU
+        - Gradients: Full gradients on each GPU
+        - Optimizer states: Full optimizer states on each GPU (FP32)
+        - Activations: Batch size dependent, split by data parallel
+
+        Returns:
+            MemoryResult with complete memory breakdown
+        """
+        # 1. Model parameters (in the specified dtype)
+        model_params_gb = calculate_parameter_memory(
+            num_params=self.model_config.num_parameters,
+            dtype=self.training_config.dtype.value,
+        )
+
+        # 2. Gradients (same precision as parameters for mixed precision)
+        gradients_gb = calculate_gradient_memory(
+            num_params=self.model_config.num_parameters,
+            dtype=self.training_config.dtype.value,
+        )
+
+        # 3. Optimizer states (always FP32 for Adam/AdamW)
+        optimizer_gb = calculate_optimizer_memory(
+            num_params=self.model_config.num_parameters,
+            optimizer=self.training_config.optimizer.value,
+        )
+
+        # 4. Activations (depends on batch size and model architecture)
+        activations_gb = calculate_activation_memory(
+            batch_size=self.training_config.batch_size,
+            seq_len=self.model_config.max_seq_len,
+            hidden_size=self.model_config.hidden_size,
+            num_layers=self.model_config.num_layers,
+            num_attention_heads=self.model_config.num_attention_heads,
+            tensor_parallel_size=self.parallelism_config.tensor_parallel_size,
+            activation_checkpointing=self.training_config.activation_checkpointing,
+            moe_enabled=self.model_config.moe_enabled,
+            num_experts=self.model_config.num_experts,
+            top_k=self.model_config.top_k,
+            expert_intermediate_size=self.model_config.expert_intermediate_size,
+        )
+
+        # 5. Calculate overhead
+        base_memory = model_params_gb + gradients_gb + optimizer_gb + activations_gb
+        overhead_gb = calculate_overhead(base_memory)
+
+        # Create breakdown
+        breakdown = MemoryBreakdown(
+            model_params_gb=model_params_gb,
+            gradients_gb=gradients_gb,
+            optimizer_states_gb=optimizer_gb,
+            activations_gb=activations_gb,
+            overhead_gb=overhead_gb,
+        )
+
+        return self._create_result(breakdown)

src/gpu_mem_calculator/exporters/__init__.py

@@ -0,0 +1,14 @@
+"""Framework configuration exporters."""
+
+from gpu_mem_calculator.exporters.accelerate import AccelerateExporter
+from gpu_mem_calculator.exporters.axolotl import AxolotlExporter
+from gpu_mem_calculator.exporters.lightning import LightningExporter
+from gpu_mem_calculator.exporters.manager import ExportFormat, ExportManager
+
+__all__ = [
+    "ExportManager",
+    "ExportFormat",
+    "AccelerateExporter",
+    "LightningExporter",
+    "AxolotlExporter",
+]