nemo_automodel.components.models.glm_moe_dsa.model

Module Contents

Classes

Name	Description
`Block`	-
`GlmMoeDsaForCausalLM`	-
`GlmMoeDsaModel`	-

Functions

Name	Description
`_uses_indexshare`	Return whether the model has layers that reuse another layer’s DSA indices.

Data

ModelClass

API

class nemo_automodel.components.models.glm_moe_dsa.model.Block(
    layer_idx: int,
    config: transformers.models.glm_moe_dsa.configuration_glm_moe_dsa.GlmMoeDsaConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

input_layernorm

mlp

= MoE(moe_config, backend)

post_attention_layernorm

self_attn

skip_topk

nemo_automodel.components.models.glm_moe_dsa.model.Block._mlp(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor

nemo_automodel.components.models.glm_moe_dsa.model.Block.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    prev_topk_indices: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]

Run the block and return (hidden_states, topk_indices).

topk_indices is this layer’s DSA selection — freshly computed on “full” layers, or prev_topk_indices passed through on “shared” layers — so the caller can thread it to subsequent shared layers (GLM IndexShare).

nemo_automodel.components.models.glm_moe_dsa.model.Block.init_weights(
    buffer_device: torch.device
)

class nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM(
    config: transformers.models.glm_moe_dsa.configuration_glm_moe_dsa.GlmMoeDsaConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

_packed_cp_attn_backends

= ('tilelang', 'cudnn')

backend

= backend or BackendConfig()

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM._is_pipeline_parallel_stage() -> bool

True when this module is a trimmed pipeline-parallel stage (not the whole model).

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    carry: torch.Tensor = (),
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | tuple[torch.Tensor, ...] | torch.Tensor

Forward pass.

Single process (no pipeline parallelism): returns :class:~transformers.modeling_outputs.CausalLMOutputWithPast, threading the IndexShare top-k internally (seeded None).

Pipeline parallelism: input_ids is the upstream hidden state on non-first stages. IndexShare models additionally use *carry for the previous stage’s running top-k selection. Non-last stages return the declared pipeline outputs and the last stage returns the logits tensor.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Token-ID tensor of shape [batch, sequence] (BSHD) or [1, tokens] (packed THD) on the first stage. On later stages, the upstream hidden-state tensor has shape [batch, sequence, hidden] or [tokens, hidden].

carry

torch.TensorDefaults to ()

Optional top-k tensor carried from the previous pipeline stage. Optimized THD backends use float32 [tokens, 1, index_topk]; dense BSHD uses float32 [batch, sequence, min(index_topk, sequence)].

position_ids

torch.Tensor | NoneDefaults to None

Optional position-ID tensor of shape [batch, sequence] or [1, tokens] for packed THD.

attention_mask

torch.Tensor | NoneDefaults to None

Optional key-mask tensor of shape [batch, sequence] or additive mask of shape [batch, 1, sequence, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Optional padding-mask tensor of shape [batch, sequence] or [1, tokens] for packed THD.

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

If 0, project all positions; else only the last logits_to_keep. A tensor value contains the one-dimensional token indices to project.

output_hidden_states

bool | NoneDefaults to None

When set (single-process), carry final hidden states on the output.

**attn_kwargs

AnyDefaults to {}

Additional attention metadata forwarded to the base model. Packed THD uses an int32 cu_seqlens tensor of shape [sequences + 1] and qkv_format="thd".

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

Single-process execution returns a causal-LM output whose logits tensor has shape

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.from_config(
    config: transformers.models.glm_moe_dsa.configuration_glm_moe_dsa.GlmMoeDsaConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.get_input_embeddings()

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.get_output_embeddings()

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

Declare PP inter-stage I/O metas, adding a top-k carry only for IndexShare models.

IndexShare models additionally receive and emit the previous “full” layer’s top-k selection so a stage that begins with a “shared” layer has the indices it needs. Models with all-full indexers need only the hidden-state channel.

Parameters:

is_first

bool

Whether this module is the first pipeline stage.

microbatch_size

int

Number of sequences in the pipeline microbatch.

seq_len

int

Sequence or packed-token length represented by the metadata.

dtype

torch.dtype

Hidden-state and logits dtype.

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

Pair of input and output metadata tuples. Optimized THD stages use hidden-state

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

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

Attach GLM DSA’s packed THD context-parallel batch sharder.

Parameters:

batch

dict[str, Any]

The batch dict.

num_chunks

intDefaults to 1

Number of chunks for load-balanced CP sharding.

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.set_input_embeddings(
    value
)

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.set_output_embeddings(
    new_embeddings
)

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.should_pack_validation_with_training() -> bool

Return whether validation must use the optimized packed THD layout.

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaForCausalLM.update_moe_gate_bias() -> None

Delegate the noaux router correction-bias update to the inner model.

class nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaModel(
    config: transformers.models.glm_moe_dsa.configuration_glm_moe_dsa.GlmMoeDsaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig | None = None,
    moe_overrides: dict | None = None
)

Bases: Module

embed_tokens

freqs

layers

= torch.nn.ModuleDict()

max_seq_len

= config.max_position_embeddings

moe_config

= moe_config or MoEConfig(**moe_defaults)

norm

qk_rope_head_dim

= config.qk_rope_head_dim

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaModel.forward(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    prev_topk_indices: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

Run the decoder stack, returning (hidden_states, topk_indices).

Parameters:

input_ids

torch.Tensor

Token IDs with shape [batch, sequence] (or packed [tokens]), or previous-stage hidden states with the matching token axes plus [hidden] when this pipeline stage has no embedding.

position_ids

torch.Tensor | NoneDefaults to None

Optional integer positions with the same token axes as input_ids. Packed THD callers supply shape [tokens].

attention_mask

torch.Tensor | NoneDefaults to None

Optional token mask with shape [batch, sequence] or additive attention mask broadcastable to [batch, heads, query, key].

padding_mask

torch.Tensor | NoneDefaults to None

Optional boolean mask with the input token axes; True marks padding. Packed THD callers supply shape [tokens].

prev_topk_indices

torch.Tensor | NoneDefaults to None

Optional previous pipeline stage’s IndexShare selection, shaped [batch, sequence, K] or packed [tokens, 1, K]. Packed values are global THD padded-storage K/V coordinates.

**attn_kwargs

AnyDefaults to {}

Attention layout metadata. Packed THD uses cu_seqlens and optional cu_seqlens_padded of shape [sequences + 1], plus optional glm_dsa_cp_query_indices of shape [tokens] containing global padded-storage query coordinates.

Returns: torch.Tensor

Hidden states with the input token axes plus [hidden], and the latest

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

nemo_automodel.components.models.glm_moe_dsa.model.GlmMoeDsaModel.update_moe_gate_bias() -> None

Update the noaux router correction bias of each local MoE layer; dense layers and disabled gates are skipped.

nemo_automodel.components.models.glm_moe_dsa.model._uses_indexshare(
    config: transformers.models.glm_moe_dsa.configuration_glm_moe_dsa.GlmMoeDsaConfig
) -> bool

Return whether the model has layers that reuse another layer’s DSA indices.

nemo_automodel.components.models.glm_moe_dsa.model.ModelClass = GlmMoeDsaForCausalLM