File size: 5,435 Bytes
c61c435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | from __future__ import annotations
import math
from pathlib import Path
from typing import Any
from .base import CONFIG_FILENAMES, MODEL_EXTENSIONS, TensorStats, dtype_size, parameter_count
def bytes_label(size: int | float | None) -> str:
if size is None:
return "-"
value = float(size)
for unit in ("B", "KB", "MB", "GB", "TB"):
if abs(value) < 1024 or unit == "TB":
return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B"
value /= 1024
return f"{value:.1f} TB"
def component_for_name(name: str) -> str:
lowered = name.casefold()
mapping = (
("down_blocks", ("down_blocks", "down.", "downsample")),
("mid_block", ("mid_block", "middle_block", "mid.")),
("up_blocks", ("up_blocks", "up.", "upsample")),
("attention", ("attn", "attention", "to_q", "to_k", "to_v", "query", "key", "value")),
("embeddings", ("embed", "embedding", "position", "token")),
("transformer blocks", ("transformer", "blocks.", "layers.", "encoder", "decoder")),
("output layers", ("out.", "output", "proj_out", "lm_head", "conv_out")),
("LoRA adapters", ("lora", "hada", "lokr", "adapter")),
("normalization", ("norm", "bn", "ln", "group_norm", "layer_norm")),
)
for component, tokens in mapping:
if any(token in lowered for token in tokens):
return component
return name.split(".", 1)[0] if "." in name else "other"
def safe_number(value: Any) -> float | None:
try:
number = float(value)
except (TypeError, ValueError, OverflowError):
return None
return number if math.isfinite(number) else None
def tensor_stats_from_torch(name: str, tensor: Any, *, sample_limit: int = 1_000_000) -> TensorStats:
shape = tuple(int(dim) for dim in getattr(tensor, "shape", ()))
dtype = str(getattr(tensor, "dtype", "unknown")).replace("torch.", "")
count = parameter_count(shape)
stat = TensorStats(
name=name,
shape=shape,
dtype=dtype,
parameter_count=count,
memory_bytes=count * dtype_size(dtype),
component=component_for_name(name),
)
if count == 0:
stat.health.append("Empty tensor")
return stat
try:
import torch
with torch.no_grad():
values = tensor.detach().to(device="cpu")
if not values.is_floating_point() and not values.is_complex():
values = values.float()
else:
values = values.float()
flat = values.reshape(-1)
if flat.numel() > sample_limit:
stride = max(1, flat.numel() // sample_limit)
flat = flat[::stride][:sample_limit]
finite = torch.isfinite(flat)
if not bool(finite.all()):
if bool(torch.isnan(flat).any()):
stat.health.append("Invalid: NaN values found")
if bool(torch.isinf(flat).any()):
stat.health.append("Invalid: Inf values found")
flat = flat[finite]
if flat.numel() == 0:
return stat
stat.minimum = safe_number(flat.min().item())
stat.maximum = safe_number(flat.max().item())
stat.mean = safe_number(flat.mean().item())
stat.std = safe_number(flat.std(unbiased=False).item()) if flat.numel() > 1 else 0.0
stat.abs_mean = safe_number(flat.abs().mean().item())
stat.l2_norm = safe_number(torch.linalg.vector_norm(flat).item())
stat.zero_percent = safe_number((flat == 0).float().mean().item() * 100)
except Exception as exc:
stat.health.append(f"Statistics unavailable: {exc}")
if stat.abs_mean is not None and stat.abs_mean > 100:
stat.health.append("Unusual: very large average weight magnitude")
if stat.maximum is not None and stat.minimum is not None and max(abs(stat.maximum), abs(stat.minimum)) > 1_000:
stat.health.append("Unusual: very large absolute weight value")
return stat
def discover_config_files(path: Path) -> list[Path]:
root = path if path.is_dir() else path.parent
files: list[Path] = []
try:
for item in root.rglob("*"):
if item.is_file() and item.name in CONFIG_FILENAMES:
files.append(item)
except OSError:
return []
return sorted(files)
def discover_checkpoint_paths(path: Path) -> list[Path]:
root = path if path.is_dir() else path.parent
candidates: list[Path] = []
try:
for item in root.rglob("*"):
if item.is_file() and item.suffix.casefold() in MODEL_EXTENSIONS:
candidates.append(item)
elif item.is_dir() and item.name.startswith("checkpoint-"):
candidates.append(item)
except OSError:
return []
return sorted(candidates, key=lambda item: (step_from_name(item.name) or -1, str(item)))
def step_from_name(name: str) -> int | None:
import re
matches = re.findall(r"(?:step|checkpoint|epoch|e|s)[-_]?(\d+)", name, flags=re.I)
if not matches:
matches = re.findall(r"(\d+)", name)
if not matches:
return None
try:
return int(matches[-1])
except ValueError:
return None
def shape_label(shape: tuple[int, ...]) -> str:
return " x ".join(str(dim) for dim in shape) if shape else "scalar"
|