mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-07-18 17:03:56 -05:00
63 lines
2.1 KiB
Plaintext
63 lines
2.1 KiB
Plaintext
# core.solver.TrainingMemoryModel { #mlsysim.core.solver.TrainingMemoryModel }
|
|
|
|
```python
|
|
core.solver.TrainingMemoryModel()
|
|
```
|
|
|
|
Decomposes per-accelerator training memory into weights, gradients, optimizer
|
|
state, activations, and communication buffers.
|
|
|
|
This model is intended for first-order training feasibility analysis. It makes
|
|
the difference between inference memory and training memory explicit without
|
|
modeling framework internals.
|
|
|
|
## Methods
|
|
|
|
| Name | Description |
|
|
| --- | --- |
|
|
| [solve](#mlsysim.core.solver.TrainingMemoryModel.solve) | Estimate per-accelerator training memory. |
|
|
|
|
### solve { #mlsysim.core.solver.TrainingMemoryModel.solve }
|
|
|
|
```python
|
|
core.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,
|
|
)
|
|
```
|
|
|
|
#### Parameters
|
|
|
|
| Name | Type | Description | Default |
|
|
| --- | --- | --- | --- |
|
|
| model | TransformerWorkload | Transformer workload to train. | _required_ |
|
|
| hardware | HardwareNode | Per-rank accelerator target. | _required_ |
|
|
| batch_size | int | Global batch size. | _required_ |
|
|
| seq_len | int | Training sequence length. | `2048` |
|
|
| precision | str | Parameter/gradient precision. | `'fp16'` |
|
|
| optimizer | str | `adam`, `adamw`, `sgd`, or `none`. | `'adam'` |
|
|
| activation_checkpointing | str | `none`, `selective`, or `full`. | `'selective'` |
|
|
| tp_size, pp_size, dp_size, ep_size | int | Parallelism degrees. | `1` |
|
|
| zero_stage | int | ZeRO stage 0--3. | `0` |
|
|
| gradient_accumulation_steps | int | Steps used to derive local microbatch. | `1` |
|
|
| trainable_fraction | float | Fraction of local parameters with gradients and optimizer state. | `1.0` |
|
|
| communication_buffer_fraction | float | Gradient bucket buffer fraction. | `0.05` |
|
|
|
|
#### Returns
|
|
|
|
`TrainingMemoryResult` with total memory, available memory, feasibility,
|
|
utilization, and a component breakdown.
|