`nemo_rl.utils.flops_formulas`#

Module Contents#

Classes#

FLOPSConfig

Contains the model hparams needed for FLOPS computations.

Functions#

`gpt3`	Model FLOPs for GPT3 family.
`llama`	Model FLOPs for llama3 family.
`nemotron`	Model FLOPs for nemotron family.
`mixtral`	Model FLOPs for mixtral family.
`qwen2`	Model FLOPs for Qwen2 family.
`qwen3`	Model FLOPs for Qwen3 family.
`bert`	Model FLOPs for BERT family.
`transformer`	Calculate FLOPs for a standard Transformer model.
`flux`	Model FLOPs for FLUX.
`_mla_projection_params`	Return the per-layer MLA projection parameter count.
`_moe_ffn_params`	Return activated dense/MoE FFN parameters for all model and MTP layers.
`_mla_moe_linear_flops`	Return MLA projection and activated FFN training FLOPs per input.
`_vocab_and_mtp_flops`	Return output projection and MTP embedding projection FLOPs per input.
`deepseekv3`	Model FLOPs for DeepSeek V3.
`_causal_topk_pairs`	Approximate causal sparse-attention query/key pairs with dense formula conventions.
`glm_moe_dsa`	Model FLOPs for GLM MoE DSA.
`_mlp_layer_flops`	Model FLOPs for MLP layer.
`_non_mla_attn_layer_flops`	Model FLOPs for attention layer.
`_mamba_layer_flops`	Model FLOPs for Mamba layer. We ignore part of the flops of scan because the chunk size is not known from model config.
`_hybrid_model_flops`	Model FLOPs for hybrid model.
`nemotronh`	Model FLOPs for NemotronH.

API#

class nemo_rl.utils.flops_formulas.FLOPSConfig[source]#

Contains the model hparams needed for FLOPS computations.

gbs: int#: None

enc_seq_len: Optional[int]#: None

hs: Optional[int]#: None

layers: Optional[int]#: None

ffn_hs: Optional[int]#: None

attention_heads: Optional[int]#: None

head_dim: Optional[int]#: None

moe_router_topk: Optional[int]#: None

query_groups: Optional[int]#: None

img_seq_len: Optional[int]#: None

img_h: Optional[int]#: None

img_w: Optional[int]#: None

in_channels: Optional[int]#: None

patch_dim: Optional[int]#: None

class_token_len: Optional[int]#: None

projector_type: Optional[str]#: None

inp_s: Optional[int]#: None

model_pattern: Optional[str]#: None

vocab_size: Optional[int]#: None

model_channels: Optional[int]#: None

vec_in_dim: Optional[int]#: None

q_lora_rank: Optional[int]#: None

kv_lora_rank: Optional[int]#: None

qk_head_dim: Optional[int]#: None

qk_pos_emb_head_dim: Optional[int]#: None

v_head_dim: Optional[int]#: None

moe_layer_freq: Optional[Union[int, List[int]]]#: None

moe_shared_expert_intermediate_size: Optional[int]#: None

moe_ffn_hidden_size: Optional[int]#: None

mtp_num_layers: Optional[int]#: None

causal_self_attn: Optional[bool]#: None

is_hybrid_model: bool#: False

hybrid_override_pattern: Optional[str]#: None

mamba_state_dim: Optional[int]#: None

mamba_head_dim: Optional[int]#: None

mamba_num_groups: Optional[int]#: None

mamba_num_heads: Optional[int]#: None

dsa_indexer_n_heads: Optional[int]#: None

dsa_indexer_head_dim: Optional[int]#: None

dsa_indexer_topk: Optional[int]#: None

dsa_indexer_compute_layers: Optional[int]#: None

nemo_rl.utils.flops_formulas.gpt3(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for GPT3 family.

nemo_rl.utils.flops_formulas.llama(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for llama3 family.

nemo_rl.utils.flops_formulas.nemotron(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for nemotron family.

nemo_rl.utils.flops_formulas.mixtral(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for mixtral family.

nemo_rl.utils.flops_formulas.qwen2(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for Qwen2 family.

nemo_rl.utils.flops_formulas.qwen3(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for Qwen3 family.

nemo_rl.utils.flops_formulas.bert(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for BERT family.

nemo_rl.utils.flops_formulas.transformer(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#

Calculate FLOPs for a standard Transformer model.

Note: This does not cover encoder-decoder models.

nemo_rl.utils.flops_formulas.flux(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for FLUX.

nemo_rl.utils.flops_formulas._mla_projection_params( config: nemo_rl.utils.flops_formulas.FLOPSConfig, ) → int[source]#: Return the per-layer MLA projection parameter count.

nemo_rl.utils.flops_formulas._moe_ffn_params( config: nemo_rl.utils.flops_formulas.FLOPSConfig, ) → int[source]#: Return activated dense/MoE FFN parameters for all model and MTP layers.

nemo_rl.utils.flops_formulas._mla_moe_linear_flops( config: nemo_rl.utils.flops_formulas.FLOPSConfig, ) → int[source]#: Return MLA projection and activated FFN training FLOPs per input.

nemo_rl.utils.flops_formulas._vocab_and_mtp_flops( config: nemo_rl.utils.flops_formulas.FLOPSConfig, ) → int[source]#: Return output projection and MTP embedding projection FLOPs per input.

nemo_rl.utils.flops_formulas.deepseekv3(config: nemo_rl.utils.flops_formulas.FLOPSConfig) → float[source]#: Model FLOPs for DeepSeek V3.

nemo_rl.utils.flops_formulas._causal_topk_pairs(seq_len: int, topk: int) → float[source]#: Approximate causal sparse-attention query/key pairs with dense formula conventions.

nemo_rl.utils.flops_formulas.glm_moe_dsa(config: nemo_rl.utils.flops_formulas.FLOPSConfig) → float[source]#

Model FLOPs for GLM MoE DSA.

GLM-MoE-DSA uses MLA projections and MoE FFNs like DeepSeek-style MoE models, but its core attention path uses a DSA indexer to select sparse top-k keys. The indexer is forward-only in the RL runs where indexer loss is disabled, so its matmuls are counted with a single 2x FMA factor rather than the 6x forward+backward training factor used for trainable model-path GEMMs.

nemo_rl.utils.flops_formulas._mlp_layer_flops(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for MLP layer.

nemo_rl.utils.flops_formulas._non_mla_attn_layer_flops( config: nemo_rl.utils.flops_formulas.FLOPSConfig, )[source]#: Model FLOPs for attention layer.

nemo_rl.utils.flops_formulas._mamba_layer_flops(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for Mamba layer. We ignore part of the flops of scan because the chunk size is not known from model config.

nemo_rl.utils.flops_formulas._hybrid_model_flops(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for hybrid model.

nemo_rl.utils.flops_formulas.nemotronh(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for NemotronH.

nemo_rl.utils.flops_formulas#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.utils.flops_formulas`#