fix: shard_end convention — inclusive (0-based) not exclusive
model_backend.py was using Python-style exclusive end (layers[start:end]) while all callers (CLI, tests, QUICKSTART) use inclusive 0-based indexing. Result: 24-layer model with shard_end=23 ran only 23 layers and never set is_tail=True, so decode_tail() was never called and responses were empty. - is_tail: == total_layers → >= total_layers - 1 - _run_layers: layers[start:end] → layers[start:end+1] - Validation: > total_layers → >= total_layers (was also wrong) Inference confirmed: Qwen2.5-0.5B-Instruct now returns real LLM output. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user