diff --git "a/modeling_molmoe.py" "b/modeling_molmoe.py" --- "a/modeling_molmoe.py" +++ "b/modeling_molmoe.py" @@ -1,25 +1,163 @@ +""" +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 +import time +from abc import abstractmethod +from collections import defaultdict +from dataclasses import replace +from functools import partial +from os.path import join +from pathlib import Path +from typing import ( + Callable, + Dict, + Iterable, + List, + NamedTuple, + Optional, + Sequence, + Set, + Tuple, + cast, + Union, +) 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 -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 .config_molmoe import MolmoeConfig -from torch.nn import functional as F +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 transformers.modeling_outputs import CausalLMOutputWithPast + +from olmo.aliases import PathOrStr +from olmo.beam_search import ( + BeamSearch, + Constraint, + FinalSequenceScorer, + Sampler +) +from olmo.config import ( + ActivationType, + BlockType, + LayerNormType, + VisionBackboneType, + ImagePooling2DType, + ImageProjectType, + AttentionType, +) + +from olmo.util import resource_path +from .config_molmoe import ( + MolmoConfig, + VisionBackboneConfig +) + +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") + +__all__ = [ + "LayerNormBase", + "LayerNorm", + "RMSLayerNorm", + "RotaryEmbedding", + "Activation", + "GELU", + "ReLU", + "SwiGLU", + "OLMoBlock", + "OLMoSequentialBlock", + "OLMo", + "OLMoOutput", + "OLMoGenerateOutput", +] log = logging.getLogger(__name__) +def activation_checkpoint_function(cfg: ModelConfig): + preserve_rng_state = not ( + (cfg.attention_dropout == 0.0) and (cfg.embedding_dropout == 0.0) and + (cfg.residual_dropout == 0.0) and (cfg.response_residual_dropout == 0.0) + ) + from torch.utils.checkpoint import checkpoint + + return partial( + checkpoint, + preserve_rng_state=True, + use_reentrant=False, + ) + + +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) + + +def activation_checkpoint_function(cfg: MolmoConfig): + preserve_rng_state = not ( + (cfg.attention_dropout == 0.0) and (cfg.embedding_dropout == 0.0) and + (cfg.residual_dropout == 0.0) and (cfg.response_residual_dropout == 0.0) + ) + from torch.utils.checkpoint import checkpoint + + return partial( + checkpoint, + preserve_rng_state=True, + use_reentrant=False, + ) + + +def vit_activation_checkpoint_function(cfg: MolmoConfig): + v_cfg = cfg.vision_backbone + preserve_rng_state = ( + (v_cfg.attention_dropout == 0.0) and (v_cfg.residual_dropout == 0.0) + ) + from torch.utils.checkpoint import checkpoint + + return partial( + checkpoint, + preserve_rng_state=preserve_rng_state, + use_reentrant=False, + ) + + +def should_checkpoint_block(strategy: Optional[ActivationCheckpointingStrategy], block_idx: int) -> bool: + if strategy is None: + return False + elif ( + (strategy == ActivationCheckpointingStrategy.whole_layer) + or (strategy == ActivationCheckpointingStrategy.one_in_two and block_idx % 2 == 0) + or (strategy == ActivationCheckpointingStrategy.one_in_three and block_idx % 3 == 0) + or (strategy == ActivationCheckpointingStrategy.one_in_four and block_idx % 4 == 0) + or (strategy == ActivationCheckpointingStrategy.two_in_three and block_idx % 3 != 0) + or (strategy == ActivationCheckpointingStrategy.three_in_four and block_idx % 4 != 0) + ): + return True + else: + return False + + class BufferCache(dict, MutableMapping[str, torch.Tensor]): """ Cache for attention biases and other things that would normally be stored as buffers. @@ -31,61 +169,210 @@ class BufferCache(dict, MutableMapping[str, torch.Tensor]): """ -class StrEnum(str, Enum): - def __str__(self) -> str: - return self.value +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") + - def __repr__(self) -> str: - return f"'{str(self)}'" +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 ImageProjectType(StrEnum): - mlp = "mlp" - mlpx2 = "2mlp" - linear = "linear" +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 -class ImagePooling2DType(StrEnum): - attention = "attention" - attention_meanq = "attention-meanq" - attention_2wide = "attention_2wide" - attention_v2 = "attention-v2" - none = "none" - stack = "stack" + 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) -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 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 ensure_finite_(x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False): +class LayerNorm(LayerNormBase): """ - 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``. + The default :class:`LayerNorm` implementation which can optionally run in low precision. """ - 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, + 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 -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 _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") + 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 RMSLayerNorm(LayerNormBase): + """ + RMS layer norm, a simplified :class:`LayerNorm` implementation + """ + 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 + + 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 class RotaryEmbedding(nn.Module): @@ -93,7 +380,7 @@ class RotaryEmbedding(nn.Module): [Rotary positional embeddings (RoPE)](https://arxiv.org/abs/2104.09864). """ - def __init__(self, config: MolmoeConfig, cache: BufferCache): + def __init__(self, config: MolmoConfig, cache: BufferCache): super().__init__() self.config = config self.__cache = cache @@ -119,10 +406,10 @@ class RotaryEmbedding(nn.Module): return pos_sin[:, :, :seq_len, :], pos_cos[:, :, :seq_len, :] with torch.autocast(device.type, enabled=False): - dim = self.config.d_model // self.config.n_heads + 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)) seq = torch.arange(seq_len, device=device, dtype=torch.float) - freqs = torch.einsum("i , j -> i j", seq, inv_freq) + freqs = einsum("i , j -> i j", seq, inv_freq) if self.config.rope_impl == "cockatoo": positions = freqs.repeat_interleave(2, dim=-1) else: @@ -189,155 +476,159 @@ class RotaryEmbedding(nn.Module): return q_.type_as(q), k_.type_as(k) -class OLMoBlock(nn.Module): - """ - A base class for transformer block implementations. - """ - - def __init__(self, layer_id: int, config: MolmoeConfig, cache: BufferCache): +class Activation(nn.Module): + def __init__(self, config: MolmoConfig): 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 - # Dropout. - self.dropout = Dropout(config.residual_dropout, mask_p=config.response_residual_dropout) + @abstractmethod + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError - # 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) + @property + @abstractmethod + def output_multiplier(self) -> float: + raise NotImplementedError - # 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 GELU(nn.GELU): + @property + def output_multiplier(self) -> float: + return 1.0 - # 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 - ) +class QuickGELU(Activation): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(1.702 * x) - if self.config.block_type != "moe": - # 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 + @property + def output_multiplier(self) -> float: + return 1.0 - # Rotary embeddings. - if self.config.rope: - self.rotary_emb = RotaryEmbedding(config, self.__cache) - self.flash_attn_func = None - if config.attention_type == "flash": - try: - from flash_attn import flash_attn_func # type: ignore +class ReLU(nn.ReLU): + @property + def output_multiplier(self) -> float: + return 1.0 - self.flash_attn_func = flash_attn_func - except ModuleNotFoundError: - pass - 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, - ) +class SiLU(nn.SiLU): + @property + def output_multiplier(self) -> float: + return 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 - 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) +class SwiGLU(Activation): + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, gate = x.chunk(2, dim=-1) + return F.silu(gate) * x - 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) + @property + def output_multiplier(self) -> float: + return 0.5 - return F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=attn_mask, - dropout_p=dropout_p, - is_causal=is_causal, + +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 + + +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 + ) + + if config.attention_layer_norm: + if config.n_kv_heads is None: + config.n_kv_heads = config.n_heads + self.q_norm = RMSLayerNorm( + config, + size=config.d_model, + eps=config.layer_norm_eps + ) + self.k_norm = RMSLayerNorm( + config, + size=config.d_model, + eps=config.layer_norm_eps ) - def attention( - self, + # 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 + + # 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) + + self.flash_attn_func = None + if self.config.attention_type == AttentionType.flash: + try: + from flash_attn import flash_attn_func + self.flash_attn_func = flash_attn_func + except ModuleNotFoundError: + pass + + def attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -348,7 +639,7 @@ class OLMoBlock(nn.Module): use_cache: bool = False, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: B, T, C = q.size() # batch size, sequence length, d_model - dtype = k.dtype + dtype = k.dtype # Optionally apply layer norm to keys and queries. if self.q_norm is not None and self.k_norm is not None: @@ -359,9 +650,9 @@ class OLMoBlock(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.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2) + k = k.view(B, T, self.config.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.effective_n_kv_heads, C // self.config.n_heads).transpose(1, 2) + v = v.view(B, T, self.config.n_kv_heads, C // self.config.n_heads).transpose(1, 2) if self.config.use_position_ids and self.config.rope: # Apply rotary embeddings @@ -408,79 +699,20 @@ class OLMoBlock(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 build(cls, layer_id: int, config: MolmoeConfig, cache: BufferCache): - if config.block_type == "sequential": - return OLMoSequentialBlock(layer_id, config, cache) - elif config.block_type == "llama": - return OLMoLlamaBlock(layer_id, config, cache) - elif config.block_type == "moe": - return OLMoEBlock(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: MolmoeConfig, 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 _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 _scaled_dot_product_attention( self, @@ -493,47 +725,29 @@ class OLMoLlamaBlock(OLMoBlock): response_dropout_p: float = 0.0, is_causal: bool = False, ) -> torch.Tensor: - # 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) + """ + 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) - assert response_dropout_p == 0.0, "Response dropout is not supported in Llama." - - 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) + 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) - att = torch.matmul(attn_weights, v) - elif self.config.attention_type == "sdpa": - att = F.scaled_dot_product_attention( + return F.scaled_dot_product_attention( q, k, v, @@ -541,761 +755,259 @@ class OLMoLlamaBlock(OLMoBlock): dropout_p=dropout_p, is_causal=is_causal, ) - else: - raise NotImplementedError(self.config.attention_type) - att = att.to(og_dtype) - return att 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: - 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. - 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) - - # 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 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 + 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: - x = self.act(x1, x2) - x = self.ff_out(x) - x = self.dropout(x, drop_mask=drop_mask) - x = og_x + x + atten_in = x - return x, cache + 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) -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). - """ + q, k, v = qkv.split(self.fused_dims, dim=-1) + + # 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.config.norm_after: + att = self.attn_norm(att) + + return att, cache - def __init__(self, layer_id: int, config: MolmoeConfig, 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.att_proj = nn.Linear( - config.d_model, sum(self.fused_dims), - bias=config.include_bias or config.qkv_bias, - device=config.init_device - ) +class MolmoMLP(nn.Module): + def __init__( + self, + config: MolmoConfig + ): # Feed-forward input projection. - self.ff_proj = nn.Linear( - config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device + 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 ) - - 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 + self.act = SwiGLU(config) + 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, ) - 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: 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) - + + def forward(self, x): if not self.config.norm_after: - 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_norm(x) x = self.ff_proj(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.act(x) x = self.ff_out(x) if 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.dropout(x, drop_mask=drop_mask) - x = og_x + x - - return x, cache - - -class OLMoEBlock(OLMoBlock): - """ - This is a transformer MoE block where the output is computed as ``MoE(LN(x + Attention(LN(x))))`` - (plus another skip connection). - """ - - def __init__(self, layer_id: int, config, cache: BufferCache): - try: - from megablocks.layers.dmoe import dMoE - from megablocks.layers.moe import MoE - except ImportError: - raise ImportError( - "To train MoEs, run `pip install git+https://github.com/Muennighoff/megablocks.git@olmoe`" - ) - from .config_molmoe import config_to_moe_args - - super().__init__(layer_id, config, cache) - - self.moe_args = config_to_moe_args(config) - self.ffn = dMoE(self.moe_args) - - 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.att_proj = nn.Linear( - config.d_model, sum(self.fused_dims), bias=config.include_bias, device=config.init_device - ) - - 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() - - if self.config.init_fn == InitFnType.normal: - attn_out_std = ff_out_std = in_std = self.config.init_std - cutoff_factor = self.config.init_cutoff_factor - elif self.config.init_fn == InitFnType.mitchell: - in_std = 1 / math.sqrt(self.config.d_model) - attn_out_std = 1 / (math.sqrt(2 * self.config.d_model * (self.layer_id + 1))) - ff_out_std = 1 / (math.sqrt(2 * self.ff_out.in_features * (self.layer_id + 1))) - cutoff_factor = self.config.init_cutoff_factor or 3.0 - elif self.config.init_fn == InitFnType.full_megatron: - in_std = self.config.init_std - attn_out_std = ff_out_std = self.config.init_std / math.sqrt(2.0 * self.config.n_layers) - cutoff_factor = self.config.init_cutoff_factor or 3.0 - else: - raise NotImplementedError(self.config.init_fn) - - init_normal(self.att_proj, std=in_std, init_cutoff_factor=cutoff_factor) - init_normal(self.attn_out, std=attn_out_std, init_cutoff_factor=cutoff_factor) - self.attn_norm.reset_parameters() - self.ff_norm.reset_parameters() - init_normal(self.ffn.experts.mlp.w1, std=in_std, init_cutoff_factor=cutoff_factor) - init_normal(self.ffn.experts.mlp.w2, std=ff_out_std, init_cutoff_factor=cutoff_factor) - if hasattr(self.ffn.experts.mlp, "v1"): - init_normal(self.ffn.experts.mlp.v1, std=in_std, init_cutoff_factor=cutoff_factor) - if self.ffn.experts.bias is not None: - torch.nn.init.zeros_(self.ffn.experts.bias) - init_normal(self.ffn.router.layer, std=in_std, init_cutoff_factor=cutoff_factor) - - 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, - max_doc_len: Optional[int] = None, - cu_doc_lens: Optional[torch.Tensor] = None, - ) -> 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: - if self._activation_checkpoint_fn is not None: - qkv = self.att_proj(self._activation_checkpoint_fn(self.attn_norm, x)) - else: - qkv = self.att_proj(self.attn_norm(x)) - else: - qkv = self.att_proj(x) - - 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, - # max_doc_len=max_doc_len, - # cu_doc_lens=cu_doc_lens, - ) - 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, - # max_doc_len=max_doc_len, - # cu_doc_lens=cu_doc_lens, - ) - - 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 self.config.norm_after: - x = self.ffn(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 og_x + self.dropout(x, drop_mask=drop_mask), cache - else: - 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) - # Activation checkpointing for the MoE FFN is not supported - return og_x + self.dropout(self.ffn(x), drop_mask=drop_mask), cache - - -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. 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_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_dim: int = 2048 - 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 FullMolmoeConfig: - 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 = 1 - mlp_hidden_size: Optional[int] = None - activation_type: str = "swiglu" - block_type: str = "moe" - 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 = True - 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: FullMolmoeConfig): - 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) + x = self.ff_norm(x) - 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: FullMolmoeConfig): +class MolmoeMLP(nn.Module): + def __init__(self, config): + from transformers.activations import ACT2FN 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, + self.d_model = config.d_model + self.hidden_size = ( + config.mlp_hidden_size if config.mlp_hidden_size is not None \ + else config.mlp_ratio * config.d_model + ) // 2 + self.gate_proj = nn.Linear(self.d_model, self.hidden_size, bias=False) + self.up_proj = nn.Linear(self.d_model, self.hidden_size, bias=False) + self.down_proj = nn.Linear(self.hidden_size, self.d_model, bias=False) + self.act_fn = ACT2FN["silu"] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + +class MolmoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.num_experts = config.moe_num_experts + self.top_k = config.moe_top_k + self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False) + self.experts = nn.ModuleList([MolmoeMLP(config) for _ in range(self.num_experts)]) + self.ff_norm = RMSLayerNorm( + config, + size=config.d_model, + eps=config.layer_norm_eps ) - 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, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.ff_norm(hidden_states) + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) - 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 + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) -class BlockCollection(nn.Module): + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) - def __init__(self, config: FullMolmoeConfig): - super().__init__() - self.config = config - self.grad_checkpointing: bool = False + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be selected + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - v_cfg = config.vision_backbone - self.resblocks = nn.ModuleList([ - ResidualAttentionBlock(config) for _ in range(v_cfg.image_num_layers) - ]) + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx]) - 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 + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits -class VisionTransformer(nn.Module): - def __init__(self, config: FullMolmoeConfig): +class MolmoDecoderLayer(nn.Module): + """ + A base class for transformer block implementations. + """ + def __init__( + self, + layer_id: int, + config: MolmoConfig, + cache: BufferCache + ): super().__init__() + self.attn = MolmoAttention(config, cache) + if getattr(config, "moe_num_experts", 0) > 0: + self.mlp = MolmoeSparseMoeBlock(config) + else: + self.mlp = MolmoMLP(config) + self.layer_id = layer_id 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.hidden_size = ( + config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model ) + self.__cache = cache + if config.head_dim is None: + assert config.d_model % config.n_heads == 0 - self.transformer = BlockCollection(config) - - @torch.jit.ignore - def set_grad_checkpointing(self, enable=True): - self.transformer.grad_checkpointing = enable + self._activation_checkpoint_fn = None - 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() + # Dropout. + self.dropout = Dropout( + config.residual_dropout, + mask_p=config.response_residual_dropout + ) - 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:] + 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]]]: + """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) + """ - pos_emb = pos_emb.reshape( - (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1]) + att, cache = 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 - (patch_num_0, patch_num_1) = patch_num + return x, cache - 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 +class MolmoOutput(NamedTuple): + attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] + """ + Attention keys and values from each block. + """ - 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 + hidden_states: Optional[Tuple[torch.Tensor]] + """ + Hidden states from each block. + """ - x = self.patch_embedding(x) + last_hidden_states: torch.Tensor - # 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) +class MOLMoGenerateOutput(NamedTuple): + token_ids: torch.LongTensor + """ + The generated token IDs, a tensor of shape `(batch_size, beam_size, max_steps)`. + These do *not* include the original input IDs. + """ - hidden_states = self.transformer(x) - return hidden_states + scores: torch.FloatTensor + """ + The scores of the generated sequences, a tensor of shape `(batch_size, beam_size)`. + """ class MultiHeadDotProductAttention(nn.Module): - def __init__(self, config: FullMolmoeConfig, use_bias: bool = True, is_vit_layer: Optional[bool] = True): + def __init__(self, config: MolmoConfig, use_bias: bool = True, is_vit_layer: Optional[bool] = True): super().__init__() self.config = config self.use_bias = use_bias - + v_cfg = config.vision_backbone self.embed_dim = v_cfg.image_emb_dim self.num_heads = v_cfg.image_num_heads @@ -1312,30 +1024,30 @@ 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)) self.residual_dropout = Dropout(v_cfg.residual_dropout) - + def reset_parameters(self): nn.init.normal_(self.wq.weight, std=self.initializer_range) nn.init.normal_(self.wk.weight, std=self.initializer_range) @@ -1352,16 +1064,15 @@ class MultiHeadDotProductAttention(nn.Module): def _merge_heads(self, hidden_states) -> torch.Tensor: 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: - + + 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 else: inputs_k = inputs_q inputs_v = inputs_q - + xq, xk, xv = self.wq(inputs_q), self.wk(inputs_k), self.wv(inputs_v) xq = self._split_heads(xq, self.num_heads) @@ -1377,15 +1088,16 @@ 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 == "direct": + if self.config.attention_type == AttentionType.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 == "sdpa": + elif self.config.attention_type == AttentionType.sdpa: attn_output = F.scaled_dot_product_attention( xq.transpose(1, 2).contiguous(), xk.transpose(1, 2).contiguous(), @@ -1406,7 +1118,7 @@ class MultiHeadDotProductAttention(nn.Module): class MultiHeadAttentionPool(nn.Module): def __init__( self, - config: FullMolmoeConfig, + config: MolmoConfig, factor: int = 1, use_bias: bool = True, dropout: bool = True, @@ -1423,7 +1135,7 @@ class MultiHeadAttentionPool(nn.Module): self.output_layer = output_layer self.mean_residual = mean_residual self.query = query - + v_cfg = config.vision_backbone input_dim = v_cfg.image_emb_dim self.embed_dim = v_cfg.image_emb_dim * factor @@ -1441,25 +1153,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, - ), + ), ) if output_layer: @@ -1468,7 +1180,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) @@ -1544,10 +1256,47 @@ class MultiHeadAttentionPool(nn.Module): return attn_output -class MLP(nn.Module): - def __init__(self, config: FullMolmoeConfig, input_dim: int, dropout: float = 0.0): +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: ModelConfig, input_dim: int, dropout: float = 0.0): + 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 ) @@ -1558,52 +1307,212 @@ 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, - ) - # Activation function. - self.act = Activation.build(config) + ) + #`MLP` assume the activation takes two inputs, so it must be a 'llama' version. + self.act = LlamaSwiGLU(config) self.dropout = Dropout(dropout) + + def reset_parameters(self): + nn.init.normal_(self.w1.weight, std=self.initializer_range) + nn.init.normal_(self.w2.weight, std=self.initializer_range) + nn.init.normal_(self.w3.weight, std=self.initializer_range) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.w2(self.act(self.w1(x), self.w3(x))) + x = self.dropout(x) + return x + + +class Residual(nn.Module): + def __init__(self, submodule: nn.Module): + super().__init__() + self.submodule = submodule + + def reset_parameters(self): + self.submodule.reset_parameters() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + 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 + self.grad_checkpointing: bool = False + self._activation_checkpoint_fn: Callable = vit_activation_checkpoint_function(self.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: + if self.grad_checkpointing and not torch.jit.is_scripting(): + x = self._activation_checkpoint_fn(r, x) + else: + 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): + 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 - def reset_parameters(self): - nn.init.normal_(self.w1.weight, std=self.initializer_range) - nn.init.normal_(self.w2.weight, std=self.initializer_range) - nn.init.normal_(self.w3.weight, std=self.initializer_range) + 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) - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.w2(self.act(self.w1(x), self.w3(x))) - x = self.dropout(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 + 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 -class Residual(nn.Module): - def __init__(self, submodule: nn.Module): - super().__init__() - self.submodule = submodule + x = self.patch_embedding(x) - def reset_parameters(self): - self.submodule.reset_parameters() + # 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) - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x + self.submodule(x) + x = self.pre_ln(x) + + hidden_states = self.transformer(x) + return hidden_states -class OLMoVisionBackbone(nn.Module): - def __init__(self, config: FullMolmoeConfig): +class MolmoVisionBackbone(nn.Module): + def __init__(self, config: VisionBackboneConfig): super().__init__() self.config = config - self.image_vit = VisionTransformer(config) - input_dim: int = None self.image_pooling_2d: nn.Module = None if config.image_pooling_2d in {ImagePooling2DType.attention, ImagePooling2DType.attention_meanq}: @@ -1640,37 +1549,23 @@ class OLMoVisionBackbone(nn.Module): input_dim = nlayers * config.vision_backbone.image_emb_dim else: raise NotImplementedError(f"Unknown image pooling 2D method: {config.image_pooling_2d}") - + self.input_dim = 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: - #import pdb; pdb.set_trace() - #mlp_config.image_mlp_dim = 2048 - mlp_config.mlp_hidden_size = 2048 - 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_projector = MLP(config, input_dim) self.image_feature_dropout = Dropout(config.image_feature_dropout) + @classmethod + def build(cls, config: MolmoConfig) -> OLMoVisionBackbone: + v_cfg = config.vision_backbone + assert v_cfg is not None + return MolmoPretrainedVisionBackbone(config) + + @abstractmethod + def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]): + raise NotImplementedError() + def reset_parameters(self): if self.image_pooling_2d is not None: self.image_pooling_2d.reset_parameters() @@ -1682,15 +1577,20 @@ class OLMoVisionBackbone(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]]: raise NotImplementedError + - -class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): - def __init__(self, config: FullMolmoeConfig): +class MolmoPretrainedVisionBackbone(MolmoVisionBackbone): + def __init__(self, config: MolmoVisionBackboneConfig): super().__init__(config) v_cfg = self.config.vision_backbone - self.grad_checkpointing = False + + 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.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" @@ -1702,7 +1602,7 @@ class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): self.input_dim, bias=False, device=config.init_device, - ) + ) self.pad_embed = None if config.image_padding_embed: @@ -1716,12 +1616,35 @@ class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): 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() if self.config.use_cls_feature: nn.init.xavier_uniform_(self.cls_projector.weight) + def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]): + self.grad_checkpointing = True + if strategy in (ActivationCheckpointingStrategy.whole_layer, ActivationCheckpointingStrategy.vit_only): + self.image_vit.set_grad_checkpointing() + def encode_image(self, images: torch.Tensor) -> torch.Tensor: """ : param images: (batch_size, num_crops, num_patch, n_pixels) @@ -1741,332 +1664,130 @@ class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): features = [] for layer in cfg.vit_layers: features.append(image_features[layer]) - image_features = torch.cat(features, dim=-1) - else: - image_features = image_features[-1] - - cls_embed: torch.Tensor = None - if self.num_prefix_tokens > 0: - cls_embed = image_features[:, 0] - image_features = image_features[:, 1:] - - image_features = image_features * mask - image_features = image_features.view(B, T, N, -1) - - cls_embed = cls_embed.view(B, T, -1) if cls_embed is not None else None - - return image_features, cls_embed - - 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) - - if cfg.image_padding_embed: - assert image_masks is not None - if cfg.image_padding_embed == "pad_embed": - all_pad = (image_masks == 0).to(dtype=torch.float32) - pad_embed = self.pad_embed[None, None, None, :] - 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) - elif cfg.image_padding_embed == "pad_and_partial_pad": - 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) - 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 = 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.image_num_patch + (-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, - (0, 0, 0, 1, 0, 1, 0, 0, 0, 0), - ) - - # image pooling - image_features = einops.rearrange( - image_features, - 'b n (h dh) (w dw) c -> (b n h w) (dh dw) c', - dh=cfg.image_pooling_h, - dw=cfg.image_pooling_w, - ) - - 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 not in {ImagePooling2DType.none, ImagePooling2DType.stack}: - 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() - image_features = image_features.reshape(batch_size, num_image, h * w, -1) - - # MLP layer to map the feature. - 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: - 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 ModuleType(str, Enum): - in_module = "in" - out_module = "out" - emb = "emb" - final_out = "final_out" - - -def init_weights( - config: FullMolmoeConfig, - 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 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: FullMolmoeConfig): - 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: FullMolmoeConfig) -> '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: MolmoeConfig, - *, - 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: FullMolmoeConfig, 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) + image_features = torch.cat(features, dim=-1) else: - raise NotImplementedError(f"Unknown LayerNorm type: '{config.layer_norm_type}'") + image_features = image_features[-1] + cls_embed: torch.Tensor = None + if self.num_prefix_tokens > 0: + cls_embed = image_features[:, 0] + image_features = image_features[:, 1:] + + image_features = image_features * mask + image_features = image_features.view(B, T, N, -1) -class RMSLayerNorm(LayerNormBase): - """ - RMS layer norm, a simplified :class:`LayerNorm` implementation - """ + cls_embed = cls_embed.view(B, T, -1) if cls_embed is not None else None - def __init__( - self, - config: FullMolmoeConfig, - size: Optional[int] = None, - elementwise_affine: Optional[bool] = None, - eps: float = 1e-5, - ): - super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) + return image_features, cls_embed + + def forward(self, images: torch.Tensor, image_masks: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + cfg = self.config - 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) + # 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) - if self.weight is not None: - if self.bias is not None: - return self.weight * x + self.bias + og_dtype = image_features.dtype + if cfg.image_padding_embed: + assert image_masks is not None + if cfg.image_padding_embed == "pad_embed": + all_pad = (image_masks == 0).to(dtype=torch.float32) + pad_embed = self.pad_embed[None, None, None, :] + 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) + 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) + 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: - return self.weight * x - else: - return x + 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,), + ) -class LayerNorm(LayerNormBase): - """ - The default :class:`LayerNorm` implementation which can optionally run in low precision. - """ + if cfg.vision_backbone.image_num_patch[0] % cfg.image_pooling_h == 1: + # Pad so we can still pool 2x2 patches + image_features = F.pad( + image_features, + (0, 0, 0, 1, 0, 1, 0, 0, 0, 0), + ) + + # image pooling + image_features = einops.rearrange( + image_features, + 'b n (h dh) (w dw) c -> (b n h w) (dh dw) c', + dh=cfg.image_pooling_h, + dw=cfg.image_pooling_w, + ) - def __init__( - self, - config: FullMolmoeConfig, - 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 + 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) - 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 - ) + 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) else: - return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps) + 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) + + # image_features: (batch_size, num_image, num_patch, d_model) + # cls_embed: (batch_size, num_image, d_model) + return image_features, cls_embed -class MOLMoE(nn.Module): - def __init__(self, config: FullMolmoeConfig, init_params: bool = True): - super().__init__() +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 MolmoModel(MolmoPretrainedModel): + def __init__( + self, + config: MolmoConfig, + init_params: bool = True + ): + super().__init__(config) self.config = config self.__cache = BufferCache() @@ -2080,6 +1801,16 @@ class MOLMoE(nn.Module): warnings.warn( "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning ) + + self.activation_checkpointing_strategy: Optional[ActivationCheckpointingStrategy] = None + self._activation_checkpoint_fn: Callable = activation_checkpoint_function(self.config) + + if not ( + 0 < self.config.block_group_size <= self.config.n_layers + and self.config.n_layers % self.config.block_group_size == 0 + ): + raise OLMoConfigurationError("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 @@ -2102,62 +1833,46 @@ class MOLMoE(nn.Module): dict( wte=wte, emb_drop=Dropout(config.embedding_dropout), - ln_f=LayerNorm.build(config), + ln_f=RMSLayerNorm( + config, + size=config.d_model, + eps=config.layer_norm_eps), ) ) - 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, - ) - } - ) - + layers = [ + MolmoDecoderLayer(i, config, self.__cache) \ + for i in range(config.n_layers) + ] + self.transformer.update({"blocks": nn.ModuleList(layers)}) + self.vision_backbone: Optional[OLMoVisionBackbone] = None if config.vision_backbone is not None: - self.vision_backbone = OLMoPretrainedVisionBackbone(config) + self.vision_backbone = MolmoVisionBackbone.build(config) - self.__num_fwd_flops: Optional[int] = None - - def reset_parameters(self): if self.vision_backbone is not None: - self.vision_backbone.reset_parameters() - self.reset_non_vision_parameters() - - 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.vision_backbone.reset_with_pretrained_weights() - 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: + def set_activation_checkpointing(self, strategy: Optional[ActivationCheckpointingStrategy]): + self.activation_checkpointing_strategy = strategy + if self.config.block_group_size != 1: + for block_group in self.transformer.block_groups: + block_group.set_activation_checkpointing(strategy) + else: for block in self.transformer.blocks: - block.reset_parameters() + block.set_activation_checkpointing(strategy) + + if self.vision_backbone is not None: + self.vision_backbone.set_activation_checkpointing(strategy) + + @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) else: - for block_group in self.transformer.block_groups: - block_group.reset_parameters() + return device + def forward( self, @@ -2176,7 +1891,7 @@ class MOLMoE(nn.Module): last_logits_only: bool = False, output_hidden_states: Optional[bool] = None, append_last_valid_logits: Optional[torch.Tensor] = None, - ) -> ModelOutput: + ) -> MolmoOutput: """ :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 @@ -2185,16 +1900,20 @@ class MOLMoE(nn.Module): 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 @@ -2223,9 +1942,13 @@ class MOLMoE(nn.Module): 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 - + if subsegment_ids is not None: assert not use_cache, "Subsegment_ids cannot be used with cache." subsegment_mask = subsegment_ids.unsqueeze(2) <= subsegment_ids.unsqueeze(1) @@ -2240,7 +1963,7 @@ class MOLMoE(nn.Module): 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) @@ -2271,7 +1994,7 @@ class MOLMoE(nn.Module): 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 ) @@ -2283,15 +2006,7 @@ class MOLMoE(nn.Module): position_ids = torch.clamp( torch.cumsum(attention_mask, dim=-1) - 1, min=0, - ).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 + ).broadcast_to((batch_size, attention_mask.shape[-1])) # Add input + positional embeddings and apply dropout. # shape: (batch_size, seq_len, d_model) @@ -2315,17 +2030,12 @@ class MOLMoE(nn.Module): 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 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: + if 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) @@ -2353,38 +2063,25 @@ class MOLMoE(nn.Module): all_hidden_states = [] # Apply blocks one-by-one. - 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 - ] + for block_idx, layer 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] + if should_checkpoint_block(self.activation_checkpointing_strategy, block_idx): + # shape: (batch_size, seq_len, d_model) + x, cache = self._activation_checkpoint_fn( + layer, x, attention_bias=attention_bias, position_ids=position_ids, drop_mask=response_mask, layer_past=layer_past, use_cache=use_cache ) - 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) + else: + # 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 images is not None and self.config.use_cls_feature: assert num_image is not None x = torch.cat( @@ -2407,81 +2104,174 @@ class MOLMoE(nn.Module): if output_hidden_states: # add final hidden state post-final-layernorm, following HuggingFace's convention all_hidden_states.append(x) + + # 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) # Get logits. # shape: (batch_size, seq_len or 1, vocab_size) - if self.config.weight_tying: - logits = F.linear(x, self.transformer.wte.weight, None) # type: ignore + return MolmoOutput( + last_hidden_states=x, + attn_key_values=attn_key_values, + hidden_states=tuple(all_hidden_states) \ + if output_hidden_states else None + ) + + def num_params(self, include_embedding: bool = True) -> int: + """ + Get the total number of parameters. + """ + params = (np for np in self.named_parameters()) + if not include_embedding: + params = filter( # type: ignore + lambda np: ".wte." not in np[0] and ".wpe." not in np[0], + params, + ) + return sum(p.numel() for _, p in params) + + @classmethod + def from_checkpoint( + cls, checkpoint_dir: PathOrStr, device: str = "cpu", + checkpoint_type: Optional[CheckpointType] = None + ) -> OLMo: + """ + Load an OLMo model from a checkpoint. + """ + raise NotImplementedError("This method is not implemented yet.") + + def _make_state_dict_compatible( + self, state_dict: Dict[str, torch.Tensor] + ) -> Tuple[Dict[str, torch.Tensor], Dict[str, Set[str]]]: + """ + Handles some cases where the state dict is valid yet may need to be transformed in order to + be loaded. + + This modifies the state dict in-place and also returns it, along with a mapping of original key + names to new key names in cases where the keys were simply renamed. That mapping can be used + to make a corresponding optimizer state dict compatible as well. + """ + import re + from fnmatch import fnmatch + + new_keys_to_og_keys: Dict[str, str] = {} + + # Remove "_fsdp_wrapped_module." prefix from all keys. We don't want this prefix when the model is + # not wrapped in FSDP. And when the model is wrapped in FSDP, loading this state dict will still work + # fine without the prefixes. This also simplifies the other steps below. + for key in list(state_dict.keys()): + state_dict[(new_key := key.replace("_fsdp_wrapped_module.", ""))] = state_dict.pop(key) + new_keys_to_og_keys[new_key] = key + + # For backwards compatibility prior to fixing https://github.com/allenai/LLM/issues/222 + if self.config.block_type == BlockType.sequential: + for key in list(state_dict.keys()): + if fnmatch(key, "transformer.*.norm.weight"): + tensor = state_dict.pop(key) + state_dict[(new_key := key.replace("norm.weight", "attn_norm.weight"))] = tensor + new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key] + state_dict[(new_key := key.replace("norm.weight", "ff_norm.weight"))] = tensor.clone() + new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key] + del new_keys_to_og_keys[key] + elif fnmatch(key, "transformer.*.norm.bias"): + tensor = state_dict.pop(key) + state_dict[(new_key := key.replace("norm.bias", "attn_norm.bias"))] = tensor + new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key] + state_dict[(new_key := key.replace("norm.bias", "ff_norm.bias"))] = tensor.clone() + new_keys_to_og_keys[new_key] = new_keys_to_og_keys[key] + del new_keys_to_og_keys[key] + + # For loading a state dict that was saved with a different `block_group_size`. + if "transformer.block_groups.0.0.attn_out.weight" in state_dict.keys(): + state_dict_block_group_size = len( + [k for k in state_dict.keys() if fnmatch(k, "transformer.block_groups.0.*.attn_out.weight")] + ) else: - logits = self.transformer.ff_out(x) # type: ignore - if self.config.scale_logits: - logits.mul_(1 / math.sqrt(self.config.d_model)) + state_dict_block_group_size = 1 + if self.config.block_group_size != state_dict_block_group_size: + log.info( + f"Regrouping state dict blocks from group size {state_dict_block_group_size} to " + f"group size {self.config.block_group_size}" + ) + # For simplicity we're first going to flatten out the block groups in the state dict (if necessary) + # and then (re-)group them into the right block sizes. + if state_dict_block_group_size > 1: + for key in list(state_dict.keys()): + if (m := re.match(r"transformer.block_groups\.(\d+)\.(\d+)\..*", key)) is not None: + group_idx, group_block_idx = int(m.group(1)), int(m.group(2)) + block_idx = (group_idx * state_dict_block_group_size) + group_block_idx + state_dict[ + ( + new_key := key.replace( + f"block_groups.{group_idx}.{group_block_idx}.", f"blocks.{block_idx}." + ) + ) + ] = state_dict.pop(key) + new_keys_to_og_keys[new_key] = new_keys_to_og_keys.pop(key) + + if self.config.block_group_size > 1: + # Group the state dict blocks into the right block size. + for key in list(state_dict.keys()): + if (m := re.match(r"transformer.blocks\.(\d+)\..*", key)) is not None: + block_idx = int(m.group(1)) + group_idx, group_block_idx = ( + block_idx // self.config.block_group_size, + block_idx % self.config.block_group_size, + ) + state_dict[ + ( + new_key := key.replace( + f"blocks.{block_idx}.", f"block_groups.{group_idx}.{group_block_idx}." + ) + ) + ] = state_dict.pop(key) + new_keys_to_og_keys[new_key] = new_keys_to_og_keys.pop(key) - 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) + og_keys_to_new: Dict[str, Set[str]] = defaultdict(set) + for new_key, og_key in new_keys_to_og_keys.items(): + og_keys_to_new[og_key].add(new_key) - 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] + return state_dict, og_keys_to_new -class MOLMoEForCausalLM(PreTrainedModel): - config_class = MolmoeConfig +class MolmoForCausalLM(PreTrainedModel): + """ + Extremely barebones HF model wrapper. + """ + config_class = MolmoConfig base_model_prefix = "model" - _no_split_modules = ["OLMoBlock"] + _no_split_modules = ["MolmoDecoderLayer"] - def __init__(self, config: MolmoeConfig, model: Optional[MOLMoE] = None, init_params: bool = False): + def __init__( + self, + config: MolmoConfig + ): 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 model: - full_config = FullMolmoeConfig( - 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=10000.0, - layer_norm_eps=1e-5, - 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, - ) + 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, ) - self.model = MOLMoE(full_config, init_params=init_params) - else: - self.model = model def forward( self, @@ -2517,7 +2307,7 @@ class MOLMoEForCausalLM(PreTrainedModel): 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.forward( + outputs = self.model( input_ids=input_ids, input_embeddings=inputs_embeds, attention_mask=attention_mask, @@ -2534,9 +2324,25 @@ class MOLMoEForCausalLM(PreTrainedModel): output_hidden_states=output_hidden_states, 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 - logits = outputs.logits - hidden_states = outputs.hidden_states + 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) loss = None if labels is not None: @@ -2580,12 +2386,68 @@ class MOLMoEForCausalLM(PreTrainedModel): loss=loss, logits=logits, past_key_values=outputs.attn_key_values, - hidden_states=hidden_states, + hidden_states=outputs.hidden_states, ) def can_generate(self) -> bool: return True + @torch.no_grad() + def generate( + self, + input_ids, + images=None, + attention_mask=None, + image_masks=None, + image_input_idx=None, + generation_config=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 + position_ids: Optional[torch.Tensor] = None + 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 + ) + 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))], + dim=1, + ) + if attention_mask is not None: + assert attention_mask.shape == (batch_size, mask_len) + + out = super().generate( + # batch["input_ids"], + input_ids, + generation_config, + attention_mask=attention_mask, + images=images, + image_masks=image_masks, + image_input_idx=image_input_idx, + position_ids=position_ids, + append_last_valid_logits=append_last_valid_logits, + **kwargs, + ) + + return out + @torch.no_grad() def generate_from_batch( self, @@ -2595,7 +2457,7 @@ class MOLMoEForCausalLM(PreTrainedModel): ): 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") @@ -2622,7 +2484,7 @@ class MOLMoEForCausalLM(PreTrainedModel): ) if attention_mask is not None: assert attention_mask.shape == (batch_size, mask_len) - + out = super().generate( batch["input_ids"], generation_config, @@ -2636,7 +2498,7 @@ class MOLMoEForCausalLM(PreTrainedModel): ) return out - + def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple]] = None, **kwargs ): @@ -2664,7 +2526,7 @@ class MOLMoEForCausalLM(PreTrainedModel): model_inputs["image_masks"] = image_masks model_inputs["image_input_idx"] = image_input_idx model_inputs["append_last_valid_logits"] = append_last_valid_logits - else: + else: model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values} model_inputs.update(kwargs) @@ -2676,7 +2538,6 @@ class MOLMoEForCausalLM(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: @@ -2687,8 +2548,9 @@ class MOLMoEForCausalLM(PreTrainedModel): del model_kwargs["images"] del model_kwargs["image_masks"] del model_kwargs["image_input_idx"] - model_kwargs = super()._update_model_kwargs_for_generation( - outputs, model_kwargs, is_encoder_decoder, standardize_cache_format, num_new_tokens) + 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 return model_kwargs # TODO: these are required to make the implementation complete. @@ -2711,22 +2573,24 @@ class MOLMoEForCausalLM(PreTrainedModel): if self.config.weight_tying: return self.model.transformer.wte else: - return self.model.transformer.ff_out + return self.lm_head def set_output_embeddings(self, value: torch.nn.Module): if self.config.weight_tying: self.model.transformer.wte = value else: - self.model.transformer.ff_out = value + self.lm_head = 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 `ff_out` layer is conditionally defined based on the `weight_tying` configuration. + - When the model is initialized, the `lm_head` 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 @@ -2736,7 +2600,9 @@ class MOLMoEForCausalLM(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 @@ -2745,12 +2611,15 @@ class MOLMoEForCausalLM(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` @@ -2778,7 +2647,3 @@ class MOLMoEForCausalLM(PreTrainedModel): self.tie_weights() return model_embeds - - -# Always register for multi-modal features -AutoModelForCausalLM.register(MolmoeConfig, MOLMoEForCausalLM) \ No newline at end of file