`nemo_rl.algorithms.sft_v2`#

Single-controller SFT with one colocated Energon loader per DP replica.

Module Contents#

Classes#

`MasterConfig`	Standalone SFTv2 configuration.
`SFTV2SaveState`	Controller state committed at an optimizer-step boundary.
`SFTV2ActorArgs`	Driver-built objects used by the SFTv2 controller actor.
`SFTSingleControllerActor`	Drive colocated loaders and the existing TQPolicy from one actor.

Functions#

`_initial_save_state`
`_restore_save_state`
`_max_train_steps`
`setup_sft_v2`	Build the V2 cluster, TQPolicy, placement, and resume state.

Data#

`LOADER_LATENCY_MAX_KEY`
`POLICY_TIME_KEY`
`TOTAL_STEP_TIME_KEY`
`VALID_TOKENS_PER_SECOND_KEY`
`SFT_V2_TEED_METRICS`
`__all__`

API#

nemo_rl.algorithms.sft_v2.LOADER_LATENCY_MAX_KEY#: ‘loader_latency_max’

nemo_rl.algorithms.sft_v2.POLICY_TIME_KEY#: ‘policy_time’

nemo_rl.algorithms.sft_v2.TOTAL_STEP_TIME_KEY#: ‘total_step_time’

nemo_rl.algorithms.sft_v2.VALID_TOKENS_PER_SECOND_KEY#: ‘valid_tokens_per_second’

nemo_rl.algorithms.sft_v2.SFT_V2_TEED_METRICS#: ()

class nemo_rl.algorithms.sft_v2.MasterConfig#

Bases: pydantic.BaseModel

Standalone SFTv2 configuration.

policy: nemo_rl.models.policy.PolicyConfig#: None

data: nemo_rl.data.DataConfig#: None

sft: nemo_rl.algorithms.sft.SFTConfig#: None

data_plane: nemo_rl.data_plane.interfaces.LocalDataPlaneConfig#: None

logger: nemo_rl.utils.logger.LoggerConfig#: None

cluster: nemo_rl.distributed.virtual_cluster.ClusterConfig#: None

checkpointing: nemo_rl.utils.checkpoint.CheckpointingConfig#: None

telemetry: Optional[nemo_rl.telemetry.config.TelemetryConfig]#: None

class nemo_rl.algorithms.sft_v2.SFTV2SaveState#

Controller state committed at an optimizer-step boundary.

total_steps: int#: None

consumed_samples: int#: None

total_valid_tokens: int#: None

placement_hash: str#: None

class nemo_rl.algorithms.sft_v2.SFTV2ActorArgs#

Driver-built objects used by the SFTv2 controller actor.

trainer: nemo_rl.models.policy.tq_policy.TQPolicy#: None

loss_fn: nemo_rl.algorithms.loss.loss_functions.NLLLossFn#: None

train_cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster#: None

placement_plan: nemo_rl.data.energon.topology.DataLoaderPlacementPlan#: None

save_state: nemo_rl.algorithms.sft_v2.SFTV2SaveState#: None

loader_states: list[dict[str, Any]] | None#: None

nemo_rl.algorithms.sft_v2._initial_save_state( placement_hash: str, ) → nemo_rl.algorithms.sft_v2.SFTV2SaveState#

nemo_rl.algorithms.sft_v2._restore_save_state( loaded: Optional[dict[str, Any]], *, placement_hash: str, ) → nemo_rl.algorithms.sft_v2.SFTV2SaveState#

nemo_rl.algorithms.sft_v2._max_train_steps( master_config: nemo_rl.algorithms.sft_v2.MasterConfig, ) → int#

class nemo_rl.algorithms.sft_v2.SFTSingleControllerActor( master_config: nemo_rl.algorithms.sft_v2.MasterConfig, actor_args: nemo_rl.algorithms.sft_v2.SFTV2ActorArgs, )#

Drive colocated loaders and the existing TQPolicy from one actor.

Initialization

run_with_trace_context() → dict[str, Any]#: Run under the driver’s trace context when called through Ray.

run() → dict[str, Any]#: Run SFT training.

_setup_loaders() → None#

_load_envelopes() → list[nemo_rl.data.energon.sft_types.StepEnvelope]#

_run_train_step() → dict[str, Any]#

static _policy_metrics( train_results: dict[str, Any], ) → dict[str, Any]#: Convert policy output into flat scalar logger metrics.

_owner_call(method_name: str) → list[Any]#

_loader_state_dicts() → list[dict[str, Any]]#

_should_save(*, save_by_timeout: bool) → bool#

_checkpoint_metric(metrics: dict[str, Any]) → dict[str, float]#

Read checkpointing.metric_name out of one step’s training metrics.

Called every step rather than only on a save so a name that no step produces fails on step 1 instead of at the first checkpoint.

_save_checkpoint(metric: dict[str, float]) → None#

_close_loaders() → None#

nemo_rl.algorithms.sft_v2.setup_sft_v2( master_config: nemo_rl.algorithms.sft_v2.MasterConfig, tokenizer_or_processor: Any, ) → nemo_rl.algorithms.sft_v2.SFTV2ActorArgs#: Build the V2 cluster, TQPolicy, placement, and resume state.

nemo_rl.algorithms.sft_v2.__all__#: [‘MasterConfig’, ‘SFTSingleControllerActor’, ‘SFTV2ActorArgs’, ‘SFTV2SaveState’, ‘setup_sft_v2’]

nemo_rl.algorithms.sft_v2#