Method load_state_dict

codegeex/megatron/model/language_model.py:406–445

Customized load.

(self, state_dict, strict=True)

        return state_dict_

    def load_state_dict(self, state_dict, strict=True):
        """Customized load."""

        # Position embedding.
        if self._top_query_embeddings_key in state_dict:
            state_dict_ = state_dict[self._top_query_embeddings_key]
        else:
            # for backward compatibility.
            state_dict_ = {}
            for key in state_dict.keys():
                if 'top_query_embeddings' in key:
                    state_dict_[key.split('top_query_embeddings.')[1]] \
                        = state_dict[key]
        pos_len = state_dict_['weight'].shape[0]
        max_seq_len = self.max_sequence_length // get_tensor_model_parallel_world_size()
        if pos_len < max_seq_len:
            print_rank_0(f"Top query embedding padded {pos_len} -> {max_seq_len}.")
            top_query_embeddings_padded = torch.nn.Embedding(
            max_seq_len - pos_len, self.hidden_size).half()
            self.init_method(top_query_embeddings_padded.weight)
            state_dict_['weight'] = torch.cat([state_dict_['weight'], top_query_embeddings_padded.weight], dim=0)
        self.top_query_embeddings.load_state_dict(state_dict_, strict=strict)

        # Tokentype embedding.
        if self.num_tokentypes > 0:
            state_dict_ = {}
            if self._tokentype_embeddings_key in state_dict:
                state_dict_ = state_dict[self._tokentype_embeddings_key]
            else:
                # for backward compatibility.
                for key in state_dict.keys():
                    if 'tokentype_embeddings' in key:
                        state_dict_[key.split('tokentype_embeddings.')[1]] \
                            = state_dict[key]
            if len(state_dict_.keys()) > 0:
                self.tokentype_embeddings.load_state_dict(state_dict_,
                                                          strict=strict)
            else:
                print('***WARNING*** expected tokentype embeddings in the '
                      'checkpoint but could not find it', flush=True)


class QueryEmbeddingPipe(QueryEmbedding):