diff --git a/exllamav3/architecture/glm4v.py b/exllamav3/architecture/glm4v.py index 70226972..1e328500 100644 --- a/exllamav3/architecture/glm4v.py +++ b/exllamav3/architecture/glm4v.py @@ -212,7 +212,10 @@ def __init__( rope_style = RopeStyle.NEOX, ), key_fused_qkv = "qkv", - key_o = "proj", + key_q="q_proj", + key_k="k_proj", + key_v="v_proj", + key_o = "o_proj", qmap = "block.attn", ), mlp_norm = RMSNorm( diff --git a/exllamav3/architecture/qwen2_5_vl.py b/exllamav3/architecture/qwen2_5_vl.py index b8466eef..16a90835 100644 --- a/exllamav3/architecture/qwen2_5_vl.py +++ b/exllamav3/architecture/qwen2_5_vl.py @@ -224,8 +224,11 @@ def __init__( head_dim = v.head_dim, rope_style = RopeStyle.NEOX, ), - key_fused_qkv = "qkv", - key_o = "proj", + key_fused_qkv = "qkv_proj", + key_q="q_proj", + key_k="k_proj", + key_v="v_proj", + key_o = "o_proj", qmap = "block.attn", use_cu_seqlens = bool(idx not in v.fullatt_block_indexes) ), diff --git a/exllamav3/architecture/qwen3_vl.py b/exllamav3/architecture/qwen3_vl.py index 24e71111..26cbaa20 100644 --- a/exllamav3/architecture/qwen3_vl.py +++ b/exllamav3/architecture/qwen3_vl.py @@ -211,8 +211,11 @@ def __init__( head_dim = v.head_dim, rope_style = RopeStyle.NEOX, ), - key_fused_qkv = "qkv", - key_o = "proj", + key_fused_qkv = "qkv_proj", + key_q="q_proj", + key_k="k_proj", + key_v="v_proj", + key_o = "o_proj", qmap = "block.attn", ), mlp_norm = LayerNorm( diff --git a/exllamav3/model/model.py b/exllamav3/model/model.py index 13fe16dd..da594cfe 100644 --- a/exllamav3/model/model.py +++ b/exllamav3/model/model.py @@ -441,6 +441,74 @@ def get_tensor_size(tensors): vram_bits = head_numel * head_bpw + sum_bits return sum_bits / sum_numel, head_bpw, vram_bits + def get_detailed_weights_info(self): + """ + Get detailed weights information per layer type. + This focuses only on layers with trainable weights. + + Returns: + dict: {layer_name: {"count": int, "size_bytes": int, "numel": int, "bits": float}} + + layer_name has numeric indices replaced with '*' for aggregation. + Example: "model.language_model.layers.0.mlp.experts.0.up_proj" + -> "model.language_model.layers.*.mlp.experts.*.up_proj" + """ + from ..modules import Linear, Embedding, Conv + + # Trainable leaf modules - modules that always have trainable weights + # Exclude normalization layers which may be unweighted (RMSNorm, LayerNorm, etc.) + TRAINABLE_LEAF_TYPES = (Linear, Embedding, Conv) + + def normalize_key(key): + """Replace numeric indices with '*' for aggregation.""" + parts = key.split(".") + return ".".join("*" if p.isdigit() else p for p in parts) + + def get_tensor_size_bytes(tensors): + return sum(t.element_size() * t.numel() for t in tensors.values()) + + stats = {} + + for module in self: + if not isinstance(module, TRAINABLE_LEAF_TYPES): + continue + + layer_name = normalize_key(module.key) + numel = module.weights_numel() + + # Throw exception if module type has no trainable weights + # and was not filtered out via TRAINABLE_LEAF_TYPES + # This only affects new modules that we forgot to config + if numel is None or numel == 0: + raise ValueError( + f"Layer type '{module.__class__.__name__}' at '{module.key}' has no trainable weights" + ) + + if module.device is not None: + size_bytes = get_tensor_size_bytes(module.get_tensors()) + else: + # Metadata only + if isinstance(module, Linear): + size_bytes = module.storage_size() + else: + size_bytes = sum(self.config.stc.get_tensor_sizes(module.key)) + + if layer_name not in stats: + stats[layer_name] = { + "count": 0, + "size_bytes": 0, + "numel": 0, + "bits": 0.0, + "leaf_module": module.__class__.__name__, + } + + stats[layer_name]["count"] += 1 + stats[layer_name]["size_bytes"] += size_bytes + stats[layer_name]["numel"] += numel + stats[layer_name]["bits"] += 8 * size_bytes + + return stats + def get_name(self): return self.__class__.__name__ diff --git a/exllamav3/modules/conv.py b/exllamav3/modules/conv.py index b3a459fb..7fd6f83b 100644 --- a/exllamav3/modules/conv.py +++ b/exllamav3/modules/conv.py @@ -71,8 +71,17 @@ def get_tensors(self): @override def weights_numel(self): - return self._numel - + if self._numel is not None: + return self._numel + + # Weight is not loaded, derive from metadata + # Conv weight shape: [out_channels, in_channels, kernel_size] + # It doesn't seem like there is a need for strided convolution + weight_numel = self.out_channels * self.in_channels + for k in self.kernel_size: + weight_numel *= k + # TODO: We do not count bias parameters, should we?. + return weight_numel @override def forward( diff --git a/util/size_estimation.py b/util/size_estimation.py index 2b9bab1f..dd449481 100644 --- a/util/size_estimation.py +++ b/util/size_estimation.py @@ -4,25 +4,52 @@ import argparse from exllamav3.loader.safetensors import SafetensorsCollection, VariantSafetensorsCollection import yaml +from tabulate import tabulate +def print_markdown_table(stats): + """Print layer statistics as a pretty-aligned Markdown table using tabulate.""" + total_bytes = sum(s["size_bytes"] for s in stats.values()) -def tsize(t): - return t.nelement() * t.element_size() + # Sort by size descending + sorted_layers = sorted(stats.items(), key=lambda x: -x[1]["size_bytes"]) + # Build table rows with raw numeric values + rows = [] + for layer_name, stat in sorted_layers: + size_mib = stat["size_bytes"] / 1024**2 + size_pct = 100 * stat["size_bytes"] / total_bytes if total_bytes > 0 else 0 + bpw = stat["bits"] / stat["numel"] -def dsize(d): - size = 0 - for _, v in d.items(): size += tsize(v) - return size + rows.append([layer_name, stat["count"], size_mib, size_pct, bpw]) + + # Print table with github format + print() + print( + tabulate( + rows, + headers=["Layer name", "Number", "Size (MiB)", "Size (%)", "Effective BPW"], + tablefmt="github", + stralign="left", + numalign="right", + floatfmt=".2f", + intfmt=",", + ) + ) + + # Print summary + total_bits = sum(s["bits"] for s in stats.values()) + total_numel = sum(s["numel"] for s in stats.values()) + avg_bpw = round(total_bits / total_numel, 2) + + print() + print(f" -- Average bitrate: {avg_bpw:.2f} bpw") + print(f" -- Size: {total_bytes / 1024**2:,.2f} MiB") + print() def main(args): - # Config/model config = Config.from_directory(args.in_dir) - model = Model.from_config(config) - - # Tensor collection stc = SafetensorsCollection(args.in_dir) # Override tensors @@ -45,12 +72,28 @@ def main(args): vstc.add_stc(o_keys, SafetensorsCollection(o_dir)) config.stc = vstc - # New bpw etc. - bpw_layer, bpw_head, vram_bits = model.get_storage_info() - bpw_layer = round(bpw_layer, 2) - bpw_head = round(bpw_head) - print(f" -- New estimated model bitrate: {bpw_layer:.2f} bpw / {bpw_head:.2f} bpw (head)") - print(f" -- VRAM: {vram_bits / 8 / 1024**3:.0f} GiB") + # Iterate over all model components (text, vision, etc.) + all_stats = {} + + for component in config.model_classes: + model = Model.from_config(config, component=component) + + # Aggregate detailed stats + component_stats = model.get_detailed_weights_info() + for layer_name, stat in component_stats.items(): + if layer_name not in all_stats: + all_stats[layer_name] = { + "count": 0, + "size_bytes": 0, + "numel": 0, + "bits": 0.0, + } + all_stats[layer_name]["count"] += stat["count"] + all_stats[layer_name]["size_bytes"] += stat["size_bytes"] + all_stats[layer_name]["numel"] += stat["numel"] + all_stats[layer_name]["bits"] += stat["bits"] + + print_markdown_table(all_stats) if __name__ == "__main__":