Skip to content

Commit a6155a8

Browse files
committed
Add compute_layer_statistics() function
1 parent be60469 commit a6155a8

File tree

1 file changed

+60
-4
lines changed

1 file changed

+60
-4
lines changed

tools/imatrix/imatrix.cpp

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -272,11 +272,11 @@ static void compute_tensor_statistics(std::vector<tensor_statistics> & tstats) {
272272
if (curr_avg.size() == prev_avg.size() && !curr_avg.empty()) {
273273
float dot_prod = 0.0f, vec1 = 0.0f, vec2 = 0.0f;
274274
for (size_t i = 0; i < curr_avg.size(); ++i) {
275-
dot_prod += curr_avg[i]*prev_avg[i];
276-
vec1 += curr_avg[i]*curr_avg[i];
277-
vec2 += prev_avg[i]*prev_avg[i];
275+
dot_prod += curr_avg[i] * prev_avg[i];
276+
vec1 += curr_avg[i] * curr_avg[i];
277+
vec2 += prev_avg[i] * prev_avg[i];
278278
}
279-
if (vec1 > 0 && vec2 > 0) ts.cossim = dot_prod / (std::sqrt(vec1)*std::sqrt(vec2));
279+
if (vec1 > 0 && vec2 > 0) ts.cossim = dot_prod / (std::sqrt(vec1) * std::sqrt(vec2));
280280
}
281281
}
282282
}
@@ -308,6 +308,62 @@ static void compute_tensor_statistics(std::vector<tensor_statistics> & tstats) {
308308
}
309309
}
310310

311+
static void compute_layer_statistics(const std::vector<tensor_statistics> & tstats,
312+
std::map<int, float> & layer_cossim,
313+
const std::unordered_map<std::string, Stats> & stats_map) {
314+
struct layer_aggregation {
315+
std::vector<float> curr_avg;
316+
std::vector<float> prev_avg;
317+
};
318+
319+
static const std::regex pattern(R"(blk\.(\d+)\.)");
320+
321+
// index tensor stats by name for quick lookup
322+
std::unordered_map<std::string, const tensor_statistics*> tidx;
323+
tidx.reserve(tstats.size());
324+
for (const auto & ts : tstats) tidx[ts.tensor] = &ts;
325+
326+
// concatenate per-layer
327+
std::map<int, layer_aggregation> taggr; // ordered by layer
328+
for (const auto & ts : tstats) {
329+
std::smatch match;
330+
if (!std::regex_search(ts.tensor, match, pattern)) continue;
331+
const int blk = std::stoi(match[1]);
332+
if (blk <= 0) continue;
333+
334+
std::string prev_lyr(ts.tensor);
335+
prev_lyr.replace(match.position(1), match.length(1), std::to_string(blk-1));
336+
337+
if (auto it_prev = tidx.find(prev_lyr); it_prev == tidx.end()) continue;
338+
339+
// use stored Stats to rebuild averages
340+
const auto curr_avg = compute_tensor_averages(stats_map.at(ts.tensor));
341+
const auto prev_avg = compute_tensor_averages(stats_map.at(prev_lyr));
342+
if (curr_avg.empty() || prev_avg.empty() || curr_avg.size() != prev_avg.size()) continue;
343+
344+
auto & [curr, prev] = taggr[blk];
345+
curr.insert(curr.end(), curr_avg.begin(), curr_avg.end());
346+
prev.insert(prev.end(), prev_avg.begin(), prev_avg.end());
347+
}
348+
349+
// compute cosine per layer
350+
for (auto & kv : taggr) {
351+
const auto & curr = kv.second.curr_avg;
352+
const auto & prev = kv.second.prev_avg;
353+
if (curr.size() != prev.size() || curr.empty()) continue;
354+
float dot_prod = 0.0, lyr1 = 0.0, lyr2 = 0.0;
355+
for (size_t i = 0; i < curr.size(); ++i) {
356+
const double a = curr[i], b = prev[i];
357+
dot_prod += a*b;
358+
lyr1 += a*a;
359+
lyr2 += b*b;
360+
}
361+
float cossim = 0.0f;
362+
if (lyr1 > 0.0 && lyr2 > 0.0) cossim = dot_prod / (std::sqrt(lyr1) * std::sqrt(lyr2));
363+
layer_cossim[kv.first] = cossim;
364+
}
365+
}
366+
311367
bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * user_data) {
312368
GGML_UNUSED(user_data);
313369

0 commit comments

Comments
 (0)