> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.training.domain_mixture

Objective weighting for named pretraining data mixtures.

## Module Contents

### Classes

| Name                                                                                            | Description                                                      |
| ----------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |
| [`DomainMixture`](#nemo_automodel-components-training-domain_mixture-DomainMixture)             | Runtime importance weights and metrics for a named data mixture. |
| [`DomainMixtureConfig`](#nemo_automodel-components-training-domain_mixture-DomainMixtureConfig) | Configuration for importance-weighted multi-domain pretraining.  |
| [`DomainWeightConfig`](#nemo_automodel-components-training-domain_mixture-DomainWeightConfig)   | Declarative sampling and objective weights for one data domain.  |

### Data

[`WEIGHTED_AGGREGATE_NAME`](#nemo_automodel-components-training-domain_mixture-WEIGHTED_AGGREGATE_NAME)

### API

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

---

```python
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`** `bool` — default: 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.

---

```python
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.

```python
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`** `int` — default: -100

Label value excluded from the count.

---

**Returns:** `torch.Tensor`

Int64 tensor of shape `[num_domains]`.

```python
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`** `bool` — default: 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

```python
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.

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

---

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

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

Build the immutable runtime domain-mixture objective.

```python
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`

---

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

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