memscale 0.1.0__py3-none-any.whl

memscale/__init__.py

@@ -0,0 +1,29 @@
+"""
+MemScale — Drop-in memory optimizer for PyTorch training.
+
+Usage:
+    import memscale
+
+    trainer = memscale.wrap(your_trainer)
+    # That's it. VRAM optimization is now active.
+"""
+
+from memscale.core.profiler import MemoryProfiler
+from memscale.core.decision_engine import DecisionEngine
+from memscale.core.executor import Executor
+from memscale.core.config import Config, OptimizationMode
+from memscale.api import wrap, optimize, detach
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "wrap",
+    "optimize",
+    "detach",
+    "Config",
+    "OptimizationMode",
+    "MemoryProfiler",
+    "DecisionEngine",
+    "Executor",
+    "__version__",
+]

memscale/api.py

@@ -0,0 +1,167 @@
+"""Public API for MemScale.
+
+This is the user-facing interface. Most users only need:
+
+    import memscale
+    trainer = memscale.wrap(your_trainer)
+
+Or for custom training loops:
+
+    with memscale.optimize(model, optimizer):
+        for batch in dataloader:
+            ...
+"""
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from memscale.core.config import Config
+from memscale.core.decision_engine import DecisionEngine
+from memscale.core.executor import Executor
+from memscale.core.profiler import MemoryProfiler
+from memscale.observability.logger import get_logger
+
+logger = logging.getLogger(__name__)
+
+
+def wrap(
+    target: Any,
+    config: Optional[Config] = None,
+    sample_input: Optional[torch.Tensor] = None,
+) -> Any:
+    """Wrap a trainer or model with MemScale optimizations.
+
+    Args:
+        target: Either a HuggingFace Trainer, Lightning Trainer, or PyTorch Module
+        config: Optional configuration. Defaults to balanced mode.
+        sample_input: Sample input tensor (improves profiling accuracy)
+
+    Returns:
+        The same trainer/model with optimization hooks attached.
+
+    Example:
+        >>> from transformers import Trainer
+        >>> import memscale
+        >>>
+        >>> trainer = Trainer(model=model, args=args, ...)
+        >>> trainer = memscale.wrap(trainer)
+        >>> trainer.train()  # VRAM optimized automatically
+    """
+    if config is None:
+        config = Config()
+
+    # Initialize observability
+    obs_logger = get_logger(config)
+    obs_logger.info("MemScale initialized", version="0.1.0", mode=config.mode.value)
+
+    # Detect what we're wrapping
+    if hasattr(target, "model") and hasattr(target, "train"):
+        # HuggingFace Trainer
+        from memscale.integrations.huggingface import wrap_hf_trainer
+
+        return wrap_hf_trainer(target, config, sample_input)
+
+    if isinstance(target, nn.Module):
+        # Raw PyTorch module — return wrapped model
+        return _wrap_module(target, config, sample_input)
+
+    raise TypeError(
+        f"Cannot wrap object of type {type(target).__name__}. "
+        f"Supported: HuggingFace Trainer, PyTorch Module, Lightning Trainer."
+    )
+
+
+def _wrap_module(
+    model: nn.Module,
+    config: Config,
+    sample_input: Optional[torch.Tensor],
+) -> nn.Module:
+    """Wrap a raw nn.Module with optimization hooks."""
+    profiler = MemoryProfiler(
+        prefer_static=config.use_static_profiling,
+        fallback_empirical=config.use_empirical_fallback,
+    )
+
+    # Profile
+    hw = profiler.detect_hardware()
+    logger.info(f"Detected hardware: {hw}")
+
+    graph = profiler.profile(model, sample_input)
+    logger.info(graph.summary())
+
+    # Decide
+    engine = DecisionEngine(config)
+    plan = engine.decide(graph, hw)
+    logger.info(plan.summary())
+
+    # Execute
+    executor = Executor(model, plan, config)
+    executor.attach()
+
+    # Attach executor reference so user can detach later
+    model._memscale_executor = executor
+    return model
+
+
+@contextmanager
+def optimize(
+    model: nn.Module,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    config: Optional[Config] = None,
+    sample_input: Optional[torch.Tensor] = None,
+):
+    """Context manager for MemScale optimization.
+
+    Use this for custom training loops where you don't have a Trainer object.
+
+    Example:
+        >>> import memscale
+        >>>
+        >>> with memscale.optimize(model, optimizer) as ms:
+        ...     for batch in dataloader:
+        ...         loss = model(batch).loss
+        ...         loss.backward()
+        ...         optimizer.step()
+        ...         optimizer.zero_grad()
+    """
+    if config is None:
+        config = Config()
+
+    profiler = MemoryProfiler(
+        prefer_static=config.use_static_profiling,
+        fallback_empirical=config.use_empirical_fallback,
+    )
+
+    hw = profiler.detect_hardware()
+    graph = profiler.profile(model, sample_input)
+    engine = DecisionEngine(config)
+    plan = engine.decide(graph, hw)
+
+    logger.info(plan.summary())
+
+    executor = Executor(model, plan, config)
+    executor.attach()
+
+    try:
+        yield executor
+    finally:
+        executor.detach()
+
+
+def detach(model: nn.Module) -> None:
+    """Remove MemScale optimization from a model.
+
+    Args:
+        model: Module previously wrapped with memscale.wrap()
+    """
+    if hasattr(model, "_memscale_executor"):
+        model._memscale_executor.detach()
+        delattr(model, "_memscale_executor")
+        logger.info("MemScale detached from model")
+    else:
+        logger.warning("Model was not wrapped with MemScale — nothing to detach")

memscale/core/__init__.py

@@ -0,0 +1,28 @@
+"""
+MemScale — Drop-in memory optimizer for PyTorch training.
+
+Usage:
+    import memscale
+
+    trainer = memscale.wrap(your_trainer)
+    # That's it. VRAM optimization is now active.
+"""
+
+from memscale.core.profiler import MemoryProfiler
+from memscale.core.decision_engine import DecisionEngine
+from memscale.core.executor import Executor
+from memscale.core.config import Config, OptimizationMode
+from memscale.api import wrap, optimize
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "wrap",
+    "optimize",
+    "Config",
+    "OptimizationMode",
+    "MemoryProfiler",
+    "DecisionEngine",
+    "Executor",
+    "__version__",
+]

memscale/core/config.py

@@ -0,0 +1,112 @@
+"""Configuration classes for MemScale.
+
+This module defines the user-facing configuration API and internal
+hardware budget representation.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+
+class OptimizationMode(str, Enum):
+    """Optimization aggressiveness mode.
+
+    - CONSERVATIVE: Only optimize when VRAM > 90% capacity. Safest, least overhead.
+    - BALANCED: Default. Optimize when VRAM > 75% capacity. Good speed-memory trade-off.
+    - AGGRESSIVE: Always optimize. Lowest VRAM usage, may have small overhead.
+    """
+
+    CONSERVATIVE = "conservative"
+    BALANCED = "balanced"
+    AGGRESSIVE = "aggressive"
+
+
+@dataclass
+class Config:
+    """User-facing configuration for MemScale.
+
+    Most users won't need to set anything — defaults work for 90% of cases.
+
+    Example:
+        config = Config(
+            mode=OptimizationMode.BALANCED,
+            enable_offloading=True,
+            max_cpu_offload_gb=64,
+        )
+    """
+
+    # Optimization aggressiveness
+    mode: OptimizationMode = OptimizationMode.BALANCED
+
+    # Which techniques to enable
+    enable_checkpointing: bool = True
+    enable_offloading: bool = True
+    enable_tiling: bool = False  # Disabled by default (more experimental)
+
+    # Resource limits
+    max_cpu_offload_gb: Optional[float] = None  # None = use available CPU RAM
+    target_gpu_utilization: float = 0.85  # Try to use this fraction of GPU memory
+
+    # Profiling
+    use_static_profiling: bool = True  # Try torch.fx first
+    use_empirical_fallback: bool = True  # Fallback to runtime profiling
+    warmup_steps: int = 2  # For empirical profiling
+
+    # Observability
+    enable_logging: bool = True
+    log_file: Optional[str] = None  # None = stdout only
+    observability_port: Optional[int] = None  # If set, expose Prometheus metrics
+
+    # Cost tracking
+    gpu_cost_per_hour: float = 2.50  # Default H100 spot rate (USD)
+
+    # Safety
+    verify_correctness: bool = False  # If True, runs baseline comparison (slow!)
+
+    def __post_init__(self) -> None:
+        if self.target_gpu_utilization <= 0 or self.target_gpu_utilization > 1.0:
+            raise ValueError(
+                f"target_gpu_utilization must be in (0, 1], got {self.target_gpu_utilization}"
+            )
+        if self.warmup_steps < 1:
+            raise ValueError(f"warmup_steps must be >= 1, got {self.warmup_steps}")
+
+
+@dataclass
+class HardwareBudget:
+    """Available hardware resources detected at runtime.
+
+    Populated by Profiler.detect_hardware() — users shouldn't construct this directly.
+    """
+
+    gpu_total_bytes: int
+    gpu_available_bytes: int
+    cpu_total_bytes: int
+    cpu_available_bytes: int
+    num_gpus: int
+    pcie_bandwidth_gbps: float = 16.0  # Reasonable default for PCIe 4.0 x16
+
+    @property
+    def gpu_total_gb(self) -> float:
+        return self.gpu_total_bytes / (1024**3)
+
+    @property
+    def gpu_available_gb(self) -> float:
+        return self.gpu_available_bytes / (1024**3)
+
+    @property
+    def cpu_total_gb(self) -> float:
+        return self.cpu_total_bytes / (1024**3)
+
+    @property
+    def cpu_available_gb(self) -> float:
+        return self.cpu_available_bytes / (1024**3)
+
+    def __str__(self) -> str:
+        return (
+            f"HardwareBudget(GPU: {self.gpu_available_gb:.1f}/{self.gpu_total_gb:.1f} GB, "
+            f"CPU: {self.cpu_available_gb:.1f}/{self.cpu_total_gb:.1f} GB, "
+            f"GPUs: {self.num_gpus})"
+        )

memscale/core/decision_engine.py

@@ -0,0 +1,277 @@
+"""Decision engine: decides which optimization technique to apply per layer.
+
+Given a memory graph and hardware budget, produce an optimization plan.
+
+Current implementation is rule-based heuristic. Future versions will use
+learned policies trained on customer data.
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, List
+
+from memscale.core.config import Config, HardwareBudget, OptimizationMode
+from memscale.core.memory_graph import LayerMemoryInfo, MemoryGraph
+
+logger = logging.getLogger(__name__)
+
+
+class Technique(str, Enum):
+    """Optimization techniques that can be applied per layer."""
+
+    KEEP = "keep"  # No optimization, keep on GPU
+    CHECKPOINT = "checkpoint"  # Activation checkpointing (recompute on backward)
+    OFFLOAD = "offload"  # CPU offloading
+    TILE = "tile"  # Activation tiling
+
+
+@dataclass
+class OptimizationPlan:
+    """Plan for which technique to apply to each layer.
+
+    Produced by DecisionEngine.decide() and consumed by Executor.
+    """
+
+    layer_techniques: Dict[str, Technique] = field(default_factory=dict)
+    estimated_memory_savings_bytes: int = 0
+    estimated_throughput_overhead_pct: float = 0.0
+    rationale: List[str] = field(default_factory=list)
+
+    def technique_for(self, layer_name: str) -> Technique:
+        return self.layer_techniques.get(layer_name, Technique.KEEP)
+
+    def layers_with_technique(self, technique: Technique) -> List[str]:
+        return [name for name, t in self.layer_techniques.items() if t == technique]
+
+    def summary(self) -> str:
+        gb = 1024**3
+        counts = {t: 0 for t in Technique}
+        for technique in self.layer_techniques.values():
+            counts[technique] += 1
+
+        lines = [
+            "OptimizationPlan:",
+            f"  Estimated VRAM savings: {self.estimated_memory_savings_bytes / gb:.2f} GB",
+            f"  Estimated throughput overhead: {self.estimated_throughput_overhead_pct:.1f}%",
+            f"  Layers per technique:",
+            f"    KEEP: {counts[Technique.KEEP]}",
+            f"    CHECKPOINT: {counts[Technique.CHECKPOINT]}",
+            f"    OFFLOAD: {counts[Technique.OFFLOAD]}",
+            f"    TILE: {counts[Technique.TILE]}",
+        ]
+        if self.rationale:
+            lines.append("  Rationale:")
+            for note in self.rationale[:5]:
+                lines.append(f"    - {note}")
+        return "\n".join(lines)
+
+
+class DecisionEngine:
+    """Decide optimization techniques for each layer based on rules.
+
+    Algorithm:
+    1. Compute total memory needed for entire model
+    2. If it fits in GPU comfortably (< target_utilization × GPU memory), no-op
+    3. Otherwise, sort layers by memory footprint, apply techniques to biggest first:
+       - High compute/memory ratio: CHECKPOINT (recompute is cheap)
+       - Low compute/memory ratio with CPU room: OFFLOAD
+       - Otherwise: TILE
+    4. Stop when projected memory fits target budget
+    """
+
+    # Heuristic thresholds (tunable based on customer data)
+    RECOMPUTE_THRESHOLD = 1e6  # FLOPs/byte ratio above which checkpoint is preferred
+    MIN_LAYER_SIZE_FOR_OPTIMIZATION = 4 * 1024 * 1024  # 4 MB minimum
+
+    def __init__(self, config: Config):
+        self.config = config
+
+    def decide(self, graph: MemoryGraph, hw: HardwareBudget) -> OptimizationPlan:
+        """Produce an optimization plan for the given memory graph and hardware."""
+        plan = OptimizationPlan()
+
+        # Determine optimization mode threshold (used both for early-exit and loop)
+        mode_thresholds = {
+            OptimizationMode.CONSERVATIVE: 0.90,
+            OptimizationMode.BALANCED: 0.75,
+            OptimizationMode.AGGRESSIVE: 0.50,
+        }
+        threshold_pct = mode_thresholds[self.config.mode]
+
+        # Calculate target VRAM budget
+        # If no GPU detected, treat target as a fraction of model memory itself
+        # (this lets users force optimization for testing/CPU scenarios)
+        total_memory_needed = graph.total_memory_bytes()
+
+        if hw.gpu_total_bytes > 0:
+            target_gpu_bytes = int(hw.gpu_total_bytes * self.config.target_gpu_utilization)
+        else:
+            # CPU-only or no GPU: use total memory needed × user's target as the budget
+            # This makes target_gpu_utilization a "how much of model memory should fit" hint
+            target_gpu_bytes = int(total_memory_needed * self.config.target_gpu_utilization)
+
+        logger.info(
+            f"Memory analysis: need {total_memory_needed / (1024**3):.3f} GB, "
+            f"target budget {target_gpu_bytes / (1024**3):.3f} GB, "
+            f"mode={self.config.mode.value}"
+        )
+
+        # Quick exit: if model fits comfortably AND user hasn't asked for aggressive opt
+        # We only skip if BOTH conditions are met:
+        # 1. Model fits well within target
+        # 2. User is in conservative or balanced mode (not explicitly requesting aggressive)
+        comfortable_fit = total_memory_needed < target_gpu_bytes * 0.7
+        is_aggressive_request = self.config.mode == OptimizationMode.AGGRESSIVE
+
+        if comfortable_fit and not is_aggressive_request:
+            logger.info("Model fits comfortably. No optimization needed.")
+            for name in graph.layers:
+                plan.layer_techniques[name] = Technique.KEEP
+            plan.rationale.append("Model fits in GPU; no optimization applied.")
+            return plan
+
+        # Initial plan: keep all
+        for name in graph.layers:
+            plan.layer_techniques[name] = Technique.KEEP
+
+        # Sort layers by total memory footprint (largest first)
+        sorted_layers = graph.layers_sorted_by_memory()
+
+        # Track running memory budget
+        projected_memory = total_memory_needed
+        cpu_offload_budget = self._compute_cpu_offload_budget(hw)
+        cpu_used = 0
+
+        # Mode-aware minimum layer size: aggressive mode optimizes smaller layers
+        if self.config.mode == OptimizationMode.AGGRESSIVE:
+            min_layer_size = self.MIN_LAYER_SIZE_FOR_OPTIMIZATION // 16  # 256 KB
+        elif self.config.mode == OptimizationMode.BALANCED:
+            min_layer_size = self.MIN_LAYER_SIZE_FOR_OPTIMIZATION // 4  # 1 MB
+        else:
+            min_layer_size = self.MIN_LAYER_SIZE_FOR_OPTIMIZATION  # 4 MB
+
+        for layer in sorted_layers:
+            # Stop optimizing once we're under target
+            if projected_memory < target_gpu_bytes * threshold_pct:
+                break
+
+            # Skip tiny layers — overhead not worth it (mode-aware threshold)
+            if layer.total_memory_bytes < min_layer_size:
+                continue
+
+            technique = self._choose_technique_for_layer(
+                layer, cpu_offload_budget - cpu_used
+            )
+
+            if technique == Technique.KEEP:
+                continue
+
+            plan.layer_techniques[layer.name] = technique
+            savings = self._estimate_savings(layer, technique)
+            projected_memory -= savings
+            plan.estimated_memory_savings_bytes += savings
+
+            if technique == Technique.OFFLOAD:
+                cpu_used += layer.activation_bytes
+
+            plan.rationale.append(
+                f"{layer.name}: {technique.value} "
+                f"(saves {savings / (1024**2):.1f} MB)"
+            )
+
+        # Estimate throughput overhead
+        plan.estimated_throughput_overhead_pct = self._estimate_overhead(plan, graph)
+
+        return plan
+
+    def _choose_technique_for_layer(
+        self, layer: LayerMemoryInfo, cpu_remaining: int
+    ) -> Technique:
+        """Pick the best technique for a single layer based on its characteristics."""
+
+        # Layers without activations don't benefit from checkpointing/offloading
+        # Mode-aware minimum activation size threshold
+        if self.config.mode == OptimizationMode.AGGRESSIVE:
+            min_activation = 64 * 1024  # 64 KB
+        elif self.config.mode == OptimizationMode.BALANCED:
+            min_activation = 512 * 1024  # 512 KB
+        else:
+            min_activation = 1024 * 1024  # 1 MB
+
+        if layer.activation_bytes < min_activation:
+            return Technique.KEEP
+
+        # If checkpointing enabled and layer is compute-bound, prefer checkpoint
+        # (cheap to recompute relative to activation memory saved)
+        if (
+            self.config.enable_checkpointing
+            and layer.compute_to_memory_ratio > self.RECOMPUTE_THRESHOLD
+        ):
+            return Technique.CHECKPOINT
+
+        # If offloading enabled and we have CPU room, offload
+        if (
+            self.config.enable_offloading
+            and cpu_remaining > layer.activation_bytes * 1.2  # 20% safety margin
+        ):
+            return Technique.OFFLOAD
+
+        # If tiling enabled, use it as fallback
+        if self.config.enable_tiling and layer.activation_bytes > 16 * 1024 * 1024:
+            return Technique.TILE
+
+        # If checkpointing is enabled but we hit here, still use it as last resort
+        if self.config.enable_checkpointing:
+            return Technique.CHECKPOINT
+
+        return Technique.KEEP
+
+    def _compute_cpu_offload_budget(self, hw: HardwareBudget) -> int:
+        """How much CPU RAM can we use for offloading?"""
+        if self.config.max_cpu_offload_gb is not None:
+            return int(self.config.max_cpu_offload_gb * 1024**3)
+        # Use 50% of available CPU RAM by default (leave room for OS, data loader)
+        return int(hw.cpu_available_bytes * 0.5)
+
+    @staticmethod
+    def _estimate_savings(layer: LayerMemoryInfo, technique: Technique) -> int:
+        """How much GPU memory does this technique save for this layer?"""
+        if technique == Technique.CHECKPOINT:
+            # Save activation memory (will recompute on backward)
+            return layer.activation_bytes
+        elif technique == Technique.OFFLOAD:
+            # Save activation + can also offload params during inactive periods
+            return layer.activation_bytes
+        elif technique == Technique.TILE:
+            # Tiling saves ~50-75% of activation peak
+            return int(layer.activation_bytes * 0.6)
+        else:
+            return 0
+
+    @staticmethod
+    def _estimate_overhead(plan: OptimizationPlan, graph: MemoryGraph) -> float:
+        """Estimate throughput overhead in percent."""
+        # Rule of thumb (will be calibrated with real benchmarks):
+        # - CHECKPOINT: ~25% extra compute on the layer (1 extra forward)
+        # - OFFLOAD: ~5% if PCIe transfer overlaps; ~30% if it doesn't
+        # - TILE: ~10% from sequential processing
+        weighted_overhead = 0.0
+        total_flops = sum(layer.flops for layer in graph.layers.values()) or 1
+
+        for layer_name, technique in plan.layer_techniques.items():
+            layer = graph.layers.get(layer_name)
+            if not layer or layer.flops == 0:
+                continue
+
+            layer_weight = layer.flops / total_flops
+
+            if technique == Technique.CHECKPOINT:
+                weighted_overhead += 25.0 * layer_weight
+            elif technique == Technique.OFFLOAD:
+                weighted_overhead += 5.0 * layer_weight
+            elif technique == Technique.TILE:
+                weighted_overhead += 10.0 * layer_weight
+
+        return min(weighted_overhead, 100.0)