calc_subblock_params_and_memory

Calculate memory usage and parameter counts for neural network subblocks.

This module provides utilities to compute memory footprints and parameter counts for different subblock types (FFN, Attention, Mamba, MoE) in large language models, considering various data types, batch sizes, and sequence lengths.

Functions

calc_subblock_active_params

Calculate the number of "active" parameters for a subblock (FFN, Attention, or MoE).

calculate_additive_metrics

Return deterministic additive bytes and phase FLOPs for one subblock.

calculate_ffn_memory

Estimate the memory usage in MiB of a feed-forward network (FFN) subblock.

calculate_mamba_memory

Calculate memory usage (MiB) for a Mamba subblock.

calculate_mamba_state_size

Calculate the total state size for a Mamba attention subblock.

calculate_non_block_memory

Estimate the memory usage in MiB of non-subblock components (e.g., embeddings, output projection).

calculate_non_block_params

Calculate the number of parameters for non-subblock components (e.g., embeddings, output projection).

calculate_subblock_memory

Calculate the memory usage of a single subblock (FFN or Attention).

calculate_subblock_params

Count parameters on one meta decoder layer.
calc_subblock_active_params(sublayer_config, model_config, descriptor, n_embd, num_params=None)

Calculate the number of “active” parameters for a subblock (FFN, Attention, or MoE).

For non-MoE subblocks, simply calls calculate_subblock_params to count all parameters. For MoE (Mixture-of-Experts) FFN subblocks, the active parameter count is deterministic: the router selects a fixed top_k experts per token, so it is the router plus the always-on shared expert plus top_k routed experts.

Parameters:

  • sublayer_config (SubblockConfig) – The subblock configuration (either FFNConfig or AttentionConfig).

  • model_config (PreTrainedConfig) – The Hugging Face model configuration.

  • descriptor (type[ModelDescriptor]) – The ModelDescriptor class corresponding to this model family.

  • n_embd (int) – The embedding size (hidden dimension).

  • num_params (int | None)

Returns:

The number of “active” parameters for the given subblock.

Return type:

int

calculate_additive_metrics(subblock_config, *, model_config, descriptor, batch_size, prefill_seq_len, generation_seq_len, n_embd, n_head, weights_dtype, kv_cache_dtype, num_params=None, active_params=None)

Return deterministic additive bytes and phase FLOPs for one subblock.

Parameters:

  • subblock_config (SubblockConfig)

  • model_config (PreTrainedConfig)

  • descriptor (type[ModelDescriptor])

  • batch_size (int)

  • prefill_seq_len (int)

  • generation_seq_len (int)

  • n_embd (int)

  • n_head (int)

  • weights_dtype (dtype)

  • kv_cache_dtype (dtype)

  • num_params (int | None)

  • active_params (int | None)

Return type:

dict[str, float | int | dict[str, str]]

calculate_ffn_memory(ffn_config, model_config, descriptor, weights_dtype, experts_dtype=None, num_params=None)

Estimate the memory usage in MiB of a feed-forward network (FFN) subblock.

Parameters:

  • ffn_config (FFNConfig | MoEConfig) – FFN configuration for the block.

  • model_config (PreTrainedConfig) – The parent model configuration.

  • descriptor (type[ModelDescriptor]) – Model descriptor class.

  • weights_dtype (dtype | str) – Data type for FFN weights.

  • experts_dtype (dtype | str | None) – Data type for expert weights (for MoE layers, if present).

  • num_params (int | None)

Returns:

Estimated FFN memory usage in mebibytes (MiB).

Return type:

float

calculate_mamba_memory(mamba_config, model_config, descriptor, batch_size, weights_dtype, kv_cache_dtype, num_params=None)

Calculate memory usage (MiB) for a Mamba subblock.

Parameters:

  • mamba_config (MambaConfig) – Mamba configuration.

  • model_config (PreTrainedConfig) – Model configuration.

  • descriptor (type[ModelDescriptor]) – Model descriptor class.

  • batch_size (int) – Batch size for memory estimate.

  • weights_dtype (dtype) – Data type for model weights.

  • kv_cache_dtype (dtype) – Data type for state/kv-cache.

  • num_params (int | None)

Returns:

Estimated memory usage in mebibytes (MiB) for the Mamba subblock.

Return type:

int

calculate_mamba_state_size(mamba_config, batch_size)

Calculate the total state size for a Mamba attention subblock.

Parameters:

  • mamba_config (MambaConfig) – Configuration object containing Mamba subblock parameters.

  • batch_size (int) – Batch size to estimate the memory/state requirements for.

Returns:

Total state size (number of elements) required for the Mamba subblock, including convolution and SSM state.

Return type:

int

calculate_non_block_memory(n_embd, vocab_size, weight_dtype)

Estimate the memory usage in MiB of non-subblock components (e.g., embeddings, output projection).

Parameters:

  • n_embd (int)

  • vocab_size (int)

  • weight_dtype (dtype)

Return type:

float

calculate_non_block_params(n_embd, vocab_size)

Calculate the number of parameters for non-subblock components (e.g., embeddings, output projection).

Parameters:

  • n_embd (int)

  • vocab_size (int)

Return type:

int

calculate_subblock_memory(subblock_config, batch_size, prefill_seq_len, generation_seq_len, n_embd, n_head, weights_dtype, kv_cache_dtype, model_config, descriptor, num_params=None)

Calculate the memory usage of a single subblock (FFN or Attention).

Given its configuration and runtime dimensions, returns bytes or a detailed dict.

Parameters:

  • subblock_config (SubblockConfig) – Subblock configuration dataclass.

  • batch_size (int) – Batch size for memory estimate.

  • prefill_seq_len (int) – Sequence length for prefill phase.

  • generation_seq_len (int) – Sequence length for generation phase (token-by-token).

  • n_embd (int) – Embedding (hidden) dimension.

  • n_head (int) – Number of attention heads (used for non-FFN).

  • weights_dtype (dtype) – PyTorch dtype for model weights.

  • kv_cache_dtype (dtype) – PyTorch dtype for KV cache.

  • model_config (PreTrainedConfig) – HuggingFace-style config instance describing the model.

  • descriptor (type[ModelDescriptor]) – Model descriptor type (for puzzletron model types).

  • num_params (int | None)

Returns:

Memory usage in bytes (float), or a dictionary by memory type.

Return type:

float | dict[str, float]

calculate_subblock_params(config, layer_config, descriptor)

Count parameters on one meta decoder layer.

Parameters:
Return type:

int