bridge.recipes.utils.nemo_run_utils

Module Contents

Functions

get_partial_fn	Creates a run.Partial object for the given target function with configuration and forward step function.
prepare_config_for_nemo_run	Prepares a pure ConfigContainer instance for use with NeMo Run by patching fields that are not directly serializable by nemo_run, such as functools.partial objects and enum values that cause YAML serialization issues.
_fix_yaml_serialization_issues	Fix common YAML serialization issues in the configuration. This function handles specific known issues like enum objects that can’t be serialized to YAML. It modifies the config in-place to use serializable representations while preserving functionality.
_check_nemo_run_available	Check if nemo_run is available and raise helpful error if not.

Data

logger

API

bridge.recipes.utils.nemo_run_utils.logger: logging.Logger

‘getLogger(…)’

bridge.recipes.utils.nemo_run_utils.get_partial_fn(

target_fn: Callable[[megatron.bridge.training.config.ConfigContainer, Callable], None],

config: megatron.bridge.training.config.ConfigContainer,

forward_step_func: Callable,

) → nemo_run.Partial

Creates a run.Partial object for the given target function with configuration and forward step function.

This is a convenience function that combines the preparation of the config for NeMo Run and the creation of the run.Partial object for any target function that expects a ConfigContainer and forward_step_func.

Parameters:

• target_fn – The target function to be wrapped in run.Partial (e.g., pretrain).

• config – The ConfigContainer dataclass instance.

• forward_step_func – The forward step function to use for training.

Returns:

A run.Partial object ready to be executed with NeMo Run.

Raises:

ImportError – If nemo_run is not installed.

bridge.recipes.utils.nemo_run_utils.prepare_config_for_nemo_run(

config: megatron.bridge.training.config.ConfigContainer,

) → megatron.bridge.training.config.ConfigContainer

Prepares a pure ConfigContainer instance for use with NeMo Run by patching fields that are not directly serializable by nemo_run, such as functools.partial objects and enum values that cause YAML serialization issues.

Parameters:

config – The ConfigContainer dataclass instance.

Returns:

The patched ConfigContainer instance compatible for execution with nemo_run.

Raises:

ImportError – If nemo_run is not installed.

bridge.recipes.utils.nemo_run_utils._fix_yaml_serialization_issues(

config: megatron.bridge.training.config.ConfigContainer,

) → None

Fix common YAML serialization issues in the configuration. This function handles specific known issues like enum objects that can’t be serialized to YAML. It modifies the config in-place to use serializable representations while preserving functionality.

Parameters:

config – The ConfigContainer to fix

bridge.recipes.utils.nemo_run_utils._check_nemo_run_available() → None

Check if nemo_run is available and raise helpful error if not.