bridge.training.llava_step#

Module Contents#

Functions#

_expand_packed_metadata_for_visual_embeddings

Adjust packed row boundaries for image placeholders expanded by LLaVA.

_validate_packed_parallelism

Reject packed LLaVA layouts whose data ownership is not implemented.

get_batch_from_iterator

Get a batch of data from the iterator.

get_batch

Generate a batch.

forward_step

Forward training step.

_create_loss_function

Create a partial loss function with the specified configuration.

Data#

logger

_PACKED_SEQ_DEVICE_KEYS

_PACKED_SEQ_HOST_KEYS

_PACKED_SEQ_PARAM_KEYS

API#

bridge.training.llava_step.logger#

‘getLogger(…)’

bridge.training.llava_step._PACKED_SEQ_DEVICE_KEYS#

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

bridge.training.llava_step._PACKED_SEQ_HOST_KEYS#

(‘max_seqlen_q’, ‘max_seqlen_kv’)

bridge.training.llava_step._PACKED_SEQ_PARAM_KEYS#

()

bridge.training.llava_step._expand_packed_metadata_for_visual_embeddings(
packed_metadata: dict[str, Any],
input_ids: torch.Tensor,
model: megatron.core.models.gpt.GPTModel,
) → None#

Adjust packed row boundaries for image placeholders expanded by LLaVA.

bridge.training.llava_step._validate_packed_parallelism(*, pg_collection) → None#

Reject packed LLaVA layouts whose data ownership is not implemented.

bridge.training.llava_step.get_batch_from_iterator(
data_iterator: Iterable,
skip_getting_attention_mask_from_dataset: bool = True,
*,
is_first_pp_stage: bool,
is_last_pp_stage: bool,
) → dict[str, Any]#

Get a batch of data from the iterator.

Parameters:

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

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

  • is_first_pp_stage – Whether this is the first pipeline parallel stage.

  • is_last_pp_stage – Whether this is the last pipeline parallel stage.

Returns:

A dictionary containing the batch data.

Return type:

dict[str, torch.Tensor]

bridge.training.llava_step.get_batch(
data_iterator: Iterable,
cfg: megatron.bridge.training.config.ConfigContainer,
*,
pg_collection,
) → tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, dict[str, Any] | None]#

Generate a batch.

Parameters:

  • data_iterator – Input data iterator

  • cfg – Configuration container

  • pg_collection – Process group collection for distributed training

Returns:

tuple of tensors containing tokens, labels, loss_mask, attention_mask, position_ids, cu_seqlens (optional), cu_seqlens_argmin (optional), max_seqlen (optional), images (optional)

bridge.training.llava_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.llava_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