nemo_automodel.components.models.inkling.layers

Expert-parallel Mixture-of-Experts layer for the Inkling model.

Only the MoE feed-forward is reimplemented here; every other Inkling submodule (attention, short convolutions, relative-position logits, norms, vision/audio towers) is reused verbatim from HuggingFace transformers. The routing math and shared-expert formulation mirror transformers.models.inkling.modeling_inkling exactly, while the routed experts run through NeMo AutoModel’s :class:~nemo_automodel.components.moe.experts.GroupedExperts, which provides grouped-GEMM compute and expert-parallel (EP) sharding.

Module Contents

Classes

Name	Description
`InklingCorrectionBias`	Own and apply Inkling’s trained fp32 router correction bias.
`InklingDenseMLP`	Dense Inkling MLP retaining the checkpoint’s fused interleaved layout.
`InklingGate`	Top-k router matching `transformers` `InklingTopkRouter`.
`InklingMoE`	Inkling MoE feed-forward: custom router + shared experts + grouped routed experts.
`InklingSharedExperts`	Always-on shared experts, matching `transformers` `InklingSharedExperts`.
`InklingShortConvolution`	Keep short-convolution weights in an existing model-owned fp32 holder.
`_InklingShortConvolutionFP32`	Run one short convolution inside a dtype-uniform fp32 FSDP unit.

Functions

Name	Description
`build_inkling_moe_config`	Build the routed-expert :class:`MoEConfig` from an `InklingTextConfig`.
`inkling_swiglu`	Apply SwiGLU to Inkling’s interleaved `[gate, up]` projection layout.

API

class nemo_automodel.components.models.inkling.layers.InklingCorrectionBias(
    num_experts: int
)

Bases: Module

Own and apply Inkling’s trained fp32 router correction bias.

e_score_correction_bias

nemo_automodel.components.models.inkling.layers.InklingCorrectionBias.forward(
    routed_scores: torch.Tensor
) -> torch.Tensor

Add the correction bias in the routing score dtype.

class nemo_automodel.components.models.inkling.layers.InklingDenseMLP(
    config
)

Bases: Module

Dense Inkling MLP retaining the checkpoint’s fused interleaved layout.

down_proj

gate_up_proj

global_scale

= nn.Parameter(torch.ones(1, dtype=model_dtype))

nemo_automodel.components.models.inkling.layers.InklingDenseMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Run the dense MLP over the fused interleaved gate/up projection.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [..., hidden] with arbitrary leading dimensions.

Returns: torch.Tensor

torch.Tensor: Tensor of shape [..., hidden]. gate_up_proj maps

nemo_automodel.components.models.inkling.layers.InklingDenseMLP.init_weights(
    init_std: float = 0.02
) -> None

class nemo_automodel.components.models.inkling.layers.InklingGate(
    config,
    gate_precision: torch.dtype | None = torch.float32
)

Bases: Module

Top-k router matching transformers InklingTopkRouter.

The gate weight bank additionally holds n_shared_experts rows whose logits participate in the softmax normalization; the resulting shared weights (gammas) scale the always-on shared experts. Selection uses sigmoid scores plus a per-expert correction bias, while the returned weights come from a log-sigmoid softmax over the selected routed logits concatenated with the shared logits, scaled by route_scale * global_scale.

_fp32_params

= InklingCorrectionBias(self.num_experts)

e_score_correction_bias

Parameter

Expose the correction bias under the Transformers-compatible name.

global_scale

= nn.Parameter(torch.ones(1, dtype=model_dtype))

hidden_dim

= config.hidden_size

n_shared_experts

= config.n_shared_experts

n_total_experts

= self.num_experts + self.n_shared_experts

num_experts

= config.n_routed_experts

route_scale

= config.route_scale

top_k

= config.num_experts_per_tok

weight

nemo_automodel.components.models.inkling.layers.InklingGate.forward(
    hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Route tokens to experts.

Parameters:

hidden_states

torch.Tensor

Input tokens of shape [num_tokens, hidden_dim].

Returns: torch.Tensor

(topk_weights, topk_indices, shared_gammas) where

nemo_automodel.components.models.inkling.layers.InklingGate.update_bias() -> None

No-op: Inkling’s correction bias is a trained parameter, not an aux-loss update.

class nemo_automodel.components.models.inkling.layers.InklingMoE(
    config,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None
)

Bases: MoE

Inkling MoE feed-forward: custom router + shared experts + grouped routed experts.

Subclasses :class:~nemo_automodel.components.moe.layers.MoE so the distributed parallelizer (which discovers MoE blocks via isinstance and shards self.experts across the expert-parallel mesh) treats it uniformly, but the generic __init__ / forward are replaced to match Inkling’s routing.

gate

= InklingGate(config, gate_precision=gate_precision)

moe_config

shared_experts

= InklingSharedExperts(config)

nemo_automodel.components.models.inkling.layers.InklingMoE.forward(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None
) -> torch.Tensor

Route tokens through the shared and routed experts and recombine.

Parameters:

hidden_states

torch.Tensor

Input of shape [..., hidden_dim] (the trailing dimension is the model dim; leading dims are flattened to tokens).

padding_mask

torch.Tensor | NoneDefaults to None

Optional boolean mask of padding positions; padded tokens are excluded from routed-expert dispatch.

Returns: torch.Tensor

torch.Tensor: Output with the same shape as hidden_states.

nemo_automodel.components.models.inkling.layers.InklingMoE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Initialize MoE parameters in-place (used for from-scratch training).

class nemo_automodel.components.models.inkling.layers.InklingSharedExperts(
    config
)

Bases: Module

Always-on shared experts, matching transformers InklingSharedExperts.

The fused gate/up projection stays in the checkpoint’s interleaved layout so distributed loading only needs transpose views. The per-token gammas (from :class:InklingGate) weight each expert’s output before summation.

down_proj

gate_up_proj

n_shared_experts

= config.n_shared_experts

nemo_automodel.components.models.inkling.layers.InklingSharedExperts.forward(
    hidden_states: torch.Tensor,
    gammas: torch.Tensor
) -> torch.Tensor

Apply the shared experts.

Parameters:

hidden_states

torch.Tensor

Input tokens of shape [num_tokens, hidden_dim].

gammas

torch.Tensor

Per-token shared weights of shape [num_tokens, n_shared_experts].

Returns: torch.Tensor

torch.Tensor: Output of shape [num_tokens, hidden_dim].

class nemo_automodel.components.models.inkling.layers.InklingShortConvolution(
    source: torch.nn.Module
)

Bases: Module

Keep short-convolution weights in an existing model-owned fp32 holder.

_fp32_params

= _InklingShortConvolutionFP32(source)

nemo_automodel.components.models.inkling.layers.InklingShortConvolution.forward(
    hidden_states: torch.Tensor,
    past_key_values: typing.Any | None = None,
    conv_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Run the short convolution in fp32 and cast the result back.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

past_key_values

Any | NoneDefaults to None

Optional cache holding the per-layer conv state.

conv_mask

torch.Tensor | NoneDefaults to None

Optional boolean padding mask of shape [batch, sequence].

Returns: torch.Tensor

torch.Tensor: Tensor of shape [batch, sequence, hidden] in the input

class nemo_automodel.components.models.inkling.layers._InklingShortConvolutionFP32(
    source: torch.nn.Module
)

Bases: Module

Run one short convolution inside a dtype-uniform fp32 FSDP unit.

conv_idx

= source.conv_idx

conv_kernel_size

= source.conv_kernel_size

layer_idx

= source.layer_idx

weight

nemo_automodel.components.models.inkling.layers._InklingShortConvolutionFP32.forward(
    hidden_states: torch.Tensor,
    past_key_values: typing.Any | None = None,
    conv_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Apply the causal short convolution with a residual connection.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden]. Internally transposed to [batch, hidden, sequence] for the depthwise conv1d and transposed back before the residual add.

past_key_values

Any | NoneDefaults to None

Optional cache holding the per-layer conv state; updated in place during incremental decoding.

conv_mask

torch.Tensor | NoneDefaults to None

Optional boolean padding mask of shape [batch, sequence] that zeroes padded positions before the convolution.

Returns: torch.Tensor

torch.Tensor: Tensor of shape [batch, sequence, hidden] (input plus

nemo_automodel.components.models.inkling.layers.build_inkling_moe_config(
    text_config,
    backend: nemo_automodel.components.models.common.utils.BackendConfig
) -> nemo_automodel.components.moe.config.MoEConfig

Build the routed-expert :class:MoEConfig from an InklingTextConfig.

Only the fields consumed by :class:GroupedExperts matter here (routing and shared experts are handled by :class:InklingGate / :class:InklingSharedExperts), so n_shared_experts is set to 0 and the gate-specific fields are left at neutral defaults.

Parameters:

text_config

HuggingFace InklingTextConfig.

backend

BackendConfig

Backend configuration selecting the expert compute/dispatch kernels.

Returns: MoEConfig

Configuration for the routed grouped experts.

nemo_automodel.components.models.inkling.layers.inkling_swiglu(
    hidden_states: torch.Tensor,
    routing_weights: torch.Tensor
) -> torch.Tensor

Apply SwiGLU to Inkling’s interleaved [gate, up] projection layout.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [..., 2 * intermediate] with arbitrary leading dimensions, whose last axis interleaves gate and up channels as [g0, u0, g1, u1, ...] (::2 selects gate, 1::2 selects up).

routing_weights

torch.Tensor

Tensor broadcastable to [..., intermediate] that scales the activated output.

Returns: torch.Tensor

torch.Tensor: Activated tensor of shape [..., intermediate] in