The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of cross-attention is added between the self-attention layers, following the architecture described in [Attention is all you need](https://arxiv.org/abs/1706.03762) by Ashish V
| 781 | CHATGLM_6B_START_DOCSTRING, |
| 782 | ) |
| 783 | class ChatGLMModel(ChatGLMPreTrainedModel): |
| 784 | """ |
| 785 | |
| 786 | The model can behave as an encoder (with only self-attention) as well |
| 787 | as a decoder, in which case a layer of cross-attention is added between |
| 788 | the self-attention layers, following the architecture described in [Attention is |
| 789 | all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, |
| 790 | Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. |
| 791 | |
| 792 | To behave as an decoder the model needs to be initialized with the |
| 793 | `is_decoder` argument of the configuration set to `True`. |
| 794 | To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` |
| 795 | argument and `add_cross_attention` set to `True`; an |
| 796 | `encoder_hidden_states` is then expected as an input to the forward pass. |
| 797 | """ |
| 798 | |
| 799 | def __init__(self, config: ChatGLMConfig): |
| 800 | super().__init__(config) |
| 801 | |
| 802 | # recording parameters |
| 803 | self.max_sequence_length = config.max_sequence_length |
| 804 | self.hidden_size = config.hidden_size |
| 805 | self.params_dtype = torch.half |
| 806 | self.num_attention_heads = config.num_attention_heads |
| 807 | self.vocab_size = config.vocab_size |
| 808 | self.num_layers = config.num_layers |
| 809 | self.layernorm_epsilon = config.layernorm_epsilon |
| 810 | self.inner_hidden_size = config.inner_hidden_size |
| 811 | self.hidden_size_per_attention_head = self.hidden_size // self.num_attention_heads |
| 812 | self.position_encoding_2d = config.position_encoding_2d |
| 813 | self.model_parallel = True |
| 814 | |
| 815 | self.word_embeddings = skip_init( |
| 816 | torch.nn.Embedding, |
| 817 | num_embeddings=self.vocab_size, embedding_dim=self.hidden_size, |
| 818 | dtype=self.params_dtype |
| 819 | ) |
| 820 | |
| 821 | def get_layer(layer_id): |
| 822 | return GLMBlock( |
| 823 | self.hidden_size, |
| 824 | self.num_attention_heads, |
| 825 | self.layernorm_epsilon, |
| 826 | layer_id, |
| 827 | inner_hidden_size=self.inner_hidden_size, |
| 828 | hidden_size_per_attention_head=self.hidden_size_per_attention_head, |
| 829 | layernorm=LayerNorm, |
| 830 | use_bias=True, |
| 831 | params_dtype=self.params_dtype, |
| 832 | position_encoding_2d=self.position_encoding_2d, |
| 833 | ) |
| 834 | |
| 835 | self.layers = torch.nn.ModuleList( |
| 836 | [get_layer(layer_id) for layer_id in range(self.num_layers)] |
| 837 | ) |
| 838 | |
| 839 | # Final layer norm before output. |
| 840 | self.final_layernorm = LayerNorm(self.hidden_size, eps=self.layernorm_epsilon) |