nemo_automodel.components.training.domain_mixture

Objective weighting for named pretraining data mixtures.

Module Contents

Classes

Name	Description
`DomainMixture`	Runtime importance weights and metrics for a named data mixture.
`DomainMixtureConfig`	Configuration for importance-weighted multi-domain pretraining.
`DomainWeightConfig`	Declarative sampling and objective weights for one data domain.

Data

WEIGHTED_AGGREGATE_NAME

API

class nemo_automodel.components.training.domain_mixture.DomainMixture(
    names: tuple[str, ...],
    sampling_weights: tuple[float, ...],
    objective_weights: tuple[float, ...],
    loss_multipliers: tuple[float, ...],
    _multiplier_cache: dict = dict()
)

Dataclass

Runtime importance weights and metrics for a named data mixture.

loss_multipliers[i] = objective_weights[i] / sampling_weights[i]. Applying that multiplier to every supervised token sampled from domain i makes the expected training objective match objective_weights without changing the sampler or creating separate optimizers.

_multiplier_cache

dict

loss_multipliers

tuple[float, ...]

names

tuple[str, ...]

objective_weights

tuple[float, ...]

sampling_weights

tuple[float, ...]

nemo_automodel.components.training.domain_mixture.DomainMixture._domain_ids(
    dataset_ids: torch.Tensor,
    validate_range: bool = True
) -> torch.Tensor

Validate and flatten per-sample dataset IDs to shape [batch].

Parameters:

dataset_ids

torch.Tensor

Integer tensor of per-sample domain IDs.

validate_range

boolDefaults to True

Bounds-check the IDs. This costs two blocking device-to-host syncs, so callers on the per-microbatch critical path pass False once a step-level pass has already validated the very same batches.

nemo_automodel.components.training.domain_mixture.DomainMixture._multiplier_tensor(
    device: torch.device
) -> torch.Tensor

Return the multiplier vector on device, building it at most once.

loss_multipliers is immutable after :meth:DomainMixtureConfig.build, so rebuilding it per microbatch only adds a pageable host-to-device copy on the critical path.

nemo_automodel.components.training.domain_mixture.DomainMixture.label_counts(
    dataset_ids: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int = -100
) -> torch.Tensor

Count supervised tokens per domain.

Parameters:

dataset_ids

torch.Tensor

Integer tensor with one domain ID per label row.

labels

torch.Tensor

Target token IDs of shape [batch, ...] or a flat single-sample tensor.

ignore_index

intDefaults to -100

Label value excluded from the count.

Returns: torch.Tensor

Int64 tensor of shape [num_domains].

nemo_automodel.components.training.domain_mixture.DomainMixture.loss_weights(
    dataset_ids: torch.Tensor,
    labels: torch.Tensor,
    validate_range: bool = True
) -> torch.Tensor

Expand per-sample importance weights to the label layout.

Parameters:

dataset_ids

torch.Tensor

Integer tensor of shape [batch] (or a scalar for a single sample). IDs follow the domain order in the config.

labels

torch.Tensor

Target token IDs of shape [batch, sequence] or [sequence] for a single flattened sample.

validate_range

boolDefaults to True

Bounds-check dataset_ids. Costs two device syncs; see :meth:_domain_ids.

Returns: torch.Tensor

Float32 tensor matching labels.shape. Every row is constant

nemo_automodel.components.training.domain_mixture.DomainMixture.weighted_validation_loss(
    losses: collections.abc.Mapping[str, float]
) -> float

Combine named per-domain validation losses using objective weights.

class nemo_automodel.components.training.domain_mixture.DomainMixtureConfig(
    domains: tuple[nemo_automodel.components.training.domain_mixture.DomainWeightConfig, ...]
)

Dataclass

Configuration for importance-weighted multi-domain pretraining.

domains

tuple[DomainWeightConfig, ...]

nemo_automodel.components.training.domain_mixture.DomainMixtureConfig.__post_init__() -> None

nemo_automodel.components.training.domain_mixture.DomainMixtureConfig.build() -> 'DomainMixture'

Build the immutable runtime domain-mixture objective.

class nemo_automodel.components.training.domain_mixture.DomainWeightConfig(
    name: str,
    sampling_weight: float,
    objective_weight: float
)

Dataclass

Declarative sampling and objective weights for one data domain.

name

str

objective_weight

float

sampling_weight

float

nemo_automodel.components.training.domain_mixture.DomainWeightConfig.__post_init__() -> None

nemo_automodel.components.training.domain_mixture.WEIGHTED_AGGREGATE_NAME = 'weighted'