nemo_automodel.components.models.qwen3.model

Dense Qwen3 implementation with packed THD context parallelism.

Module Contents

Classes

Name	Description
`Qwen3Attention`	HuggingFace Qwen3 attention extended with a packed THD TE path.
`Qwen3DecoderLayer`	Qwen3 decoder layer using the packed-aware attention implementation.
`Qwen3ForCausalLM`	Dense Qwen3 causal LM with packed THD context parallelism.
`Qwen3Model`	Dense Qwen3 decoder supporting padded BSHD and packed THD layouts.
`Qwen3PreTrainedModel`	Base class for the dense Qwen3 implementation.

Data

ModelClass

__all__

check_model_inputs

API

class nemo_automodel.components.models.qwen3.model.Qwen3Attention(
    config: transformers.Qwen3Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: HFQwen3Attention

HuggingFace Qwen3 attention extended with a packed THD TE path.

rope_fusion

= backend.rope_fusion

nemo_automodel.components.models.qwen3.model.Qwen3Attention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

Run Qwen3 attention over padded BSHD or packed THD states.

Parameters:

hidden_states

torch.Tensor

Hidden states [B, S, H] or packed local states [T, H]. B is batch, S is sequence, T is local total tokens, and H is hidden size.

position_embeddings

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

RoPE tensors (cos, sin) or fused (cos, sin, freqs_cis). Local tensors use [B, S, D] or [T, D] and the fused raw table uses [S, 1, 1, D].

attention_mask

torch.Tensor | None

Padded attention mask for BSHD; THD uses cumulative document lengths from kwargs.

past_key_values

Cache | NoneDefaults to None

Optional BSHD KV cache; unsupported for THD.

**kwargs

Unpack[TransformersKwargs]Defaults to {}

THD requires qkv_format='thd' and cu_seqlens [N + 1]; CP additionally supplies cp_size and cp_rank.

Returns: torch.Tensor

Attention output shaped like hidden_states and optional BSHD

class nemo_automodel.components.models.qwen3.model.Qwen3DecoderLayer(
    config: transformers.Qwen3Config,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: HFQwen3DecoderLayer

Qwen3 decoder layer using the packed-aware attention implementation.

attention_type

= config.layer_types[layer_idx]

self_attn

class nemo_automodel.components.models.qwen3.model.Qwen3ForCausalLM(
    config: transformers.Qwen3Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: HFCheckpointingMixin, Qwen3PreTrainedModel, GenerationMixin

Dense Qwen3 causal LM with packed THD context parallelism.

_keep_in_fp32_modules

= ['rotary_emb']

_pp_plan

= {'lm_head': (['hidden_states'], ['logits'])}

_tied_weights_keys

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

_tp_plan

= {'lm_head': 'colwise_rep'}

backend

= backend or BackendConfig()

lm_head

model

= Qwen3Model(config=config, backend=(self.backend))

state_dict_adapter

= Qwen3StateDictAdapter(config=(self.config))

tie_word_embeddings_support

TieSupport = TieSupport.BOTH

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3.model.Qwen3ForCausalLM.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

Run causal LM projection for BSHD or packed THD inputs.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Token IDs [B, S] or packed local IDs [T].

attention_mask

torch.Tensor | NoneDefaults to None

Optional padded mask. THD uses cu_seqlens.

position_ids

torch.LongTensor | NoneDefaults to None

Position IDs [B, S] or packed local IDs [T].

past_key_values

Cache | NoneDefaults to None

Optional BSHD KV cache; unsupported for THD.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional hidden inputs [B, S, H] or [T, H].

labels

torch.LongTensor | NoneDefaults to None

Optional labels [B, S] or packed [T].

use_cache

bool | NoneDefaults to None

Whether to update the BSHD KV cache.

output_attentions

bool | NoneDefaults to None

Whether to request attention outputs.

output_hidden_states

bool | NoneDefaults to None

Whether to return per-layer hidden states.

return_dict

bool | NoneDefaults to None

Whether to return CausalLMOutputWithPast.

cache_position

torch.LongTensor | NoneDefaults to None

Optional BSHD cache positions [S].

logits_to_keep

int | torch.TensorDefaults to 0

Positions to project from hidden size H to vocabulary size V.

**kwargs

Unpack[TransformersKwargs]Defaults to {}

THD metadata. cu_seqlens is [N + 1] and CP adds cp_size and cp_rank.

Returns: CausalLMOutputWithPast

Causal LM output with BSHD logits [B, S, V]. Packed local

nemo_automodel.components.models.qwen3.model.Qwen3ForCausalLM.get_input_embeddings() -> torch.nn.Module

nemo_automodel.components.models.qwen3.model.Qwen3ForCausalLM.get_output_embeddings() -> torch.nn.Module

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

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

nemo_automodel.components.models.qwen3.model.Qwen3ForCausalLM.tie_weights(
    _args: object = (),
    _kwargs: object = {}
) -> None

class nemo_automodel.components.models.qwen3.model.Qwen3Model(
    config: transformers.Qwen3Config,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Qwen3PreTrainedModel

Dense Qwen3 decoder supporting padded BSHD and packed THD layouts.

embed_tokens

has_sliding_layers

= 'sliding_attention' in self.config.layer_types

layers

norm

padding_idx

= config.pad_token_id

rotary_emb

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3.model.Qwen3Model.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: transformers.cache_utils.Cache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,
    return_dict: bool | None = None,
    cache_position: torch.LongTensor | None = None,
    kwargs: transformers.processing_utils.Unpack[transformers.utils.TransformersKwargs] = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Run the Qwen3 decoder in padded BSHD or packed THD layout.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Token IDs [B, S] or packed local IDs [T].

attention_mask

torch.Tensor | NoneDefaults to None

Optional padded mask [B, S]. THD uses packed document boundaries instead.

position_ids

torch.LongTensor | NoneDefaults to None

Position IDs [B, S] or packed local IDs [T].

past_key_values

Cache | NoneDefaults to None

Optional BSHD generation cache; unsupported for THD.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Alternative hidden inputs [B, S, H] or [T, H].

use_cache

bool | NoneDefaults to None

Whether to update the BSHD KV cache.

output_attentions

bool | NoneDefaults to None

Whether to request attention outputs.

output_hidden_states

bool | NoneDefaults to None

Whether to retain per-layer hidden states.

return_dict

bool | NoneDefaults to None

Whether to return BaseModelOutputWithPast.

cache_position

torch.LongTensor | NoneDefaults to None

Optional BSHD cache positions [S].

**kwargs

Unpack[TransformersKwargs]Defaults to {}

THD metadata including qkv_format, cu_seqlens, max_seqlen, cp_size, and cp_rank.

Returns: BaseModelOutputWithPast

Decoder output with final states [B, S, H] or packed local

class nemo_automodel.components.models.qwen3.model.Qwen3PreTrainedModel()

Bases: PreTrainedModel

Base class for the dense Qwen3 implementation.

_can_record_outputs

_no_split_modules

= ['Qwen3DecoderLayer']

_skip_keys_device_placement

= ['past_key_values']

base_model_prefix

= 'model'

nemo_automodel.components.models.qwen3.model.ModelClass = Qwen3ForCausalLM

nemo_automodel.components.models.qwen3.model.__all__ = ['Qwen3ForCausalLM']

nemo_automodel.components.models.qwen3.model.check_model_inputs = get_check_model_inputs_decorator()