Skip to main content

models.types.TransformerWorkload

models.types.TransformerWorkload()

Workload representation of an autoregressive Transformer (e.g., LLMs).

Attributes

Name Description
head_width Return the number of components in one attention head.
kv_width Return the expanded key or value feature width per layer.
queries_per_kv_head Return the query-head sharing ratio for grouped attention.

Methods

Name Description
get_kv_cache_size Count cached attention state for a batch of sequences.
lower Lower one Transformer forward work unit and its resident weights.
training_memory Estimate per-device weights, gradients, optimizer state, and activations.

get_kv_cache_size

models.types.TransformerWorkload.get_kv_cache_size(
    seq_len,
    batch_size,
    precision=BYTES_FP16,
)

Count cached attention state for a batch of sequences.

Parameters

Name Type Description Default
seq_len int Nonnegative cached tokens per sequence, including any shared prefix separately in every request; this method does not deduplicate pages. required
batch_size int Nonnegative number of resident sequences. required
precision Quantity or str Bytes per stored element, or a registered format such as ‘fp16’. A byte quantity specifies the width directly; it must be positive. BYTES_FP16

Returns

Name Type Description
Quantity Cache storage in bytes for all layers and sequences.

See Also

mlsysim.physics.memory.calc_kv_cache_size : Standard KV tensor accounting. mlsysim.physics.memory.calc_mla_cache_size : Compressed latent accounting.

Notes

Equation: M_KV = 2 * L * H_KV * d_head * S * B * b for standard multi-head or grouped-query attention. When both MLA ranks are present, M_MLA = L * (kv_lora_rank + qk_rope_head_dim) * S * B * b.

Assumptions: all layers share the same attention geometry. Missing hidden_dim and heads fall back to 4096 and 32 here; an explicit head_dim overrides hidden_dim // heads. Missing kv_heads means full multi-head attention. Both MLA dimensions must be present to select the latent path.

Excludes: query activations, weight storage, page rounding, cache sharing, allocator metadata, and inference-engine workspaces.

References

Derivation: standard attention stores one key and one value vector per KV head, layer, token, and sequence. MLA stores a shared latent vector and the rotary-position key components instead of expanded head tensors. The arithmetic is delegated to the two named cache-accounting functions.

Examples

>>> from mlsysim import Models
>>> model = Models.Language.Llama3_8B
>>> model.get_kv_cache_size(128, 2) == 2 * model.get_kv_cache_size(128, 1)
True

lower

models.types.TransformerWorkload.lower(precision=BYTES_FP16)

Lower one Transformer forward work unit and its resident weights.

Parameters

Name Type Description Default
precision Quantity or str Bytes per stored element, or a registered format such as ‘fp16’. A byte quantity specifies the width directly; it must be positive. BYTES_FP16

Returns

Name Type Description
ComputationGraph Work in FLOPs, resident weight bytes, parameter count, arithmetic intensity in FLOP/byte, and the available layer count.

Notes

Equation: F = inference_flops if specified, otherwise 2 * P; W = size_in_bytes(precision); I = F / W.

Assumptions: the registered FLOP count defines the work unit when supplied; the fallback represents parameter-dominated dense work for one token. Layer count is copied into the graph for later overhead modeling.

Excludes: attention over cached context, KV storage, backward work, activation traffic, and device timing. Do not interpret this graph alone as a complete decode model.

References

Derivation: multiply parameter count by storage width to obtain weight bytes, then divide the declared or approximated work by those bytes. This is a demand-accounting projection; it does not predict device time.

Examples

>>> from mlsysim import Models
>>> model = Models.Language.Llama3_8B
>>> graph = model.lower('fp16')
>>> graph.weight_bytes == model.size_in_bytes('fp16')
True

training_memory

models.types.TransformerWorkload.training_memory(
    batch_size,
    seq_len,
    precision='fp16',
    optimizer='adam',
    strategy='selective',
    zero_stage=0,
    dp_size=1,
    is_lora=False,
    lora_rank=16,
)

Estimate per-device weights, gradients, optimizer state, and activations.

Parameters

Name Type Description Default
batch_size int Positive sequences resident on this device for activation accounting. required
seq_len int Positive tokens per sequence. required
precision str Registered weight and gradient format; optimizer storage uses its own mixed-precision convention in optimizer_state_bytes. 'fp16'
optimizer str Optimizer convention: ‘adam’, ‘adamw’, ‘sgd’, or ‘none’. 'adam'
strategy str Activation recomputation mode: ‘none’, ‘selective’, or ‘full’, delegated to calc_activation_memory. 'selective'
zero_stage int ZeRO stage: 0 leaves state local; 1 shards optimizer state; 2 also shards gradients; 3 also shards weights. 0
dp_size int Positive number of data-parallel ranks sharing the sharded state. 1
is_lora bool Restrict trainable gradient and optimizer state to an approximate adapter parameter fraction. Full frozen weights remain resident. False
lora_rank int Positive adapter rank when is_lora is true. The model assumes adapters on Q, K, V, and O projections in every layer. 16

Returns

Name Type Description
Quantity Modeled resident memory per device in GiB. This method does not compare the result with a device capacity or declare feasibility.

See Also

mlsysim.physics.memory.calc_activation_memory : Activation assumptions. mlsysim.physics.memory.optimizer_state_bytes : Optimizer storage convention.

Notes

Equation: M = M_weights + M_gradients + M_optimizer + M_activations. Weights and gradients initially cost P * b each; optimizer state costs P * optimizer_state_bytes(optimizer). LoRA multiplies gradient and optimizer storage by min(8 * rank * hidden_dim * layers / P, 1). ZeRO divides each selected state category by dp_size.

Assumptions: homogeneous data-parallel shards and the selected activation formula. Missing hidden_dim uses 4096. The adapter estimate covers four square attention projections, not every linear layer or architecture. Optimizer state follows the shared mixed-precision byte convention.

Excludes: communication buffers, allocator reserve, runtime workspaces, fragmentation, tensor-parallel activation partitioning, and peak transient all-gather state. ZeRO sharding is a storage estimate, not a latency model.

References

Rajbhandari et al. (2020), “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.” https://doi.org/10.1109/sc41405.2020.00024 (bibliography key: rajbhandari2020). The LoRA projection count and frozen-weight accounting are explicit implementation assumptions; the ZeRO citation supports state partitioning.

Examples

>>> from mlsysim import Models
>>> model = Models.Language.Llama3_8B
>>> local = model.training_memory(1, 128, zero_stage=0)
>>> sharded = model.training_memory(1, 128, zero_stage=3, dp_size=2)
>>> sharded < local
True
Back to top