nemo_automodel.components.models.glm_moe_dsa.model

Module Contents

Classes

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

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

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: typing.Optional[bool] = 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 and *carry holds the previous stage’s running top-k selection. Non-last stages return (hidden_states, topk_indices) and the last stage returns the logits tensor.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Token IDs (BSHD [B, S] / THD [1, T]) on the first stage, or the upstream hidden state on later pipeline stages.

carry

torch.TensorDefaults to ()

Optional (topk_indices,) carried from the previous pipeline stage.

position_ids / attention_mask / padding_mask

Optional masks / positions.

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

If 0, project all positions; else only the last logits_to_keep.

output_hidden_states

Optional[bool]Defaults to None

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

**attn_kwargs

AnyDefaults to {}

Additional arguments forwarded to the base model.

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, threading the IndexShare top-k as a carry tensor.

Non-first stages additionally receive the previous “full” layer’s top-k selection, and non-last stages emit the running selection, so a stage that begins with a “shared” layer has the top-k it needs (correct at any sequence length).

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

GLM DSA TileLang kernels require validation to use the THD packed 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).

prev_topk_indices seeds the IndexShare running selection (used under pipeline parallelism, where an earlier “full” layer lives on the previous stage); it is None in the single-process path. The returned topk_indices is the running selection at the end of this stage’s layers, so it can be carried to the next pipeline stage.

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.ModelClass = GlmMoeDsaForCausalLM