ReferenceTrainingStepScenario
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.
One optimizer step of a reference pretraining job, fixed once for every chapter that quotes it.
The layout is replica (tensor parallelism across one node, the shallowest pipeline that fits unsharded training state) repeated across fleet for data parallelism. One step is one optimizer update over the global batch, split into microbatches of sequences_per_microbatch sequences, so the data-parallel gradient reduction happens once per step. The global batch is the cited token batch divided by the context length, rounded down to a multiple of batch_divisor so every data-parallel degree used in the book divides it. Step time is modeled by DistributedModel, not measured.
Notes
Equation: global sequences = rounded token batch / context; data degree = fleet devices / replica devices; step time = DistributedModel.solve(…).
Assumptions: The declared replica layout and global batch admit whole microbatches. The local efficiency and communication-overlap choice are inputs.
Excludes: A measured step-time claim, convergence prediction, and incomplete replica or microbatch allocations.
References
Derivation: This configuration forwards its stored inputs to mlsysim.solvers.DistributedModel.solve. The solver contract records the equations and their provenance.
Examples
The named scenario supplies illustrative inputs; the output is calculated.
>>> from mlsysim import Scenarios
>>> case = Scenarios.DenseLLMPretrainingStep
>>> round(float(case.tensor_parallel), 3)
8.0Attributes
| Name | Description |
|---|---|
| batch_divisor | |
| batch_tokens | |
| context_length | Tokens per sequence taken from the replica’s workload specification. |
| data_parallel | Whole replicas that fit in the fleet; incomplete replicas are discarded. |
| fleet | |
| global_batch_sequences | Global sequence count rounded down to a multiple of batch_divisor. |
| global_batch_tokens | Actual token count per optimizer step after sequence-count rounding. |
| local_efficiency | |
| microbatches_per_replica | Whole microbatches in each replica’s sequence allocation. |
| overlap_comm | |
| pipeline_depth | Pipeline stage count inherited from the self-contained replica layout. |
| precision | |
| replica | |
| result | Modeled result from mlsysim.solvers.DistributedModel.solve using this |
| sequences_per_microbatch | |
| state_per_gpu | Unsharded training-state memory per accelerator, excluding activations and |
| step_time | Modeled optimizer-step latency from DistributedModel.solve, carrying time |
| tensor_parallel | Tensor-parallel degree inherited from the self-contained replica layout. |
replica
Type: SelfContainedReplicaScenario.
Declared default: required.
fleet
Type: Fleet.
Declared default: required.
batch_tokens
Type: float.
Declared default: required.
batch_divisor
Type: int.
Declared default: required.
sequences_per_microbatch
Type: int.
Declared default: required.
precision
Type: str.
Declared default: required.
local_efficiency
Type: float.
Declared default: required.
overlap_comm
Type: bool.
Declared default: required.
tensor_parallel
engine.scenario_types.ReferenceTrainingStepScenario.tensor_parallel
Tensor-parallel degree inherited from the self-contained replica layout.
Equation: return replica.tensor_parallel.
pipeline_depth
engine.scenario_types.ReferenceTrainingStepScenario.pipeline_depth
Pipeline stage count inherited from the self-contained replica layout.
Equation: return replica.pipeline_depth.
data_parallel
engine.scenario_types.ReferenceTrainingStepScenario.data_parallel
Whole replicas that fit in the fleet; incomplete replicas are discarded.
Equation: return fleet.total_accelerators // replica.replica_gpus.
context_length
engine.scenario_types.ReferenceTrainingStepScenario.context_length
Tokens per sequence taken from the replica’s workload specification.
Equation: return replica.workload.context_length.
global_batch_sequences
engine.scenario_types.ReferenceTrainingStepScenario.global_batch_sequences
Global sequence count rounded down to a multiple of batch_divisor.
Equation: return int(batch_tokens / context_length) // batch_divisor * batch_divisor.
global_batch_tokens
engine.scenario_types.ReferenceTrainingStepScenario.global_batch_tokens
Actual token count per optimizer step after sequence-count rounding.
Equation: return global_batch_sequences * context_length.
microbatches_per_replica
engine.scenario_types.ReferenceTrainingStepScenario.microbatches_per_replica
Whole microbatches in each replica’s sequence allocation.
Equation: return global_batch_sequences // (data_parallel * sequences_per_microbatch).
state_per_gpu
engine.scenario_types.ReferenceTrainingStepScenario.state_per_gpu
Unsharded training-state memory per accelerator, excluding activations and workspace.
Equation: return replica.workload.parameters.magnitude * replica.bytes_per_param / replica.replica_gpus.
result
engine.scenario_types.ReferenceTrainingStepScenario.result
Modeled result from mlsysim.solvers.DistributedModel.solve using this scenario’s declared inputs; see that method for units, limits, and feasibility.
Equation: Delegates the declared fields to mlsysim.solvers.DistributedModel.solve; returned quantities and feasibility follow that contract.
step_time
engine.scenario_types.ReferenceTrainingStepScenario.step_time
Modeled optimizer-step latency from DistributedModel.solve, carrying time units.
Equation: return result.step_latency_total.