physics.memory

physics.memory

Model and activation memory accounting.

Functions

Name	Description
calc_activation_memory	Estimate saved Transformer training activations per device.
calc_capped_exponential_scale	Solve for the exponential scale of a context-capped length model.
calc_checkpoint_memory_units	Estimate evenly spaced checkpoint memory in layer-equivalent units.
calc_checkpoint_size	Calculate persistent training-checkpoint storage.
calc_csr_bytes	Calculate compressed sparse row matrix storage.
calc_decoder_matmul_weights	Count GPT-style decoder matrix weights by location.
calc_expected_paged_kv_tokens	Estimate page allocation under a capped exponential request model.
calc_kv_attention_design	Classify a cache geometry as MLA, GQA, or MHA.
calc_kv_block_count	Count fixed-size blocks needed by a token sequence.
calc_kv_cache_size	Calculate resident key/value cache storage.
calc_kv_resident_fit	Count complete equal-sized sequences that fit a memory pool.
calc_low_rank_update_params	Count trainable weights in a low-rank matrix update.
calc_mla_cache_size	Calculate resident Multi-Head Latent Attention cache storage.
calc_paged_kv_cache_size	Calculate page-rounded KV storage and internal fragmentation.
calc_residual_adapter_params	Count weights and biases in a bottleneck residual adapter.
calc_shared_prefix_fit	Count requests that share one physical prefix allocation.
calc_sharing_gain_limit	Calculate the asymptotic residency gain from prefix sharing.
calc_sparse_breakeven_zero_share	Find the zero fraction at which sparse and dense storage tie.
calc_sparse_bytes_per_nonzero	Calculate sparse value-plus-index storage per nonzero.
calc_update_step_traffic	Account for tensor traffic in one full-update training step.
model_memory	Calculate weight storage in a requested memory unit.
optimizer_state_bytes	Return optimizer storage per trainable parameter, excluding gradients.

calc_activation_memory

physics.memory.calc_activation_memory(
    n_layers,
    seq_len,
    batch_size,
    hidden_dim,
    n_heads=None,
    precision_bytes=2,
    strategy='selective',
)

Estimate saved Transformer training activations per device.

Parameters

Name	Type	Description	Default
n_layers	int	Positive number of layers resident on the device after pipeline partitioning.	required
seq_len	int	Positive tokens per sequence.	required
batch_size	int	Positive microbatch sequence count per device.	required
hidden_dim	int	Positive hidden width in elements.	required
n_heads	int	Positive attention-head count; required for strategy=“none” and otherwise unused.	`None`
precision_bytes	float	Positive bytes per activation element; default 2 scales the FP16 accounting.	`2`
strategy	str	Saved-state policy: none, selective (default), or full recomputation.	`'selective'`

Returns

Name	Type	Description
	Quantity	Total saved activation storage in bytes.

Notes

Equation: M = L * (b/2) * A. At FP16, A = 34*S*B*H + 5*N_heads*S**2*B for none, 34*S*B*H for selective, and 2*S*B*H for full.

Assumptions: Equal layer shapes and the FP16 saved-tensor convention of Korthikanti et al. Selective recomputation discards attention matrices; full recomputation retains layer inputs. Other precisions scale all coefficients uniformly.

Excludes: Weights, gradients, optimizer state, kernel scratch space, allocator reserve, sequence/tensor-parallel sharding, and schedule-dependent simultaneous microbatch residency.

The attention coefficient counts softmax probabilities, a dropout mask, and dropped probabilities. The FP16 byte widths are already included in the coefficients; multiplying them by precision_bytes without dividing by two would count precision twice.

References

Korthikanti et al. (2023), “Reducing Activation Recomputation in Large Transformer Models,” https://arxiv.org/abs/2205.05198 (bibliography key: korthikanti2023). The linear precision rescaling is a modeling convention.

Examples

>>> full = calc_activation_memory(2, 8, 1, 16, strategy="full")
>>> selective = calc_activation_memory(2, 8, 1, 16, strategy="selective")
>>> bool(full < selective)
True

calc_capped_exponential_scale

physics.memory.calc_capped_exponential_scale(mean_tokens, max_tokens)

Solve for the exponential scale of a context-capped length model.

Parameters

Name	Type	Description	Default
mean_tokens	float	Mean resident tokens after capping, in (0, max_tokens].	required
max_tokens	float	Positive context cap in tokens.	required

Returns

Name	Type	Description
	float	Uncapped exponential scale in tokens; infinity when every request reaches the cap.

Notes

Equation: mean_tokens = mu * (1 - exp(-max_tokens/mu)).

Assumptions: Length is min(X, max_tokens), where X is a continuous exponential variable. Bisection solves the monotone mean relation; the shape is assumed, not fitted to a trace.

Excludes: Discrete token effects and measured request-length distributions.

References

Derivation: integrate the survival function P(X > s) = exp(-s/mu) from zero to max_tokens. The capped mean rises monotonically from zero to the cap.

Examples

>>> math.isinf(calc_capped_exponential_scale(16, 16))
True
>>> mu = calc_capped_exponential_scale(8, 32)
>>> math.isclose(mu * (1 - math.exp(-32 / mu)), 8)
True

calc_checkpoint_memory_units

physics.memory.calc_checkpoint_memory_units(n_layers, checkpoints)

Estimate evenly spaced checkpoint memory in layer-equivalent units.

Parameters

Name	Type	Description	Default
n_layers	int	Positive depth L of an equal-sized layer chain.	required
checkpoints	int	Positive checkpoint count k; use k <= L for a physical partition.	required

Returns

Name	Type	Description
	float	Activation footprint measured in units of one layer, not bytes.

Notes

Equation: M_layers = k + L/k.

Assumptions: The forward pass keeps k checkpoints while backward recomputes one segment of L/k equal-sized layers at a time.

Excludes: Integer segment rounding, unequal layer sizes, tensor parallelism, and recomputation runtime.

References

Chen et al. (2016), “Training Deep Nets with Sublinear Memory Cost,” https://arxiv.org/abs/1604.06174 (bibliography key: chen2016training). The equal-segment approximation yields k + L/k.

Examples

>>> calc_checkpoint_memory_units(16, 4) == 4 + 16 / 4
True

calc_checkpoint_size

physics.memory.calc_checkpoint_size(n_params, bytes_per_param=14)

Calculate persistent training-checkpoint storage.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Nonnegative parameter count; quantities must be dimensionless/count-compatible.	required
bytes_per_param	Quantity or float	Stored bytes per parameter, default 14: an FP16 weight plus an FP32 master weight and two FP32 moments.	`14`

Returns

Name	Type	Description
	Quantity	Checkpoint storage in bytes.

Notes

Equation: M_checkpoint = P * b_checkpoint.

Assumptions: Every parameter has the same saved-state footprint; the default stores mixed-precision Adam state.

Excludes: Gradients, activations, dataset position, random-number state, serialization metadata, filesystem replication, and compression.

References

Derivation: multiply parameter count by the sum of persisted tensor widths. The default is 2 + 4 + 4 + 4 bytes per parameter, not a measured file size.

Examples

>>> calc_checkpoint_size(1000).m_as("byte") == 1000 * 14
True

calc_csr_bytes

physics.memory.calc_csr_bytes(nonzeros, rows, value_bytes, index_bytes)

Calculate compressed sparse row matrix storage.

Parameters

Name	Type	Description	Default
nonzeros	int	Nonnegative number of stored matrix entries, including any explicitly stored zeros.	required
rows	int	Nonnegative matrix row count.	required
value_bytes	Quantity or float	Nonnegative bytes per value, in units compatible with index_bytes.	required
index_bytes	Quantity or float	Nonnegative bytes per column index and row pointer.	required

Returns

Name	Type	Description
	Quantity or float	Total value, column-index, and row-pointer storage, preserving input byte units.

Notes

Equation: M = nnz*(b_value+b_index) + (rows+1)*b_index.

Assumptions: Each stored entry has one value and one column index; row boundaries use rows+1 pointers of the same width as indices.

Excludes: Object headers, alignment, capacity over-allocation, and separate indexing-width choices for columns and row offsets.

References

Derivation: sum the three CSR arrays: nnz values, nnz column indices, and rows+1 row pointers.

Examples

>>> calc_csr_bytes(5, 3, 2, 4) == 5 * (2 + 4) + (3 + 1) * 4
True

calc_decoder_matmul_weights

physics.memory.calc_decoder_matmul_weights(
    layers,
    hidden_dim,
    vocab_size,
    mlp_ratio=4,
)

Count GPT-style decoder matrix weights by location.

Parameters

Name	Type	Description	Default
layers	int	Positive decoder-block count L.	required
hidden_dim	int	Positive hidden width d.	required
vocab_size	int	Positive output-vocabulary size V.	required
mlp_ratio	float	Positive intermediate-width multiple r; default 4.	`4`

Returns

Name	Type	Description
	dict	Counts: block weights across all layers, block_matrices across all layers, and output projection weights.

Notes

Equation: block = L*(4*d**2 + 2*r*d**2); block_matrices = 4*L; output = V*d.

Assumptions: Each block has one fused QKV matrix, one attention output matrix, and two MLP matrices. The output projection is reported separately.

Excludes: Biases, normalization scales, position embeddings, GQA/MLA geometry, and gated-MLP extra matrices.

References

Derivation: sum the rectangular matrix element counts d*(3*d), d*d, d*(r*d), (r*d)*d, and V*d. Fusing QKV changes matrix count, not parameter count.

Examples

>>> result = calc_decoder_matmul_weights(2, 8, 16)
>>> result["block"] == 2 * (4 + 2 * 4) * 8**2
True

calc_expected_paged_kv_tokens

physics.memory.calc_expected_paged_kv_tokens(
    mean_request_tokens,
    max_seq_len,
    page_size_tokens=16,
)

Estimate page allocation under a capped exponential request model.

Parameters

Name	Type	Description	Default
mean_request_tokens	float	Positive mean resident tokens, at most max_seq_len.	required
max_seq_len	int	Positive cap on resident tokens per request.	required
page_size_tokens	int	Positive tokens per page; default 16.	`16`

Returns

Name	Type	Description
	tuple	Expected allocated tokens per request and expected unused fraction of that allocation, both floats.

Notes

Equation: E[pages] = sum(exp(-j*p/mu), j=0..K-1) = (1-r**K)/(1-r), with r=exp(-p/mu), K=ceil(S_max/p). Allocation is p*E[pages], and waste is 1-mean/allocation.

Assumptions: Resident length is min(X, S_max), with an exponential X calibrated only to the supplied capped mean. At the fixed-length limit, allocate ceil(S_max/p)*p tokens.

Excludes: External fragmentation, block metadata, time-varying occupancy, prefix sharing, and departures from the assumed length distribution.

References

Derivation: the tail-sum identity for the integer page count gives the finite geometric series. Kwon et al. (2023), “Efficient Memory Management for Large Language Model Serving with PagedAttention,” https://arxiv.org/abs/2309.06180 (bibliography key: kwon2023), motivates page allocation; the exponential request model is an additional assumption.

Examples

>>> allocation, waste = calc_expected_paged_kv_tokens(16, 16, page_size_tokens=8)
>>> allocation == 16 and waste == 0
True

calc_kv_attention_design

physics.memory.calc_kv_attention_design(heads, kv_heads=None, kv_lora_rank=None)

Classify a cache geometry as MLA, GQA, or MHA.

Parameters

Name	Type	Description	Default
heads	int	Positive query-head count for a physical attention configuration.	required
kv_heads	int	Key/value-head count; a positive count below heads selects GQA.	`None`
kv_lora_rank	int	Compressed KV latent width; any nonzero value selects MLA before checking head counts.	`None`

Returns

Name	Type	Description
	str	MLA, GQA, or MHA; absent geometry defaults to MHA.

Notes

Equation: design = MLA if kv_lora_rank is truthy, else GQA if 0 < kv_heads < heads, else MHA.

Assumptions: This is a routing convention based on declared geometry, not an inference from weights. Inputs describe an already valid model.

Excludes: Architectural validity checks, cache sizing, and distinctions between single-head MQA and other GQA configurations.

References

Derivation: classify whether cache storage uses a compressed latent, grouped key/value heads, or one key/value head per query head. The precedence is a local API convention.

Examples

>>> calc_kv_attention_design(8, kv_heads=2)
'GQA'

calc_kv_block_count

physics.memory.calc_kv_block_count(tokens, block_tokens)

Count fixed-size blocks needed by a token sequence.

Parameters

Name	Type	Description	Default
tokens	int	Nonnegative resident sequence length in tokens.	required
block_tokens	int	Positive tokens per allocation block.	required

Returns

Name	Type	Description
	int	Number of whole blocks, with zero blocks for an empty sequence.

Notes

Equation: N_blocks = ceil(tokens / block_tokens).

Assumptions: A block has identical token capacity for every sequence.

Excludes: Metadata bytes, sharing, and future token reservations.

References

Derivation: round the fractional block requirement upward because a partial block still occupies one block.

Examples

>>> calc_kv_block_count(17, 8) == math.ceil(17 / 8)
True

calc_kv_cache_size

physics.memory.calc_kv_cache_size(
    n_layers,
    n_heads,
    head_dim,
    seq_len,
    batch_size,
    bytes_per_elem=2,
    kv_precision_bytes=None,
)

Calculate resident key/value cache storage.

Parameters

Name	Type	Description	Default
n_layers	int	Positive number of decoder layers.	required
n_heads	int	Positive key/value-head count, not query-head count; use the GQA/MQA count.	required
head_dim	int	Positive elements per key/value head.	required
seq_len	int	Nonnegative resident token count per request, including prompt and generated tokens.	required
batch_size	int	Positive number of requests with this same resident length.	required
bytes_per_elem	Quantity or float	Bytes per cache element; default 2 for FP16/BF16.	`2`
kv_precision_bytes	Quantity or float	Cache-specific byte width overriding bytes_per_elem when supplied.	`None`

Returns

Name	Type	Description
	Quantity	Combined K and V storage in bytes.

Notes

Equation: M_KV = 2 * L * H_KV * d_head * S * B * b.

Assumptions: Each layer stores separate K and V vectors for every resident token and KV head. Both vectors use the effective byte width.

Excludes: Page rounding, quantization scales, shared-prefix deduplication, allocator reserve, and runtime workspaces.

References

Derivation: multiply the tensor dimensions for K and V and add their equal footprints. The leading two counts tensors, not bytes of precision.

Examples

>>> small = calc_kv_cache_size(2, 4, 8, 16, 1)
>>> calc_kv_cache_size(2, 4, 8, 32, 1) == 2 * small
True

calc_kv_resident_fit

physics.memory.calc_kv_resident_fit(pool, bytes_per_sequence)

Count complete equal-sized sequences that fit a memory pool.

Parameters

Name	Type	Description	Default
pool	Quantity	Nonnegative available memory in byte-compatible units.	required
bytes_per_sequence	Quantity	Positive total cache bytes per sequence, including any caller-selected reserve.	required

Returns

Name	Type	Description
	int	Whole sequence count; integer conversion equals floor for nonnegative ratios.

Notes

Equation: B = floor(M_pool / M_sequence).

Assumptions: Equal resident length and storage per request; the supplied pool is available exclusively for these caches.

Excludes: Block rounding, allocator headroom, other tensors, and output growth unless the caller includes them in the inputs.

References

Derivation: the capacity inequality B*M_sequence <= M_pool is solved for the largest nonnegative integer B.

Examples

>>> calc_kv_resident_fit(10 * ureg.byte, 3 * ureg.byte) == 10 // 3
True

calc_low_rank_update_params

physics.memory.calc_low_rank_update_params(in_dim, out_dim, rank)

Count trainable weights in a low-rank matrix update.

Parameters

Name	Type	Description	Default
in_dim	int	Positive input width d_in of the frozen matrix.	required
out_dim	int	Positive output width d_out of the frozen matrix.	required
rank	int	Positive factor rank r; a low-rank use normally has r below both matrix dimensions.	required

Returns

Name	Type	Description
	int	Trainable parameter count in the two update factors.

Notes

Equation: P = r*d_in + d_out*r.

Assumptions: The update is a product of an r-by-d_in down projection and a d_out-by-r up projection.

Excludes: Frozen base weights, biases, optimizer state, and any scale parameter.

References

Derivation: add the element counts of the two rectangular factors. This is the factor-size accounting used by LoRA; it does not predict adaptation quality.

Examples

>>> calc_low_rank_update_params(8, 16, 2) == 2 * (8 + 16)
True

calc_mla_cache_size

physics.memory.calc_mla_cache_size(
    n_layers,
    kv_lora_rank,
    qk_rope_head_dim,
    seq_len,
    batch_size,
    bytes_per_elem=2,
)

Calculate resident Multi-Head Latent Attention cache storage.

Parameters

Name	Type	Description	Default
n_layers	int	Positive decoder-layer count.	required
kv_lora_rank	int	Positive elements in the compressed key/value latent per token and layer.	required
qk_rope_head_dim	int	Positive elements in the decoupled rotary key per token and layer.	required
seq_len	int	Nonnegative resident tokens per request.	required
batch_size	int	Positive request count with the same resident length.	required
bytes_per_elem	Quantity or float	Bytes per cached element, default 2.	`2`

Returns

Name	Type	Description
	Quantity	Compressed latent plus rotary-key cache in bytes.

Notes

Equation: M_MLA = L * (r_KV + d_RoPE) * S * B * b.

Assumptions: The latent and rotary key are the only resident cache vectors. K and V are reconstructed/absorbed at use time, so no separate leading factor of two applies.

Excludes: Reconstruction computation, quantization metadata, page rounding, prefix sharing, and attention workspace.

References

DeepSeek-AI (2024), “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model,” https://arxiv.org/abs/2405.04434 (bibliography key: deepseek2024v2).

Examples

>>> calc_mla_cache_size(2, 8, 4, 16, 1).m_as("byte") == 2 * (8 + 4) * 16 * 2
True

calc_paged_kv_cache_size

physics.memory.calc_paged_kv_cache_size(
    n_layers,
    n_heads,
    head_dim,
    seq_len,
    batch_size,
    page_size_tokens=16,
    bytes_per_elem=2,
)

Calculate page-rounded KV storage and internal fragmentation.

Parameters

Name	Type	Description	Default
n_layers	int	Positive decoder-layer count.	required
n_heads	int	Positive key/value-head count.	required
head_dim	int	Positive elements per key/value head.	required
seq_len	int	Nonnegative resident tokens per request.	required
batch_size	int	Positive request count, each with the same length.	required
page_size_tokens	int	Positive tokens per physical page, default 16.	`16`
bytes_per_elem	Quantity or float	Bytes per cache element, default 2.	`2`

Returns

Name	Type	Description
	tuple	Allocated bytes as a Quantity and unused fraction of allocated bytes as a float; empty sequences return zero for both.

Notes

Equation: S_alloc = ceil(S/p)*p; M = 2*L*H_KV*d*S_alloc*B*b; f = (S_alloc-S)/S_alloc for nonempty allocation.

Assumptions: Each request receives whole pages and its final page is the only source of internal slack.

Excludes: External fragmentation, block-table metadata, prefix sharing, and reserve for future output tokens.

A single length divisible by the page size has no internal waste. Use calc_expected_paged_kv_tokens when the question concerns a distribution of lengths.

References

Kwon et al. (2023), “Efficient Memory Management for Large Language Model Serving with PagedAttention,” https://arxiv.org/abs/2309.06180 (bibliography key: kwon2023). Page rounding here is deterministic tensor accounting.

Examples

>>> size, waste = calc_paged_kv_cache_size(2, 4, 8, 16, 1, page_size_tokens=8)
>>> size == calc_kv_cache_size(2, 4, 8, 16, 1) and waste == 0
True

calc_residual_adapter_params

physics.memory.calc_residual_adapter_params(width, bottleneck)

Count weights and biases in a bottleneck residual adapter.

Parameters

Name	Type	Description	Default
width	int	Positive input/output width d.	required
bottleneck	int	Positive hidden width r of the adapter.	required

Returns

Name	Type	Description
	int	Trainable parameter count for both affine projections.

Notes

Equation: P = d*r + r + r*d + d.

Assumptions: A down projection maps d to r and an up projection maps r to d; each has a bias.

Excludes: Layer-normalization parameters, gating, activation storage, and the frozen surrounding model.

References

Derivation: count both projection matrices and their output biases. Residual addition introduces no learned parameter.

Examples

>>> calc_residual_adapter_params(8, 2) == 2 * 8 * 2 + 2 + 8
True

calc_shared_prefix_fit

physics.memory.calc_shared_prefix_fit(
    pool_blocks,
    shared_blocks,
    private_blocks,
)

Count requests that share one physical prefix allocation.

Parameters

Name	Type	Description	Default
pool_blocks	int	Total nonnegative physical block capacity; must be at least shared_blocks for a feasible pool.	required
shared_blocks	int	Nonnegative prefix blocks stored once for the group.	required
private_blocks	int	Positive unique suffix blocks needed by each request.	required

Returns

Name	Type	Description
	int	Whole request count; a negative value signals that the supplied prefix already exceeds the pool.

Notes

Equation: B = floor((pool_blocks - shared_blocks) / private_blocks).

Assumptions: Every admitted request reuses the same prefix and has the same private suffix footprint.

Excludes: Copy-on-write duplication, reference-count metadata, and unequal suffix lengths.

References

Derivation: solve shared_blocks + B*private_blocks <= pool_blocks for integer B.

Examples

>>> calc_shared_prefix_fit(20, 4, 3) == (20 - 4) // 3
True

calc_sharing_gain_limit

physics.memory.calc_sharing_gain_limit(shared_fraction)

Calculate the asymptotic residency gain from prefix sharing.

Parameters

Name	Type	Description	Default
shared_fraction	float	Fraction f of each request footprint shared across requests, in [0, 1).	required

Returns

Name	Type	Description
	float	Dimensionless large-pool request-capacity multiplier.

Notes

Equation: gain = 1 / (1 - f).

Assumptions: The shared prefix is stored once and becomes negligible as pool/request count grows.

Excludes: Finite-pool prefix cost, integer block rounding, and metadata overhead.

References

Derivation: the per-request private footprint is (1-f) times the unshared footprint; capacity scales inversely with footprint.

Examples

>>> calc_sharing_gain_limit(0) == 1
True

calc_sparse_breakeven_zero_share

physics.memory.calc_sparse_breakeven_zero_share(
    dense_bytes_per_value,
    bytes_per_nonzero,
)

Find the zero fraction at which sparse and dense storage tie.

Parameters

Name	Type	Description	Default
dense_bytes_per_value	Quantity or float	Positive bytes per dense entry.	required
bytes_per_nonzero	Quantity or float	Positive bytes per sparse nonzero, in units compatible with dense_bytes_per_value.	required

Returns

Name	Type	Description
	float	Dimensionless zero-share threshold; sparse wins above it. A negative threshold means sparse entries are already cheaper before sparsity.

Notes

Equation: z_star = 1 - b_dense/b_nonzero.

Assumptions: Only surviving nonzeros consume sparse storage; both representations cover the same logical element count.

Excludes: Per-row pointers, format headers, sparse compute overhead, and alignment.

References

Derivation: equate N*b_dense with N*(1-z)*b_nonzero and solve for z.

Examples

>>> calc_sparse_breakeven_zero_share(2, 6) == 1 - 2 / 6
True

calc_sparse_bytes_per_nonzero

physics.memory.calc_sparse_bytes_per_nonzero(
    value_bytes,
    index_bytes,
    indices=1,
)

Calculate sparse value-plus-index storage per nonzero.

Parameters

Name	Type	Description	Default
value_bytes	Quantity or float	Nonnegative bytes per stored value; use the same units as index_bytes.	required
index_bytes	Quantity or float	Nonnegative bytes per index.	required
indices	int	Nonnegative indices stored per nonzero: default 1 for CSR, 2 for COO.	`1`

Returns

Name	Type	Description
	Quantity or float	Per-nonzero storage in the input byte units.

Notes

Equation: b_nonzero = b_value + indices*b_index.

Assumptions: Every stored nonzero carries its value and the declared fixed number of indices.

Excludes: CSR row pointers, headers, padding, and compressed/variable-width index coding.

References

Derivation: add the storage widths of the fields in one sparse entry.

Examples

>>> calc_sparse_bytes_per_nonzero(2, 4, indices=2) == 2 + 2 * 4
True

calc_update_step_traffic

physics.memory.calc_update_step_traffic(
    weights,
    gradients,
    optimizer_state,
    activations,
)

Account for tensor traffic in one full-update training step.

Parameters

Name	Type	Description	Default
weights	Quantity or float	Nonnegative weight footprint; use bytes or byte-compatible quantities consistently.	required
gradients	Quantity or float	Nonnegative gradient footprint in the same units as weights.	required
optimizer_state	Quantity or float	Nonnegative persistent optimizer-state footprint in the same units.	required
activations	Quantity or float	Nonnegative saved-activation footprint in the same units.	required

Returns

Name	Type	Description
	tuple	Total step traffic and traffic from the three weight reads; both preserve the input memory units.

Notes

Equation: V = (W+A) + (A+W+G) + (W+G+S) + (W+S) = 4*W+2*A+2*G+2*S; weight_reads = 3*W.

Assumptions: Forward reads weights and writes saved activations. Backward reads both and writes gradients. The optimizer reads weights, gradients, state, then writes weights and state.

Excludes: Cache reuse, fused-kernel savings, temporary workspace, communication, and activation recomputation.

References

Derivation: sum each read and write at the memory boundary exactly once for the declared unfused execution sequence.

Examples

>>> total, reads = calc_update_step_traffic(2, 3, 4, 5)
>>> total == 4*2 + 2*3 + 2*4 + 2*5 and reads == 3*2
True

model_memory

physics.memory.model_memory(params, bytes_per_param, unit=MB)

Calculate weight storage in a requested memory unit.

Parameters

Name	Type	Description	Default
params	Quantity or int	Nonnegative parameter count, dimensionless or in param/count units.	required
bytes_per_param	Quantity or float	Nonnegative bytes per stored parameter; quantities must have byte units.	required
unit	pint.Unit	Memory unit for the returned magnitude; defaults to decimal MB.	`MB`

Returns

Name	Type	Description
	float	Weight-memory magnitude in unit; unlike most physics helpers, this API returns no attached unit.

Notes

Equation: M = P * b, converted from bytes to unit.

Assumptions: Each parameter occupies the same storage width; the caller supplies any quantized effective width.

Excludes: Activations, gradients, optimizer state, KV cache, metadata, padding, and packing overhead.

References

Derivation: summing P equal-width stored elements gives P times the bytes per element.

Examples

>>> model_memory(1000, 2, unit=ureg.byte) == 1000 * 2
True

optimizer_state_bytes

physics.memory.optimizer_state_bytes(optimizer)

Return optimizer storage per trainable parameter, excluding gradients.

Parameters

Name	Type	Description	Default
optimizer	str	Case-insensitive storage convention: adam, adamw, sgd, or none.	required

Returns

Name	Type	Description
	Quantity	Bytes per parameter; multiply by a parameter count to obtain total storage.

Notes

Equation: B = c * b_FP32, with c = 3 for Adam/AdamW, 1 for SGD, and 0 for none.

Assumptions: Adam retains an FP32 master weight and two FP32 moments. SGD retains a master weight but no momentum; this is an explicit mixed-precision convention.

Excludes: Model weights at execution precision, gradients, activations, allocator reserve, and optimizer workspaces.

References

Derivation: count the stored FP32 tensors per parameter. The Adam convention follows the model-state accounting in Rajbhandari et al. (2020), “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” https://arxiv.org/abs/1910.02054 (bibliography key: rajbhandari2020). The SGD convention is declared here, not inferred from every SGD implementation.

Examples

>>> optimizer_state_bytes("adamw") == 3 * BYTES_FP32
True