nemo_automodel.components.models.inkling2.layers

Differentiable Inkling2 primitives and the common MoE integration.

Module Contents

Classes

Name	Description
`Inkling2LMHead`	Bias-free output projection with FP32 arithmetic and ordinary FSDP hooks.
`Inkling2MoE`	Common Gate/GroupedExperts with Inkling2’s activation callable.
`L2Clip`	Per-head L2 clipping, optionally preceded by learned per-channel scales.
`ResidualQuery`	Callable FP32 owner for one set of depth-mixture queries.
`ResidualState`	Immutable depth state; every tensor has shape [batch, sequence, hidden].
`ShortConvolution`	Causal width-W convolution with static or token-generated per-head taps.

Functions

Name	Description
`capped_swiglu`	Apply the interleaved capped activation, with routing before the down GEMM.
`residual_mix`	Mix depth sources independently at each token in FP32.

API

class nemo_automodel.components.models.inkling2.layers.Inkling2LMHead()

Bases: Linear

Bias-free output projection with FP32 arithmetic and ordinary FSDP hooks.

nemo_automodel.components.models.inkling2.layers.Inkling2LMHead.forward(
    input: torch.Tensor
) -> torch.Tensor

Project through the module boundary so FSDP gathers its weight.

Parameters:

input

torch.Tensor

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

Returns: torch.Tensor

FP32 tensor of shape […, vocab]. Weight storage may be FP32 or BF16.

class nemo_automodel.components.models.inkling2.layers.Inkling2MoE(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)

Bases: MoE

Common Gate/GroupedExperts with Inkling2’s activation callable.

class nemo_automodel.components.models.inkling2.layers.L2Clip(
    heads: int,
    width: int,
    threshold: float,
    learned_scale: bool
)

Bases: Module

Per-head L2 clipping, optionally preceded by learned per-channel scales.

scale

nemo_automodel.components.models.inkling2.layers.L2Clip.forward(
    x: torch.Tensor,
    tau: torch.Tensor | None = None
) -> torch.Tensor

Clip head vectors, then apply optional length scaling.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim].

tau

torch.Tensor | NoneDefaults to None

Optional FP32 tensor of shape [batch, sequence], applied after clipping.

Returns: torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim], in x’s dtype.

class nemo_automodel.components.models.inkling2.layers.ResidualQuery(
    hidden_size: int,
    count: int,
    normalize_values: bool
)

Bases: Module

Callable FP32 owner for one set of depth-mixture queries.

weight

nemo_automodel.components.models.inkling2.layers.ResidualQuery.forward(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState
) -> torch.Tensor

Mix immutable depth state.

Parameters:

state

ResidualState

Completed and partial blocks, each [batch, sequence, hidden].

Returns: torch.Tensor

FP32 tensor of shape [queries, batch, sequence, hidden].

class nemo_automodel.components.models.inkling2.layers.ResidualState(
    blocks: tuple[torch.Tensor, ...],
    partial: torch.Tensor
)

Dataclass

Immutable depth state; every tensor has shape [batch, sequence, hidden].

blocks

tuple[Tensor, ...]

Completed depth blocks, including the initial embedding block.

partial

Tensor

Current incomplete block. Neither field is detached or mutated.

class nemo_automodel.components.models.inkling2.layers.ShortConvolution(
    channels: int,
    width: int,
    dtype: torch.dtype,
    dynamic_input_size: int = 0,
    heads: int = 1
)

Bases: Module

Causal width-W convolution with static or token-generated per-head taps.

cp_mesh

DeviceMesh | None = None

static_w

weight_proj

nemo_automodel.components.models.inkling2.layers.ShortConvolution.forward(
    x: torch.Tensor,
    valid: torch.Tensor,
    dynamic_input: torch.Tensor | None = None,
    positions: torch.Tensor | None = None
) -> torch.Tensor

Apply convolution with exchanged CP history, zero document history and masked padding.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, channels].

valid

torch.Tensor

Boolean tensor of shape [batch, sequence], true for real tokens.

dynamic_input

torch.Tensor | NoneDefaults to None

Tensor of shape [batch, sequence, hidden] for dynamic taps.

positions

torch.Tensor | NoneDefaults to None

Logical document positions [batch, sequence], including CP offsets. Taps cannot read before the current document start.

Returns: torch.Tensor

Tensor of shape [batch, sequence, channels], in x’s dtype; padding is zero.

nemo_automodel.components.models.inkling2.layers.ShortConvolution.reset_parameters(
    init_std: float
) -> None

Initialize every static and dynamic coefficient without checkpoint data.

nemo_automodel.components.models.inkling2.layers.capped_swiglu(
    gate_up: torch.Tensor,
    weights: torch.Tensor,
    gate_softcap: float,
    up_softcap: float
) -> torch.Tensor

Apply the interleaved capped activation, with routing before the down GEMM.

Parameters:

gate_up

torch.Tensor

Tensor of shape [tokens, 2 * intermediate], interleaved gate/up.

weights

torch.Tensor

Tensor of shape [tokens, 1], containing selected expert weights.

gate_softcap

float

Positive gate tanh cap; sigmoid still receives the raw gate.

up_softcap

float

Positive up tanh cap.

Returns: torch.Tensor

Tensor of shape [tokens, intermediate], in gate_up’s dtype.

nemo_automodel.components.models.inkling2.layers.residual_mix(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState,
    queries: torch.Tensor,
    normalize_values: bool
) -> torch.Tensor

Mix depth sources independently at each token in FP32.

Parameters:

state

ResidualState

Completed blocks and partial, each [batch, sequence, hidden].

queries

torch.Tensor

Tensor of shape [queries, hidden], with one or three pseudo-queries.

normalize_values

bool

Normalize values as well as scoring keys.

Returns: torch.Tensor

Tensor of shape [queries, batch, sequence, hidden], without mutating state.