ActivationMemoryScenario

Training scenarios · Implementation: mlsysim/engine/scenario_definitions/training.py.

Construct this scenario with keyword inputs. The fields below show the types and defaults declared by this class.

Activation memory of workload at a microbatch, under each recomputation strategy.

Notes

Equation: Compose mlsysim.physics.memory.KORTHIKANTI_PER_TOKEN_BYTES, mlsysim.physics.memory.KORTHIKANTI_SCORE_BYTES, mlsysim.physics.memory.calc_checkpoint_memory_units, mlsysim.physics.memory.calc_activation_memory, mlsysim.physics.transformer.calc_transformer_training_flops, mlsysim.physics.transformer.calc_transformer_decode_flops.

Assumptions: Transformer layers share geometry. Aggregate Korthikanti-style coefficients and explicit tensor line items are separate approximations; segment sizes use integer division.

Excludes: Allocator reserves, runtime workspaces, unequal layer shapes, and leftover segment layers are absent from the simple segment estimate.

References

Derivation: this scenario binds its input fields to the referenced implementation contracts: mlsysim.physics.memory.KORTHIKANTI_PER_TOKEN_BYTES, mlsysim.physics.memory.KORTHIKANTI_SCORE_BYTES, mlsysim.physics.memory.calc_checkpoint_memory_units, mlsysim.physics.memory.calc_activation_memory, mlsysim.physics.transformer.calc_transformer_training_flops, mlsysim.physics.transformer.calc_transformer_decode_flops. The methods below identify any additional arithmetic.

Examples

>>> from mlsysim import Scenarios
>>> case = Scenarios.Gpt3ActivationsMicrobatch4
>>> round(case.sequence, 4)
2048

Attributes

Name	Description
checkpoint_segments
checkpointed_activations	Boundary storage plus live-segment activations. Equation: `segment_inputs + \| \| [ffn_activations_per_layer](#mlsysim.engine.scenario_types.ActivationMemoryScenario.ffn_activations_per_layer) \| Illustrative expanded feed-forward activation tensor. Equation:`_elements(ffn_multiple
ffn_multiple
forward_equivalents	Training arithmetic expressed in forward-pass work units. Equation:
line_items_per_layer	Sum of the explicit projection, feed-forward, score, and optional normalization tensors.
line_items_total	Explicit line-item footprint across the model. Equation: line_items_per_layer * \| \| [live_segment](#mlsysim.engine.scenario_types.ActivationMemoryScenario.live_segment) \| Linear activation footprint of the currently live segment. Equation: full linear \| \| [micro_batch](#mlsysim.engine.scenario_types.ActivationMemoryScenario.micro_batch) \| \| \| [normalization_per_layer](#mlsysim.engine.scenario_types.ActivationMemoryScenario.normalization_per_layer) \| \| \| [per_sequence_activations](#mlsysim.engine.scenario_types.ActivationMemoryScenario.per_sequence_activations) \| Selective-recomputation activation footprint per sequence. Equation: \| \| [per_token_activations](#mlsysim.engine.scenario_types.ActivationMemoryScenario.per_token_activations) \| Linear-in-sequence activation footprint with attention-score recomputation. Equation: \| \| [per_token_bytes_coefficient](#mlsysim.engine.scenario_types.ActivationMemoryScenario.per_token_bytes_coefficient) \| Linear activation coefficient scaled to the requested element width. Equation: \| \| [precision_bytes](#mlsysim.engine.scenario_types.ActivationMemoryScenario.precision_bytes) \| \| \| [projection_activations_per_layer](#mlsysim.engine.scenario_types.ActivationMemoryScenario.projection_activations_per_layer) \| Illustrative saved state for four hidden-width projection tensors. Equation:4 *
recompute_forward_equivalents	Full-recomputation training arithmetic in forward-pass work units. Equation:
recompute_overhead	Relative arithmetic increase from full activation recomputation. Equation: checkpointed
score_activations	Stored attention-score contribution removed by selective recomputation. Equation:
score_activations_per_layer	Average stored attention-score footprint per layer. Equation: score_activations / \| \| [score_bytes_coefficient](#mlsysim.engine.scenario_types.ActivationMemoryScenario.score_bytes_coefficient) \| Attention-score coefficient expressed relative to the linear hidden-state term. \| \| [score_share](#mlsysim.engine.scenario_types.ActivationMemoryScenario.score_share) \| Attention-score fraction of the unrecomputed activation estimate. Equation: score \| \| [segment_input](#mlsysim.engine.scenario_types.ActivationMemoryScenario.segment_input) \| One segment boundary's retained hidden-state tensor. Equation: batch times sequence \| \| [segment_inputs](#mlsysim.engine.scenario_types.ActivationMemoryScenario.segment_inputs) \| Retained inputs across all checkpoint segments. Equation:checkpoint_segments *
segment_layers	Integer number of layers within each checkpoint segment. Equation: `workload.layers // \| \| [seq_len](#mlsysim.engine.scenario_types.ActivationMemoryScenario.seq_len) \| \| \| [sequence](#mlsysim.engine.scenario_types.ActivationMemoryScenario.sequence) \| Context length used by the activation model. Equation: explicit`seq_len`` when
workload

Methods

Name	Description
activations	Activation footprint for the named recomputation strategy. Equation:
checkpoint_memory_units	Normalized checkpoint storage at the requested checkpoint count. Equation:

activations

engine.scenario_types.ActivationMemoryScenario.activations(strategy)

Activation footprint for the named recomputation strategy. Equation: mlsysim.physics.memory.calc_activation_memory with scenario layers, sequence, microbatch, width, heads, and element bytes. Units: bytes; the helper owns the approximation and valid strategy names.

checkpoint_memory_units

engine.scenario_types.ActivationMemoryScenario.checkpoint_memory_units(
    checkpoints,
)

Normalized checkpoint storage at the requested checkpoint count. Equation: mlsysim.physics.memory.calc_checkpoint_memory_units(workload.layers, checkpoints). Units: layer-equivalent activation storage.

Activation memory with checkpoints evenly spaced checkpoints, in layers: k + L/k.

workload

Type: Workload.

Declared default: required.

micro_batch

Type: int.

Declared default: required.

seq_len

Type: Optional[int].

Declared default: None.

precision_bytes

Type: float.

Declared default: 2.

normalization_per_layer

Type: Optional[Quantity].

Declared default: None.

ffn_multiple

Type: int.

Declared default: 4.

checkpoint_segments

Type: Optional[int].

Declared default: None.

sequence

engine.scenario_types.ActivationMemoryScenario.sequence

Context length used by the activation model. Equation: explicit seq_len when present, otherwise workload.context_length. Units: tokens per sequence.

per_token_bytes_coefficient

engine.scenario_types.ActivationMemoryScenario.per_token_bytes_coefficient

Linear activation coefficient scaled to the requested element width. Equation: KORTHIKANTI_PER_TOKEN_BYTES * precision_bytes / 2. Units: bytes per hidden element; the reference coefficient is normalized to two-byte storage.

score_bytes_coefficient

engine.scenario_types.ActivationMemoryScenario.score_bytes_coefficient

Attention-score coefficient expressed relative to the linear hidden-state term. Equation: KORTHIKANTI_SCORE_BYTES * heads * sequence / hidden_dim * precision_bytes / 2. Units: bytes per hidden element under the same activation accounting convention.

score_share

engine.scenario_types.ActivationMemoryScenario.score_share

Attention-score fraction of the unrecomputed activation estimate. Equation: score coefficient divided by the sum of linear and score coefficients. Units: dimensionless; all shared shape factors cancel.

per_sequence_activations

engine.scenario_types.ActivationMemoryScenario.per_sequence_activations

Selective-recomputation activation footprint per sequence. Equation: per_token_activations / micro_batch. Units: GB; the property name refers to sequence allocation, not a single token.

Per-token activations of one sequence (scores not stored): the slope against batch size.

segment_layers

engine.scenario_types.ActivationMemoryScenario.segment_layers

Integer number of layers within each checkpoint segment. Equation: workload.layers // checkpoint_segments. Units: layers; leftover layers are not represented by this floor- division estimate.

segment_input

engine.scenario_types.ActivationMemoryScenario.segment_input

One segment boundary’s retained hidden-state tensor. Equation: batch times sequence times hidden width times element bytes via _elements. Units: MB; other boundary state is excluded.

segment_inputs

engine.scenario_types.ActivationMemoryScenario.segment_inputs

Retained inputs across all checkpoint segments. Equation: checkpoint_segments * segment_input. Units: GB; the same tensor shape is assumed at each boundary.

live_segment

engine.scenario_types.ActivationMemoryScenario.live_segment

Linear activation footprint of the currently live segment. Equation: full linear footprint divided by model layers times segment layers. Units: GB; segments are treated as equal-sized.

One segment’s per-token activations, rebuilt in the backward pass (scores not stored).

checkpointed_activations

engine.scenario_types.ActivationMemoryScenario.checkpointed_activations

Boundary storage plus live-segment activations. Equation: segment_inputs + live_segment. Units: GB; runtime workspaces and uneven segment effects are excluded.

Segment inputs plus one live rebuilt segment (its input counted twice: conservative).

per_token_activations

engine.scenario_types.ActivationMemoryScenario.per_token_activations

Linear-in-sequence activation footprint with attention-score recomputation. Equation: activations('selective'). Units: GB; despite the name this is the whole microbatch, not per-token storage.

Activations that grow in proportion to sequence length (everything but the attention scores).

score_activations

engine.scenario_types.ActivationMemoryScenario.score_activations

Stored attention-score contribution removed by selective recomputation. Equation: activations('none') - activations('selective'). Units: GB; both branches use identical workload geometry.

projection_activations_per_layer

engine.scenario_types.ActivationMemoryScenario.projection_activations_per_layer

Illustrative saved state for four hidden-width projection tensors. Equation: 4 * _elements(hidden_dim) * precision_bytes. Units: MB per layer; this line-item view is distinct from the aggregate coefficient model.

ffn_activations_per_layer

engine.scenario_types.ActivationMemoryScenario.ffn_activations_per_layer

Illustrative expanded feed-forward activation tensor. Equation: _elements(ffn_multiple * hidden_dim) * precision_bytes. Units: MB per layer; the supplied expansion factor sets the intermediate width.

score_activations_per_layer

engine.scenario_types.ActivationMemoryScenario.score_activations_per_layer

Average stored attention-score footprint per layer. Equation: score_activations / workload.layers. Units: MB; layers share the same geometry.

line_items_per_layer

engine.scenario_types.ActivationMemoryScenario.line_items_per_layer

Sum of the explicit projection, feed-forward, score, and optional normalization tensors. Equation: add these line items, using zero for unspecified normalization. Units: MB per layer; this is not forced to equal the separate aggregate coefficient estimate.

line_items_total

engine.scenario_types.ActivationMemoryScenario.line_items_total

Explicit line-item footprint across the model. Equation: line_items_per_layer * workload.layers. Units: GB; identical layer shapes are assumed.

forward_equivalents

engine.scenario_types.ActivationMemoryScenario.forward_equivalents

Training arithmetic expressed in forward-pass work units. Equation: calc_transformer_training_flops(P, 1) / calc_transformer_decode_flops(P, 1). Units: dimensionless; parameter-dominated accounting omits context attention.

recompute_forward_equivalents

engine.scenario_types.ActivationMemoryScenario.recompute_forward_equivalents

Full-recomputation training arithmetic in forward-pass work units. Equation: forward_equivalents * (1 + recompute_overhead). Units: dimensionless; the extra work is a recomputed forward pass.

recompute_overhead

engine.scenario_types.ActivationMemoryScenario.recompute_overhead

Relative arithmetic increase from full activation recomputation. Equation: checkpointed training FLOPs divided by baseline training FLOPs, minus one. Units: dimensionless; this is a work ratio, not measured runtime overhead.