infer_engine.cpp 4.42 KB
Newer Older
1
#include "infer_engine.hpp"
Ceng's avatar
Ceng committed
2
#include "spdlog/spdlog.h"
3
4
5
6
7
8
9

namespace infinilm::engine {

//------------------------------------------------------
// Constructor
//------------------------------------------------------
InferEngine::InferEngine(
Jiacheng Huang's avatar
Jiacheng Huang committed
10
    const InfinilmModel::Config &config,
11
    const distributed::DistConfig &distributed_config,
12
13
    infinicore::Device::Type device_type,
    const cache::CacheConfig &cache_config) // Changed parameter
14
    : communication_group_(distributed_config, device_type),
15
16
17
      model_config_(config),
      cache_config_(cache_config) {

18
    spdlog::info("Launch InferEngine with {}", std::string(distributed_config));
19
20
21
22
23
24
25
    spdlog::info("Cache configuration: type={}, layers={}, max_kv_cache_length={}",
                 static_cast<int>(cache_config_.type),
                 cache_config_.num_layers,
                 cache_config_.max_kv_cache_length);

    // Try to extract model configuration to override default cache parameters if needed
    try {
Jiacheng Huang's avatar
Jiacheng Huang committed
26
27
        if (const auto llama_config_ptr = dynamic_cast<const models::llama::LlamaConfig *>(&config)) {
            const auto &llama_config = *llama_config_ptr;
28
29
30
31
32
33
34
35
36
37
38

            cache_config_.num_layers = llama_config.num_hidden_layers;
            cache_config_.max_kv_cache_length = llama_config.max_position_embeddings;

            spdlog::info("Updated cache config from model: layers={}, max_kv_cache_length={}",
                         cache_config_.num_layers, cache_config_.max_kv_cache_length);
        }
    } catch (...) {
        spdlog::warn("Could not extract model config, using provided CacheConfig");
    }

39
40
41
42
    // Create one RankWorker per rank
    int world_size = communication_group_.get_world_size();
    workers_.reserve(world_size);
    for (int r = 0; r < world_size; ++r) {
43
44
45
46
        workers_.emplace_back(std::make_unique<RankWorker>(
            model_config_,
            communication_group_.get_rank_info(r),
            cache_config_));
47
48
49
50
51
52
53
54
55
56
57
58
    }
}

//------------------------------------------------------
// load_param
//------------------------------------------------------
void InferEngine::load_param(const std::string &name, const infinicore::Tensor &param) {
    // Load the parameter on all workers
    for (auto &worker : workers_) {
        worker->load_param(name, param);
    }
}
59

60
61
62
//------------------------------------------------------
// state_dict
//------------------------------------------------------
63
64
std::vector<std::unordered_map<std::string, infinicore::nn::Parameter>> InferEngine::state_dict() {
    std::vector<std::unordered_map<std::string, infinicore::nn::Parameter>> results;
65
66
67
    if (0 == workers_.size()) {
        throw std::runtime_error(" Model object not found. ");
    }
68
69
70
71
72

    for (auto &worker : workers_) {
        results.push_back(worker->state_dict());
    }
    return results;
73
74
}

75
//------------------------------------------------------
76
// forward
77
//------------------------------------------------------
78
79
80
InferEngine::Output InferEngine::forward(const InferEngine::Input &input) {
    const auto &[input_ids, position_ids] = input;

81
82
    // Trigger each worker to run inference
    for (auto &worker : workers_) {
83
        worker->run({input_ids, position_ids});
84
    }
PanZezhong's avatar
PanZezhong committed
85
86
87
88
    // Wait for all workers
    for (auto &worker : workers_) {
        worker->wait();
    }
89

90
    return {workers_[0]->get_output().logits};
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
}

//------------------------------------------------------
// Destructor
//------------------------------------------------------
InferEngine::~InferEngine() {
    // Close all workers
    for (auto &worker : workers_) {
        worker->close();
    }
}

const distributed::DistConfig &InferEngine::get_dist_config() const {
    return communication_group_.get_dist_config();
}

Ceng's avatar
Ceng committed
107
108
109
//------------------------------------------------------
// reset_cache
//------------------------------------------------------
110
111
112
113
114
115
116
117
118
119
120
121
122
123
void InferEngine::reset_cache(size_t pos) {
    for (auto &worker : workers_) {
        worker->reset_cache(pos);
    }
    for (auto &worker : workers_) {
        worker->wait();
    }
}

//------------------------------------------------------
// reset_cache (overloaded with CacheConfig)
//------------------------------------------------------
void InferEngine::reset_cache(const cache::CacheConfig &new_config, size_t pos) {
    cache_config_ = new_config;
Ceng's avatar
Ceng committed
124
    for (auto &worker : workers_) {
125
126
127
128
        worker->reset_cache(new_config, pos);
    }
    for (auto &worker : workers_) {
        worker->wait();
Ceng's avatar
Ceng committed
129
130
131
    }
}

132
} // namespace infinilm::engine