lvkaokao commited on
Commit
592a33a
Β·
1 Parent(s): abf5bac

update scheme.

Browse files

Signed-off-by: lkk12014402 <kaokao.lv@intel.com>

app.py CHANGED
@@ -119,7 +119,7 @@ from src.app_helpers.pipeline_table import (
119
  )
120
  from src.app_helpers.queues import filter_failed_quant_df, refresh_pipeline_leaderboard, refresh_queue_tables
121
  from src.app_helpers.sidebar import render_sidebar_menu
122
- from src.app_helpers.submissions import submit_model, submit_quant, analyze_model
123
  from src.app_helpers.zip_stream import register_routes as register_zip_route
124
 
125
 
@@ -520,9 +520,9 @@ with demo:
520
  model_name_textbox = gr.Textbox(label="Model name", placeholder="org/model-name", elem_id="model-input-quant")
521
  with gr.Column(scale=2):
522
  compute_type = gr.Dropdown(
523
- choices=["MXFP4", "NVFP4", "INT4 (W4A16)"],
524
  label="Quantization Scheme",
525
- info='Follows the AutoRound approach. <a href="https://github.com/intel/auto-round" target="_blank" rel="noopener">See AutoRound for details β†—</a>',
526
  multiselect=False,
527
  value="INT4 (W4A16)",
528
  interactive=True,
@@ -698,6 +698,18 @@ with demo:
698
  outputs=[analysis_result_quant, ignore_layers_quant, layer_config_quant],
699
  )
700
 
 
 
 
 
 
 
 
 
 
 
 
 
701
  with gr.Column():
702
  with gr.Row(elem_classes=["filter-row"]):
703
  my_submissions_quant_cb = gr.Checkbox(
 
119
  )
120
  from src.app_helpers.queues import filter_failed_quant_df, refresh_pipeline_leaderboard, refresh_queue_tables
121
  from src.app_helpers.sidebar import render_sidebar_menu
122
+ from src.app_helpers.submissions import submit_model, submit_quant, analyze_model, guard_method_scheme
123
  from src.app_helpers.zip_stream import register_routes as register_zip_route
124
 
125
 
 
520
  model_name_textbox = gr.Textbox(label="Model name", placeholder="org/model-name", elem_id="model-input-quant")
521
  with gr.Column(scale=2):
522
  compute_type = gr.Dropdown(
523
+ choices=["MXFP4", "MXFP8", "NVFP4", "INT4 (W4A16)"],
524
  label="Quantization Scheme",
525
+ info='Global weight precision. Combine MXFP8 (global) + a Layer Config that overrides experts to MXFP4 for mixed precision. <a href="https://github.com/intel/auto-round" target="_blank" rel="noopener">See AutoRound for details β†—</a>',
526
  multiselect=False,
527
  value="INT4 (W4A16)",
528
  interactive=True,
 
698
  outputs=[analysis_result_quant, ignore_layers_quant, layer_config_quant],
699
  )
700
 
701
+ # Keep Method compatible with Scheme (Model-Free only for weight-only schemes).
702
+ compute_type.change(
703
+ fn=guard_method_scheme,
704
+ inputs=[compute_type, quant_method],
705
+ outputs=[quant_method],
706
+ )
707
+ quant_method.change(
708
+ fn=guard_method_scheme,
709
+ inputs=[compute_type, quant_method],
710
+ outputs=[quant_method],
711
+ )
712
+
713
  with gr.Column():
714
  with gr.Row(elem_classes=["filter-row"]):
715
  my_submissions_quant_cb = gr.Checkbox(
src/app_helpers/submissions.py CHANGED
@@ -24,6 +24,10 @@ from src.submission.model_analysis import analyze_model_structure, render_analys
24
 
25
  logger = logging.getLogger(__name__)
26
 
 
 
 
 
27
 
28
  def _parse_model_size_input(model_size_input):
29
  """Parse an optional user-provided param count.
@@ -159,14 +163,32 @@ def analyze_model(model, ignore_current=None, layercfg_current=None,
159
  The recommended ignore/layer_config are pre-filled ONLY when the target
160
  textbox is currently empty, so a user's own edits are never clobbered.
161
  Read-only and safe; gated in the UI to whitelisted users.
 
 
 
162
  """
163
  if not model or not model.strip():
164
  return gr.update(value="Please enter a model name first.", visible=True), gr.update(), gr.update()
165
 
166
  token = oauth_token.token if oauth_token else None
167
  try:
168
- res = analyze_model_structure(model.strip(), token=token)
 
 
 
169
  md = render_analysis_markdown(res)
 
 
 
 
 
 
 
 
 
 
 
 
170
  except Exception as e: # never let the button crash the app
171
  logger.warning("[analyze_model] failed for %s: %s", model, e)
172
  return gr.update(value=f"Analysis error: {e}", visible=True), gr.update(), gr.update()
@@ -181,9 +203,24 @@ def analyze_model(model, ignore_current=None, layercfg_current=None,
181
  return gr.update(value=md, visible=True), ign_upd, lc_upd
182
 
183
 
184
- def _empty_result():
185
- """Five outputs that signal "no progress" to the streaming UI."""
186
- return "", gr.update(), gr.update(interactive=True), gr.update(), gr.update()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
187
 
188
 
189
  def _bail_with_error(message):
 
24
 
25
  logger = logging.getLogger(__name__)
26
 
27
+ # Wall-clock guard for the "Analyze Structure" button so the UI never hangs on a
28
+ # pathologically large / slow-to-fetch safetensors index.
29
+ ANALYZE_TIMEOUT_SEC = 90
30
+
31
 
32
  def _parse_model_size_input(model_size_input):
33
  """Parse an optional user-provided param count.
 
163
  The recommended ignore/layer_config are pre-filled ONLY when the target
164
  textbox is currently empty, so a user's own edits are never clobbered.
165
  Read-only and safe; gated in the UI to whitelisted users.
166
+
167
+ A wall-clock timeout guards against pathologically large models whose
168
+ safetensors index is huge/slow to fetch (the UI must never hang).
169
  """
170
  if not model or not model.strip():
171
  return gr.update(value="Please enter a model name first.", visible=True), gr.update(), gr.update()
172
 
173
  token = oauth_token.token if oauth_token else None
174
  try:
175
+ import concurrent.futures
176
+ with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
177
+ fut = ex.submit(analyze_model_structure, model.strip(), "main", token)
178
+ res = fut.result(timeout=ANALYZE_TIMEOUT_SEC)
179
  md = render_analysis_markdown(res)
180
+ except concurrent.futures.TimeoutError:
181
+ return (
182
+ gr.update(
183
+ value=(
184
+ f"⏱️ Analysis timed out after {ANALYZE_TIMEOUT_SEC}s β€” this model's "
185
+ "safetensors index is very large or the Hub is slow. Please retry, or "
186
+ "define Ignore Layers / Layer Config manually."
187
+ ),
188
+ visible=True,
189
+ ),
190
+ gr.update(), gr.update(),
191
+ )
192
  except Exception as e: # never let the button crash the app
193
  logger.warning("[analyze_model] failed for %s: %s", model, e)
194
  return gr.update(value=f"Analysis error: {e}", visible=True), gr.update(), gr.update()
 
203
  return gr.update(value=md, visible=True), ign_upd, lc_upd
204
 
205
 
206
+ def guard_method_scheme(scheme_choice, method_choice):
207
+ """Proactively keep the Method compatible with the Scheme.
208
+
209
+ Model-Free only supports weight-only schemes (W4A16/MXFP4/MXFP8). If the user
210
+ picks an incompatible scheme (e.g. NVFP4) while Method is Model-Free, reset the
211
+ method to RTN and surface a non-blocking notice. Returns a ``gr.update`` for the
212
+ method dropdown. The authoritative check still runs server-side in ``submit_quant``.
213
+ """
214
+ if method_choice and "Model-Free" in method_choice and not is_model_free_supported_scheme(scheme_choice):
215
+ try:
216
+ gr.Info(f"Model-Free doesn't support {scheme_choice}; switched Method to RTN.")
217
+ except Exception:
218
+ pass
219
+ return gr.update(value="RTN (Round-To-Nearest)")
220
+ return gr.update()
221
+
222
+
223
+
224
 
225
 
226
  def _bail_with_error(message):
src/submission/check_validity.py CHANGED
@@ -526,6 +526,15 @@ SUPPORTED_QUANT_SCHEMES: dict[str, QuantSchemeSpec] = {
526
  script="ITREX",
527
  detector=lambda cfg: (False, ""), # input models are FP β€” no pre-quant detection
528
  ),
 
 
 
 
 
 
 
 
 
529
  # ── Add future schemes here ──────────────────────────────────────
530
  # "INT8 (W8A16)": QuantSchemeSpec(
531
  # name="INT8 (W8A16)", precision="8bit", weight_dtype="int8",
 
526
  script="ITREX",
527
  detector=lambda cfg: (False, ""), # input models are FP β€” no pre-quant detection
528
  ),
529
+ "MXFP8": QuantSchemeSpec(
530
+ name="MXFP8",
531
+ precision="8bit",
532
+ weight_dtype="mxfp8",
533
+ bits=8,
534
+ hardware="gpu",
535
+ script="ITREX",
536
+ detector=lambda cfg: (False, ""), # input models are FP β€” no pre-quant detection
537
+ ),
538
  # ── Add future schemes here ──────────────────────────────────────
539
  # "INT8 (W8A16)": QuantSchemeSpec(
540
  # name="INT8 (W8A16)", precision="8bit", weight_dtype="int8",
src/submission/model_analysis.py CHANGED
@@ -10,32 +10,40 @@ Given a HuggingFace model id, this inspects the model's ``config.json`` and the
10
  * recommended ``ignore_layers`` / mixed-precision ``layer_config`` presets a user
11
  can drop straight into the advanced submission fields.
12
 
13
- The heavy lifting uses ``huggingface_hub.get_safetensors_metadata`` which reads
14
- only the safetensors headers (tensor name + shape + dtype), so even trillion-param
15
- models are analyzed without downloading weights.
 
16
  """
17
  from __future__ import annotations
18
 
 
19
  import logging
20
  import re
21
  from dataclasses import dataclass, field
22
 
23
- from huggingface_hub import get_safetensors_metadata
 
24
  from transformers import AutoConfig
25
 
26
  logger = logging.getLogger(__name__)
27
 
28
- # Collapse per-layer indices so ``...layers.12.self_attn.q_proj.weight`` and
29
- # ``...layers.0.self_attn.q_proj.weight`` bucket together.
 
30
  _LAYER_IDX_RE = re.compile(r"\.\d+\.")
 
 
 
31
 
32
 
33
  @dataclass
34
- class CategoryStat:
35
- name: str
36
- params: int = 0
37
- tensors_2d: int = 0
38
- tensors_total: int = 0
 
39
 
40
 
41
  @dataclass
@@ -49,13 +57,15 @@ class ModelAnalysis:
49
  num_layers: int | None = None
50
  num_experts: int | None = None
51
  vocab_size: int | None = None
52
- total_params: int = 0
53
  is_moe: bool = False
54
  has_shared_experts: bool = False
55
  has_attn_indexer: bool = False
56
  has_vision: bool = False
57
- approximate: bool = False
58
- categories: list[CategoryStat] = field(default_factory=list)
 
 
 
59
  recommended_ignore_layers: str = ""
60
  recommended_layer_config: str = ""
61
 
@@ -148,91 +158,215 @@ def analyze_model_structure(model_id: str, revision: str = "main", token: str |
148
  except Exception as e:
149
  logger.warning("[analyze] config load failed for %s: %s", model_id, e)
150
 
151
- # 2. Per-tensor metadata from safetensors headers (no weight download).
152
- try:
153
- meta = get_safetensors_metadata(model_id, revision=revision, token=token)
154
- except Exception as e:
 
155
  res.ok = False
156
- res.error = (
157
- f"Could not read safetensors metadata: {e}. "
158
- "The model may lack safetensors (e.g. GGUF/pytorch_model.bin) or be gated/private."
159
- )
160
  return res
161
 
162
- # Map filename β†’ {tensor_name: TensorInfo} for shapes/dtypes/param counts.
163
- cats: dict[str, CategoryStat] = {c: CategoryStat(c) for c in _CATEGORY_ORDER}
164
- total = 0
165
- for fname, fmeta in (meta.files_metadata or {}).items():
166
- for tname, tinfo in (fmeta.tensors or {}).items():
167
- pc = int(getattr(tinfo, "parameter_count", 0) or 0)
168
- shape = list(getattr(tinfo, "shape", []) or [])
169
- norm = _LAYER_IDX_RE.sub(".N.", tname)
170
- cat = _categorize(norm)
171
- cs = cats[cat]
172
- cs.params += pc
173
- cs.tensors_total += 1
174
- if len(shape) == 2:
175
- cs.tensors_2d += 1
176
- total += pc
177
-
178
- res.total_params = total
179
- res.categories = [cats[c] for c in _CATEGORY_ORDER if cats[c].tensors_total > 0]
180
-
181
- # 3. Derived traits.
182
- res.is_moe = cats["moe_experts"].tensors_total > 0 or bool(res.num_experts)
183
- res.has_shared_experts = cats["shared_experts"].tensors_total > 0
184
- res.has_attn_indexer = cats["attn_indexer"].tensors_total > 0
185
- res.has_vision = cats["vision"].tensors_total > 0
186
-
187
- # 4. Recommendations.
188
- res.recommended_ignore_layers, res.recommended_layer_config = _recommend(res, cats)
 
189
  return res
190
 
191
 
192
- def _recommend(res: ModelAnalysis, cats: dict[str, CategoryStat]) -> tuple[str, str]:
193
- """Produce a recommended ignore_layers list + mixed-precision layer_config.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
194
 
195
  Heuristic (from the reference mixed-precision docs):
196
- * Always skip/ignore: lm_head, embeddings (type-skipped), router/gate, norms (1D).
197
- * Vision tower, projectors, attention indexer β†’ ignore (keep bf16).
198
- * MoE routed experts β†’ low-bit (MXFP4); shared experts β†’ higher precision
199
- via layer_config (MXFP8) when present.
200
  """
201
  ignore: list[str] = []
202
- if cats["lm_head"].tensors_total:
203
- ignore.append("lm_head")
204
- if cats["router_gate"].tensors_total:
205
- # Common router names; users can trim to their model's exact name.
206
- ignore.append("gate")
207
- if cats["vision"].tensors_total:
208
- ignore.append("vision_tower")
209
- if cats["projector"].tensors_total:
210
- ignore.append("multi_modal_projector")
211
- ignore.append("patch_merge")
212
- if cats["attn_indexer"].tensors_total:
213
- ignore.append("indexer")
214
  # Deduplicate, preserve order.
215
- seen = set()
216
  ignore = [x for x in ignore if not (x in seen or seen.add(x))]
217
 
218
  layer_config = ""
219
- if res.is_moe and cats["moe_experts"].tensors_total:
220
- # Route the routed experts to MXFP4; leave the rest at the global scheme.
221
- # 'experts' matches routed experts without touching 'shared_experts'.
222
- layer_config = "{experts:{bits:4,data_type:mx_fp}}"
223
 
224
  return ",".join(ignore), layer_config
225
 
226
 
227
  # ── Markdown rendering ─────────────────────────────────────────────────────
228
- def _fmt_b(params: int) -> str:
229
- if params >= 1e9:
230
- return f"{params / 1e9:.2f} B"
231
- if params >= 1e6:
232
- return f"{params / 1e6:.2f} M"
233
- return str(params)
234
-
235
-
236
  def render_analysis_markdown(res: ModelAnalysis) -> str:
237
  if not res.ok:
238
  return f"### ⚠️ Model structure analysis failed\n\n{res.error}"
@@ -247,7 +381,6 @@ def render_analysis_markdown(res: ModelAnalysis) -> str:
247
  lines.append(f"| Architecture | `{', '.join(res.architectures)}` |")
248
  if res.model_type:
249
  lines.append(f"| model_type | `{res.model_type}` |")
250
- lines.append(f"| Total parameters (measured) | **{_fmt_b(res.total_params)}** |")
251
  if res.num_layers is not None:
252
  lines.append(f"| Layers | {res.num_layers} |")
253
  if res.num_experts:
@@ -256,6 +389,7 @@ def render_analysis_markdown(res: ModelAnalysis) -> str:
256
  lines.append(f"| hidden_size | {res.hidden_size} |")
257
  if res.vocab_size is not None:
258
  lines.append(f"| vocab_size | {res.vocab_size} |")
 
259
  traits = []
260
  if res.is_moe:
261
  traits.append("MoE")
@@ -269,27 +403,33 @@ def render_analysis_markdown(res: ModelAnalysis) -> str:
269
  lines.append(f"| Traits | {', '.join(traits)} |")
270
  lines.append("")
271
 
272
- # Parameter distribution
273
- lines.append("#### Parameter distribution by module category\n")
274
- lines.append("| Category | Params | Share | 2D-linear tensors |")
275
- lines.append("|---|---:|---:|---:|")
276
- total = res.total_params or 1
277
- for cs in sorted(res.categories, key=lambda c: c.params, reverse=True):
278
- share = 100.0 * cs.params / total
279
- label = _CATEGORY_LABELS.get(cs.name, cs.name)
280
- lines.append(f"| {label} | {_fmt_b(cs.params)} | {share:.2f}% | {cs.tensors_2d} |")
 
 
 
 
 
281
  lines.append("")
282
 
283
- # Recommendations
284
- lines.append("#### Recommended quantization controls\n")
285
  if res.recommended_ignore_layers:
286
  lines.append(f"- **Ignore Layers:** `{res.recommended_ignore_layers}`")
287
  else:
288
- lines.append("- **Ignore Layers:** *(defaults are fine β€” nothing extra to skip)*")
289
  if res.recommended_layer_config:
290
  lines.append(f"- **Layer Config (mixed precision):** `{res.recommended_layer_config}`")
291
  lines.append(" - routes MoE routed experts to MXFP4 while the rest stay at the global scheme.")
292
  lines.append("")
293
- lines.append("> These are suggestions β€” review against your model's exact module names before submitting. "
294
  "Norms (1D) and embeddings are skipped automatically by AutoRound.")
295
  return "\n".join(lines)
 
 
10
  * recommended ``ignore_layers`` / mixed-precision ``layer_config`` presets a user
11
  can drop straight into the advanced submission fields.
12
 
13
+ The structure comes from the **safetensors index** (``model.safetensors.index.json``) β€”
14
+ a single small JSON that lists every tensor name and its shard. No weights, shapes, or
15
+ dtypes are downloaded, so even 100+ shard / trillion-param models are analyzed in seconds.
16
+ Module types are inferred from tensor names (Linear / Embedding / Norm / Bias).
17
  """
18
  from __future__ import annotations
19
 
20
+ import json
21
  import logging
22
  import re
23
  from dataclasses import dataclass, field
24
 
25
+ from huggingface_hub import HfApi, hf_hub_download
26
+ from huggingface_hub.utils import EntryNotFoundError
27
  from transformers import AutoConfig
28
 
29
  logger = logging.getLogger(__name__)
30
 
31
+ # Collapse per-layer AND per-expert indices so
32
+ # ``...layers.12.mlp.experts.37.gate_proj.weight`` and
33
+ # ``...layers.0.mlp.experts.5.gate_proj.weight`` bucket to one module.
34
  _LAYER_IDX_RE = re.compile(r"\.\d+\.")
35
+ # Also collapse a trailing ``.<digits>`` (e.g. ``oe_embed_proj0`` style is rare;
36
+ # handles names like ``...experts.37`` with no suffix).
37
+ _TRAIL_IDX_RE = re.compile(r"\.\d+$")
38
 
39
 
40
  @dataclass
41
+ class ModuleStat:
42
+ """One normalized module (layer/expert indices collapsed to N)."""
43
+ name: str # e.g. "...layers.N.block_sparse_moe.experts.N.gate_proj"
44
+ category: str = ""
45
+ kind: str = "" # "Linear" | "Embedding" | "Norm (1D)" | "Bias (1D)"
46
+ count: int = 0 # how many real tensors collapsed into this row
47
 
48
 
49
  @dataclass
 
57
  num_layers: int | None = None
58
  num_experts: int | None = None
59
  vocab_size: int | None = None
 
60
  is_moe: bool = False
61
  has_shared_experts: bool = False
62
  has_attn_indexer: bool = False
63
  has_vision: bool = False
64
+ num_tensors: int = 0
65
+ modules: list[ModuleStat] = field(default_factory=list)
66
+ recommended_ignore_layers: str = ""
67
+ recommended_layer_config: str = ""
68
+ modules: list[ModuleStat] = field(default_factory=list)
69
  recommended_ignore_layers: str = ""
70
  recommended_layer_config: str = ""
71
 
 
158
  except Exception as e:
159
  logger.warning("[analyze] config load failed for %s: %s", model_id, e)
160
 
161
+ # 2. Tensor NAMES only β€” from the safetensors index (one small JSON), so even
162
+ # 194-shard / trillion-param models are analyzed in seconds. No weights, no
163
+ # shapes/dtypes downloaded. Falls back to a single-file header for unsharded models.
164
+ names, err = _list_tensor_names(model_id, revision, token)
165
+ if err:
166
  res.ok = False
167
+ res.error = err
 
 
 
168
  return res
169
 
170
+ res.num_tensors = len(names)
171
+ mods: dict[str, ModuleStat] = {}
172
+ cat_present: set[str] = set()
173
+ for tname in names:
174
+ # Collapse layer + expert indices: any ".<digits>." and a trailing ".<digits>".
175
+ norm = _TRAIL_IDX_RE.sub(".N", _LAYER_IDX_RE.sub(".N.", tname))
176
+ cat = _categorize(norm)
177
+ cat_present.add(cat)
178
+ mkey = _module_key(norm)
179
+ ms = mods.get(mkey)
180
+ if ms is None:
181
+ ms = ModuleStat(name=mkey, category=cat, kind=_kind_from_name(mkey))
182
+ mods[mkey] = ms
183
+ ms.count += 1
184
+
185
+ # Order modules by category (structural grouping), then name β€” reads like the
186
+ # reference "model structure" table rather than a params ranking.
187
+ cat_rank = {c: i for i, c in enumerate(_CATEGORY_ORDER)}
188
+ res.modules = sorted(mods.values(), key=lambda m: (cat_rank.get(m.category, 99), m.name))
189
+
190
+ # 3. Derived traits (from names + config).
191
+ res.is_moe = ("moe_experts" in cat_present) or bool(res.num_experts)
192
+ res.has_shared_experts = "shared_experts" in cat_present
193
+ res.has_attn_indexer = "attn_indexer" in cat_present
194
+ res.has_vision = "vision" in cat_present
195
+
196
+ # 4. Suggested controls (optional; derived from REAL module names β†’ precise substrings).
197
+ res.recommended_ignore_layers, res.recommended_layer_config = _recommend(res, cat_present, mods)
198
  return res
199
 
200
 
201
+ def _list_tensor_names(model_id: str, revision: str, token: str | None):
202
+ """Return ``(names, error)`` β€” the model's tensor names without downloading weights.
203
+
204
+ Strategy (fast β†’ fallback):
205
+ 1. ``model.safetensors.index.json`` (sharded) β†’ ``weight_map`` keys. One file.
206
+ 2. a single ``*.safetensors`` file β†’ read just its header via the HfApi.
207
+ Never downloads weight bytes.
208
+ """
209
+ api = HfApi(token=token)
210
+ # 1. Sharded index (covers the vast majority of large models).
211
+ for index_name in ("model.safetensors.index.json", "pytorch_model.bin.index.json"):
212
+ try:
213
+ path = hf_hub_download(model_id, index_name, revision=revision, token=token)
214
+ with open(path) as f:
215
+ data = json.load(f)
216
+ wm = data.get("weight_map") or {}
217
+ if wm:
218
+ return list(wm.keys()), None
219
+ except EntryNotFoundError:
220
+ continue
221
+ except Exception as e:
222
+ logger.warning("[analyze] index read failed for %s (%s): %s", model_id, index_name, e)
223
+
224
+ # 2. Single-file safetensors β†’ header only.
225
+ try:
226
+ files = api.list_repo_files(model_id, revision=revision)
227
+ st_files = [f for f in files if f.endswith(".safetensors")]
228
+ if st_files:
229
+ meta = api.parse_safetensors_file_metadata(model_id, st_files[0], revision=revision)
230
+ return list((meta.tensors or {}).keys()), None
231
+ # No safetensors at all.
232
+ return [], (
233
+ "No safetensors found (model may be GGUF / pytorch_model.bin only). "
234
+ "Structure analysis needs a safetensors checkpoint."
235
+ )
236
+ except Exception as e:
237
+ return [], (
238
+ f"Could not read model index/metadata: {e}. "
239
+ "The model may be gated/private (log in) or lack safetensors."
240
+ )
241
+
242
+
243
+ def _module_key(norm_name: str) -> str:
244
+ """Normalized tensor name β†’ module key (drop trailing .weight/.bias/.scale)."""
245
+ for suf in (".weight", ".bias", ".weight_scale", ".weight_packed", ".scale", ".g_idx", ".qweight", ".qzeros", ".scales"):
246
+ if norm_name.endswith(suf):
247
+ return norm_name[: -len(suf)]
248
+ return norm_name
249
+
250
+
251
+ def _kind_from_name(tname: str) -> str:
252
+ """Infer the module type from its name only (no shape needed).
253
+
254
+ Mirrors the human 'type' column in the reference structure tables.
255
+ """
256
+ n = tname.lower()
257
+ if "lm_head" in n:
258
+ return "Linear (head)"
259
+ if "embed" in n or "wte" in n or "word_embeddings" in n:
260
+ return "Embedding"
261
+ if n.endswith("_bias") or n.endswith(".bias") or "correction_bias" in n:
262
+ return "Bias (1D)"
263
+ if "norm" in n or "layernorm" in n or "ln_f" in n or "ln_1" in n or "ln_2" in n:
264
+ return "Norm (1D)"
265
+ # Everything else that carries weights is a Linear/projection in these models
266
+ # (q/k/v/o_proj, gate, experts w1/w2/w3, fc1/fc2, router classifier, …).
267
+ return "Linear"
268
+
269
+
270
+ def _rel_name(module_key: str) -> str:
271
+ """Strip the model wrapper + the ``layers.N.`` prefix to get a substring a user
272
+ can paste into ignore_layers. e.g.
273
+ ``language_model.model.layers.N.block_sparse_moe.gate`` β†’ ``block_sparse_moe.gate``.
274
+ """
275
+ key = module_key
276
+ for pre in ("language_model.model.", "language_model.", "model.model.", "model.", "transformer."):
277
+ if key.startswith(pre):
278
+ key = key[len(pre):]
279
+ break
280
+ # Drop everything up to and including the first ``layers.N.``
281
+ key = re.sub(r"^.*?layers\.N\.", "", key)
282
+ return key
283
+
284
+
285
+ def _char_lcp(strings: list[str]) -> str:
286
+ """Character-level longest common prefix (used to fold sibling leaves into one
287
+ precise substring, e.g. self_attn.index_q_proj + self_attn.index_k_proj β†’
288
+ ``self_attn.index_``)."""
289
+ if not strings:
290
+ return ""
291
+ s1, s2 = min(strings), max(strings)
292
+ i = 0
293
+ while i < len(s1) and i < len(s2) and s1[i] == s2[i]:
294
+ i += 1
295
+ return s1[:i]
296
+
297
+
298
+ def _ignore_token_for(category_keys: list[str]) -> str | None:
299
+ """Turn a category's real module keys into ONE precise ignore substring.
300
+
301
+ Uses the relative names (after ``layers.N.``); if there are several siblings,
302
+ a character-level common prefix yields a safe substring (never the bare last
303
+ segment like ``gate`` which would also hit ``gate_proj``).
304
+ """
305
+ rels = sorted({_rel_name(k) for k in category_keys})
306
+ rels = [r for r in rels if r]
307
+ if not rels:
308
+ return None
309
+ if len(rels) == 1:
310
+ return rels[0]
311
+ lcp = _char_lcp(rels)
312
+ # Require a reasonably specific prefix; otherwise just return the shortest name.
313
+ if len(lcp) >= 4:
314
+ return lcp
315
+ return min(rels, key=len)
316
+
317
+
318
+ def _experts_layer_config_key(expert_keys: list[str]) -> str:
319
+ """Derive the precise ``layer_config`` key for routed experts from real names.
320
+
321
+ e.g. ``block_sparse_moe.experts.N.gate_proj`` β†’ ``block_sparse_moe.experts``;
322
+ ``mlp.experts.N.gate_proj`` β†’ ``mlp.experts``. Using ``<parent>.experts``
323
+ (not bare ``experts``) keeps it precise and never touches ``shared_experts``.
324
+ """
325
+ for k in expert_keys:
326
+ rel = _rel_name(k)
327
+ segs = rel.split(".")
328
+ for i, s in enumerate(segs):
329
+ if s == "experts":
330
+ return ".".join(segs[max(0, i - 1):i + 1]) if i > 0 else "experts"
331
+ return "experts"
332
+
333
+
334
+ def _keys_in(mods: dict, category: str) -> list[str]:
335
+ return [k for k, m in mods.items() if m.category == category]
336
+
337
+
338
+ def _recommend(res: ModelAnalysis, cat_present: set, mods: dict) -> tuple[str, str]:
339
+ """Produce suggested ignore_layers + mixed-precision layer_config using the
340
+ model's REAL module names, so the substrings are precise and safe.
341
+
342
+ Critical: never emit a bare last segment like ``gate`` β€” under auto-round's
343
+ substring matching that would also hit ``gate_proj`` (a dense-MLP projection).
344
+ We derive ``<parent>.gate`` / ``self_attn.index_`` etc. from actual names.
345
 
346
  Heuristic (from the reference mixed-precision docs):
347
+ * ignore: lm_head, router/gate, vision tower, projectors, attention indexer.
348
+ * embeddings + norms are auto-skipped by AutoRound (not listed).
349
+ * MoE routed experts β†’ MXFP4 via layer_config; the rest stay at the global scheme.
 
350
  """
351
  ignore: list[str] = []
352
+ for cat in ("lm_head", "router_gate", "vision", "projector", "attn_indexer"):
353
+ if cat in cat_present:
354
+ tok = _ignore_token_for(_keys_in(mods, cat))
355
+ if tok:
356
+ ignore.append(tok)
 
 
 
 
 
 
 
357
  # Deduplicate, preserve order.
358
+ seen: set[str] = set()
359
  ignore = [x for x in ignore if not (x in seen or seen.add(x))]
360
 
361
  layer_config = ""
362
+ if res.is_moe and "moe_experts" in cat_present:
363
+ key = _experts_layer_config_key(_keys_in(mods, "moe_experts"))
364
+ layer_config = f"{{{key}:{{bits:4,data_type:mx_fp}}}}"
 
365
 
366
  return ",".join(ignore), layer_config
367
 
368
 
369
  # ── Markdown rendering ─────────────────────────────────────────────────────
 
 
 
 
 
 
 
 
370
  def render_analysis_markdown(res: ModelAnalysis) -> str:
371
  if not res.ok:
372
  return f"### ⚠️ Model structure analysis failed\n\n{res.error}"
 
381
  lines.append(f"| Architecture | `{', '.join(res.architectures)}` |")
382
  if res.model_type:
383
  lines.append(f"| model_type | `{res.model_type}` |")
 
384
  if res.num_layers is not None:
385
  lines.append(f"| Layers | {res.num_layers} |")
386
  if res.num_experts:
 
389
  lines.append(f"| hidden_size | {res.hidden_size} |")
390
  if res.vocab_size is not None:
391
  lines.append(f"| vocab_size | {res.vocab_size} |")
392
+ lines.append(f"| Tensors (total) | {res.num_tensors} |")
393
  traits = []
394
  if res.is_moe:
395
  traits.append("MoE")
 
403
  lines.append(f"| Traits | {', '.join(traits)} |")
404
  lines.append("")
405
 
406
+ # Module schema β€” the model's composition, grouped by category. Layer & expert
407
+ # indices are shown as `N`. These are the exact substrings a user pastes into
408
+ # Ignore Layers / Layer Config.
409
+ lines.append("#### Model structure β€” modules (normalized: layer & expert indices = `N`)\n")
410
+ lines.append("Use these names to craft **Ignore Layers** / **Layer Config** substrings.\n")
411
+ lines.append("| Module | Type | Category | Count |")
412
+ lines.append("|---|---|---|---:|")
413
+ _MAX_ROWS = 80
414
+ shown = res.modules[:_MAX_ROWS]
415
+ for ms in shown:
416
+ label = _CATEGORY_LABELS.get(ms.category, ms.category)
417
+ lines.append(f"| `{ms.name}` | {ms.kind} | {label} | {ms.count} |")
418
+ if len(res.modules) > _MAX_ROWS:
419
+ lines.append(f"| … | | | *(+{len(res.modules) - _MAX_ROWS} more)* |")
420
  lines.append("")
421
 
422
+ # Suggested controls (optional; users can define their own from the table above)
423
+ lines.append("#### Suggested quantization controls (optional)\n")
424
  if res.recommended_ignore_layers:
425
  lines.append(f"- **Ignore Layers:** `{res.recommended_ignore_layers}`")
426
  else:
427
+ lines.append("- **Ignore Layers:** *(nothing extra suggested β€” defaults are fine)*")
428
  if res.recommended_layer_config:
429
  lines.append(f"- **Layer Config (mixed precision):** `{res.recommended_layer_config}`")
430
  lines.append(" - routes MoE routed experts to MXFP4 while the rest stay at the global scheme.")
431
  lines.append("")
432
+ lines.append("> Suggestions only β€” review against the module table above before submitting. "
433
  "Norms (1D) and embeddings are skipped automatically by AutoRound.")
434
  return "\n".join(lines)
435
+