nemo_automodel.components.models.qwen3_5.model

Qwen3.5 dense causal LM with Megatron-style MTP support.

Module Contents

Classes

Name	Description
`Fp32SafeQwen3_5TextRotaryEmbedding`	Ensure inv_freq stays in float32 across `.to(dtype)` calls.
`Qwen3_5CausalLMOutputWithPast`	Qwen3.5 causal-LM output extended with MTP auxiliary hidden states.
`Qwen3_5DenseBlock`	Qwen3.5 dense decoder block on top of the Qwen3-Next `Block`.
`Qwen3_5DenseMTPSublayer`	One full-attention Qwen3.5 dense MTP sublayer.
`Qwen3_5DenseTextBackbone`	Qwen3.5 dense text decoder rebuilt on the Qwen3-Next `Block`.
`Qwen3_5ForCausalLM`	Qwen3.5 dense causal LM with optional Megatron-style MTP head.
`Qwen3_5ForConditionalGeneration`	Qwen3.5/Qwen3.6 dense VLM with optional Megatron-style MTP head.
`Qwen3_5Model`	Thin VLM wrapper exposing `language_model` internals as properties and
`_Qwen3_5Attention`	Qwen3.5 full attention with packed block-diagonal CP dispatch.

Functions

Name	Description
`_default_init_device`	-
`_dense_moe_config`	Trivial MoEConfig for the dense Qwen3.5 backbone.
`_make_full_attention_config`	-
`_mask_mtp_embed_inputs`	Zero invalid future-token embeddings in the local CP layout.
`_mtp_block_causal_mask`	Build a 4D block-causal attention mask from an indexed packing mask.
`_qwen3_5_backend`	Return a Qwen3.5 backend with TE fused RoPE disabled.
`_resolve_mtp_num_layers`	-
`_rolled_embed_inputs`	-
`_splice_multimodal_features`	Replace multimodal placeholder rows without dynamic-size indexing.
`_split_qwen3_5_position_ids`	-
`build_mtp_config_from_hf`	Build Qwen3.5 MTP runtime config from HF-style config fields.
`build_qwen3_5_dense_mtp`	Construct dense Qwen3.5 MTP blocks.

Data

ModelClass

API

class nemo_automodel.components.models.qwen3_5.model.Fp32SafeQwen3_5TextRotaryEmbedding()

Bases: Qwen3_5TextRotaryEmbedding

Ensure inv_freq stays in float32 across .to(dtype) calls.

nemo_automodel.components.models.qwen3_5.model.Fp32SafeQwen3_5TextRotaryEmbedding._apply(
    fn: typing.Any,
    recurse: bool = True
)

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5CausalLMOutputWithPast(
    rope_deltas: torch.Tensor | None = None,
    mtp_per_depth_h: list[torch.Tensor] | None = None,
    mtp_loss_scaling_factor: float | None = None
)

Dataclass

Bases: CausalLMOutputWithPast

Qwen3.5 causal-LM output extended with MTP auxiliary hidden states.

mtp_loss_scaling_factor

float | None = None

mtp_per_depth_h

list[Tensor] | None = None

rope_deltas

Tensor | None = None

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseBlock(
    layer_idx,
    config,
    moe_config,
    backend
)

Bases: Block

Qwen3.5 dense decoder block on top of the Qwen3-Next Block.

Identical to Qwen3_5MoeBlock except the MLP degrades to a dense MLP (no experts). The CP-aware GatedDeltaNet is built natively for linear-attention layers, and the forward threads NEAT-packing kwargs.

input_layernorm

linear_attn

= CPAwareGatedDeltaNet(config, layer_idx)

mlp

post_attention_layernorm

self_attn

= _Qwen3_5Attention(config, layer_idx, backend)

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseBlock._post_attention(
    x: torch.Tensor,
    attn_out: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor

Add the attention output to the residual stream, then apply the MLP sub-block.

The residual add is fused into the post-attention norm.

Parameters:

torch.Tensor

Residual stream of shape [batch, sequence, hidden], or [tokens, hidden] for the unpadded TE stream.

attn_out

torch.Tensor

Attention or GatedDeltaNet output with x’s shape.

padding_mask

torch.Tensor | None

Optional padding mask of shape [batch, sequence].

Returns: torch.Tensor

Hidden states with x’s shape.

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseBlock.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    packed_gdn_metadata: nemo_automodel.components.models.qwen3_5.packing.GatedDeltaPackedMetadata | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Run one dense Qwen3.5 decoder block.

Parameters:

torch.Tensor

Hidden states of shape [batch, sequence, hidden].

freqs_cis

torch.Tensor

Rotary cosine/sine values of shape [batch, sequence, rotary_dim].

attention_mask

torch.Tensor | NoneDefaults to None

Optional validity, indexed document, or backend mask of shape [batch, sequence] or [batch, 1, sequence, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Optional padding mask of shape [batch, sequence].

position_ids

torch.Tensor | NoneDefaults to None

Optional positions of shape [batch, sequence] or [axes, batch, sequence].

packed_gdn_metadata

GatedDeltaPackedMetadata | NoneDefaults to None

Optional model-forward-owned packing metadata; tensor layouts are documented by GatedDeltaPackedMetadata.

**attn_kwargs

AnyDefaults to {}

Packed or THD metadata: _packed_seq_ids has shape [batch, sequence] with 1-based document IDs and zero padding. packed_token_indices has shape [batch, sequence] with row-local positions and -1 padding, or [tokens] indexing flattened batch and sequence axes. cu_seqlens has shape [batch, max_documents + 1] with row-local boundaries and -1 padding, or [documents + 1] for flattened/THD inputs. max_seqlen is an integer document length.

Returns: torch.Tensor

Hidden states of shape [batch, sequence, hidden].

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseBlock.init_weights(
    buffer_device: torch.device
)

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseMTPSublayer(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    layer_idx: int,
    has_fusion: bool = False,
    has_final_norm: bool = False,
    dtype: torch.dtype = torch.bfloat16
)

Bases: Qwen3_5DecoderLayer

One full-attention Qwen3.5 dense MTP sublayer.

eh_proj

enorm

final_layernorm

hnorm

layer_type

= self.block_type

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseMTPSublayer.forward(
    hidden_states: torch.Tensor,
    embed_input: torch.Tensor | None = None,
    rotary_emb: torch.nn.Module,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    past_key_values: typing.Any | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseMTPSublayer.init_weights(
    buffer_device: torch.device | None = None
) -> None

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseTextBackbone(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Qwen3.5 dense text decoder rebuilt on the Qwen3-Next Block.

Native counterpart of Qwen3_5MoeTextModelBackend for the dense model: reuses the same blocks/GatedDeltaNet/norm/rotary so dense and MoE share one code path, with the fp32 SSMGate built at construction (no runtime patch).

embed_tokens

layers

norm

padding_idx

= getattr(config, 'pad_token_id', None)

rotary_emb

= Fp32SafeQwen3_5TextRotaryEmbedding(config=config)

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseTextBackbone.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    cache_position: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    past_key_values: typing.Any | None = None,
    use_cache: bool | None = None,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Decode text tokens with model-forward-owned packed GDN metadata.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Optional token IDs of shape [batch, sequence].

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional token embeddings of shape [batch, sequence, hidden].

attention_mask

torch.Tensor | NoneDefaults to None

Optional validity, indexed document, or backend mask of shape [batch, sequence] or [batch, 1, sequence, sequence].

position_ids

torch.Tensor | NoneDefaults to None

Optional positions of shape [batch, sequence] or [axes, batch, sequence].

cache_position

torch.Tensor | NoneDefaults to None

Optional token positions of shape [sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Optional padding mask of shape [batch, sequence].

past_key_values

Any | NoneDefaults to None

Unsupported recurrent or KV cache.

use_cache

bool | NoneDefaults to None

Whether to use a cache; only False or None is supported.

output_hidden_states

bool | NoneDefaults to None

Accepted for Hugging Face compatibility and ignored.

**attn_kwargs

AnyDefaults to {}

Backend-specific attention arguments, including optional _packed_seq_ids of shape [batch, sequence], packed_token_indices of shape [batch, sequence] or [tokens], cu_seqlens of shape [batch, max_documents + 1] or [documents + 1], and integer max_seqlen. Batch-major metadata uses -1 padding. Explicit token metadata is supported without context parallelism. The multimodal wrapper temporarily prefixes cu_seqlens and max_seqlen with _text_ while crossing the HF vision path; their layouts and values are unchanged.

Returns: BaseModelOutputWithPast

Model output whose last_hidden_state has shape [batch, sequence,

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseTextBackbone.get_input_embeddings() -> torch.nn.Module

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseTextBackbone.init_weights(
    buffer_device: torch.device | None = None
) -> None

nemo_automodel.components.models.qwen3_5.model.Qwen3_5DenseTextBackbone.set_input_embeddings(
    value: torch.nn.Module
) -> None

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    mtp_loss_scaling_factor: float = 0.1,
    num_nextn_predict_layers: int | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, Module

Qwen3.5 dense causal LM with optional Megatron-style MTP head.

_tied_weights_keys

= {'lm_head.weight': 'model.embed_tokens.weight'}

backend

= _qwen3_5_backend(backend)

lm_head

model

= Qwen3_5DenseTextBackbone(config, self.backend)

mtp

mtp_config

packed_mask_type

PackedMaskType

Request compact NEAT document IDs when TE consumes packed offsets.

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.BOTH

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: typing.Any | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.qwen3_5.model.Qwen3_5CausalLMOutputWithPast

Compute causal-LM logits and optional multi-token prediction states.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Optional token IDs of shape [batch, sequence].

attention_mask

torch.Tensor | NoneDefaults to None

Optional document/validity mask of shape [batch, sequence] or block-causal mask of shape [batch, 1, sequence, sequence].

position_ids

torch.LongTensor | NoneDefaults to None

Optional text positions of shape [batch, sequence] or multi-axis positions of shape [3, batch, sequence] or [4, batch, sequence] including a leading text-position axis.

past_key_values

Any | NoneDefaults to None

Optional HF cache; cached generation is unsupported.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional embeddings of shape [batch, sequence, hidden].

labels

torch.LongTensor | NoneDefaults to None

Optional targets of shape [batch, sequence], ignored here.

use_cache

bool | NoneDefaults to None

Whether to cache states; only False or None is supported.

logits_to_keep

int | torch.TensorDefaults to 0

Trailing token count (zero means all), or position indices of shape [selected_tokens] indexing the sequence axis.

**kwargs

AnyDefaults to {}

Backbone arguments, including _packed_seq_ids of shape [batch, sequence], packed_token_indices of shape [batch, sequence] or [tokens], cu_seqlens of shape [batch, max_documents + 1] or [documents + 1], and integer max_seqlen. Batch-major metadata uses -1 padding and is removed before the separate MTP sublayers.

Returns: Qwen3_5CausalLMOutputWithPast

Logits of shape [batch, selected_sequence, vocab], final hidden states

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.from_config(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

classmethod

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
)

classmethod

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.get_input_embeddings() -> torch.nn.Module

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.get_output_embeddings() -> torch.nn.Module

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.set_input_embeddings(
    value: torch.nn.Module
) -> None

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.set_output_embeddings(
    new_embeddings: torch.nn.Module
) -> None

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM.tie_weights(
    _args: set[str] | bool | None = (),
    _kwargs: set[str] | bool | None = {}
) -> None

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    mtp_loss_scaling_factor: float = 0.1,
    num_nextn_predict_layers: int | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, HFQwen3_5ForConditionalGeneration

Qwen3.5/Qwen3.6 dense VLM with optional Megatron-style MTP head.

The base VLM retains upstream image/video and M-RoPE semantics, with model-owned host metadata and media insertion paths to avoid CUDA syncs. MTP is added as an auxiliary train-time module over the final language hidden states, matching the dense text-only MTP architecture.

_keep_in_fp32_modules

list[str] = keep_fp32

_pp_keep_self_forward

bool = True

_tied_weights_keys

backend

= _qwen3_5_backend(backend)

lm_head

= self.lm_head.to(dtype)

mtp

mtp_config

packed_mask_type

PackedMaskType

Request compact NEAT document IDs when TE consumes packed offsets.

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.BOTH

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration._embed_and_splice_for_cp(
    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
) -> torch.Tensor

Embed token ids and splice image/video features into the embeddings.

The VLM->LM multimodal scatter runs on the full (unsharded) sequence inside the forward before the CP sequence shard, so it is identical to the pre-CP-refactor pre-embed. Vision features may be frame-sharded across the CP group before the fixed-size splice inserts them into the full input_ids sequence.

Parameters:

input_ids

torch.Tensor

Token ids [batch, sequence] (full, unsharded).

pixel_values

torch.Tensor | None

Optional packed image patches for HF vision encoding.

pixel_values_videos

torch.Tensor | None

Optional packed video patches.

image_grid_thw

torch.Tensor | None

Per-image [num_images, 3] T/H/W grid.

video_grid_thw

torch.Tensor | None

Per-video [num_videos, 3] T/H/W grid.

Returns: torch.Tensor

inputs_embeds of shape [batch, sequence, hidden] with image /

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration._encode_vision_for_cp(
    pixel_values: torch.Tensor,
    grid_thw: torch.Tensor,
    is_video: bool
) -> torch.Tensor

Encode one modality’s patches into flat, entry-ordered visual tokens under CP.

Under context parallelism the vision tower otherwise runs redundantly on the full, un-sharded patch set on every CP rank. When the CP vision-frame sharding group is published (see set_cp_vision_group), route the single self.model.visual call through maybe_distribute_visual, which shards the frames across the CP group and all-gathers the per-frame embeds in original entry order. Its pooler_output is the flat, already-concatenated tensor, identical to torch.cat(get_image_features(...).pooler_output, dim=0) (HF get_image_features splits that same flat tensor per entry). When sharding is disabled / inactive this is the exact replicated get_image_features / get_video_features path.

Parameters:

pixel_values

torch.Tensor

Patch rows of shape [total_patch_rows, patch_dim] for ALL entries of one modality, frame-contiguous in entry order.

grid_thw

torch.Tensor

[num_entries, 3] tensor of per-entry (t, h, w) grid sizes.

is_video

bool

Select the video (True) vs image (False) replicated fallback.

Returns: torch.Tensor

[total_patch_rows / spatial_merge_size**2, hidden] flat merged-token embeddings

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration._pop_staged_vlm_media(
    input_ids: torch.Tensor | None,
    kwargs: dict[str, typing.Any]
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

Resolve explicit media or consume the next staged pipeline chunk.

Staged media advances once per microbatch, including text-only inputs. Only explicitly passed media is filtered using token-presence metadata.

Parameters:

input_ids

torch.Tensor | None

Token IDs of shape [batch, sequence], or None.

kwargs

dict[str, Any]

Forward arguments with optional image and video patch tensors of shape [patches, patch_features] and grids of shape [entries, 3].

Returns: torch.Tensor | None

Image patches [image_patches, patch_features], video patches

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: typing.Any | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    pixel_values: torch.Tensor | None = None,
    pixel_values_videos: torch.FloatTensor | None = None,
    image_grid_thw: torch.LongTensor | None = None,
    video_grid_thw: torch.LongTensor | None = None,
    mm_token_type_ids: torch.IntTensor | None = None,
    use_cache: bool | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    padding_mask: torch.Tensor | None = None,
    mtp_per_depth_input_ids: tuple[torch.LongTensor, ...] | None = None,
    mtp_per_depth_position_ids: tuple[torch.LongTensor, ...] | None = None,
    mtp_per_depth_valid_masks: tuple[torch.BoolTensor, ...] | None = None,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.qwen3_5.model.Qwen3_5CausalLMOutputWithPast

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.from_config(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

classmethod

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
)

classmethod

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Per-stage input/output meta tensors for the PP schedule’s shape inference.

Matches the framework default (first stage consumes full token ids [mb, seq]; later stages consume hidden states; the last stage owning lm_head emits logits, earlier stages emit hidden states) except that under context parallelism the first stage embeds the full sequence and shards it in forward, so every stage output and later-stage input carries the LOCAL (padded-to-2*cp then //cp) sequence length while the first-stage input stays full-length. At cp_size == 1 this reduces to the default symmetric shapes.

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.prepare_host_batch_metadata(
    batch: dict[str, typing.Any]
) -> dict[str, typing.Any]

Retain host metadata outside distributed tensor input transfers.

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]

Return a sharder-only CP backend plus the full-sequence mRoPE positions.

Embedding and the VLM->LM multimodal scatter now run inside forward per microbatch (see _embed_and_splice_for_cp), so this hook only (a) computes the mRoPE position_ids on the full (unsharded) sequence via get_rope_index when the batch has none, and (b) returns the ContextParallelSharder. Packed batches (_packed_seq_ids or a 4D mask) get block-diagonal CP over contiguous chunks so attention and the GatedDeltaNet keep per-document boundaries; other batches get round-robin CP. input_ids and the media inputs are left in the batch for the forward; mm_token_type_ids is consumed here (only get_rope_index needs it) so the sharded forward never sees a full-length copy.

Parameters:

batch

dict[str, Any]

The batch dict (with input_ids and optional multimodal keys). input_ids is [batch, sequence]; _packed_seq_ids, when present, is [batch, sequence] with 1-based document IDs and zero padding.

num_chunks

intDefaults to 1

Accepted for hook-signature parity; unused.

Raises:

ValueError: Packed batches with a full-attention backend other than SDPA.
NotImplementedError: Packed batches with an MTP head.

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.prepare_mtp_inputs_for_cp(
    batch: dict[str, typing.Any],
    ignore_index: int = -100
) -> nemo_automodel.components.models.common.mtp.MTPContextParallelInputs

Prepare dense Qwen VLM future-token streams before CP sharding.

Call this after prepare_model_inputs_for_cp has materialized full-sequence mRoPE positions.

Parameters:

batch

dict[str, Any]

Full-sequence batch whose input_ids and labels tensors have shape [batch, sequence] and whose position_ids tensor has shape [axes, batch, sequence].

ignore_index

intDefaults to -100

Fill value for invalid targets at trailing and packed document-boundary positions.

Returns: MTPContextParallelInputs

Global per-depth token IDs, mRoPE position IDs, targets, and validity

nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration.tie_weights(
    _args: set[str] | bool | None = (),
    _kwargs: set[str] | bool | None = {}
) -> None

Tie lm_head to the active VLM text embedding when requested.

class nemo_automodel.components.models.qwen3_5.model.Qwen3_5Model()

Bases: HFQwen3_5Model

Thin VLM wrapper exposing language_model internals as properties and routing the forward: vision encoding and a fixed-size splice when media and token IDs are present, mask-based splicing for embedded media, and the NeMo dense backbone for text only.

nemo_automodel.components.models.qwen3_5.model.Qwen3_5Model._scatter_media_embeddings(
    inputs_embeds: torch.Tensor,
    media_embeds: torch.Tensor,
    media_mask: torch.Tensor
) -> torch.Tensor

staticmethod

Replace selected token embeddings without dynamic-size backward ops.

Parameters:

inputs_embeds

torch.Tensor

Token embeddings of shape [batch, sequence, hidden].

media_embeds

torch.Tensor

Media embeddings of shape [media_tokens, hidden], in the order of true entries in media_mask.

media_mask

torch.Tensor

Boolean mask of shape [batch, sequence, 1].

Returns: torch.Tensor

Combined token embeddings of shape [batch, sequence, hidden].

nemo_automodel.components.models.qwen3_5.model.Qwen3_5Model.forward(
    input_ids = None,
    attention_mask = None,
    position_ids = None,
    past_key_values = None,
    inputs_embeds = None,
    pixel_values = None,
    pixel_values_videos = None,
    image_grid_thw = None,
    video_grid_thw = None,
    cache_position = None,
    kwargs = {}
)

Route multimodal inputs without exposing text boundaries to vision.

Parameters:

input_ids

Defaults to None

Token IDs of shape [batch, sequence], or floating-point pipeline hidden states of shape [batch, sequence, hidden].

attention_mask

Defaults to None

Padding/document IDs of shape [batch, sequence] or a dense mask of shape [batch, 1, sequence, sequence].

position_ids

Defaults to None

Text positions of shape [batch, sequence] or mRoPE coordinates of shape [3 or 4, batch, sequence].

past_key_values

Defaults to None

Unsupported recurrent or KV cache.

inputs_embeds

Defaults to None

Optional embeddings of shape [batch, sequence, hidden].

pixel_values

Defaults to None

Optional image patches of shape [patches, patch_features].

pixel_values_videos

Defaults to None

Optional video patches of shape [patches, patch_features].

image_grid_thw

Defaults to None

Image grids of shape [images, 3] in temporal/height/width order.

video_grid_thw

Defaults to None

Video grids of shape [videos, 3] in temporal/height/width order.

cache_position

Defaults to None

Optional token positions of shape [sequence].

**kwargs

Defaults to {}

Text attention metadata, including cu_seqlens of shape [batch, max_documents + 1] or [documents + 1], integer max_seqlen, and packed_token_indices of shape [batch, sequence] or [tokens]. Image/video-prefixed metadata retains the inherited HF precomputed-vision contract.

Returns:

Model output with last_hidden_state of shape [batch, sequence,

nemo_automodel.components.models.qwen3_5.model.Qwen3_5Model.get_image_features(
    pixel_values: torch.Tensor,
    image_grid_thw: torch.Tensor | None = None,
    host_image_grid_thw: nemo_automodel.components.models.qwen3_5.packing.HostTensor | None = None,
    kwargs: typing.Any = {}
) -> tuple | transformers.modeling_outputs.BaseModelOutputWithPooling

Encode image patches, keeping SDPA segment lengths on the host.

Parameters:

pixel_values

torch.Tensor

Tensor of shape [patches, patch_features], in image order.

image_grid_thw

torch.Tensor | NoneDefaults to None

Integer tensor of shape [images, 3], with temporal, height, and width grid sizes in that order.

host_image_grid_thw

HostTensor | NoneDefaults to None

Matching CPU grid retained before the distributed wrapper moves tensor inputs to CUDA.

**kwargs

AnyDefaults to {}

Additional Hugging Face vision-forward arguments.

Returns: tuple | BaseModelOutputWithPooling

Vision output whose last_hidden_state has shape [patches, hidden]

nemo_automodel.components.models.qwen3_5.model.Qwen3_5Model.get_placeholder_mask(
    input_ids: torch.Tensor | None,
    inputs_embeds: torch.Tensor,
    image_features: torch.Tensor | None = None,
    video_features: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]

Return image and video masks while checking token counts on device.

Parameters:

input_ids

torch.Tensor | None

Token IDs of shape [batch, sequence], or None when only embeddings are available.

inputs_embeds

torch.Tensor

Token embeddings of shape [batch, sequence, hidden].

image_features

torch.Tensor | NoneDefaults to None

Optional image features of shape [image_tokens, hidden].

video_features

torch.Tensor | NoneDefaults to None

Optional video features of shape [video_tokens, hidden].

Returns: tuple[torch.Tensor, torch.Tensor]

Image and video boolean masks, each of shape [batch, sequence, 1].

class nemo_automodel.components.models.qwen3_5.model._Qwen3_5Attention(
    args: typing.Any = (),
    kwargs: typing.Any = {}
)

Bases: Qwen3NextAttention

Qwen3.5 full attention with packed block-diagonal CP dispatch.

_base_attn_func

= self.attn_func

attn_func

= self._dispatch_attention

nemo_automodel.components.models.qwen3_5.model._Qwen3_5Attention._dispatch_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

Route preprocessed QKV through packed CP or the configured backend.

Parameters:

query

torch.Tensor

Query states of shape [batch, heads, local_sequence, head_dim].

key

torch.Tensor

Key states of shape [batch, kv_heads, local_sequence, head_dim].

value

torch.Tensor

Value states with the same layout as key.

**attn_kwargs

AnyDefaults to {}

Keyword arguments produced by the parent attention’s backend preprocessing.

Returns: torch.Tensor

Attention output of shape [batch, heads, local_sequence, head_dim].

nemo_automodel.components.models.qwen3_5.model._default_init_device() -> torch.device

nemo_automodel.components.models.qwen3_5.model._dense_moe_config(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    dtype: torch.dtype
) -> nemo_automodel.components.moe.layers.MoEConfig

Trivial MoEConfig for the dense Qwen3.5 backbone.

The dense model has no experts (num_experts is 0/absent), so Block builds a dense MLP and never consults this config; it is only required to satisfy Block.__init__’s signature.

nemo_automodel.components.models.qwen3_5.model._make_full_attention_config(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    layer_idx: int
) -> transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig

nemo_automodel.components.models.qwen3_5.model._mask_mtp_embed_inputs(
    embed_inputs: tuple[torch.Tensor, ...],
    valid_masks: tuple[torch.BoolTensor, ...]
) -> tuple[torch.Tensor, ...]

Zero invalid future-token embeddings in the local CP layout.

Parameters:

embed_inputs

tuple[torch.Tensor, ...]

Per-depth embeddings of shape [batch, sequence, hidden].

valid_masks

tuple[torch.BoolTensor, ...]

Per-depth masks of shape [batch, sequence] in the same local CP token layout as embed_inputs.

Returns: torch.Tensor

Per-depth embeddings with invalid trailing or packed-boundary positions

nemo_automodel.components.models.qwen3_5.model._mtp_block_causal_mask(
    packing_mask: torch.Tensor,
    inputs_embeds: torch.Tensor
) -> torch.Tensor

Build a 4D block-causal attention mask from an indexed packing mask.

packing_mask is [B, S] with the 1-based document index per token (0 = padding). The returned bool mask [B, 1, S, S] (True = attend) keeps attention causal and within each packed document, matching the backbone’s packed-sequence semantics. Used for the MTP sublayers, which run SDPA self-attention over the same packed batch (NVBugs 6330129).

nemo_automodel.components.models.qwen3_5.model._qwen3_5_backend(
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.common.BackendConfig

Return a Qwen3.5 backend with TE fused RoPE disabled.

Qwen3.5 VLM training can feed full-attention layers in packed/THD shape via the shared Qwen3-Next attention block. TE fused RoPE expects 4D inputs there, so keep the non-fused RoPE path while preserving the rest of the backend selection (TE Linear, attention backend, etc.).

nemo_automodel.components.models.qwen3_5.model._resolve_mtp_num_layers(
    config: typing.Any,
    override: int | None = None
) -> int

nemo_automodel.components.models.qwen3_5.model._rolled_embed_inputs(
    inputs_embeds: torch.Tensor,
    num_depths: int
) -> tuple[torch.Tensor, ...]

nemo_automodel.components.models.qwen3_5.model._splice_multimodal_features(
    inputs_embeds: torch.Tensor,
    input_ids: torch.Tensor,
    features: torch.Tensor,
    token_id: int,
    modality: str
) -> torch.Tensor

Replace multimodal placeholder rows without dynamic-size indexing.

masked_scatter implements its source gradient through masked_select. On CUDA that dynamic-size selection synchronizes the stream before allocating its result. The feature row count is already known, so use fixed-size nonzero_static indices and index_copy instead. Its backward is the asynchronous index_fill/index_select pair.

Parameters:

inputs_embeds

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading token dimensions.

input_ids

torch.Tensor

Tensor of shape […], matching the leading dimensions of inputs_embeds and on the same device.

features

torch.Tensor

Tensor of shape [num_placeholder_tokens, hidden], with the same device and dtype as inputs_embeds. Rows follow the flattened placeholder order.

token_id

int

Token ID marking the positions to replace.

modality

str

Modality name used in the count-mismatch error.

Returns: torch.Tensor

Tensor of shape […, hidden] with placeholder rows replaced by features.

Raises:

ValueError: The placeholder count differs from num_placeholder_tokens.

nemo_automodel.components.models.qwen3_5.model._split_qwen3_5_position_ids(
    position_ids: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    device: torch.device,
    past_key_values: typing.Any | None = None
) -> tuple[torch.Tensor, torch.Tensor | None]

nemo_automodel.components.models.qwen3_5.model.build_mtp_config_from_hf(
    config: typing.Any,
    loss_scaling_factor: float = 0.1,
    num_nextn_predict_layers: int | None = None
) -> nemo_automodel.components.models.common.mtp.MTPConfig

Build Qwen3.5 MTP runtime config from HF-style config fields.

nemo_automodel.components.models.qwen3_5.model.build_qwen3_5_dense_mtp(
    config: transformers.models.qwen3_5.configuration_qwen3_5.Qwen3_5TextConfig,
    mtp_config: nemo_automodel.components.models.common.mtp.MTPConfig,
    dtype: torch.dtype
) -> nemo_automodel.components.models.common.mtp.MTPModule

Construct dense Qwen3.5 MTP blocks.

Qwen3.5 MTP follows Megatron Bridge: each depth is one full-attention Qwen3.5 decoder block, regardless of the backbone’s GatedDeltaNet layers.

nemo_automodel.components.models.qwen3_5.model.ModelClass = Qwen3_5ForCausalLM