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

namespace infinilm::engine {

//------------------------------------------------------
// Constructor
//------------------------------------------------------
9
10
11
12
13
14
15
16
17
18
19
20
/**
 * @deprecated This function is deprecated and will be REMOVED in the next major release (v0.2.0).
 *
 * ⚠️ DEVELOPMENT POLICY:
 *   - NO new development or feature additions permitted on this interface
 *   - Only critical bug fixes (security/stability) allowed until removal
 *   - All new code MUST migrate to the polymorphic overload below
 *
 * Replacement: Use the polymorphic overload of this same function name with updated signature
 * Reason: Legacy signature lacks support for dynamic quantization modes.
 * Removal target: v0.2.0 (Q2 2026)
 */
21
InferEngine::InferEngine(
Jiacheng Huang's avatar
Jiacheng Huang committed
22
    const InfinilmModel::Config &config,
23
    const distributed::DistConfig &distributed_config,
24
    infinicore::Device::Type device_type,
25
26
    const cache::CacheConfig *cache_config,
    bool enable_graph_compiling) // Changed parameter
27
    : communication_group_(distributed_config, device_type),
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
      legacy_model_config_(config) {
    if (cache_config != nullptr) {
        cache_config_ = cache_config->unique_copy();
    }
    // Create one RankWorker per rank
    int world_size = communication_group_.get_world_size();
    barrier_ = std::make_unique<RankBarrier>((size_t)world_size);
    workers_.reserve(world_size);
    for (int r = 0; r < world_size; ++r) {
        workers_.emplace_back(std::make_unique<RankWorker>(
            legacy_model_config_,
            communication_group_.get_rank_info(r),
            cache_config_ != nullptr ? cache_config_.get() : nullptr,
            barrier_.get(),
            enable_graph_compiling));
    }

    // Compile the model on all workers
    this->compile();
}
48

49
50
51
52
53
54
55
InferEngine::InferEngine(
    const std::string &model_path,
    const distributed::DistConfig &distributed_config,
    infinicore::Device::Type device_type,
    const cache::CacheConfig *cache_config,
    bool enable_graph_compiling) // Changed parameter
    : communication_group_(distributed_config, device_type) {
PanZezhong's avatar
PanZezhong committed
56
57
    if (cache_config != nullptr) {
        cache_config_ = cache_config->unique_copy();
58
    }
59
60
61

    // Load model config if model_path is provided, model_path must be valid, and config.json exists
    this->model_config_ = std::make_shared<infinilm::config::ModelConfig>(model_path + "/config.json");
62
63
    // Create one RankWorker per rank
    int world_size = communication_group_.get_world_size();
64
    barrier_ = std::make_unique<RankBarrier>((size_t)world_size);
65
66
    workers_.reserve(world_size);
    for (int r = 0; r < world_size; ++r) {
67
68
69
        workers_.emplace_back(std::make_unique<RankWorker>(
            model_config_,
            communication_group_.get_rank_info(r),
70
            cache_config_ != nullptr ? cache_config_.get() : nullptr,
71
            barrier_.get(),
72
            enable_graph_compiling));
73
    }
74
75
    // Compile the model on all workers
    this->compile();
76
77
78
79
80
81
82
83
84
85
86
}

//------------------------------------------------------
// 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);
    }
}
87

88
89
90
//------------------------------------------------------
// state_dict
//------------------------------------------------------
91
92
std::vector<std::unordered_map<std::string, infinicore::nn::Parameter>> InferEngine::state_dict() {
    std::vector<std::unordered_map<std::string, infinicore::nn::Parameter>> results;
93
94
95
    if (0 == workers_.size()) {
        throw std::runtime_error(" Model object not found. ");
    }
96
97
98
99
100

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

103
//------------------------------------------------------
104
// forward
105
//------------------------------------------------------
106
107
infinilm::InfinilmModel::Input
InferEngine::Input::to_model_input(infinicore::Device device) const {
108

109
110
111
112
    auto to_device = [&](const std::optional<infinicore::Tensor> &t)
        -> std::optional<infinicore::Tensor> {
        return t.has_value() ? t.value()->to(device) : t;
    };
113
114

    return {
115
        to_device(input_ids), // @todo: on device in the future
116
        to_device(position_ids),
117
        to_device(past_sequence_lengths), // @todo: on device in the future
118
119
120
121
122
        to_device(total_sequence_lengths),
        to_device(input_offsets),
        to_device(block_tables),
        to_device(slot_mapping),
    };
PanZezhong's avatar
PanZezhong committed
123
}
124

PanZezhong's avatar
PanZezhong committed
125
InferEngine::Output InferEngine::forward(const InferEngine::Input &input) {
126
127
    // Trigger each worker to run inference
    for (auto &worker : workers_) {
128
        worker->run(input);
129
    }
PanZezhong's avatar
PanZezhong committed
130
131
132
133
    // Wait for all workers
    for (auto &worker : workers_) {
        worker->wait();
    }
134

135
    return workers_[0]->get_output();
136
137
}

138
139
140
141
142
143
144
145
146
147
void InferEngine::compile() {
    for (auto &worker : workers_) {
        worker->compile();
    }
    // Wait for all workers
    for (auto &worker : workers_) {
        worker->wait();
    }
}

148
149
150
151
152
153
154
155
156
157
158
159
160
161
//------------------------------------------------------
// 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();
}

162
163
164
//------------------------------------------------------
// reset_cache (overloaded with CacheConfig)
//------------------------------------------------------
PanZezhong's avatar
PanZezhong committed
165
void InferEngine::reset_cache(const cache::CacheConfig *new_config) {
Ceng's avatar
Ceng committed
166
    for (auto &worker : workers_) {
PanZezhong's avatar
PanZezhong committed
167
        worker->reset_cache(new_config);
168
169
170
    }
    for (auto &worker : workers_) {
        worker->wait();
Ceng's avatar
Ceng committed
171
    }
172
173

    this->compile();
Ceng's avatar
Ceng committed
174
175
}

176
} // namespace infinilm::engine