Method __init__

TD_Pipe/model_executor/models/mixtral.py:305–330

(
        self,
        config: MixtralConfig,
        parallel_config: ParallelConfig,
        attention_backend: str,
        linear_method: Optional[LinearMethodBase] = None,
    )

class MixtralModel(nn.Module):

    def __init__(
        self,
        config: MixtralConfig,
        parallel_config: ParallelConfig,
        attention_backend: str,
        linear_method: Optional[LinearMethodBase] = None,
    ) -> None:
        super().__init__()
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.parallel_config = parallel_config
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
        self.layers = nn.ModuleList()
        for i in range(self.parallel_config.start,
                       self.parallel_config.end + 1):
            self.layers.add_module(
                f"{i}",
                MixtralDecoderLayer(config,
                                    attention_backend,
                                    linear_method=linear_method))

        if self.parallel_config.is_last:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,