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