MlpTrainingStepScenario

Optimization scenarios · Implementation: mlsysim/engine/scenario_definitions/optimization.py.

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

One eager training step of a one-hidden-layer MLP, operator by operator.

Forward: x @ W1, ReLU, h @ W2, and a combined cross-entropy loss (loss_flops_per_logit operations per logit). Backward: the weight gradients of both layers and the hidden-activation gradient (raw inputs need no gradient), the ReLU derivative, and the loss derivative, which reads saved probabilities and integer labels of label_bytes each. Bias adds are folded into their matrix multiplies. Every value is bytes_per_value wide. Times are ideal lower bounds on hardware; the host issues one Python entry per forward operator plus one for backward(), and launches one kernel per operator.

Notes

Equation: Each dense multiply counts two FLOPs per MAC; operator tables separately sum work and tensor movement.

Assumptions: A dense two-layer network, separate eager operators, and the explicit operation-count conventions above.

Excludes: Cross-operator fusion, cache reuse, optimizer FLOPs, and measured utilization.

References

Derivation: Each dense multiply counts two FLOPs per MAC; operator tables separately sum work and tensor movement. Individual property contracts identify the exact accounting expression or reusable calculation. Registry entries supply the scenario inputs; this class does not turn them into measured outcomes.

Examples

Evaluate a named scenario using its declared inputs.

>>> from mlsysim import Scenarios
>>> scenario = Scenarios.MnistTwoLayerTrainingStepOnH100
>>> round(float(scenario.op_flops['matmul1'].magnitude), 4)
6422528.0

Attributes

Name	Description
backward_bytes	Bytes summed over the backward operators in `op_bytes` without fusion.
backward_flops	FLOPs summed over the declared backward operators in `op_flops`.
backward_ops
batch_size
bytes_per_value
compute_time	Ideal microseconds for forward and backward work at the selected peak FLOP/s.
dense_training_ratio	Dimensionless total training-matmul work divided by forward-matmul work; raw inputs do
dispatch
forward_bytes	Bytes summed over the forward operators in `op_bytes` without cross-operator cache
forward_flops	FLOPs summed over the declared forward operators in `op_flops`.
forward_ops
hardware
hidden_width
host	Host-dispatch scenario combining forward Python entries and forward/backward kernel
input_width
label_bytes
launch
loss_flops_per_logit
memory_time	Ideal microseconds for eager tensor traffic at peak memory bandwidth.
op_bytes	Per-operator bytes read and written, including saved activations and integer labels;
op_flops	Per-operator FLOP counts: dense matrix multiplication costs twice its MAC count; ReLU
optimizer_step_bytes	Bytes moved by one optimizer update, using the supplied values-moved count per
optimizer_values_moved
output_width
parameters	Trainable scalar count including both dense weight matrices and their bias vectors.
peak_rate	Hardware peak FLOP/s for the selected precision, without an efficiency derating.
precision
total_bytes	Forward plus backward bytes per batch; optimizer traffic is accounted separately.
total_flops	Forward plus backward FLOPs per batch; optimizer update work is not included.

Equation: For matrix shapes (m,k) and (k,n), FLOPs=2*m*k*n; elementwise counts use the batch and feature dimensions stated in the operator table.

op_bytes

engine.scenario_types.MlpTrainingStepScenario.op_bytes

Per-operator bytes read and written, including saved activations and integer labels; tensors are counted separately for each eager operator.

Equation: For each operator, bytes=value_width*(input values+output values); loss backward additionally reads batch_size*label_bytes.

forward_flops

engine.scenario_types.MlpTrainingStepScenario.forward_flops

FLOPs summed over the declared forward operators in op_flops.

Equation: _sum(op_flops, _MLP_FORWARD_OPS).

backward_flops

engine.scenario_types.MlpTrainingStepScenario.backward_flops

FLOPs summed over the declared backward operators in op_flops.

Equation: _sum(op_flops, _MLP_BACKWARD_OPS).

forward_bytes

engine.scenario_types.MlpTrainingStepScenario.forward_bytes

Bytes summed over the forward operators in op_bytes without cross-operator cache reuse.

Equation: _sum(op_bytes, _MLP_FORWARD_OPS).

backward_bytes

engine.scenario_types.MlpTrainingStepScenario.backward_bytes

Bytes summed over the backward operators in op_bytes without fusion.

Equation: _sum(op_bytes, _MLP_BACKWARD_OPS).

total_flops

engine.scenario_types.MlpTrainingStepScenario.total_flops

Forward plus backward FLOPs per batch; optimizer update work is not included.

Equation: forward_flops + backward_flops.

total_bytes

engine.scenario_types.MlpTrainingStepScenario.total_bytes

Forward plus backward bytes per batch; optimizer traffic is accounted separately.

Equation: forward_bytes + backward_bytes.

dense_training_ratio

engine.scenario_types.MlpTrainingStepScenario.dense_training_ratio

Dimensionless total training-matmul work divided by forward-matmul work; raw inputs do not require gradients.

Matrix-multiply work of the whole step over the forward matrix multiplies.

Equation: (forward matmul FLOPs+weight-gradient FLOPs+hidden-gradient FLOPs)/forward matmul FLOPs.

parameters

engine.scenario_types.MlpTrainingStepScenario.parameters

Trainable scalar count including both dense weight matrices and their bias vectors.

Equation: i, h, o = (input_width, hidden_width, output_width); result = i * h + h + h * o + o.

optimizer_step_bytes

engine.scenario_types.MlpTrainingStepScenario.optimizer_step_bytes

Bytes moved by one optimizer update, using the supplied values-moved count per parameter.

Equation: (parameters * optimizer_values_moved * bytes_per_value).to('byte').

peak_rate

engine.scenario_types.MlpTrainingStepScenario.peak_rate

Hardware peak FLOP/s for the selected precision, without an efficiency derating.

Equation: hardware.compute.precision_flops[precision].

compute_time

engine.scenario_types.MlpTrainingStepScenario.compute_time

Ideal microseconds for forward and backward work at the selected peak FLOP/s.

Equation: (total_flops / peak_rate).to('microsecond').

memory_time

engine.scenario_types.MlpTrainingStepScenario.memory_time

Ideal microseconds for eager tensor traffic at peak memory bandwidth.

Equation: (total_bytes / hardware.memory.bandwidth).to('microsecond').

host

engine.scenario_types.MlpTrainingStepScenario.host

Host-dispatch scenario combining forward Python entries and forward/backward kernel launches; device time is the sum of the compute and memory floors.

The host’s dispatch and launch work against the step’s ideal device time.

Equation: Python entries=forward_ops+1; kernel launches=forward_ops+backward_ops; ideal device time=compute_time+memory_time.