bridge.models.muse_spark_open.modeling_muse_spark_open.model

Hybrid language decoder with directly invoked HF perception encoders.

Module Contents

Classes

MuseSparkModel

Place HF perception towers on the embedding stage of an MCore HybridModel.

API

class bridge.models.muse_spark_open.modeling_muse_spark_open.model.MuseSparkModel(

config: megatron.bridge.models.muse_spark_open.muse_spark_open_config.MuseSparkModelConfig,

*,

vocab_size: int,

pg_collection: megatron.core.process_groups_config.ProcessGroupCollection,

pre_process: bool,

post_process: bool,

)

Bases: megatron.core.models.hybrid.hybrid_model.HybridModel

Place HF perception towers on the embedding stage of an MCore HybridModel.

Initialization

_build_perception() → None

_freeze_requested_modules() → None

_normalize_embeddings(embeddings: torch.Tensor) → torch.Tensor

_vision_features(

pixel_values: torch.Tensor,

grid_thw: torch.Tensor,

) → torch.Tensor

_restore_hf_vision_rope_buffers() → None

Restore HF’s nonpersistent FP32 frequencies after model dtype conversions.

_audio_features(

input_values: torch.Tensor,

input_values_mask: torch.Tensor | None,

) → torch.Tensor

static _insert_features(

embeddings: torch.Tensor,

input_ids: torch.Tensor,

features: torch.Tensor,

token_id: int,

) → torch.Tensor

forward(

input_ids: torch.Tensor | None = None,

position_ids: torch.Tensor | None = None,

attention_mask: torch.Tensor | None = None,

*,

inputs_embeds: torch.Tensor | None = None,

pixel_values: torch.Tensor | None = None,

image_grid_thw: torch.Tensor | None = None,

pixel_values_videos: torch.Tensor | None = None,

video_grid_thw: torch.Tensor | None = None,

input_values: torch.Tensor | None = None,

input_values_mask: torch.Tensor | None = None,

labels: torch.Tensor | None = None,

loss_mask: torch.Tensor | None = None,

padding_mask: torch.Tensor | None = None,

packed_seq_params: Any = None,

**kwargs: Any,

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

Encode media with HF, insert features, then execute the Hybrid decoder.

inputs_embeds is batch-first and already normalized, as in HF. The wrapper receives full sequences and slices for CP after media insertion.