nemo_automodel.components.datasets.llm.megatron.megatron_utils#

Module Contents#

Functions#

get_blend_from_list

Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list

compile_helper

Compile helper function ar runtime. Make sure this is invoked on a single process.

Data#

API#

nemo_automodel.components.datasets.llm.megatron.megatron_utils.logger#

‘getLogger(…)’

nemo_automodel.components.datasets.llm.megatron.megatron_utils.get_blend_from_list(
blend: Optional[List[str]],
) Optional[Tuple[List[str], Optional[List[float]]]]#

Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list

Parameters:

blend (Optional[List[str]]) – The blend list, which can be either (1) a list of prefixes, e.g. [“path/to/dataset_1_prefix”, “path/to/dataset_2_prefix”], or (2) a flattened, zipped list of weights and prefixes, e.g. [“30”, “path/to/dataset_1_prefix”, “70”, “path/to/dataset_2_prefix”]

Returns:

The blend, consisting of a list of dataset prefixes and optionally a list of dataset weights, e.g. [[“path/to/dataset_1_prefix”, “path/to/dataset_2_prefix”], [30.0, 70.0]].

Return type:

Optional[Tuple[List[str], Optional[List[float]]]]

nemo_automodel.components.datasets.llm.megatron.megatron_utils.compile_helper()#

Compile helper function ar runtime. Make sure this is invoked on a single process.