vllm - 💡(How to fix) Fix Early Stopping for Ouro models [1 participants]

## ouro.py

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = self.embed_input_ids(input_ids)

        # self.early_exit_threshold = getattr(config, "early_exit_threshold", None)
        # self.use_early_exit = self.early_exit_threshold is not None and 0<=self.early_exit_threshold<=1

        gate_probs = None
        unfinished = torch.full(fill_value = True, 
                                size = (hidden_states.shape[0],), 
                                dtype=torch.bool, device=hidden_states.device)

        early_exist_hidden_state = hidden_states

        for current_ut in range(self.total_ut_steps):
            residual = None
            for layer in self.layers[self.start_layer : self.end_layer]:
                hidden_states, residual = layer(
                    positions, hidden_states, current_ut, residual
                )
            hidden_states, _ = self.norm(hidden_states, residual)

            early_exist_hidden_state[unfinished] = hidden_states[unfinished]

            if self.use_early_exit:
                ut_logits, _ = self.early_exit_gate(hidden_states)
                ut_probs = torch.sigmoid(ut_logits.squeeze(-1))                
                gate_probs = ut_probs if gate_probs is None else gate_probs + (1 - gate_probs) * ut_probs
                unfinished[(gate_probs>=self.early_exit_threshold)] = False
   
        return early_exist_hidden_state