nemo_automodel.components.models.qwen3_8_flash_next.model

Trainable Qwen3.8-Flash-Next conditional-generation model.

This implementation uses the checkpoint’s compressed-block QSA router and a FlexAttention sparse-GQA path for CUDA BF16 long-sequence SFT, with a PyTorch oracle for CPU and numerical parity. Pipeline and tensor parallelism remain unsupported. Context parallelism uses a model-owned contiguous sequence shard for QSA, GDN, and PLE, and composes with sequence packing.

Module Contents

Classes

Name	Description
`Qwen3_8_FlashNextCausalLMOutput`	Causal-LM output with optional per-layer HC states for parity capture.
`Qwen3_8_FlashNextForConditionalGeneration`	Trainable language-only Qwen3.8-Flash-Next causal-LM wrapper.
`Qwen3_8_FlashNextModel`	Language-only Qwen3.8-Flash-Next decoder wrapper.
`Qwen3_8_FlashNextTextModelBackend`	Qwen3.8-Flash-Next text decoder with four HC streams and one PLE layer.

Functions

Name	Description
`_default_owner_group`	-
`_qwen3_8_flash_next_backend`	Return a backend whose rotary path supports text and multimodal layouts.
`_resolve_model_dtype`	-

Data

ModelClass

API

class nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextCausalLMOutput()

Dataclass

Bases: CausalLMOutputWithPast

Causal-LM output with optional per-layer HC states for parity capture.

class nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration(
    config: nemo_automodel.components.models.qwen3_8_flash_next.config.Qwen3_8_FlashNextConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    engram_process_group: torch.distributed.ProcessGroup | None = None,
    engram_table_config: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Trainable language-only Qwen3.8-Flash-Next causal-LM wrapper.

_keep_in_fp32_modules_strict

= ['_fp32_params']

_packed_cp_attn_backends

= ('flex',)

backend

= _qwen3_8_flash_next_backend(backend)

lm_head

model

moe_config

= self.model.language_model.moe_config

pad_token_id

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vocab_size

= config.text_config.vocab_size

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration._nemo_prepare_model_owned_dtensors(
    fsdp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> set[torch.nn.Parameter]

Convert PLE owner shards to global DTensors before FSDP wrapping.

Parameters:

fsdp_mesh

DeviceMesh

Flattened data/context shard mesh whose rank order is identical to the PLE owner process group.

Returns: set[nn.Parameter]

The exact registered parameter identities that FSDP must ignore

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    past_key_values: object | None = None,
    use_cache: bool | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool | None = None,
    _qwen3_8_flash_next_cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextCausalLMOutput

Run language-only generation and project final HC-mixed states.

Parameters:

input_ids

torch.Tensor

Raw tokenizer IDs of shape [batch, sequence].

attention_mask

torch.Tensor | NoneDefaults to None

Optional token-validity mask of shape [batch, sequence] or backend-specific attention mask.

position_ids

torch.Tensor | NoneDefaults to None

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

inputs_embeds

torch.Tensor | NoneDefaults to None

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

labels

torch.Tensor | NoneDefaults to None

Optional labels of shape [batch, sequence]. They are accepted for recipe compatibility; loss is computed externally.

past_key_values

object | NoneDefaults to None

Cache state; unsupported.

use_cache

bool | NoneDefaults to None

Cache request; True is unsupported.

logits_to_keep

int | torch.TensorDefaults to 0

0 for all positions, a positive trailing count, or an integer tensor of shape [kept_sequence] containing explicit sequence indices.

output_hidden_states

bool | NoneDefaults to None

Explicit True returns the embedding, every per-layer HC state, and the final state. When omitted, config.output_hidden_states=True returns only the final state for fused linear cross entropy without retaining every decoder activation.

_qwen3_8_flash_next_cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Internal contiguous CP metadata with replicated raw-ID/padding tensors of shape [batch, global_sequence].

**kwargs

AnyDefaults to {}

Text-attention backend metadata.

Returns: Qwen3_8_FlashNextCausalLMOutput

Causal-LM output with logits

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.from_config(
    config: nemo_automodel.components.models.qwen3_8_flash_next.config.Qwen3_8_FlashNextConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration

classmethod

Construct from a parsed Qwen3.8-Flash-Next configuration.

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration

classmethod

Construct architecture from a local/HF config before checkpoint load.

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.get_input_embeddings() -> torch.nn.Module

Return the text token embedding module.

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.get_output_embeddings() -> torch.nn.Module

Return the untied LM projection.

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

Initialize checkpoint-free model weights.

Parameters:

buffer_device

torch.device | NoneDefaults to None

Target device for backend initializers.

dtype

torch.dtypeDefaults to torch.bfloat16

Final model parameter dtype, excluding intrinsic fp32 GDN state.

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

Return Qwen3.8-Flash-Next’s contiguous model-owned CP batch sharder.

Parameters:

batch

dict[str, Any]

Full-sequence batch. Token-aligned tensors have shape [batch, global_sequence, ...] and remain unchanged until the returned sharder is invoked.

num_chunks

intDefaults to 1

Accepted for the framework hook contract; Qwen3.8-Flash-Next shards each non-packed batch directly and does not use it.

Returns: dict[str, Any]

Mapping containing one unresolved cp_sharder. Once invoked it

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

Replace the text token embedding module.

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextForConditionalGeneration.set_output_embeddings(
    value: torch.nn.Module
) -> None

Replace the untied LM projection.

class nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextModel(
    config: nemo_automodel.components.models.qwen3_8_flash_next.config.Qwen3_8_FlashNextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    moe_overrides: dict[str, typing.Any] | None = None,
    engram_process_group: torch.distributed.ProcessGroup | None = None,
    engram_table_config: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig | None = None
)

Bases: Module

Language-only Qwen3.8-Flash-Next decoder wrapper.

language_model

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextModel.forward(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    past_key_values: object | None = None,
    output_hidden_states: bool | None = None,
    _qwen3_8_flash_next_cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Run the language-only Qwen3.8-Flash-Next decoder.

Parameters:

input_ids

torch.Tensor

Raw IDs of shape [batch, sequence]; required for PLE hashing.

attention_mask

torch.Tensor | NoneDefaults to None

Optional token-validity mask of shape [batch, sequence] or backend-specific attention mask.

position_ids

torch.Tensor | NoneDefaults to None

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

inputs_embeds

torch.Tensor | NoneDefaults to None

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

past_key_values

object | NoneDefaults to None

Cache state; unsupported for training.

output_hidden_states

bool | NoneDefaults to None

Capture decoder HC states.

_qwen3_8_flash_next_cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Internal contiguous CP metadata with replicated raw-ID/padding tensors of shape [batch, global_sequence].

**kwargs

AnyDefaults to {}

Text-attention backend arguments.

Returns: BaseModelOutputWithPast

Base-model output whose final text states have shape

class nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextTextModelBackend(
    config: nemo_automodel.components.models.qwen3_8_flash_next.config.Qwen3_8_FlashNextTextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    moe_overrides: dict[str, typing.Any] | None = None,
    engram_process_group: torch.distributed.ProcessGroup | None = None,
    engram_table_config: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig | None = None
)

Bases: Module

Qwen3.8-Flash-Next text decoder with four HC streams and one PLE layer.

Parameters:

config

Qwen3_8_FlashNextTextConfig

Text architecture configuration.

backend

BackendConfig

Native attention, linear, and MoE backend configuration.

moe_config

MoEConfig | NoneDefaults to None

Optional native MoE configuration override.

moe_overrides

dict[str, Any] | NoneDefaults to None

Optional fields merged into the default MoE config.

engram_process_group

dist.ProcessGroup | NoneDefaults to None

Group that owns contiguous PLE table row shards. With the supported no-PP/no-TP topology this is the EP/world group.

engram_table_config

Qwen3_8_FlashNextEngramTableConfig | NoneDefaults to None

Optional tiny-table override for unit tests. The released checkpoint uses [320001536, 160] globally.

embed_tokens

hyper_connection_mixer

layers

= nn.ModuleDict()

model_dtype

= _resolve_model_dtype(config)

moe_config

= moe_config or MoEConfig(**moe_defaults)

padding_idx

= config.pad_token_id

rotary_emb

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextTextModelBackend.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,
    padding_mask: torch.Tensor | None = None,
    past_key_values: object | None = None,
    use_cache: bool | None = None,
    output_hidden_states: bool | None = None,
    _qwen3_8_flash_next_cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Run the HC decoder and final HC mixer.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Raw tokenizer IDs of shape [batch, sequence]. They remain required when inputs_embeds is supplied because PLE hashes the raw IDs.

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional precomputed token/vision embeddings of shape [batch, sequence, hidden_size].

attention_mask

torch.Tensor | NoneDefaults to None

Optional mask of shape [batch, sequence] or a backend-specific attention mask.

position_ids

torch.Tensor | NoneDefaults to None

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

padding_mask

torch.Tensor | NoneDefaults to None

Optional mask [batch, sequence] where True is padding.

past_key_values

object | NoneDefaults to None

KV/SSM cache; unsupported by this training backend.

use_cache

bool | NoneDefaults to None

Cache request; True is unsupported.

output_hidden_states

bool | NoneDefaults to None

Include embedding, per-layer HC, and final collapsed states for parity diagnostics.

_qwen3_8_flash_next_cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Internal contiguous CP metadata. Its replicated raw-ID and padding tensors have shape [batch, global_sequence]; model activations remain local [batch, sequence, hidden] tensors.

**attn_kwargs

AnyDefaults to {}

Attention backend metadata.

Returns: BaseModelOutputWithPast

A base-model output whose last_hidden_state has shape

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextTextModelBackend.get_input_embeddings() -> torch.nn.Module

Return the raw-token embedding table.

nemo_automodel.components.models.qwen3_8_flash_next.model.Qwen3_8_FlashNextTextModelBackend.init_weights(
    buffer_device: torch.device
) -> None

Initialize decoder weights for training from scratch.

Parameters:

buffer_device

torch.device

Device used by layer initializers.

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

Replace the raw-token embedding table.

Parameters:

value

nn.Module

Module mapping [batch, sequence] IDs to [batch, sequence, hidden_size] embeddings.

nemo_automodel.components.models.qwen3_8_flash_next.model._default_owner_group() -> torch.distributed.ProcessGroup | None

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

Return a backend whose rotary path supports text and multimodal layouts.

nemo_automodel.components.models.qwen3_8_flash_next.model._resolve_model_dtype(
    config: object
) -> torch.dtype

nemo_automodel.components.models.qwen3_8_flash_next.model.ModelClass = Qwen3_8_FlashNextForConditionalGeneration