pymemora 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
pymemora-0.1.0/LICENSE ADDED
@@ -0,0 +1,18 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ Copyright 2026 Memora Authors
6
+
7
+ Licensed under the Apache License, Version 2.0 (the "License");
8
+ you may not use this file except in compliance with the License.
9
+ You may obtain a copy of the License at
10
+
11
+ http://www.apache.org/licenses/LICENSE-2.0
12
+
13
+ Unless required by applicable law or agreed to in writing, software
14
+ distributed under the License is distributed on an "AS IS" BASIS,
15
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ See the License for the specific language governing permissions and
17
+ limitations under the License.
18
+
@@ -0,0 +1,179 @@
1
+ Metadata-Version: 2.4
2
+ Name: pymemora
3
+ Version: 0.1.0
4
+ Summary: A memory-aware training optimization framework for PyTorch across any OS & hardware.
5
+ Author: Memora Authors
6
+ License: Apache-2.0
7
+ Classifier: Development Status :: 3 - Alpha
8
+ Classifier: Intended Audience :: Developers
9
+ Classifier: Intended Audience :: Science/Research
10
+ Classifier: License :: OSI Approved :: Apache Software License
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
17
+ Classifier: Operating System :: OS Independent
18
+ Requires-Python: >=3.9
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: torch>=2.0.0
22
+ Requires-Dist: psutil>=5.9.0
23
+ Requires-Dist: pydantic>=2.0.0
24
+ Requires-Dist: rich>=13.0.0
25
+ Requires-Dist: typer>=0.9.0
26
+ Provides-Extra: dev
27
+ Requires-Dist: torchvision>=0.15.0; extra == "dev"
28
+ Requires-Dist: ruff>=0.1.0; extra == "dev"
29
+ Provides-Extra: gpu
30
+ Requires-Dist: nvidia-ml-py>=12.535.0; extra == "gpu"
31
+ Dynamic: license-file
32
+
33
+ # Memora 🧠⚡
34
+
35
+ [![PyPI version](https://img.shields.io/pypi/v/pymemora.svg?color=blue)](https://pypi.org/project/pymemora/)
36
+ [![Python Version](https://img.shields.io/pypi/pyversions/pymemora.svg)](https://pypi.org/project/pymemora/)
37
+ [![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](https://opensource.org/licenses/Apache-2.0)
38
+
39
+ > **A memory-aware training optimization framework for PyTorch across any OS & hardware.**
40
+
41
+ Memora helps PyTorch developers use limited GPU memory (and CPU resources) more effectively during model training. It inspects hardware, benchmarks memory requirements safely, and automatically configures batch sizes, mixed precision, gradient accumulation, and activation checkpointing through a single unified developer API.
42
+
43
+ ---
44
+
45
+ ## 🚀 Key Features
46
+
47
+ - **🌐 Works on Any OS & System**: Seamlessly runs on Linux, macOS (Apple Silicon MPS), and Windows. Detects CUDA GPUs (e.g. RTX 3050, T4, A100), Apple MPS, and CPU with graceful fallbacks.
48
+ - **🔍 Safe Micro-Batch Search**: Probes candidate batch sizes with isolated trial runs, catching CUDA Out-Of-Memory (OOM) errors and enforcing configurable memory safety margins without crashing your workflow.
49
+ - **⚡ Adaptive Mixed Precision (AMP)**: Automatically selects BF16 on modern Ampere+ GPUs (Compute Capability >= 8.0) for zero-scaler speed, FP16 with `torch.amp.GradScaler` on Volta/Turing/Pascal GPUs, or FP32 on CPU.
50
+ - **⚖️ Gradient Accumulation**: Preserves a target effective batch size regardless of hardware VRAM constraints by automatically computing micro-batches and accumulation steps with normalized loss gradients.
51
+ - **💾 Activation Checkpointing**: Automatically advises or wraps layers with `torch.utils.checkpoint` to conserve 40–60% activation memory when training under constrained VRAM.
52
+ - **📊 Transparent Profiling & Reporting**: Generates structured JSON and Markdown execution reports comparing baseline PyTorch vs. optimized configurations (peak memory, step times, throughput, accuracy).
53
+ - **🖥️ Rich CLI**: Instant hardware discovery and report viewer via `memora hardware` and `memora report`.
54
+
55
+ ---
56
+
57
+ ## 📦 Installation
58
+
59
+ Install the official package from **PyPI**:
60
+
61
+ ```bash
62
+ pip install pymemora
63
+ ```
64
+
65
+ Or install from source:
66
+
67
+ ```bash
68
+ git clone https://github.com/Premkumar499/memora.git
69
+ cd memora
70
+ pip install -e .
71
+ ```
72
+
73
+ ---
74
+
75
+ ## 🛠️ Quickstart
76
+
77
+ ```python
78
+ import torch
79
+ import torch.nn as nn
80
+ from torch.utils.data import TensorDataset
81
+ from memora import Memora
82
+
83
+ # 1. Define your PyTorch model and dataset
84
+ model = nn.Sequential(
85
+ nn.Linear(64, 128),
86
+ nn.ReLU(),
87
+ nn.Linear(128, 10),
88
+ )
89
+
90
+ dataset = TensorDataset(
91
+ torch.randn(256, 64),
92
+ torch.randint(0, 10, (256,)),
93
+ )
94
+
95
+ # 2. Instantiate Memora with automatic memory optimization
96
+ trainer = Memora(
97
+ model=model,
98
+ dataset=dataset,
99
+ batch_size="auto", # Automatically find optimal batch size
100
+ precision="auto", # Automatically select BF16/FP16/FP32
101
+ gradient_accumulation="auto", # Compute accumulation steps
102
+ activation_checkpointing="auto", # Enable checkpointing if VRAM is tight
103
+ device="auto", # 'cuda', 'mps', or 'cpu'
104
+ )
105
+
106
+ # 3. Preview recommendations (optional)
107
+ trainer.optimize()
108
+
109
+ # 4. Train model and collect performance report
110
+ report = trainer.fit(epochs=3)
111
+
112
+ # 5. Display execution report
113
+ report.print_summary()
114
+ ```
115
+
116
+ ---
117
+
118
+ ## 💻 Command Line Interface (CLI)
119
+
120
+ Memora includes a built-in CLI powered by Typer and Rich:
121
+
122
+ ### 1. Hardware Discovery
123
+ Inspect available compute devices, memory limits, and compute capabilities:
124
+ ```bash
125
+ memora hardware
126
+ ```
127
+
128
+ ### 2. View Training Run Reports
129
+ Pretty-print detailed memory and throughput metrics from a saved run JSON:
130
+ ```bash
131
+ memora report <path-to-report.json>
132
+ ```
133
+
134
+ ### 3. Version Check
135
+ ```bash
136
+ memora version
137
+ ```
138
+
139
+ ---
140
+
141
+ ## 📁 Repository Structure
142
+
143
+ ```text
144
+ framework/
145
+ ├── memora/
146
+ │ ├── __init__.py # Public exports
147
+ │ ├── config.py # Pydantic schemas & configuration
148
+ │ ├── engine.py # Core Memora trainer API
149
+ │ ├── cli.py # Rich Typer CLI
150
+ │ ├── hardware/
151
+ │ │ ├── gpu_detector.py # CUDA, MPS, CPU & NVML telemetry
152
+ │ │ └── memory_monitor.py # Peak memory tracking & snapshots
153
+ │ ├── optimizer/
154
+ │ │ ├── batch_optimizer.py # Exponential & binary batch search
155
+ │ │ ├── precision_optimizer.py # Autocast & GradScaler manager
156
+ │ │ ├── gradient_optimizer.py # Effective batch accumulation
157
+ │ │ └── checkpoint_optimizer.py # Activation checkpointing
158
+ │ ├── profiler/
159
+ │ │ └── training_profiler.py # Step timings & throughput
160
+ │ ├── reporting/
161
+ │ │ └── report.py # RunReport & baseline comparison
162
+ │ └── utils/
163
+ │ └── common.py # Formatting, device resolution & seeds
164
+ ├── scripts/
165
+ │ └── deploy_check.py # Deployment health check smoke test
166
+ ├── docs/
167
+ │ └── INTEGRATION_GUIDE.md # Cross-computer integration guide
168
+ ├── pyproject.toml # PEP 517/621 packaging metadata
169
+ ├── README.md # Project documentation
170
+ ├── LICENSE # Apache 2.0 license
171
+ ├── idea.md # Concept & architectural motivation
172
+ └── plan.md # Development milestones & roadmap
173
+ ```
174
+
175
+ ---
176
+
177
+ ## 📄 License
178
+
179
+ Apache License 2.0. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,147 @@
1
+ # Memora 🧠⚡
2
+
3
+ [![PyPI version](https://img.shields.io/pypi/v/pymemora.svg?color=blue)](https://pypi.org/project/pymemora/)
4
+ [![Python Version](https://img.shields.io/pypi/pyversions/pymemora.svg)](https://pypi.org/project/pymemora/)
5
+ [![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](https://opensource.org/licenses/Apache-2.0)
6
+
7
+ > **A memory-aware training optimization framework for PyTorch across any OS & hardware.**
8
+
9
+ Memora helps PyTorch developers use limited GPU memory (and CPU resources) more effectively during model training. It inspects hardware, benchmarks memory requirements safely, and automatically configures batch sizes, mixed precision, gradient accumulation, and activation checkpointing through a single unified developer API.
10
+
11
+ ---
12
+
13
+ ## 🚀 Key Features
14
+
15
+ - **🌐 Works on Any OS & System**: Seamlessly runs on Linux, macOS (Apple Silicon MPS), and Windows. Detects CUDA GPUs (e.g. RTX 3050, T4, A100), Apple MPS, and CPU with graceful fallbacks.
16
+ - **🔍 Safe Micro-Batch Search**: Probes candidate batch sizes with isolated trial runs, catching CUDA Out-Of-Memory (OOM) errors and enforcing configurable memory safety margins without crashing your workflow.
17
+ - **⚡ Adaptive Mixed Precision (AMP)**: Automatically selects BF16 on modern Ampere+ GPUs (Compute Capability >= 8.0) for zero-scaler speed, FP16 with `torch.amp.GradScaler` on Volta/Turing/Pascal GPUs, or FP32 on CPU.
18
+ - **⚖️ Gradient Accumulation**: Preserves a target effective batch size regardless of hardware VRAM constraints by automatically computing micro-batches and accumulation steps with normalized loss gradients.
19
+ - **💾 Activation Checkpointing**: Automatically advises or wraps layers with `torch.utils.checkpoint` to conserve 40–60% activation memory when training under constrained VRAM.
20
+ - **📊 Transparent Profiling & Reporting**: Generates structured JSON and Markdown execution reports comparing baseline PyTorch vs. optimized configurations (peak memory, step times, throughput, accuracy).
21
+ - **🖥️ Rich CLI**: Instant hardware discovery and report viewer via `memora hardware` and `memora report`.
22
+
23
+ ---
24
+
25
+ ## 📦 Installation
26
+
27
+ Install the official package from **PyPI**:
28
+
29
+ ```bash
30
+ pip install pymemora
31
+ ```
32
+
33
+ Or install from source:
34
+
35
+ ```bash
36
+ git clone https://github.com/Premkumar499/memora.git
37
+ cd memora
38
+ pip install -e .
39
+ ```
40
+
41
+ ---
42
+
43
+ ## 🛠️ Quickstart
44
+
45
+ ```python
46
+ import torch
47
+ import torch.nn as nn
48
+ from torch.utils.data import TensorDataset
49
+ from memora import Memora
50
+
51
+ # 1. Define your PyTorch model and dataset
52
+ model = nn.Sequential(
53
+ nn.Linear(64, 128),
54
+ nn.ReLU(),
55
+ nn.Linear(128, 10),
56
+ )
57
+
58
+ dataset = TensorDataset(
59
+ torch.randn(256, 64),
60
+ torch.randint(0, 10, (256,)),
61
+ )
62
+
63
+ # 2. Instantiate Memora with automatic memory optimization
64
+ trainer = Memora(
65
+ model=model,
66
+ dataset=dataset,
67
+ batch_size="auto", # Automatically find optimal batch size
68
+ precision="auto", # Automatically select BF16/FP16/FP32
69
+ gradient_accumulation="auto", # Compute accumulation steps
70
+ activation_checkpointing="auto", # Enable checkpointing if VRAM is tight
71
+ device="auto", # 'cuda', 'mps', or 'cpu'
72
+ )
73
+
74
+ # 3. Preview recommendations (optional)
75
+ trainer.optimize()
76
+
77
+ # 4. Train model and collect performance report
78
+ report = trainer.fit(epochs=3)
79
+
80
+ # 5. Display execution report
81
+ report.print_summary()
82
+ ```
83
+
84
+ ---
85
+
86
+ ## 💻 Command Line Interface (CLI)
87
+
88
+ Memora includes a built-in CLI powered by Typer and Rich:
89
+
90
+ ### 1. Hardware Discovery
91
+ Inspect available compute devices, memory limits, and compute capabilities:
92
+ ```bash
93
+ memora hardware
94
+ ```
95
+
96
+ ### 2. View Training Run Reports
97
+ Pretty-print detailed memory and throughput metrics from a saved run JSON:
98
+ ```bash
99
+ memora report <path-to-report.json>
100
+ ```
101
+
102
+ ### 3. Version Check
103
+ ```bash
104
+ memora version
105
+ ```
106
+
107
+ ---
108
+
109
+ ## 📁 Repository Structure
110
+
111
+ ```text
112
+ framework/
113
+ ├── memora/
114
+ │ ├── __init__.py # Public exports
115
+ │ ├── config.py # Pydantic schemas & configuration
116
+ │ ├── engine.py # Core Memora trainer API
117
+ │ ├── cli.py # Rich Typer CLI
118
+ │ ├── hardware/
119
+ │ │ ├── gpu_detector.py # CUDA, MPS, CPU & NVML telemetry
120
+ │ │ └── memory_monitor.py # Peak memory tracking & snapshots
121
+ │ ├── optimizer/
122
+ │ │ ├── batch_optimizer.py # Exponential & binary batch search
123
+ │ │ ├── precision_optimizer.py # Autocast & GradScaler manager
124
+ │ │ ├── gradient_optimizer.py # Effective batch accumulation
125
+ │ │ └── checkpoint_optimizer.py # Activation checkpointing
126
+ │ ├── profiler/
127
+ │ │ └── training_profiler.py # Step timings & throughput
128
+ │ ├── reporting/
129
+ │ │ └── report.py # RunReport & baseline comparison
130
+ │ └── utils/
131
+ │ └── common.py # Formatting, device resolution & seeds
132
+ ├── scripts/
133
+ │ └── deploy_check.py # Deployment health check smoke test
134
+ ├── docs/
135
+ │ └── INTEGRATION_GUIDE.md # Cross-computer integration guide
136
+ ├── pyproject.toml # PEP 517/621 packaging metadata
137
+ ├── README.md # Project documentation
138
+ ├── LICENSE # Apache 2.0 license
139
+ ├── idea.md # Concept & architectural motivation
140
+ └── plan.md # Development milestones & roadmap
141
+ ```
142
+
143
+ ---
144
+
145
+ ## 📄 License
146
+
147
+ Apache License 2.0. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,24 @@
1
+ """Memora: A memory-aware training optimization framework for PyTorch."""
2
+
3
+ from memora.config import MemoraConfig, OptimizationRecommendation
4
+ from memora.engine import Memora
5
+ from memora.hardware.gpu_detector import DeviceInfo, GPUDetector, HardwareReport
6
+ from memora.hardware.memory_monitor import MemoryMonitor, MemorySnapshot, MemoryTracker
7
+ from memora.reporting.report import RunReport, compare_runs
8
+
9
+ __version__ = "0.1.0"
10
+
11
+ __all__ = [
12
+ "Memora",
13
+ "MemoraConfig",
14
+ "OptimizationRecommendation",
15
+ "GPUDetector",
16
+ "DeviceInfo",
17
+ "HardwareReport",
18
+ "MemoryMonitor",
19
+ "MemorySnapshot",
20
+ "MemoryTracker",
21
+ "RunReport",
22
+ "compare_runs",
23
+ ]
24
+
@@ -0,0 +1,104 @@
1
+ """Command line interface for Memora using Typer and Rich."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import typer
7
+ from rich.console import Console
8
+ from rich.panel import Panel
9
+ from rich.table import Table
10
+
11
+ from memora import __version__
12
+ from memora.hardware.gpu_detector import GPUDetector
13
+ from memora.utils.common import format_bytes
14
+
15
+ app = typer.Typer(help="Memora: PyTorch Memory-Aware Optimization CLI")
16
+ console = Console()
17
+
18
+
19
+ @app.command()
20
+ def hardware() -> None:
21
+ """Detect and display system hardware, CPU RAM, and GPU memory."""
22
+ report = GPUDetector.get_hardware_report()
23
+
24
+ console.print(Panel.fit(
25
+ f"[bold cyan]Memora Hardware Discovery[/bold cyan]\n"
26
+ f"OS: {report.platform_system} ({report.platform_release})\n"
27
+ f"Python: {report.python_version} | PyTorch: {report.torch_version}\n"
28
+ f"CUDA Available: [{'green' if report.cuda_available else 'yellow'}]{report.cuda_available}[/]"
29
+ + (f" (v{report.cuda_version})" if report.cuda_version else "")
30
+ + f"\nSystem RAM: {report.to_dict()['system_ram_available_str']} free / {report.to_dict()['system_ram_total_str']} total",
31
+ title="System Environment",
32
+ ))
33
+
34
+ table = Table(title="Detected Compute Devices", header_style="bold magenta")
35
+ table.add_column("Index", style="cyan")
36
+ table.add_column("Type", style="bold")
37
+ table.add_column("Device Name", style="white")
38
+ table.add_column("Total Memory", style="green")
39
+ table.add_column("Compute Cap.", style="yellow")
40
+ table.add_column("FP16 / BF16", style="blue")
41
+
42
+ for dev in report.devices:
43
+ half_support = f"FP16: {dev.supports_fp16} | BF16: {dev.supports_bf16}"
44
+ table.add_row(
45
+ str(dev.device_index),
46
+ dev.device_type.upper(),
47
+ dev.name,
48
+ format_bytes(dev.total_memory),
49
+ dev.compute_capability or "N/A",
50
+ half_support,
51
+ )
52
+
53
+ console.print(table)
54
+
55
+
56
+ @app.command()
57
+ def report(file_path: str = typer.Argument(..., help="Path to Memora report JSON file")) -> None:
58
+ """Read and format a Memora training run report."""
59
+ try:
60
+ with open(file_path, "r", encoding="utf-8") as f:
61
+ data = json.load(f)
62
+
63
+ console.print(Panel(
64
+ f"[bold]Run ID:[/bold] {data.get('run_id')}\n"
65
+ f"[bold]Timestamp:[/bold] {data.get('timestamp')}\n"
66
+ f"[bold]Device:[/bold] {data.get('hardware_info', {}).get('default_device', {}).get('name')}",
67
+ title="[bold green]Report Summary[/bold green]",
68
+ ))
69
+
70
+ opts = data.get("optimization_decisions", {})
71
+ opt_table = Table(title="Optimization Decisions", header_style="bold cyan")
72
+ opt_table.add_column("Parameter", style="bold")
73
+ opt_table.add_column("Value", style="green")
74
+
75
+ for k, v in opts.items():
76
+ if k != "reasons":
77
+ opt_table.add_row(k, str(v))
78
+ console.print(opt_table)
79
+
80
+ prof = data.get("profiler_summary", {})
81
+ if prof:
82
+ prof_table = Table(title="Performance Metrics", header_style="bold yellow")
83
+ prof_table.add_column("Metric", style="bold")
84
+ prof_table.add_column("Value", style="cyan")
85
+ for k, v in prof.items():
86
+ prof_table.add_row(k, str(v))
87
+ console.print(prof_table)
88
+
89
+ except Exception as e:
90
+ console.print(f"[bold red]Error loading report:[/bold red] {e}")
91
+ raise typer.Exit(code=1)
92
+
93
+
94
+ @app.command()
95
+ def version() -> None:
96
+ """Show Memora version."""
97
+ import torch
98
+ console.print(f"Memora version: [bold green]{__version__}[/bold green]")
99
+ console.print(f"PyTorch version: [bold cyan]{torch.__version__}[/bold cyan]")
100
+
101
+
102
+ if __name__ == "__main__":
103
+ app()
104
+
@@ -0,0 +1,126 @@
1
+ """Configuration models and schemas for Memora."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any, Dict, Literal, Optional, Union
6
+ from pydantic import BaseModel, Field, field_validator
7
+
8
+
9
+ PrecisionType = Literal["auto", "fp32", "fp16", "bf16"]
10
+
11
+
12
+ class MemoraConfig(BaseModel):
13
+ """Configuration for Memora optimization and training run."""
14
+
15
+ # Device & resource management
16
+ device: str = Field(
17
+ default="auto",
18
+ description="Target compute device: 'auto', 'cuda', 'cuda:0', 'mps', or 'cpu'.",
19
+ )
20
+ gpu_memory_limit: Union[float, int, Literal["auto"]] = Field(
21
+ default="auto",
22
+ description="VRAM memory limit. If float <= 1.0, fraction of total VRAM. If int > 1, bytes limit. 'auto' uses default safety headroom.",
23
+ )
24
+ safety_margin: float = Field(
25
+ default=0.15,
26
+ ge=0.01,
27
+ le=0.50,
28
+ description="Fraction of memory reserved as a safety buffer against unexpected memory spikes (default: 15%).",
29
+ )
30
+
31
+ # Batch search parameters
32
+ batch_size: Union[int, Literal["auto"]] = Field(
33
+ default="auto",
34
+ description="Micro-batch size for training. 'auto' benchmarks and selects optimal size.",
35
+ )
36
+ min_batch_size: int = Field(
37
+ default=1,
38
+ ge=1,
39
+ description="Minimum batch size allowed during search.",
40
+ )
41
+ max_batch_size: int = Field(
42
+ default=512,
43
+ ge=1,
44
+ description="Maximum batch size allowed during search.",
45
+ )
46
+
47
+ # Gradient accumulation parameters
48
+ effective_batch_size: Optional[int] = Field(
49
+ default=None,
50
+ ge=1,
51
+ description="Target effective batch size. When provided with small batch_size, gradient accumulation is automatically configured.",
52
+ )
53
+ gradient_accumulation_steps: Union[int, Literal["auto"]] = Field(
54
+ default="auto",
55
+ description="Gradient accumulation steps. 'auto' computes steps to match effective_batch_size or 1.",
56
+ )
57
+
58
+ # Precision
59
+ precision: PrecisionType = Field(
60
+ default="auto",
61
+ description="Numerical precision: 'auto', 'fp32', 'fp16', or 'bf16'.",
62
+ )
63
+
64
+ # Activation checkpointing
65
+ activation_checkpointing: Union[bool, Literal["auto"]] = Field(
66
+ default="auto",
67
+ description="Activation checkpointing toggle: 'auto' enables it when VRAM is tight, or explicit boolean.",
68
+ )
69
+
70
+ # Profiling & reporting
71
+ enable_profiling: bool = Field(
72
+ default=True,
73
+ description="Whether to monitor and profile memory and step timings.",
74
+ )
75
+ report_output_dir: Optional[str] = Field(
76
+ default="./memora_reports",
77
+ description="Directory to persist optimization reports.",
78
+ )
79
+ seed: int = Field(
80
+ default=42,
81
+ description="Random seed for deterministic runs.",
82
+ )
83
+
84
+ @field_validator("batch_size")
85
+ @classmethod
86
+ def validate_batch_size(cls, v: Any) -> Any:
87
+ if isinstance(v, str):
88
+ if v.lower() != "auto":
89
+ raise ValueError("batch_size string value must be 'auto'")
90
+ return "auto"
91
+ if isinstance(v, int):
92
+ if v < 1:
93
+ raise ValueError("batch_size must be >= 1")
94
+ return v
95
+ raise ValueError(f"Invalid batch_size type: {type(v)}")
96
+
97
+ @field_validator("gradient_accumulation_steps")
98
+ @classmethod
99
+ def validate_grad_accum(cls, v: Any) -> Any:
100
+ if isinstance(v, str):
101
+ if v.lower() != "auto":
102
+ raise ValueError("gradient_accumulation_steps string must be 'auto'")
103
+ return "auto"
104
+ if isinstance(v, int):
105
+ if v < 1:
106
+ raise ValueError("gradient_accumulation_steps must be >= 1")
107
+ return v
108
+ raise ValueError(f"Invalid gradient_accumulation_steps type: {type(v)}")
109
+
110
+ def to_dict(self) -> Dict[str, Any]:
111
+ """Convert config to dictionary."""
112
+ return self.model_dump()
113
+
114
+
115
+ class OptimizationRecommendation(BaseModel):
116
+ """Structured recommendation produced by the optimization engine."""
117
+
118
+ recommended_batch_size: int
119
+ recommended_precision: Literal["fp32", "fp16", "bf16"]
120
+ recommended_grad_accum_steps: int
121
+ recommended_checkpointing: bool
122
+ estimated_peak_memory_bytes: int
123
+ memory_limit_bytes: int
124
+ reasons: Dict[str, str] = Field(default_factory=dict)
125
+ probe_results: Dict[str, Any] = Field(default_factory=dict)
126
+