`nemo_rl.utils.flops_tracker`#

Module Contents#

Classes#

FLOPTracker

Functions#

`get_default_hf_config`	Get the default Hugging Face config for a model.
`convert_config_to_flops_config`	Convert a pretrained config to a tuple containing a FLOPSConfig and a flops formula.
`get_theoretical_tflops`	Get the theoretical total flops for a device name.

Data#

THEORETICAL_TFLOPS

API#

nemo_rl.utils.flops_tracker.get_default_hf_config( model_name: str, ) → transformers.configuration_utils.PretrainedConfig#

Get the default Hugging Face config for a model.

Both the DTensor and MCore paths use the same default config, we initialize the model config here to allow computation of theoretical flops which is agnostic to the backend.

nemo_rl.utils.flops_tracker.convert_config_to_flops_config( model_name: str, 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.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#