> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.glm_moe_dsa.model

## Module Contents

### Classes

| Name                                                                                               | Description |
| -------------------------------------------------------------------------------------------------- | ----------- |
| [`Block`](#nemo_automodel-components-models-glm_moe_dsa-model-Block)                               | -           |
| [`GlmMoeDsaForCausalLM`](#nemo_automodel-components-models-glm_moe_dsa-model-GlmMoeDsaForCausalLM) | -           |
| [`GlmMoeDsaModel`](#nemo_automodel-components-models-glm_moe_dsa-model-GlmMoeDsaModel)             | -           |

### Functions

| Name                                                                                       | Description                                                                 |
| ------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------- |
| [`_uses_indexshare`](#nemo_automodel-components-models-glm_moe_dsa-model-_uses_indexshare) | Return whether the model has layers that reuse another layer's DSA indices. |

### Data

[`ModelClass`](#nemo_automodel-components-models-glm_moe_dsa-model-ModelClass)

### API

```python
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`**

---

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

```python
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).

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

```python
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](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-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`

---

```python
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).

```python
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 | None` — default: 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.Tensor` — default: ()

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 | None` — default: None

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

---

**`attention_mask`** `torch.Tensor | None` — default: None

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

---

**`padding_mask`** `torch.Tensor | None` — default: None

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

---

**`logits_to_keep`** `Union[int, torch.Tensor]` — default: 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 | None` — default: None

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

---

**`**attn_kwargs`** `Any` — default: \{}

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

```python
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

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

classmethod

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

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

```python
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

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

```python
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`** `int` — default: 1

Number of chunks for load-balanced CP sharding.

---

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

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

```python
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.

```python
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.

```python
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`

---

```python
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 | None` — default: None

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

---

**`attention_mask`** `torch.Tensor | None` — default: None

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

---

**`padding_mask`** `torch.Tensor | None` — default: None

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

---

**`prev_topk_indices`** `torch.Tensor | None` — default: 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`** `Any` — default: \{}

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

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

```python
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.

```python
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.

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