`nemo_rl.data.energon.sft_dataloader`#

Module Contents#

Classes#

`SFTDataLoader`	Iterator and state methods consumed by the SFT algorithm.
`EnergonSFTDataLoader`	Expose Energon rank state through NeMo-RL’s dataloader interface.

Functions#

`compact_sample_error_handler`	Log a sample-processing error and let Energon request another sample.
`_identity_fingerprint`
`_brief`
`_identity_differences`	List the identity entries that changed, as `key old -> new` lines.
`_source_config`
`_loader_config`
`_v2_topology`	Describe the DP shard a V2 loader state belongs to.
`_loader_identity`	Describe what a restored loader must still agree with.
`_worker_config`
`_task_encoder`
`build_energon_sft_loader`	Build one loader for an explicit logical data shard and split.

Data#

`_V2_STATE_FORMAT_VERSION`
`__all__`

API#

nemo_rl.data.energon.sft_dataloader._V2_STATE_FORMAT_VERSION#: 2

nemo_rl.data.energon.sft_dataloader.compact_sample_error_handler( exception: Exception, sample: Any | list[Any], sources: list[megatron.energon.SourceInfo] | None = None, ) → None#: Log a sample-processing error and let Energon request another sample.

class nemo_rl.data.energon.sft_dataloader.SFTDataLoader#

Bases: typing.Protocol

Iterator and state methods consumed by the SFT algorithm.

__iter__() → Iterator[nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]]#

__len__() → int#

state_dict() → dict[str, Any]#

load_state_dict(state: dict[str, Any]) → None#

nemo_rl.data.energon.sft_dataloader._identity_fingerprint( identity: Mapping[str, Any], ) → str#

nemo_rl.data.energon.sft_dataloader._brief(value: Any) → str#

nemo_rl.data.energon.sft_dataloader._identity_differences( saved: Any, current: Any, path: str = '', ) → list[str]#: List the identity entries that changed, as key old -> new lines.

class nemo_rl.data.energon.sft_dataloader.EnergonSFTDataLoader( loader: Any, *, identity: Mapping[str, Any], )#

Expose Energon rank state through NeMo-RL’s dataloader interface.

Initialization

__iter__() → Iterator[nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]]#

__len__() → int#

state_dict() → dict[str, Any]#

load_state_dict(state: dict[str, Any]) → None#

nemo_rl.data.energon.sft_dataloader._source_config( value: Any, *, name: str, ) → nemo_rl.data.energon.config.EnergonSourceConfig#

nemo_rl.data.energon.sft_dataloader._loader_config( value: Any, ) → nemo_rl.data.energon.config.EnergonLoaderConfig#

nemo_rl.data.energon.sft_dataloader._v2_topology( *, loader_config: nemo_rl.data.energon.config.EnergonLoaderConfig, placement_fingerprint: str, logical_rank: int, logical_world_size: int, ) → dict[str, Any]#

Describe the DP shard a V2 loader state belongs to.

This is what makes one DP shard refuse another shard’s state, so it is hashed into the loader identity by :func:_loader_identity.

nemo_rl.data.energon.sft_dataloader._loader_identity( *, source: nemo_rl.data.energon.config.EnergonSourceConfig, loader_config: nemo_rl.data.energon.config.EnergonLoaderConfig, adapter_fingerprint: str, split_role: str, batch_size: int, shuffle: bool | None, topology: dict[str, Any], packing_algorithm: str | None, max_sequences_per_bin: int | None, sequence_length_pad_multiple: int, only_unmask_final: bool, ) → dict[str, Any]#: Describe what a restored loader must still agree with.

nemo_rl.data.energon.sft_dataloader._worker_config( config: nemo_rl.data.energon.config.EnergonLoaderConfig, *, logical_rank: int, logical_world_size: int, ) → megatron.energon.WorkerConfig#

nemo_rl.data.energon.sft_dataloader._task_encoder( *, loader_config: nemo_rl.data.energon.config.EnergonLoaderConfig, adapter: Any, include_source_ids: bool, packing_algorithm: str | None, max_sequences_per_bin: int | None, max_sequence_length: int, sequence_length_pad_multiple: int, tokenizer: Any, only_unmask_final: bool, ) → nemo_rl.data.energon.multimodal.task_encoders.base.BaseSFTTaskEncoder#

nemo_rl.data.energon.sft_dataloader.build_energon_sft_loader( *, data_config: Mapping[str, Any], source: Mapping[str, Any] | nemo_rl.data.energon.config.EnergonSourceConfig, processor: Any, batch_size: int, max_sequence_length: int, split_role: Literal[train, validation], logical_rank: int, logical_world_size: int, placement_fingerprint: str, packing_algorithm: str | None, max_sequences_per_bin: int | None, sequence_length_pad_multiple: int, only_unmask_final: bool, ) → nemo_rl.data.energon.sft_dataloader.EnergonSFTDataLoader#: Build one loader for an explicit logical data shard and split.

nemo_rl.data.energon.sft_dataloader.__all__#: [‘EnergonSFTDataLoader’, ‘SFTDataLoader’, ‘build_energon_sft_loader’]

nemo_rl.data.energon.sft_dataloader#