Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsUtilsnemo_automodel.components.utils.flops_utils

nemo_automodel.components.utils.flops_utils

View as MarkdownOpen in Claude

Module Contents

Classes

NameDescription
_AttentionConfigTransformer model with multi-head attention.
_AttentionLayerConfigMulti-head attention layer.
_ClipViTConfigNeMo FLOPSConfig fields for CLIP ViT; a missing img_seq_len is derived and stored.
_CompositeConfigVision-language config that keeps its text backbone’s fields on text_config.
_DeepseekV3ConfigDeepSeek-V3 MoE width and router fan-out.
_DenseConfigTransformer model with multi-head attention and a dense MLP width.
_FluxConfigNeMo FLOPSConfig fields for FLUX; layers holds the joint and single block counts.
_GatedDeltaNetConfigQwen3.5 hybrid model: Gated DeltaNet (linear attention) head shapes.
_GptOssConfigGPT-OSS router fan-out.
_MLAConfigMulti-Latent Attention (DeepSeek-V3 style) head shapes.
_MambaLayerConfigMamba-2 layer head width.
_MlpLayerConfigDense MLP layer.
_MoELayerConfigNemotron-H MoE layer.
_ModelConfigHidden width, read by every model and layer FLOPs formula.
_NemotronHConfigNemotron-H hybrid model; per-layer fields are read only for the layer types its pattern uses.
_NevaProjectionConfigNeMo FLOPSConfig fields for the NeVA projector.
_TransformerConfigDepth and vocabulary of a transformer model.

Functions

NameDescription
_attrReturn field name of config, or None when its config family omits that field.
_build_moe_layer_patternBuild a list of 0/1 indicating dense(0) vs MoE(1) per layer.
_gdn_attention_per_layer_flopsFLOPs for a single Gated DeltaNet (GDN / linear attention) layer.
_has_text_backboneReturn whether config is a composite config whose text backbone lives on text_config.
_hybrid_model_flopsModel FLOPs for hybrid model
_int_attrReturn integer field name, or default when it is omitted or unset.
_is_dense_compositeNarrow a dense-transformer argument to its composite (vision-language) form.
_is_gated_delta_net_compositeNarrow a Qwen3.5 argument to its composite (vision-language) form.
_is_mla_compositeNarrow an MLA argument to its composite (vision-language) form.
_kda_attention_per_layer_flopsPer-layer FLOPs for Kimi Delta Attention (KDA, gated delta-rule linear attention).
_mamba_layer_flopsModel FLOPs for Mamba layer.
_mla_attention_per_layer_flopsPer-layer FLOPs for Multi-Latent Attention (MLA).
_mla_moe_model_flopsFLOPs for MLA + MoE transformer models (DeepSeek-V3 style).
_moe_layer_freq_attrReturn moe_layer_freq: a MoE layer interval, a per-layer 0/1 list, or None when omitted.
_nemotronh_mlp_layer_flopsModel FLOPs for MLP layer. Assume gated linear unit.
_nemotronh_moe_layer_flopsModel FLOPs for a MoE layer in Nemotron V3/Super V3 (hybrid Mamba/Attention/MoE).
_nemotronh_mtp_flopsModel FLOPs for the Multi-Token-Prediction (MTP) head of Nemotron-3 Super/Ultra.
_non_mla_attn_layer_flopsModel FLOPs for attention layer
_optional_bool_attrReturn boolean field name, or None when it is omitted or unset.
_optional_int_attrReturn integer field name, or None when it is omitted or unset.
_optional_int_list_attrReturn integer-list field name, or None when it is omitted or unset.
_optional_str_list_attrReturn string-list field name, or None when it is omitted or unset.
_sum_min_floor_divsum_{j=1..seq_len} min(cap, floor(j / ratio)) in closed form (cap=None disables the cap).
_sum_min_windowsum_{i=0..seq_len-1} min(i + 1, window): causal sliding-window keys per query, summed.
_sum_modsum_{j=1..seq_len} (j mod ratio) in closed form.
attention_flops_calculatorCalculate the flops for the attention part.
bert_flopsModel FLOPs for BERT family - accepts either AutoConfig or normalized config
calculate_mfuCalculate Model FLOPs Utilization (MFU).
clip_vit_l_flopsModel FLOPs for CLIP ViT
deepseek_v41_flopsModel FLOPs for DeepSeek-V4.1 (CSA2 sparse attention, single-pass mHC, Engram, MoE).
deepseekv3_flopsModel FLOPs for DeepSeek V3 - accepts either AutoConfig or normalized config
flux_flopsModel FLOPs for FLUX
get_flops_formula_for_hf_configGet the appropriate FLOPs formula function for a given HuggingFace config.
glm4_moe_flopsEstimate FLOPs for GLM4 MoE model configurations.
gpt3_flopsModel FLOPs for GPT3 family - accepts either AutoConfig or normalized config
gpt_oss_flopsModel FLOPs for GPT-OSS
gpt_oss_flops_calculatorCalculate the flops for the GPT-OSS model
kimi_k3_flopsModel FLOPs for Kimi K3 (hybrid KDA / gated-MLA attention + latent MoE with a SiTU shared expert).
llama2_flopsModel FLOPs for llama2 family - accepts either AutoConfig or normalized config
llama3_flopsModel FLOPs for llama3 family - accepts either AutoConfig or normalized config
loss_flops_calculatorCalculate the flops for the loss
minimax_m2_flopsModel FLOPs for MiniMax-M2 family - accepts either AutoConfig or normalized config.
mixtral_flopsModel FLOPs for mixtral family - accepts either AutoConfig or normalized config
mla_moe_flopsModel FLOPs for MLA + MoE models (Kimi K2, GLM-5, Mistral Small 4, etc.).
moe_mlp_flops_calculatorCalculate the flops for the MLP
nemotron_flopsModel FLOPs for nemotron family - accepts either AutoConfig or normalized config
nemotronh_flopsModel FLOPs for NemotronH
neva_projection_flopsModel FLOPs for NeVA Projection
qwen3_5_flopsModel FLOPs for Qwen3.5 family (MoE and Dense) with hybrid GDN/full attention.
qwen3_8_flash_next_flopsModel FLOPs for Qwen3.8-Flash-Next (GDN + compressed-block QSA, HyperConnections, Engram PLE, MoE).
qwen3_flopsModel FLOPs for Qwen3 family - accepts either AutoConfig or normalized config
step3_5_flash_flopsModel FLOPs for Step3.5-Flash (GQA + sliding-window / full attention + MoE).
transformer_flopsCalculate FLOPs for a standard Transformer model - accepts either AutoConfig or normalized config.

Data

_ConfigT_co

API

class nemo_automodel.components.utils.flops_utils._AttentionConfig()
Protocol

Bases: _TransformerConfig, _AttentionLayerConfig

Transformer model with multi-head attention.

class nemo_automodel.components.utils.flops_utils._AttentionLayerConfig()
Protocol

Bases: _ModelConfig

Multi-head attention layer.

num_attention_heads
int
class nemo_automodel.components.utils.flops_utils._ClipViTConfig()
Protocol

NeMo FLOPSConfig fields for CLIP ViT; a missing img_seq_len is derived and stored.

class_token_len
int
gbs
int
hs
int
img_h
int
img_seq_len
float | None
img_w
int
in_channels
int
layers
int
patch_dim
int
class nemo_automodel.components.utils.flops_utils._CompositeConfig()
Protocol

Bases: Protocol[_ConfigT_co]

Vision-language config that keeps its text backbone’s fields on text_config.

text_config
_ConfigT_co
class nemo_automodel.components.utils.flops_utils._DeepseekV3Config()
Protocol

Bases: _MLAConfig, _MlpLayerConfig

DeepSeek-V3 MoE width and router fan-out.

moe_intermediate_size
int
num_experts_per_tok
int
class nemo_automodel.components.utils.flops_utils._DenseConfig()
Protocol

Bases: _AttentionConfig, _MlpLayerConfig

Transformer model with multi-head attention and a dense MLP width.

class nemo_automodel.components.utils.flops_utils._FluxConfig()
Protocol

NeMo FLOPSConfig fields for FLUX; layers holds the joint and single block counts.

gbs
int
hs
int
in_channels
int
inp_s
int
layers
Sequence[int]
model_channels
int
vec_in_dim
int
class nemo_automodel.components.utils.flops_utils._GatedDeltaNetConfig()
Protocol

Bases: _AttentionConfig

Qwen3.5 hybrid model: Gated DeltaNet (linear attention) head shapes.

linear_key_head_dim
int
linear_num_key_heads
int
linear_num_value_heads
int
linear_value_head_dim
int
class nemo_automodel.components.utils.flops_utils._GptOssConfig()
Protocol

Bases: _DenseConfig

GPT-OSS router fan-out.

num_experts_per_tok
int
class nemo_automodel.components.utils.flops_utils._MLAConfig()
Protocol

Bases: _AttentionConfig

Multi-Latent Attention (DeepSeek-V3 style) head shapes.

kv_lora_rank
int
qk_nope_head_dim
int
qk_rope_head_dim
int
v_head_dim
int
class nemo_automodel.components.utils.flops_utils._MambaLayerConfig()
Protocol

Bases: _ModelConfig

Mamba-2 layer head width.

mamba_head_dim
int
class nemo_automodel.components.utils.flops_utils._MlpLayerConfig()
Protocol

Bases: _ModelConfig

Dense MLP layer.

intermediate_size
int
class nemo_automodel.components.utils.flops_utils._MoELayerConfig()
Protocol

Bases: _ModelConfig

Nemotron-H MoE layer.

moe_intermediate_size
int
moe_shared_expert_intermediate_size
int
n_routed_experts
int
num_experts_per_tok
int
class nemo_automodel.components.utils.flops_utils._ModelConfig()
Protocol

Hidden width, read by every model and layer FLOPs formula.

hidden_size
int
class nemo_automodel.components.utils.flops_utils._NemotronHConfig()
Protocol

Bases: _AttentionLayerConfig, _MambaLayerConfig, _MlpLayerConfig, _MoELayerConfig

Nemotron-H hybrid model; per-layer fields are read only for the layer types its pattern uses.

hybrid_override_pattern
str
vocab_size
int
class nemo_automodel.components.utils.flops_utils._NevaProjectionConfig()
Protocol

NeMo FLOPSConfig fields for the NeVA projector.

ffn_hs
int
gbs
int
hs
int
img_seq_len
int
inp_s
int
projector_type
str
class nemo_automodel.components.utils.flops_utils._TransformerConfig()
Protocol

Bases: _ModelConfig

Depth and vocabulary of a transformer model.

num_hidden_layers
int
vocab_size
int
nemo_automodel.components.utils.flops_utils._attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str
) -> object

Return field name of config, or None when its config family omits that field.

nemo_automodel.components.utils.flops_utils._build_moe_layer_pattern(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    layers: int
) -> list[int]

Build a list of 0/1 indicating dense(0) vs MoE(1) per layer.

Handles multiple config styles: first_k_dense_replace + moe_layer_freq, mlp_layer_types list, etc.

nemo_automodel.components.utils.flops_utils._gdn_attention_per_layer_flops(
    gbs: int,
    seq_len: int,
    hidden_size: int,
    linear_key_head_dim: int,
    linear_value_head_dim: int,
    linear_num_key_heads: int,
    linear_num_value_heads: int,
    linear_conv_kernel_dim: int
) -> float

FLOPs for a single Gated DeltaNet (GDN / linear attention) layer.

Based on the GDN FLOPs calculator from Megatron-Bridge PR #2925.

Return whether config is a composite config whose text backbone lives on text_config.

nemo_automodel.components.utils.flops_utils._hybrid_model_flops(
    config: nemo_automodel.components.utils.flops_utils._NemotronHConfig,
    gbs: int,
    seq_len: int
) -> float

Model FLOPs for hybrid model

nemo_automodel.components.utils.flops_utils._int_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str,
    default: int
) -> int

Return integer field name, or default when it is omitted or unset.

Narrow a dense-transformer argument to its composite (vision-language) form.

Narrow a Qwen3.5 argument to its composite (vision-language) form.

Narrow an MLA argument to its composite (vision-language) form.

nemo_automodel.components.utils.flops_utils._kda_attention_per_layer_flops(
    gbs: int,
    seq_len: int,
    hs: int,
    head_dim: int,
    num_heads: int,
    conv_kernel_size: int,
    use_full_rank_gate: bool,
    chunk_size: int = 64
) -> float

Per-layer FLOPs for Kimi Delta Attention (KDA, gated delta-rule linear attention).

Projections follow the KDA layer definition (Kimi Linear, arXiv:2510.26692): q/k/v with short causal convolutions, the low-rank forget gate (f_a/f_b), the beta projection, the output gate (full-rank g_proj or low-rank g_a/g_b) and o_proj. The chunkwise kernel is costed with the paper’s own count, FLOPs_KDA(T; C, d_h) = 6 T d_h^2 + 3 T C d_h + T C^2 per head (forward, chunk size C), halved to MAC-equivalents so the shared 6 x (forward + backward) factor of this module applies.

nemo_automodel.components.utils.flops_utils._mamba_layer_flops(
    config: nemo_automodel.components.utils.flops_utils._MambaLayerConfig,
    gbs: int,
    seq_len: int
) -> float

Model FLOPs for Mamba layer.

Three components:

  • in_proj: input projections (x_proj, z_proj, dt_proj, B_proj, C_proj)
  • scan: SSM scan kernel (7x factor accounts for the full SSD scan cost)
  • out_proj: output projection back to hidden_size Multiplied by 6 (3x fwd+bwd * 2x FMA) for in_proj/out_proj (standard GEMMs), and 7 * 3 = 21 for scan (non-GEMM kernel, higher op count per element).
nemo_automodel.components.utils.flops_utils._mla_attention_per_layer_flops(
    gbs: int,
    seq_len: int,
    hs: int,
    attention_heads: int,
    q_lora_rank: int | None,
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    qk_nope_head_dim: int,
    v_head_dim: int,
    index_topk: int | None = None,
    index_n_heads: int = 0,
    index_head_dim: int = 0
) -> float

Per-layer FLOPs for Multi-Latent Attention (MLA).

Shared by DeepSeek V3, Kimi K2.5, Mistral Small 4, GLM-5, etc.

When index_topk is set (DSA / sparse attention), accounts for:

  • Sparse main attention BMM: S * index_topk instead of 0.5 * S^2
  • DSA indexer overhead: Q/K/weights projections + full S^2 indexer BMM
nemo_automodel.components.utils.flops_utils._mla_moe_model_flops(
    gbs: int,
    seq_len: int,
    hs: int,
    layers: int,
    attention_heads: int,
    vocab_size: int,
    q_lora_rank: int | None,
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    qk_nope_head_dim: int,
    v_head_dim: int,
    dense_ffn_hs: int,
    moe_ffn_hs: int,
    moe_router_topk: int,
    moe_shared_expert_hs: int,
    moe_layer_pattern: collections.abc.Sequence[int],
    mtp_num_layers: int = 0,
    index_topk: int | None = None,
    index_n_heads: int = 0,
    index_head_dim: int = 0
) -> float

FLOPs for MLA + MoE transformer models (DeepSeek-V3 style).

Parameters:

moe_layer_pattern
Sequence[int]

List of 0/1 per layer (0=dense, 1=MoE).

moe_shared_expert_hs
int

Total intermediate size for all shared experts combined.

index_topk
int | NoneDefaults to None

If set, use DSA sparse attention with this many selected positions.

index_n_heads
intDefaults to 0

Number of heads in the DSA indexer.

index_head_dim
intDefaults to 0

Head dimension of the DSA indexer.

nemo_automodel.components.utils.flops_utils._moe_layer_freq_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig
) -> int | list[int] | None

Return moe_layer_freq: a MoE layer interval, a per-layer 0/1 list, or None when omitted.

nemo_automodel.components.utils.flops_utils._nemotronh_mlp_layer_flops(
    config: nemo_automodel.components.utils.flops_utils._MlpLayerConfig,
    gbs: int,
    seq_len: int
) -> float

Model FLOPs for MLP layer. Assume gated linear unit.

nemo_automodel.components.utils.flops_utils._nemotronh_moe_layer_flops(
    config: nemo_automodel.components.utils.flops_utils._MoELayerConfig,
    gbs: int,
    seq_len: int
) -> float

Model FLOPs for a MoE layer in Nemotron V3/Super V3 (hybrid Mamba/Attention/MoE).

Nemotron V3 uses relu2 (non-gated) for both routed and shared experts, so each expert has 2 linear projections (up_proj + down_proj), not 3.

When moe_latent_size is set (Super V3), routed experts operate in a reduced latent space with additional projection layers (fc1_latent_proj, fc2_latent_proj). The shared expert and gate always operate in the full hidden_size dimension.

Accounts for:

  1. Routed experts: only num_experts_per_tok activated per token.
  2. Shared expert: always active for every token (full hidden_size).
  3. Router/gate: linear projection hidden_size -> n_routed_experts.
  4. Latent projections (if moe_latent_size is set): down and up projections.
nemo_automodel.components.utils.flops_utils._nemotronh_mtp_flops(
    config: nemo_automodel.components.utils.flops_utils._NemotronHConfig,
    gbs: int,
    seq_len: int,
    num_mtp_layers: int,
    mtp_block_types: collections.abc.Sequence[str],
    use_repeated_layer: bool
) -> float

Model FLOPs for the Multi-Token-Prediction (MTP) head of Nemotron-3 Super/Ultra.

The head predicts num_mtp_layers (N) additional tokens. Each of the N depths runs the MTP block pattern once, plus a depth-fusion projection (eh_proj: cat[enorm(embed), hnorm(hidden)] of size 2*hidden -> hidden) and a vocab projection (weight-tied lm_head).

Repeated layer: when use_repeated_layer is True the model builds a SINGLE physical depth (mtp_block_types lists its sublayers) and reuses it across all N depths — but it still EXECUTES once per depth, so the block compute is N x the physical sublayers. That N x is the repeated layer’s FLOPs. When False, mtp_block_types already spans all N physical depths, so the block runs once.

These settings are NOT fully recoverable from the HF config (it retains only the physical depth count and omits the block pattern), so callers pass the effective values read from the built model: num_mtp_layers = model.mtp_config.num_layers and mtp_block_types = [s.block_type for s in model.mtp.layers].

Parameters:

num_mtp_layers
int

effective MTP depths actually run (model.mtp_config.num_layers).

mtp_block_types
Sequence[str]

block types (“mamba”/“attention”/“mlp”/“moe”) of the physical MTP sublayers (model.mtp.layers).

use_repeated_layer
bool

True if the physical depth is reused across the N depths.

nemo_automodel.components.utils.flops_utils._non_mla_attn_layer_flops(
    config: nemo_automodel.components.utils.flops_utils._AttentionLayerConfig,
    gbs: int,
    seq_len: int
) -> float

Model FLOPs for attention layer

nemo_automodel.components.utils.flops_utils._optional_bool_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str
) -> bool | None

Return boolean field name, or None when it is omitted or unset.

nemo_automodel.components.utils.flops_utils._optional_int_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str
) -> int | None

Return integer field name, or None when it is omitted or unset.

nemo_automodel.components.utils.flops_utils._optional_int_list_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str
) -> list[int] | None

Return integer-list field name, or None when it is omitted or unset.

nemo_automodel.components.utils.flops_utils._optional_str_list_attr(
    config: nemo_automodel.components.utils.flops_utils._ModelConfig,
    name: str
) -> list[str] | None

Return string-list field name, or None when it is omitted or unset.

nemo_automodel.components.utils.flops_utils._sum_min_floor_div(
    seq_len: int,
    ratio: int,
    cap: int | None
) -> int

sum_{j=1..seq_len} min(cap, floor(j / ratio)) in closed form (cap=None disables the cap).

floor(j / ratio) is the number of complete compressed groups visible to the query at zero-based position j - 1; the cap is the sparse top-k. Closed form so million-token sequences cost nothing to evaluate.

nemo_automodel.components.utils.flops_utils._sum_min_window(
    seq_len: int,
    window: int
) -> int

sum_{i=0..seq_len-1} min(i + 1, window): causal sliding-window keys per query, summed.

nemo_automodel.components.utils.flops_utils._sum_mod(
    seq_len: int,
    ratio: int
) -> int

sum_{j=1..seq_len} (j mod ratio) in closed form.

j mod ratio is the length of the incomplete compressed tail visible to the query at zero-based position j - 1.

nemo_automodel.components.utils.flops_utils.attention_flops_calculator(
    seqlen: int,
    hidden_size: int,
    num_attention_heads: int,
    num_query_groups: int,
    kv_channels: int | None = None,
    is_swa: bool = False,
    swa_window_size: int = 128
) -> float

Calculate the flops for the attention part.

nemo_automodel.components.utils.flops_utils.bert_flops(
    config: nemo_automodel.components.utils.flops_utils._TransformerConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for BERT family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.calculate_mfu(
    tflops: float,
    world_size: int,
    time_seconds: float,
    reference_mfu: float | None = None
) -> float

Calculate Model FLOPs Utilization (MFU).

Parameters:

tflops
float

Total TFLOPs across all devices for the measured step.

world_size
int

Total number of GPUs.

time_seconds
float

Time taken for computation.

reference_mfu
float | NoneDefaults to None

Peak TFLOPs/s per device for the training precision. The legacy default is the H100 dense-FP8 peak.

Returns: float

MFU as a percentage.

nemo_automodel.components.utils.flops_utils.clip_vit_l_flops(
    config: nemo_automodel.components.utils.flops_utils._ClipViTConfig
) -> float

Model FLOPs for CLIP ViT

nemo_automodel.components.utils.flops_utils.deepseek_v41_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for DeepSeek-V4.1 (CSA2 sparse attention, single-pass mHC, Engram, MoE).

Accepts DeepseekV41TextConfig or the multimodal DeepseekV41Config wrapper (its text_config is used; the vision tower is not counted, as for other VL entries). The module shapes follow nemo_automodel/components/models/deepseek_v41:

  • attention linears per layer: wq_a (hidden x q_lora_rank), wq_b (q_lora_rank x headshead_dim), wkv (hidden x head_dim, one shared latent), grouped wo_a (headshead_dim -> o_groupso_lora_rank) and wo_b (o_groupso_lora_rank x hidden);
  • compressors on kv_source_layer_ids: wkv and, for ratio > 1, wgate (hidden x head_dim each, applied to every token before pooling);
  • sparse attention BMMs: every query attends to min(i+1, sliding_window) local keys plus min(index_topk, floor((i+1)/ratio)) selected compressed keys (SWA-only layers have compress_ratios[i] == 0); QK^T and PV each cost heads * head_dim MACs per key;
  • the indexer is frozen (_Indexer.requires_grad_(False); hard top-k has no gradient), so its projections and its causal scoring over the compressed positions are counted forward only (2 x MACs) — the implementation scores the full [S, S/ratio] block before masking, which is executed work but not model work;
  • MoE per layer: router GEMM (hidden x n_routed_experts) plus num_experts_per_tok + n_shared_experts experts of 3 * hidden * moe_intermediate_size;
  • single-pass mHC: two coefficient projections per layer (hc_mult*(hc_mult+2) x hc_mult*hidden); the stream collapse/expand are elementwise and not counted;
  • Engram on engram_layer_ids: the fused key/value projection ((engram_max_ngram_size-1)*engram_n_heads*engram_head_dim x hidden*(hc_mult+1)); the table gather, hashing, norms, RoPE and quantize/dequantize boundaries are not GEMMs;
  • the FP32 language-model head (hidden x vocab).

Trainable GEMMs and attention BMMs cost 6 x MACs (forward + backward); activation recomputation is excluded (that is HFU). For the released Flash configuration the formula counts 16.09B active GEMM parameters per token (model card: 16B activated during decode) and 105.6 GFLOPs/token at 4096 tokens, 1.095x the dense-FFN identity 6 x active.

Parameters:

config
Any

DeepseekV41TextConfig, or a DeepseekV41Config wrapper whose text_config is used.

gbs
intDefaults to 1

Number of sequences per step.

seq_len
int | NoneDefaults to None

Tokens per sequence; defaults to config.max_position_embeddings.

Returns: float

Model FLOPs for one training step of gbs sequences of seq_len tokens.

nemo_automodel.components.utils.flops_utils.deepseekv3_flops(
    config: nemo_automodel.components.utils.flops_utils._DeepseekV3Config,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for DeepSeek V3 - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.flux_flops(
    config: nemo_automodel.components.utils.flops_utils._FluxConfig
) -> float

Model FLOPs for FLUX

nemo_automodel.components.utils.flops_utils.get_flops_formula_for_hf_config(
    config: typing.Any
) -> typing.Callable | None

Get the appropriate FLOPs formula function for a given HuggingFace config.

Parameters:

config
Any

HuggingFace model config object

Returns: Callable | None

The appropriate FLOPs formula function, or None for an unregistered

nemo_automodel.components.utils.flops_utils.glm4_moe_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Estimate FLOPs for GLM4 MoE model configurations.

nemo_automodel.components.utils.flops_utils.gpt3_flops(
    config: nemo_automodel.components.utils.flops_utils._TransformerConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for GPT3 family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.gpt_oss_flops(
    config: nemo_automodel.components.utils.flops_utils._GptOssConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for GPT-OSS

nemo_automodel.components.utils.flops_utils.gpt_oss_flops_calculator(
    gbs: int,
    num_layers: int,
    seqlen: int,
    hidden_size: int,
    num_attention_heads: int,
    num_query_groups: int,
    moe_ffn_hidden_size: int,
    moe_router_topk: int,
    vocab_size: int,
    kv_channels: int | None = None,
    swa_window_size: int = 128,
    window_attn_skip_freq: int = 2
) -> float

Calculate the flops for the GPT-OSS model

nemo_automodel.components.utils.flops_utils.kimi_k3_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for Kimi K3 (hybrid KDA / gated-MLA attention + latent MoE with a SiTU shared expert).

Accepts KimiK3TextConfig or the multimodal KimiK3Config / KimiK31Config wrapper (its text_config is used; the vision tower is not counted, as for other VL entries). Kimi K3.1 full-attention layers listed in dsa_config["dsa_layers"] are costed as sparse MLA plus the DSA indexer.

The layer pattern is read from the config rather than assumed: 1-based linear_attn_config["kda_layers"] marks KDA layers and every other layer is MLA (with an extra output-gate projection when mla_use_output_gate); first_k_dense_replace / moe_layer_freq mark dense vs MoE MLPs, using the same rule as KimiDecoderLayer. Routed experts run in a latent space of routed_expert_hidden_size behind shared down/up projections, and the shared expert is one SiTU MLP of moe_intermediate_size * num_shared_experts. The router GEMM (hidden_size x num_experts, about 0.6% of the total), norms and the scalar attention-residual projections are omitted, matching the other MoE formulas in this module.

The router top-k field on this config is num_experts_per_token (not the num_experts_per_tok that the generic transformer_flops fallback probes), which is why K3 was previously costed as a dense 93-layer transformer at well under half of its real FLOPs. For the released checkpoint’s configuration (96 attention heads, dense intermediate_size 33792) the formula counts 104.0B active matmul parameters per token and 2.78T total, matching the model card (“104B activated / 2.8T”) and the tech report’s Table 1 (104.2B / 2.78T) to within 0.2%, which is inside the accounting difference for the non-GEMM parameters (router, norms, gate biases). The report’s single MTP layer is not shipped in the checkpoint (num_nextn_predict_layers is 0) and is not modelled by Automodel, so it is not counted.

nemo_automodel.components.utils.flops_utils.llama2_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for llama2 family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.llama3_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for llama3 family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.loss_flops_calculator(
    seqlen: int,
    hidden_size: int,
    vocab_size: int
) -> float

Calculate the flops for the loss

nemo_automodel.components.utils.flops_utils.minimax_m2_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for MiniMax-M2 family - accepts either AutoConfig or normalized config.

Architecture: GQA attention (Q/K/V/O separate projections, head_dim may differ from hidden_size // num_heads) + MoE with SwiGLU (no shared experts by default). Optionally includes MTP (Multi-Token Prediction) modules gated by use_mtp.

nemo_automodel.components.utils.flops_utils.mixtral_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for mixtral family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.mla_moe_flops(
    config: nemo_automodel.components.utils.flops_utils._MLAConfig | nemo_automodel.components.utils.flops_utils._CompositeConfig[nemo_automodel.components.utils.flops_utils._MLAConfig],
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for MLA + MoE models (Kimi K2, GLM-5, Mistral Small 4, etc.).

Handles VL wrappers by extracting text_config if present.

nemo_automodel.components.utils.flops_utils.moe_mlp_flops_calculator(
    seqlen: int,
    hidden_size: int,
    moe_ffn_hidden_size: int,
    moe_router_topk: int,
    gated_linear_unit: bool = True
) -> float

Calculate the flops for the MLP

nemo_automodel.components.utils.flops_utils.nemotron_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for nemotron family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.nemotronh_flops(
    config: nemo_automodel.components.utils.flops_utils._NemotronHConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for NemotronH

nemo_automodel.components.utils.flops_utils.neva_projection_flops(
    config: nemo_automodel.components.utils.flops_utils._NevaProjectionConfig
) -> float

Model FLOPs for NeVA Projection

Model FLOPs for Qwen3.5 family (MoE and Dense) with hybrid GDN/full attention.

Qwen3.5 uses a hybrid attention pattern: 75% GDN (linear attention) layers and 25% standard GQA (full attention) layers (full_attention_interval=4). Supports both the MoE variant (Qwen3.5-35B-A3B) and Dense variant (Qwen3.5-27B).

nemo_automodel.components.utils.flops_utils.qwen3_8_flash_next_flops(
    config: typing.Any,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for Qwen3.8-Flash-Next (GDN + compressed-block QSA, HyperConnections, Engram PLE, MoE).

Accepts Qwen3_8_FlashNextTextConfig or the multimodal Qwen3_8_FlashNextConfig wrapper (its text_config is used; the vision tower is not loaded by the training path). The module shapes follow nemo_automodel/components/models/qwen3_8_flash_next:

  • GatedDeltaNet layers (layers_block_type == "linear_attention"): the shared _gdn_attention_per_layer_flops term (QKV/Z/B/A projections, causal conv, chunked delta-rule recurrence, output projection);
  • QSA full-attention layers: gated q_proj (hidden x 2headshead_dim), k_proj/v_proj (hidden x kv_heads*head_dim) and o_proj; the sparse GQA BMMs where the query at zero-based position t attends to ratio * min(indexer_budget / ratio, floor((t+1)/ratio)) routed tokens plus the (t+1) mod ratio tokens of its incomplete causal tail, QK^T and PV each costing heads * head_dim MACs per key;
  • the QSA indexer is frozen (requires_grad_(False); hard top-k has no gradient), so its index_qk_proj and its causal scoring of indexer_n_heads queries against the floor((t+1)/ratio) compressed keys are counted forward only (2 x MACs);
  • MoE per layer: router GEMM (hidden x num_experts), num_experts_per_tok routed plus one shared expert of 3 * hidden * moe_intermediate_size / shared_expert_intermediate_size, and the shared-expert gate (hidden x 1);
  • HyperConnections: two mixers per layer (attention and MoE), each with input_mix_weight_down (hc_counthidden x hc_lowrank), input_mix_weight_up (hc_lowrank x hc_counthidden) and block_inject_weight (hc_count*hidden x hc_count); the final read mixer has no inject weight;
  • Engram PLE on ple_layer_ids: key_proj (ple_embed_dim x hc_counthidden), value_proj (ple_embed_dim x hidden) and the depthwise causal convolution (hc_counthidden x kernel); the table gather, hashing and norms are not GEMMs;
  • the untied language-model head (hidden x vocab). The checkpoint’s MTP head is not loaded.

Trainable GEMMs and attention BMMs cost 6 x MACs (forward + backward); activation recomputation is excluded (that is HFU). On the released configuration the formula counts 6.14B active GEMM parameters per token excluding the LM head (model card: 6B activated), implies a 125.8B backbone excluding the 51.2B Engram table (model card: 125B), and gives 42.0 GFLOPs/token at 4096 tokens, 1.034x the dense identity 6 x active.

Parameters:

config
Any

Qwen3_8_FlashNextTextConfig, or a Qwen3_8_FlashNextConfig wrapper whose text_config is used.

gbs
intDefaults to 1

Number of sequences per step.

seq_len
int | NoneDefaults to None

Tokens per sequence; defaults to config.max_position_embeddings.

Returns: float

Model FLOPs for one training step of gbs sequences of seq_len tokens.

nemo_automodel.components.utils.flops_utils.qwen3_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig | nemo_automodel.components.utils.flops_utils._CompositeConfig[nemo_automodel.components.utils.flops_utils._DenseConfig],
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for Qwen3 family - accepts either AutoConfig or normalized config

nemo_automodel.components.utils.flops_utils.step3_5_flash_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Model FLOPs for Step3.5-Flash (GQA + sliding-window / full attention + MoE).

Architecture: hybrid full/SWA attention with different head counts per type, MoE with shared expert on most layers, first few layers dense, SwiGLU.

nemo_automodel.components.utils.flops_utils.transformer_flops(
    config: nemo_automodel.components.utils.flops_utils._DenseConfig,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Calculate FLOPs for a standard Transformer model - accepts either AutoConfig or normalized config. Note: This does not cover encoder-decoder models.

nemo_automodel.components.utils.flops_utils._ConfigT_co = TypeVar('_ConfigT_co', covariant=True)