nemo_automodel.components.utils.flops_utils
nemo_automodel.components.utils.flops_utils
Module Contents
Functions
API
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.
FLOPs for a single Gated DeltaNet (GDN / linear attention) layer.
Based on the GDN FLOPs calculator from Megatron-Bridge PR #2925.
Model FLOPs for hybrid model
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.
Model FLOPs for Mamba layer.
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).
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
FLOPs for MLA + MoE transformer models (DeepSeek-V3 style).
Parameters:
List of 0/1 per layer (0=dense, 1=MoE).
Total intermediate size for all shared experts combined.
If set, use DSA sparse attention with this many selected positions.
Number of heads in the DSA indexer.
Head dimension of the DSA indexer.
Model FLOPs for MLP layer. Assume gated linear unit.
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.
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:
effective MTP depths actually run (model.mtp_config.num_layers).
block types (“mamba”/“attention”/“mlp”/“moe”) of the physical MTP sublayers (model.mtp.layers).
True if the physical depth is reused across the N depths.
Model FLOPs for attention layer
Calculate the flops for the attention part.
Model FLOPs for BERT family - accepts either AutoConfig or normalized config
Calculate Model FLOPs Utilization (MFU).
Parameters:
Total TFLOPs across all devices for the measured step.
Total number of GPUs.
Time taken for computation.
Peak TFLOPs/s per device for the training precision. The legacy default is the H100 dense-FP8 peak.
Returns: float
MFU as a percentage.
Model FLOPs for CLIP ViT
Model FLOPs for DeepSeek V3 - accepts either AutoConfig or normalized config
Model FLOPs for FLUX
Get the appropriate FLOPs formula function for a given HuggingFace config.
Parameters:
HuggingFace model config object
Returns: Callable | None
The appropriate FLOPs formula function, or None if model type is not supported
Estimate FLOPs for GLM4 MoE model configurations.
Model FLOPs for GPT3 family - accepts either AutoConfig or normalized config
Model FLOPs for GPT-OSS
Calculate the flops for the GPT-OSS model
Model FLOPs for Kimi K3 (hybrid KDA / gated-MLA attention + latent MoE with a SiTU shared expert).
Accepts KimiK3TextConfig or the multimodal KimiK3Config wrapper (its
text_config is used; the vision tower is not counted, as for other VL entries).
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.
Model FLOPs for llama2 family - accepts either AutoConfig or normalized config
Model FLOPs for llama3 family - accepts either AutoConfig or normalized config
Calculate the flops for the loss
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.
Model FLOPs for mixtral family - accepts either AutoConfig or normalized config
Model FLOPs for MLA + MoE models (Kimi K2, GLM-5, Mistral Small 4, etc.).
Handles VL wrappers by extracting text_config if present.
Calculate the flops for the MLP
Model FLOPs for nemotron family - accepts either AutoConfig or normalized config
Model FLOPs for NemotronH
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).
Model FLOPs for Qwen3 family - accepts either AutoConfig or normalized config
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.
Calculate FLOPs for a standard Transformer model - accepts either AutoConfig or normalized config. Note: This does not cover encoder-decoder models.