`nemo_automodel.components.utils.model_utils`#

Module Contents#

Functions#

`_supports_logits_to_keep`	Check if the model supports logits_to_keep.
`_supports_seq_lens`	Check if the model supports seq_lens.
`_get_model_param_stats`	Get the number of trainable parameters and the L2 norm of the model.
`print_trainable_parameters`	Print the number of trainable parameters in the model.
`_freeze_module_by_attribute_and_patterns`	Helper function to freeze parameters by attribute name and name patterns.
`apply_parameter_freezing`	Apply parameter freezing based on configuration.

Data#

logger

API#

nemo_automodel.components.utils.model_utils.logger#: ‘getLogger(…)’

nemo_automodel.components.utils.model_utils._supports_logits_to_keep(model: torch.nn.Module) → bool#

Check if the model supports logits_to_keep.

Parameters:: model (nn.Module) – The model to check.
Returns:: True if the model supports logits_to_keep, False otherwise.
Return type:: bool

nemo_automodel.components.utils.model_utils._supports_seq_lens(model: torch.nn.Module) → bool#: Check if the model supports seq_lens.

nemo_automodel.components.utils.model_utils._get_model_param_stats( model: torch.nn.Module, ) → tuple[int, int, float]#

Get the number of trainable parameters and the L2 norm of the model.

Parameters:: model – Model to analyze
Returns:: int trainable_params: int local_sq_norm: float
Return type:: total_params

nemo_automodel.components.utils.model_utils.print_trainable_parameters(model: torch.nn.Module) → tuple[int, int]#

Print the number of trainable parameters in the model.

Parameters:: model – Model to analyze
Returns:: int total_params: int
Return type:: trainable_params

nemo_automodel.components.utils.model_utils._freeze_module_by_attribute_and_patterns( model, attribute_name, name_patterns, )#

Helper function to freeze parameters by attribute name and name patterns.

Parameters:

model – The model to apply freezing to.
attribute_name – Name of the model attribute to freeze (e.g., ‘vision_tower’).
name_patterns – List of patterns to match in module names.

nemo_automodel.components.utils.model_utils.apply_parameter_freezing(model, freeze_config)#

Apply parameter freezing based on configuration.