nemo_rl.data_plane.schema#

Shared constants and type aliases for the data-plane meta contract.

Module Contents#

Functions#

fields_with_optional_routed_experts

Return fields plus routed experts when router replay is enabled.

fields_with_optional_opd_full

Return fields plus the full-vocabulary MOPD teacher payload column(s).

Data#

API#

nemo_rl.data_plane.schema.Layout#

None

nemo_rl.data_plane.schema.MICRO_BATCH_INDICES#

‘micro_batch_indices’

nemo_rl.data_plane.schema.MICRO_BATCH_LENGTHS#

‘micro_batch_lengths’

nemo_rl.data_plane.schema.ELEM_COUNTS_PER_GB#

‘elem_counts_per_gb’

nemo_rl.data_plane.schema.GLOBAL_FORWARD_PAD_SEQLEN#

‘global_forward_pad_seqlen’

nemo_rl.data_plane.schema.ROLLOUT_METRICS#

‘rollout_metrics’

nemo_rl.data_plane.schema.GROUP_ID_TAG#

‘group_id’

nemo_rl.data_plane.schema.PER_TOKEN_STAGING_FIELDS#

(‘token_ids_delta’, ‘token_mask_delta’, ‘generation_logprobs_delta’)

nemo_rl.data_plane.schema.INPUT_IDS#

‘input_ids’

nemo_rl.data_plane.schema.INPUT_LENGTHS#

‘input_lengths’

nemo_rl.data_plane.schema.SAMPLE_MASK#

‘sample_mask’

nemo_rl.data_plane.schema.MASK_SAMPLE#

‘mask_sample’

nemo_rl.data_plane.schema.TRUNCATED#

‘truncated’

nemo_rl.data_plane.schema.META_IDX#

‘meta_idx’

nemo_rl.data_plane.schema.INVALID_TOOL_CALL_MASK#

‘invalid_tool_call_mask’

nemo_rl.data_plane.schema.MALFORMED_THINKING_MASK#

‘malformed_thinking_mask’

nemo_rl.data_plane.schema.DP_TRAIN_FIELDS#

(‘input_ids’, ‘input_lengths’, ‘generation_logprobs’, ‘prev_logprobs’, ‘reference_policy_logprobs’, …

nemo_rl.data_plane.schema.OPD_FULL_HIDDEN_STATES_FIELD#

‘teacher_full_hidden_states’

nemo_rl.data_plane.schema.OPD_FULL_LOGITS_FIELD#

‘teacher_full_logits’

nemo_rl.data_plane.schema.OPD_FULL_FIELDS#

()

nemo_rl.data_plane.schema.OPD_FULL_TEACHER_INDEX_FIELD#

‘opd_full_teacher_index’

nemo_rl.data_plane.schema.SC_ROLLOUT_SCHEMA_FIELDS#

()

nemo_rl.data_plane.schema.LP_SEED_FIELDS#

(‘input_ids’, ‘input_lengths’, ‘token_mask’, ‘sample_mask’)

nemo_rl.data_plane.schema.TEACHER_LP_FIELDS#

()

nemo_rl.data_plane.schema.PPO_VALUE_FIELDS#

(‘values’, ‘returns’)

nemo_rl.data_plane.schema.DP_VALUE_TRAIN_FIELDS#

(‘input_ids’, ‘input_lengths’, ‘token_mask’, ‘sample_mask’)

nemo_rl.data_plane.schema.VALUE_SEED_FIELDS#

None

nemo_rl.data_plane.schema.DP_CALIB_INPUT_FIELDS#

()

nemo_rl.data_plane.schema.ROUTED_EXPERTS_FIELD#

‘routed_experts’

nemo_rl.data_plane.schema.ROUTED_LEN_FIELD#

‘routed_len’

nemo_rl.data_plane.schema.ROUTED_EXPERTS_ENCODING_FIELD#

‘routed_experts_encoding’

nemo_rl.data_plane.schema.ROUTED_EXTRAS_METADATA_FIELD#

‘extras_metadata_json’

nemo_rl.data_plane.schema.ROUTE_ENCODING_NONE#

0

nemo_rl.data_plane.schema.ROUTE_ENCODING_ENVELOPE#

1

nemo_rl.data_plane.schema.ROUTE_ENCODING_LIST#

2

nemo_rl.data_plane.schema.ROUTE_PLAN_TAG#

‘route_assembly_plan’

nemo_rl.data_plane.schema.ROUTE_PASSTHROUGH_FLAG#

‘route_passthrough’

nemo_rl.data_plane.schema.fields_with_optional_routed_experts(
fields: Sequence[str],
*,
enabled: bool,
) → list[str]#

Return fields plus routed experts when router replay is enabled.

nemo_rl.data_plane.schema.fields_with_optional_opd_full(
fields: Sequence[str],
*,
field: Optional[str],
teacher_index_field: Optional[str] = None,
) → list[str]#

Return fields plus the full-vocabulary MOPD teacher payload column(s).

Added only when the run configures them: a GRPO run requesting a column nobody wrote would error on read, hence not folded into DP_TRAIN_FIELDS.

Parameters:
  • fields – Base field list.

  • field – Payload column name, or None when opd_full is off.

  • teacher_index_field – Per-sample teacher-identity column name (hidden- state path only, see OPD_FULL_TEACHER_INDEX_FIELD), or None when opd_full is off or using the logits payload.

Returns:

The field list, with the payload/index columns appended when applicable.