code/models/experimental/pi0/reference/torch_prefix.py

pi05-base-p150 / code /models /experimental /pi0 /reference /torch_prefix.py

pi0.5 multi-config megakernel (cameras 1-4, H 1..64, N 1..10, prompt buckets 32-224; 3 programs per call, 4 with 3-4 cameras)

ab447b7 verified about 1 month ago

Raw History Blame Contribute Delete

5.96 kB

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Prefix Embedding module - PyTorch Reference Implementation.
This module handles embedding of images and language tokens to create
the prefix part of the sequence for VLM transformer processing.
Components:
    - Vision tower: SigLIP for image encoding (via embed_image_fn callback)
    - Language embedding: Token embeddings (via embed_language_fn callback)
"""

import math
from typing import List, Tuple

import torch

from models.experimental.pi0.common.configs import PrefixConfig


def safe_cat(tensors: list, dim: int = -1) -> torch.Tensor:
    """Safely concatenate tensors with dtype handling."""
    if len(tensors) == 0:
        raise ValueError("Cannot concatenate empty list of tensors")

    target_dtype = tensors[0].dtype
    converted = [t.to(dtype=target_dtype) if t.dtype != target_dtype else t for t in tensors]
    return torch.cat(converted, dim=dim)


class PrefixEmbedding:
    """
    PyTorch implementation of prefix embedding.
    Combines image and language embeddings for the VLM backbone.
    Uses callback functions for embedding operations.
    """

    def __init__(
        self,
        config: PrefixConfig,
        embed_image_fn=None,
        embed_language_fn=None,
    ):
        """
        Initialize prefix embedding.
        Args:
            config: Prefix configuration
            embed_image_fn: Function to embed images (from SigLIP)
            embed_language_fn: Function to embed language tokens (from Gemma)
        """
        self.config = config
        self.embed_image_fn = embed_image_fn
        self.embed_language_fn = embed_language_fn

    def embed_images(
        self,
        images: List[torch.Tensor],
        img_masks: List[torch.Tensor],
    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
        """
        Embed multiple images.
        Args:
            images: List of image tensors (batch_size, channels, height, width)
            img_masks: List of boolean masks (batch_size,) indicating valid images
        Returns:
            Tuple of (image_embeddings, expanded_masks):
                - image_embeddings: List of (batch_size, num_tokens, hidden_dim)
                - expanded_masks: List of (batch_size, num_tokens)
        """
        if self.embed_image_fn is None:
            raise RuntimeError("embed_image_fn not set")

        image_embs = []
        expanded_masks = []

        for img, mask in zip(images, img_masks):
            # Embed image through vision tower
            img_emb = self.embed_image_fn(img)
            image_embs.append(img_emb)

            # Expand mask to match token count
            batch_size, num_tokens = img_emb.shape[:2]
            expanded_mask = mask[:, None].expand(batch_size, num_tokens)
            expanded_masks.append(expanded_mask)

        return image_embs, expanded_masks

    def embed_language(
        self,
        lang_tokens: torch.Tensor,
        lang_masks: torch.Tensor,
    ) -> torch.Tensor:
        """
        Embed language tokens.
        Args:
            lang_tokens: (batch_size, seq_len) token IDs
            lang_masks: (batch_size, seq_len) boolean validity mask
        Returns:
            (batch_size, seq_len, hidden_dim) scaled embeddings
        """
        if self.embed_language_fn is None:
            raise RuntimeError("embed_language_fn not set")

        lang_emb = self.embed_language_fn(lang_tokens)

        # Scale by sqrt(hidden_dim) as per standard practice
        hidden_dim = lang_emb.shape[-1]
        return lang_emb * math.sqrt(hidden_dim)

    def embed_prefix(
        self,
        images: List[torch.Tensor],
        img_masks: List[torch.Tensor],
        lang_tokens: torch.Tensor,
        lang_masks: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Main embedding function for prefix (images + language).
        Args:
            images: List of image tensors
            img_masks: List of image validity masks
            lang_tokens: Language token IDs
            lang_masks: Language token masks
        Returns:
            Tuple of (prefix_embs, prefix_pad_masks, prefix_att_masks):
                - prefix_embs: (batch_size, prefix_len, hidden_dim)
                - prefix_pad_masks: (batch_size, prefix_len) padding mask
                - prefix_att_masks: (batch_size, prefix_len) attention pattern
                    where 0 = bidirectional attention (prefix behavior)
        """
        embs = []
        pad_masks = []
        att_masks = []

        # Process images
        if images and self.embed_image_fn is not None:
            image_embs, img_pad_masks = self.embed_images(images, img_masks)
            for img_emb, img_pad_mask in zip(image_embs, img_pad_masks):
                embs.append(img_emb)
                pad_masks.append(img_pad_mask)
                # Images use bidirectional attention (0)
                num_img_tokens = img_emb.shape[1]
                att_masks.extend([0] * num_img_tokens)

        # Process language
        if self.embed_language_fn is not None:
            lang_emb = self.embed_language(lang_tokens, lang_masks)
            embs.append(lang_emb)
            pad_masks.append(lang_masks)
            # Language uses bidirectional attention within prefix (0)
            num_lang_tokens = lang_emb.shape[1]
            att_masks.extend([0] * num_lang_tokens)

        # Concatenate all embeddings
        prefix_embs = safe_cat(embs, dim=1)
        prefix_pad_masks = torch.cat(pad_masks, dim=1)

        # Create attention mask tensor
        batch_size = prefix_embs.shape[0]
        device = prefix_embs.device
        att_masks_tensor = torch.tensor(att_masks, dtype=torch.bool, device=device)
        prefix_att_masks = att_masks_tensor.unsqueeze(0).expand(batch_size, -1)

        return prefix_embs, prefix_pad_masks, prefix_att_masks

	# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
	# SPDX-License-Identifier: Apache-2.0

	"""
	Prefix Embedding module - PyTorch Reference Implementation.

	This module handles embedding of images and language tokens to create
	the prefix part of the sequence for VLM transformer processing.

	Components:
	- Vision tower: SigLIP for image encoding (via embed_image_fn callback)
	- Language embedding: Token embeddings (via embed_language_fn callback)
	"""

	import math
	from typing import List, Tuple

	import torch

	from models.experimental.pi0.common.configs import PrefixConfig


	def safe_cat(tensors: list, dim: int = -1) -> torch.Tensor:
	"""Safely concatenate tensors with dtype handling."""
	if len(tensors) == 0:
	raise ValueError("Cannot concatenate empty list of tensors")

	target_dtype = tensors[0].dtype
	converted = [t.to(dtype=target_dtype) if t.dtype != target_dtype else t for t in tensors]
	return torch.cat(converted, dim=dim)


	class PrefixEmbedding:
	"""
	PyTorch implementation of prefix embedding.

	Combines image and language embeddings for the VLM backbone.
	Uses callback functions for embedding operations.
	"""

	def __init__(
	self,
	config: PrefixConfig,
	embed_image_fn=None,
	embed_language_fn=None,
	):
	"""
	Initialize prefix embedding.

	Args:
	config: Prefix configuration
	embed_image_fn: Function to embed images (from SigLIP)
	embed_language_fn: Function to embed language tokens (from Gemma)
	"""
	self.config = config
	self.embed_image_fn = embed_image_fn
	self.embed_language_fn = embed_language_fn

	def embed_images(
	self,
	images: List[torch.Tensor],
	img_masks: List[torch.Tensor],
	) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
	"""
	Embed multiple images.

	Args:
	images: List of image tensors (batch_size, channels, height, width)
	img_masks: List of boolean masks (batch_size,) indicating valid images

	Returns:
	Tuple of (image_embeddings, expanded_masks):
	- image_embeddings: List of (batch_size, num_tokens, hidden_dim)
	- expanded_masks: List of (batch_size, num_tokens)
	"""
	if self.embed_image_fn is None:
	raise RuntimeError("embed_image_fn not set")

	image_embs = []
	expanded_masks = []

	for img, mask in zip(images, img_masks):
	# Embed image through vision tower
	img_emb = self.embed_image_fn(img)
	image_embs.append(img_emb)

	# Expand mask to match token count
	batch_size, num_tokens = img_emb.shape[:2]
	expanded_mask = mask[:, None].expand(batch_size, num_tokens)
	expanded_masks.append(expanded_mask)

	return image_embs, expanded_masks

	def embed_language(
	self,
	lang_tokens: torch.Tensor,
	lang_masks: torch.Tensor,
	) -> torch.Tensor:
	"""
	Embed language tokens.

	Args:
	lang_tokens: (batch_size, seq_len) token IDs
	lang_masks: (batch_size, seq_len) boolean validity mask

	Returns:
	(batch_size, seq_len, hidden_dim) scaled embeddings
	"""
	if self.embed_language_fn is None:
	raise RuntimeError("embed_language_fn not set")

	lang_emb = self.embed_language_fn(lang_tokens)

	# Scale by sqrt(hidden_dim) as per standard practice
	hidden_dim = lang_emb.shape[-1]
	return lang_emb * math.sqrt(hidden_dim)

	def embed_prefix(
	self,
	images: List[torch.Tensor],
	img_masks: List[torch.Tensor],
	lang_tokens: torch.Tensor,
	lang_masks: torch.Tensor,
	) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	"""
	Main embedding function for prefix (images + language).

	Args:
	images: List of image tensors
	img_masks: List of image validity masks
	lang_tokens: Language token IDs
	lang_masks: Language token masks

	Returns:
	Tuple of (prefix_embs, prefix_pad_masks, prefix_att_masks):
	- prefix_embs: (batch_size, prefix_len, hidden_dim)
	- prefix_pad_masks: (batch_size, prefix_len) padding mask
	- prefix_att_masks: (batch_size, prefix_len) attention pattern
	where 0 = bidirectional attention (prefix behavior)
	"""
	embs = []
	pad_masks = []
	att_masks = []

	# Process images
	if images and self.embed_image_fn is not None:
	image_embs, img_pad_masks = self.embed_images(images, img_masks)
	for img_emb, img_pad_mask in zip(image_embs, img_pad_masks):
	embs.append(img_emb)
	pad_masks.append(img_pad_mask)
	# Images use bidirectional attention (0)
	num_img_tokens = img_emb.shape[1]
	att_masks.extend([0] * num_img_tokens)

	# Process language
	if self.embed_language_fn is not None:
	lang_emb = self.embed_language(lang_tokens, lang_masks)
	embs.append(lang_emb)
	pad_masks.append(lang_masks)
	# Language uses bidirectional attention within prefix (0)
	num_lang_tokens = lang_emb.shape[1]
	att_masks.extend([0] * num_lang_tokens)

	# Concatenate all embeddings
	prefix_embs = safe_cat(embs, dim=1)
	prefix_pad_masks = torch.cat(pad_masks, dim=1)

	# Create attention mask tensor
	batch_size = prefix_embs.shape[0]
	device = prefix_embs.device
	att_masks_tensor = torch.tensor(att_masks, dtype=torch.bool, device=device)
	prefix_att_masks = att_masks_tensor.unsqueeze(0).expand(batch_size, -1)

	return prefix_embs, prefix_pad_masks, prefix_att_masks