`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`
`_energon_fanout_attributes`	Per-rank Energon worker count and queue depth; both 0 when num_workers=0.
`_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`
`ENERGON_WORKERS_PER_RANK_ATTR`
`ENERGON_IN_FLIGHT_ATTR`
`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.ENERGON_WORKERS_PER_RANK_ATTR#: ‘rl.energon.workers_per_rank’

nemo_rl.algorithms.sft_v2.ENERGON_IN_FLIGHT_ATTR#: ‘rl.energon.batches_in_flight_per_rank’

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

setup_timing: nemo_rl.algorithms.metric_utils.SetupTimingMetrics#: 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._energon_fanout_attributes( data_config: Mapping[str, Any], ) → dict[str, int]#: Per-rank Energon worker count and queue depth; both 0 when num_workers=0.

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, trace_carrier: Optional[dict[str, str]] = None, )#

Drive colocated loaders and the existing TQPolicy from one actor.

Initialization

_defer_checkpoint_finalize: bool#: False

_pending_checkpoint_path: Optional[str]#: None

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.

_log_setup_timing( setup_timing: nemo_rl.algorithms.metric_utils.SetupTimingMetrics, ) → None#: Log the setup phases; total_setup is the one that cannot double-count.

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

_finish_pending_save(*, blocking: bool) → None#

Finalize and publish an async save deferred by defer_async_checkpoint_finalize.

Runs between steps, and once more at teardown so a run that raises after a save still publishes it. A non-blocking poll finalizes the save on every policy worker only once all ranks’ background writes are done, so the write overlaps training rather than holding back the next step. Finalization runs collectives on the training communicator, so it stays on the workers’ step boundaries rather than in a background thread. tmp_step_N becomes step_N only after it.

_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#