mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-07-21 17:57:24 -05:00
Vijay Janapa Reddi
cb0ae4082f
Run quartodoc build against the refactored package (engine/ extracted from core/, infra->infrastructure): regenerate the API reference so every stub points at the new paths. Adds engine.*/infrastructure.* stubs, removes the 41 stale core.solver.* / core.engine.* / core.scenarios.* / core.config.* / core.evaluation.* / infra*.qmd orphans, and refreshes the package pages (core now primitives-only) + index. Verified: all 47 documented symbols resolve in the new package; every index link resolves; zero stale mlsysim.core.<engine-mod> / mlsysim.infra references anywhere in docs/. (objects.json inventory is gitignored -- regenerated at build.)
55 lines
1.7 KiB
Plaintext
55 lines
1.7 KiB
Plaintext
# engine.solver.TrainingMemoryModel { #mlsysim.engine.solver.TrainingMemoryModel }
|
|
|
|
```python
|
|
engine.solver.TrainingMemoryModel()
|
|
```
|
|
|
|
Decomposes per-accelerator training memory into teachable components.
|
|
|
|
This model answers a different question than ``SingleNodeModel``. Roofline
|
|
feasibility asks whether a workload's inference weights fit; training
|
|
feasibility must also account for gradients, optimizer state, activations,
|
|
and communication buffers. The accounting follows the common mixed-precision
|
|
state breakdown used by Megatron-LM and ZeRO.
|
|
|
|
Literature Source:
|
|
1. Shoeybi et al. (2019), "Megatron-LM" (tensor/pipeline parallel state).
|
|
2. Rajbhandari et al. (2020), "ZeRO" (data-parallel state sharding).
|
|
3. Korthikanti et al. (2023), activation recomputation accounting.
|
|
|
|
## Methods
|
|
|
|
| Name | Description |
|
|
| --- | --- |
|
|
| [solve](#mlsysim.engine.solver.TrainingMemoryModel.solve) | Estimate per-accelerator training memory. |
|
|
|
|
### solve { #mlsysim.engine.solver.TrainingMemoryModel.solve }
|
|
|
|
```python
|
|
engine.solver.TrainingMemoryModel.solve(
|
|
model,
|
|
hardware,
|
|
batch_size,
|
|
seq_len=2048,
|
|
precision='fp16',
|
|
optimizer='adam',
|
|
activation_checkpointing='selective',
|
|
tp_size=1,
|
|
pp_size=1,
|
|
dp_size=1,
|
|
ep_size=1,
|
|
zero_stage=0,
|
|
gradient_accumulation_steps=1,
|
|
trainable_fraction=1.0,
|
|
communication_buffer_fraction=0.05,
|
|
)
|
|
```
|
|
|
|
Estimate per-accelerator training memory.
|
|
|
|
``batch_size`` is the global batch. The activation term uses the local
|
|
microbatch implied by data parallelism and gradient accumulation. Model
|
|
states are sharded by tensor, pipeline, and expert parallelism first;
|
|
ZeRO then shards optimizer, gradient, and parameter states across the
|
|
data-parallel group according to its stage.
|