Buckets:
| from __future__ import annotations | |
| from typing import Callable, Iterable, TYPE_CHECKING | |
| import torch | |
| if TYPE_CHECKING: | |
| from torch import Tensor | |
| from .base import ModelBase, TextModel, gguf, logger | |
| class BeamModel(TextModel): | |
| model_arch = gguf.MODEL_ARCH.BEAM | |
| # reference NoPE query-scale constants; not in config.json | |
| NOPE_SCALE_POWER = 0.1 | |
| NOPE_SCALE_LENGTH = 2048 | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| hp = self.hparams | |
| rope = hp.get("rope_parameters") or {} | |
| if rope.get("full_attention") is not None: | |
| raise ValueError("Beam: expected NoPE full-attention layers (rope_parameters.full_attention = null)") | |
| if "sliding_attention" in rope: | |
| swa_rope = rope["sliding_attention"] or {} | |
| else: | |
| # flat rope_theta / rope_scaling, already normalized by TextModel | |
| swa_rope = {k: v for k, v in self.rope_parameters.items() if k not in ("sliding_attention", "full_attention", "type")} | |
| swa_rope.update(self.rope_parameters.get("sliding_attention") or {}) | |
| if swa_rope.get("rope_type", "default") != "default": | |
| raise ValueError(f"Beam: unsupported sliding-attention rope type {swa_rope.get('rope_type')!r}") | |
| # full layers are NoPE, so mirror the sliding rope params for the base class | |
| swa_rope = {k: v for k, v in swa_rope.items() if not (k == "rope_type" and v == "default")} | |
| self.rope_parameters = {"sliding_attention": dict(swa_rope), "full_attention": dict(swa_rope)} | |
| layer_types = hp["layer_types"] | |
| mlp_types = hp["mlp_layer_types"] | |
| n_layer = hp["num_hidden_layers"] | |
| if len(layer_types) != n_layer or len(mlp_types) != n_layer: | |
| raise ValueError("Beam: layer_types / mlp_layer_types must have num_hidden_layers entries") | |
| if any(t not in ("sliding_attention", "full_attention") for t in layer_types): | |
| raise ValueError(f"Beam: unknown layer type in {layer_types}") | |
| if "sliding_attention" not in layer_types: | |
| raise ValueError("Beam: expected at least one sliding_attention layer") | |
| n_dense = 0 | |
| while n_dense < n_layer and mlp_types[n_dense] == "dense": | |
| n_dense += 1 | |
| if any(t != "sparse" for t in mlp_types[n_dense:]): | |
| raise ValueError("Beam: llama.cpp supports dense MLP layers only as a leading block") | |
| self._n_dense_lead = n_dense | |
| if hp.get("attention_bias", False): | |
| raise ValueError("Beam: attention_bias=True is not supported") | |
| self._experts: dict[tuple[int, str], dict[int, Tensor]] = {} | |
| def set_vocab(self): | |
| self._set_vocab_gpt2() | |
| # otherwise EOT is guessed from a hash-ordered scan | |
| from transformers import AutoTokenizer | |
| tok = AutoTokenizer.from_pretrained(self.dir_model) | |
| ret_id = tok.convert_tokens_to_ids("<|return|>") # ty: ignore[unresolved-attribute] | |
| if isinstance(ret_id, int) and ret_id >= 0: | |
| self.gguf_writer.add_eot_token_id(ret_id) | |
| def set_gguf_parameters(self): | |
| super().set_gguf_parameters() | |
| hp = self.hparams | |
| head_dim = hp["head_dim"] | |
| self.gguf_writer.add_rope_dimension_count(head_dim) | |
| self.gguf_writer.add_sliding_window(hp["sliding_window"]) | |
| self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hp["layer_types"]]) | |
| self.gguf_writer.add_attn_temperature_scale(self.NOPE_SCALE_POWER) | |
| self.gguf_writer.add_attn_temperature_length(self.NOPE_SCALE_LENGTH) | |
| self.gguf_writer.add_leading_dense_block_count(self._n_dense_lead) | |
| self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"]) | |
| self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) | |
| self.gguf_writer.add_expert_weights_norm(bool(hp.get("norm_topk_prob", True))) | |
| # equivalent to the reference's post-aggregation scaling | |
| self.gguf_writer.add_expert_weights_scale(float(hp["routed_scaling_factor"])) | |
| if (softcap := hp.get("final_logit_softcapping")) is not None: | |
| self.gguf_writer.add_final_logit_softcapping(float(softcap)) | |
| def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: | |
| name, gen = item | |
| if name.endswith(".self_attn.gate_bias"): | |
| name = name.replace(".self_attn.gate_bias", ".self_attn.gate_proj.bias") | |
| return super().filter_tensors((name, gen)) | |
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | |
| hp = self.hparams | |
| if name.endswith(".self_attn.gate_proj.bias"): | |
| assert tuple(data_torch.shape) == (hp["num_attention_heads"],), f"{name}: one gate bias per head expected, got {tuple(data_torch.shape)}" | |
| elif name.endswith(".self_attn.gate_proj.weight"): | |
| assert tuple(data_torch.shape) == (hp["num_attention_heads"] * hp["head_dim"], hp["hidden_size"]), f"{name}: {tuple(data_torch.shape)}" | |
| elif name.endswith(".mlp.gate.e_score_correction.bias"): | |
| assert tuple(data_torch.shape) == (hp["num_experts"],), f"{name}: {tuple(data_torch.shape)}" | |
| elif name.endswith((".self_attn.q_norm.weight", ".self_attn.k_norm.weight")): | |
| assert tuple(data_torch.shape) == (hp["head_dim"],), f"{name}: per-head q/k norm expected, got {tuple(data_torch.shape)}" | |
| elif bid is not None and bid < self._n_dense_lead and ".mlp.experts." in name: | |
| raise ValueError(f"Beam: expert tensor {name} in a dense layer") | |
| if ".mlp.experts." in name: | |
| assert bid is not None | |
| n_experts = self.hparams["num_experts"] | |
| parts = name.split(".") | |
| xid = int(parts[parts.index("experts") + 1]) | |
| proj = parts[parts.index("experts") + 2] | |
| if proj not in ("gate_proj", "up_proj", "down_proj"): | |
| raise ValueError(f"Beam: unexpected expert tensor {name}") | |
| bucket = self._experts.setdefault((bid, proj), {}) | |
| if xid in bucket: | |
| raise ValueError(f"Beam: duplicate expert tensor {name}") | |
| bucket[xid] = data_torch | |
| if len(bucket) < n_experts: | |
| return | |
| del self._experts[(bid, proj)] | |
| data = torch.stack([bucket[e] for e in range(n_experts)], dim=0) | |
| yield from super().modify_tensors(data, f"model.layers.{bid}.mlp.experts.{proj}.weight", bid) | |
| return | |
| yield from super().modify_tensors(data_torch, name, bid) | |
| def prepare_tensors(self): | |
| super().prepare_tensors() | |
| if self._experts: | |
| missing = sorted(f"layer {b} {p} ({len(v)} experts)" for (b, p), v in self._experts.items()) | |
| raise ValueError(f"Beam: unprocessed experts: {missing}") | |
| logger.info("Beam: all routed experts stacked") | |
Xet Storage Details
- Size:
- 7.03 kB
- Xet hash:
- a18acb0e5e0e459a0681c160ed592c1ad4a7ec2aa7bc2eaf3d68acb9f325ce5d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.