diff --git "a/modeling_molmo.py" "b/modeling_molmo.py" --- "a/modeling_molmo.py" +++ "b/modeling_molmo.py" @@ -1,84 +1,25 @@ -""" -Adapted from -[MosaiclML](https://github.com/mosaicml/examples.git) and -[minGPT](https://github.com/karpathy/minGPT.git) -""" - -from __future__ import annotations - import logging import math -import sys -from abc import abstractmethod -from pathlib import Path -from typing import ( - Callable, - Dict, - List, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, - Any, -) from copy import deepcopy +from dataclasses import fields, dataclass, replace +from enum import Enum +from typing import List, Optional, Tuple, Union, Dict, Any, Sequence, Callable, cast, MutableMapping + import torch -import torch.backends.cuda -import torch.nn as nn -import torch.nn.functional as F -from torch import einsum -import einops -from transformers import PreTrainedModel +from einops import einsum, einops +from transformers import PreTrainedModel, GenerationConfig +from transformers.cache_utils import Cache from transformers.modeling_outputs import CausalLMOutputWithPast, ModelOutput +from transformers.models.auto import AutoModelForCausalLM +from torch import nn -# from olmo.util import resource_path -from .config_molmo import ( - MolmoConfig, - VisionBackboneConfig, - VisionBackboneType, - ImagePooling2DType, - ImageProjectType, - AttentionType, - MolmoConfigurationError, -) - -if sys.version_info.minor > 8: - from collections.abc import MutableMapping -elif sys.version_info.minor == 8: - from typing import MutableMapping -else: - raise SystemExit("This script supports Python 3.8 or higher") +from .config_molmo import MolmoConfig +from torch.nn import functional as F log = logging.getLogger(__name__) -def resource_path( - folder: Union[str, Path], - fname: str, - local_cache: Optional[Union[str, Path]] = None, -) -> Path: - if local_cache is not None and (local_path := Path(local_cache) / fname).is_file(): - log.info(f"Found local cache of {fname} at {local_path}") - return local_path - else: - from cached_path import cached_path - - return cached_path(f"{str(folder).rstrip('/')}/{fname}") - - -def ensure_finite_(x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False): - """ - Modify ``x`` in place to replace ``float("-inf")`` with the minimum value of the dtype when ``check_neg_inf`` - is ``True`` and to replace ``float("inf")`` with the maximum value of the dtype when ``check_pos_inf`` is ``True``. - """ - if check_neg_inf: - x.masked_fill_(x == float("-inf"), torch.finfo(x.dtype).min) - if check_pos_inf: - x.masked_fill_(x == float("inf"), torch.finfo(x.dtype).max) - - class BufferCache(dict, MutableMapping[str, torch.Tensor]): """ Cache for attention biases and other things that would normally be stored as buffers. @@ -90,210 +31,61 @@ class BufferCache(dict, MutableMapping[str, torch.Tensor]): """ -def _non_meta_init_device(config: MolmoConfig) -> torch.device: - if config.init_device is not None and config.init_device != "meta": - return torch.device(config.init_device) - else: - return torch.device("cuda" if torch.cuda.is_available() else "cpu") - - -class Embedding(nn.Module): - def __init__( - self, - num_embeddings: int, - num_new_embeddings: int, - features: int, - device: Union[str, torch.device], - initializer_range: float = 0.02, - new_embed_initializer_range: float = 0.02, - ): - super().__init__() - self.initializer_range = initializer_range - self.new_embed_initializer_range = new_embed_initializer_range - self.embedding = nn.Parameter( - torch.zeros(num_embeddings, features, device=device), - ) - self.new_embedding = nn.Parameter( - torch.zeros(num_new_embeddings, features, device=device), - ) - - def reset_parameters(self): - nn.init.normal_(self.embedding, std=self.initializer_range) - nn.init.normal_(self.new_embedding, std=self.new_embed_initializer_range) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0)) - - -class Dropout(nn.Dropout): - def __init__( - self, - p: float = 0.5, - inplace: bool = False, - mask_p: float = 0, - broadcast_dims: Sequence[int] = (), - ): - super().__init__(p, inplace) - self.mask_p = mask_p - self.broadcast_dims = broadcast_dims - - def forward(self, input: torch.Tensor, drop_mask: Optional[torch.Tensor] = None) -> torch.Tensor: - """ - :param input: A tensor of shape `(batch_size, seq_len, embed_dim)` - :param drop_mask: A tensor of shape `(batch_size, seq_len)` with values of zero or one. - """ - if self.p == 0.0 and (self.mask_p is None or self.mask_p == 0.0): - return input - else: - if self.mask_p > 0.0 and self.training: - assert drop_mask is not None - drop_mask = drop_mask.to(input.dtype) - keep_prob = 1.0 - self.p - keep_prob2 = 1.0 - self.mask_p - keep_prob = drop_mask * keep_prob2 + (1 - drop_mask) * keep_prob - keep_prob = keep_prob.unsqueeze(-1) - dropout_shape = list(input.shape) - keep_prob = keep_prob.broadcast_to(dropout_shape) - multiplier = input.new_empty(dropout_shape).bernoulli_(keep_prob) - multiplier.div_(keep_prob) - return input * multiplier - elif self.p > 0.0 and len(self.broadcast_dims) > 0 and self.training: - keep_prob = 1.0 - self.p - dropout_shape = list(input.shape) - for dim in self.broadcast_dims: - dropout_shape[dim] = 1 - keep = input.new_empty(dropout_shape).bernoulli_(keep_prob) - multiplier = keep.broadcast_to(input.shape) - multiplier.div_(keep_prob) - input = input * multiplier - else: - return F.dropout(input, self.p, self.training, self.inplace) - +class StrEnum(str, Enum): + def __str__(self) -> str: + return self.value -class LayerNormBase(nn.Module): - def __init__( - self, - config: MolmoConfig, - *, - size: Optional[int] = None, - elementwise_affine: Optional[bool] = True, - eps: float = 1e-05, - weight_initializer: Optional[Callable] = torch.ones, - bias_initializer: Optional[Callable] = torch.zeros, - ): - super().__init__() - self.config = config - self.eps = self.config.layer_norm_eps or eps - self.normalized_shape = (size or config.d_model,) - if elementwise_affine or (elementwise_affine is None and self.config.layer_norm_with_affine): - self.weight = nn.Parameter(weight_initializer(self.normalized_shape, device=config.init_device)) - use_bias = self.config.bias_for_layer_norm - if use_bias is None: - use_bias = self.config.include_bias - if use_bias: - self.bias = nn.Parameter(bias_initializer(self.normalized_shape, device=config.init_device)) - else: - self.register_parameter("bias", None) - else: - self.register_parameter("bias", None) - self.register_parameter("weight", None) + def __repr__(self) -> str: + return f"'{str(self)}'" -class LayerNorm(LayerNormBase): - """ - The default :class:`LayerNorm` implementation which can optionally run in low precision. - """ +class ImageProjectType(StrEnum): + mlp = "mlp" + mlpx2 = "2mlp" + linear = "linear" - def __init__( - self, - config: MolmoConfig, - size: Optional[int] = None, - low_precision: bool = False, - elementwise_affine: Optional[bool] = None, - eps: float = 1e-05, - ): - super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) - self.low_precision = low_precision - def _cast_if_autocast_enabled(self, tensor: torch.Tensor, dtype: Optional[torch.dtype] = None) -> torch.Tensor: - # NOTE: `is_autocast_enabled()` only checks for CUDA autocast, so we use the separate function - # `is_autocast_cpu_enabled()` for CPU autocast. - # See https://github.com/pytorch/pytorch/issues/110966. - if tensor.device.type == "cuda" and torch.is_autocast_enabled(): - return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_gpu_dtype()) - elif tensor.device.type == "cpu" and torch.is_autocast_cpu_enabled(): - return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_cpu_dtype()) - else: - return tensor +class ImagePooling2DType(StrEnum): + attention = "attention" + attention_meanq = "attention-meanq" + attention_2wide = "attention_2wide" + attention_v2 = "attention-v2" + none = "none" + stack = "stack" - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.low_precision: - module_device = x.device - downcast_x = self._cast_if_autocast_enabled(x) - downcast_weight = ( - self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight - ) - downcast_bias = self._cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias - with torch.autocast(enabled=False, device_type=module_device.type): - return F.layer_norm( - downcast_x, self.normalized_shape, weight=downcast_weight, bias=downcast_bias, eps=self.eps - ) - else: - return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps) - def reset_parameters(self): - if self.weight is not None: - torch.nn.init.ones_(self.weight) # type: ignore - if self.bias is not None: - torch.nn.init.zeros_(self.bias) # type: ignore +class ActivationType(StrEnum): + quick_gelu = "quick_gelu" + gelu = "gelu" + gelu_tanh = "gelu_tanh" + relu = "relu" + silu = "silu" + llama_geglu = "llama_geglu" + llama_geglu_tanh = "llama_geglu_tanh" + llama_swiglu = "llama_swiglu" + swiglu = "swiglu" -class RMSLayerNorm(LayerNormBase): +def ensure_finite_(x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False): """ - RMS layer norm, a simplified :class:`LayerNorm` implementation + Modify ``x`` in place to replace ``float("-inf")`` with the minimum value of the dtype when ``check_neg_inf`` + is ``True`` and to replace ``float("inf")`` with the maximum value of the dtype when ``check_pos_inf`` is ``True``. """ + if check_neg_inf: + x.masked_fill_(x == float("-inf"), torch.finfo(x.dtype).min) + if check_pos_inf: + x.masked_fill_(x == float("inf"), torch.finfo(x.dtype).max) - def __init__( - self, - config: MolmoConfig, - size: Optional[int] = None, - elementwise_affine: Optional[bool] = None, - eps: float = 1e-5, - ): - super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - with torch.autocast(enabled=False, device_type=x.device.type): - og_dtype = x.dtype - x = x.to(torch.float32) - variance = x.pow(2).mean(-1, keepdim=True) - x = x * torch.rsqrt(variance + self.eps) - x = x.to(og_dtype) - if self.weight is not None: - if self.bias is not None: - return self.weight * x + self.bias - else: - return self.weight * x - else: - return x +class OLMoConfigurationError(Exception): + pass - def _cast_if_autocast_enabled(self, tensor: torch.Tensor, dtype: Optional[torch.dtype] = None) -> torch.Tensor: - # NOTE: `is_autocast_enabled()` only checks for CUDA autocast, so we use the separate function - # `is_autocast_cpu_enabled()` for CPU autocast. - # See https://github.com/pytorch/pytorch/issues/110966. - if tensor.device.type == "cuda" and torch.is_autocast_enabled(): - return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_gpu_dtype()) - elif tensor.device.type == "cpu" and torch.is_autocast_cpu_enabled(): - return tensor.to(dtype=dtype if dtype is not None else torch.get_autocast_cpu_dtype()) - else: - return tensor - def reset_parameters(self): - if self.weight is not None: - torch.nn.init.ones_(self.weight) # type: ignore - if self.bias is not None: - torch.nn.init.zeros_(self.bias) # type: ignore +def _non_meta_init_device(config) -> torch.device: + if config.init_device is not None and config.init_device != "meta": + return torch.device(config.init_device) + else: + return torch.device("cuda" if torch.cuda.is_available() else "cpu") class RotaryEmbedding(nn.Module): @@ -307,7 +99,8 @@ class RotaryEmbedding(nn.Module): self.__cache = cache # Warm up cache. self.get_rotary_embedding( - config.max_position_embeddings or config.max_sequence_length, _non_meta_init_device(config) + config.max_position_embeddings or config.max_sequence_length, + _non_meta_init_device(config) ) def get_rotary_embedding(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: @@ -326,16 +119,10 @@ class RotaryEmbedding(nn.Module): return pos_sin[:, :, :seq_len, :], pos_cos[:, :, :seq_len, :] with torch.autocast(device.type, enabled=False): - dim = ( - self.config.head_dim - if self.config.head_dim is not None - else self.config.d_model // self.config.n_heads - ) - inv_freq = 1.0 / ( - self.config.rope_theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float) / dim) - ) + dim = self.config.d_model // self.config.n_heads + inv_freq = 1.0 / (self.config.rope_theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float) / dim)) seq = torch.arange(seq_len, device=device, dtype=torch.float) - freqs = einsum("i , j -> i j", seq, inv_freq) + freqs = torch.einsum("i , j -> i j", seq, inv_freq) if self.config.rope_impl == "cockatoo": positions = freqs.repeat_interleave(2, dim=-1) else: @@ -365,7 +152,10 @@ class RotaryEmbedding(nn.Module): return ((t * pos_cos) + (self.rotate_half(t) * pos_sin)).to(t.dtype) def forward( - self, q: torch.Tensor, k: torch.Tensor, position_ids: Optional[torch.Tensor] = None + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: if self.config.rope_full_precision: q_, k_ = q.float(), k.float() @@ -376,7 +166,7 @@ class RotaryEmbedding(nn.Module): batch_size = q_.shape[0] query_len, key_len = q_.shape[-2], k_.shape[-2] # could be different if layer_past not None if position_ids is not None: - freqs_cis_len = self.config.max_position_embeddings or self.config.max_sequence_length + freqs_cis_len = (self.config.max_position_embeddings or self.config.max_sequence_length) else: freqs_cis_len = key_len pos_sin, pos_cos = self.get_rotary_embedding(freqs_cis_len, q_.device) @@ -384,8 +174,12 @@ class RotaryEmbedding(nn.Module): pos_cos = pos_cos.type_as(q_) if position_ids is not None: assert query_len == key_len, "Query and key lengths must be equal when using position IDs." - pos_sin = pos_sin[0, 0][position_ids].view((batch_size, 1, key_len, pos_sin.shape[-1])) - pos_cos = pos_cos[0, 0][position_ids].view((batch_size, 1, key_len, pos_cos.shape[-1])) + pos_sin = pos_sin[0, 0][position_ids].view( + (batch_size, 1, key_len, pos_sin.shape[-1]) + ) + pos_cos = pos_cos[0, 0][position_ids].view( + (batch_size, 1, key_len, pos_cos.shape[-1]) + ) q_ = self.apply_rotary_pos_emb( pos_sin[:, :, key_len - query_len : key_len, :], pos_cos[:, :, key_len - query_len : key_len, :], @@ -395,142 +189,150 @@ class RotaryEmbedding(nn.Module): return q_.type_as(q), k_.type_as(k) -class Activation(nn.Module): - def __init__(self, config: MolmoConfig): +class OLMoBlock(nn.Module): + """ + A base class for transformer block implementations. + """ + + def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): super().__init__() + self.layer_id = layer_id self.config = config + self.hidden_size = ( + config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model + ) + self.__cache = cache + self._activation_checkpoint_fn = None - @abstractmethod - def forward(self, x: torch.Tensor) -> torch.Tensor: - raise NotImplementedError - - @property - @abstractmethod - def output_multiplier(self) -> float: - raise NotImplementedError + # Dropout. + self.dropout = Dropout(config.residual_dropout, mask_p=config.response_residual_dropout) + # Layer norms. + self.k_norm: Optional[LayerNormBase] = None + self.q_norm: Optional[LayerNormBase] = None + if config.attention_layer_norm: + assert config.effective_n_kv_heads is not None + self.k_norm = LayerNormBase.build( + config, + size=(config.d_model // config.n_heads) * config.effective_n_kv_heads, + elementwise_affine=config.attention_layer_norm_with_affine, + ) + self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine) -class GELU(nn.GELU): - @property - def output_multiplier(self) -> float: - return 1.0 + # Make sure QKV clip coefficient is positive, otherwise it's not well-defined. + if config.clip_qkv is not None: + assert config.clip_qkv > 0 + # Activation function. + self.act = Activation.build(config) + assert (self.act.output_multiplier * self.hidden_size) % 1 == 0 -class QuickGELU(Activation): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x * torch.sigmoid(1.702 * x) + # Attention output projection. + input_dim = config.d_model + self.attn_out = nn.Linear( + input_dim, config.d_model, + bias=config.include_bias, + device=config.init_device + ) - @property - def output_multiplier(self) -> float: - return 1.0 + # Feed-forward output projection. + self.ff_out = nn.Linear( + int(self.act.output_multiplier * self.hidden_size), + config.d_model, + bias=config.include_bias, + device=config.init_device, + ) + self.ff_out._is_residual = True # type: ignore + # Rotary embeddings. + if self.config.rope: + self.rotary_emb = RotaryEmbedding(config, self.__cache) -class ReLU(nn.ReLU): - @property - def output_multiplier(self) -> float: - return 1.0 - - -class SiLU(nn.SiLU): - @property - def output_multiplier(self) -> float: - return 1.0 - - -class SwiGLU(Activation): - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, gate = x.chunk(2, dim=-1) - return F.silu(gate) * x - - @property - def output_multiplier(self) -> float: - return 0.5 - - -class LlamaSwiGLU(Activation): - def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: - return F.silu(x1) * x2 - - @property - def output_multiplier(self) -> float: - return 0.5 - - -def causal_attention_bias(seq_len: int, device: torch.device) -> torch.FloatTensor: - att_bias = torch.triu( - torch.ones(seq_len, seq_len, device=device, dtype=torch.float), - diagonal=1, - ) - att_bias.masked_fill_(att_bias == 1, torch.finfo(att_bias.dtype).min) - return att_bias.view(1, 1, seq_len, seq_len) # type: ignore - - -def get_causal_attention_bias(cache: BufferCache, seq_len: int, device: torch.device) -> torch.Tensor: - if (causal_bias := cache.get("causal_attention_bias")) is not None and causal_bias.shape[-1] >= seq_len: - if causal_bias.device != device: - causal_bias = causal_bias.to(device) - cache["causal_attention_bias"] = causal_bias - return causal_bias - with torch.autocast(device.type, enabled=False): - causal_bias = causal_attention_bias(seq_len, device) - cache["causal_attention_bias"] = causal_bias - return causal_bias + self.flash_attn_func = None + if config.attention_type == "flash": + try: + from flash_attn import flash_attn_func # type: ignore + self.flash_attn_func = flash_attn_func + except ModuleNotFoundError: + pass -class MolmoAttention(nn.Module): - def __init__(self, config: MolmoConfig, cache: BufferCache): - super().__init__() - self.config = config - self.__cache = cache - self.rotary_emb = RotaryEmbedding(config, self.__cache) - self.k_norm: Optional[LayerNormBase] = None - self.q_norm: Optional[LayerNormBase] = None - self.hidden_size = ( - config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model + def reset_parameters(self): + if self.k_norm is not None: + self.k_norm.reset_parameters() + if self.q_norm is not None: + self.q_norm.reset_parameters() + init_weights( + self.config, + self.attn_out, + d=self.config.d_model, + layer_id=self.layer_id, + type_of_module=ModuleType.out_module, + ) + init_weights( + self.config, + self.ff_out, + d=self.ff_out.in_features, + layer_id=self.layer_id, + type_of_module=ModuleType.out_module, ) - if config.attention_layer_norm: - assert config.n_kv_heads is not None - self.k_norm = LayerNormBase.build( - config, - size=(config.d_model // config.n_heads) * config.n_kv_heads, - elementwise_affine=config.attention_layer_norm_with_affine, - ) - self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine) - - # Make sure QKV clip coefficient is positive, otherwise it's not well-defined. - if config.clip_qkv is not None: - assert config.clip_qkv > 0 - - # Activation function - self.act = SwiGLU(config) - assert (self.act.output_multiplier * self.hidden_size) % 1 == 0 + @classmethod + def _cast_attn_bias(cls, bias: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor: + target_dtype = input_dtype + # NOTE: `is_autocast_enabled()` only checks for CUDA autocast, so we use the separate function + # `is_autocast_cpu_enabled()` for CPU autocast. + # See https://github.com/pytorch/pytorch/issues/110966. + if bias.device.type == "cuda" and torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + elif bias.device.type == "cpu" and torch.is_autocast_cpu_enabled(): + target_dtype = torch.get_autocast_cpu_dtype() + if bias.dtype != target_dtype: + bias = bias.to(target_dtype) + ensure_finite_(bias, check_neg_inf=True, check_pos_inf=False) + return bias - # Attention output projection. - input_dim = config.head_dim * config.n_heads if config.head_dim is not None else config.d_model - head_dim = config.d_model // config.n_heads - self.fused_dims = ( - config.d_model, - config.n_kv_heads * head_dim, - config.n_kv_heads * head_dim, - ) - self.att_proj = nn.Linear( - config.d_model, - sum(self.fused_dims), - bias=config.include_bias or config.qkv_bias, - device=config.init_device, - ) - self.attn_out = nn.Linear(input_dim, config.d_model, bias=config.include_bias, device=config.init_device) - self.attn_norm = RMSLayerNorm(config, size=config.d_model, eps=config.layer_norm_eps) + def _scaled_dot_product_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + drop_mask: Optional[torch.Tensor] = None, + dropout_p: float = 0.0, + response_dropout_p: float = 0.0, + is_causal: bool = False, + ) -> torch.Tensor: + """ + Computes scaled dot product attention on query, key and value tensors, using an optional + attention mask if passed, and applying dropout if a probability greater than 0.0 is specified. + """ + if attn_mask is not None: + attn_mask = attn_mask.to(q.device) - self.flash_attn_func = None - if self.config.attention_type == AttentionType.flash: - try: - from flash_attn import flash_attn_func + if self.flash_attn_func is not None and attn_mask is None: + r = self.flash_attn_func( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=is_causal + ) + return r.transpose(1, 2) + else: + # torch's sdpa doesn't support GQA, so we're doing this + assert k.size(1) == v.size(1) + num_kv_heads = k.size(1) + num_q_heads = q.size(1) + if num_q_heads != num_kv_heads: + assert num_q_heads % num_kv_heads == 0 + k = k.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) + v = v.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) - self.flash_attn_func = flash_attn_func - except ModuleNotFoundError: - pass + return F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + ) def attention( self, @@ -555,9 +357,9 @@ class MolmoAttention(nn.Module): # shape: (B, nh, T, hs) q = q.view(B, T, self.config.n_heads, C // self.config.n_heads).transpose(1, 2) # shape: (B, n_kv_h, T, hs) - k = k.view(B, T, self.config.n_kv_heads, C // self.config.n_heads).transpose(1, 2) + k = k.view(B, T, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2) # shape: (B, n_kv_h, T, hs) - v = v.view(B, T, self.config.n_kv_heads, C // self.config.n_heads).transpose(1, 2) + v = v.view(B, T, self.config.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2) if self.config.use_position_ids and self.config.rope: # Apply rotary embeddings @@ -604,20 +406,77 @@ class MolmoAttention(nn.Module): # Apply output projection. return self.attn_out(att), present + def forward( + self, + x: torch.Tensor, + attention_bias: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.Tensor] = None, + drop_mask: Optional[torch.Tensor] = None, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + raise NotImplementedError + @classmethod - def _cast_attn_bias(cls, bias: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor: - target_dtype = input_dtype - # NOTE: `is_autocast_enabled()` only checks for CUDA autocast, so we use the separate function - # `is_autocast_cpu_enabled()` for CPU autocast. - # See https://github.com/pytorch/pytorch/issues/110966. - if bias.device.type == "cuda" and torch.is_autocast_enabled(): - target_dtype = torch.get_autocast_gpu_dtype() - elif bias.device.type == "cpu" and torch.is_autocast_cpu_enabled(): - target_dtype = torch.get_autocast_cpu_dtype() - if bias.dtype != target_dtype: - bias = bias.to(target_dtype) - ensure_finite_(bias, check_neg_inf=True, check_pos_inf=False) - return bias + def build(cls, layer_id: int, config: MolmoConfig, cache: BufferCache): + if config.block_type == "sequential": + return OLMoSequentialBlock(layer_id, config, cache) + elif config.block_type == "llama": + return OLMoLlamaBlock(layer_id, config, cache) + else: + raise NotImplementedError(f"Unknown block type: '{config.block_type}'") + + +class OLMoLlamaBlock(OLMoBlock): + """ + This is a transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). This block is similar to `OLMoSequentialBlock` + but some operations have slightly different implementations to imitate the + behavior of Llama. + """ + + def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): + super().__init__(layer_id, config, cache) + # Layer norms. + self.attn_norm = LayerNorm.build(config) + self.ff_norm = LayerNorm.build(config) + self.__cache = cache + + # Attention input projection. Projects x -> (q, k, v) + q_proj_out_dim = config.d_model + k_proj_out_dim = config.effective_n_kv_heads * (config.d_model // config.n_heads) + v_proj_out_dim = config.effective_n_kv_heads * (config.d_model // config.n_heads) + + self.q_proj = nn.Linear( + config.d_model, q_proj_out_dim, bias=config.qkv_bias, device=config.init_device + ) + self.k_proj = nn.Linear( + config.d_model, k_proj_out_dim, bias=config.qkv_bias, device=config.init_device + ) + self.v_proj = nn.Linear( + config.d_model, v_proj_out_dim, bias=config.qkv_bias, device=config.init_device + ) + + # Feed-forward input projection. + self.ff_proj1 = nn.Linear( + config.d_model, self.hidden_size // 2, bias=False, device=config.init_device + ) + self.ff_proj2 = nn.Linear( + config.d_model, self.hidden_size // 2, bias=False, device=config.init_device + ) + if self.config.norm_after: + raise NotImplementedError() + + def reset_parameters(self): + super().reset_parameters() + self.attn_norm.reset_parameters() + self.ff_norm.reset_parameters() + # NOTE: the standard deviation for these weights does not depend on the layer. + init_weights(self.config, self.q_proj, d=self.config.d_model, layer_id=None) + init_weights(self.config, self.k_proj, d=self.config.d_model, layer_id=None) + init_weights(self.config, self.v_proj, d=self.config.d_model, layer_id=None) + init_weights(self.config, self.ff_proj1, d=self.config.d_model, layer_id=None) + init_weights(self.config, self.ff_proj2, d=self.config.d_model, layer_id=None) def _scaled_dot_product_attention( self, @@ -630,29 +489,47 @@ class MolmoAttention(nn.Module): response_dropout_p: float = 0.0, is_causal: bool = False, ) -> torch.Tensor: - """ - Computes scaled dot product attention on query, key and value tensors, using an optional - attention mask if passed, and applying dropout if a probability greater than 0.0 is specified. - """ + # For GQA + assert k.size(1) == v.size(1) + num_kv_heads = k.size(1) + num_q_heads = q.size(1) + if num_q_heads != num_kv_heads: + assert num_q_heads % num_kv_heads == 0 + k = k.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) + v = v.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) + + og_dtype = q.dtype + k = k.to(q.device) + v = v.to(q.device) if attn_mask is not None: attn_mask = attn_mask.to(q.device) - if self.flash_attn_func is not None and attn_mask is None: - r = self.flash_attn_func( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=dropout_p, causal=is_causal - ) - return r.transpose(1, 2) - else: - # torch's sdpa doesn't support GQA, so we're doing this - assert k.size(1) == v.size(1) - num_kv_heads = k.size(1) - num_q_heads = q.size(1) - if num_q_heads != num_kv_heads: - assert num_q_heads % num_kv_heads == 0 - k = k.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) - v = v.repeat_interleave(num_q_heads // num_kv_heads, dim=1, output_size=num_q_heads) + assert response_dropout_p == 0.0, "Response dropout is not supported in Llama." - return F.scaled_dot_product_attention( + if self.config.float32_attention: + q, k = q.to(torch.float), k.to(torch.float) + + if self.config.attention_type == "direct": + attn_weights = torch.matmul(q, k.transpose(-2, -1)) / (q.shape[-1] ** 0.5) + + if is_causal: + assert attn_mask is None + + query_len, key_len = q.shape[-2], k.shape[-2] # could be different if layer_past not None + attn_bias = get_causal_attention_bias(self.__cache, key_len, q.device)[:, :, :query_len, :key_len] + elif attn_mask is not None: + attn_bias = attn_mask + else: + attn_bias = torch.zeros_like(attn_weights) + + attn_weights += attn_bias + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout_p, training=self.training).to(v.dtype) + + att = torch.matmul(attn_weights, v) + elif self.config.attention_type == "sdpa": + att = F.scaled_dot_product_attention( q, k, v, @@ -660,144 +537,597 @@ class MolmoAttention(nn.Module): dropout_p=dropout_p, is_causal=is_causal, ) - - def forward(self, x, attention_bias, position_ids, drop_mask, layer_past, use_cache): - if not self.config.norm_after: - atten_in = self.attn_norm(x) else: - atten_in = x + raise NotImplementedError(self.config.attention_type) + att = att.to(og_dtype) + return att - qkv = self.att_proj(atten_in) + def forward( + self, + x: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + drop_mask: Optional[torch.Tensor] = None, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + # Get query, key, value projections. + # shape: + # - for regular attn q, k, v: (batch_size, seq_len, d_model) + # - for multi-query attn q: (batch_size, seq_len, d_model) + # k, v: (batch_size, seq_len, d_model // n_heads) + x_normed = self.attn_norm(x) + q = self.q_proj(x_normed) + k = self.k_proj(x_normed) + v = self.v_proj(x_normed) if self.config.clip_qkv is not None: - qkv.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) - - q, k, v = qkv.split(self.fused_dims, dim=-1) + q.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + k.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + v.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) # Get attention scores. - att, cache = self.attention( - q, - k, - v, - attention_bias, - position_ids=position_ids, - drop_mask=drop_mask, - layer_past=layer_past, - use_cache=use_cache, - ) + if self._activation_checkpoint_fn is not None: + att, cache = self._activation_checkpoint_fn( # type: ignore + self.attention, q, k, v, attention_bias, position_ids=position_ids, drop_mask=drop_mask, layer_past=layer_past, use_cache=use_cache + ) + else: + att, cache = self.attention(q, k, v, attention_bias, position_ids=position_ids, drop_mask=drop_mask, layer_past=layer_past, use_cache=use_cache) - if self.config.norm_after: - att = self.attn_norm(att) + # Add attention scores. + # shape: (B, T, C) + x = x + self.dropout(att, drop_mask=drop_mask) - return att, cache + # Add feed-forward projection. + # shape: (batch_size, seq_len, d_model) + og_x = x + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore + else: + x = self.ff_norm(x) + x1 = self.ff_proj1(x) + x2 = self.ff_proj2(x) + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.act, x1, x2) # type: ignore + else: + x = self.act(x1, x2) + x = self.ff_out(x) + x = self.dropout(x, drop_mask=drop_mask) + x = og_x + x + return x, cache -class MolmoMLP(nn.Module): - def __init__(self, config: MolmoConfig): - # Feed-forward input projection. - super().__init__() - self.config = config - self.hidden_size = ( - config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model + +class OLMoSequentialBlock(OLMoBlock): + """ + This is a typical transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). + """ + + def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): + super().__init__(layer_id, config, cache) + # Layer norms. + self.attn_norm = LayerNorm.build(config) + self.ff_norm = LayerNorm.build(config) + # Attention input projection. Projects x -> (q, k, v) + + head_dim = config.d_model // config.n_heads + self.fused_dims = ( + config.d_model, + config.effective_n_kv_heads * head_dim, + config.effective_n_kv_heads * head_dim, ) - self.act = SwiGLU(config) + self.att_proj = nn.Linear( + config.d_model, sum(self.fused_dims), + bias=config.include_bias or config.qkv_bias, + device=config.init_device + ) + # Feed-forward input projection. self.ff_proj = nn.Linear( config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device ) - self.ff_out = nn.Linear( - int(self.act.output_multiplier * self.hidden_size), - config.d_model, - bias=config.include_bias, - device=config.init_device, + + def reset_parameters(self): + super().reset_parameters() + self.attn_norm.reset_parameters() + self.ff_norm.reset_parameters() + # NOTE: the standard deviation for these weights does not depend on the layer. + init_weights( + self.config, self.att_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module + ) + init_weights( + self.config, self.ff_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module ) - self.ff_norm = RMSLayerNorm(config, size=config.d_model, eps=config.layer_norm_eps) - def forward(self, x): + def forward( + self, + x: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + drop_mask: Optional[torch.Tensor] = None, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + # Get query, key, value projections. + # shape: + # - for regular attn q, k, v: (batch_size, seq_len, d_model) + # - for multi-query attn q: (batch_size, seq_len, d_model) + # k, v: (batch_size, seq_len, d_model // n_heads) + # - for group query attn q: (batch_size, seq_len, d_model) + # k, v: (batch_size, seq_len, d_model // n_kv_heads) + if not self.config.norm_after: - x = self.ff_norm(x) + if self._activation_checkpoint_fn is not None: + atten_in = self._activation_checkpoint_fn(self.attn_norm, x) + else: + atten_in = self.attn_norm(x) + else: + atten_in = x + qkv = self.att_proj(atten_in) + + if self.config.clip_qkv is not None: + qkv.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv) + + q, k, v = qkv.split(self.fused_dims, dim=-1) + + # Get attention scores. + if self._activation_checkpoint_fn is not None: + att, cache = self._activation_checkpoint_fn( # type: ignore + self.attention, q, k, v, attention_bias, position_ids=position_ids, drop_mask=drop_mask, layer_past=layer_past, use_cache=use_cache + ) + else: + att, cache = self.attention(q, k, v, attention_bias, position_ids=position_ids, drop_mask=drop_mask, layer_past=layer_past, use_cache=use_cache) + + if self.config.norm_after: + if self._activation_checkpoint_fn is not None: + att = self._activation_checkpoint_fn(self.attn_norm, att) + else: + att = self.attn_norm(att) + + # Add attention scores. + # shape: (B, T, C) + x = x + self.dropout(att, drop_mask=drop_mask) + + # Add feed-forward projection. + # shape: (batch_size, seq_len, d_model) + og_x = x + + if not self.config.norm_after: + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore + else: + x = self.ff_norm(x) x = self.ff_proj(x) - x = self.act(x) + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.act, x) # type: ignore + else: + x = self.act(x) x = self.ff_out(x) if self.config.norm_after: - x = self.ff_norm(x) + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore + else: + x = self.ff_norm(x) - return x + x = self.dropout(x, drop_mask=drop_mask) + x = og_x + x + return x, cache -class MolmoDecoderLayer(nn.Module): - """ - A base class for transformer block implementations. - """ - def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): +class Embedding(nn.Module): + def __init__( + self, + num_embeddings: int, + num_new_embeddings: int, + features: int, + device: Union[str, torch.device], + initializer_range: float = 0.02, + new_embed_initializer_range: float = 0.02, + ): super().__init__() - self.self_attn = MolmoAttention(config, cache) - self.mlp = MolmoMLP(config) - self.layer_id = layer_id - self.config = config - self.hidden_size = ( - config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model + self.initializer_range = initializer_range + self.new_embed_initializer_range = new_embed_initializer_range + self.embedding = nn.Parameter( + torch.zeros(num_embeddings, features, device=device), + ) + self.new_embedding = nn.Parameter( + torch.zeros(num_new_embeddings, features, device=device), ) - self.__cache = cache - if config.head_dim is None: - assert config.d_model % config.n_heads == 0 - # Dropout. - self.dropout = Dropout(config.residual_dropout, mask_p=config.response_residual_dropout) + def reset_parameters(self): + nn.init.normal_(self.embedding, std=self.initializer_range) + nn.init.normal_(self.new_embedding, std=self.new_embed_initializer_range) - def forward( + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0)) + + +class Dropout(nn.Dropout): + def __init__( self, - x: torch.Tensor, - attention_bias: Optional[torch.FloatTensor] = None, - position_ids: Optional[torch.Tensor] = None, - drop_mask: Optional[torch.Tensor] = None, - layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - use_cache: bool = False, - ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: - """Get query, key, value projections. - shape: - for regular attn q, k, v: (batch_size, seq_len, d_model) - for multi-query attn q: (batch_size, seq_len, d_model) - k, v: (batch_size, seq_len, d_model // n_heads) - for group query attn q: (batch_size, seq_len, d_model) - k, v: (batch_size, seq_len, d_model // n_kv_heads) + p: float = 0.5, + inplace: bool = False, + mask_p: float = 0, + broadcast_dims: Sequence[int] = (), + ): + super().__init__(p, inplace) + self.mask_p = mask_p + self.broadcast_dims = broadcast_dims + + def forward(self, input: torch.Tensor, drop_mask: Optional[torch.Tensor] = None) -> torch.Tensor: """ + :param input: A tensor of shape `(batch_size, seq_len, embed_dim)` + :param drop_mask: A tensor of shape `(batch_size, seq_len)` with values of zero or one. + """ + if self.p == 0.0 and (self.mask_p is None or self.mask_p == 0.0): + return input + else: + if self.mask_p > 0. and self.training: + assert drop_mask is not None + drop_mask = drop_mask.to(input.dtype) + keep_prob = 1.0 - self.p + keep_prob2 = 1.0 - self.mask_p + keep_prob = drop_mask * keep_prob2 + (1 - drop_mask) * keep_prob + keep_prob = keep_prob.unsqueeze(-1) + dropout_shape = list(input.shape) + keep_prob = keep_prob.broadcast_to(dropout_shape) + multiplier = input.new_empty(dropout_shape).bernoulli_(keep_prob) + multiplier.div_(keep_prob) + return input * multiplier + elif self.p > 0. and len(self.broadcast_dims) > 0 and self.training: + keep_prob = 1.0 - self.p + dropout_shape = list(input.shape) + for dim in self.broadcast_dims: + dropout_shape[dim] = 1 + keep = input.new_empty(dropout_shape).bernoulli_(keep_prob) + multiplier = keep.broadcast_to(input.shape) + multiplier.div_(keep_prob) + input = input * multiplier + else: + return F.dropout(input, self.p, self.training, self.inplace) + + +@dataclass +class VisionBackboneConfig: + image_model_type: str = "openai" + image_default_input_size: Tuple[int, int] = (336, 336) + image_patch_size: int = 14 + image_pos_patch_size: int = 14 + image_emb_dim: int = 1024 + image_num_heads: int = 16 + image_num_key_value_heads: int = 16 + image_num_layers: int = 24 + image_head_dim: int = 64 + image_mlp_dim: int = 4096 + image_mlp_activations: str = "gelu" + image_dropout_rate: float = 0.0 + image_num_pos: int = 577 + image_norm_eps: float = 1e-5 + attention_dropout: float = 0.0 + residual_dropout: float = 0.0 + initializer_range: float = 0.02 + fsdp_wrap: bool = False + resize_mode: str = "default" + + def __post_init__(self): + self.image_default_input_size = tuple(self.image_default_input_size) # type: ignore[assignment] + + @property + def image_num_patch(self): + h, w = self.image_default_input_size + return h // self.image_patch_size, w // self.image_patch_size + + +@dataclass +class FullMolmoConfig: + d_model: int = 768 + n_heads: int = 12 + head_dim: int = 64 + n_kv_heads: Optional[int] = None + qkv_bias: bool = False + clip_qkv: Optional[float] = None + n_layers: int = 12 + mlp_ratio: int = 4 + mlp_hidden_size: Optional[int] = None + activation_type: str = "swiglu" + block_type: str = "sequential" + block_group_size: int = 1 + alibi: bool = False + alibi_bias_max: float = 8.0 + rope: bool = False + rope_full_precision: bool = True + rope_theta: float = 10000. + rope_impl: str = "cockatoo" + vision_backbone: Optional[VisionBackboneConfig] = None + vit_load_path: Optional[str] = None + llm_load_path: Optional[str] = None + attention_type: str = "sdpa" + float32_attention: bool = True + attention_dropout: float = 0.1 + response_attention_dropout: float = 0.0 + multi_query_attention: Optional[bool] = None + attention_layer_norm: bool = False + residual_dropout: float = 0.1 + response_residual_dropout: float = 0.0 + embedding_dropout: float = 0.1 + layer_norm_type: str = "default" + layer_norm_with_affine: bool = True + layer_norm_eps: Optional[float] = None + attention_layer_norm_with_affine: bool = True + max_sequence_length: int = 1024 + max_position_embeddings: Optional[int] = None + include_bias: bool = True + bias_for_layer_norm: Optional[bool] = None + scale_logits: bool = False + vocab_size: int = 50257 + embedding_size: Optional[int] = 50304 + additional_vocab_size: Optional[int] = None + new_embedding_init_range: float = 0.02 + weight_tying: bool = True + pad_token_id: int = -1 + init_device: Optional[str] = None + init_std: float = 0.02 + init_cutoff_factor: Optional[float] = None + norm_after: bool = False + precision: Optional[str] = None + max_crops: int = 12 + crop_mode: str = "patchify-v2-and-resize-c2" + do_random_scale: bool = True + use_col_tokens: bool = True + image_padding_embed: Optional[str] = None + vit_layers: Tuple = (-1,) + image_pooling_h: int = 2 + image_pooling_w: int = 2 + image_pooling_2d: str = "attention" + image_projector: str = "mlp" + image_feature_dropout: float = 0.0 + use_cls_feature: bool = False + initializer_range: float = 0.02 + pad_tokenizer: bool = False + normalize_input_embeds: bool = False + use_position_ids: bool = True + query_pre_attn_scalar: int = 224 + + @property + def effective_n_kv_heads(self) -> int: + if self.n_kv_heads is None: + if self.multi_query_attention is True: + return 1 + else: + return self.n_heads + else: + if self.multi_query_attention is None: + return self.n_kv_heads + if self.multi_query_attention: + n_kv_heads_should_be = 1 + else: + n_kv_heads_should_be = self.n_heads + if self.n_kv_heads == n_kv_heads_should_be: + return n_kv_heads_should_be + else: + raise OLMoConfigurationError( + "You can't set `multi_query_attention` and `n_kv_heads` at the same time." + ) + + @property + def image_num_patch(self): + assert self.vision_backbone is not None + return self.vision_backbone.image_num_patch + + @property + def image_patch_size(self): + assert self.vision_backbone is not None + return self.visoin_backbone.image_patch_size + + def llm_patches_per_crop(self): + h, w = self.image_num_patch + # Round up in case we need to pad the image features for pooling + h = (h + self.image_pooling_h - 1) // self.image_pooling_h + w = (w + self.image_pooling_w - 1) // self.image_pooling_w + return h, w + + +def _expand_token(token, batch_size: int): + return token.view(1, 1, -1).expand(batch_size, -1, -1) + + +class LayerNormFp32(nn.LayerNorm): + """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back). + Derived from https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/transformer.py. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + orig_type = x.dtype + x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps) + return x.to(orig_type) + + +class ViTMLP(nn.Module): + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + v_cfg = config.vision_backbone + + self.w1 = nn.Linear( + v_cfg.image_emb_dim, + v_cfg.image_mlp_dim, + bias=True, + device=config.init_device, + ) + # Activation function. + cfg = deepcopy(config) + cfg.activation_type = v_cfg.image_mlp_activations + self.act = Activation.build(cfg) + self.w2 = nn.Linear( + v_cfg.image_mlp_dim, + v_cfg.image_emb_dim, + bias=True, + device=config.init_device, + ) + + def reset_parameters(self): + v_cfg = self.config.vision_backbone + nn.init.trunc_normal_(self.w1.weight, std=math.sqrt(1 / v_cfg.image_emb_dim), a=-2.0, b=2.0) + nn.init.trunc_normal_(self.w2.weight, std=math.sqrt(1 / v_cfg.image_mlp_dim), a=-2.0, b=2.0) + nn.init.zeros_(self.w1.bias) + nn.init.zeros_(self.w2.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.w1(x) + x = self.act(x) + x = self.w2(x) + return x + + + +class ResidualAttentionBlock(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + v_cfg = config.vision_backbone + self.attention = MultiHeadDotProductAttention(config) + self.feed_forward = ViTMLP(config) + self.attention_norm = nn.LayerNorm( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + device=config.init_device, + ) + self.ffn_norm = nn.LayerNorm( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + device=config.init_device, + ) + + def reset_parameters(self): + self.attention.reset_parameters() + self.feed_forward.reset_parameters() + self.attention_norm.reset_parameters() + self.ffn_norm.reset_parameters() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attention(self.attention_norm(x)) + x = x + self.feed_forward(self.ffn_norm(x)) + return x + + +class BlockCollection(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + self.grad_checkpointing: bool = False + + v_cfg = config.vision_backbone + self.resblocks = nn.ModuleList([ + ResidualAttentionBlock(config) for _ in range(v_cfg.image_num_layers) + ]) + + def reset_parameters(self): + for r in self.resblocks: + r.reset_parameters() + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + hidden_states = [] + for r in self.resblocks: + x = r(x) + hidden_states.append(x) + return hidden_states + + +class VisionTransformer(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + v_cfg = config.vision_backbone + # class embeddings and positional embeddings + self.scale = v_cfg.image_emb_dim ** -0.5 + self.class_embedding = nn.Parameter( + torch.zeros(v_cfg.image_emb_dim, device=config.init_device), + ) + self.num_prefix_tokens: int = 1 + self.positional_embedding = nn.Parameter( + torch.zeros(v_cfg.image_num_pos, v_cfg.image_emb_dim, device=config.init_device), + ) + + image_patch_size = v_cfg.image_patch_size + self.patch_embedding = nn.Linear( + image_patch_size * image_patch_size * 3, + v_cfg.image_emb_dim, + bias=False, + device=config.init_device, + ) + + self.pre_ln = LayerNormFp32( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + device=config.init_device, + ) + + self.transformer = BlockCollection(config) + + @torch.jit.ignore + def set_grad_checkpointing(self, enable=True): + self.transformer.grad_checkpointing = enable + + def reset_parameters(self): + nn.init.normal_(self.class_embedding, std=self.scale) + nn.init.normal_(self.positional_embedding, std=self.scale) + nn.init.normal_(self.patch_embedding.weight, std=0.02) + self.pre_ln.reset_parameters() + self.transformer.reset_parameters() + + def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: + cls_emb = self.positional_embedding[0:1] + pos_emb = self.positional_embedding[1:] + + pos_emb = pos_emb.reshape( + (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1]) + ) + + (patch_num_0, patch_num_1) = patch_num + + if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1: + # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + # antialias: default True in jax.image.resize + pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2) + pos_emb = F.interpolate( + pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True, + ) + pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0) - att, cache = self.self_attn( - x, - attention_bias=attention_bias, - position_ids=position_ids, - drop_mask=drop_mask, - layer_past=layer_past, - use_cache=use_cache, - ) - x = x + self.dropout(att, drop_mask=drop_mask) - og_x = x - x = self.mlp(x) - x = self.dropout(x, drop_mask=drop_mask) - x = og_x + x + pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1]) + x = x + torch.cat([cls_emb[None, :, :], pos_emb[None, :, :]], dim=1).to(x.dtype) + return x - return x, cache + def forward(self, x: torch.Tensor, patch_num: int = None) -> List[torch.Tensor]: + """ + : param x: (batch_size, num_patch, n_pixels) + """ + if patch_num is None: + patch_num = self.config.vision_backbone.image_num_patch + B, N, D = x.shape + x = self.patch_embedding(x) -class MolmoOutput(NamedTuple): - attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] - """ - Attention keys and values from each block. - """ + # class embeddings and positional embeddings + x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1) + x = self.add_pos_emb(x, patch_num) - hidden_states: Optional[Tuple[torch.Tensor]] - """ - Hidden states from each block. - """ + x = self.pre_ln(x) - last_hidden_states: torch.Tensor + hidden_states = self.transformer(x) + return hidden_states class MultiHeadDotProductAttention(nn.Module): - def __init__(self, config: MolmoConfig, use_bias: bool = True, is_vit_layer: Optional[bool] = True): + def __init__(self, config: FullMolmoConfig, use_bias: bool = True, is_vit_layer: Optional[bool] = True): super().__init__() self.config = config self.use_bias = use_bias @@ -818,25 +1148,25 @@ class MultiHeadDotProductAttention(nn.Module): self.num_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) self.wk = nn.Linear( nlayers * self.embed_dim, self.num_key_value_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) self.wv = nn.Linear( nlayers * self.embed_dim, self.num_key_value_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) self.wo = nn.Linear( self.num_heads * self.head_dim, self.embed_dim, bias=use_bias, device=config.init_device, - ) + ) self.attention_dropout: Optional[Dropout] = None if v_cfg.attention_dropout > 0: self.attention_dropout = Dropout(v_cfg.attention_dropout, broadcast_dims=(0, 1)) @@ -860,6 +1190,7 @@ class MultiHeadDotProductAttention(nn.Module): return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,)) def forward(self, inputs_q: torch.Tensor, inputs_kv: Optional[torch.Tensor] = None) -> torch.Tensor: + if inputs_kv is not None: inputs_k = inputs_kv inputs_v = inputs_kv @@ -882,22 +1213,21 @@ class MultiHeadDotProductAttention(nn.Module): if self.config.float32_attention: xq = xq.to(torch.float) xk = xk.to(torch.float) - xv = xv.to(torch.float) - if self.config.attention_type == AttentionType.direct: + if self.config.attention_type == "direct": attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk) attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype) if self.attention_dropout is not None: attn_weights = self.attention_dropout(attn_weights) attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv) - elif self.config.attention_type == AttentionType.sdpa: + elif self.config.attention_type == "sdpa": attn_output = F.scaled_dot_product_attention( xq.transpose(1, 2).contiguous(), xk.transpose(1, 2).contiguous(), xv.transpose(1, 2).contiguous(), is_causal=False, - dropout_p=self.config.vision_backbone.attention_dropout, + dropout_p=self.config.vision_backbone.attention_dropout ).transpose(1, 2) else: raise NotImplementedError(self.config.attention_type) @@ -912,14 +1242,14 @@ class MultiHeadDotProductAttention(nn.Module): class MultiHeadAttentionPool(nn.Module): def __init__( self, - config: MolmoConfig, + config: FullMolmoConfig, factor: int = 1, use_bias: bool = True, dropout: bool = True, output_layer: bool = True, mean_residual: bool = False, query: str = "mean", - is_vit_layer: Optional[bool] = True, + is_vit_layer: Optional[bool] = True ): super().__init__() self.config = config @@ -947,27 +1277,25 @@ class MultiHeadAttentionPool(nn.Module): self.num_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) self.wk = nn.Linear( nlayers * input_dim, self.num_key_value_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) self.wv = nn.Linear( nlayers * input_dim, self.num_key_value_heads * self.head_dim, bias=use_bias, device=config.init_device, - ) + ) if query == "vector": self.attention_query = nn.Parameter( torch.zeros( - 1, - self.num_key_value_heads * self.head_dim, - device=config.init_device, - ), + 1, self.num_key_value_heads * self.head_dim, device=config.init_device, + ), ) if output_layer: @@ -976,7 +1304,7 @@ class MultiHeadAttentionPool(nn.Module): self.embed_dim, bias=use_bias, device=config.init_device, - ) + ) self.attention_dropout = Dropout(v_cfg.attention_dropout, broadcast_dims=(0, 1)) if dropout: self.residual_dropout = Dropout(v_cfg.residual_dropout) @@ -1005,6 +1333,7 @@ class MultiHeadAttentionPool(nn.Module): return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,)) def forward(self, inputs_kv: torch.Tensor) -> torch.Tensor: + xk, xv = self.wk(inputs_kv), self.wv(inputs_kv) if self.query == "mean": @@ -1051,45 +1380,8 @@ class MultiHeadAttentionPool(nn.Module): return attn_output -class ViTMLP(nn.Module): - def __init__(self, config: MolmoConfig): - super().__init__() - self.config = config - v_cfg = config.vision_backbone - - self.w1 = nn.Linear( - v_cfg.image_emb_dim, - v_cfg.image_mlp_dim, - bias=True, - device=config.init_device, - ) - # Activation function. - cfg = deepcopy(config) - cfg.activation_type = v_cfg.image_mlp_activations - self.act = QuickGELU(cfg) - self.w2 = nn.Linear( - v_cfg.image_mlp_dim, - v_cfg.image_emb_dim, - bias=True, - device=config.init_device, - ) - - def reset_parameters(self): - v_cfg = self.config.vision_backbone - nn.init.trunc_normal_(self.w1.weight, std=math.sqrt(1 / v_cfg.image_emb_dim), a=-2.0, b=2.0) - nn.init.trunc_normal_(self.w2.weight, std=math.sqrt(1 / v_cfg.image_mlp_dim), a=-2.0, b=2.0) - nn.init.zeros_(self.w1.bias) - nn.init.zeros_(self.w2.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.w1(x) - x = self.act(x) - x = self.w2(x) - return x - - class MLP(nn.Module): - def __init__(self, config: MolmoConfig, input_dim: int, dropout: float = 0.0): + def __init__(self, config: FullMolmoConfig, input_dim: int, dropout: float = 0.0): super().__init__() self.config = config self.hidden_size = ( @@ -1102,21 +1394,21 @@ class MLP(nn.Module): self.hidden_size // 2, bias=False, device=config.init_device, - ) + ) self.w2 = nn.Linear( self.hidden_size // 2, config.d_model, bias=False, device=config.init_device, - ) + ) self.w3 = nn.Linear( input_dim, self.hidden_size // 2, bias=False, device=config.init_device, - ) - # `MLP` assume the activation takes two inputs, so it must be a 'llama' version. - self.act = LlamaSwiGLU(config) + ) + # Activation function. + self.act = Activation.build(config) self.dropout = Dropout(dropout) def reset_parameters(self): @@ -1142,165 +1434,12 @@ class Residual(nn.Module): return x + self.submodule(x) -class LayerNormFp32(nn.LayerNorm): - """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back). - Derived from https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/transformer.py. - """ - - def forward(self, x: torch.Tensor) -> torch.Tensor: - orig_type = x.dtype - if self.training: - x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps) - else: - x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) - return x.to(orig_type) - - -class ResidualAttentionBlock(nn.Module): - def __init__(self, config: MolmoConfig): - super().__init__() - self.config = config - - v_cfg = config.vision_backbone - self.attention = MultiHeadDotProductAttention(config) - self.feed_forward = ViTMLP(config) - self.attention_norm = nn.LayerNorm( - v_cfg.image_emb_dim, - eps=v_cfg.image_norm_eps, - device=config.init_device, - ) - self.ffn_norm = nn.LayerNorm( - v_cfg.image_emb_dim, - eps=v_cfg.image_norm_eps, - device=config.init_device, - ) - - def reset_parameters(self): - self.attention.reset_parameters() - self.feed_forward.reset_parameters() - self.attention_norm.reset_parameters() - self.ffn_norm.reset_parameters() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x + self.attention(self.attention_norm(x)) - x = x + self.feed_forward(self.ffn_norm(x)) - return x - - -class BlockCollection(nn.Module): - def __init__(self, config: MolmoConfig): - super().__init__() - self.config = config - - v_cfg = config.vision_backbone - self.resblocks = nn.ModuleList([ResidualAttentionBlock(config) for _ in range(v_cfg.image_num_layers)]) - - def reset_parameters(self): - for r in self.resblocks: - r.reset_parameters() - - def forward(self, x: torch.Tensor) -> List[torch.Tensor]: - hidden_states = [] - for r in self.resblocks: - x = r(x) - hidden_states.append(x) - return hidden_states - - -def _expand_token(token, batch_size: int): - return token.view(1, 1, -1).expand(batch_size, -1, -1) - - -class VisionTransformer(nn.Module): - def __init__(self, config: MolmoConfig): +class OLMoVisionBackbone(nn.Module): + def __init__(self, config: FullMolmoConfig): super().__init__() self.config = config + self.image_vit = VisionTransformer(config) - v_cfg = config.vision_backbone - # class embeddings and positional embeddings - self.scale = v_cfg.image_emb_dim**-0.5 - self.class_embedding = nn.Parameter( - torch.zeros(v_cfg.image_emb_dim, device=config.init_device), - ) - self.num_prefix_tokens: int = 1 - self.positional_embedding = nn.Parameter( - torch.zeros(v_cfg.image_num_pos, v_cfg.image_emb_dim, device=config.init_device), - ) - - image_patch_size = v_cfg.image_patch_size - self.patch_embedding = nn.Linear( - image_patch_size * image_patch_size * 3, - v_cfg.image_emb_dim, - bias=False, - device=config.init_device, - ) - - self.pre_ln = LayerNormFp32( - v_cfg.image_emb_dim, - eps=v_cfg.image_norm_eps, - device=config.init_device, - ) - - self.transformer = BlockCollection(config) - - def reset_parameters(self): - nn.init.normal_(self.class_embedding, std=self.scale) - nn.init.normal_(self.positional_embedding, std=self.scale) - nn.init.normal_(self.patch_embedding.weight, std=0.02) - self.pre_ln.reset_parameters() - self.transformer.reset_parameters() - - def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: - cls_emb = self.positional_embedding[0:1] - pos_emb = self.positional_embedding[1:] - - pos_emb = pos_emb.reshape( - (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1]) - ) - - (patch_num_0, patch_num_1) = patch_num - - if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1: - # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py - # antialias: default True in jax.image.resize - pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2) - pos_emb = F.interpolate( - pos_emb, - size=(patch_num_0, patch_num_1), - mode="bicubic", - align_corners=False, - antialias=True, - ) - pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0) - - pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1]) - x = x + torch.cat([cls_emb[None, :, :], pos_emb[None, :, :]], dim=1).to(x.dtype) - return x - - def forward(self, x: torch.Tensor, patch_num: int = None) -> List[torch.Tensor]: - """ - : param x: (batch_size, num_patch, n_pixels) - """ - if patch_num is None: - patch_num = self.config.vision_backbone.image_num_patch - B, N, D = x.shape - - x = self.patch_embedding(x) - - # class embeddings and positional embeddings - x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1) - x = self.add_pos_emb(x, patch_num) - - x = self.pre_ln(x) - - hidden_states = self.transformer(x) - return hidden_states - - -class MolmoVisionBackbone(nn.Module): - def __init__(self, config: VisionBackboneConfig): - super().__init__() - self.config = config input_dim: int = None self.image_pooling_2d: nn.Module = None if config.image_pooling_2d in {ImagePooling2DType.attention, ImagePooling2DType.attention_meanq}: @@ -1340,16 +1479,31 @@ class MolmoVisionBackbone(nn.Module): self.input_dim = input_dim - self.image_projector = MLP(config, input_dim) + # `MLP` assume the activation takes two inputs, so it must be a 'llama' version + if config.activation_type == ActivationType.swiglu: + mlp_config = replace(config, activation_type=ActivationType.llama_swiglu) + elif config.activation_type == ActivationType.gelu: + mlp_config = replace(config, activation_type=ActivationType.llama_geglu) + else: + mlp_config = config + if config.image_projector == ImageProjectType.mlpx2: + self.image_projector = nn.ModuleList( + [MLP(mlp_config, input_dim), Residual(MLP(config, input_dim))] + ) + elif config.image_projector == ImageProjectType.mlp: + self.image_projector = MLP(mlp_config, input_dim) + elif config.image_projector == ImageProjectType.linear: + self.image_projector = nn.Linear( + input_dim, + config.d_model, + bias=False, + device=config.init_device, + ) + else: + raise NotImplementedError(f"Unknown image projector: {config.image_projector}") self.image_feature_dropout = Dropout(config.image_feature_dropout) - @classmethod - def build(cls, config: MolmoConfig) -> MolmoVisionBackbone: - v_cfg = config.vision_backbone - assert v_cfg is not None - return MolmoPretrainedVisionBackbone(config) - def reset_parameters(self): if self.image_pooling_2d is not None: self.image_pooling_2d.reset_parameters() @@ -1361,22 +1515,15 @@ class MolmoVisionBackbone(nn.Module): else: self.image_projector.reset_parameters() - @abstractmethod - def forward( - self, images: torch.Tensor, image_masks: torch.Tensor - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, images: torch.Tensor, image_masks: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: raise NotImplementedError -class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): - def __init__(self, config: MolmoConfig): +class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): + def __init__(self, config: FullMolmoConfig): super().__init__(config) v_cfg = self.config.vision_backbone - - if v_cfg.image_model_type == VisionBackboneType.openai: - self.image_vit = VisionTransformer(config) - else: - raise NotImplementedError(f"Unknown image model type: {v_cfg.image_model_type}") + self.grad_checkpointing = False self.num_prefix_tokens = self.image_vit.num_prefix_tokens assert self.num_prefix_tokens in {0, 1}, "Only 0 or 1 prefix tokens are supported" @@ -1388,37 +1535,20 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): self.input_dim, bias=False, device=config.init_device, - ) + ) self.pad_embed = None if config.image_padding_embed: - image_dim = v_cfg.image_emb_dim * len(self.config.vit_layers) + image_dim = v_cfg.image_emb_dim*len(self.config.vit_layers) if config.image_padding_embed in ["pad_embed", "regress"]: - self.pad_embed = nn.Parameter(torch.zeros((image_dim,), device=config.init_device)) + self.pad_embed = nn.Parameter( + torch.zeros((image_dim,), device=config.init_device)) elif config.image_padding_embed == "pad_and_partial_pad": - self.pad_embed = nn.Parameter(torch.zeros((2, image_dim), device=config.init_device)) + self.pad_embed = nn.Parameter( + torch.zeros((2, image_dim), device=config.init_device)) else: raise ValueError(config.image_padding_embed) - def reset_with_pretrained_weights(self): - super().reset_parameters() # resets the connector - if self.config.vit_load_path: - vit_load_path = Path(self.config.vit_load_path) - state_dict_path = resource_path( - vit_load_path.parent, - vit_load_path.name, - local_cache=vit_load_path.parent, - ) - assert state_dict_path.is_file(), f"Model file {str(state_dict_path)} not found" - state_dict = torch.load(state_dict_path, map_location="cpu") - self.image_vit.load_state_dict(state_dict) - else: - self.image_vit.reset_parameters() - if self.config.use_cls_feature: - nn.init.xavier_uniform_(self.cls_projector.weight) - if self.pad_embed is not None: - nn.init.zeros_(self.pad_embed) - def reset_parameters(self): super().reset_parameters() self.image_vit.reset_parameters() @@ -1433,7 +1563,7 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): v_cfg = self.config.vision_backbone B, T, N, D = images.shape - mask = torch.logical_not(torch.all(images.view(B * T, N, D) == -1, dim=(1, 2), keepdim=True)) + mask = torch.all(images.view(B * T, N, D) != -1, dim=(1, 2), keepdim=True) # Output all hidden states # n_layers x (batch_num_crops, (1+)n_tokens, image_emb_dim) @@ -1460,16 +1590,13 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): return image_features, cls_embed - def forward( - self, images: torch.Tensor, image_masks: torch.Tensor - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + def forward(self, images: torch.Tensor, image_masks: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: cfg = self.config # image_features: (batch_size, num_crops(=num_image), num_patch, nximage_emb_dim) batch_size, num_image = images.shape[:2] image_features, cls_embed = self.encode_image(images) - og_dtype = image_features.dtype if cfg.image_padding_embed: assert image_masks is not None if cfg.image_padding_embed == "pad_embed": @@ -1478,32 +1605,26 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): image_features = image_features + pad_embed * torch.unsqueeze(all_pad, -1) elif cfg.image_padding_embed == "regress": pad_embed = self.pad_embed[None, None, None, :] - image_features = image_features + pad_embed * torch.unsqueeze( - torch.maximum(image_masks, torch.zeros_like(image_masks)), -1 - ) + image_features = image_features + pad_embed * torch.unsqueeze(torch.maximum(image_masks, torch.zeros_like(image_masks)), -1) elif cfg.image_padding_embed == "pad_and_partial_pad": - og_dtype = image_features.dtype pad_embed = self.pad_embed[:, None, None, None, :] all_pad = image_masks == 0 - partial_pad = torch.logical_and(image_masks < 1, torch.logical_not(all_pad)).to( - dtype=torch.float32 - ) + partial_pad = torch.logical_and(image_masks < 1, torch.logical_not(all_pad)).to(dtype=torch.float32) all_pad = all_pad.to(dtype=torch.float32) image_features = image_features + pad_embed[0] * torch.unsqueeze(all_pad, -1) image_features = image_features + pad_embed[1] * torch.unsqueeze(partial_pad, -1) else: raise ValueError(cfg.image_padding_embed) - image_features = image_features.to(og_dtype) image_features = self.image_feature_dropout(image_features) if cls_embed is not None: cls_embed = self.image_feature_dropout(cls_embed) image_features = image_features.reshape( - (batch_size, num_image) + cfg.vision_backbone.image_num_patch + (-1,), - ) + (batch_size, num_image) + cfg.image_num_patch + (-1,), + ) - if cfg.vision_backbone.image_num_patch[0] % cfg.image_pooling_h == 1: + if cfg.image_num_patch[0] % cfg.image_pooling_h == 1: # Pad so we can still pool 2x2 patches image_features = F.pad( image_features, @@ -1513,7 +1634,7 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): # image pooling image_features = einops.rearrange( image_features, - "b n (h dh) (w dw) c -> (b n h w) (dh dw) c", + 'b n (h dh) (w dw) c -> (b n h w) (dh dw) c', dh=cfg.image_pooling_h, dw=cfg.image_pooling_w, ) @@ -1521,75 +1642,277 @@ class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): if cfg.image_pooling_2d == ImagePooling2DType.attention_meanq: query = image_features.mean(-2, keepdim=True) image_features = self.image_pooling_2d(query, image_features) - elif cfg.image_pooling_2d == ImagePooling2DType.attention_v2: - image_features = self.image_pooling_2d(image_features) elif cfg.image_pooling_2d not in {ImagePooling2DType.none, ImagePooling2DType.stack}: - image_features = self.image_pooling_2d(image_features[:, :1, :], image_features) + if self.grad_checkpointing: + from torch.utils.checkpoint import checkpoint + image_features = checkpoint(self.image_pooling_2d, image_features[:, :1, :], image_features, use_reentrant=False) + else: + image_features = self.image_pooling_2d(image_features[:, :1, :], image_features) - h, w = cfg.llm_patches_per_crop + h, w = cfg.llm_patches_per_crop() image_features = image_features.reshape(batch_size, num_image, h * w, -1) # MLP layer to map the feature. - if cfg.image_projector == ImageProjectType.mlpx2: - for module in self.image_projector: - image_features = module(image_features) + if self.grad_checkpointing: + from torch.utils.checkpoint import checkpoint + image_features = checkpoint(self.image_projector, image_features, use_reentrant=False) else: image_features = self.image_projector(image_features) if self.config.use_cls_feature: - cls_embed = self.cls_projector(cls_embed) - if cfg.image_projector == ImageProjectType.mlpx2: - for module in self.image_projector: - cls_embed = module(cls_embed) - else: - cls_embed = self.image_projector(cls_embed) + raise NotImplementedError() # image_features: (batch_size, num_image, num_patch, d_model) # cls_embed: (batch_size, num_image, d_model) return image_features, cls_embed -class MolmoPretrainedModel(PreTrainedModel): - config_class = MolmoConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["MolmoDecoderLayer"] - _skip_keys_device_placement = ["past_key_values"] - _supports_flash_attn_2 = True - _supports_sdpa = True - _supports_cache_class = True - _supports_quantized_cache = True - _supports_static_cache = True - - def _init_weights(self, module): - if self.vision_backbone is not None: - self.vision_backbone.reset_parameters() - self.reset_non_vision_parameters() +class ModuleType(str, Enum): + in_module = "in" + out_module = "out" + emb = "emb" + final_out = "final_out" + + +def init_weights( + config: FullMolmoConfig, + module: Union[nn.Linear, nn.Embedding], + d: Optional[int] = None, + layer_id: Optional[int] = None, + std_factor: float = 1.0, + type_of_module: Optional[ModuleType] = None, +) -> None: + d = d if d is not None else config.d_model + std = config.init_std * std_factor + if config.init_cutoff_factor is not None: + cutoff_value = config.init_cutoff_factor * std + nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value) + else: + nn.init.normal_(module.weight, mean=0.0, std=std) -class MolmoModel(MolmoPretrainedModel): - def __init__(self, config: MolmoConfig, init_params: bool = True): - super().__init__(config) +class LlamaSwiGLU(nn.Module): + def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: + return F.silu(x1) * x2 + + @property + def output_multiplier(self) -> float: + return 0.5 + + +class SwiGLU(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, gate = x.chunk(2, dim=-1) + return F.silu(gate) * x + + @property + def output_multiplier(self) -> float: + return 0.5 + + +class Activation(nn.Module): + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + @property + def output_multiplier(self) -> float: + raise NotImplementedError + + @classmethod + def build(cls, config: FullMolmoConfig) -> 'Activation': + if config.activation_type == "quick_gelu": + return QuickGELU(config) + elif config.activation_type == "gelu": + return cast(Activation, GELU(approximate="none")) + elif config.activation_type == "gelu_tanh": + return cast(Activation, GELU(approximate="tanh")) + elif config.activation_type == "relu": + return cast(Activation, ReLU(inplace=False)) + elif config.activation_type == "silu": + return cast(Activation, SiLU(inplace=False)) + # elif config.activation_type == "llama_geglu": + # return LlamaGEGLU(config) + # elif config.activation_type == "llama_geglu_tanh": + # return LlamaGEGLUTanh(config) + elif config.activation_type == "llama_swiglu": + return LlamaSwiGLU() + elif config.activation_type == "swiglu": + return SwiGLU() + else: + raise NotImplementedError(f"Unknown activation: '{config.activation_type}'") + + +class QuickGELU(Activation): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(1.702 * x) + + @property + def output_multiplier(self) -> float: + return 1.0 + + +class GELU(nn.GELU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +class ReLU(nn.ReLU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +class SiLU(nn.SiLU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +def causal_attention_bias(seq_len: int, device: torch.device) -> torch.FloatTensor: + att_bias = torch.triu( + torch.ones(seq_len, seq_len, device=device, dtype=torch.float), + diagonal=1, + ) + att_bias.masked_fill_(att_bias == 1, torch.finfo(att_bias.dtype).min) + return att_bias.view(1, 1, seq_len, seq_len) # type: ignore + + +def get_causal_attention_bias(cache: BufferCache, seq_len: int, device: torch.device) -> torch.Tensor: + if (causal_bias := cache.get("causal_attention_bias")) is not None and causal_bias.shape[-1] >= seq_len: + if causal_bias.device != device: + causal_bias = causal_bias.to(device) + cache["causal_attention_bias"] = causal_bias + return causal_bias + with torch.autocast(device.type, enabled=False): + causal_bias = causal_attention_bias(seq_len, device) + cache["causal_attention_bias"] = causal_bias + return causal_bias + + +class LayerNormBase(nn.Module): + def __init__( + self, + config: MolmoConfig, + *, + size: Optional[int] = None, + elementwise_affine: Optional[bool] = True, + eps: float = 1e-05, + weight_initializer: Optional[Callable] = torch.ones, + bias_initializer: Optional[Callable] = torch.zeros, + ): + super().__init__() + self.config = config + self.eps = self.config.layer_norm_eps or eps + self.normalized_shape = (size or config.d_model,) + if elementwise_affine or (elementwise_affine is None and self.config.layer_norm_with_affine): + self.weight = nn.Parameter(weight_initializer(self.normalized_shape, device=config.init_device)) + use_bias = self.config.bias_for_layer_norm + if use_bias is None: + use_bias = self.config.include_bias + if use_bias: + self.bias = nn.Parameter(bias_initializer(self.normalized_shape, device=config.init_device)) + else: + self.register_parameter("bias", None) + else: + self.register_parameter("bias", None) + self.register_parameter("weight", None) + + @classmethod + def build(cls, config: FullMolmoConfig, size: Optional[int] = None, **kwargs): + if config.layer_norm_type == "default": + return LayerNorm(config, size=size, low_precision=False, **kwargs) + elif config.layer_norm_type == "low_precision": + return LayerNorm(config, size=size, low_precision=True, **kwargs) + elif config.layer_norm_type == "rms": + return RMSLayerNorm(config, size=size, **kwargs) + else: + raise NotImplementedError(f"Unknown LayerNorm type: '{config.layer_norm_type}'") + + +class RMSLayerNorm(LayerNormBase): + """ + RMS layer norm, a simplified :class:`LayerNorm` implementation + """ + + def __init__( + self, + config: FullMolmoConfig, + size: Optional[int] = None, + elementwise_affine: Optional[bool] = None, + eps: float = 1e-5, + ): + super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + with torch.autocast(enabled=False, device_type=x.device.type): + og_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + x = x.to(og_dtype) + + if self.weight is not None: + if self.bias is not None: + return self.weight * x + self.bias + else: + return self.weight * x + else: + return x + + +class LayerNorm(LayerNormBase): + """ + The default :class:`LayerNorm` implementation which can optionally run in low precision. + """ + + def __init__( + self, + config: FullMolmoConfig, + size: Optional[int] = None, + low_precision: bool = False, + elementwise_affine: Optional[bool] = None, + eps: float = 1e-05, + ): + super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) + self.low_precision = low_precision + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.low_precision: + module_device = x.device + downcast_x = self._cast_if_autocast_enabled(x) + downcast_weight = ( + self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight + ) + downcast_bias = self._cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias + with torch.autocast(enabled=False, device_type=module_device.type): + return F.layer_norm( + downcast_x, self.normalized_shape, weight=downcast_weight, bias=downcast_bias, eps=self.eps + ) + else: + return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps) + + +class MOLMo(nn.Module): + def __init__(self, config: FullMolmoConfig, init_params: bool = True): + super().__init__() self.config = config self.__cache = BufferCache() # Validate config. if self.config.embedding_size is not None and self.config.embedding_size != self.config.vocab_size: if self.config.embedding_size < self.config.vocab_size: - raise MolmoConfigurationError("embedding size should be at least as big as vocab size") + raise OLMoConfigurationError("embedding size should be at least as big as vocab size") elif self.config.embedding_size % 128 != 0: import warnings warnings.warn( "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning ) - - if not ( - 0 < self.config.block_group_size <= self.config.n_layers - and self.config.n_layers % self.config.block_group_size == 0 - ): - raise MolmoConfigurationError("n layers must be divisible by block group size") - torch.backends.cuda.enable_flash_sdp(True) torch.backends.cuda.enable_mem_efficient_sdp(False) # this is super slow so make sure torch won't use it @@ -1601,10 +1924,10 @@ class MolmoModel(MolmoPretrainedModel): config.d_model, device=config.init_device, initializer_range=config.initializer_range, - new_embed_initializer_range=config.new_embedding_init_range, + new_embed_initializer_range=config.new_embedding_init_range ) else: - wte = nn.Embedding( + wte=nn.Embedding( config.embedding_size or config.vocab_size, config.d_model, device=config.init_device ) @@ -1612,27 +1935,62 @@ class MolmoModel(MolmoPretrainedModel): dict( wte=wte, emb_drop=Dropout(config.embedding_dropout), - ln_f=RMSLayerNorm(config, size=config.d_model, eps=config.layer_norm_eps), + ln_f=LayerNorm.build(config), ) ) - layers = [MolmoDecoderLayer(i, config, self.__cache) for i in range(config.n_layers)] - self.transformer.update({"layers": nn.ModuleList(layers)}) + blocks = [OLMoBlock.build(i, config, self.__cache) for i in range(config.n_layers)] + if self.config.block_group_size > 1: + raise NotImplementedError() + else: + self.transformer.update({"blocks": nn.ModuleList(blocks)}) + + if not (self.config.alibi or self.config.rope): + self.transformer.update( + {"wpe": nn.Embedding(config.max_sequence_length, config.d_model, device=config.init_device)} + ) + if not config.weight_tying: + self.transformer.update( + { + "ff_out": nn.Linear( + config.d_model, + config.embedding_size or config.vocab_size, + bias=config.include_bias, + device=config.init_device, + ) + } + ) - self.vision_backbone: Optional[MolmoVisionBackbone] = None + self.vision_backbone: Optional[OLMoVisionBackbone] = None if config.vision_backbone is not None: - self.vision_backbone = MolmoVisionBackbone.build(config) + self.vision_backbone = OLMoPretrainedVisionBackbone(config) + + self.__num_fwd_flops: Optional[int] = None + def reset_parameters(self): if self.vision_backbone is not None: - self.vision_backbone.reset_with_pretrained_weights() + self.vision_backbone.reset_parameters() + self.reset_non_vision_parameters() - @property - def device(self) -> torch.device: - device: torch.device = self.transformer.wte.weight.device # type: ignore - if device.type == "meta": - return _non_meta_init_device(self.config) + def reset_non_vision_parameters(self): + self.transformer.wte.reset_parameters() + if hasattr(self.transformer.wte, "new_embedding"): + nn.init.normal_(self.transformer.wte.new_embedding, std=self.config.new_embedding_init_range) + + if hasattr(self.transformer, "wpe"): + nn.init.normal_(self.transformer.wpe, mean=0.0, std=1.0) + + self.transformer.ln_f.reset_parameters() # type: ignore + + if hasattr(self.transformer, "ff_out"): + nn.init.normal_(self.transformer.ff_out, mean=0.0, std=0.02) + + if self.config.block_group_size == 1: + for block in self.transformer.blocks: + block.reset_parameters() else: - return device + for block_group in self.transformer.block_groups: + block_group.reset_parameters() def forward( self, @@ -1651,7 +2009,7 @@ class MolmoModel(MolmoPretrainedModel): last_logits_only: bool = False, output_hidden_states: Optional[bool] = None, append_last_valid_logits: Optional[torch.Tensor] = None, - ) -> MolmoOutput: + ) -> ModelOutput: """ :param input_ids: A tensor of shape `(batch_size, seq_len)`. :param input_embeddings: A tensor of shape `(batch_size, seq_len, d_model)` with input @@ -1660,20 +2018,16 @@ class MolmoModel(MolmoPretrainedModel): which input IDs are masked. A `1` value in the mask means that the corresponding input ID should *not* be ignored. A `0` means that the corresponding input ID is masked. - This has the same meaning as the `attention_mask` in HuggingFace's `transformers` library. :param attention_bias: A tensor of shape `(batch_size, 1, seq_len, seq_len)`, `(1, 1, seq_len, seq_len)`, or `(seq_len, seq_len)`. This is used to introduce causal or other biases. - If the tensor is a bool or byte tensor, a `True` or `1` at `attention_bias[:, :, i, j]` indicates that the i-th element in the sequence is allowed to attend to the j-th element in the sequence. - If the tensor is a float tensor, it will just be added to the attention scores before the softmax. - The default is causal, which corresponds to a lower-diagonal byte matrix of ones. :param response_mask: A tensor of shape `(batch_size, seq_len)` that indicates the response mask. A `1` value in the mask means that the corresponding token @@ -1694,9 +2048,7 @@ class MolmoModel(MolmoPretrainedModel): has_image = images is not None assert not (has_image and input_embeddings is not None), "Cannot provide both images and input embeddings." - assert not ( - has_image and past_key_values is not None - ), "Cached key and values should not be used with images." + assert not (has_image and past_key_values is not None), "Cached key and values should not be used with images." batch_size, seq_len = input_ids.size() if input_embeddings is None else input_embeddings.size()[:2] if past_key_values is None: @@ -1704,10 +2056,6 @@ class MolmoModel(MolmoPretrainedModel): else: past_length = past_key_values[0][0].size(-2) - if self.config.unconditioned and input_embeddings is None: - images = None - image_input_idx = None - if self.config.use_position_ids and attention_mask is None: attention_mask = input_ids != -1 @@ -1715,18 +2063,17 @@ class MolmoModel(MolmoPretrainedModel): assert not use_cache, "Subsegment_ids cannot be used with cache." subsegment_mask = subsegment_ids.unsqueeze(2) <= subsegment_ids.unsqueeze(1) attention_mask = ( - subsegment_mask.to(attention_mask.dtype) - * attention_mask.unsqueeze(2) - * attention_mask.unsqueeze(1) - ) + subsegment_mask.to(attention_mask.dtype) * + attention_mask.unsqueeze(2) * + attention_mask.unsqueeze(1)) if position_ids is None: - raise ValueError("Positioned ids must be given if using subsegment_ids") + raise ValueError(f"Positioned ids must be given if using subsegment_ids") else: if self.config.use_position_ids and position_ids is None: position_ids = torch.clamp( torch.cumsum(attention_mask.to(torch.int32), dim=-1) - 1, min=0, - ).broadcast_to((batch_size, attention_mask.shape[-1])) + ).broadcast_to((batch_size, attention_mask.shape[-1])) # Get embeddings of input. # shape: (batch_size, seq_len, d_model) @@ -1758,7 +2105,9 @@ class MolmoModel(MolmoPretrainedModel): if self.config.use_cls_feature: x = torch.cat([x[:, :1], cls_embed, x[:, 1:-num_image]], dim=1) - valid_images = torch.any((image_input_idx >= 0).view(batch_size, num_image, num_patch), dim=-1) + valid_images = torch.any( + (image_input_idx >= 0).view(batch_size, num_image, num_patch), dim=-1 + ) valid_images = valid_images.to(attention_mask.dtype) attention_mask = torch.cat( [attention_mask[:, :1], valid_images, attention_mask[:, 1:-num_image]], @@ -1767,7 +2116,15 @@ class MolmoModel(MolmoPretrainedModel): position_ids = torch.clamp( torch.cumsum(attention_mask, dim=-1) - 1, min=0, - ).broadcast_to((batch_size, attention_mask.shape[-1])) + ).broadcast_to((batch_size, attention_mask.shape[-1])) + + if not (self.config.alibi or self.config.rope): + # Get positional embeddings. + # shape: (1, seq_len) + pos = torch.arange(past_length, past_length + seq_len, dtype=torch.long, device=x.device).unsqueeze(0) + # shape: (1, seq_len, d_model) + pos_emb = self.transformer.wpe(pos) # type: ignore + x = pos_emb + x # Add input + positional embeddings and apply dropout. # shape: (batch_size, seq_len, d_model) @@ -1775,13 +2132,13 @@ class MolmoModel(MolmoPretrainedModel): # normalized if self.config.normalize_input_embeds: - x = x * (self.config.d_model**0.5) + x = x * (self.config.d_model ** 0.5) # Transform the attention mask into what the blocks expect. if attention_mask is not None: # shape: (batch_size, 1, 1, seq_len) if len(attention_mask.shape) == 2: - attention_mask = attention_mask[:, : past_length + seq_len] + attention_mask = attention_mask[:, :past_length + seq_len] attention_mask = attention_mask.to(dtype=torch.float).view(batch_size, -1)[:, None, None, :] else: attention_mask = attention_mask.unsqueeze(1).to(dtype=torch.float) @@ -1791,12 +2148,17 @@ class MolmoModel(MolmoPretrainedModel): if ( attention_bias is not None or attention_mask is not None + or self.config.alibi # NOTE (epwalsh): we need to initialize the attn bias in order for attn to work properly # with key+value cache. Otherwise `F.scaled_dot_product_attention()` doesn't seem to compute # scores correctly. or past_key_values is not None ): - if attention_bias is None: + if attention_bias is None and self.config.alibi: + attention_bias = get_causal_attention_bias( + self.__cache, past_length + seq_len, x.device + ) + self.get_alibi_attention_bias(past_length + seq_len, x.device) + elif attention_bias is None: attention_bias = get_causal_attention_bias(self.__cache, past_length + seq_len, x.device) elif attention_bias.dtype in (torch.int8, torch.bool): attention_bias = attention_bias.to(dtype=torch.float) @@ -1824,30 +2186,42 @@ class MolmoModel(MolmoPretrainedModel): all_hidden_states = [] # Apply blocks one-by-one. - for block_idx, layer in enumerate(self.transformer.layers): - if output_hidden_states: - # add hidden states - all_hidden_states.append(x) - - layer_past = None if past_key_values is None else past_key_values[block_idx] - # shape: (batch_size, seq_len, d_model) - x, cache = layer( - x, - attention_bias=attention_bias, - position_ids=position_ids, - drop_mask=response_mask, - layer_past=layer_past, - use_cache=use_cache, - ) - - if attn_key_values is not None: - assert cache is not None - attn_key_values.append(cache) + if self.config.block_group_size == 1: + for block_idx, block in enumerate(self.transformer.blocks): + if output_hidden_states: + # add hidden states + all_hidden_states.append(x) + + layer_past = None if past_key_values is None else past_key_values[block_idx] + x, cache = block(x, attention_bias=attention_bias, position_ids=position_ids, drop_mask=response_mask, layer_past=layer_past, use_cache=use_cache) + + if attn_key_values is not None: + assert cache is not None + attn_key_values.append(cache) + else: + for group_idx, block_group in enumerate(self.transformer.block_groups): + if output_hidden_states: + # add hidden states + all_hidden_states.append(x) + + layers_past = ( + None + if past_key_values is None + else past_key_values[ + group_idx * self.config.block_group_size : (group_idx + 1) * self.config.block_group_size + ] + ) + x, cache = block_group( + x, attention_bias=attention_bias, position_ids=position_ids, drop_mask=response_mask, layers_past=layers_past, use_cache=use_cache + ) + if attn_key_values is not None: + assert cache is not None + attn_key_values.extend(cache) if images is not None and self.config.use_cls_feature: assert num_image is not None x = torch.cat( - [x[:, :1], x[:, num_image + 1 :], torch.zeros_like(x[:, :num_image])], + [x[:, :1], x[:, num_image+1:], torch.zeros_like(x[:, :num_image])], dim=1, ) @@ -1855,8 +2229,7 @@ class MolmoModel(MolmoPretrainedModel): # shape: (batch_size, 1, d_model) if append_last_valid_logits is not None: last_valid_output = x[ - torch.arange(x.shape[0], device=x.device), append_last_valid_logits.to(x.device) - ] + torch.arange(x.shape[0], device=x.device), append_last_valid_logits.to(x.device)] x = last_valid_output.unsqueeze(1) else: x = x[:, -1, :].unsqueeze(1) @@ -1870,40 +2243,78 @@ class MolmoModel(MolmoPretrainedModel): # Get logits. # shape: (batch_size, seq_len or 1, vocab_size) - return MolmoOutput( - last_hidden_states=x, - attn_key_values=attn_key_values, - hidden_states=tuple(all_hidden_states) if output_hidden_states else None, - ) + if self.config.weight_tying: + logits = F.linear(x, self.transformer.wte.weight, None) # type: ignore + else: + logits = self.transformer.ff_out(x) # type: ignore + if self.config.scale_logits: + logits.mul_(1 / math.sqrt(self.config.d_model)) + + if not last_logits_only and append_last_valid_logits is not None: + last_valid_logit = logits[ + torch.arange(logits.shape[0], device=logits.device), append_last_valid_logits] + logits = torch.cat([logits[:, :-1], last_valid_logit[:, None]], dim=1) + return ModelOutput(logits=logits, attn_key_values=attn_key_values, hidden_states=tuple(all_hidden_states) if output_hidden_states else None) # type: ignore[arg-type] -class MolmoForCausalLM(PreTrainedModel): - """ - Extremely barebones HF model wrapper. - """ +class MOLMoForCausalLM(PreTrainedModel): config_class = MolmoConfig base_model_prefix = "model" - _no_split_modules = ["MolmoDecoderLayer"] + _no_split_modules = ["OLMoBlock"] - def __init__(self, config: MolmoConfig): + def __init__(self, config: MolmoConfig, model: Optional[MOLMo] = None, init_params: bool = False): super().__init__(config) - # model_config = create_model_config_from_pretrained_config(config) - # Initialize model (always on CPU to start with so we don't run out of GPU memory). - config.init_device = "cpu" - v_cfg = config.vision_backbone - if v_cfg is not None: - v_cfg = VisionBackboneConfig(**v_cfg) - config.vision_backbone = v_cfg - self.model = MolmoModel(config) - if not config.weight_tying: - self.lm_head = nn.Linear( - config.d_model, - config.embedding_size or config.vocab_size, - bias=config.include_bias, - device=config.init_device, + if not model: + full_config = FullMolmoConfig( + rope_impl="llama", + vocab_size=config.vocab_size, + max_sequence_length=config.max_position_embeddings, + qkv_bias=config.qkv_bias, + embedding_size=config.embedding_size, + attention_type="sdpa", + embedding_dropout=0, + response_residual_dropout=0, + attention_dropout=0, + residual_dropout=0, + rope=True, + weight_tying=False, + include_bias=False, + d_model=config.hidden_size, + mlp_hidden_size=config.intermediate_size, + n_layers=config.num_hidden_layers, + additional_vocab_size=128, + n_heads=config.num_attention_heads, + n_kv_heads=config.num_key_value_heads, + rope_theta=1000000.0, + layer_norm_eps=1e-6, + layer_norm_type="rms", + pad_tokenizer=True, + vit_layers=[-2, -9], + vision_backbone=VisionBackboneConfig( + image_model_type="openai", + image_default_input_size=(336, 336), + image_patch_size=14, + image_pos_patch_size=14, + image_emb_dim=1024, + image_num_heads=16, + image_num_key_value_heads=16, + image_num_layers=23, + image_head_dim=64, + image_mlp_dim=4096, + image_mlp_activations="quick_gelu", + image_dropout_rate=0.0, + image_num_pos=577, + image_norm_eps=1e-5, + attention_dropout=0.0, + residual_dropout=0.0, + initializer_range=0.02, + ) ) + self.model = MOLMo(full_config, init_params=init_params) + else: + self.model = model def forward( self, @@ -1927,19 +2338,19 @@ class MolmoForCausalLM(PreTrainedModel): append_last_valid_logits: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None, cache_position: Optional[ - torch.Tensor + Cache ] = None, # This is a hack mitigation of an issue in transformers `4.39.x` https://github.com/huggingface/transformers/issues/29426 ) -> Union[Tuple, CausalLMOutputWithPast]: if use_cache is None: use_cache = self.config.use_cache if output_attentions: - raise ValueError("output_attentions is not yet supported in Molmo") + raise ValueError("output_attentions is not yet supported in OLMo") return_dict = return_dict if return_dict is not None else self.config.use_return_dict # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( + outputs = self.model.forward( input_ids=input_ids, input_embeddings=inputs_embeds, attention_mask=attention_mask, @@ -1957,25 +2368,8 @@ class MolmoForCausalLM(PreTrainedModel): append_last_valid_logits=append_last_valid_logits, ) - x = outputs.last_hidden_states - if self.config.weight_tying: - logits = F.linear(x, self.model.transformer.wte.weight, None) # type: ignore - else: - logits = self.lm_head(x) # type: ignore - - if self.config.scale_logits: - logits.mul_(1 / math.sqrt(self.config.d_model)) - - if self.config.final_logit_softcapping is not None: - logits = logits / self.config.final_logit_softcapping - logits = torch.tanh(logits) - logits = logits * self.config.final_logit_softcapping - - if not last_logits_only and append_last_valid_logits is not None: - last_valid_logit = logits[ - torch.arange(logits.shape[0], device=logits.device), append_last_valid_logits - ] - logits = torch.cat([logits[:, :-1], last_valid_logit[:, None]], dim=1) + logits = outputs.logits + hidden_states = outputs.hidden_states loss = None if labels is not None: @@ -1986,7 +2380,7 @@ class MolmoForCausalLM(PreTrainedModel): labels.masked_fill_(~(loss_masks > 0), -100) labels = labels.view(-1) logits_for_loss = logits.to(torch.float32).view(-1, logits.size(-1)) - loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction="none") + loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none') loss = loss_fct(logits_for_loss, labels) loss = loss.view(input_ids.shape[0], -1) loss = loss * loss_masks @@ -2019,28 +2413,30 @@ class MolmoForCausalLM(PreTrainedModel): loss=loss, logits=logits, past_key_values=outputs.attn_key_values, - hidden_states=outputs.hidden_states, + hidden_states=hidden_states, ) def can_generate(self) -> bool: return True @torch.no_grad() - def generate( + def generate_from_batch( self, - input_ids, - images=None, - attention_mask=None, - image_masks=None, - image_input_idx=None, - generation_config=None, + batch: Dict[str, Any], + generation_config: Optional[GenerationConfig] = None, **kwargs, ): if generation_config is not None: assert generation_config.use_cache + images = batch.get("images") + image_masks = batch.get("image_masks") + image_input_idx = batch.get("image_input_idx") + # Validate inputs. + input_ids = batch["input_ids"] batch_size, seq_len = input_ids.shape + attention_mask = batch.get("attention_mask", None) max_new_tokens = generation_config.max_new_tokens assert max_new_tokens is not None mask_len = seq_len + max_new_tokens if self.config.use_position_ids else seq_len @@ -2048,7 +2444,10 @@ class MolmoForCausalLM(PreTrainedModel): append_last_valid_logits: Optional[torch.Tensor] = None if self.config.use_position_ids and attention_mask is None: attention_mask = input_ids != -1 - position_ids = torch.clamp(torch.cumsum(attention_mask.to(torch.int32), dim=-1) - 1, min=0) + position_ids = torch.clamp( + torch.cumsum(attention_mask.to(torch.int32), dim=-1) - 1, + min=0 + ) append_last_valid_logits = attention_mask.long().sum(dim=-1) - 1 attention_mask = torch.cat( [attention_mask, attention_mask.new_ones((batch_size, max_new_tokens))], @@ -2058,7 +2457,7 @@ class MolmoForCausalLM(PreTrainedModel): assert attention_mask.shape == (batch_size, mask_len) out = super().generate( - input_ids, + batch["input_ids"], generation_config, attention_mask=attention_mask, images=images, @@ -2110,6 +2509,7 @@ class MolmoForCausalLM(PreTrainedModel): outputs: ModelOutput, model_kwargs: Dict[str, Any], is_encoder_decoder: bool = False, + standardize_cache_format: bool = False, num_new_tokens: int = 1, ) -> Dict[str, Any]: if self.config.use_position_ids: @@ -2120,9 +2520,8 @@ class MolmoForCausalLM(PreTrainedModel): del model_kwargs["images"] del model_kwargs["image_masks"] del model_kwargs["image_input_idx"] - cache_name, cache = super()._extract_past_from_model_output(outputs) - model_kwargs[cache_name] = cache - model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens + model_kwargs = super()._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder, standardize_cache_format, num_new_tokens) return model_kwargs # TODO: these are required to make the implementation complete. @@ -2145,24 +2544,22 @@ class MolmoForCausalLM(PreTrainedModel): if self.config.weight_tying: return self.model.transformer.wte else: - return self.lm_head + return self.model.transformer.ff_out def set_output_embeddings(self, value: torch.nn.Module): if self.config.weight_tying: self.model.transformer.wte = value else: - self.lm_head = value + self.model.transformer.ff_out = value def tie_weights(self): """ This function is intentionally left as a no-op. - Weight tying is handled as follows: - - When the model is initialized, the `lm_head` layer is conditionally defined based on the `weight_tying` configuration. + - When the model is initialized, the `ff_out` layer is conditionally defined based on the `weight_tying` configuration. See: `if not config.weight_tying: self.transformer.update(...)` in `olmo/model.py`. - When computing logits, the `wte` weights are used directly if `weight_tying` is enabled. See: `if self.config.weight_tying: logits = F.linear(x, self.transformer.wte.weight, None)` in the `forward` method. - Therefore, there is no need to explicitly tie the weights in this function. """ pass @@ -2172,9 +2569,7 @@ class MolmoForCausalLM(PreTrainedModel): ) -> torch.nn.Embedding: """ Resizes input token embeddings matrix of the model if `new_num_tokens != config.embedding_size`. - Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. - Arguments: new_num_tokens (`int`, *optional*): The new number of tokens in the embedding matrix. Increasing the size will add newly initialized @@ -2183,19 +2578,16 @@ class MolmoForCausalLM(PreTrainedModel): pad_to_multiple_of (`int`, *optional*): If set will pad the embedding matrix to a multiple of the provided value. If `new_num_tokens` is set to `None` will just pad the embedding to a multiple of `pad_to_multiple_of`. - This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more details about this, or help on choosing the correct value for resizing, refer to this guide: https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc - Return: `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. - Note: This method differs from the base class implementation by resizing the `embedding_size` attribute of the model configuration instead of the `vocab_size`. It also includes a warning if the resized `embedding_size` - is less than the `vocab_size`. In Molmo, `embedding_size` refers to the dimensionality of the model's token + is less than the `vocab_size`. In OLMo, `embedding_size` refers to the dimensionality of the model's token embeddings, while `vocab_size` refers to the number of unique tokens in the vocabulary. """ model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of) @@ -2219,3 +2611,7 @@ class MolmoForCausalLM(PreTrainedModel): self.tie_weights() return model_embeds + + +# Always register for multi-modal features +AutoModelForCausalLM.register(MolmoConfig, MOLMoForCausalLM) \ No newline at end of file