| 90 | |
| 91 | |
| 92 | class CohereModel(Module): |
| 93 | |
| 94 | def __init__(self, config: CohereConfig) -> None: |
| 95 | super().__init__() |
| 96 | |
| 97 | self.mapping = config.mapping |
| 98 | if self.mapping.is_first_pp_rank(): |
| 99 | self.vocab_embedding = Embedding(config.vocab_size, |
| 100 | config.hidden_size, |
| 101 | dtype=config.dtype, |
| 102 | tp_group=config.mapping.tp_group, |
| 103 | tp_size=config.mapping.tp_size, |
| 104 | tp_rank=config.mapping.tp_rank) |
| 105 | |
| 106 | self.layers = DecoderLayerList(CohereDecoderLayer, config) |
| 107 | |
| 108 | if self.mapping.is_last_pp_rank(): |
| 109 | self.ln_f = LayerNorm(normalized_shape=config.hidden_size, |
| 110 | eps=config.norm_epsilon, |
| 111 | bias=False, |
| 112 | dtype=config.dtype) |
| 113 | |
| 114 | def forward( |
| 115 | self, |
| 116 | input_ids=None, |
| 117 | position_ids=None, |
| 118 | use_cache=False, |
| 119 | attention_mask=None, |
| 120 | spec_decoding_params=None, |
| 121 | kv_cache_params=None, |
| 122 | attention_params=None, |
| 123 | hidden_states=None, |
| 124 | ): |
| 125 | if self.mapping.is_first_pp_rank(): |
| 126 | hidden_states = self.vocab_embedding(input_ids) |
| 127 | else: |
| 128 | hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) |
| 129 | |
| 130 | hidden_states = self.layers.forward( |
| 131 | hidden_states, |
| 132 | use_cache=use_cache, |
| 133 | attention_mask=attention_mask, |
| 134 | kv_cache_params=kv_cache_params, |
| 135 | attention_params=attention_params, |
| 136 | spec_decoding_params=spec_decoding_params) |
| 137 | |
| 138 | if use_cache: |
| 139 | hidden_states, presents = hidden_states |
| 140 | |
| 141 | if self.mapping.is_last_pp_rank(): |
| 142 | hidden_states = self.ln_f(hidden_states) |
| 143 | else: |
| 144 | hidden_states = send(hidden_states, self.mapping.next_pp_rank()) |
| 145 | |
| 146 | if use_cache: |
| 147 | return (hidden_states, presents) |
| 148 | return hidden_states |
| 149 | |