mirror of
https://github.com/AlexsJones/llmfit.git
synced 2026-10-21 21:54:55 +00:00
fix(scraper): never size hybrid SSM models by the estimator; cap by declared size (#1063)
* fix(scraper): never size hybrid SSM models by the estimator; cap by declared size #1057 gated the packed-count correction on the repo being quantized, to stop the estimator overriding the exact 31.6B count of the BF16 Nemotron-3-Nano with 101.6B. The first weekly run on that code (#1062, not merged) showed the gate was wrong in both directions: - Quantized Nemotron-H repos still went through the estimator, which prices every hybrid Mamba layer as a MoE transformer layer: nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 (1.2M downloads) 17.8B -> 101.6B, and 23 siblings. - Unquantized repos lost the rescue the old rule gave a placeholder safetensors total: ornith-ai/Ornith-1.0-35B (2M downloads) 34.1B -> 0.0B, NousResearch/Hermes-4-14B -> 0.0B, unsloth/Qwen3-235B-A22B-GGUF 235B -> 16.4B. The rule is now about what the estimator can size, not about quantization: - Hybrid SSM architectures (nemotron_h, zamba*, jamba, bamba, mamba*, falcon_h1, granitemoehybrid, lfm2*, ..., or any config carrying hybrid_override_pattern / layers_block_type) are never estimated. An understated count falls back to the size the repo name declares. - Everything else keeps the original 'estimate > 2x reported' rescue, capped at the name-declared size when the estimate exceeds it by 1.5x. - A speculative-decoding draft head keeps its own count: it is named after its target model but is a fraction of its size. Of the 52 entries that moved in #1062 and declare a size in their name, the number within 0.6-1.5x of it: main 32, #1062 10, this change 52. * fix(scraper): match hybrid SSM architectures by family Review follow-up on #1063: exact model_type matching missed variants the catalog already holds (nemotron_h_puzzle, hybrid_mamba_attn, rwkv7_native, lfm2_vl, qwen3_mamba3, ...). Over-matching is the safe direction: it keeps the reported count or the declared size instead of an estimate. --------- Co-authored-by: Alex Jones <axjns@example.com>
This commit is contained in:
Alex Jones
GitHub
parent
44de3655e0
commit
2dd06281bc
2 changed files with 168 additions and 25 deletions
|
|
@ -1446,23 +1446,92 @@ def is_prequantized_repo(repo_id: str, config: dict | None) -> bool:
|
|||
return _name_says_prequantized(repo_id)
|
||||
|
||||
|
||||
# Architectures estimate_params_from_arch cannot size: it prices every layer
|
||||
# as attention + (MoE) MLP, but these interleave Mamba/SSM or linear-attention
|
||||
# layers. It reported 101.6B for the 31.6B Nemotron-3-Nano and 17.0B for the
|
||||
# 8.1B Nemotron-H-8B.
|
||||
_HYBRID_SSM_MODEL_TYPE = re.compile(
|
||||
r"mamba|rwkv|zamba|jamba|bamba|hyena|nemotron_h|falcon_h1|granitemoehybrid"
|
||||
r"|lfm2|plamo2|recurrent_gemma|hybrid|(?:^|[_-])ssm(?:$|[_-])"
|
||||
)
|
||||
|
||||
# A speculative-decoding draft head is named after its target model but is a
|
||||
# fraction of its size; the safetensors count is right and the name is not.
|
||||
_DRAFT_HEAD_NAME = re.compile(
|
||||
r"(?<![a-z])(dflash|dspark|eagle\d?|medusa|speculator|draft)(?![a-z])", re.IGNORECASE)
|
||||
|
||||
_DECLARED_SIZE = re.compile(r"(?<![A-Za-z0-9.])(\d+(?:\.\d+)?)([BT])(?![A-Za-z0-9])")
|
||||
|
||||
|
||||
def name_declared_params(repo_id: str) -> int | None:
|
||||
"""Total parameters the repo name declares ("Qwen3-235B-A22B" -> 235e9).
|
||||
|
||||
None when the name declares nothing usable: no size token, an "NxMB"
|
||||
expert product (Mixtral-8x7B is 46.7B, not 7B or 56B), or a draft head.
|
||||
Active counts ("A22B") and effective sizes ("E4B") are not matched.
|
||||
"""
|
||||
base = repo_id.split("/")[-1]
|
||||
if _DRAFT_HEAD_NAME.search(base) or re.search(r"\d+x\d", base, re.IGNORECASE):
|
||||
return None
|
||||
# "1.58B" is BitNet's bit width, not a size.
|
||||
found = [(float(v), u) for v, u in _DECLARED_SIZE.findall(base.upper()) if v != "1.58"]
|
||||
billions = [v * 1e9 for v, u in found if u == "B"]
|
||||
# Next to a B size, a T figure is training tokens (bitnet-b1.58-2B-4T).
|
||||
# On its own it is the size (Qwen3.8-2.4T-A95B).
|
||||
sizes = billions or [v * 1e12 for v, u in found if u == "T"]
|
||||
return int(max(sizes)) if sizes else None
|
||||
|
||||
|
||||
def _is_hybrid_ssm(config: dict | None) -> bool:
|
||||
cfg = config or {}
|
||||
for src in (cfg, cfg.get("text_config") or {}):
|
||||
# Families, not exact names: the catalog already holds variants such
|
||||
# as nemotron_h_puzzle, hybrid_mamba_attn, rwkv7_native and lfm2_vl.
|
||||
# Over-matching is the safe direction, since it only means keeping the
|
||||
# reported count or the declared size instead of an estimate.
|
||||
if _HYBRID_SSM_MODEL_TYPE.search(str(src.get("model_type", "")).lower()):
|
||||
return True
|
||||
if "hybrid_override_pattern" in src or "layers_block_type" in src:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def correct_packed_param_count(repo_id: str, total_params: int,
|
||||
config: dict | None) -> int:
|
||||
"""Replace a packed element count with the architecture estimate.
|
||||
"""Replace an understated safetensors count with a better figure.
|
||||
|
||||
safetensors.total counts tensor elements, so int4 weights packed into
|
||||
int32 report ~1/6 of the real parameters (#1045). That only happens to
|
||||
quantized repos: for a full-precision checkpoint the count is exact, and
|
||||
the estimate must not override it. The estimate over-counts hybrid models
|
||||
(it prices every Nemotron-H layer as a MoE transformer layer, 101.6B for
|
||||
a 31.6B BF16 checkpoint), which the 2x margin alone does not catch.
|
||||
int32 report ~1/6 of the real parameters (#1045), and some repos publish
|
||||
a placeholder (ornith-ai/Ornith-1.0-35B reports ~0). The architecture
|
||||
estimate rescues both when it is more than twice the reported count.
|
||||
|
||||
Two guards, because the estimate is not always right:
|
||||
- hybrid SSM architectures are never estimated (see
|
||||
_HYBRID_SSM_MODEL_TYPE); an understated count falls back to the size
|
||||
the name declares instead.
|
||||
- any replacement more than 1.5x the name-declared size is capped to it.
|
||||
A draft head keeps its own count: its name describes another model.
|
||||
"""
|
||||
if not is_prequantized_repo(repo_id, config):
|
||||
base = repo_id.split("/")[-1]
|
||||
if _DRAFT_HEAD_NAME.search(base):
|
||||
return total_params
|
||||
declared = name_declared_params(repo_id)
|
||||
|
||||
if _is_hybrid_ssm(config):
|
||||
# A quantized repo's count is packed, so it only ever understates:
|
||||
# NVFP4 stores two weights per element (17.8B for a 30B model), which
|
||||
# the 2x margin misses. FP8 counts are exact and stay within 1.25x.
|
||||
margin = 1.25 if is_prequantized_repo(repo_id, config) else 2.0
|
||||
if declared and declared > total_params * margin:
|
||||
return declared
|
||||
return total_params
|
||||
|
||||
arch_params = estimate_params_from_arch(config)
|
||||
if arch_params and arch_params > total_params * 2:
|
||||
return arch_params
|
||||
return total_params
|
||||
if not arch_params or arch_params <= total_params * 2:
|
||||
return total_params
|
||||
if declared and arch_params > declared * 1.5:
|
||||
return declared
|
||||
return arch_params
|
||||
|
||||
|
||||
def revalidation_priority(model: dict) -> int | None:
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from scrape_hf_models import ( # noqa: E402
|
|||
REVALIDATION_COOLDOWN_DAYS,
|
||||
correct_packed_param_count,
|
||||
is_prequantized_repo,
|
||||
name_declared_params,
|
||||
RATE_LIMIT_MAX_RETRIES,
|
||||
RATE_LIMIT_STATS,
|
||||
detect_moe,
|
||||
|
|
@ -446,29 +447,99 @@ _QWEN38_27B_TEXT = {
|
|||
}
|
||||
|
||||
|
||||
def test_packed_count_is_corrected_only_for_quantized_repos():
|
||||
_NEMOTRON_H_30B = {
|
||||
"model_type": "nemotron_h", "hidden_size": 2688, "num_hidden_layers": 52,
|
||||
"vocab_size": 131072, "num_attention_heads": 32, "num_key_value_heads": 2,
|
||||
"head_dim": 128, "n_routed_experts": 128, "num_experts_per_tok": 6,
|
||||
"moe_intermediate_size": 1856,
|
||||
}
|
||||
|
||||
|
||||
def test_packed_count_is_rescued_by_the_architecture_estimate():
|
||||
# #1045: int4 packed into int32 reports 7.8B for a 27B-class model.
|
||||
quantized = {"text_config": _QWEN38_27B_TEXT,
|
||||
"quantization_config": {"quant_method": "compressed-tensors"}}
|
||||
fixed = correct_packed_param_count(
|
||||
"TelperionAI/Qwen3.8-27B-INT4-AWQ-GPTQ", 7_839_289_360, quantized)
|
||||
assert 20e9 < fixed < 32e9, fixed
|
||||
# The name alone is enough when config.json does not declare it.
|
||||
# Not only quantized repos: ornith-ai/Ornith-1.0-35B publishes a
|
||||
# placeholder total, and gating the rescue on quantization zeroed it.
|
||||
fixed = correct_packed_param_count(
|
||||
"RedHatAI/Qwen3-32B-quantized.w4a16", 7_839_289_360,
|
||||
{"text_config": _QWEN38_27B_TEXT})
|
||||
assert fixed > 20e9
|
||||
|
||||
# A full-precision checkpoint's count is exact. The estimator prices every
|
||||
# Nemotron-H layer as a MoE transformer layer and said 101.6B for this one.
|
||||
nemotron_h = {"hidden_size": 2688, "num_hidden_layers": 52, "vocab_size": 131072,
|
||||
"num_attention_heads": 32, "num_key_value_heads": 2, "head_dim": 128,
|
||||
"n_routed_experts": 128, "num_experts_per_tok": 6,
|
||||
"moe_intermediate_size": 1856}
|
||||
assert estimate_params_from_arch(nemotron_h) > 2 * 31_577_937_344
|
||||
"org/Placeholder-27B", 1_000, {"text_config": _QWEN38_27B_TEXT})
|
||||
assert 20e9 < fixed < 32e9, fixed
|
||||
# A sound count is left alone.
|
||||
assert correct_packed_param_count(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", 31_577_937_344, nemotron_h
|
||||
"org/Model-27B", 26_900_000_000, {"text_config": _QWEN38_27B_TEXT}
|
||||
) == 26_900_000_000
|
||||
|
||||
|
||||
def test_hybrid_ssm_models_are_never_sized_by_the_estimator():
|
||||
# It prices every Nemotron-H layer as a MoE transformer layer.
|
||||
assert estimate_params_from_arch(_NEMOTRON_H_30B) > 3 * 30e9
|
||||
# Full precision: the safetensors count is exact and stays.
|
||||
assert correct_packed_param_count(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", 31_577_937_344, _NEMOTRON_H_30B
|
||||
) == 31_577_937_344
|
||||
# Packed NVFP4 (17.8B reported, 1.2M downloads): the weekly run wrote
|
||||
# 101.6B for this. The name is the best figure available.
|
||||
assert correct_packed_param_count(
|
||||
"nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", 17_800_000_000,
|
||||
_NEMOTRON_H_30B) == 30_000_000_000
|
||||
# FP8 stores one element per weight, so its count is already right.
|
||||
assert correct_packed_param_count(
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", 31_600_000_000,
|
||||
_NEMOTRON_H_30B) == 31_600_000_000
|
||||
assert correct_packed_param_count(
|
||||
"lmstudio-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4bit", 4_900_000_000,
|
||||
_NEMOTRON_H_30B) == 30_000_000_000
|
||||
# A hybrid with no size in its name keeps what safetensors said.
|
||||
assert correct_packed_param_count(
|
||||
"org/mystery-hybrid-awq", 4_900_000_000, _NEMOTRON_H_30B) == 4_900_000_000
|
||||
# Families rather than exact names: variants already in the catalog.
|
||||
for variant in ("nemotron_h_puzzle", "hybrid_mamba_attn", "rwkv7_native",
|
||||
"lfm2_vl", "qwen3_mamba3", "qu_ssm_moe", "bailing_hybrid"):
|
||||
cfg = {**_NEMOTRON_H_30B, "model_type": variant}
|
||||
assert correct_packed_param_count("org/Variant-30B-4bit", 4_900_000_000, cfg) \
|
||||
== 30_000_000_000, variant
|
||||
# Ordinary transformers still reach the estimator.
|
||||
for plain in ("qwen3_5", "llama", "deepseek_v4", "glm_moe_dsa", "gpt_oss"):
|
||||
cfg = {"text_config": {**_QWEN38_27B_TEXT, "model_type": plain}}
|
||||
assert correct_packed_param_count("org/Plain-27B-AWQ", 7_800_000_000, cfg) > 20e9, plain
|
||||
# Detected by layout fields too, not only model_type.
|
||||
layout = {**_NEMOTRON_H_30B, "model_type": "new_thing",
|
||||
"hybrid_override_pattern": "MEMEM*E"}
|
||||
assert correct_packed_param_count(
|
||||
"org/New-30B-4bit", 4_900_000_000, layout) == 30_000_000_000
|
||||
|
||||
|
||||
def test_estimate_is_capped_by_the_size_the_name_declares():
|
||||
# An estimate far above the declared size is the estimator's error.
|
||||
inflated = {**_NEMOTRON_H_30B, "model_type": "some_moe"}
|
||||
assert correct_packed_param_count(
|
||||
"org/Thing-30B-A3B-AWQ", 5_000_000_000, inflated) == 30_000_000_000
|
||||
# A draft head is named after its target but is a fraction of its size.
|
||||
assert correct_packed_param_count(
|
||||
"z-lab/Qwen3.6-35B-A3B-DFlash", 400_000_000,
|
||||
{"text_config": _QWEN38_27B_TEXT}) == 400_000_000
|
||||
|
||||
|
||||
def test_name_declared_params():
|
||||
cases = {
|
||||
"Qwen/Qwen3-235B-A22B": 235e9,
|
||||
"nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4": 30e9,
|
||||
"Qwen/Qwen3.8-2.4T-A95B": 2.4e12,
|
||||
"meta-llama/Llama-3.1-8B-Instruct": 8e9,
|
||||
"ornith-ai/Ornith-1.0-35B": 35e9,
|
||||
"microsoft/bitnet-b1.58-2B-4T": 2e9, # 4T is training tokens
|
||||
"tzervas/qwen2.5-coder-32b-bitnet-1.58b": 32e9,
|
||||
"mistralai/Mixtral-8x7B-Instruct-v0.1": None, # 46.7B, not 7B
|
||||
"google/gemma-4-E4B-it": None, # effective size
|
||||
"z-lab/Qwen3.6-35B-A3B-DFlash": None, # draft head
|
||||
"deepseek-ai/DeepSeek-V4-Flash-0731": None,
|
||||
}
|
||||
for name, expected in cases.items():
|
||||
got = name_declared_params(name)
|
||||
assert got == (int(expected) if expected else None), (name, got)
|
||||
|
||||
|
||||
def test_revalidation_keeps_the_retained_entry_on_a_sharp_parameter_drop():
|
||||
|
|
@ -628,7 +699,10 @@ if __name__ == "__main__":
|
|||
test_revalidation_ranks_uncorrectable_packed_counts_first,
|
||||
test_revalidation_flags_suspect_context_and_missing_date,
|
||||
test_revalidation_selection_respects_budget_and_skips_fresh_entries,
|
||||
test_packed_count_is_corrected_only_for_quantized_repos,
|
||||
test_packed_count_is_rescued_by_the_architecture_estimate,
|
||||
test_hybrid_ssm_models_are_never_sized_by_the_estimator,
|
||||
test_estimate_is_capped_by_the_size_the_name_declares,
|
||||
test_name_declared_params,
|
||||
test_revalidation_keeps_the_retained_entry_on_a_sharp_parameter_drop,
|
||||
test_revalidation_cooldown_stops_failures_holding_the_budget,
|
||||
test_unquantized_in_the_name_is_not_prequantized,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue