# Copyright (c) Meta Platforms, Inc. and affiliates. import argparse import logging import os import time from dataclasses import dataclass, field from typing import List, Optional logging.basicConfig(level=logging.INFO) import torch from huggingface_hub import snapshot_download from omegaconf import OmegaConf from PIL import Image from torch import nn from torch.nn import functional as F from torch.nn.attention.flex_attention import create_block_mask from tqdm import tqdm from apps.plm.tokenizer import PLMTokenizer, Tokenizer, build_tokenizer from apps.plm.transformer import LMTransformer, LMTransformerArgs from core.args import dataclass_from_dict from core.checkpoint import load_consolidated_checkpoint from core.transformer import (Attention, causal_mask, generate_doc_mask_mod, lengths_to_local_ids, lengths_to_start_ids) from core.transforms.image_transform import get_image_transform from core.transforms.video_transform import get_video_transform def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor: probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) probs_sum = torch.cumsum(probs_sort, dim=-1) mask = probs_sum - probs_sort > p probs_sort[mask] = 0.0 next_token = torch.multinomial(probs_sort, num_samples=1) next_token = torch.gather(probs_idx, -1, next_token) return next_token def sample_top_k(probs, k): topk_value, _ = torch.topk(probs, k) # batch_sz x topk min_value_top_k = topk_value[:, [-1]] probs[probs < min_value_top_k] = 0.0 probs.div_(probs.sum(dim=-1, keepdim=True)) next_token = torch.multinomial(probs, num_samples=1) return next_token def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None): shape = logits.shape logits = logits.flatten(end_dim=-2) if temperature > 0.0: probs = torch.softmax(logits / temperature, dim=-1) if top_p is not None: next_token = sample_top_p(probs, top_p) elif top_k is not None: next_token = sample_top_k(probs, top_k) else: next_token = torch.multinomial(probs, num_samples=1) else: next_token = torch.argmax(logits, dim=-1) return next_token.view(shape[:-1]) def pack_prompts(prompts: List[int]): res = [] lengths = [] for i, p in enumerate(prompts): p = torch.tensor(p, dtype=torch.long) l = p.size(0) res.append(p) lengths.append(l) lengths = torch.tensor(lengths, dtype=torch.long) res = torch.cat(res) return res, lengths def batch_prompts(prompts, max_elements, lengths=None): batches = [] current_batch = [] current_count = 0 for i in range(len(prompts)): prt = prompts[i] prompt_size = len(prt) if lengths is None else lengths[i] if current_count + prompt_size <= max_elements: current_batch.append(prt) current_count += prompt_size else: if current_batch: # Add the current batch to batches batches.append(current_batch) # Start a new batch with the current prompt current_batch = [prt] current_count = prompt_size # Add the last batch if it contains any prompts if current_batch: batches.append(current_batch) return batches class KVCache(nn.Module): def __init__(self, bsz, seqlen, n_heads, head_dim, dtype, device): super().__init__() shape = (bsz, seqlen, n_heads, head_dim) self.register_buffer("k_cache", torch.zeros(shape, dtype=dtype, device=device)) self.register_buffer("v_cache", torch.zeros(shape, dtype=dtype, device=device)) self.offset = 0 def reset(self): self.k_cache.zero_() self.v_cache.zero_() self.offset = 0 def update(self, k_val, v_val, tok_idx): # input_pos: [B], k_val: [B, S, H, D] self.k_cache.index_copy_(1, self.offset + tok_idx, k_val) self.v_cache.index_copy_(1, self.offset + tok_idx, v_val) return self.k_cache, self.v_cache @dataclass class PackedCausalTransformerGeneratorArgs: temperature: float = 0.0 top_p: Optional[float] = None top_k: Optional[float] = None max_gen_len: int = 256 max_tokens: int = 11264 until: List[str] = field(default_factory=list) compile_prefilling: bool = False reduce_generation_overhead: bool = False show_progress: bool = False dtype: Optional[str] = "bf16" device: Optional[str] = "cuda" class PackedCausalTransformerGenerator: def __init__( self, cfg: PackedCausalTransformerGeneratorArgs, model: nn.Module, tokenizer: Tokenizer, ): """ This class wraps a causal transformer model with its corresponding tokenizer and provides an efficient way to pack prompts together and do generation on the packed sequence. For example, if we had the prompts "Hello, I am a " and "Initiating calibration " Then this class will concatenate those sequence (pack them together) "Hello, I am a Initiating calibration" And make the necessary attention masks such that a sequence only attends to itself during prefilling and generation. This class creates a fixed size cache of size max_tokens or sum of prompt sizes + the max number of generated tokens per sequence. """ self.model = model self.tokenizer = tokenizer self.temperature = cfg.temperature self.top_p = cfg.top_p self.top_k = cfg.top_k self.max_gen_len = cfg.max_gen_len self.max_tokens = cfg.max_tokens self.until = cfg.until self.max_until_size = max([len(e) for e in self.until]) if self.until else 1 self.device = cfg.device # Compile if necessary self.prefill = torch.compile(self.prefill, disable=not cfg.compile_prefilling) self.generate_next_token = torch.compile( self.generate_next_token, backend="inductor", fullgraph=True, mode="reduce-overhead", # Other availabel mode is "max-autotune" disable=not cfg.reduce_generation_overhead, ) self.show_progress = cfg.show_progress self.dtype = dict(fp32=torch.float32, bf16=torch.bfloat16)[cfg.dtype] self.prefill_doc_id, self.prefill_tok_id = None, None self.padded_doc_id, self.padded_tok_id = None, None self.current_doc_id, self.current_tok_id = None, None self.padded_doc_start = None self.prefill_mask = None def clear_cache(self, offset): for module in self.model.modules(): if isinstance(module, Attention): if not hasattr(module, "kv_cache"): module.kv_cache = KVCache( 1, self.max_tokens, module.n_kv_heads, module.head_dim, self.dtype, self.device, ) module.kv_cache.offset = offset @torch.compiler.disable def setup_prefilling(self, lengths: torch.Tensor): # The KV cache is a fixed size tensor of size max_tokens that we need # to update in order to do correct autoregressive generation. # Here we will generate token by token but on multiple sequences # at once. To do so, we need to have an attention mask that makes # each sequence independent. # Each sequence will write to its allocated space in the KV Cache. # We allocate len(seq) + max_gen_len to each sequence in the cache. # We will generate max_gen_len for each document padded_lengths = lengths + self.max_gen_len max_tokens = self.max_tokens or padded_lengths.sum().item() # The last document might have more padding to fill up to max_tokens padded_lengths[-1] += max_tokens - padded_lengths.sum() # This is the start index in the cache for each document self.padded_doc_start = lengths_to_start_ids(padded_lengths) # For example with ab--123--cdef-- # this would be 0, 4, 9 if max_gen_len is 2 # We repeat interleave to align with tokens for prefilling # Ex: ab--123--cdef-- # 000044444999999 prefill_offset = torch.repeat_interleave(self.padded_doc_start, lengths) # This offset will make sure the tokens are written to the # correct positions in the cache during prefilling # We either init the cache or clear it by resetting the offset to prefill_offset self.clear_cache(prefill_offset) # The prefilling mask looks like the following for # the two packed sequences ab and 123 : ab123 # Where spaces are empty cache positions # keys # ab---123--- # queries a 10000000000 # b 11000000000 # 1 00000100000 # 2 00000110000 # 3 00000111000 # We make sure to skip the empty cache positions # and only attend to positions within the same sequence doc_mask_mod = generate_doc_mask_mod(causal_mask, lengths, padded_lengths) self.prefill_mask = create_block_mask( doc_mask_mod, 1, None, lengths.sum(), max_tokens ) # This creates the prefilling token ids which look like # the following for the packed sequence abcdefg1234 # abcdefg1234 # 01234560123 # The token id gives us the position within each sequence # This is used to compute ROPE and to update the cache # At each forward pass the current tokens are written to # offset + tok_id self.prefill_doc_id, self.prefill_tok_id = lengths_to_local_ids(lengths) # This creates the padded token and document ids # which look like the following for the packed sequence ab123 # ab---123--- ab---123--- # padded_doc_id 00000111111 padded_tok_id 01234012345 # This will later be useful for the attention mask at generation self.padded_doc_id, self.padded_tok_id = lengths_to_local_ids(padded_lengths) @torch.compiler.disable def setup_generation(self, lengths): # KV Cache offset is set to the start of the padded documents for module in self.model.modules(): if isinstance(module, Attention): module.kv_cache.offset = self.padded_doc_start # The token ids during generations correspond to the lengths of each doc # current_tok_id will be incremented during generation self.current_tok_id = lengths.clone() # Since we're generating one token per document # the document id is just an arange self.current_doc_id = torch.arange(lengths.size(0), device=lengths.device) @torch.compiler.disable def prepare_media_inputs( self, tokens, lengths, images, image_patch_text_ids, num_image_chunks ): image_pos_index = None num_chunks = [] if images is not None and len(images) > 0: assert image_patch_text_ids is not None and len( image_patch_text_ids ) == len(images) assert num_image_chunks is not None and len(num_image_chunks) == len(images) image_pos_index = torch.full(tokens.shape, -1, dtype=torch.int).to( self.device ) assert tokens.shape[0] == 1 offsets = torch.roll(lengths.cpu(), shifts=1, dims=-1).numpy() offsets[0] = 0 num_chunks_seq = 0 image_id_offset = 0 for image_id, offset in enumerate(offsets): num_image_tokens = len(image_patch_text_ids[image_id]) image_indices = ( torch.arange(num_image_tokens, dtype=torch.int).to(self.device) + image_id_offset ) text_indices = [i + offset for i in image_patch_text_ids[image_id]] image_pos_index[0, text_indices] = image_indices image_id_offset += num_image_tokens num_chunks_seq += num_image_chunks[image_id] num_chunks.append(num_chunks_seq) # Move images to the same device and dtype as model parameters model_param = next(self.model.parameters()) images = torch.cat(images).to(model_param) else: images = None return images, image_pos_index, num_chunks # From here on some methods for generation def prefill( self, tokens: torch.Tensor, lengths: torch.Tensor, images: Optional[List[torch.Tensor]] = None, image_patch_text_ids: Optional[List[List[int]]] = None, num_image_chunks: Optional[List[int]] = None, ): # Prefilling is done by taking multiple packed sequences and # doing block diagonal attention on them so they remain independent self.setup_prefilling(lengths=lengths) images, image_pos_index, num_chunks = self.prepare_media_inputs( tokens, lengths, images, image_patch_text_ids, num_image_chunks ) prefill_out = self.model.forward( tokens, tok_idx=self.prefill_tok_id, mask="causal", images=images, image_pos_index=image_pos_index, num_chunks=num_chunks, attn_impl="sdpa", ) self.setup_generation(lengths=lengths) return prefill_out def generate_next_token(self, current_token): # Since we're doing generation with multiple sequences at once # we need to ignore tokens and cache entries from other sequences # or in the future. # Example mask : # keys # abc--1234-- # queries c 11100000000 # 4 00000111100 # mask shape : (n_seqs, cache_size) doc_mask = self.current_doc_id.unsqueeze(1) == self.padded_doc_id.unsqueeze(0) caus_mask = self.current_tok_id.unsqueeze(1) >= self.padded_tok_id.unsqueeze(0) mask = doc_mask & caus_mask out = self.model.forward( current_token, tok_idx=self.current_tok_id, # n_seqs mask=mask, attn_impl="sdpa", ) self.current_tok_id += 1 return out @torch.inference_mode() def generate(self, prompts): images = [] image_patch_text_ids = [] num_image_chunks = [] # Tokenize if isinstance(self.tokenizer, PLMTokenizer): encoded_prompts = [] for p in prompts: assert isinstance(p, tuple) assert len(p) == 2 question, image = p images.append(image) text_ids, image_pos = self.tokenizer._tokenize_for_generation( question, image ) num_chunks = image.size(0) encoded_prompts.append(text_ids) image_patch_text_ids.append(image_pos) num_image_chunks.append(num_chunks) prompts = encoded_prompts else: prompts = [ self.tokenizer.encode(p, add_bos=False, add_eos=False) for p in prompts ] # Account for the generation in lengths padded_lengths = [len(p) + self.max_gen_len for p in prompts] generation = [] loglikelihood = [] greedy = [] it = batch_prompts(prompts, self.max_tokens, lengths=padded_lengths) if self.show_progress: it = tqdm(it) for batch in it: n_seqs = len(batch) generated_tokens = [[] for _ in range(n_seqs)] is_done = [False for _ in range(n_seqs)] packed_batch, lengths = pack_prompts(batch) packed_batch, lengths = packed_batch.cuda(), lengths.cuda() n_seqs = lengths.size(0) current_images = images[:n_seqs] current_image_patch_text_ids = image_patch_text_ids[:n_seqs] current_num_image_chunks = num_image_chunks[:n_seqs] images = images[n_seqs:] image_patch_text_ids = image_patch_text_ids[n_seqs:] num_image_chunks = num_image_chunks[n_seqs:] # Prefilling cache prompt_logits = self.prefill( packed_batch.unsqueeze(0), lengths, images=current_images, image_patch_text_ids=current_image_patch_text_ids, num_image_chunks=current_num_image_chunks, ) # Selecting last token in each prompt all_tokens = sample_tokens( prompt_logits, self.temperature, self.top_p, self.top_k ) start_token = all_tokens[:, lengths.cumsum(0) - 1] for seq_id, tok in enumerate(start_token.squeeze(0).tolist()): generated_tokens[seq_id].append(tok) current_token = start_token for i in range(1, self.max_gen_len): next_logits = self.generate_next_token(current_token) next_token = sample_tokens( next_logits.clone(), self.temperature, self.top_p, self.top_k ) for seq_id, tok in enumerate(next_token.squeeze(0).tolist()): if not is_done[seq_id]: generated_tokens[seq_id].append(tok) current_end_str = self.tokenizer.decode( generated_tokens[seq_id][-self.max_until_size :] ) contains_end_string = any( [e in current_end_str for e in self.until] ) is_done[seq_id] = ( contains_end_string or tok == self.tokenizer.eot_id or tok == self.tokenizer.eos_id ) if all(is_done): break current_token = next_token generation.extend([self.tokenizer.decode(g) for g in generated_tokens]) for p, logit in zip( batch, prompt_logits.squeeze(0).split(lengths.tolist()) ): x = logit[:-1] y = torch.tensor(p[1:], device=x.device) loglikelihood.append(-F.cross_entropy(x, y, reduction="none").cpu()) greedy.append((x.argmax(dim=-1) == y).cpu()) generation = [ response.replace("<|eot_id|>", "").replace("<|end_of_text|>", "") for response in generation ] return generation, loglikelihood, greedy def load_consolidated_model_and_tokenizer(ckpt): # Download the model from Hugging Face if not available locally. if os.path.exists(ckpt): ckpt_path = ckpt # It's a local path else: try: print(f"Downloading {ckpt} from Hugging Face Hub...") ckpt_path = snapshot_download(ckpt) ckpt_path = os.path.join(ckpt_path, "original") print(f"Downloaded to: {ckpt_path}") except Exception as e: # Handle exceptions, such as model not found on HF Hub print(f"An error occurred while downloading {ckpt}: {e}") return # Load params from model config config = os.path.join(ckpt_path, "params.json") config = OmegaConf.load(config) # Build tokenizer tokenizer = build_tokenizer( config.data.tokenizer_name, ( config.data.tokenizer_path if os.path.exists(config.data.tokenizer_path) else os.path.join(ckpt_path, config.data.tokenizer_path) ), pooling_ratio=config.model.pooling_ratio, patch_size=config.model.vision_model.patch_size, ) # Build model and load the consolidated checkpoints model_args = dataclass_from_dict(LMTransformerArgs, config.model, strict=False) model = LMTransformer(model_args) load_consolidated_checkpoint(model, ckpt_path) param_dtype = dict(fp32=torch.float32, fp16=torch.float16, bf16=torch.bfloat16)[ config.distributed.model_dtype ] model = model.cuda().eval() for param in model.parameters(): param.data = param.data.to(dtype=param_dtype) return model, tokenizer, config def main(args): # Load model and tokenizer model, tokenizer, config = load_consolidated_model_and_tokenizer(args.ckpt) media_type = args.media_type media_path = args.media_path question = args.question prompts = [] if media_type == "image": transform = get_image_transform( vision_input_type=config.data.vision_input_type, image_res=model.vision_model.image_size, max_num_tiles=config.data.max_num_tiles, ) image = Image.open(media_path).convert("RGB") image, _ = transform(image) prompts.append((question, image)) elif media_type == "video": transform = get_video_transform( image_res=model.vision_model.image_size, ) video_info = (media_path, config.data.max_video_frames, None, None, None) frames, _ = transform(video_info) prompts.append((question, frames)) else: raise NotImplementedError( f"The provided generate function only supports image and video." ) # Create generator gen_cfg = dataclass_from_dict( PackedCausalTransformerGeneratorArgs, {}, strict=False ) generator = PackedCausalTransformerGenerator(gen_cfg, model, tokenizer) # Run generation start_time = time.time() generation, loglikelihood, greedy = generator.generate(prompts) end_time = time.time() for i, gen in enumerate(generation): # Calculate tokens per second total_tokens = sum( len(tokenizer.encode(gen, False, False)) for gen in generation ) tokens_per_second = total_tokens / (end_time - start_time) print("==============================================") print(f"\nPrompt {i+1}: {prompts[i][0]}") print(f"Generated Text: {gen}") print(f"Tokens per second: {tokens_per_second:.2f}") print("==============================================") def parse_args(): parser = argparse.ArgumentParser(description="Run a multimodal language model") # Define CLI arguments that will be used to override YAML config parser.add_argument( "--ckpt", type=str, required=True, help="Path to the checkpoint directory." ) parser.add_argument( "--media_type", type=str, choices=["image", "video"], required=True, help="Type of media (image or video)", ) parser.add_argument( "--media_path", type=str, required=True, help="Path to the media file" ) parser.add_argument( "--question", type=str, required=True, help="Question or prompt for the model." ) return parser.parse_args() if __name__ == "__main__": args = parse_args() main(args)