Azrail/smallm_70_rope

1import torch2from torch import nn3import torch.nn.functional as F4from transformers import PreTrainedModel, GenerationMixin5from transformers.cache_utils import Cache, DynamicCache6from transformers.modeling_outputs import (7    BaseModelOutputWithPast,8    CausalLMOutputWithPast,9)10from .config import SmalLmConfig11from typing import Optional12import logging13from einops import rearrange14from transformers.modeling_attn_mask_utils import AttentionMaskConverter15from einops._torch_specific import allow_ops_in_compiled_graph16 17allow_ops_in_compiled_graph()18from transformers.utils import is_flash_attn_2_available19 20if is_flash_attn_2_available():21    from flash_attn import flash_attn_varlen_func22    from flash_attn.bert_padding import unpad_input, pad_input23 24 25 26logger = logging.getLogger(__name__)27 28 29class SwiGLU(nn.Module):30    def __init__(31        self, input_size: int, hidden_size: int, bias: bool = False, *args, **kwargs32    ):33        super().__init__(*args, **kwargs)34        self.input_size = input_size35        self.hidden_size = hidden_size36        self.up_proj = nn.Linear(input_size, hidden_size * 2, bias=bias)37        self.down_proj = nn.Linear(hidden_size, input_size, bias=bias)38 39    def forward(self, x):40        up_gate = self.up_proj(x)41        up, gate = rearrange(up_gate, "... (d span) -> span ... d", d=self.hidden_size)42        down = F.silu(gate) * up43        return self.down_proj(down)44 45 46class Router(nn.Module):47    def __init__(self, config: SmalLmConfig, *args, **kwargs):48        49        super().__init__(*args, **kwargs)50        self.config = config51        self.experts_to_select = self.config.token_experts - self.config.shared_experts52        self.gate = nn.Linear(config.hidden_size, config.routed_experts, bias=False)53        self.gate_noise = (54            nn.Linear(config.hidden_size, config.routed_experts, bias=False)55            if config.noisy_experts is True56            else None57        )58        self.bias_coef = config.balancing_coef59        self.register_buffer(60            "bias", torch.zeros(config.routed_experts), persistent=True61        )62        self.register_buffer(63            "expert_counts", torch.zeros(config.routed_experts), persistent=False64        )65 66    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]:67        # calculating with fp32 for stability68        # num_tokens n_shared_experts69        gate_logits = self.gate(x)70        if self.gate_noise is not None:71            gate_logits_noise = F.softplus(self.gate_noise(x))72            gate_logits_noise = torch.randn_like(gate_logits_noise) * gate_logits_noise73            gate_logits = gate_logits + gate_logits_noise74 75        gate_weights = gate_logits.sigmoid()76        original_weights = gate_weights77 78        gate_weights = gate_weights + self.bias79 80        _, top_experts_idx = torch.topk(gate_weights, self.experts_to_select, dim=-1)81        counts = torch.bincount(82            top_experts_idx.flatten(), minlength=self.config.routed_experts83        ).detach()84        if self.training:85            self.expert_counts += counts86        top_experts_weights = original_weights.gather(1, top_experts_idx)87        top_experts_weights = top_experts_weights / top_experts_weights.sum(88            dim=-1, keepdim=True89        )90        return top_experts_idx, top_experts_weights.type_as(x), counts.tolist()91 92    def update_bias(self):93        mean = self.expert_counts.float().mean()94        delta = self.bias_coef * torch.sign(mean - self.expert_counts)95        self.bias += delta96        self.expert_counts.zero_()97 98 99class MoE(nn.Module):100    def __init__(self, config: SmalLmConfig, *args, **kwargs):101        super().__init__(*args, **kwargs)102        self.config = config103        self.shared_experts = SwiGLU(104            config.hidden_size,105            config.shared_experts * config.expert_size,106            config.moe_bias,107        )108        self.routed_experts = nn.ModuleList(109            [110                SwiGLU(config.hidden_size, config.expert_size, config.moe_bias)111                for _ in range(config.routed_experts)112            ]113        )114        self.router = Router(config)115 116    def forward(self, x: torch.Tensor) -> torch.Tensor:117        shape = x.size()118        x = x.view(-1, self.config.hidden_size)119        experts_idx, experts_weights, counts = self.router(x)120        out = torch.zeros_like(x)121        for i, expert in enumerate(self.routed_experts):122            if counts[i] == 0:123                continue124            idx, pos = torch.where(experts_idx == i)125            out[idx] += expert(x[idx]) * experts_weights[idx, pos, None]126        shared_out = self.shared_experts(x)127        return (out + shared_out).view(shape)128 129 130def build_alibi_bias(config: SmalLmConfig) -> torch.Tensor:131    """Build ALiBi for specified number of heads:132 133    Returns:134        Tensor with ALiBi biases, shape: [num heads]135    """136    bias = (137        2**-8138        / config.num_attention_heads139        * torch.arange(1, config.num_attention_heads + 1).float()140    )141    return bias142 143 144def calc_rotation(num_rotaitions, dim, base, seq_len):145    return (146        dim147        * torch.log(torch.tensor(seq_len).float() / (num_rotaitions * 2 * torch.pi))148        / torch.log(torch.tensor(base))149    )150 151 152def get_ramp_interpolation(min_idx, max_idx, thetas_dim, eps=1e-6):153    if min_idx == max_idx:154        max_idx += eps155    mult = (torch.arange(thetas_dim) - min_idx) / (max_idx - min_idx)156    mult = torch.clamp(mult, 0, 1)157    return 1 - mult158 159 160def build_rope_bias(config: SmalLmConfig) -> torch.Tensor:161    dim = config.head_size162 163    theta = 1.0 / (config.rope_base ** (torch.arange(0, dim, 2).float() / dim))164 165    # neural tangent kernel by part korrection166    if config.max_seq_len > config.original_seq_len:167        scale = config.max_seq_len / config.original_seq_len168        # from idea that lambda = 2pi / theta_i and lmbad = seq_len / num_rotations, lambda - wavelen169        low_interpolation_idx = max(170            0,171            torch.ceil(172                calc_rotation(173                    config.high_rotations,174                    dim,175                    config.rope_base,176                    config.original_seq_len,177                )178            ).item(),179        )180        high_interpolation_idx = min(181            dim - 1,182            torch.floor(183                calc_rotation(184                    config.low_rotations, dim, config.rope_base, config.original_seq_len185                )186            ).item(),187        )188        interpolation_mult = get_ramp_interpolation(189            low_interpolation_idx, high_interpolation_idx, dim // 2190        )191        theta = (1 - interpolation_mult) * theta / scale + interpolation_mult * theta192 193    seq_idx = torch.arange(config.max_seq_len)194    seq_theta = torch.outer(seq_idx, theta)195    bias = torch.polar(torch.ones_like(seq_theta), seq_theta)196    return bias197 198 199def apply_rope_bias(x: torch.Tensor, precompute_bias: torch.Tensor) -> torch.Tensor:200    ini_dtype = x.dtype201    # for stbility to fp32, also need for torch202    x = rearrange(x.float(), "b n s (d i) -> b n s d i", i=2).contiguous()203    x = torch.view_as_complex(x)204    x = x * precompute_bias205    x = torch.view_as_real(x)206    x = rearrange(x, "b n s d i -> b n s (d i)")207    return x.to(ini_dtype)208 209 210def flash_attention_forward(211        module: nn.Module,212        x: torch.Tensor,213        query: torch.Tensor,214        key: torch.Tensor,215        value: torch.Tensor,216        attention_mask: torch.Tensor,217        alibi_slope: Optional[torch.Tensor]218) -> torch.Tensor:219        query = rearrange(query, "b n s d -> b s n d")220        key = rearrange(key, "b n s d -> b s n d")221        value = rearrange(value, "b n s d -> b s n d")222        query, idx_q, cu_seqlens_q, max_seqlen_q, _ = unpad_input(query, attention_mask)223        key, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(key, attention_mask)224        value, _, _, _, _ = unpad_input(value, attention_mask)225 226        key = key.contiguous()227        value = value.contiguous()228        query = query.contiguous()229 230        attention_probs = flash_attn_varlen_func(231            query,232            key,233            value,234            cu_seqlens_q=cu_seqlens_q,235            cu_seqlens_k=cu_seqlens_k,236            max_seqlen_q=max_seqlen_q,237            max_seqlen_k=max_seqlen_k,238            dropout_p=module.config.attention_dropout if module.training else 0.0,239            causal=True,240            alibi_slopes=alibi_slope if module.config.attention_bias == "alibi" else None,241        )242        attention_probs = pad_input(attention_probs, idx_q, x.size(0), x.size(1))243        out = rearrange(attention_probs, "b s n d -> b s (n d)")244        return out, None245 246 247def sdpa_attention_forward(248    module: nn.Module,249    x: torch.Tensor,250    query: torch.Tensor,251    key: torch.Tensor,252    value: torch.Tensor,253    attention_mask: torch.Tensor,254    alibi_slope: Optional[torch.Tensor]255) -> torch.Tensor:256    is_causal = attention_mask is None and query.size(-2) > 1257 258    attention_probs = F.scaled_dot_product_attention(259        query,260        key,261        value,262        attn_mask=attention_mask,263        enable_gqa=True,264        is_causal=is_causal,265        dropout_p=module.config.attention_dropout if module.training else 0.0,266    )267    out = rearrange(attention_probs, "b n s d -> b s (n d)")268    269    return out, None270 271def eager_attention_forward(272    module: nn.Module,273    x: torch.Tensor,274    query: torch.Tensor,275    key: torch.Tensor,276    value: torch.Tensor,277    attention_mask: torch.Tensor,278    alibi_slope: Optional[torch.Tensor]279) -> torch.Tensor:280    query = rearrange(query, 'b (kv group) s d -> b kv group s d', kv=module.config.num_kv_heads, group=module.head_per_group)281    key = rearrange(key, 'b kv s d -> b kv 1 s d')282    value = rearrange(283        value, 'b kv s d -> b kv 1 s d'284        )285    attention_weights = query @ key.transpose(-1, -2)286    attention_probs = F.dropout(attention_weights / torch.sqrt(287        torch.tensor(value.size(-1), device=x.device)288        ),289        p=module.config.attention_dropout if module.training else 0.0290        )291    if alibi_slope is not None:292        alibi_slope = rearrange(293            alibi_slope, 'b n s s -> b kv group s s', kv=module.config.num_kv_heads, group=module.head_per_group294            )295        attention_probs = attention_probs + alibi_slope296    elif alibi_slope is None and attention_mask is not None:297        attention_mask = attention_mask.expand(-1, module.config.num_attention_heads, -1, -1)298        attention_mask = rearrange(299            attention_mask, 'b (kv group) s1 s2 -> b kv group s1 s2', kv=module.config.num_kv_heads, group=module.head_per_group300            )301        attention_probs = attention_probs + attention_mask302    attention_probs = F.softmax(attention_probs, dim=-1)303    attention_probs = attention_probs @ value304    out = rearrange(attention_probs, "b kv group s d -> b s (kv group d)")305    return out, attention_weights306 307 308ALL_ATTENTION_FUNCTIONS = {309    "eager": eager_attention_forward,310    "sdpa": sdpa_attention_forward,311    "flash_attention_2": flash_attention_forward,312}313 314 315class CausalSelfAttention(nn.Module):316    def __init__(self, config: SmalLmConfig, layer_idx: int, *args, **kwargs):317        super().__init__(*args, **kwargs)318        if config.num_attention_heads % config.num_kv_heads != 0:319            raise ValueError("Num attention heads should divided by num kv heads")320 321        self.config = config322        self.layer_idx = layer_idx323        self.head_per_group = config.num_attention_heads // config.num_kv_heads324        self.q_proj = nn.Linear(325            config.hidden_size,326            config.head_size * config.num_attention_heads,327            bias=config.attention_bias,328        )329        self.kv_proj = nn.Linear(330            config.hidden_size,331            config.head_size * config.num_kv_heads * 2,332            bias=config.attention_bias,333        )334        self.out_proj = nn.Linear(335            config.head_size * config.num_attention_heads,336            config.hidden_size,337            bias=config.attention_bias,338        )339 340    def forward(341        self,342        x: torch.Tensor,343        attention_mask: torch.Tensor,344        past_key_values: Optional[Cache | torch.FloatTensor],345        cache_position: Optional[torch.LongTensor],346        bias: torch.Tensor,347    ):348        q = self.q_proj(x)349        kv = self.kv_proj(x)350        q = rearrange(q, "b s (n d) -> b n s d", n=self.config.num_attention_heads)351        k, v = rearrange(kv, "b s (n d q) -> q b n s d", q=2, d=self.config.head_size)352 353        if self.config.positional_bias_type == "rope":354            k = apply_rope_bias(k, bias)355            q = apply_rope_bias(q, bias)356 357        if past_key_values is not None:358            # for static cache359            cach_kwargs = {"cache_position": cache_position}360            k, v = past_key_values.update(361                key_states=k,362                value_states=v,363                layer_idx=self.layer_idx,364                cache_kwargs=cach_kwargs,365            )366 367        attention_interface = eager_attention_forward368        if self.config._attn_implementation != "eager":369            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]370 371        out, attention_weights = attention_interface(372            self,373            x,374            q,375            k,376            v,377            attention_mask,378            bias if self.config.positional_bias_type == "alibi" else None379        )380 381        out = self.out_proj(out)382        return out, attention_weights383 384 385class WeightedResidual(nn.Module):386    def __init__(self, config: SmalLmConfig, *args, **kwargs):387        super().__init__(*args, **kwargs)388        self.weight = nn.Parameter(389            torch.ones(config.hidden_size), requires_grad=config.static_residual390        )391 392    def forward(self, short, long):393        return self.weight * short + long394 395 396class Block(nn.Module):397    def __init__(self, config: SmalLmConfig, layer_idx: int, *args, **kwargs):398        super().__init__(*args, **kwargs)399        self.attn_norm = nn.RMSNorm(400            config.hidden_size,401            eps=config.rms_norm_eps,402            elementwise_affine=config.rms_affine,403        )404        self.ffn_norm = nn.RMSNorm(405            config.hidden_size,406            eps=config.rms_norm_eps,407            elementwise_affine=config.rms_affine,408        )409        self.dropout1 = nn.Dropout(config.layer_dropout)410        self.dropout2 = nn.Dropout(config.layer_dropout)411        self.attention = CausalSelfAttention(config, layer_idx)412        self.mlp = (413            MoE(config)414            if (415                config.use_moe416                and layer_idx % config.moe_period == 0417                and layer_idx > config.no_moe_layers418            )419            else SwiGLU(config.hidden_size, config.intermediate_size, config.mlp_bias)420        )421        self.attention_residual = WeightedResidual(config)422        self.ffn_residual = WeightedResidual(config)423 424    def forward(425        self,426        inputs_embeds: torch.Tensor,427        attention_mask: torch.Tensor,428        past_key_values: Optional[Cache | torch.FloatTensor],429        output_attentions: bool,430        cache_position: Optional[torch.LongTensor],431        bias: torch.Tensor,432    ) -> tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:433        identity = inputs_embeds434 435        # attention block436        out = self.attn_norm(inputs_embeds)437        out, attention_probs = self.attention(438            out, attention_mask, past_key_values, cache_position, bias439        )440        out = self.dropout1(out)441        identity = self.attention_residual(identity, out)442 443        # swiglu / MoE block444        out = self.dropout2(self.mlp(self.ffn_norm(identity)))445        out = self.ffn_residual(identity, out)446        if output_attentions:447            return out, attention_probs448        return (out,)449 450 451class SmalLmPreTrainedModel(PreTrainedModel):452    config_class = SmalLmConfig453    base_model_prefix = "model"454    supports_gradient_checkpointing = True455    _no_split_modules = ["Block"]456    _skip_keys_device_placement = "past_key_values"457    _supports_sdpa = True458    _supports_flash_attn_2 = True459    def __init__(self, *inputs, **kwargs):460        super().__init__(*inputs, **kwargs)461 462    def _init_weights(self, module):463        std = self.config.initializer_range464        if isinstance(module, nn.Linear):465            torch.nn.init.normal_(module.weight, mean=0.0, std=std)466            if module.bias is not None:467                torch.nn.init.zeros_(module.bias)468        elif isinstance(module, nn.Embedding):469            torch.nn.init.normal_(module.weight, mean=0.0, std=std)470            module.weight.data[self.pad_idx].zero_()471 472 473class SmalLmModel(SmalLmPreTrainedModel):474    def __init__(self, config: SmalLmConfig, *args, **kwargs):475        super().__init__(config, *args, **kwargs)476        self.config = config477        self.pad_idx = config.pad_token_id478        self.pad_token_id = config.pad_token_id479        self.vocab_size = config.vocab_size480        self.config = config481        precompute_bias = (482            build_alibi_bias(config)483            if config.positional_bias_type == "alibi"484            else build_rope_bias(config)485        )486        self.register_buffer("precompute_bias", precompute_bias, persistent=False)487        # не забыть про sharing weights на output голове self.embedding.weight = self.output.weight488        self.embedding = nn.Embedding(489            self.vocab_size, config.hidden_size, padding_idx=config.pad_token_id490        )491        self.embedding_dropout = nn.Dropout(config.embedding_dropout)492        self.layers = nn.ModuleList(493            [Block(config, idx) for idx in range(1, config.num_hidden_layers + 1)]494        )495        self.out_norm = nn.RMSNorm(496            config.hidden_size,497            eps=config.rms_norm_eps,498            elementwise_affine=config.rms_affine,499        )500 501        self.gradient_checkpointing = False502        self.post_init()503 504    def get_input_embeddings(self):505        return self.embedding506 507    def set_input_embeddings(self, value):508        self.embedding = value509 510    def forward(511        self,512        # input options513        input_ids: torch.LongTensor = None,514        attention_mask: Optional[torch.Tensor] = None,515        inputs_embeds: Optional[torch.FloatTensor] = None,516        # output options517        output_attentions: Optional[bool] = None,518        output_hidden_states: Optional[bool] = None,519        return_dict: Optional[bool] = None,520        # cache options521        use_cache: Optional[bool] = None,522        past_key_values: Optional[Cache | torch.FloatTensor] = None,523        cache_position: Optional[torch.LongTensor] = None,524        position_ids: Optional[torch.LongTensor] = None,525        **kwargs,526    ) -> tuple | BaseModelOutputWithPast:527        # check additional parameters528        output_hidden_states = (529            output_hidden_states530            if output_hidden_states is not None531            else self.config.output_hidden_states532        )533        use_cache = (534            use_cache535            if use_cache is not None536            else (False if self.training else self.config.use_cache)537        )538        return_dict = (539            return_dict if return_dict is not None else self.config.return_dict540        )541 542        if input_ids is not None and inputs_embeds is not None:543            raise ValueError(544                "You must specify only input_ids or inputs_embeds, not both"545            )546 547        if self.training and use_cache:548            use_cache = False549 550        if inputs_embeds is None:551            inputs_embeds = self.embedding(input_ids)552 553        if use_cache and past_key_values is None:554            past_key_values = DynamicCache()555 556        # calculating position for StaticCache557        if cache_position is None:558            last_position = (559                past_key_values.get_seq_length() if past_key_values is not None else 0560            )561            cache_position = torch.arange(562                last_position,563                last_position + inputs_embeds.size(1),564                device=inputs_embeds.device,565            )566 567        causal_mask = self._get_causal_masks(568            attention_mask, inputs_embeds, past_key_values, cache_position569        )570        if self.config.positional_bias_type == "rope":571            end_pos = (572                inputs_embeds.size(1)573                if past_key_values is None574                else cache_position[-1] + 1575            )576            start_pos = 0 if past_key_values is None else cache_position[0]577            bias = self.precompute_bias[start_pos:end_pos]578 579        elif self.config.positional_bias_type == "alibi":580            if self.config._attn_implementation == "flash_attention_2":581                bias = self.precompute_bias582            else:583                i = torch.arange(584                    (585                        inputs_embeds.size(1)586                        if past_key_values is None587                        else cache_position[-1] + 1588                    ),589                    device=inputs_embeds.device,590                )591                bias = i[:, None] - i[None, :]592                bias = torch.tril(bias).expand(593                    inputs_embeds.size(0), self.config.num_attention_heads, -1, -1594                ) * rearrange(self.precompute_bias, "n -> 1 n 1 1")595                if causal_mask is not None:596                    causal_mask = causal_mask + bias597                else:598                    causal_mask = bias599 600        hidden_state = inputs_embeds601        hidden_states = [hidden_state] if output_hidden_states else None602        attentions = [] if output_attentions else None603        for idx, layer in enumerate(self.layers, 1):604            if self.gradient_checkpointing:605                # for details see:606                # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3107607                # https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3149608                layer_out = self._gradient_checkpointing_func(609                    layer.__call__,610                    hidden_state,611                    causal_mask,612                    past_key_values,613                    output_attentions,614                    cache_position,615                    bias,616                )617            else:618                layer_out = layer(619                    hidden_state,620                    causal_mask,621                    past_key_values,622                    output_attentions,623                    cache_position,624                    bias,625                )626            hidden_state = layer_out[0]627            if output_hidden_states:628                hidden_states.append(hidden_state)629            if output_attentions:630                attentions.append(layer_out[1])631 632        hidden_state = self.out_norm(hidden_state)633        out = BaseModelOutputWithPast(634            last_hidden_state=hidden_state,635            past_key_values=past_key_values if use_cache else None,636            hidden_states=tuple(hidden_states) if hidden_states is not None else None,637            attentions=tuple(attentions) if attentions is not None else None,638        )639        return out if return_dict else out.to_tuple()640 641    def _get_causal_masks(642        self,643        attention_mask: Optional[torch.Tensor],644        inputs_embeds: torch.Tensor,645        past_key_values: Optional[torch.Tensor],646        cache_position: Optional[torch.Tensor],647    ):648        if self.config._attn_implementation == "flash_attention_2":649            if attention_mask is None:650                attention_mask = torch.ones(651                    (inputs_embeds.size(0), inputs_embeds.size(1)), device=inputs_embeds.device652                    ).long()653            return attention_mask654        dtype, device = inputs_embeds.dtype, inputs_embeds.device655        past_token = (656            past_key_values.get_seq_length() if past_key_values is not None else 0657        )658        if attention_mask is not None and torch.all(attention_mask == 0.0):659            return None660        if AttentionMaskConverter._ignore_causal_mask_sdpa(661            attention_mask=attention_mask,662            inputs_embeds=inputs_embeds,663            past_key_values_length=past_token,664            is_training=self.training,665        ):666            return None667 668        sequence_length = inputs_embeds.size(1)669        target_length = (670            attention_mask.size(-1)671            if isinstance(attention_mask, torch.Tensor)672            else past_token + sequence_length + 1673        )674 675        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(676            attention_mask=attention_mask,677            sequence_length=sequence_length,678            target_length=target_length,679            dtype=dtype,680            device=device,681            cache_position=cache_position,682            batch_size=inputs_embeds.size(0),683        )684 685        min_dtype = torch.finfo(dtype).min686        causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)687        return causal_mask688 689    @staticmethod690    def _prepare_4d_causal_attention_mask_with_cache_position(691        attention_mask: Optional[torch.Tensor],692        sequence_length: int,693        target_length: int,694        dtype: torch.dtype,695        device: torch.device,696        cache_position: Optional[torch.Tensor],697        batch_size: int,698    ):699        if attention_mask is not None and attention_mask.dim() == 4:700            causal_mask = attention_mask701        else:702            min_dtype = torch.finfo(dtype).min703            causal_mask = torch.full(704                (sequence_length, target_length),705                fill_value=min_dtype,706                dtype=dtype,707                device=device,708            )709            if sequence_length != 1:710                causal_mask = torch.triu(causal_mask, diagonal=1)711            causal_mask *= torch.arange(712                target_length, device=device713            ) > cache_position.reshape(-1, 1)714            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)715            if attention_mask is not None:716                causal_mask = causal_mask.clone()717                mask_length = attention_mask.shape[-1]718                padding_mask = (719                    causal_mask[:, :, :, :mask_length]720                    + attention_mask[:, None, None, :]721                )722                padding_mask = padding_mask == 0723                causal_mask[:, :, :, :mask_length] = causal_mask[724                    :, :, :, :mask_length725                ].masked_fill(padding_mask, min_dtype)726        return causal_mask727 728 729class SmalLmForCausalLM(SmalLmPreTrainedModel, GenerationMixin):730    _tied_weights_keys = ["lm_head.weight"]731 732    def __init__(self, config: SmalLmConfig, *args, **kwargs):733        super().__init__(config, *args, **kwargs)734        self.config = config735        self.model = SmalLmModel(config)736        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)737        self.post_init()738 739    def get_output_embeddings(self):740        return self.lm_head741 742    def set_output_embeddings(self, new_embeddings):743        self.lm_head = new_embeddings744 745    def forward(746        self,747        # input options748        input_ids: torch.LongTensor = None,749        attention_mask: Optional[torch.Tensor] = None,750        inputs_embeds: Optional[torch.FloatTensor] = None,751        # output options752        output_attentions: Optional[bool] = None,753        output_hidden_states: Optional[bool] = None,754        return_dict: Optional[bool] = None,755        # cache options756        use_cache: Optional[bool] = None,757        past_key_values: Optional[Cache | torch.FloatTensor] = None,758        cache_position: Optional[torch.LongTensor] = None,759        # generation options760        labels: Optional[torch.Tensor] = None,761        logits_to_keep: int | torch.Tensor = 0,762        **kwargs,763    ) -> tuple | CausalLMOutputWithPast:764        output_attentions = (765            output_attentions766            if output_attentions is not None767            else self.config.output_attentions768        )769        output_hidden_states = (770            output_hidden_states771            if output_hidden_states is not None772            else self.config.output_hidden_states773        )774        use_cache = use_cache if use_cache is not None else self.config.use_cache775        return_dict = (776            return_dict if return_dict is not None else self.config.return_dict777        )778 779        model_outputs = self.model(780            input_ids=input_ids,781            attention_mask=attention_mask,782            past_key_values=past_key_values,783            inputs_embeds=inputs_embeds,784            use_cache=use_cache,785            output_attentions=output_attentions,786            output_hidden_states=output_hidden_states,787            return_dict=return_dict,788            cache_position=cache_position,789            **kwargs,790        )791 792        hidden_states = model_outputs[0]793        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep794        logits = self.lm_head(hidden_states[:, slice_indices, :])795 796        loss = None797        if labels is not None:798            loss = self.loss_function(799                logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs800                )801 802        if not return_dict:803            output = (logits, model_outputs[1:])804            return (loss, output) if loss is not None else output805 806        return CausalLMOutputWithPast(807            loss=loss,808            logits=logits,809            past_key_values=model_outputs.past_key_values,810            hidden_states=model_outputs.hidden_states,811            attentions=model_outputs.attentions,812        )813 814 815__all__ = ["SmalLmForCausalLM", "SmalLmModel", "SmalLmPreTrainedModel"]816