nemo_rl.algorithms.mpo#

Mixed Preference Optimization built on the maintained preference trainer.

Module Contents#

Classes#

MPOConfig

MasterConfig

MPOSaveState

MPOValMetrics

Functions#

_initial_mpo_save_state

_validate_pair_safe_packing

_make_loss_fn

_sum_metric

_update_reward_shift

Synchronize the BCO shift once per completed optimizer step.

setup

Set up MPO without reintroducing the legacy Omni collapse/expand path.

mpo_train

Run MPO with driver-owned, checkpointed reward-shift updates.

API#

class nemo_rl.algorithms.mpo.MPOConfig#

Bases: nemo_rl.algorithms.dpo.DPOConfig

bco_loss_weight: float#

1.0

quality_average_log_probs: bool#

False

reward_shift_momentum: float#

0.99

reward_shift: float#

0.0

class nemo_rl.algorithms.mpo.MasterConfig#

Bases: pydantic.BaseModel

policy: nemo_rl.models.policy.PolicyConfig#

None

data: nemo_rl.data.DataConfig#

None

mpo: nemo_rl.algorithms.mpo.MPOConfig#

None

logger: nemo_rl.utils.logger.LoggerConfig#

None

cluster: nemo_rl.distributed.virtual_cluster.ClusterConfig#

None

checkpointing: nemo_rl.utils.checkpoint.CheckpointingConfig#

None

property dpo: nemo_rl.algorithms.mpo.MPOConfig#

Compatibility view consumed by the shared preference trainer.

class nemo_rl.algorithms.mpo.MPOSaveState#

Bases: nemo_rl.algorithms.dpo.DPOSaveState

reward_shift: float#

0.0

reward_shift_num_updates: int#

0

class nemo_rl.algorithms.mpo.MPOValMetrics#

Bases: nemo_rl.algorithms.dpo.DPOValMetrics

bco_loss: float#

None

bco_rewards_chosen_mean: float#

None

bco_rewards_rejected_mean: float#

None

nemo_rl.algorithms.mpo._initial_mpo_save_state(
config: nemo_rl.algorithms.mpo.MPOConfig,
) → nemo_rl.algorithms.mpo.MPOSaveState#
nemo_rl.algorithms.mpo._validate_pair_safe_packing(
policy_config: nemo_rl.models.policy.PolicyConfig,
cluster_config: nemo_rl.distributed.virtual_cluster.ClusterConfig,
) → None#
nemo_rl.algorithms.mpo._make_loss_fn(
config: nemo_rl.algorithms.mpo.MPOConfig,
policy_config: nemo_rl.models.policy.PolicyConfig,
save_state: nemo_rl.algorithms.mpo.MPOSaveState,
) → nemo_rl.algorithms.loss.MPOLossFn#
nemo_rl.algorithms.mpo._sum_metric(values: list[Any]) → float#
nemo_rl.algorithms.mpo._update_reward_shift(
train_results: dict[str, Any],
loss_fn: nemo_rl.algorithms.loss.MPOLossFn,
save_state: nemo_rl.algorithms.mpo.MPOSaveState,
) → None#

Synchronize the BCO shift once per completed optimizer step.

nemo_rl.algorithms.mpo.setup(
master_config: nemo_rl.algorithms.mpo.MasterConfig,
tokenizer: transformers.AutoTokenizer,
train_dataset: nemo_rl.data.datasets.AllTaskProcessedDataset,
val_dataset: dict[str, nemo_rl.data.datasets.AllTaskProcessedDataset],
) → tuple[nemo_rl.models.policy.lm_policy.Policy, nemo_rl.distributed.virtual_cluster.RayVirtualCluster, torchdata.stateful_dataloader.StatefulDataLoader, dict[str, torchdata.stateful_dataloader.StatefulDataLoader], nemo_rl.algorithms.loss.MPOLossFn, nemo_rl.utils.logger.Logger, nemo_rl.utils.checkpoint.CheckpointManager, nemo_rl.algorithms.mpo.MPOSaveState, nemo_rl.algorithms.mpo.MasterConfig]#

Set up MPO without reintroducing the legacy Omni collapse/expand path.

nemo_rl.algorithms.mpo.mpo_train(
policy: nemo_rl.models.policy.lm_policy.Policy,
train_dataloader: torchdata.stateful_dataloader.StatefulDataLoader,
val_dataloader: dict[str, torchdata.stateful_dataloader.StatefulDataLoader],
tokenizer: transformers.AutoTokenizer,
loss_fn: nemo_rl.algorithms.loss.MPOLossFn,
master_config: nemo_rl.algorithms.mpo.MasterConfig,
logger: nemo_rl.utils.logger.Logger,
checkpointer: nemo_rl.utils.checkpoint.CheckpointManager,
mpo_save_state: nemo_rl.algorithms.mpo.MPOSaveState,
) → None#

Run MPO with driver-owned, checkpointed reward-shift updates.