nemo_automodel.components.models.qwen3_5_moe.model
nemo_automodel.components.models.qwen3_5_moe.model
Qwen3.5-MoE (VL) NeMo Automodel support.
Module Contents
Classes
Functions
Data
API
Bases: Qwen3_5MoeTextRotaryEmbedding
Ensure inv_freq stays in float32 across .to(dtype) calls.
Bases: Qwen3_5MoeVisionRotaryEmbedding
Ensure the vision rotary inv_freq buffer remains float32.
Bases: Block
Block that uses the Qwen3.5-MoE native GatedDeltaNet (separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a)
Mirror :meth:Block.forward but thread NEAT-packing kwargs into
CPAwareGatedDeltaNet.
The parent Block.forward calls linear_attn with only
hidden_states and attention_mask; for packed sequences the
gated_delta_rule kernel additionally needs cu_seqlens /
indices to reset state at document boundaries (issue #2131).
Derived once per forward from the indexed attention mask.
Bases: CausalLMOutputWithPast
Qwen3.5-MoE output extended with MTP auxiliary hidden states.
Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin
Text-only Qwen3.5-MoE causal language model.
Run the text-only causal LM forward pass.
Parameters:
Token ids of shape [batch, sequence]. On a non-first
pipeline stage where embed_tokens is absent, this argument
instead carries hidden states of shape
[batch, sequence, hidden] and is routed as inputs_embeds.
Optional mask of shape [batch, sequence] or
[batch, 1, sequence, sequence].
Optional M-RoPE positions of shape
[axes, batch, sequence] or [batch, sequence].
Optional KV-cache state; caching is not supported
by this backend and a non-None value raises.
Optional hidden inputs of shape
[batch, sequence, hidden].
Optional labels of shape [batch, sequence]. Accepted for
Hugging Face compatibility; loss is computed outside this model.
Whether to use a KV cache; True is unsupported.
0 to project every position, a positive integer
to keep the last positions, or indices of shape
[kept_sequence].
Whether to include final decoder states in the output for the fused loss path.
Optional token positions of shape [sequence].
Optional padding mask of shape [batch, sequence]
where True marks padding.
Additional attention and packed-sequence metadata forwarded to the decoder and optional MTP layers.
Returns: Qwen3_5MoeCausalLMOutputWithPast
A causal-LM output whose logits have shape
Bases: HFCheckpointingMixin, HFQwen3_5MoeForConditionalGeneration, MoEFSDPSyncMixin
Qwen3.5-MoE VL conditional generation model using NeMo backend components.
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, identical to the pre-CP-refactor pre-embed.
Parameters:
Token ids [batch, sequence] (full, unsharded).
Optional packed image patches for HF vision encoding.
Optional packed video patches.
Per-image [num_images, 3] T/H/W grid.
Per-video [num_videos, 3] T/H/W grid.
Returns: torch.Tensor
inputs_embeds of shape [batch, sequence, hidden] with image /
Encode one modality into flat, entry-ordered visual tokens under CP.
Parameters:
Patch rows of shape [total_patch_rows, patch_dim]
for all image or video entries, frame-contiguous in entry order.
Per-entry temporal/height/width grids of shape
[num_entries, 3].
Whether the replicated fallback uses the video feature helper instead of the image feature helper.
Returns: torch.Tensor
Flat merged-token embeddings of shape [visual_tokens, hidden]
Run Qwen3.5-MoE VLM generation with model-owned CP embedding.
Parameters:
Token ids of shape [batch, sequence]. Under CP this
remains full-length until multimodal features are spliced.
Optional mRoPE positions of shape
[axes, batch, sequence]; under CP, sequence is local.
Optional mask of shape [batch, sequence] or
[batch, 1, sequence, sequence].
Optional local padding mask of shape
[batch, sequence] where True marks padding.
Optional embeddings of shape
[batch, sequence, hidden].
Optional token positions of shape [sequence].
Number of trailing logits to retain, or indices of
shape [kept_sequence].
Whether to expose final decoder states for the fused loss path.
Model inputs including image/video patch rows of shape
[total_patch_rows, patch_dim], their grids of shape
[num_entries, 3], and optional packed document ids of shape
[batch, sequence].
Returns:
A causal-LM output whose logits have shape
Return model-owned PP microbatch split dims for keyword inputs.
Qwen3.5-MoE mRoPE position ids can arrive as [T/H/W, batch, seq] or
[text/T/H/W, batch, seq]. The PP schedule must split that tensor on
the batch axis, while all standard batch-major kwargs keep the framework
default.
Per-stage input/output meta tensors for the PP schedule’s shape inference.
Matches the framework default 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.
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 :meth:_embed_and_splice_for_cp), so this hook only
computes the mRoPE position_ids on the full (unsharded) sequence via
get_rope_index and returns them with a
:class:ContextParallelSharder. Packed SDPA batches select contiguous
block-diagonal CP; ordinary batches retain the round-robin sharder.
input_ids and media stay in the batch for the forward, while
mm_token_type_ids is consumed here (only get_rope_index needs it).
Parameters:
The full-sequence batch. Tensor fields consumed here are
input_ids of shape [batch, sequence], optional
attention_mask of shape [batch, sequence] or
[batch, 1, sequence, sequence], optional position_ids
of shape [axes, batch, sequence], and image/video grids of
shape [num_entries, 2 or 3].
Accepted for hook-signature parity; unused.
Returns: dict[str, Any]
A mapping containing the selected cp_sharder, full-sequence
Bases: HFQwen3_5MoeModel
Thin wrapper that exposes language_model internals as properties
expected by the NeMo training loop (e.g. model.layers).
Bases: Module
Qwen3.5-MoE text decoder rebuilt on top of the Qwen3-Next Block.
Bases: Qwen3NextAttention
Qwen3.5-MoE full attention with packed block-diagonal CP dispatch.
Route preprocessed QKV through packed CP or the configured backend.
Parameters:
Query states of shape [batch, heads, local_sequence, head_dim].
Key states of shape [batch, kv_heads, local_sequence, head_dim].
Value states with the same layout as key.
Keyword arguments produced by the parent attention’s backend preprocessing.
Returns: torch.Tensor
Attention output of shape
Return a Qwen3.5-MoE backend with TE fused RoPE disabled.
The Qwen3.5 full-attention blocks reuse Qwen3-Next attention, and VLM/packed execution can present THD-shaped q/k tensors. TE fused RoPE expects 4D inputs in this path, so use non-fused RoPE while preserving the rest of the backend.
Build Qwen3.5-MoE MTP runtime config from HF-style config fields.
Construct Qwen3.5-MoE MTP blocks.