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(),
    paths: tuple[str | None, ...] = (),
    blend: dataclasses.InitVar[tuple[collections.abc.Sequence[str], collections.abc.Sequence[float]] | None] = None
)

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, ...]

paths

tuple[str | None, ...] = field(default=(), kw_only=True)

sampling_weights

tuple[float, ...]

nemo_automodel.components.training.domain_mixture.DomainMixture.__post_init__(
    blend: tuple[collections.abc.Sequence[str], collections.abc.Sequence[float]] | None
) -> None

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 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 _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(
    blend: tuple[collections.abc.Sequence[str], collections.abc.Sequence[float]] | None = None
) -> 'DomainMixture'

Build the immutable runtime domain-mixture objective.

Parameters:

blend

tuple[Sequence[str], Sequence[float]] | NoneDefaults to None

Parsed training corpus prefixes and explicit sampling weights, in dataset ID order. When provided, construction validates the declared domain bindings against this blend. Omit for standalone objective calculations.

Returns: 'DomainMixture'

The domain-mixture objective with validated bindings when supplied.

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

Dataclass

Declarative sampling and objective weights for one data domain.

name

str

Domain name. The position in DomainMixtureConfig.domains must match the dataset_id assigned by the blended dataset.

objective_weight

float

Relative contribution of the domain to the training and validation objectives.

path

str | None = None

Explicit Megatron corpus prefix for this domain. Required by the training recipe to verify blend order; optional for standalone objective calculations. Names are never inferred from this path.

sampling_weight

float

Relative probability used to sample the domain.

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

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