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 +18 -0
- pymemora-0.1.0/PKG-INFO +179 -0
- pymemora-0.1.0/README.md +147 -0
- pymemora-0.1.0/memora/__init__.py +24 -0
- pymemora-0.1.0/memora/cli.py +104 -0
- pymemora-0.1.0/memora/config.py +126 -0
- pymemora-0.1.0/memora/engine.py +481 -0
- pymemora-0.1.0/memora/hardware/__init__.py +14 -0
- pymemora-0.1.0/memora/hardware/gpu_detector.py +221 -0
- pymemora-0.1.0/memora/hardware/memory_monitor.py +181 -0
- pymemora-0.1.0/memora/optimizer/__init__.py +17 -0
- pymemora-0.1.0/memora/optimizer/batch_optimizer.py +278 -0
- pymemora-0.1.0/memora/optimizer/checkpoint_optimizer.py +88 -0
- pymemora-0.1.0/memora/optimizer/gradient_optimizer.py +69 -0
- pymemora-0.1.0/memora/optimizer/precision_optimizer.py +113 -0
- pymemora-0.1.0/memora/profiler/__init__.py +6 -0
- pymemora-0.1.0/memora/profiler/training_profiler.py +122 -0
- pymemora-0.1.0/memora/reporting/__init__.py +6 -0
- pymemora-0.1.0/memora/reporting/report.py +142 -0
- pymemora-0.1.0/memora/utils/__init__.py +6 -0
- pymemora-0.1.0/memora/utils/common.py +72 -0
- pymemora-0.1.0/pymemora.egg-info/PKG-INFO +179 -0
- pymemora-0.1.0/pymemora.egg-info/SOURCES.txt +27 -0
- pymemora-0.1.0/pymemora.egg-info/dependency_links.txt +1 -0
- pymemora-0.1.0/pymemora.egg-info/entry_points.txt +2 -0
- pymemora-0.1.0/pymemora.egg-info/requires.txt +12 -0
- pymemora-0.1.0/pymemora.egg-info/top_level.txt +1 -0
- pymemora-0.1.0/pyproject.toml +49 -0
- pymemora-0.1.0/setup.cfg +4 -0
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
|
+
|
pymemora-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
+
[](https://pypi.org/project/pymemora/)
|
|
36
|
+
[](https://pypi.org/project/pymemora/)
|
|
37
|
+
[](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.
|
pymemora-0.1.0/README.md
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
1
|
+
# Memora 🧠⚡
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/pymemora/)
|
|
4
|
+
[](https://pypi.org/project/pymemora/)
|
|
5
|
+
[](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
|
+
|