modeling_neuron_lm.py

from __future__ import annotations

import math
from copy import copy
from typing import Any, cast

import torch
from torch import Tensor, nn
from torch.utils.checkpoint import checkpoint
from transformers import Cache, DynamicCache, PreTrainedModel
from transformers.generation.utils import GenerationMixin
from transformers.masking_utils import create_causal_mask
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from .attention import NeuronLMAttention
from .configuration_neuron_lm import NeuronLMConfig
from .layers import RMSNorm, SwiGLU
from .rotary import RotaryEmbedding

__all__ = [
    "NeuronLMDecoderLayer",
    "NeuronLMPreTrainedModel",
    "NeuronLMModel",
    "NeuronLMForCausalLM",
]

# Marks the linear projection that writes a residual branch back into the
# residual stream. ``_init_weights`` scales these down by
# ``1 / sqrt(2 * num_hidden_layers)`` so residual-stream variance stays
# roughly constant with depth at initialization (GPT-2 / OLMo convention).
# Set through ``setattr`` because ``nn.Module.__setattr__`` is typed for
# parameters, buffers, and submodules only.
RESIDUAL_PROJECTION_FLAG = "_neuron_lm_residual_projection"


def _cache_seq_length(cache: Cache | None, layer_idx: int = 0) -> int:
    if cache is None:
        return 0
    length = cache.get_seq_length(layer_idx)
    if isinstance(length, Tensor):
        # ``.item()`` is a graph break under torch.compile. Callers only reach
        # this path when they did not supply position_ids, and generation
        # always supplies them.
        length = length.item()
    return int(length)


class NeuronLMDecoderLayer(nn.Module):
    def __init__(
        self,
        config: NeuronLMConfig,
        layer_idx: int,
    ) -> None:
        super().__init__()

        if type(layer_idx) is not int or layer_idx < 0:
            raise ValueError(
                f"layer_idx must be a non-negative integer, got {layer_idx!r}"
            )

        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx

        self.input_layernorm = RMSNorm(
            hidden_size=config.hidden_size,
            eps=config.rms_norm_eps,
        )

        self.self_attn = NeuronLMAttention(
            config,
            layer_idx=layer_idx,
        )

        self.post_attention_layernorm = RMSNorm(
            hidden_size=config.hidden_size,
            eps=config.rms_norm_eps,
        )

        self.mlp = SwiGLU(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
        )

        self.residual_dropout = nn.Dropout(
            p=config.residual_dropout,
        )

        # Both branches of this layer write through these two projections.
        setattr(self.self_attn.out_proj, RESIDUAL_PROJECTION_FLAG, True)
        setattr(self.mlp.down_proj, RESIDUAL_PROJECTION_FLAG, True)

    def forward(
        self,
        hidden_states: Tensor,
        position_embeddings: tuple[Tensor, Tensor],
        attention_mask: Tensor | None = None,
        past_key_values: Cache | None = None,
        output_attentions: bool = False,
    ) -> Tensor | tuple[Tensor, Tensor | None]:
        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)
        attention_outputs = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            output_attentions=output_attentions,
        )
        if output_attentions:
            hidden_states, attention_weights = attention_outputs
        else:
            hidden_states = attention_outputs
            attention_weights = None
        hidden_states = residual + self.residual_dropout(hidden_states)

        residual = hidden_states

        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + self.residual_dropout(hidden_states)

        if output_attentions:
            return hidden_states, attention_weights
        return hidden_states


class NeuronLMPreTrainedModel(PreTrainedModel):
    config_class = NeuronLMConfig
    base_model_prefix = "model"

    supports_gradient_checkpointing = True

    _no_split_modules = ["NeuronLMDecoderLayer"]

    # Backends verified against the SDPA reference in tests/test_attention.py.
    # `_supports_flash_attn` stays unset: flash-attn is not installed here, so
    # the claim cannot be tested, and SDPA already dispatches flash kernels on
    # recent hardware. FlexAttention is what intra-document masking compiles
    # its BlockMask through.
    _supports_sdpa = True
    _supports_flex_attn = True

    # The forward is free of data-dependent control flow, so Transformers may
    # use its compiled generation path. Regression coverage:
    # tests/test_modeling.py::test_forward_compiles_as_a_full_graph.
    _can_compile_fullgraph = True

    # _tp_plan is intentionally unset. The fused qkv_proj packs three blocks
    # whose sizes follow the GQA head counts (num_attention_heads,
    # num_key_value_heads, num_key_value_heads), while Transformers'
    # "packed_colwise" style assumes two equally sized blocks -- it would cut
    # through the K block. Supporting tensor parallelism here needs both a
    # custom sharding style and a _project_qkv that splits on per-rank head
    # counts.
    #
    # That work is not on the critical path: FSDP2 (configs/accelerate/
    # fsdp2.yaml) shards an 8B AdamW run to roughly 30 GiB per GPU across
    # four devices, so memory is not the binding constraint at the sizes this
    # model targets. Revisit if serving latency or a much larger model makes
    # tensor parallelism necessary; TrainingArguments.parallelism_config is
    # the entry point.

    def residual_initializer_std(self) -> float:
        """Initialization std for projections feeding the residual stream.

        Scaling by ``1 / sqrt(2 * num_hidden_layers)`` keeps the variance of
        the residual stream from growing with depth. There are two residual
        branches per decoder layer, hence the factor of two.
        """

        depth_scale = math.sqrt(2.0 * self.config.num_hidden_layers)
        return self.config.initializer_range / depth_scale

    def _init_weights(self, module: nn.Module) -> None:

        if isinstance(module, nn.Linear):
            if getattr(module, RESIDUAL_PROJECTION_FLAG, False):
                std = self.residual_initializer_std()
            else:
                std = self.config.initializer_range

            module.weight.data.normal_(
                mean=0.0,
                std=std,
            )

            if module.bias is not None:
                module.bias.data.zero_()

        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(
                mean=0.0,
                std=self.config.initializer_range,
            )

            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

        elif isinstance(module, nn.RMSNorm):
            if module.elementwise_affine:
                module.weight.data.fill_(1.0)

        elif isinstance(module, RotaryEmbedding):
            module.reset_parameters()


class NeuronLMModel(NeuronLMPreTrainedModel):
    def __init__(self, config: NeuronLMConfig) -> None:
        super().__init__(config)

        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(
            num_embeddings=config.vocab_size,
            embedding_dim=config.hidden_size,
            padding_idx=config.pad_token_id,
        )

        self.layers = nn.ModuleList(
            [
                NeuronLMDecoderLayer(
                    config=config,
                    layer_idx=layer_idx,
                )
                for layer_idx in range(config.num_hidden_layers)
            ]
        )

        self.norm = RMSNorm(
            hidden_size=config.hidden_size,
            eps=config.rms_norm_eps,
        )

        # RoPE frequencies are computed once per model forward and shared by
        # all decoder layers.
        self.rotary_emb = RotaryEmbedding(
            head_dim=config.head_dim,
            base=config.rope_theta,
        )

        # PreTrainedModel.gradient_checkpointing_enable() updates this flag
        # and assigns self._gradient_checkpointing_func.
        self.gradient_checkpointing = False

        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed_tokens

    def set_input_embeddings(
        self,
        value: nn.Embedding,
    ) -> None:
        self.embed_tokens = value

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> BaseModelOutputWithPast | tuple[Tensor, ...]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.return_dict
        )
        use_cache = (
            use_cache
            if use_cache is not None
            else (self.config.use_cache or past_key_values is not None)
        )

        if kwargs:
            unsupported = ", ".join(sorted(kwargs))
            raise TypeError(f"Unsupported model forward arguments: {unsupported}")

        if past_key_values is not None and not isinstance(
            past_key_values,
            Cache,
        ):
            raise TypeError(
                "past_key_values must be a Hugging Face Cache instance; "
                "legacy tuple caches are not supported"
            )

        # Cache mutation is incompatible with recomputation during backward.
        # This mirrors the behavior of the current Transformers decoder
        # layers while keeping the public forward API convenient.
        if self.gradient_checkpointing and self.training:
            use_cache = False
            past_key_values = None

        if use_cache:
            if past_key_values is None:
                past_key_values = DynamicCache(config=self.config)
        elif past_key_values is not None:
            raise ValueError("past_key_values can only be used when use_cache=True")

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Specify exactly one of input_ids or inputs_embeds")

        if input_ids is not None:
            if input_ids.ndim != 2:
                raise ValueError(
                    "input_ids must have shape "
                    "(batch_size, sequence_length), "
                    f"got shape={tuple(input_ids.shape)}"
                )

            inputs_embeds = self.embed_tokens(input_ids)

        assert inputs_embeds is not None

        if inputs_embeds.ndim != 3:
            raise ValueError(
                "inputs_embeds must have shape "
                "(batch_size, sequence_length, hidden_size), "
                f"got shape={tuple(inputs_embeds.shape)}"
            )

        batch_size, sequence_length, hidden_size = inputs_embeds.shape

        if hidden_size != self.config.hidden_size:
            raise ValueError(
                f"Expected hidden_size={self.config.hidden_size}, "
                f"got hidden_size={hidden_size}"
            )

        if sequence_length == 0:
            raise ValueError("sequence_length must be greater than zero")

        past_key_length = (
            _cache_seq_length(past_key_values) if past_key_values is not None else 0
        )
        hidden_states = inputs_embeds

        if position_ids is None:
            position_ids = torch.arange(
                past_key_length,
                past_key_length + sequence_length,
                dtype=torch.long,
                device=hidden_states.device,
            ).unsqueeze(0)
        else:
            if position_ids.ndim not in {1, 2}:
                raise ValueError(
                    "position_ids must have shape (sequence_length,) or "
                    "(batch_size, sequence_length), "
                    f"got shape={tuple(position_ids.shape)}"
                )
            if position_ids.shape[-1] != sequence_length:
                raise ValueError(
                    "The final position_ids dimension must equal the "
                    f"sequence length {sequence_length}, got "
                    f"{position_ids.shape[-1]}"
                )
            position_ids = position_ids.to(
                device=hidden_states.device,
                dtype=torch.long,
            )
            if position_ids.ndim == 1:
                position_ids = position_ids.unsqueeze(0)

        # Transformers derives packed-document boundaries from gaps in
        # position_ids, and that detection requires a 2D tensor.
        # See create_causal_mask / find_packed_sequence_indices.

        # Reading position_ids.max() is data-dependent control flow, which
        # torch.compile cannot trace in a full graph. The bound is a static
        # property of the config, so the eager check is sufficient: any shape
        # that would trip it also trips it before compilation warms up.
        if (
            not torch.compiler.is_compiling()
            and position_ids.numel() > 0
            and position_ids.max() >= self.config.max_position_embeddings
        ):
            raise ValueError(
                "position_ids contain a position at or beyond "
                f"max_position_embeddings={self.config.max_position_embeddings}"
            )

        position_embeddings = self.rotary_emb(
            hidden_states,
            position_ids=position_ids,
        )

        mask_config = self.config
        if output_attentions and self.config._attn_implementation != "eager":
            mask_config = copy(self.config)
            mask_config._attn_implementation = "eager"

        causal_attention_mask = create_causal_mask(
            config=mask_config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )

        all_hidden_states: tuple[Tensor, ...] | None = (
            () if output_hidden_states else None
        )
        all_self_attentions: tuple[Tensor, ...] | None = (
            () if output_attentions else None
        )

        for decoder_layer in self.layers:
            decoder_layer = cast(NeuronLMDecoderLayer, decoder_layer)
            if all_hidden_states is not None:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:

                def custom_forward(
                    states: Tensor,
                    layer: NeuronLMDecoderLayer = decoder_layer,
                ) -> Tensor | tuple[Tensor, Tensor | None]:
                    return layer(
                        hidden_states=states,
                        position_embeddings=position_embeddings,
                        attention_mask=causal_attention_mask,
                        past_key_values=None,
                        output_attentions=output_attentions,
                    )

                checkpointing_function = getattr(
                    self,
                    "_gradient_checkpointing_func",
                    None,
                )

                if checkpointing_function is None:
                    layer_outputs = checkpoint(
                        custom_forward,
                        hidden_states,
                        use_reentrant=False,
                    )
                else:
                    layer_outputs = checkpointing_function(
                        custom_forward,
                        hidden_states,
                    )
            else:
                layer_outputs = decoder_layer(
                    hidden_states=hidden_states,
                    position_embeddings=position_embeddings,
                    attention_mask=causal_attention_mask,
                    past_key_values=past_key_values,
                    output_attentions=output_attentions,
                )

            if output_attentions:
                hidden_states, attention_weights = layer_outputs
                assert all_self_attentions is not None
                assert attention_weights is not None
                all_self_attentions += (attention_weights,)
            else:
                hidden_states = layer_outputs

        hidden_states = self.norm(hidden_states)

        if all_hidden_states is not None:
            all_hidden_states += (hidden_states,)

        if not return_dict:
            outputs: tuple[Any, ...] = (hidden_states,)

            if use_cache:
                outputs += (past_key_values,)

            if output_hidden_states:
                outputs += (all_hidden_states,)

            if output_attentions:
                outputs += (all_self_attentions,)

            return outputs

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=cast(Any, all_hidden_states),
            attentions=cast(Any, all_self_attentions),
        )


class NeuronLMForCausalLM(
    NeuronLMPreTrainedModel,
    GenerationMixin,
):
    """NeuronLM decoder with a causal language-modeling head."""

    _tied_weights_keys = {
        "lm_head.weight": "model.embed_tokens.weight",
    }

    def __init__(self, config: NeuronLMConfig) -> None:
        super().__init__(config)

        self.model = NeuronLMModel(config)
        self.vocab_size = config.vocab_size

        self.lm_head = nn.Linear(
            in_features=config.hidden_size,
            out_features=config.vocab_size,
            bias=False,
        )

        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

    def set_input_embeddings(
        self,
        value: nn.Embedding,
    ) -> None:
        self.model.embed_tokens = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(
        self,
        value: nn.Linear,
    ) -> None:
        self.lm_head = value

    def get_decoder(self) -> NeuronLMModel:
        return self.model

    def set_decoder(
        self,
        decoder: NeuronLMModel,
    ) -> None:
        self.model = decoder

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        labels: Tensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        num_items_in_batch: Tensor | int | None = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple[Tensor, ...]:
        return_dict = (
            return_dict if return_dict is not None else self.config.return_dict
        )

        model_outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )

        if return_dict:
            hidden_states = model_outputs.last_hidden_state
        else:
            hidden_states = model_outputs[0]

        logits = self.lm_head(hidden_states)

        loss: Tensor | None = None

        if labels is not None:
            if labels.ndim != 2:
                raise ValueError(
                    "labels must have shape "
                    "(batch_size, sequence_length), "
                    f"got shape={tuple(labels.shape)}"
                )

            expected_shape = hidden_states.shape[:2]

            if tuple(labels.shape) != tuple(expected_shape):
                raise ValueError(
                    f"labels must have shape {tuple(expected_shape)}, "
                    f"got {tuple(labels.shape)}"
                )

            if labels.shape[1] < 2:
                raise ValueError(
                    "At least two sequence positions are required "
                    "to compute causal language-modeling loss"
                )

            labels = labels.to(device=logits.device)

            loss = self.loss_function(
                logits=logits,
                labels=labels,
                vocab_size=self.config.vocab_size,
                num_items_in_batch=num_items_in_batch,
            )

        if not return_dict:
            output = (logits,) + model_outputs[1:]

            if loss is not None:
                return (loss,) + output

            return output

        return CausalLMOutputWithPast(
            loss=cast(Any, loss),
            logits=logits,
            past_key_values=model_outputs.past_key_values,
            hidden_states=model_outputs.hidden_states,
            attentions=model_outputs.attentions,
        )