diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index d3af530..b04000d 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -107,12 +107,13 @@ class TorchModelShard: self.tokenizer = AutoTokenizer.from_pretrained(model_id) self.layers = _model_layers(self.model) self.total_layers = len(self.layers) - if shard_end > self.total_layers: + # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention. + if shard_end >= self.total_layers: raise ValueError( - f"shard_end {shard_end} exceeds total layer count {self.total_layers}" + f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}" ) self.is_head = shard_start == 0 - self.is_tail = shard_end == self.total_layers + self.is_tail = shard_end >= self.total_layers - 1 self.hidden_size = int( getattr(self.model.config, "hidden_size", 0) or getattr(self.model.config, "n_embd", 0) @@ -168,10 +169,60 @@ class TorchModelShard: token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) return self.tokenizer.decode([token_id], skip_special_tokens=True) - def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any: + def generate_text(self, prompt: str, max_new_tokens: int = 16) -> str: + """Generate text locally when this process owns the full model.""" + if not self.is_head or not self.is_tail: + raise ModelBackendError("local generation requires a full head+tail shard") + if hasattr(self.tokenizer, "apply_chat_template"): + try: + encoded = self.tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + except Exception: + encoded = self.tokenizer(prompt, return_tensors="pt") + else: + encoded = self.tokenizer(prompt, return_tensors="pt") + input_ids = encoded["input_ids"].to(self.device) + attention_mask = encoded.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(self.device) + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_token_id is None: + pad_token_id = getattr(self.tokenizer, "eos_token_id", None) with self.torch.inference_mode(): - for layer in self.layers[self.shard_start:self.shard_end]: - hidden_states = _call_layer(layer, hidden_states, attention_mask, position_ids) + generated = self.model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=max(1, int(max_new_tokens)), + do_sample=False, + pad_token_id=pad_token_id, + ) + new_tokens = generated[0, input_ids.shape[-1]:] + return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + + def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any: + position_embeddings = _rotary_position_embeddings( + self.model, + hidden_states, + position_ids, + ) + layer_attention_mask = _decoder_attention_mask( + attention_mask, + hidden_states, + self.torch, + ) + with self.torch.inference_mode(): + for layer in self.layers[self.shard_start:self.shard_end + 1]: + hidden_states = _call_layer( + layer, + hidden_states, + layer_attention_mask, + position_ids, + position_embeddings, + ) return hidden_states.to(self.torch.bfloat16) def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload: @@ -236,8 +287,60 @@ def _position_ids(attention_mask: Any, torch: Any) -> Any: return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long) -def _call_layer(layer: Any, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any: +def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any: + """Build a causal additive mask for decoder layers called outside model.forward.""" + if attention_mask is None: + return None + if len(getattr(attention_mask, "shape", ())) != 2: + return attention_mask + batch_size, seq_len = attention_mask.shape + if seq_len <= 1: + return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype) + + min_value = torch.finfo(hidden_states.dtype).min + causal = torch.full( + (seq_len, seq_len), + min_value, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + causal = torch.triu(causal, diagonal=1) + causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone() + + padding = attention_mask.to(device=hidden_states.device) + if not bool(padding.all()): + causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value) + return causal + + +def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None: + """Return model-level rotary embeddings required by newer HF decoder layers.""" + if position_ids is None: + return None + rotary = None + if hasattr(model, "model") and hasattr(model.model, "rotary_emb"): + rotary = model.model.rotary_emb + elif hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"): + rotary = model.transformer.rotary_emb + if rotary is None: + return None + return rotary(hidden_states, position_ids) + + +def _call_layer( + layer: Any, + hidden_states: Any, + attention_mask: Any, + position_ids: Any, + position_embeddings: Any | None = None, +) -> Any: attempts = ( + { + "attention_mask": attention_mask, + "position_ids": position_ids, + "position_embeddings": position_embeddings, + "use_cache": False, + }, { "attention_mask": attention_mask, "position_ids": position_ids, @@ -272,7 +375,7 @@ def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> An def _int_tensor_header(tensor: Any) -> str: - data = tensor.detach().cpu().to(tensor.int64).contiguous() + data = tensor.detach().cpu().long().contiguous() raw = data.numpy().tobytes() shape = ",".join(str(dim) for dim in data.shape) encoded = base64.b64encode(raw).decode("ascii")