""" Hugging Face-compatible Laneformer implementation. This file defines: - LaneformerPreTrainedModel - LaneformerModel (backbone) - LaneformerForCausalLM (head + loss) It mirrors the TorchTitan Laneformer structure so that state_dict keys largely match, which makes converting checkpoints trivial (often identity mapping). """ from __future__ import annotations import inspect import math from dataclasses import dataclass from enum import Enum from functools import partial from typing import Optional, Tuple, Union, Unpack import torch import torch.nn.functional as F from torch import nn from transformers.cache_utils import Cache, DynamicCache from transformers.generation import GenerationMixin from transformers.masking_utils import ( create_causal_mask, create_sliding_window_causal_mask, ) from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.utils import auto_docstring, can_return_tuple, TransformersKwargs from .configuration_laneformer import LaneformerConfig # ============================================================ # Rotary helpers (mirrors your implementation) # ============================================================ @dataclass class RoPEScalingArgs: scaling_factor: float = 1.0 low_freq_factor: float = 1.0 high_freq_factor: float = 1.0 original_max_position_embeddings: int = 4096 def _rope_scaling_from_config( config: LaneformerConfig, ) -> RoPEScalingArgs | None: if config.rope_scaling is None: return None try: args = RoPEScalingArgs(**config.rope_scaling) except TypeError as exc: raise ValueError(f"Invalid rope_scaling configuration: {exc}") from exc if args.scaling_factor <= 0: raise ValueError("rope_scaling.scaling_factor must be > 0") if args.low_freq_factor <= 0 or args.high_freq_factor <= 0: raise ValueError("rope_scaling frequency factors must be > 0") if args.high_freq_factor == args.low_freq_factor: raise ValueError( "rope_scaling.high_freq_factor and low_freq_factor must differ" ) if args.original_max_position_embeddings <= 0: raise ValueError("rope_scaling.original_max_position_embeddings must be > 0") return args def compute_rope_freqs( dim: int, theta: float = 10000.0, scaling_args: RoPEScalingArgs | None = None, device: torch.device | None = None, ) -> torch.Tensor: freqs = 1.0 / ( theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim) ) if scaling_args is not None: scaling_factor = scaling_args.scaling_factor low_freq_factor = scaling_args.low_freq_factor high_freq_factor = scaling_args.high_freq_factor original_max_position_embeddings = scaling_args.original_max_position_embeddings wavelen = 2 * math.pi / freqs high_freq_wavelen = original_max_position_embeddings / high_freq_factor low_freq_wavelen = original_max_position_embeddings / low_freq_factor freqs = torch.where(wavelen > low_freq_wavelen, freqs / scaling_factor, freqs) smooth_factor = ( original_max_position_embeddings / wavelen - low_freq_factor ) / (high_freq_factor - low_freq_factor) smoothed_freqs = ( 1 - smooth_factor ) * freqs / scaling_factor + smooth_factor * freqs is_medium_freqs = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) freqs = torch.where(is_medium_freqs, smoothed_freqs, freqs) return freqs def precompute_freqs_cis( dim: int, end: int, theta: float = 10000.0, scaling_args: RoPEScalingArgs | None = None, ) -> torch.Tensor: freqs = compute_rope_freqs(dim, theta, scaling_args) t = torch.arange(end, device=freqs.device) freqs = torch.outer(t, freqs).float() freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis def freqs_cis_for_position_ids( position_ids: torch.Tensor, dim: int, theta: float = 10000.0, scaling_args: RoPEScalingArgs | None = None, ) -> torch.Tensor: freqs = compute_rope_freqs( dim, theta, scaling_args, device=position_ids.device, ) angles = position_ids.float().unsqueeze(-1) * freqs return torch.polar(torch.ones_like(angles), angles) def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: ndim = x.ndim assert ndim > 1 seqlen = x.shape[1] freqs_cis = freqs_cis[0:seqlen] assert freqs_cis.shape == (seqlen, x.shape[-1]) shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] return freqs_cis.view(*shape) def apply_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # freqs_cis = reshape_for_broadcast(freqs_cis, xq_) xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: bs, slen, n_kv_heads, head_dim = x.shape if n_rep == 1: return x return ( torch.unsqueeze(x, dim=3) .expand(bs, slen, n_kv_heads, n_rep, head_dim) .reshape(bs, slen, n_kv_heads * n_rep, head_dim) ) class LaneLinear(nn.Module): def __init__(self, in_features: int, out_features: int, num_lanes: int): super().__init__() self.num_lanes = num_lanes self.in_features = in_features self.out_features = out_features self.weight = torch.empty(self.get_shape()) self.reset_parameters() def reset_parameters(self): raise NotImplementedError def get_shape(self) -> Tuple[int, int, int]: raise NotImplementedError def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.einsum( "bsli,loi->bslo", x, self.weight ) # [L, out_features, 1] @ [B*S, L, in_features // L, 1] -> [B*S, L, out_features, 1] class LaneColumnLinear(LaneLinear): def __init__( self, in_features: int, out_features: int, num_lanes: int, ): super().__init__( in_features=in_features, out_features=out_features, num_lanes=num_lanes ) def get_shape(self) -> Tuple[int, int, int]: return (self.num_lanes, self.out_features, self.in_features // self.num_lanes) def reset_parameters(self): w = torch.empty(self.out_features, self.in_features) nn.init.kaiming_uniform_(w, a=math.sqrt(5)) self.weight = nn.Parameter( w.reshape( self.out_features, self.num_lanes, self.in_features // self.num_lanes ) .permute(1, 0, 2) .contiguous() ) class LaneRowLinear(LaneLinear): def __init__( self, in_features: int, out_features: int, num_lanes: int, ): super().__init__( in_features=in_features, out_features=out_features, num_lanes=num_lanes ) def get_shape(self) -> Tuple[int, int, int]: return (self.num_lanes, self.out_features // self.num_lanes, self.in_features) def reset_parameters(self): w = torch.empty(self.out_features, self.in_features) nn.init.kaiming_uniform_(w, a=math.sqrt(5)) self.weight = nn.Parameter(w.reshape(self.weight.shape).contiguous()) class LaneLMHead(nn.Module): def __init__(self, in_features: int, out_features: int, num_lanes: int): super().__init__() self.num_lanes = num_lanes self.linear = nn.Linear( in_features=num_lanes * in_features, out_features=out_features, bias=False ) def forward(self, x: torch.Tensor) -> torch.Tensor: B, S, L, d = x.shape return self.linear(x.view(B, S, L * d)) class LaneRMSNorm(nn.Module): def __init__(self, dim: int, num_lanes: int, eps: float = 1e-8): super().__init__() self.rms_norm = nn.RMSNorm(dim, eps=eps, elementwise_affine=False) self.dim = dim self.eps = eps self.scale = nn.Parameter(torch.ones(num_lanes, dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.rms_norm(x) * self.scale def reset_parameters(self): nn.init.ones_(self.scale) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: Optional[float] = None, dropout: float = 0.0, **kwargs, ): num_key_value_groups = getattr(module, "num_key_value_groups", 1) key_states = repeat_kv(key.transpose(1, 2), num_key_value_groups).transpose(1, 2) value_states = repeat_kv(value.transpose(1, 2), num_key_value_groups).transpose( 1, 2 ) if scaling is None: scaling = query.shape[-1] ** -0.5 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( query.dtype ) attn_weights = nn.functional.dropout( attn_weights, p=dropout, training=module.training ) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights def _attention_interface(config: LaneformerConfig): attn_implementation = getattr(config, "_attn_implementation", "eager") if attn_implementation == "eager": return eager_attention_forward get_interface = getattr(ALL_ATTENTION_FUNCTIONS, "get_interface", None) if get_interface is not None: return get_interface(attn_implementation, eager_attention_forward) return ALL_ATTENTION_FUNCTIONS[attn_implementation] _CAUSAL_MASK_USES_INPUTS_EMBEDS = ( "inputs_embeds" in inspect.signature(create_causal_mask).parameters ) _SLIDING_MASK_USES_INPUTS_EMBEDS = ( "inputs_embeds" in inspect.signature(create_sliding_window_causal_mask).parameters ) def _call_attention_mask_factory( mask_factory, *, uses_inputs_embeds: bool, config: LaneformerConfig, inputs_embeds: torch.Tensor, attention_mask: Optional[torch.Tensor], cache_position: torch.Tensor, past_key_values: Optional[Cache], position_ids: Optional[torch.Tensor], ): if uses_inputs_embeds: return mask_factory( config=config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) return mask_factory( config=config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) def _attention_mask_needs_sliding_window_overlay( attention_mask: Optional[torch.Tensor], ) -> bool: return ( isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 4 and attention_mask.shape[-1] == attention_mask.shape[-2] ) def _apply_sliding_window_to_attention_mask( attention_mask: Optional[torch.Tensor], *, sliding_window: int, ) -> Optional[torch.Tensor]: if not _attention_mask_needs_sliding_window_overlay(attention_mask): return attention_mask sequence_length = int(attention_mask.shape[-1]) if sequence_length <= sliding_window: return attention_mask positions = torch.arange(sequence_length, device=attention_mask.device) outside_window = positions[None, :] <= positions[:, None] - int(sliding_window) if not bool(outside_window.any()): return attention_mask sliding_window_attention_mask = attention_mask.clone() if attention_mask.dtype == torch.bool: fill_value = False elif torch.is_floating_point(attention_mask): fill_value = torch.finfo(attention_mask.dtype).min else: fill_value = 0 sliding_window_attention_mask.masked_fill_( outside_window[None, None, :, :], fill_value, ) return sliding_window_attention_mask def _layer_uses_sliding_window_attention( config: LaneformerConfig, layer_index: int, ) -> bool: sliding_window = getattr(config, "sliding_window", None) if sliding_window is None or int(sliding_window) <= 0: return False swa_layers = getattr(config, "swa_layers", None) if swa_layers is not None: return layer_index in {int(layer) for layer in swa_layers} layer_types = getattr(config, "layer_types", None) if isinstance(layer_types, (list, tuple)) and layer_index < len(layer_types): return layer_types[layer_index] == "sliding_attention" sliding_window_n_layers = getattr(config, "sliding_window_n_layers", None) if sliding_window_n_layers is not None: return layer_index < int(sliding_window_n_layers) return True def _layer_attention_masks_from_provided_mask( config: LaneformerConfig, attention_mask: Optional[torch.Tensor], ) -> Optional[list[Optional[torch.Tensor]]]: if not _attention_mask_needs_sliding_window_overlay(attention_mask): return None sliding_window = getattr(config, "sliding_window", None) if sliding_window is None or int(sliding_window) <= 0: return [attention_mask] * config.num_hidden_layers sliding_window_attention_mask = _apply_sliding_window_to_attention_mask( attention_mask, sliding_window=int(sliding_window), ) return [ sliding_window_attention_mask if _layer_uses_sliding_window_attention(config, layer_id) else attention_mask for layer_id in range(config.num_hidden_layers) ] class ReduceMode(Enum): NO_REDUCE = "no_reduce" PRESENT = "present" PAST = "past" class LaneModule: _NO_REDUCE_MODE = 0 _PRESENT_MODE = 1 _PAST_MODE = 2 def __init__(self, no_reduce_scale: float, reduce_mode: ReduceMode): super().__init__() self.no_reduce_scale = no_reduce_scale self.reduce_mode = reduce_mode if reduce_mode is ReduceMode.NO_REDUCE: self.reduce_mode_code = self._NO_REDUCE_MODE elif reduce_mode is ReduceMode.PRESENT: self.reduce_mode_code = self._PRESENT_MODE elif reduce_mode is ReduceMode.PAST: self.reduce_mode_code = self._PAST_MODE else: raise ValueError(f"Unknown reduce mode: {reduce_mode}") def reduce_lanes(self, x: torch.Tensor, past: torch.Tensor | None) -> torch.Tensor: if self.reduce_mode_code == self._NO_REDUCE_MODE: return x * self.no_reduce_scale elif self.reduce_mode_code == self._PRESENT_MODE: return x.sum(dim=2, keepdim=True) elif self.reduce_mode_code == self._PAST_MODE: assert past is not None sum_past = torch.sum(past, dim=-2, keepdim=True) - past return x + sum_past else: raise ValueError(f"Unknown reduce mode code: {self.reduce_mode_code}") class LaneformerAttention(LaneModule, nn.Module): def __init__( self, config: LaneformerConfig, layer_idx: int, reduce_mode: ReduceMode = ReduceMode.PRESENT, ): super().__init__( no_reduce_scale=math.sqrt(config.num_lanes), reduce_mode=reduce_mode ) self.config = config self.layer_idx = layer_idx self.n_heads = config.num_attention_heads self.num_lanes = config.num_lanes self.n_kv_heads = ( config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads ) self.n_rep = config.num_attention_heads // self.n_kv_heads self.head_dim = config.hidden_size // config.num_attention_heads # Scaling & causal flag for backend attention APIs self.scaling = self.head_dim**-0.5 self.is_causal = True # decoder-only model self.num_key_value_groups = self.n_rep self.attention_dropout = getattr(config, "attention_dropout", 0.0) self.wq = LaneRowLinear( in_features=config.hidden_size, out_features=self.n_heads * self.head_dim, num_lanes=config.num_lanes, ) self.wk = LaneRowLinear( in_features=config.hidden_size, out_features=self.n_kv_heads * self.head_dim, num_lanes=config.num_lanes, ) self.wv = LaneRowLinear( in_features=config.hidden_size, out_features=self.n_kv_heads * self.head_dim, num_lanes=config.num_lanes, ) self.wo = LaneColumnLinear( in_features=self.n_heads * self.head_dim, out_features=config.hidden_size, num_lanes=config.num_lanes, ) def forward( self, x: torch.Tensor, freqs_cis: torch.Tensor, attention_mask: Optional[torch.Tensor], *, past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ): """ Forward pass of the attention module. Args: x (torch.Tensor): Input tensor (B, S, L, D) freqs_cis (torch.Tensor): Precomputed frequency tensor. Returns: torch.Tensor: Output tensor after attention (B, S, L, D) """ bs, seqlen, _, _ = x.shape xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) xq = xq.reshape(bs, seqlen, -1, self.head_dim) # (B, T, L, n_kv_heads / L, head_dim) xk = xk.reshape(bs, seqlen, -1, self.head_dim) # (B, T, L, n_kv_heads / L, head_dim) xv = xv.reshape(bs, seqlen, -1, self.head_dim) xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) xq = xq.transpose(1, 2) # (B, n_heads, T, head_dim) xk = xk.transpose(1, 2) # (B, n_kv_heads, T, head_dim) xv = xv.transpose(1, 2) # (B, n_kv_heads, T, head_dim) # Cache update BEFORE replication if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": None, "cos": None, "cache_position": cache_position} xk, xv = past_key_values.update(xk, xv, self.layer_idx, cache_kwargs) attention_interface = _attention_interface(self.config) output, _ = attention_interface( self, query=xq, key=xk, value=xv, attention_mask=attention_mask, scaling=self.scaling, dropout=0.0 if not self.training else self.attention_dropout, **kwargs, ) output = output.contiguous().view( bs, seqlen, self.num_lanes, self.n_heads * self.head_dim // self.num_lanes, ) output = self.wo(output) return output class LaneformerMLP(LaneModule, nn.Module): def __init__( self, config: LaneformerConfig, reduce_mode: ReduceMode = ReduceMode.PRESENT ): super().__init__( no_reduce_scale=math.sqrt(config.num_lanes), reduce_mode=reduce_mode ) ffn_dim = config.intermediate_size dim = config.hidden_size num_lanes = config.num_lanes self.w1 = LaneRowLinear(dim, ffn_dim, num_lanes=num_lanes) self.w2 = LaneColumnLinear(ffn_dim, dim, num_lanes=num_lanes) self.w3 = LaneRowLinear(dim, ffn_dim, num_lanes=num_lanes) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class LaneformerDecoderLayer(nn.Module): def __init__( self, config: LaneformerConfig, layer_idx: int, attention_reduce_mode: ReduceMode, mlp_reduce_mode: ReduceMode, broadcast_attention_to_future: bool, broadcast_mlp_to_future: bool, ): super().__init__() self.n_heads = config.num_attention_heads self.dim = config.hidden_size self.attention = LaneformerAttention( config, layer_idx=layer_idx, reduce_mode=attention_reduce_mode ) self.feed_forward = LaneformerMLP(config, reduce_mode=mlp_reduce_mode) norm_fun = ( nn.RMSNorm if config.replicated_rmsn_scale else partial(LaneRMSNorm, num_lanes=config.num_lanes) ) self.attention_norm = norm_fun(config.hidden_size, eps=config.norm_eps) self.ffn_norm = norm_fun(config.hidden_size, eps=config.norm_eps) self.num_lanes = config.num_lanes self.broadcast_attention_to_future = broadcast_attention_to_future self.broadcast_mlp_to_future = broadcast_mlp_to_future def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, past_attention: Optional[torch.Tensor] = None, past_mlp: Optional[torch.Tensor] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, freqs_cis: Optional[torch.Tensor] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: residual = hidden_states hidden_states = self.attention_norm(hidden_states) hidden_states = self.attention( x=hidden_states, freqs_cis=freqs_cis, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, **kwargs, ) future_attention = hidden_states if self.broadcast_attention_to_future else None hidden_states = self.attention.reduce_lanes(hidden_states, past=past_attention) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.ffn_norm(hidden_states) hidden_states = self.feed_forward(hidden_states) future_mlp = hidden_states if self.broadcast_mlp_to_future else None hidden_states = self.feed_forward.reduce_lanes(hidden_states, past=past_mlp) hidden_states = residual + hidden_states return hidden_states, future_attention, future_mlp # ============================================================ # Backbone & HF wrappers # ============================================================ class LaneformerPreTrainedModel(PreTrainedModel): config_class = LaneformerConfig config: LaneformerConfig base_model_prefix = "model" supports_gradient_checkpointing = False _no_split_modules = [ "LaneformerModel", "LaneformerDecoderLayer", "LaneformerAttention", "LaneformerMLP", ] _skip_keys_device_placement = ["past_key_values"] # Advertise support for Transformers attention backends (vLLM / flash / sdpa routing) _supports_attention_backend = True _supports_sdpa = True class LaneformerModel(LaneformerPreTrainedModel): def __init__(self, config: LaneformerConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.config = config self.vocab_size = config.vocab_size self.n_layers = config.num_hidden_layers self.eos_id = config.eos_token_id self.num_lanes = config.num_lanes self.broadcast_delay = config.broadcast_delay self.use_attention_comm = getattr(config, "use_attention_comm", config.use_comm) self.use_mlp_comm = getattr(config, "use_mlp_comm", config.use_comm) self.use_early_comm = config.use_early_comm self.pre_norm_lane_agg = config.pre_norm_lane_agg self.lm_head_type = config.lm_head_type self.replicated_rmsn_scale = config.replicated_rmsn_scale self.swa_layers = set(config.swa_layers) self.rope_dim = config.hidden_size // config.num_attention_heads self.rope_theta = config.rope_theta self.rope_scaling_args = _rope_scaling_from_config(config) self.tok_embeddings = nn.Embedding( config.vocab_size, config.hidden_size, self.padding_idx ) # Build lane blocks self.layers = nn.ModuleList() for layer_id in range(config.num_hidden_layers): if self.broadcast_delay == 0 or layer_id < self.broadcast_delay: attention_reduce_mode = ( ReduceMode.PRESENT if self.use_early_comm else ReduceMode.NO_REDUCE ) mlp_reduce_mode = ( ReduceMode.PRESENT if self.use_early_comm else ReduceMode.NO_REDUCE ) else: attention_reduce_mode = ( ReduceMode.NO_REDUCE if not (self.use_attention_comm) else ReduceMode.PAST ) mlp_reduce_mode = ( ReduceMode.NO_REDUCE if not (self.use_mlp_comm) else ReduceMode.PAST ) broadcast_attention_to_future = ( config.num_hidden_layers - self.broadcast_delay > layer_id ) and self.use_attention_comm broadcast_mlp_to_future = ( config.num_hidden_layers - self.broadcast_delay > layer_id ) and self.use_mlp_comm self.layers.append( LaneformerDecoderLayer( config=config, layer_idx=layer_id, attention_reduce_mode=attention_reduce_mode, mlp_reduce_mode=mlp_reduce_mode, broadcast_attention_to_future=broadcast_attention_to_future, broadcast_mlp_to_future=broadcast_mlp_to_future, ) ) if self.pre_norm_lane_agg and not self.lm_head_type == "replicate": raise ValueError( "pre_norm_lane_agg is only supported with lm_head_type='replicate'" ) norm_fun = ( nn.RMSNorm if config.replicated_rmsn_scale or config.pre_norm_lane_agg else partial(LaneRMSNorm, num_lanes=self.num_lanes) ) self.norm = norm_fun(config.hidden_size, eps=config.norm_eps) self.gradient_checkpointing = False self.post_init() # Required by vllm integration def get_input_embeddings(self): return self.tok_embeddings # Required by vllm integration def set_input_embeddings(self, new_emb): self.tok_embeddings = new_emb def _compute_freqs_cis(self, position_ids: torch.Tensor) -> torch.Tensor: return freqs_cis_for_position_ids( position_ids, self.rope_dim, self.rope_theta, self.rope_scaling_args, ) def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, **kwargs, ): if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You must specify exactly one of input_ids or inputs_embeds" ) if inputs_embeds is None: inputs_embeds: torch.Tensor = self.tok_embeddings(input_ids) use_cache = use_cache if use_cache is not None else self.config.use_cache past_key_values_for_forward = past_key_values if use_cache else None if use_cache and past_key_values_for_forward is None: past_key_values_for_forward = DynamicCache(config=self.config) if cache_position is None: past_seen_tokens = ( past_key_values_for_forward.get_seq_length() if past_key_values_for_forward is not None else 0 ) cache_position: torch.Tensor = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) layer_types = list(getattr(self.config, "layer_types", [])) if len(layer_types) != self.config.num_hidden_layers: layer_types = [ "sliding_attention" if layer_id in self.swa_layers else "full_attention" for layer_id in range(self.config.num_hidden_layers) ] unknown_layer_types = set(layer_types) - {"full_attention", "sliding_attention"} if unknown_layer_types: raise ValueError( f"Unknown attention layer types: {sorted(unknown_layer_types)}" ) has_sliding_layers = "sliding_attention" in layer_types if self.config.sliding_window is None and ( self.swa_layers or has_sliding_layers ): raise ValueError("sliding_window must be set for layers in swa_layers") causal_mask_mapping = {} if "full_attention" in layer_types: causal_mask_mapping["full_attention"] = _call_attention_mask_factory( create_causal_mask, uses_inputs_embeds=_CAUSAL_MASK_USES_INPUTS_EMBEDS, config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values_for_forward, position_ids=position_ids, ) if has_sliding_layers: sliding_attention_mask = _apply_sliding_window_to_attention_mask( attention_mask, sliding_window=int(self.config.sliding_window), ) causal_mask_mapping["sliding_attention"] = _call_attention_mask_factory( create_sliding_window_causal_mask, uses_inputs_embeds=_SLIDING_MASK_USES_INPUTS_EMBEDS, config=self.config, inputs_embeds=inputs_embeds, attention_mask=sliding_attention_mask, cache_position=cache_position, past_key_values=past_key_values_for_forward, position_ids=position_ids, ) freqs_cis = self._compute_freqs_cis(position_ids).unsqueeze(-2) hidden_states = inputs_embeds hidden_states = hidden_states[:, :, None, :].expand(-1, -1, self.num_lanes, -1) past_attentions = [] past_mlps = [] for i in range(self.config.num_hidden_layers): if self.broadcast_delay == 0 or i < self.broadcast_delay: past_attention = None past_mlp = None else: past_attention = past_attentions[i - self.broadcast_delay] past_mlp = past_mlps[i - self.broadcast_delay] hidden_states, future_attention, future_mlp = self.layers[i]( hidden_states=hidden_states, attention_mask=causal_mask_mapping[layer_types[i]], position_ids=position_ids, past_key_values=past_key_values_for_forward, past_attention=past_attention, past_mlp=past_mlp, use_cache=use_cache, cache_position=cache_position, freqs_cis=freqs_cis, **kwargs, ) past_attentions.append(future_attention) past_mlps.append(future_mlp) if self.pre_norm_lane_agg: # (B, S, L, D) -> (B, S, D) hidden_states = hidden_states.sum(dim=-2) hidden_states = self.norm(hidden_states) elif self.lm_head_type == "replicate": # and naturally not pre_lane_agg # (B, S, L, D) -> (B, S, D) hidden_states = self.norm(hidden_states) hidden_states = hidden_states.mean(dim=-2) else: # (B, S, L, D) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values_for_forward if use_cache else None, ) @auto_docstring class LaneformerForCausalLM(LaneformerPreTrainedModel, GenerationMixin): _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config: LaneformerConfig): super().__init__(config) if config.tie_word_embeddings: raise ValueError("tie_word_embeddings is not supported for Laneformer") self.model = LaneformerModel(config) self.vocab_size = config.vocab_size self.lm_head_type = config.lm_head_type if self.lm_head_type == "replicate": self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) elif self.lm_head_type == "lane": self.lm_head = LaneLMHead( in_features=config.hidden_size, out_features=config.vocab_size, num_lanes=config.num_lanes, ) elif self.lm_head_type == "vocab_parallel": self.lm_head = LaneRowLinear( in_features=config.hidden_size, out_features=config.vocab_size, num_lanes=config.num_lanes, ) # Initialize weights and apply final processing self.post_init() def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def resize_token_embeddings( self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None, mean_resizing: bool = True, ) -> nn.Embedding: if self.lm_head_type != "replicate" and ( new_num_tokens is not None or pad_to_multiple_of is not None ): raise NotImplementedError( "resize_token_embeddings is only supported for " "lm_head_type='replicate'" ) return super().resize_token_embeddings( new_num_tokens=new_num_tokens, pad_to_multiple_of=pad_to_multiple_of, mean_resizing=mean_resizing, ) @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: r""" Example with a tiny randomly initialized Laneformer model: ```python >>> import torch >>> from torchtitan.models.laneformer.hf import ( ... LaneformerConfig, ... LaneformerForCausalLM, ... ) >>> config = LaneformerConfig( ... hidden_size=16, ... num_hidden_layers=2, ... num_attention_heads=4, ... num_key_value_heads=2, ... intermediate_size=32, ... vocab_size=64, ... num_lanes=2, ... ) >>> model = LaneformerForCausalLM(config) >>> input_ids = torch.tensor([[1, 2, 3]]) >>> logits = model(input_ids=input_ids, use_cache=False).logits >>> logits.shape torch.Size([1, 3, 64]) ```""" outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if # we are not computing the loss. slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) if self.lm_head_type == "replicate": logits = self.lm_head(hidden_states[:, slice_indices, :]) else: logits = self.lm_head(hidden_states[:, slice_indices, :, :]) if self.lm_head_type == "vocab_parallel": # (B, S, L, D) -> (B, S, D) B, S, L, D = logits.shape logits = logits.reshape(B, S, L * D) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs, ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )