# --------------------------------------------------------
# InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import os
import warnings
from typing import List, Optional, Tuple, Union
import torch.utils.checkpoint
import torch.nn.functional as F
import transformers
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import GenerationConfig
from transformers.cache_utils import DynamicCache
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging
from transformers import LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM, Qwen3MoeForCausalLM
from .configuration_internvl_chat import InternVLChatConfig
from .conversation import get_conv_template
from .modeling_intern_vit import InternVisionModel, has_flash_attn
logger = logging.get_logger(__name__)
def version_cmp(v1, v2, op='eq'):
import operator
from packaging import version
op_func = getattr(operator, op)
return op_func(version.parse(v1), version.parse(v2))
class HPDRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
class _HPDGatedMLP(nn.Module):
def __init__(self, hidden_size, intermediate_size):
super().__init__()
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
def forward(self, x):
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class _HPDMTPLayer(nn.Module):
def __init__(self, hidden_size, intermediate_size):
super().__init__()
self.mlp = _HPDGatedMLP(hidden_size, intermediate_size)
class HPDMultiTokenPredictor(nn.Module):
"""Single-block Progressive Multi-Token Prediction head (P-MTP).
Mirrors the vLLM ``Qwen3_5MultiTokenPredictor`` module (loaded there as a
``medusa`` speculative model) so its weights, stored under
``language_model.mtp.*`` in the checkpoint, load without remapping. One MTP
block is applied autoregressively to draft several future tokens; each step
is conditioned on the previous step's predicted token. The token embedding
and the output head are shared with the backbone (the checkpoint contains no
``mtp.embed_tokens``/``mtp.lm_head``), so ``step`` takes an already-embedded
``prev_embeds`` tensor.
"""
def __init__(self, hidden_size, intermediate_size, rms_norm_eps=1e-6):
super().__init__()
self.fc = nn.Linear(hidden_size * 2, hidden_size, bias=False)
self.pre_fc_norm_hidden = HPDRMSNorm(hidden_size, rms_norm_eps)
self.pre_fc_norm_embedding = HPDRMSNorm(hidden_size, rms_norm_eps)
self.layers = nn.ModuleList([_HPDMTPLayer(hidden_size, intermediate_size)])
self.norm = HPDRMSNorm(hidden_size, rms_norm_eps)
def step(self, hidden_states, prev_embeds):
# hidden_states: [B, H] backbone hidden that produced prev_token;
# prev_embeds: [B, H] backbone embedding of that token.
embeds = self.pre_fc_norm_embedding(prev_embeds)
hidden = self.pre_fc_norm_hidden(hidden_states)
x = self.fc(torch.cat([hidden, embeds], dim=-1))
mlp_out = self.layers[0].mlp(x)
return self.norm(mlp_out + x)
class InternVLChatModel(PreTrainedModel):
config_class = InternVLChatConfig
main_input_name = 'pixel_values'
base_model_prefix = 'language_model'
_supports_flash_attn_2 = True
supports_gradient_checkpointing = True
_no_split_modules = [
"InternVisionModel",
"Qwen3DecoderLayer",
]
# support transformers 4.51.+
_tp_plan = ''
def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
super().__init__(config)
assert version_cmp(transformers.__version__, '4.37.0', 'ge')
image_size = config.force_image_size or config.vision_config.image_size
patch_size = config.vision_config.patch_size
self.patch_size = patch_size
self.select_layer = config.select_layer
self.template = config.template
self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
self.downsample_ratio = config.downsample_ratio
self.ps_version = config.ps_version
use_flash_attn = use_flash_attn if has_flash_attn else False
config.vision_config.use_flash_attn = True if use_flash_attn else False
config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
logger.info(f'num_image_token: {self.num_image_token}')
logger.info(f'ps_version: {self.ps_version}')
if vision_model is not None:
self.vision_model = vision_model
else:
self.vision_model = InternVisionModel(config.vision_config)
if language_model is not None:
self.language_model = language_model
else:
architecture: str = config.llm_config.architectures[0]
if architecture == 'LlamaForCausalLM':
self.language_model = LlamaForCausalLM(config.llm_config)
elif architecture == 'Qwen2ForCausalLM':
self.language_model = Qwen2ForCausalLM(config.llm_config)
elif architecture == 'Qwen3MoeForCausalLM':
self.language_model = Qwen3MoeForCausalLM(config.llm_config)
elif architecture == 'Qwen3ForCausalLM':
self.language_model = Qwen3ForCausalLM(config.llm_config)
else:
raise NotImplementedError(f'{architecture} is not implemented.')
vit_hidden_size = config.vision_config.hidden_size
llm_hidden_size = config.llm_config.hidden_size
self.mlp1 = nn.Sequential(
nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
nn.GELU(),
nn.Linear(llm_hidden_size, llm_hidden_size)
)
self.img_context_token_id = None
self.conv_template = get_conv_template(self.template)
self.system_message = self.conv_template.system_message
# P-MTP head + hierarchical fork decoding metadata (see generate_hpd).
self.fork_token_id = getattr(config, 'fork_token_id', None)
self.child_token_id = getattr(config, 'child_token_id', None)
self._mtp_weights_loaded = False
llm_cfg = config.llm_config
self.language_model.mtp = HPDMultiTokenPredictor(
hidden_size=llm_cfg.hidden_size,
intermediate_size=llm_cfg.intermediate_size,
rms_norm_eps=getattr(llm_cfg, 'rms_norm_eps', 1e-6),
)
def forward(
self,
pixel_values: torch.FloatTensor,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
image_flags: Optional[torch.LongTensor] = None,
past_key_values: Optional[List[torch.FloatTensor]] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
image_flags = image_flags.squeeze(-1)
input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()
vit_embeds = self.extract_feature(pixel_values)
vit_embeds = vit_embeds[image_flags == 1]
vit_batch_size = pixel_values.shape[0]
B, N, C = input_embeds.shape
input_embeds = input_embeds.reshape(B * N, C)
# if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
# print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
input_ids = input_ids.reshape(B * N)
selected = (input_ids == self.img_context_token_id)
try:
input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
except Exception as e:
vit_embeds = vit_embeds.reshape(-1, C)
print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
f'vit_embeds.shape={vit_embeds.shape}')
n_token = min(selected.sum(), vit_embeds.size(0))
input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
input_embeds = input_embeds.reshape(B, N, C)
outputs = self.language_model(
inputs_embeds=input_embeds,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
logits = outputs.logits
loss = None
if labels is not None:
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
def pixel_shuffle(self, x, scale_factor=0.5):
n, w, h, c = x.size()
# N, W, H, C --> N, W, H * scale, C // scale
x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
# N, W, H * scale, C // scale --> N, H * scale, W, C // scale
x = x.permute(0, 2, 1, 3).contiguous()
# N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
x = x.view(n, int(h * scale_factor), int(w * scale_factor),
int(c / (scale_factor * scale_factor)))
if self.ps_version == 'v1':
warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
'which results in a transposed image.')
else:
x = x.permute(0, 2, 1, 3).contiguous()
return x
def extract_feature(self, pixel_values):
if self.select_layer == -1:
vit_embeds = self.vision_model(
pixel_values=pixel_values,
output_hidden_states=False,
return_dict=True).last_hidden_state
else:
vit_embeds = self.vision_model(
pixel_values=pixel_values,
output_hidden_states=True,
return_dict=True).hidden_states[self.select_layer]
vit_embeds = vit_embeds[:, 1:, :]
h = w = int(vit_embeds.shape[1] ** 0.5)
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
vit_embeds = self.mlp1(vit_embeds)
return vit_embeds
def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,
history=None, return_history=False, IMG_START_TOKEN='![]()
', IMG_END_TOKEN='',
IMG_CONTEXT_TOKEN='', verbose=False, image_counts=None):
if history is not None or return_history:
print('Now multi-turn chat is not supported in batch_chat.')
raise NotImplementedError
if image_counts is not None:
num_patches_list = image_counts
print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')
img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
self.img_context_token_id = img_context_token_id
if verbose and pixel_values is not None:
image_bs = pixel_values.shape[0]
print(f'dynamic ViT batch size: {image_bs}')
queries = []
for idx, num_patches in enumerate(num_patches_list):
question = questions[idx]
if pixel_values is not None and '' not in question:
question = '\n' + question
template = get_conv_template(self.template)
template.system_message = self.system_message
template.append_message(template.roles[0], question)
template.append_message(template.roles[1], None)
query = template.get_prompt()
image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
query = query.replace('', image_tokens, 1)
queries.append(query)
tokenizer.padding_side = 'left'
model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
input_ids = model_inputs['input_ids'].to(self.device)
attention_mask = model_inputs['attention_mask'].to(self.device)
eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
generation_config['eos_token_id'] = eos_token_id
generation_output = self.generate(
pixel_values=pixel_values,
input_ids=input_ids,
attention_mask=attention_mask,
**generation_config
)
responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)
responses = [response.split(template.sep.strip())[0].strip() for response in responses]
return responses
def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,
num_patches_list=None, IMG_START_TOKEN='![]()
', IMG_END_TOKEN='', IMG_CONTEXT_TOKEN='',
verbose=False):
if history is None and pixel_values is not None and '' not in question:
question = '\n' + question
if num_patches_list is None:
num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
assert pixel_values is None or len(pixel_values) == sum(num_patches_list)
img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
self.img_context_token_id = img_context_token_id
template = get_conv_template(self.template)
template.system_message = self.system_message
eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
history = [] if history is None else history
for (old_question, old_answer) in history:
template.append_message(template.roles[0], old_question)
template.append_message(template.roles[1], old_answer)
template.append_message(template.roles[0], question)
template.append_message(template.roles[1], None)
query = template.get_prompt()
if verbose and pixel_values is not None:
image_bs = pixel_values.shape[0]
print(f'dynamic ViT batch size: {image_bs}')
for num_patches in num_patches_list:
image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
query = query.replace('', image_tokens, 1)
model_inputs = tokenizer(query, return_tensors='pt')
input_ids = model_inputs['input_ids'].to(self.device)
attention_mask = model_inputs['attention_mask'].to(self.device)
generation_config['eos_token_id'] = eos_token_id
generation_output = self.generate(
pixel_values=pixel_values,
input_ids=input_ids,
attention_mask=attention_mask,
**generation_config
)
response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
response = response.split(template.sep.strip())[0].strip()
history.append((question, response))
if return_history:
return response, history
else:
query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '')
if verbose:
print(query_to_print, response)
return response
@torch.no_grad()
def generate(
self,
pixel_values: Optional[torch.FloatTensor] = None,
input_ids: Optional[torch.FloatTensor] = None,
attention_mask: Optional[torch.LongTensor] = None,
visual_features: Optional[torch.FloatTensor] = None,
generation_config: Optional[GenerationConfig] = None,
output_hidden_states: Optional[bool] = None,
**generate_kwargs,
) -> torch.LongTensor:
assert self.img_context_token_id is not None
if pixel_values is not None:
if visual_features is not None:
vit_embeds = visual_features
else:
vit_embeds = self.extract_feature(pixel_values)
input_embeds = self.language_model.get_input_embeddings()(input_ids)
B, N, C = input_embeds.shape
input_embeds = input_embeds.reshape(B * N, C)
input_ids = input_ids.reshape(B * N)
selected = (input_ids == self.img_context_token_id)
assert selected.sum() != 0
input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
input_embeds = input_embeds.reshape(B, N, C)
else:
input_embeds = self.language_model.get_input_embeddings()(input_ids)
outputs = self.language_model.generate(
inputs_embeds=input_embeds,
attention_mask=attention_mask,
generation_config=generation_config,
output_hidden_states=output_hidden_states,
use_cache=True,
**generate_kwargs,
)
return outputs
def load_mtp_weights(self, path=None, prefix='language_model.mtp.'):
"""Load P-MTP head weights.
If ``path`` is None the weights are assumed to already be present (e.g.
loaded from the main checkpoint via ``from_pretrained`` because the keys
``language_model.mtp.*`` match this module tree) and this call only marks
MTP as usable. Otherwise ``path`` may be a ``.safetensors`` file or a
directory containing one; keys are stripped of ``prefix`` (or a leading
``mtp.``) before loading.
"""
if path is None:
self._mtp_weights_loaded = True
return
from safetensors.torch import load_file
if os.path.isdir(path):
files = [os.path.join(path, f) for f in sorted(os.listdir(path))
if f.endswith('.safetensors')]
else:
files = [path]
remapped = {}
for fp in files:
for k, v in load_file(fp).items():
name = k
if name.startswith(prefix):
name = name[len(prefix):]
elif name.startswith('mtp.'):
name = name[len('mtp.'):]
else:
continue
remapped[name] = v
if not remapped:
raise RuntimeError(f'No MTP tensors (prefix {prefix!r}) found under {path}')
missing, unexpected = self.language_model.mtp.load_state_dict(remapped, strict=False)
self.language_model.mtp.to(device=self.device, dtype=self.dtype)
self._mtp_weights_loaded = True
return missing, unexpected
def _embed_ids(self, input_ids, vit_embeds=None):
embeds = self.language_model.get_input_embeddings()(input_ids)
if vit_embeds is not None:
B, N, C = embeds.shape
flat = embeds.reshape(B * N, C)
selected = (input_ids.reshape(B * N) == self.img_context_token_id)
flat[selected] = vit_embeds.reshape(-1, C).to(flat.dtype)
embeds = flat.reshape(B, N, C)
return embeds
def _lm_forward(self, inputs_embeds, cache, need_hidden):
past_len = cache.get_seq_length()
seq_len = inputs_embeds.shape[1]
cache_position = torch.arange(past_len, past_len + seq_len, device=inputs_embeds.device)
out = self.language_model(
inputs_embeds=inputs_embeds,
past_key_values=cache,
use_cache=True,
cache_position=cache_position,
position_ids=cache_position.unsqueeze(0),
output_hidden_states=need_hidden,
return_dict=True,
)
hidden = out.hidden_states[-1] if need_hidden else None
return out.logits, hidden, out.past_key_values
@torch.no_grad()
def _decode_children_batched(self, parent_cache, prompt_len, fork_indices,
eos_ids, max_new_tokens):
"""Decode all fork children concurrently as one batch (greedy AR).
Each child ``i`` inherits the parent prefix up to its ```` (real
length ``Li = prompt_len + fork_indices[i]``). The parent KV is copied
into a left-padded batched cache (row ``i`` right-aligned to hold its
``Li`` prefix positions), then all children step in lockstep until every
row hits EOS. Per-row ``position_ids`` carry each child's true positions
so RoPE matches the parent context. Returns a list of token-id lists.
Children run without the P-MTP head (batched ragged acceptance would
desync row lengths); the greedy output is identical to the serial path.
"""
device = self.device
B = len(fork_indices)
Ls = [prompt_len + f for f in fork_indices]
Lmax = max(Ls)
cache = DynamicCache()
for idx, layer in enumerate(parent_cache.layers):
K, V = layer.keys, layer.values # [1, h, Pseq, d]
h, d = K.shape[1], K.shape[3]
Kb = K.new_zeros((B, h, Lmax, d))
Vb = V.new_zeros((B, h, Lmax, d))
for i, Li in enumerate(Ls):
Kb[i, :, Lmax - Li:, :] = K[0, :, :Li, :]
Vb[i, :, Lmax - Li:, :] = V[0, :, :Li, :]
cache.update(Kb, Vb, idx)
mask = torch.zeros((B, Lmax), dtype=torch.long, device=device)
for i, Li in enumerate(Ls):
mask[i, Lmax - Li:] = 1
real_pos = torch.tensor(Ls, device=device) # sits at position Li
cur_embeds = self.language_model.get_input_embeddings()(
torch.tensor([[self.child_token_id]], device=device)).expand(B, 1, -1).contiguous()
outputs = [[] for _ in range(B)]
done = torch.zeros(B, dtype=torch.bool, device=device)
steps = 0
while steps < max_new_tokens and not bool(done.all()):
cache_len = cache.get_seq_length()
full_mask = torch.cat([mask, torch.ones((B, 1), dtype=torch.long, device=device)], dim=1)
out = self.language_model(
inputs_embeds=cur_embeds,
past_key_values=cache,
use_cache=True,
attention_mask=full_mask,
position_ids=real_pos.unsqueeze(1),
cache_position=torch.tensor([cache_len], device=device),
return_dict=True,
)
cache = out.past_key_values
nxt = out.logits[:, -1, :].argmax(-1) # [B]
mask = full_mask
real_pos = real_pos + 1
for i in range(B):
if not done[i]:
tid = int(nxt[i])
outputs[i].append(tid)
if tid in eos_ids or len(outputs[i]) >= max_new_tokens:
done[i] = True
cur_embeds = self.language_model.get_input_embeddings()(nxt.unsqueeze(1))
steps += 1
return outputs
@torch.no_grad()
def _decode_from(self, cache, first_embeds, eos_ids, max_new_tokens, use_mtp, k, allow_fork):
"""Greedily decode one branch (batch size 1) starting from ``cache``.
``first_embeds`` primes the branch (the prompt for the parent, or the
single ```` token embedding for a child that inherits the parent
KV cache). With ``use_mtp`` the P-MTP head drafts ``k`` tokens per step
and the backbone verifies them in a single forward (greedy accept: keep
the longest matching prefix, then one correction/bonus token), rolling
the cache back with ``crop``. Returns ``(real_tokens, fork_indices,
cache)``; ``fork_indices`` marks positions of ```` tokens (only
when ``allow_fork``).
"""
device = first_embeds.device
logits, hidden, cache = self._lm_forward(first_embeds, cache, use_mtp)
q = cache.get_seq_length() - 1 # cache holds positions [0..q]
prod_hidden = hidden[:, -1, :] if use_mtp else None
t = int(logits[0, -1].argmax(-1)) # confirmed token at position q+1
real = []
fork_indices = []
def emit(tok):
real.append(tok)
if allow_fork and self.fork_token_id is not None and tok == self.fork_token_id:
fork_indices.append(len(real) - 1)
return tok in eos_ids
stopped = False
while len(real) < max_new_tokens and not stopped:
if use_mtp:
drafts = []
h = prod_hidden
prev = t
for _ in range(k):
prev_embeds = self.language_model.get_input_embeddings()(
torch.tensor([prev], device=device))
h = self.language_model.mtp.step(h, prev_embeds)
prev = int(self.language_model.lm_head(h).argmax(-1))
drafts.append(prev)
verify_ids = torch.tensor([[t] + drafts], device=device)
vlogits, vhidden, cache = self._lm_forward(self._embed_ids(verify_ids), cache, True)
target = vlogits[0].argmax(-1).tolist() # target[i] = argmax after verify_ids[i]
emitted = []
j = 0
all_match = True
for i in range(k):
emitted.append(target[i])
if drafts[i] != target[i]:
all_match = False
break
j += 1
if all_match:
emitted.append(target[k]) # bonus token
confirmed = [t] + emitted[:-1]
next_t = emitted[-1]
cache.crop((q + 1) + 1 + j) # keep [0..q+1+j]
prod_hidden = vhidden[0, j:j + 1, :]
q = q + 1 + j
t = next_t
for tok in confirmed:
if emit(tok):
stopped = True
break
else:
vlogits, _, cache = self._lm_forward(
self._embed_ids(torch.tensor([[t]], device=device)), cache, False)
if emit(t):
stopped = True
q = q + 1
t = int(vlogits[0, -1].argmax(-1))
if not stopped and len(real) < max_new_tokens:
emit(t)
return real, fork_indices, cache
@torch.no_grad()
def generate_hpd(self, tokenizer, pixel_values, question, generation_config=None,
use_mtp=False, num_speculative_tokens=6, num_patches_list=None,
IMG_START_TOKEN='![]()
', IMG_END_TOKEN='',
IMG_CONTEXT_TOKEN='', return_token_ids=False,
batch_children=False, verbose=False):
"""Hierarchical parallel decoding (````/````) with optional P-MTP.
Reproduces the vLLM decoding paradigm in plain transformers for a single
image / single prompt (batch size 1):
- The parent "layout" branch is decoded greedily; each ```` token
spawns a content branch. The parent keeps the real ```` in its
context and continues.
- Each child branch inherits the parent KV cache (shared-prefix reuse,
like vLLM). By default children are decoded serially, reusing the
parent cache in place (crop to the tokens before each ````);
this is the fastest, lowest-memory option in transformers. Set
``batch_children=True`` to instead decode all children concurrently in
one left-padded batch (``_decode_children_batched``) -- identical
output, but for many short children over a long prefix it is slower
and much more memory-hungry (no paged prefix sharing like vLLM).
- Outputs are spliced: each ```` in the parent stream is rewritten
to ```` and the matching child's tokens are inserted after it,
in fork order.
Use ``use_mtp=True`` only after ``load_mtp_weights(...)`` has been called.
"""
if self.fork_token_id is None or self.child_token_id is None:
raise RuntimeError('fork_token_id/child_token_id missing from config; cannot fork-decode.')
if use_mtp and not self._mtp_weights_loaded:
raise RuntimeError('P-MTP weights not loaded. Call model.load_mtp_weights(path) '
'or set use_mtp=False.')
generation_config = generation_config or {}
max_new_tokens = generation_config.get('max_new_tokens', 8000)
k = num_speculative_tokens
if pixel_values is not None and '' not in question:
question = '\n' + question
if num_patches_list is None:
num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
self.img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
template = get_conv_template(self.template)
template.system_message = self.system_message
template.append_message(template.roles[0], question)
template.append_message(template.roles[1], None)
query = template.get_prompt()
for num_patches in num_patches_list:
image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
query = query.replace('', image_tokens, 1)
prompt_ids = tokenizer(query, return_tensors='pt')['input_ids'].to(self.device)
prompt_len = prompt_ids.shape[1]
eos_ids = set()
sep_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
if isinstance(sep_id, int) and sep_id >= 0:
eos_ids.add(sep_id)
if self.config.eos_token_id is not None:
eos_ids.add(self.config.eos_token_id)
vit_embeds = self.extract_feature(pixel_values) if pixel_values is not None else None
prompt_embeds = self._embed_ids(prompt_ids, vit_embeds)
parent_real, fork_indices, parent_cache = self._decode_from(
DynamicCache(), prompt_embeds, eos_ids, max_new_tokens, use_mtp, k, allow_fork=True)
# Once the parent finishes, every fork point and its shared prefix is
# fixed. Children can then be decoded either serially (reusing the parent
# cache in place, low memory) or concurrently in one padded batch
# (``batch_children=True``). NOTE: for many short children over a long
# shared prefix, batching is slower and far more memory-hungry in stock
# transformers (the prefix KV is duplicated per row and short children
# still step until the longest one finishes), so serial is the default.
children = []
if fork_indices and batch_children:
children = self._decode_children_batched(
parent_cache, prompt_len, fork_indices, eos_ids, max_new_tokens)
elif fork_indices:
child_embeds = self.language_model.get_input_embeddings()(
torch.tensor([[self.child_token_id]], device=self.device))
children = [None] * len(fork_indices)
# Reverse fork order: crop the shared parent cache to the tokens
# before each and reuse it in place. A child only appends
# beyond its crop point, which the next (smaller) crop discards, so
# the parent prefix KV stays intact and no clone is needed.
for i in range(len(fork_indices) - 1, -1, -1):
parent_cache.crop(prompt_len + fork_indices[i])
child_real, _, _ = self._decode_from(
parent_cache, child_embeds, eos_ids, max_new_tokens, use_mtp, k, allow_fork=False)
children[i] = child_real
final_ids = []
child_ptr = 0
for tok in parent_real:
if tok == self.fork_token_id:
final_ids.append(self.child_token_id)
if child_ptr < len(children):
final_ids.extend(children[child_ptr])
child_ptr += 1
else:
final_ids.append(tok)
text = tokenizer.decode(final_ids, skip_special_tokens=True)
if verbose:
print(f'[HPD] parent tokens={len(parent_real)}, forks={len(fork_indices)}, '
f'children_tokens={[len(c) for c in children]}')
if return_token_ids:
return text, final_ids
return text
@property
def lm_head(self):
return self.language_model.get_output_embeddings()
def get_output_embeddings(self):
return self.language_model.get_output_embeddings()
def get_input_embeddings(self):
return self.language_model.get_input_embeddings()
def set_input_embeddings(self, value):
return self.language_model.set_input_embeddings(value)
def set_output_embeddings(self, value):
return self.language_model.set_output_embeddings(value)