bridge.training.gpt_step#

Module Contents#

Functions#

get_packed_seq_params

Extract packed sequence parameters from the batch.

get_batch_from_iterator

Get a batch of data from the iterator.

get_batch_on_this_tp_rank

Get a batch from the data iterator, handling TP broadcasting.

get_batch

Generate a batch.

forward_step

Forward training step.

_create_loss_function

Create a partial loss function with the specified configuration.

Data#

logger

API#

bridge.training.gpt_step.logger#

‘getLogger(…)’

bridge.training.gpt_step.get_packed_seq_params(
batch: dict[str, torch.Tensor],
) → megatron.core.packed_seq_params.PackedSeqParams#

Extract packed sequence parameters from the batch.

Creates and returns a PackedSeqParams object with appropriate parameters for packed sequence processing.

Parameters:

batch – Input batch containing packed sequence information

Returns:

Parameters for packed sequence processing

Return type:

PackedSeqParams

bridge.training.gpt_step.get_batch_from_iterator(
data_iterator: Iterable,
use_mtp: bool = False,
skip_getting_attention_mask_from_dataset: bool = True,
) → dict[str, torch.Tensor]#

Get a batch of data from the iterator.

Parameters:

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

  • use_mtp – Whether Multi-Token Prediction layers are enabled.

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

Returns:

A dictionary containing the batch data.

Return type:

dict[str, torch.Tensor]

bridge.training.gpt_step.get_batch_on_this_tp_rank(
data_iterator: Iterable,
cfg: megatron.bridge.training.config.ConfigContainer,
use_mtp: bool = False,
) → dict[str, torch.Tensor]#

Get a batch from the data iterator, handling TP broadcasting.

On TP rank 0, it fetches the next batch from the iterator and broadcasts the necessary tensors to other TP ranks based on the pipeline stage. On other TP ranks, it allocates tensors and receives the broadcasted data.

Parameters:

  • data_iterator – The data iterator.

  • cfg – The configuration container.

  • use_mtp – Whether Multi-Token Prediction layers are enabled.

Returns:

A dictionary containing the batch data for the current rank.

bridge.training.gpt_step.get_batch(
data_iterator: Iterable,
cfg: megatron.bridge.training.config.ConfigContainer,
use_mtp: bool = False,
) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]#

Generate a batch.

Parameters:

  • data_iterator – Input data iterator

  • cfg – Configuration container

  • use_mtp – Whether Multi-Token Prediction layers are enabled

Returns:

tuple of tensors containing tokens, labels, loss_mask, attention_mask, position_ids, cu_seqlens, cu_seqlens_argmin, and max_seqlen

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