LlamaWeight.cc 4.38 KB
Newer Older
Li Zhang's avatar
Li Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
/*
 * Copyright (c) OpenMMLab. All rights reserved.
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

Li Zhang's avatar
Li Zhang committed
18
// Modified from
lvhan028's avatar
lvhan028 committed
19
// https://github.com/NVIDIA/FasterTransformer/blob/main/src/turbomind/models/multi_gpu_gpt/ParallelGptWeight.cc
Li Zhang's avatar
Li Zhang committed
20

lvhan028's avatar
lvhan028 committed
21
#include "src/turbomind/models/llama/LlamaWeight.h"
Li Zhang's avatar
Li Zhang committed
22

lvhan028's avatar
lvhan028 committed
23
namespace turbomind {
Li Zhang's avatar
Li Zhang committed
24
25

template<typename T>
26
27
28
LlamaWeight<T>::LlamaWeight(size_t     head_num,
                            size_t     kv_head_num,
                            size_t     size_per_head,
Li Zhang's avatar
Li Zhang committed
29
30
31
                            size_t     inter_size,
                            size_t     vocab_size,
                            size_t     num_layer,
Li Zhang's avatar
Li Zhang committed
32
                            bool       attn_bias,
33
34
                            WeightType weight_type,
                            int        group_size,
Li Zhang's avatar
Li Zhang committed
35
                            size_t     tensor_para_size,
36
                            size_t     tensor_para_rank):
37
    hidden_units_(head_num * size_per_head),
Li Zhang's avatar
Li Zhang committed
38
39
    inter_size_(inter_size),
    vocab_size_(vocab_size),
40
    vocab_size_padded_(vocab_size),
Li Zhang's avatar
Li Zhang committed
41
42
43
    num_layer_(num_layer),
    weight_type_(weight_type),
    tensor_para_size_(tensor_para_size),
44
    tensor_para_rank_(tensor_para_rank)
Li Zhang's avatar
Li Zhang committed
45
{
46
47
48
49
    if (vocab_size_padded_ % tensor_para_size_ != 0) {
        vocab_size_padded_ = (vocab_size_padded_ + tensor_para_size_ - 1) / tensor_para_size_ * tensor_para_size_;
        TM_LOG_WARNING("pad vocab size from %d to %d", vocab_size_, vocab_size_padded_);
    }
Li Zhang's avatar
Li Zhang committed
50
51
    decoder_layer_weights.reserve(num_layer_);
    for (unsigned l = 0; l < num_layer_; ++l) {
52
53
54
55
56
        decoder_layer_weights.push_back(new LlamaDecoderLayerWeight<T>(head_num,
                                                                       kv_head_num,
                                                                       size_per_head,
                                                                       inter_size_,
                                                                       weight_type_,
57
                                                                       group_size,
58
59
60
                                                                       attn_bias,
                                                                       tensor_para_size_,
                                                                       tensor_para_rank_));
Li Zhang's avatar
Li Zhang committed
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
    }

    mallocWeights();
}

template<typename T>
LlamaWeight<T>::~LlamaWeight()
{
    cudaFree((void*)pre_decoder_embedding_table);
    cudaFree((void*)output_norm_weight);
    cudaFree((void*)post_decoder_embedding_kernel);

    pre_decoder_embedding_table   = nullptr;
    post_decoder_embedding_kernel = nullptr;
}

template<typename T>
void LlamaWeight<T>::mallocWeights()
{
80
    deviceMalloc((T**)&pre_decoder_embedding_table, vocab_size_padded_ * hidden_units_);
Li Zhang's avatar
Li Zhang committed
81
    deviceMalloc((T**)&output_norm_weight, hidden_units_);
82
    deviceMalloc((T**)&post_decoder_embedding_kernel, hidden_units_ * vocab_size_padded_);
Li Zhang's avatar
Li Zhang committed
83
84
85
86
87
88
89
90
91
}

template<typename T>
void LlamaWeight<T>::loadModel(std::string dir_path)
{
    FtCudaDataType model_file_type = FtCudaDataType::FP16;
    dir_path += '/';

    loadWeightFromBin((T*)pre_decoder_embedding_table,
92
                      {vocab_size_padded_ * hidden_units_},
Li Zhang's avatar
Li Zhang committed
93
94
95
96
97
                      dir_path + "tok_embeddings.weight",
                      model_file_type);

    loadWeightFromBin((T*)output_norm_weight, {hidden_units_}, dir_path + "norm.weight", model_file_type);

98
99
100
101
    loadWeightFromBin((T*)post_decoder_embedding_kernel,
                      {hidden_units_ * vocab_size_padded_},
                      dir_path + "output.weight",
                      model_file_type);
Li Zhang's avatar
Li Zhang committed
102
103
104
105
106
107
108
109
110

    for (unsigned layer = 0; layer < num_layer_; ++layer) {
        decoder_layer_weights[layer]->loadModel(dir_path + "layers." + std::to_string(layer), model_file_type);
    }
}

template struct LlamaWeight<float>;
template struct LlamaWeight<half>;

lvhan028's avatar
lvhan028 committed
111
}  // namespace turbomind