nemo_automodel.components.datasets.llm.megatron.megatron_utils

Module Contents

Functions

Name	Description
`compile_helper`	Load the CPU dataset helpers from a locked, node-local build cache.
`get_blend_from_list`	Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list

Data

_HELPER_CACHE_ROOT

API

nemo_automodel.components.datasets.llm.megatron.megatron_utils.compile_helper() -> types.ModuleType

Load the CPU dataset helpers from a locked, node-local build cache.

Requires a C++ compiler and Ninja. Sources and build artifacts live in a per-user directory under /tmp, independent of TORCH_EXTENSIONS_DIR and TMPDIR. Each source and Python/PyTorch environment has its own cache entry. An exclusive file lock covers source generation, compilation, and import, so ranks on the same node can call this without a distributed barrier. The loaded module is retained for the lifetime of this process.

Returns: ModuleType

The compiled module containing the dataset indexing functions.

nemo_automodel.components.datasets.llm.megatron.megatron_utils.get_blend_from_list(
    blend: typing.List[str] | None
) -> typing.Tuple[typing.List[str], typing.List[float] | None] | None

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: Tuple[List[str], List[float] | None] | None

Optional[Tuple[List[str], Optional[List[float]]]]: 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]].

nemo_automodel.components.datasets.llm.megatron.megatron_utils._HELPER_CACHE_ROOT = Path('/tmp') / f'nemo_automodel_{os.getuid()}'