bridge.training.gpt_step

Module Contents

Functions

_trim_padded_cu_seqlens_for_cp	Trim padded THD cu_seqlens without introducing a CUDA sync.
_has_packed_sequence_metadata	Return whether a dataloader batch contains packed-sequence metadata.
_packed_metadata_for_forward	Extract packed-sequence metadata accepted by get_packed_seq_params.
_cu_seqlens_for_cp_partition	Return the cu-seqlens tensor TE should use to partition packed THD tokens.
_uses_packed_sequence_metadata	Return whether the dataset is expected to provide packed sequence metadata.
_middle_pp_stage_needs_batch	Return whether middle PP stages need batch metadata for attention.
_layout_stage_has_mtp	Return whether a parsed or raw pipeline layout stage owns MTP layers.
_current_stage_has_mtp_from_layout	Return whether the current PP/VPP stage owns the configured MTP block, derived from layout.
_current_stage_needs_mtp_inputs_from_layout	Return whether this stage needs token ids for MTP embedding lookup, derived from layout.
_partition_packed_batch_for_cp	Partition THD/packed batches across context-parallel ranks.
get_batch_from_iterator	Get a batch of data from the iterator.
get_batch	Generate a batch.
_forward_step_common	Forward training step.
forward_step	Forward training step.
_create_loss_function	Create a partial loss function with the specified configuration.
forward_step_modelopt	Forward training step with ModelOpt required modifications.
_create_loss_function_modelopt	Create a partial loss function with the specified configuration.

Data

logger
_CURRENT_PACKED_SEQ_DEVICE_KEYS
_CURRENT_PACKED_SEQ_HOST_KEYS
_CURRENT_PACKED_SEQ_PARAM_KEYS
_LEGACY_PACKED_SEQ_DEVICE_KEYS
_LEGACY_PACKED_SEQ_HOST_KEYS
_LEGACY_PACKED_SEQ_PARAM_KEYS
_PackedMetadataValue

API

bridge.training.gpt_step.logger

‘getLogger(…)’

bridge.training.gpt_step._CURRENT_PACKED_SEQ_DEVICE_KEYS

(‘cu_seqlens_q’, ‘cu_seqlens_kv’, ‘cu_seqlens_q_padded’, ‘cu_seqlens_kv_padded’)

bridge.training.gpt_step._CURRENT_PACKED_SEQ_HOST_KEYS

(‘max_seqlen_q’, ‘max_seqlen_kv’)

bridge.training.gpt_step._CURRENT_PACKED_SEQ_PARAM_KEYS

()

bridge.training.gpt_step._LEGACY_PACKED_SEQ_DEVICE_KEYS

(‘cu_seqlens’, ‘cu_seqlens_unpadded’)

bridge.training.gpt_step._LEGACY_PACKED_SEQ_HOST_KEYS

(‘cu_seqlens_argmin’, ‘max_seqlen’, ‘cu_seqlens_unpadded_argmin’)

bridge.training.gpt_step._LEGACY_PACKED_SEQ_PARAM_KEYS

()

bridge.training.gpt_step._PackedMetadataValue

None

bridge.training.gpt_step._trim_padded_cu_seqlens_for_cp(

cu_seqlens: torch.Tensor,

cu_seqlens_argmin: torch.Tensor | None,

) → torch.Tensor

Trim padded THD cu_seqlens without introducing a CUDA sync.

bridge.training.gpt_step._has_packed_sequence_metadata(batch: dict[str, torch.Tensor]) → bool

Return whether a dataloader batch contains packed-sequence metadata.

bridge.training.gpt_step._packed_metadata_for_forward(

batch: dict[str, torch.Tensor],

) → dict[str, bridge.training.gpt_step._PackedMetadataValue] | None

Extract packed-sequence metadata accepted by get_packed_seq_params.

bridge.training.gpt_step._cu_seqlens_for_cp_partition(

batch: dict[str, torch.Tensor],

) → torch.Tensor

Return the cu-seqlens tensor TE should use to partition packed THD tokens.

bridge.training.gpt_step._uses_packed_sequence_metadata(

cfg: megatron.bridge.training.config.ConfigContainer,

) → bool

Return whether the dataset is expected to provide packed sequence metadata.

bridge.training.gpt_step._middle_pp_stage_needs_batch(

cfg: megatron.bridge.training.config.ConfigContainer,

) → bool

Return whether middle PP stages need batch metadata for attention.

bridge.training.gpt_step._layout_stage_has_mtp(

layout,

*,

pp_rank: int,

pp_size: int,

vp_stage: int,

) → bool

Return whether a parsed or raw pipeline layout stage owns MTP layers.

bridge.training.gpt_step._current_stage_has_mtp_from_layout(

cfg: megatron.bridge.training.config.ConfigContainer,

*,

pg_collection,

vp_stage: int | None = None,

) → bool

Return whether the current PP/VPP stage owns the configured MTP block, derived from layout.

bridge.training.gpt_step._current_stage_needs_mtp_inputs_from_layout(

cfg: megatron.bridge.training.config.ConfigContainer,

*,

pg_collection,

is_last: bool,

vp_stage: int | None = None,

) → bool

Return whether this stage needs token ids for MTP embedding lookup, derived from layout.

bridge.training.gpt_step._partition_packed_batch_for_cp(

batch: dict[str, torch.Tensor],

cp_group: torch.distributed.ProcessGroup,

) → dict[str, torch.Tensor]

Partition THD/packed batches across context-parallel ranks.

Uses MCore’s packed-sequence partitioning to slice sequence dimensions aligned with packed cu_seqlens.

bridge.training.gpt_step.get_batch_from_iterator(

data_iterator: Iterable,

include_mtp_inputs: bool = False,

skip_getting_attention_mask_from_dataset: bool = True,

*,

is_first_pp_stage: bool,

is_last_pp_stage: bool,

include_full_batch_fields: bool = False,

) → dict[str, torch.Tensor]

Get a batch of data from the iterator.

Parameters:

• data_iterator – The data iterator to get the batch from.

• include_mtp_inputs – Whether this PP stage needs Multi-Token Prediction input tensors.

• skip_getting_attention_mask_from_dataset – If set, the dataset will pass a None attention mask.

• include_full_batch_fields – Whether to include all standard training tensors regardless of PP stage.

Returns:

A dictionary containing the batch data.

Return type:

dict[str, torch.Tensor]

bridge.training.gpt_step.get_batch(

data_iterator: Iterable,

cfg: megatron.bridge.training.config.ConfigContainer,

use_mtp: bool = False,

*,

pg_collection,

vp_stage: int | None = None,

) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, bridge.training.gpt_step._PackedMetadataValue] | None]

Generate a batch.

Parameters:

• data_iterator – Input data iterator

• cfg – Configuration container

• use_mtp – Whether Multi-Token Prediction layers are enabled

• vp_stage – Virtual pipeline stage for the current model chunk.

Returns:

tuple of tensors containing tokens, labels, loss_mask, attention_mask, position_ids, and optional packed-sequence metadata.

bridge.training.gpt_step._forward_step_common(

state: megatron.bridge.training.state.GlobalState,

data_iterator: Iterable,

model: megatron.core.models.gpt.GPTModel,

return_schedule_plan: bool = False,

) → tuple[torch.Tensor, torch.Tensor]

Forward training step.

Parameters:

• state – Global state for the run

• data_iterator – Input data iterator

• model – The GPT Model

• return_schedule_plan (bool) – Whether to return the schedule plan instead of the output tensor

Returns:

tuple containing the output tensor and loss mask

bridge.training.gpt_step.forward_step(

state: megatron.bridge.training.state.GlobalState,

data_iterator: Iterable,

model: megatron.core.models.gpt.GPTModel,

return_schedule_plan: bool = False,

) → tuple[torch.Tensor, functools.partial]

Forward training step.

Parameters:

• state – Global state for the run

• data_iterator – Input data iterator

• model – The GPT Model

• return_schedule_plan (bool) – Whether to return the schedule plan instead of the output tensor

Returns:

tuple containing the output tensor and the loss function

bridge.training.gpt_step._create_loss_function(

loss_mask: torch.Tensor,

check_for_nan_in_loss: bool,

check_for_spiky_loss: bool,

) → functools.partial

Create a partial loss function with the specified configuration.

Parameters:

• loss_mask – Used to mask out some portions of the loss

• check_for_nan_in_loss – Whether to check for NaN values in the loss

• check_for_spiky_loss – Whether to check for spiky loss values

Returns:

A partial function that can be called with output_tensor to compute the loss

bridge.training.gpt_step.forward_step_modelopt(

state: megatron.bridge.training.state.GlobalState,

data_iterator: Iterable,

model: megatron.core.models.gpt.GPTModel,

return_schedule_plan: bool = False,

) → tuple[torch.Tensor, functools.partial]

Forward training step with ModelOpt required modifications.

Parameters:

• state – Global state for the run

• data_iterator – Input data iterator

• model – The GPT Model

• return_schedule_plan (bool) – Whether to return the schedule plan instead of the output tensor

Returns:

tuple containing the output tensor and the loss function

bridge.training.gpt_step._create_loss_function_modelopt(

loss_mask: torch.Tensor,

model: megatron.core.models.gpt.GPTModel,

check_for_nan_in_loss: bool,

check_for_spiky_loss: bool,

) → functools.partial

Create a partial loss function with the specified configuration.

Kept here for backward compatibility with tests and callers that patch megatron.bridge.training.gpt_step.masked_next_token_loss.

Parameters:

• loss_mask – Used to mask out some portions of the loss

• model – The GPT Model

• check_for_nan_in_loss – Whether to check for NaN values in the loss

• check_for_spiky_loss – Whether to check for spiky loss values

Returns:

A partial function that can be called with output_tensor to compute the loss