DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsLing V2

nemo_automodel.components.models.ling_v2.model

View as Markdown

BailingMoeV2 model (Ling 2.0 family).

Architecture summary (from the public inclusionAI/Ling-{mini,flash,1T}-2.0 checkpoints):

  • GQA attention with per-head QK-RMSNorm and partial RoPE (rotates the first head_dim * partial_rotary_factor channels only).
  • first_k_dense_replace dense MLP layers at the start of the stack; the remaining layers are sigmoid-routed grouped MoE with shared experts and an aux-loss-free per-expert bias (DeepSeek-V3-style routing).
  • Single shared expert with intermediate size moe_intermediate_size.
  • MTP heads (num_nextn_predict_layers) are disabled in all published checkpoints and intentionally not modeled here.

Example (YAML):

model:
  _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
  pretrained_model_name_or_path: inclusionAI/Ling-mini-2.0

Module Contents

Classes

NameDescription
BailingMoeV2ForCausalLMCausal-LM head wrapping BailingMoeV2Model.
BailingMoeV2ModelEmbedding + decoder stack + final norm. No LM head.
BlockSingle transformer block: attention + (dense MLP or MoE) + residuals.

Data

ModelClass

API

class nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM(
    config: nemo_automodel.components.models.ling_v2.config.BailingMoeV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Causal-LM head wrapping BailingMoeV2Model.

_keep_in_fp32_modules_strict
= ['e_score_correction_bias']
_pp_keep_self_forward
bool = True
backend
= backend or BackendConfig()
lm_head
model
state_dict_adapter
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.forward(
    input_ids: 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

Forward pass returning CausalLMOutputWithPast.

Supports BSHD (input_ids shape [B, S]) and THD (squeezed to [T] when attn_kwargs["qkv_format"] == "thd") formats.

Parameters:

input_ids
torch.Tensor

Input token IDs.

position_ids
torch.Tensor | NoneDefaults to None

Optional position indices.

attention_mask
torch.Tensor | NoneDefaults to None

Optional 2D padding mask.

padding_mask
torch.Tensor | NoneDefaults to None

Optional padding mask used by the THD squeeze helper.

logits_to_keep
Union[int, torch.Tensor]Defaults to 0

If > 0, only compute logits for the last logits_to_keep positions (avoids materialising the full logit matrix during generation / fused-CE training). 0 computes all positions.

output_hidden_states
Optional[bool]Defaults to None

Whether to return the final hidden states (the input to lm_head) on the output. Required by the fused cross-entropy (cut-CE) training path.

**attn_kwargs
AnyDefaults to {}

Additional arguments forwarded to the base model.

Returns: CausalLMOutputWithPast

class:~transformers.modeling_outputs.CausalLMOutputWithPast with

nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.from_config(
    config: nemo_automodel.components.models.ling_v2.config.BailingMoeV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
classmethod
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)
classmethod
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.get_capabilities(
    config
) -> nemo_automodel._transformers.model_capabilities.ModelCapabilities
classmethod

Return parallelism capabilities for a specific Ling/Bailing-MoE config.

Three checkpoint variants share this class:

  1. inclusionAI/Ling-1T — 1T-param MoE, requires PP. Demonstrated by examples/llm_finetune/ling/ling_1t_sft.yaml (pp=4) and ling_1t_lora_pp.yaml (pp=8); both with ep_size>=8.
  2. inclusionAI/Ling-flash-2.0 — mid-size MoE, single-rank EP only. Demonstrated by ling_flash_2_0_sft.yaml / ling_flash_2_0_lora.yaml (pp=1, ep=8-32).
  3. inclusionAI/Ling-mini-2.0 — small MoE, single-rank EP only. Demonstrated by ling_mini_2_0_{hellaswag,sft,squad}.yaml (pp=1, ep=4-8).

Dispatch is on num_hidden_layers since Ling-1T (~80 layers) is well separated from Ling-flash-2.0 (~32) and Ling-mini-2.0 (~20).

nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.get_input_embeddings()
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.get_output_embeddings()
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.set_input_embeddings(
    value
)
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.set_output_embeddings(
    new_embeddings
)
nemo_automodel.components.models.ling_v2.model.BailingMoeV2ForCausalLM.update_moe_gate_bias() -> None
class nemo_automodel.components.models.ling_v2.model.BailingMoeV2Model(
    config: nemo_automodel.components.models.ling_v2.config.BailingMoeV2Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)

Bases: Module

Embedding + decoder stack + final norm. No LM head.

embed_tokens
head_dim
= config.head_dim
layers
= torch.nn.ModuleDict()
max_seq_len
= config.max_position_embeddings
moe_config
= moe_config or MoEConfig(**moe_defaults)
norm
rotary_emb
nemo_automodel.components.models.ling_v2.model.BailingMoeV2Model.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
nemo_automodel.components.models.ling_v2.model.BailingMoeV2Model.init_weights(
    buffer_device: torch.device | None = None
) -> None
nemo_automodel.components.models.ling_v2.model.BailingMoeV2Model.update_moe_gate_bias() -> None

No-op for SFT; published Ling checkpoints freeze the expert_bias buffer.

class nemo_automodel.components.models.ling_v2.model.Block(
    layer_idx: int,
    config: nemo_automodel.components.models.ling_v2.config.BailingMoeV2Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Single transformer block: attention + (dense MLP or MoE) + residuals.

input_layernorm
mlp
post_attention_layernorm
self_attn
= BailingMoeV2Attention(config, backend)
nemo_automodel.components.models.ling_v2.model.Block._mlp(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None
) -> torch.Tensor
nemo_automodel.components.models.ling_v2.model.Block.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor
nemo_automodel.components.models.ling_v2.model.Block.init_weights(
    buffer_device: torch.device
) -> None
nemo_automodel.components.models.ling_v2.model.ModelClass = BailingMoeV2ForCausalLM