nemo_rl.utils.flops_tracker

Contenu du module

Classes

Nom	Description
`FLOPTracker`	Aucun

Fonctions

Nom	Description
`get_default_hf_config`	Obtenir la configuration Hugging Face par défaut pour un modèle.
`convert_config_to_flops_config`	Convertir une configuration préentraînée en un tuple contenant une configuration FLOPSConfig et une formule de FLOPS.
`get_theoretical_tflops`	Obtenir le nombre total théorique de FLOPS pour un nom de périphérique.

Données

THEORETICAL_TFLOPS

API

nemo_rl.utils.flops_tracker.get_default_hf_config(model_name: str) -> transformers.configuration_utils.PretrainedConfig

Obtenir la configuration Hugging Face par défaut pour un modèle.

Les chemins DTensor et MCore utilisent la même configuration par défaut, nous initialisons ici la configuration du modèle pour permettre le calcul des FLOPS théoriques, indépendamment du backend.

nemo_rl.utils.flops_tracker.convert_config_to_flops_config(config: transformers.configuration_utils.PretrainedConfig) -> tuple[nemo_rl.utils.flops_formulas.FLOPSConfig, typing.Callable]

Convertir une configuration préentraînée en un tuple contenant une configuration FLOPSConfig et une formule de FLOPS.

nemo_rl.utils.flops_tracker.THEORETICAL_TFLOPS

Valeur: None

nemo_rl.utils.flops_tracker.get_theoretical_tflops(
    device_name: str, model_dtype: torch.dtype
) -> float

Obtenir le nombre total théorique de FLOPS pour un nom de périphérique.

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

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])

Suivre les FLOPS pour un lot de séquences.

reset()