File size: 8,167 Bytes
054df1d 0392ab3 054df1d 0392ab3 054df1d | 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """MinSparkForCausalLM: thin Transformers wrapper around the vendored Meiosis.
Exact semantics: identical to the bundled generate.py's generation
loop (EOS prefix once, truncate to last max_seq_len, effort -> loop count).
No KV cache (min-spark 1.1); right-padding is scoring-only; generation is
single-sequence (enforced in prepare_inputs_for_generation).
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
try:
from .configuration_minspark import MinSparkConfig # remote-code: sibling in cache package
except ImportError:
from configuration_minspark import MinSparkConfig # direct import with staging on sys.path
try:
from .meiosis import Meiosis, build_rope_cache # remote-code: vendored sibling
except ImportError:
from meiosis import Meiosis, build_rope_cache # direct import with staging on sys.path
EFFORT_MAP = {"low": 2, "medium": 3, "high": 4}
class MinSparkForCausalLM(PreTrainedModel, GenerationMixin):
config_class = MinSparkConfig
base_model_prefix = "model"
main_input_name = "input_ids"
supports_gradient_checkpointing = False
_no_split_modules: list[str] = []
def __init__(self, config: MinSparkConfig):
super().__init__(config)
self.model = Meiosis(config.to_meiosis())
self.post_init() # ties weights (no-op: output == input embedding)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed
def set_input_embeddings(self, value: nn.Embedding) -> None:
self.model.embed = value
def get_output_embeddings(self) -> nn.Embedding:
return self.model.embed # tied: unembed reads embed.weight
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
effort: str | None = None,
loops: int | None = None,
past_key_values=None,
use_cache: bool | None = None,
output_attentions: bool = False,
output_hidden_states: bool = False,
return_dict: bool = True,
) -> CausalLMOutputWithPast:
if past_key_values is not None or use_cache:
raise NotImplementedError(
"KV cache is not implemented in min-spark; it arrives in 1.1. "
"Set use_cache=False (the default)."
)
if input_ids.ndim != 2:
raise ValueError(f"input_ids must be (B, T), got shape {tuple(input_ids.shape)}")
if input_ids.shape[1] > self.config.max_seq_len:
raise ValueError(
f"seq_len {input_ids.shape[1]} > max {self.config.max_seq_len}; "
"truncate the context or use generate (which truncates)."
)
self._validate_attention_mask(attention_mask, input_ids)
loop_count = self._resolve_loops(effort, loops)
self._ensure_buffers()
logits = self.model(input_ids, loops=loop_count)
loss = None
if labels is not None:
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100,
)
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=None,
hidden_states=None,
attentions=None,
)
def _resolve_loops(self, effort: str | None, loops: int | None) -> int:
if loops is not None:
if not isinstance(loops, int) or not (1 <= loops <= self.config.max_loops):
raise ValueError(
f"loops must be an int in [1, {self.config.max_loops}], got {loops!r}"
)
return loops
if effort is not None:
if effort not in EFFORT_MAP:
raise ValueError(f"effort must be one of {sorted(EFFORT_MAP)}, got {effort!r}")
return EFFORT_MAP[effort]
return EFFORT_MAP[self.config.effort]
def _ensure_buffers(self) -> None:
"""Rebuild Meiosis's non-persistent buffers on first forward.
from_pretrained constructs the model on torch.device('meta'), so
build_rope_cache runs on meta tensors and yields garbage; transformers
then restores only the persistent weights, never these non-persistent
buffers. The garbage is not reliably non-finite (meta memory can be
finite-but-wrong, e.g. 1e-21), so check the actual first-row value
rather than finiteness, and rebuild unconditionally on first forward.
Idempotent: runs once per instance."""
if getattr(self, "_buffers_ok", False):
return
m = self.model
cos, sin = build_rope_cache(m.config, m.config.max_seq_len)
m.rope_cos.copy_(cos)
m.rope_sin.copy_(sin)
m.last_loop_rms.zero_()
self._buffers_ok = True
def _validate_attention_mask(
self, attention_mask: torch.Tensor | None, input_ids: torch.Tensor
) -> None:
if attention_mask is None:
return
if tuple(attention_mask.shape) != tuple(input_ids.shape):
raise ValueError(
f"attention_mask shape {tuple(attention_mask.shape)} != "
f"input_ids shape {tuple(input_ids.shape)}"
)
mask = attention_mask.bool()
# leading zeros = left padding (any row whose FIRST position is masked out)
if mask.shape[1] >= 1 and (~mask[:, 0]).any():
raise ValueError(
"left-padded batches are not supported; pad to the right or run single-sequence"
)
# interior gap: a 0 followed later by a 1
if mask.shape[1] >= 2 and (mask[:, 1:].long() - mask[:, :-1].long() > 0).any():
raise ValueError(
"attention_mask must be ones or a contiguous ones-then-zeros suffix; "
"interior gaps are not supported"
)
def prepare_inputs_for_generation(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None = None,
effort: str | None = None,
loops: int | None = None,
**kwargs,
):
"""Build the next forward's inputs. Returns exactly these four keys so
generation machinery (cache_position, position_ids, use_cache) is never
echoed into forward, which has no **kwargs. EOS is prepended BEFORE
truncation (generate.py parity — it drops off prompts >max_seq_len);
a supplied attention_mask is extended/truncated in lockstep so its length
always matches the returned input_ids (forward validates mask shape)."""
if input_ids.shape[0] != 1:
raise ValueError(
"batched generation is not supported; run single-sequence generation "
"or right-padded scoring through forward"
)
ids = input_ids
mask = attention_mask
if self.config.doc_mask_eos is not None:
eos = self.config.doc_mask_eos
ids = torch.cat(
[torch.full((1, 1), eos, dtype=ids.dtype, device=ids.device), ids], dim=1
)
if mask is not None:
# The prepended EOS is a real position: keep the mask in sync.
mask = torch.cat(
[torch.ones((1, 1), dtype=mask.dtype, device=mask.device), mask], dim=1
)
if ids.shape[1] > self.config.max_seq_len:
ids = ids[:, -self.config.max_seq_len:]
if mask is not None:
mask = mask[:, -self.config.max_seq_len:]
return {
"input_ids": ids,
"attention_mask": mask,
"effort": effort,
"loops": loops,
}
def _reorder_cache(self, past_key_values, beam_idx):
return past_key_values # no cache
|