bridge.models.exaone.exaone45.modelling_exaone45.utils

Module Contents

Classes

Exaone45VisionPatchEmbed	Vision Patch Embed for Exaone45 vision model.
Exaone45VisionRotaryEmbedding	Vision Rotary Embedding for Exaone45 vision model.
PatchMergerSubmodules	Patch Merger Submodules for Exaone45 vision model.
Exaone45VisionPatchMerger	Vision Patch Merger for Exaone45 vision model.
AllGatherVisionEmbeddings	AllGatherVisionEmbeddings for Exaone45 vision model.

Functions

find_vision_id_index	Find the vision id index for Exaone45 vision model.
reorganize_inputs	Reorganize the inputs for Exaone45 vision model.
split_data_cp_rank	Split the data by CP rank for EXAONE 4.5 vision model, using zigzag pattern.
expand_thw	Expand the THW for Exaone45 vision model.
collapse_thw	Collapse the THW for Exaone45 vision model.
exaone45_pad_and_split	Split the pixel values and image grid thws for Exaone45 vision model.
Exaone45_cp_split	Split the pixel values and image grid thws for Exaone45 vision model.
get_vision_cp_data	Get vision data and grid_thw for context parallelism.
preprocess_packed_seqs	Preprocess packed sequences CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets second and second last chunks, and so on), this is for load balancing with causal masking. See https://github.com/NVIDIA/TransformerEngine/issues/1368

API

class bridge.models.exaone.exaone45.modelling_exaone45.utils.Exaone45VisionPatchEmbed(

config: megatron.bridge.models.exaone.exaone45.modelling_exaone45.transformer_config.Exaone45TransformerConfig,

)

Bases: torch.nn.Module

Vision Patch Embed for Exaone45 vision model.

Initialization

forward(hidden_states: torch.Tensor) → torch.Tensor

class bridge.models.exaone.exaone45.modelling_exaone45.utils.Exaone45VisionRotaryEmbedding(dim: int, theta: float = 10000.0)

Bases: torch.nn.Module

Vision Rotary Embedding for Exaone45 vision model.

Initialization

forward(seqlen: int) → torch.Tensor

class bridge.models.exaone.exaone45.modelling_exaone45.utils.PatchMergerSubmodules

Patch Merger Submodules for Exaone45 vision model.

patch_norm: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]

None

linear_fc1: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]

None

linear_fc2: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]

None

class bridge.models.exaone.exaone45.modelling_exaone45.utils.Exaone45VisionPatchMerger(

config: megatron.bridge.models.exaone.exaone45.modelling_exaone45.transformer_config.Exaone45TransformerConfig,

submodules: bridge.models.exaone.exaone45.modelling_exaone45.utils.PatchMergerSubmodules,

use_postshuffle_norm=False,

tp_group: Optional[torch.distributed.ProcessGroup] = None,

)

Bases: megatron.core.transformer.module.MegatronModule

Vision Patch Merger for Exaone45 vision model.

Initialization

forward(hidden_states)

bridge.models.exaone.exaone45.modelling_exaone45.utils.find_vision_id_index(

input_ids: torch.Tensor,

image_token_id: int,

video_token_id: int,

)

Find the vision id index for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.reorganize_inputs(

input_ids: torch.Tensor,

pixel_values: torch.Tensor = None,

pixel_values_videos: torch.Tensor = None,

image_grid_thw: torch.Tensor = None,

video_grid_thw: torch.Tensor = None,

image_input_mask: torch.Tensor = None,

video_input_mask: torch.Tensor = None,

image_token_id: int = 151655,

video_token_id: int = 151656,

square_merge_size: int = 4,

)

Reorganize the inputs for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.split_data_cp_rank(

val: torch.Tensor,

cp_size: int,

seq_dim: int,

cp_rank: int = None,

)

Split the data by CP rank for EXAONE 4.5 vision model, using zigzag pattern.

bridge.models.exaone.exaone45.modelling_exaone45.utils.expand_thw(thw: torch.Tensor) → torch.Tensor

Expand the THW for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.collapse_thw(expanded: torch.Tensor) → torch.Tensor

Collapse the THW for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.exaone45_pad_and_split(

cp_size: int,

hw_factor: int,

pixel_values: list[torch.Tensor],

image_grid_thws: list[torch.Tensor],

)

Split the pixel values and image grid thws for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.Exaone45_cp_split(

cp_size: int,

pixel_values: torch.Tensor,

image_grid_thw: torch.Tensor,

)

Split the pixel values and image grid thws for Exaone45 vision model.

bridge.models.exaone.exaone45.modelling_exaone45.utils.get_vision_cp_data(

vision_data: torch.Tensor,

vision_grid_thw: torch.Tensor,

square_merge_size: int,

cp_img_num: list[int],

images_padded: list[bool],

cp_rank: int,

cp_size: int,

)

Get vision data and grid_thw for context parallelism.

Returns:

Vision data of shape [total_thw_size, n_features]. vision_grid_thw (torch.Tensor): Vision grid_thw of shape [total_thw_size, 3]. seqlens_list (list of torch.Tensor): List of seqlens of the vision data in each context parallel rank, for the all gather after vision encoder.

Return type:

vision_data (torch.Tensor)

class bridge.models.exaone.exaone45.modelling_exaone45.utils.AllGatherVisionEmbeddings

Bases: torch.autograd.Function

AllGatherVisionEmbeddings for Exaone45 vision model.

static forward(

ctx,

input,

seqlens_on_cp_ranks,

cp_group: torch.distributed.ProcessGroup,

)

Forward pass for AllGatherVisionEmbeddings.

static backward(ctx, grad_output)

Backward pass for AllGatherVisionEmbeddings.

bridge.models.exaone.exaone45.modelling_exaone45.utils.preprocess_packed_seqs(

input_ids: torch.Tensor,

attention_mask: torch.Tensor,

pre_process: bool = True,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

) → tuple[torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams]

Preprocess packed sequences CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets second and second last chunks, and so on), this is for load balancing with causal masking. See https://github.com/NVIDIA/TransformerEngine/issues/1368