bridge.models.qwen_omni.qwen3_omni_step#

Qwen3-Omni thinker training step helpers.

Module Contents#

Functions#

get_batch_from_iterator

Get a thinker-training batch from the iterator.

_normalize_multimodal_inputs

Normalize multimodal batch tensors for Qwen3-Omni model forward.

get_batch

Generate a minimal thinker-training batch.

pad_batch_sequences_for_context_parallel

Pad dense sequence tensors before Megatron’s CP zigzag split.

_get_dense_batch_on_this_cp_rank

Slice dense CP tensors, including 2D attention masks from VLM datasets.

forward_step

Forward training step for Qwen3-Omni thinker.

Data#

_MULTIMODAL_KEYS

API#

bridge.models.qwen_omni.qwen3_omni_step._MULTIMODAL_KEYS#

(‘pixel_values’, ‘image_grid_thw’, ‘pixel_values_videos’, ‘video_grid_thw’, ‘video_second_per_grid’,…

bridge.models.qwen_omni.qwen3_omni_step.get_batch_from_iterator(
data_iterator: Iterable,
use_mtp: bool = False,
skip_getting_attention_mask_from_dataset: bool = True,
*,
is_first_pp_stage: bool,
is_last_pp_stage: bool,
) → dict[str, Any]#

Get a thinker-training batch from the iterator.

bridge.models.qwen_omni.qwen3_omni_step._normalize_multimodal_inputs(
batch: dict[str, Any],
) → dict[str, torch.Tensor]#

Normalize multimodal batch tensors for Qwen3-Omni model forward.

bridge.models.qwen_omni.qwen3_omni_step.get_batch(
data_iterator: Iterable,
cfg: megatron.bridge.training.config.ConfigContainer,
use_mtp: bool = False,
*,
pg_collection,
) → tuple[...]#

Generate a minimal thinker-training batch.

bridge.models.qwen_omni.qwen3_omni_step.pad_batch_sequences_for_context_parallel(
tokens: torch.Tensor,
labels: torch.Tensor | None,
loss_mask: torch.Tensor | None,
attention_mask: torch.Tensor | None,
position_ids: torch.Tensor | None,
pg_collection,
*,
force_to_seq_length: bool = False,
seq_length: int | None = None,
) → tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]#

Pad dense sequence tensors before Megatron’s CP zigzag split.

Dense CP partitions each sequence into 2 * cp_size chunks. Padding here keeps the step-level tensors compatible with Megatron’s CP slicing while the full input_ids tensor remains available for model-internal mRoPE.

bridge.models.qwen_omni.qwen3_omni_step._get_dense_batch_on_this_cp_rank(
batch: dict[str, Any],
cp_group,
) → dict[str, Any]#

Slice dense CP tensors, including 2D attention masks from VLM datasets.

bridge.models.qwen_omni.qwen3_omni_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 for Qwen3-Omni thinker.