`nemo_rl.utils.flops_tracker`#

Module Contents#

Classes#

FLOPTracker

Functions#

`get_hf_config`	Get the effective Hugging Face config for a model.
`_get_glm_moe_layer_pattern`	Return the GLM dense/MoE layer pattern from the model config.
`_get_glm_index_compute_layers`	Return how many GLM layers compute, rather than reuse, DSA indices.
`convert_config_to_flops_config`	Convert a pretrained config to a tuple containing a FLOPSConfig and a flops formula.
`is_using_tf32`	Check if the current device is using TF32.
`get_theoretical_tflops`	Get the theoretical total flops for a device name.

Data#

THEORETICAL_TFLOPS

API#

nemo_rl.utils.flops_tracker.get_hf_config(

model_name: str,

**overrides: Any,

) → transformers.configuration_utils.PretrainedConfig#

Get the effective Hugging Face config for a model.

Both the DTensor and MCore paths use the same config, which allows backend-agnostic theoretical FLOPs computation. Policy overrides must be included so the tracker describes the model that is actually trained.

nemo_rl.utils.flops_tracker._get_glm_moe_layer_pattern( config: transformers.configuration_utils.PretrainedConfig, ) → list[int]#: Return the GLM dense/MoE layer pattern from the model config.

nemo_rl.utils.flops_tracker._get_glm_index_compute_layers( config: transformers.configuration_utils.PretrainedConfig, ) → int#: Return how many GLM layers compute, rather than reuse, DSA indices.

nemo_rl.utils.flops_tracker.convert_config_to_flops_config( config: transformers.configuration_utils.PretrainedConfig, ) → tuple[nemo_rl.utils.flops_formulas.FLOPSConfig, Callable]#: Convert a pretrained config to a tuple containing a FLOPSConfig and a flops formula.

nemo_rl.utils.flops_tracker.is_using_tf32() → bool#: Check if the current device is using TF32.

nemo_rl.utils.flops_tracker.THEORETICAL_TFLOPS#: None

nemo_rl.utils.flops_tracker.get_theoretical_tflops( device_name: str, model_dtype: torch.dtype, ) → float#: Get the theoretical total flops for a device name.

class nemo_rl.utils.flops_tracker.FLOPTracker( model_name: str, base_config: nemo_rl.utils.flops_formulas.FLOPSConfig | None = None, flops_formula: Callable[[nemo_rl.utils.flops_formulas.FLOPSConfig], float] | None = None, )#

Initialization

classmethod from_config( model_name: str, config: transformers.configuration_utils.PretrainedConfig, ) → nemo_rl.utils.flops_tracker.FLOPTracker#

track(n_samples: int, padded_seq_len: int)#

track_batch(sequence_lengths: list[int])#: Track the flops for a batch of sequences.

reset()#

nemo_rl.utils.flops_tracker#