tencent/Sequential-Hidden-Decoding-8B-n8-Instruct

1"""Qwen3 with scaled sequence length via embedding replication.2 3Extends Qwen3Model/Qwen3ForCausalLM with scale_seq_times additional4embedding tables. During forward, the original token sequence of length L5is expanded to (1 + scale_seq_times) * L via interleaved multi-stream6embedding, then processed by the standard Qwen3 transformer body.7 8Architecture overview (n = 1 + scale_seq_times):9  - n Embedding tables: E_0 (original), E_1, ..., E_{n-1} (new)10  - Interleaved layout: [E_0(t1), E_1(t1), ..., E_0(t2), E_1(t2), ...]11  - RoPE positions: 0, 1, 2, ..., n*L - 1 (continuous)12  - Standard causal attention over all n*L positions13  - Contraction: only the last stream's hidden_state per token goes through14    lm_head (the stream with the richest context), matching v4dev behavior.15 16See: Scale_SeqLen_via_Embedding_Replication.md17"""18 19from typing import Optional, Tuple, Union20 21import torch22from torch import nn23from transformers import Qwen3ForCausalLM, Qwen3Model24from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast25from transformers.processing_utils import Unpack26from transformers.utils import TransformersKwargs, can_return_tuple27 28from .configuration_qwen3_scale_seq import Qwen3ScaleSeqConfig29 30 31class Qwen3ScaleSeqModel(Qwen3Model):32    """Qwen3Model extended with multi-stream embedding for sequence scaling."""33 34    config_class = Qwen3ScaleSeqConfig35 36    def __init__(self, config: Qwen3ScaleSeqConfig):37        super().__init__(config)38        self.scale_seq_times = getattr(config, "scale_seq_times", 0)39 40        if self.scale_seq_times > 0:41            self.scale_seq_embed_tokens_list = nn.ModuleList(42                [43                    nn.Embedding(44                        config.vocab_size,45                        config.hidden_size,46                        self.padding_idx,47                    )48                    for _ in range(self.scale_seq_times)49                ]50            )51 52        self.post_init()53 54    def _expand_scale_seq(55        self,56        input_ids: torch.LongTensor,57        hidden_states: torch.FloatTensor,58    ) -> torch.FloatTensor:59        """Expand hidden_states from (B, T, D) to (B, T * scale, D).60 61        Layout per original token i:62          [main_emb_i, scale_seq_1_emb_i, ..., scale_seq_N_emb_i]63 64        Args:65            input_ids: (batch, seq_len) original token ids.66            hidden_states: (batch, seq_len, hidden) main embedding output.67 68        Returns:69            Expanded tensor of shape (batch, seq_len * scale, hidden).70        """71        device = hidden_states.device72        B, T, D = hidden_states.shape73 74        # (B, T, D) -> (B, T, 1, D)75        parts = [hidden_states.unsqueeze(2)]76 77        for s in range(self.scale_seq_times):78            emb_module = self.scale_seq_embed_tokens_list[s]79            hs_s = emb_module(input_ids.to(emb_module.weight.device)).to(device)80            parts.append(hs_s.unsqueeze(2))  # (B, T, 1, D)81 82        # (B, T, scale, D) -> (B, T * scale, D)83        expanded = torch.cat(parts, dim=2)84        return expanded.reshape(B, T * (self.scale_seq_times + 1), D)85 86    def forward(87        self,88        input_ids: Optional[torch.LongTensor] = None,89        attention_mask: Optional[torch.Tensor] = None,90        position_ids: Optional[torch.LongTensor] = None,91        past_key_values=None,92        inputs_embeds: Optional[torch.FloatTensor] = None,93        use_cache: Optional[bool] = None,94        output_attentions: Optional[bool] = None,95        output_hidden_states: Optional[bool] = None,96        return_dict: Optional[bool] = None,97        cache_position: Optional[torch.LongTensor] = None,98        **kwargs,99    ) -> Union[Tuple, BaseModelOutputWithPast]:100        if (101            self.scale_seq_times > 0102            and input_ids is not None103            and inputs_embeds is None104        ):105            scale = self.scale_seq_times + 1106 107            # Compute main embedding, then expand with scale_seq streams108            inputs_embeds = self.embed_tokens(input_ids)109            inputs_embeds = self._expand_scale_seq(input_ids, inputs_embeds)110 111            B = inputs_embeds.shape[0]112            T_expanded = inputs_embeds.shape[1]113 114            # Recompute cache_position and position_ids in expanded space115            past_seen_tokens = (116                past_key_values.get_seq_length()117                if past_key_values is not None else 0118            )119            cache_position = torch.arange(120                past_seen_tokens, past_seen_tokens + T_expanded,121                device=inputs_embeds.device,122            )123            position_ids = cache_position.unsqueeze(0).expand(B, -1)124 125            # Expand attention_mask to match expanded sequence length126            if attention_mask is not None:127                attention_mask = attention_mask.repeat_interleave(scale, dim=1)128 129            input_ids = None  # avoid double embedding lookup in super().forward()130 131        return super().forward(132            input_ids=input_ids,133            attention_mask=attention_mask,134            position_ids=position_ids,135            past_key_values=past_key_values,136            inputs_embeds=inputs_embeds,137            use_cache=use_cache,138            output_attentions=output_attentions,139            output_hidden_states=output_hidden_states,140            return_dict=return_dict,141            cache_position=cache_position,142            **kwargs,143        )144 145 146class Qwen3ScaleSeqForCausalLM(Qwen3ForCausalLM):147    """Qwen3ForCausalLM with multi-stream embedding for sequence scaling.148 149    Contraction: after the transformer body produces (B, T*scale, D),150    select only the last stream per token (the one with richest context)151    before applying lm_head, producing (B, T, vocab_size).152    """153 154    config_class = Qwen3ScaleSeqConfig155    _tied_weights_keys = ["lm_head.weight"]156 157    def __init__(self, config: Qwen3ScaleSeqConfig):158        super().__init__(config)159        # Replace the inner model with our scaled version160        self.model = Qwen3ScaleSeqModel(config)161        self.post_init()162 163    @can_return_tuple164    def forward(165        self,166        input_ids: Optional[torch.LongTensor] = None,167        attention_mask: Optional[torch.Tensor] = None,168        position_ids: Optional[torch.LongTensor] = None,169        past_key_values=None,170        inputs_embeds: Optional[torch.FloatTensor] = None,171        labels: Optional[torch.LongTensor] = None,172        use_cache: Optional[bool] = None,173        output_attentions: Optional[bool] = None,174        output_hidden_states: Optional[bool] = None,175        return_dict: Optional[bool] = None,176        cache_position: Optional[torch.LongTensor] = None,177        logits_to_keep: Union[int, torch.Tensor] = 0,178        **kwargs,179    ) -> CausalLMOutputWithPast:180        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions181        output_hidden_states = (182            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states183        )184        return_dict = return_dict if return_dict is not None else self.config.use_return_dict185 186        outputs = self.model(187            input_ids=input_ids,188            attention_mask=attention_mask,189            position_ids=position_ids,190            past_key_values=past_key_values,191            inputs_embeds=inputs_embeds,192            use_cache=use_cache,193            output_attentions=output_attentions,194            output_hidden_states=output_hidden_states,195            return_dict=return_dict,196            cache_position=cache_position,197            **kwargs,198        )199 200        hidden_states = outputs[0]201 202        # ---- scale_seq contraction ----203        # Contract expanded hidden_states (B, T*scale, D) back to logical204        # token space (B, T, D) by selecting the last stream per token group205        # (the stream with the richest context), matching v4dev behavior.206        if self.model.scale_seq_times > 0:207            scale = self.model.scale_seq_times + 1208            hidden_states = hidden_states[:, scale - 1::scale, :]209 210        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep211        logits = self.lm_head(hidden_states[:, slice_indices, :])212 213        loss = None214        if labels is not None:215            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)216 217        return CausalLMOutputWithPast(218            loss=loss,219            logits=logits,220            past_key_values=outputs.past_key_values if use_cache else None,221            hidden_states=outputs.hidden_states,222            attentions=outputs.attentions,223        )224 225 226__all__ = ["Qwen3ScaleSeqModel", "Qwen3ScaleSeqForCausalLM"]227