SLIP-Llama / multimodal_gemma.py

Add SLIP-Llama (Llama-3.2-1B backbone) base checkpoint + model card

d5b7daa verified 5 months ago

6.36 kB

"""Gemma3 decoder with SLIP cross-attention blocks at layers >= split_layer.
Refactored from HealthSensorSLM-Bench/model_factory/multimodal_gemma.py. The
key change is how sensor context reaches the cross-attention blocks: instead
of caching ``self.vis_x`` as module state (which is unsafe when the same
decoder is called multiple times per step with different micro-batch shapes,
as CODI does), each call threads ``vis_x`` through a ``contextvars.ContextVar``
scoped by a context manager. Wrapped decoder layers read the ContextVar on
every forward — so the context is always fresh and batch-size-correct.
The ``.to(torch.bfloat16)`` hard-cast from the original is also removed; dtype
is taken from the surrounding model so autocast / fp32 runs work unchanged.
"""

import os
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Optional

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModelForCausalLM

from .ts_transformer import CrossAttention


_SLIP_VIS_X: "ContextVar[Optional[torch.Tensor]]" = ContextVar("slip_vis_x", default=None)


@contextmanager
def slip_vis_x_ctx(vis_x: Optional[torch.Tensor]):
    """Bind ``vis_x`` as the cross-attention context for the enclosed forward.
    Safe to nest; the outer binding is restored on exit. Reset on exceptions.
    """
    token = _SLIP_VIS_X.set(vis_x)
    try:
        yield
    finally:
        _SLIP_VIS_X.reset(token)


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, *args, **kwargs):
        return self.fn(x, *args, **kwargs) + x


class Gemma3MultimodalLayer(nn.Module):
    """Wraps a Gemma3 decoder layer with a post-hoc cross-attention block.
    Reads the current sensor context from the module-level ``_SLIP_VIS_X``
    ContextVar rather than from ``self``, so the context stays accurate across
    consecutive forwards with changing batch shapes.
    """

    _xattn_debug_calls: int = 0
    _XATTN_DEBUG_MAX_CALLS: int = 60

    def __init__(self, original_layer, cross_attn_block):
        super().__init__()
        self.original_layer = original_layer
        self.cross_attn_block = cross_attn_block

    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            try:
                original_layer = super().__getattr__("original_layer")
            except AttributeError:
                raise AttributeError(name)
            return getattr(original_layer, name)

    def forward(self, hidden_states, **kwargs):
        vis_x = _SLIP_VIS_X.get()
        assert vis_x is not None, (
            "SLIP cross-attention invoked without a sensor context. "
            "Wrap the model forward in `slip_vis_x_ctx(sensor_hidden)` first."
        )

        outputs = self.original_layer(hidden_states, **kwargs)
        if isinstance(outputs, tuple):
            hidden_states = outputs[0]
        else:
            hidden_states = outputs

        vis_x_cast = vis_x.to(dtype=hidden_states.dtype, device=hidden_states.device)

        cls = Gemma3MultimodalLayer
        _debug = (
            os.environ.get("TLR_TS_DEBUG") == "1"
            and cls._xattn_debug_calls < cls._XATTN_DEBUG_MAX_CALLS
        )
        if _debug:
            pre = hidden_states
            post = self.cross_attn_block(pre, context=vis_x_cast)
            pre_f = pre.detach().float()
            post_f = post.detach().float()
            base = pre_f.norm().clamp_min(1e-6)
            delta = (post_f - pre_f).norm()
            print(
                f"[XATTN] call={cls._xattn_debug_calls} "
                f"vis_x_norm={vis_x_cast.detach().float().norm().item():.3f} "
                f"pre_norm={pre_f.norm().item():.3f} "
                f"post_norm={post_f.norm().item():.3f} "
                f"delta_over_pre={(delta / base).item():.4e}",
                flush=True,
            )
            cls._xattn_debug_calls += 1
            hidden_states = post
        else:
            hidden_states = self.cross_attn_block(hidden_states, context=vis_x_cast)

        if isinstance(outputs, tuple):
            return (hidden_states,) + outputs[1:]
        return hidden_states


class Gemma3MultimodalModel(nn.Module):
    """Standard HF Gemma3 with SLIP cross-attention inserted at/after ``split_layer``."""

    def __init__(
        self,
        model_id: str = "google/gemma-3-270m",
        post_train: bool = False,
        split_layer: int = 12,
        num_heads: Optional[int] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        super().__init__()
        load_kwargs = dict(attn_implementation="sdpa", trust_remote_code=True)
        if dtype is not None:
            load_kwargs["dtype"] = dtype

        # post_train=False is the right default when loading SLIP via
        # SlipForCausalLM.from_pretrained(...): the SLIP safetensors contains the
        # full post-trained Gemma weights already, so downloading the vanilla
        # google/gemma-3-270m checkpoint just to overwrite it is wasted work (and
        # produces a misleading "lm_head.weight newly initialized" warning —
        # Gemma-3 ties lm_head to embed_tokens, so it's never in the checkpoint).
        # Flip to True only when you want a randomly-initialised cross-attention
        # head on top of actually-pretrained Gemma weights, not overwritten by an
        # outer load.
        if post_train:
            self.model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
        else:
            config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
            self.model = AutoModelForCausalLM.from_config(config, **load_kwargs)

        self.split_layer = split_layer
        hidden_size = self.model.config.hidden_size
        self.hidden_size = hidden_size
        if num_heads is None:
            num_heads = self.model.config.num_attention_heads

        for i in range(split_layer, len(self.model.model.layers)):
            cross_attn = CrossAttention(
                dim=hidden_size, context_dim=hidden_size, num_heads=num_heads, dropout_rate=0.1
            )
            original_layer = self.model.model.layers[i]
            self.model.model.layers[i] = Gemma3MultimodalLayer(original_layer, Residual(cross_attn))

	"""Gemma3 decoder with SLIP cross-attention blocks at layers >= split_layer.

	Refactored from HealthSensorSLM-Bench/model_factory/multimodal_gemma.py. The
	key change is how sensor context reaches the cross-attention blocks: instead
	of caching ``self.vis_x`` as module state (which is unsafe when the same
	decoder is called multiple times per step with different micro-batch shapes,
	as CODI does), each call threads ``vis_x`` through a ``contextvars.ContextVar``
	scoped by a context manager. Wrapped decoder layers read the ContextVar on
	every forward — so the context is always fresh and batch-size-correct.

	The ``.to(torch.bfloat16)`` hard-cast from the original is also removed; dtype
	is taken from the surrounding model so autocast / fp32 runs work unchanged.
	"""

	import os
	from contextlib import contextmanager
	from contextvars import ContextVar
	from typing import Optional

	import torch
	import torch.nn as nn
	from transformers import AutoConfig, AutoModelForCausalLM

	from .ts_transformer import CrossAttention


	_SLIP_VIS_X: "ContextVar[Optional[torch.Tensor]]" = ContextVar("slip_vis_x", default=None)


	@contextmanager
	def slip_vis_x_ctx(vis_x: Optional[torch.Tensor]):
	"""Bind ``vis_x`` as the cross-attention context for the enclosed forward.

	Safe to nest; the outer binding is restored on exit. Reset on exceptions.
	"""
	token = _SLIP_VIS_X.set(vis_x)
	try:
	yield
	finally:
	_SLIP_VIS_X.reset(token)


	class Residual(nn.Module):
	def __init__(self, fn):
	super().__init__()
	self.fn = fn

	def forward(self, x, args, *kwargs):
	return self.fn(x, args, *kwargs) + x


	class Gemma3MultimodalLayer(nn.Module):
	"""Wraps a Gemma3 decoder layer with a post-hoc cross-attention block.

	Reads the current sensor context from the module-level ``_SLIP_VIS_X``
	ContextVar rather than from ``self``, so the context stays accurate across
	consecutive forwards with changing batch shapes.
	"""

	_xattn_debug_calls: int = 0
	_XATTN_DEBUG_MAX_CALLS: int = 60

	def __init__(self, original_layer, cross_attn_block):
	super().__init__()
	self.original_layer = original_layer
	self.cross_attn_block = cross_attn_block

	def __getattr__(self, name):
	try:
	return super().__getattr__(name)
	except AttributeError:
	try:
	original_layer = super().__getattr__("original_layer")
	except AttributeError:
	raise AttributeError(name)
	return getattr(original_layer, name)

	def forward(self, hidden_states, **kwargs):
	vis_x = _SLIP_VIS_X.get()
	assert vis_x is not None, (
	"SLIP cross-attention invoked without a sensor context. "
	"Wrap the model forward in `slip_vis_x_ctx(sensor_hidden)` first."
	)

	outputs = self.original_layer(hidden_states, **kwargs)
	if isinstance(outputs, tuple):
	hidden_states = outputs[0]
	else:
	hidden_states = outputs

	vis_x_cast = vis_x.to(dtype=hidden_states.dtype, device=hidden_states.device)

	cls = Gemma3MultimodalLayer
	_debug = (
	os.environ.get("TLR_TS_DEBUG") == "1"
	and cls._xattn_debug_calls < cls._XATTN_DEBUG_MAX_CALLS
	)
	if _debug:
	pre = hidden_states
	post = self.cross_attn_block(pre, context=vis_x_cast)
	pre_f = pre.detach().float()
	post_f = post.detach().float()
	base = pre_f.norm().clamp_min(1e-6)
	delta = (post_f - pre_f).norm()
	print(
	f"[XATTN] call={cls._xattn_debug_calls} "
	f"vis_x_norm={vis_x_cast.detach().float().norm().item():.3f} "
	f"pre_norm={pre_f.norm().item():.3f} "
	f"post_norm={post_f.norm().item():.3f} "
	f"delta_over_pre={(delta / base).item():.4e}",
	flush=True,
	)
	cls._xattn_debug_calls += 1
	hidden_states = post
	else:
	hidden_states = self.cross_attn_block(hidden_states, context=vis_x_cast)

	if isinstance(outputs, tuple):
	return (hidden_states,) + outputs[1:]
	return hidden_states


	class Gemma3MultimodalModel(nn.Module):
	"""Standard HF Gemma3 with SLIP cross-attention inserted at/after ``split_layer``."""

	def __init__(
	self,
	model_id: str = "google/gemma-3-270m",
	post_train: bool = False,
	split_layer: int = 12,
	num_heads: Optional[int] = None,
	dtype: Optional[torch.dtype] = None,
	):
	super().__init__()
	load_kwargs = dict(attn_implementation="sdpa", trust_remote_code=True)
	if dtype is not None:
	load_kwargs["dtype"] = dtype

	# post_train=False is the right default when loading SLIP via
	# SlipForCausalLM.from_pretrained(...): the SLIP safetensors contains the
	# full post-trained Gemma weights already, so downloading the vanilla
	# google/gemma-3-270m checkpoint just to overwrite it is wasted work (and
	# produces a misleading "lm_head.weight newly initialized" warning —
	# Gemma-3 ties lm_head to embed_tokens, so it's never in the checkpoint).
	# Flip to True only when you want a randomly-initialised cross-attention
	# head on top of actually-pretrained Gemma weights, not overwritten by an
	# outer load.
	if post_train:
	self.model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs)
	else:
	config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
	self.model = AutoModelForCausalLM.from_config(config, **load_kwargs)

	self.split_layer = split_layer
	hidden_size = self.model.config.hidden_size
	self.hidden_size = hidden_size
	if num_heads is None:
	num_heads = self.model.config.num_attention_heads

	for i in range(split_layer, len(self.model.model.layers)):
	cross_attn = CrossAttention(
	dim=hidden_size, context_dim=hidden_size, num_heads=num_heads, dropout_rate=0.1
	)
	original_layer = self.model.model.layers[i]
	self.model.model.layers[i] = Gemma3MultimodalLayer(original_layer, Residual(cross_attn))