MCPcopy Create free account
hub / github.com/evilsocket/cake / estimate_layer_size

Function estimate_layer_size

cake-core/src/cake/sharding/default.rs:225–307  ·  view source on GitHub ↗

Estimate average transformer layer size in bytes from safetensors files. For sharded models, reads each shard's header to compute exact per-tensor byte sizes, then sums only tensors matching `layer_prefix`. This excludes non-layer weights (visual encoder, MTP heads, embeddings, lm_head) which can be significant — e.g. Qwen3.5-27B-FP8 has ~6 GB of non-layer data.

(model_path: &Path, num_layers: usize, layer_prefix: &str)

Source from the content-addressed store, hash-verified

223/// non-layer weights (visual encoder, MTP heads, embeddings, lm_head) which
224/// can be significant — e.g. Qwen3.5-27B-FP8 has ~6 GB of non-layer data.
225pub fn estimate_layer_size(model_path: &Path, num_layers: usize, layer_prefix: &str) -> u64 {
226 if num_layers == 0 {
227 return 0;
228 }
229
230 let layer_dot = format!("{}.", layer_prefix);
231
232 // Try sharded model first
233 let index_path = model_path.join("model.safetensors.index.json");
234 if let Ok(data) = std::fs::read_to_string(&index_path) {
235 if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) {
236 if let Some(weight_map) = json.get("weight_map").and_then(|v| v.as_object()) {
237 let shards: HashSet<&str> =
238 weight_map.values().filter_map(|v| v.as_str()).collect();
239
240 // Try reading safetensors headers for exact tensor sizes
241 let mut layer_bytes: u64 = 0;
242 let mut total_bytes: u64 = 0;
243 let mut headers_ok = true;
244
245 for shard in &shards {
246 let shard_path = model_path.join(shard);
247 if let Some(tensors) = read_safetensors_tensor_sizes(&shard_path) {
248 for (name, size) in &tensors {
249 total_bytes += size;
250 if name.starts_with(&layer_dot) {
251 layer_bytes += size;
252 }
253 }
254 } else {
255 headers_ok = false;
256 break;
257 }
258 }
259
260 if headers_ok && layer_bytes > 0 {
261 let non_layer = total_bytes - layer_bytes;
262 if non_layer > 0 {
263 log::info!(
264 "model weights: {} total, {} layers, {} non-layer ({:.0}% excluded)",
265 human_bytes::human_bytes(total_bytes as f64),
266 human_bytes::human_bytes(layer_bytes as f64),
267 human_bytes::human_bytes(non_layer as f64),
268 non_layer as f64 / total_bytes as f64 * 100.0,
269 );
270 }
271 return layer_bytes / num_layers as u64;
272 }
273
274 // Fallback: raw file size division
275 let total: u64 = shards
276 .iter()
277 .filter_map(|s| std::fs::metadata(model_path.join(s)).ok())
278 .map(|m| m.len())
279 .sum();
280 return total / num_layers as u64;
281 }
282 }

Calls 2

getMethod · 0.45