`nemo_rl.utils.flops_formulas`#

Module Contents#

Classes#

FLOPSConfig

Contains the model hparams needed for FLOPS computations.

Functions#

`gpt3`	Model FLOPs for GPT3 family.
`llama2`	Model FLOPs for llama2 family.
`llama3`	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.
`deepseekv3`	Model FLOPs for DeepSeek V3.
`_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

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

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

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

nemo_rl.utils.flops_formulas.llama3(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.deepseekv3(config: nemo_rl.utils.flops_formulas.FLOPSConfig)[source]#: Model FLOPs for DeepSeek V3.

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`#