nemo_automodel.components.models.ling_v2

Submodules

Package Contents

Classes

Name	Description
`BailingMoeV2Config`	Configuration class for the BailingMoeV2 model (Ling 2.0).
`BailingMoeV2ForCausalLM`	Causal-LM head wrapping `BailingMoeV2Model`.

API

class nemo_automodel.components.models.ling_v2.config.BailingMoeV2Config(
    vocab_size: int = 157184,
    hidden_size: int = 2048,
    intermediate_size: int = 5120,
    num_hidden_layers: int = 20,
    num_attention_heads: int = 16,
    num_key_value_heads: int = 4,
    hidden_act: str = 'silu',
    use_qkv_bias: bool = False,
    use_bias: bool = False,
    rms_norm_eps: float = 1e-06,
    tie_word_embeddings: bool = False,
    embedding_dropout: float = 0.0,
    attention_dropout: float = 0.0,
    output_dropout: float = 0.0,
    initializer_range: float = 0.02,
    max_position_embeddings: int = 32768,
    rope_theta: float = 600000.0,
    use_cache: bool = True,
    max_window_layers: int = 20,
    rope_scaling: dict | None = None,
    pad_token_id: int = 156892,
    eos_token_id: int = 156892,
    num_experts: int = 256,
    num_shared_experts: int = 1,
    num_experts_per_tok: int = 8,
    n_group: int = 8,
    topk_group: int = 4,
    moe_intermediate_size: int = 512,
    first_k_dense_replace: int = 1,
    head_dim: int = 128,
    output_router_logits: bool = False,
    use_qk_norm: bool = True,
    partial_rotary_factor: float = 1.0,
    num_nextn_predict_layers: int = 0,
    mtp_loss_scaling_factor: float = 0,
    moe_router_enable_expert_bias: bool = True,
    routed_scaling_factor: float = 1.0,
    norm_topk_prob: bool = True,
    score_function: str = 'sigmoid',
    rotary_dim: int | None = None,
    kwargs = {}
)

Bases: PretrainedConfig

Configuration class for the BailingMoeV2 model (Ling 2.0).

The defaults reflect the Ling-mini-2.0 (16B-A1.4B) variant. Larger variants (Ling-flash-2.0 100B-A6B and Ling-1T 1T-A50B) override sizing knobs but share the same architecture: GQA attention with per-head QK-RMSNorm, partial RoPE, sigmoid-routed grouped MoE with shared experts, and first_k_dense_replace dense MLP layers at the start.

head_dim

= head_dim or hidden_size // num_attention_heads

keys_to_ignore_at_inference

= ['past_key_values']

model_type

= 'bailing_moe'

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', 'expert_bias']

_pp_keep_self_forward

bool = True

lm_head

model

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

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: bool | None = 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

bool | NoneDefaults 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:

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