From 18569a4c63f28b3962b624a17df32b9935f216fa Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Thu, 23 Jan 2025 15:18:00 -0800 Subject: [PATCH 1/7] Adding molmo code locally --- pdelfin/train/molmo/__init__.py | 0 pdelfin/train/molmo/config_molmo.py | 60 + pdelfin/train/molmo/image_processing_molmo.py | 546 ++++ pdelfin/train/molmo/modeling_molmo.py | 2367 +++++++++++++++++ pdelfin/train/molmo/preprocessing_molmo.py | 192 ++ 5 files changed, 3165 insertions(+) create mode 100644 pdelfin/train/molmo/__init__.py create mode 100644 pdelfin/train/molmo/config_molmo.py create mode 100644 pdelfin/train/molmo/image_processing_molmo.py create mode 100644 pdelfin/train/molmo/modeling_molmo.py create mode 100644 pdelfin/train/molmo/preprocessing_molmo.py diff --git a/pdelfin/train/molmo/__init__.py b/pdelfin/train/molmo/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pdelfin/train/molmo/config_molmo.py b/pdelfin/train/molmo/config_molmo.py new file mode 100644 index 0000000..2322810 --- /dev/null +++ b/pdelfin/train/molmo/config_molmo.py @@ -0,0 +1,60 @@ +from typing import List + +from transformers import PretrainedConfig, AutoTokenizer + + +class MolmoConfig(PretrainedConfig): + model_type = "molmo" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=50304, + embedding_size=50304, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + max_position_embeddings=2048, + initializer_range=0.02, + use_cache=True, + layer_norm_eps: float = 1e-5, + rope_theta=10000.0, + clip_qkv=None, + qkv_bias: bool = False, + weight_tying: bool = False, + use_position_ids: bool=True, + tie_word_embeddings: bool=True, + attention_layer_norm: bool=False, + norm_after: bool = False, + layer_norm_type: str="rms", + **kwargs, + ): + self.vocab_size = vocab_size + self.embedding_size = embedding_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.layer_norm_eps = layer_norm_eps + self.weight_tying = weight_tying + self.use_position_ids = use_position_ids + self.attention_layer_norm = attention_layer_norm + self.num_key_value_heads = num_key_value_heads + self.initializer_range = initializer_range + self.use_cache = use_cache + self.rope_theta = rope_theta + self.clip_qkv = clip_qkv + self.qkv_bias = qkv_bias + self.norm_after = norm_after + self.tie_word_embeddings = tie_word_embeddings + self.layer_norm_type = layer_norm_type + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + +MolmoConfig.register_for_auto_class() \ No newline at end of file diff --git a/pdelfin/train/molmo/image_processing_molmo.py b/pdelfin/train/molmo/image_processing_molmo.py new file mode 100644 index 0000000..ef787ec --- /dev/null +++ b/pdelfin/train/molmo/image_processing_molmo.py @@ -0,0 +1,546 @@ +"""Image processor class for Molmo""" +from typing import List, Optional, Union, Mapping + +import numpy as np +import einops +import torch +import torchvision.transforms +from torchvision.transforms import InterpolationMode +from torchvision.transforms.functional import convert_image_dtype + +from transformers.image_utils import ( + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, + ImageInput, + is_valid_image, +) +from transformers.processing_utils import ImagesKwargs +from transformers.image_processing_utils import BaseImageProcessor +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +def pad_to_bounding_box( + image, offset_height, offset_width, target_height, + target_width, value=0 +): + height, width = image.shape[:2] + after_padding_width = target_width - offset_width - width + after_padding_height = target_height - offset_height - height + return np.pad(image, [ + [offset_height, after_padding_height], + [offset_width, after_padding_width], + [0, 0] + ], constant_values=value) + + +def normalize_image(image, offset, scale): + image -= np.array(offset, dtype=np.float32)[None, None, :] + image /= np.array(scale, dtype=np.float32)[None, None, :] + return image + + +def resize_and_pad( + image, + desired_output_size, + resize_method="torch-bilinear", + pad_value=0, + normalize=True, + image_mean=OPENAI_CLIP_MEAN, + image_std=OPENAI_CLIP_STD, +): + desired_height, desired_width = desired_output_size + height, width = image.shape[:2] + + # Cast into float32 since the training code did this in float32 and it (very rarely) effects + # the results after rounding. + image_scale_y = np.array(desired_height, np.float32) / np.array(height, np.float32) + image_scale_x = np.array(desired_width, np.float32) / np.array(width, np.float32) + image_scale = min(image_scale_x, image_scale_y) + scaled_height = int(np.array(height, np.float32) * image_scale) + scaled_width = int(np.array(width, np.float32) * image_scale) + + if resize_method == "tensorflow": + # This how the original training code did resizing, it can produce slightly different + # results then using torch resize so we keep it just in case + import tensorflow as tf + image = tf.image.convert_image_dtype(tf.constant(image), dtype=tf.float32) + image = tf.image.resize( + image, + [scaled_height, scaled_width], + method=tf.image.ResizeMethod.BILINEAR, + antialias=True, + ) + image = tf.clip_by_value(image, 0.0, 1.0) + image = image.numpy() + elif resize_method == "torch-bilinear": + image = torch.permute(torch.from_numpy(image), [2, 0, 1]) + image = convert_image_dtype(image) # resize in float32 to match the training code + image = torchvision.transforms.Resize( + [scaled_height, scaled_width], InterpolationMode.BILINEAR, antialias=True + )(image) + image = torch.clip(image, 0.0, 1.0) + image = torch.permute(image, [1, 2, 0]).numpy() + else: + raise NotImplementedError(resize_method) + + top_pad = (desired_height - scaled_height) // 2 + left_pad = (desired_width - scaled_width) // 2 + padding = [ + [top_pad, desired_height - scaled_height - top_pad], + [left_pad, desired_width - scaled_width - left_pad], + [0, 0] + ] + image_mask = np.pad(np.ones_like(image[:, :, 0], dtype=bool), padding[:2]) + image = np.pad(image, padding, constant_values=pad_value) + if normalize: + image = normalize_image(image, offset=image_mean, scale=image_std) + return image, image_mask + + +def select_tiling(h, w, patch_size, max_num_patches): + """Decide how best to divide in image of size [w, h] in up to max_num_patches of size patch_size""" + original_size = np.stack([h, w]) # [1, 2] + original_res = h * w + tilings = [] + for i in range(1, max_num_patches+1): + for j in range(1, max_num_patches+1): + if i*j <= max_num_patches: + tilings.append((i, j)) + # sort so argmin and argmax favour smaller tilings in the event of a tie + tilings.sort(key=lambda x: (x[0]*x[1], x[0])) + candidate_tilings = np.array(tilings, dtype=np.int32) # [n_resolutions, 2] + candidate_resolutions = candidate_tilings * patch_size # [n_resolutions, 2] + + # How much we would need to scale the image to fit exactly in each tiling + original_size = np.stack([h, w], dtype=np.float32) # [1, 2] + required_scale_d = candidate_resolutions.astype(np.float32) / original_size + required_scale = np.min(required_scale_d, axis=-1, keepdims=True) # [n_resolutions, 1] + if np.all(required_scale < 1): + # We are forced to downscale, so try to minimize the amount of downscaling + ix = np.argmax(required_scale) + else: + # Pick the resolution that required the least upscaling so that it most closely fits the image + required_scale = np.where(required_scale < 1.0, 10e9, required_scale) + ix = np.argmin(required_scale) + return candidate_tilings[ix] + + +class MolmoImagesKwargs(ImagesKwargs, total=False): + max_crops: Optional[int] + overlap_margins: Optional[List[int]] + base_image_input_size: Optional[List[int]] + image_token_length_w: Optional[int] + image_token_length_h: Optional[int] + image_patch_size: Optional[int] + image_padding_mask: Optional[bool] + + +class MolmoImageProcessor(BaseImageProcessor): + """Preprocess images and multi-model inputs""" + + def __init__( + self, + max_crops: int = 12, + overlap_margins: List[int] = (4, 4), + base_image_input_size: List[int] = (336, 336), + image_token_length_w: int = 12, + image_token_length_h: int = 12, + image_patch_size: int = 14, + image_padding_mask: bool = True, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + **kwargs, + ): + super().__init__(**kwargs) + self.max_crops = max_crops + self.overlap_margins = overlap_margins + self.base_image_input_size = base_image_input_size + self.image_token_length_w = image_token_length_w + self.image_token_length_h = image_token_length_h + self.image_patch_size = image_patch_size + self.image_padding_mask = image_padding_mask + self.do_normalize = do_normalize + self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN + self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD + + def image_to_patches_and_tokens( + self, + image: ImageInput, + image_patch_token_id: int, + image_col_token_id: int, + image_start_token_id: int, + image_end_token_id: int, + max_crops: Optional[int] = None, + overlap_margins: Optional[List[int]] = None, + base_image_input_size: Optional[Union[int, List[int]]] = None, + image_token_length_w: Optional[int] = None, + image_token_length_h: Optional[int] = None, + image_patch_size: Optional[int] = None, + ): + if isinstance(base_image_input_size, int): + base_image_input_size = (base_image_input_size, base_image_input_size) + + base_image_input_d = image_patch_size + tokens_per_image = image_token_length_w * image_token_length_h + image_base_patch_w = base_image_input_size[1] // base_image_input_d + image_base_patch_h = base_image_input_size[0] // base_image_input_d + + original_image_h, original_image_w = image.shape[:2] + crop_size = base_image_input_size[0] + + # Discard this many patches from the (left/top, right/bottom) of crops + left_margin, right_margin = overlap_margins + # left_margin, right_margin = 2, 2 + assert left_margin % 2 == 0 # Required for compatibility with 2x2 pooling + total_margin_pixels = base_image_input_d*(right_margin + left_margin) # pixels removed per dim + crop_patches = base_image_input_size[0] // base_image_input_d # patches per crop dim + crop_window_patches = crop_patches - (right_margin + left_margin) # usable patches + crop_window_size = crop_window_patches * base_image_input_d + tiling = select_tiling( + original_image_h - total_margin_pixels, + original_image_w - total_margin_pixels, + crop_window_size, + max_crops + ) + src, img_mask = resize_and_pad( + image, + [tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels] + ) + + # Now we have to split the image into crops, while keeping track of how each patch in the + # each crop should be ordered in the global image, this require a lot of tricky booking + n_crops = tiling[0] * tiling[1] + patches_arr = [] + mask_arr = [] + patch_ordering_arr = [] + + # We assume 2x2 pooling, but can allow padding the right/bottom with extra + # patches if the number of patches per side is not even + assert (crop_patches+1)//2 == image_token_length_h + assert (crop_patches+1)//2 == image_token_length_w + on = 0 + on_patch = 0 + for i in range(tiling[0]): + y0 = i*crop_window_size + if i == 0: + crop_y0 = 0 + else: + crop_y0 = left_margin // 2 + + crop_h = image_base_patch_h - (right_margin + left_margin) + if i == 0: + crop_h += left_margin + if i == (tiling[0]-1): + crop_h += right_margin + for j in range(tiling[1]): + x0 = j*crop_window_size + if j == 0: + crop_x0 = 0 + else: + crop_x0 = left_margin // 2 + + crop_w = image_base_patch_w - (right_margin + left_margin) + if j == 0: + crop_w += left_margin + if j == (tiling[1]-1): + crop_w += right_margin + + pooled_w = (crop_w + 1) // 2 + pooled_h = (crop_h + 1) // 2 + patch_ordering_arr.append( + pad_to_bounding_box( + np.reshape(np.arange(on, on+pooled_h*pooled_w, dtype=np.int32), (pooled_h, pooled_w, 1)), + crop_y0, crop_x0, image_token_length_h, image_token_length_w, value=-1 + )[:, :, 0] + ) + patches_arr.append(src[y0:y0+crop_size, x0:x0+crop_size]) + mask_arr.append(img_mask[y0:y0+crop_size, x0:x0+crop_size]) + + on += pooled_h*pooled_w + on_patch += 1 + patches = np.stack(patches_arr) + patch_ordering = np.stack(patch_ordering_arr) + img_mask = np.stack(mask_arr) + + # Switch to [n_crops, n_patches, pixels_per_patch] format + image_layout_impatch_w, image_layout_impatch_h = tiling[0], tiling[1] + patches = einops.rearrange( + patches, 'p (h dh) (w dw) c -> p (h w) (dh dw c)', + dh=base_image_input_d, + dw=base_image_input_d, + h=image_base_patch_h, + w=image_base_patch_w + ) + img_mask = einops.rearrange( + img_mask, 'p (h dh) (w dw) -> p (h w) (dh dw)', + dh=base_image_input_d, + dw=base_image_input_d, + h=image_base_patch_h, + w=image_base_patch_w + ) + + img_mask = img_mask.astype(np.float32).mean(axis=-1) + patch_ordering = np.reshape(patch_ordering, [-1]) + valid = patch_ordering >= 0 + + # Transpose order, to get left-to-right order instead of crop-by-crop order + patch_ordering_rh = np.reshape( + patch_ordering, + [tiling[0], tiling[1], image_token_length_h, image_token_length_w] + ) + patch_ordering_rh = np.transpose(patch_ordering_rh, [0, 2, 1, 3]) + patch_ordering_rh = np.reshape(patch_ordering_rh, [-1]) + + # The transpose will screw up which patches are masked, project the + # new order into sparse structure of `patch_ordering` to fix this + patch_ordering[valid] = patch_ordering_rh[patch_ordering_rh >= 0] + + # Now build the output tokens + h = tiling[0] * crop_window_patches + (right_margin+left_margin) + w = tiling[1] * crop_window_patches + (right_margin+left_margin) + per_row = np.full( + ((w+1)//2,), + image_patch_token_id, + ) + per_row = np.concatenate([per_row, [image_col_token_id]], 0) + + joint = np.tile(per_row, [(h+1)//2]) + joint = [ + [image_start_token_id], + joint, + [image_end_token_id] + ] + + # Finally do the same for the global image + resized, _ = resize_and_pad(image, base_image_input_size) + resized = einops.rearrange( + resized, '(h dh) (w dw) c -> (h w) (dh dw c)', + dh=base_image_input_d, + dw=base_image_input_d, + h=image_base_patch_h, + w=image_base_patch_w + ) + patches = np.concatenate([np.expand_dims(resized, 0), patches], 0) + + # Global image goes first, so the order of patches in previous crops gets increased + patch_ordering = np.where( + patch_ordering >= 0, + patch_ordering + tokens_per_image, + -1 + ) + patch_ordering = np.concatenate([np.arange(0, tokens_per_image), patch_ordering], 0) + per_row = np.full( + (image_token_length_w,), + image_patch_token_id, + ) + per_row = np.concatenate([per_row, [image_col_token_id]], 0) + extra_tokens = np.tile(per_row, [image_token_length_h]) + joint = [ + [image_start_token_id], + extra_tokens, + [image_end_token_id], + ] + joint + + joint = np.concatenate(joint, 0) + img_mask = np.pad(img_mask, [[0, 1], [0, 0]], constant_values=-1) + return patches, joint, patch_ordering, img_mask + + def build_image_input_idx( + self, + image_tokens: np.ndarray, + patch_order: np.ndarray, + image_patch_token_id: int, + no_image: Optional[bool] = None, + image_token_length_w: Optional[int] = None, + image_token_length_h: Optional[int] = None, + ): + """Converts `patch_order` into a mapping of token_id -> patch_id""" + + tokens_per_image = image_token_length_w * image_token_length_h + if no_image is not None and no_image: + return np.zeros((0, tokens_per_image), np.int32) + + # Indices to insert the patches + image_input_idx = image_tokens == image_patch_token_id + image_input_idx = np.nonzero(image_input_idx)[0].astype(np.int32) + + if patch_order is not None: + n_tokens = image_input_idx.shape[0] + patch_order = np.reshape(patch_order, [-1]) + n_patches = patch_order.shape[0] + + valid = patch_order >= 0 + n_valid_patches = valid.sum() + assert len(image_input_idx) == n_valid_patches + + sorted_patch_ixs = np.zeros([n_tokens], np.int32) + sorted_patch_ixs[patch_order[valid]] = np.arange(n_valid_patches, dtype=np.int32) + + # Project the inverted mapping into same sparse structure + sorted_patch_ixs_ex = np.full(np.shape(patch_order), -1) + sorted_patch_ixs_ex[valid] = sorted_patch_ixs + + # Do the gather and then re-masked outputs that were masked in `sorted_patch_ixs` + valid = (sorted_patch_ixs_ex >= 0).astype(np.int32) + image_input_idx = image_input_idx[sorted_patch_ixs_ex*valid] + image_input_idx = image_input_idx*valid - 100*(1 - valid) + image_input_idx = np.reshape(image_input_idx, [-1, tokens_per_image]) + return image_input_idx + + def preprocess( + self, + image: np.ndarray, + image_patch_token_id: int, + image_col_token_id: int, + image_start_token_id: int, + image_end_token_id: int, + max_crops: Optional[int] = None, + overlap_margins: Optional[List[int]] = None, + base_image_input_size: Optional[Union[int, List[int]]] = None, + image_token_length_w: Optional[int] = None, + image_token_length_h: Optional[int] = None, + image_patch_size: Optional[int] = None, + **kwargs, + ): + """Preprocesses an image + + Returns: + crops: (n_crops, n_patches, patch_dim) individual crops, `n_crops` might + change between images but the other dimension are fixed + tokens: (n_tokens,) int32 tokens, pad tokens indicate where to insert the + patch features, might include other special tokens as well + image_idx: (n_crops, n_patches) index in `tokens` to put the patch features from the + crops after pooling, negative values indicates patches features to exclude + padding_mask: (n_crops, n_patches) what percent of each crop is padding, can be None + if the image mask is not being used. + """ + + max_crops = max_crops or self.max_crops + overlap_margins = overlap_margins or self.overlap_margins + base_image_input_size = base_image_input_size or self.base_image_input_size + image_token_length_w = image_token_length_w or self.image_token_length_w + image_token_length_h = image_token_length_h or self.image_token_length_h + image_patch_size = image_patch_size or self.image_patch_size + + crops, image_tokens, patch_ordering, img_mask = self.image_to_patches_and_tokens( + image, + image_patch_token_id, + image_col_token_id, + image_start_token_id, + image_end_token_id, + max_crops, + overlap_margins, + base_image_input_size, + image_token_length_w, + image_token_length_h, + image_patch_size, + ) + patch_idx = self.build_image_input_idx( + image_tokens, + patch_ordering, + image_patch_token_id, + image_token_length_w=image_token_length_w, + image_token_length_h=image_token_length_h, + ) + return crops, image_tokens, patch_idx, img_mask + + def multimodal_preprocess( + self, + images: np.ndarray, + tokens: List[int], + image_idx: np.ndarray, + sequence_length: int, + image_patch_token_id: int, + image_col_token_id: int, + image_start_token_id: int, + image_end_token_id: int, + **kwargs, + ): + """Merge images and text tokens into multi-modal features for the model + + :param images: images to use as input + :param tokens: input text tokens + :param image_idx: where to insert the images into `tokens` + :params image_patch_token_id: id to use of tokens that will contain image features + :params image_col_token_id: token id for image column special tokens + :params image_start_token_id: token id for image start special tokens + :params image_end_token_id: token id for image end special tokens + :params kwargs: override preprocessor default args + """ + max_total_crops = kwargs.get("max_crops") or self.max_crops + image_token_length_w = kwargs.get("image_token_length_w") or self.image_token_length_w + image_token_length_h = kwargs.get("image_token_length_h") or self.image_token_length_h + image_patch_size = kwargs.get("image_patch_size") or self.image_patch_size + base_image_input_size = kwargs.get("base_image_input_size") or self.base_image_input_size + image_num_patch = ( + base_image_input_size[0] // image_patch_size, + base_image_input_size[1] // image_patch_size, + ) + image_padding_mask = kwargs.get("image_padding_mask") or self.image_padding_mask + + tokens_per_image = image_token_length_w * image_token_length_h + n_pixels = image_patch_size * image_patch_size * 3 + n_patches = image_num_patch[0] * image_num_patch[1] + + if images is None: + return { + "input_ids": tokens, + } + else: + n = len(images) + all_crops = [] + all_image_idx = [] + out_tokens = [] + all_crop_masks = [] + + for ix in range(n): + token_ix = image_idx[ix] + crops, image_tokens, patch_idx, img_mask = self.preprocess( + images[ix], + image_patch_token_id, + image_col_token_id, + image_start_token_id, + image_end_token_id, + **kwargs, + ) + + if token_ix == -1: # -1 is an image inserted at the very start + start = 0 + token_ix = 0 + end = 0 + else: + start = 0 if ix == 0 else image_idx[ix-1] + 1 + end = token_ix + 1 + + all_image_idx.append(patch_idx + token_ix) + all_crops.append(crops) + out_tokens.append(tokens[start:token_ix]) + out_tokens.append(image_tokens) + if ix == (n - 1): + out_tokens.append(tokens[end:]) + if image_padding_mask: + all_crop_masks.append(img_mask) + + input_ids = np.concatenate(out_tokens, 0) + images = np.concatenate(all_crops, 0) + image_input_idx = np.concatenate(all_image_idx, 0) + if image_padding_mask: + image_masks = np.concatenate(all_crop_masks, 0) + else: + image_masks = None + + out = { + "input_ids": input_ids, + "images": images, + "image_input_idx": image_input_idx + } + if image_masks is not None: + out["image_masks"] = image_masks + return out + + +MolmoImageProcessor.register_for_auto_class() \ No newline at end of file diff --git a/pdelfin/train/molmo/modeling_molmo.py b/pdelfin/train/molmo/modeling_molmo.py new file mode 100644 index 0000000..d68da04 --- /dev/null +++ b/pdelfin/train/molmo/modeling_molmo.py @@ -0,0 +1,2367 @@ +import logging +import math +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_molmo import MolmoConfig +from torch.nn import functional as F + + +log = logging.getLogger(__name__) + + +class BufferCache(dict, MutableMapping[str, torch.Tensor]): + """ + Cache for attention biases and other things that would normally be stored as buffers. + We avoid using buffers because we've run into various issues doing so with FSDP. + In general it appears the way FSDP handles buffers is not well-defined. + It doesn't shard them but apparently it does synchronize them across processes, which we want to avoid + since (A) it isn't necessary, and (B) we sometimes have `-inf` in these biases which might get turned into + NaNs when they're synchronized due to casting or some other issue. + """ + + +class StrEnum(str, Enum): + def __str__(self) -> str: + return self.value + + def __repr__(self) -> str: + return f"'{str(self)}'" + + +class ImageProjectType(StrEnum): + mlp = "mlp" + mlpx2 = "2mlp" + linear = "linear" + + +class ImagePooling2DType(StrEnum): + attention = "attention" + attention_meanq = "attention-meanq" + attention_2wide = "attention_2wide" + attention_v2 = "attention-v2" + none = "none" + stack = "stack" + + +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" + + +def ensure_finite_(x: torch.Tensor, check_neg_inf: bool = True, check_pos_inf: bool = False): + """ + Modify ``x`` in place to replace ``float("-inf")`` with the minimum value of the dtype when ``check_neg_inf`` + is ``True`` and to replace ``float("inf")`` with the maximum value of the dtype when ``check_pos_inf`` is ``True``. + """ + if check_neg_inf: + x.masked_fill_(x == float("-inf"), torch.finfo(x.dtype).min) + if check_pos_inf: + x.masked_fill_(x == float("inf"), torch.finfo(x.dtype).max) + + +class MolmoConfigurationError(Exception): + pass + + +def _non_meta_init_device(config) -> torch.device: + if config.init_device is not None and config.init_device != "meta": + return torch.device(config.init_device) + else: + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class RotaryEmbedding(nn.Module): + """ + [Rotary positional embeddings (RoPE)](https://arxiv.org/abs/2104.09864). + """ + + def __init__(self, config: MolmoConfig, cache: BufferCache): + super().__init__() + self.config = config + self.__cache = cache + # Warm up cache. + self.get_rotary_embedding( + config.max_position_embeddings or config.max_sequence_length, + _non_meta_init_device(config) + ) + + def get_rotary_embedding(self, seq_len: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: + if ( + (pos_sin := self.__cache.get("rope_pos_sin")) is not None + and (pos_cos := self.__cache.get("rope_pos_cos")) is not None + and pos_sin.shape[-2] >= seq_len + and pos_cos.shape[-2] >= seq_len + ): + if pos_sin.device != device: + pos_sin = pos_sin.to(device) + self.__cache["rope_pos_sin"] = pos_sin + if pos_cos.device != device: + pos_cos = pos_cos.to(device) + self.__cache["rope_pos_cos"] = pos_cos + 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 + 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) + if self.config.rope_impl == "interleave": + positions = freqs.repeat_interleave(2, dim=-1) + else: + positions = torch.cat((freqs, freqs), dim=-1) + pos_sin, pos_cos = positions.sin()[None, None, :, :], positions.cos()[None, None, :, :] + self.__cache["rope_pos_sin"] = pos_sin + self.__cache["rope_pos_cos"] = pos_cos + return pos_sin, pos_cos + + def rotate_half(self, x: torch.Tensor) -> torch.Tensor: + B, nh, T, hs = x.size() + x = x.view(B, nh, T, 2, hs // 2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + def rotate_every_two(self, x: torch.Tensor) -> torch.Tensor: + B, nh, T, hs = x.size() + x = x.view(B, nh, T, hs // 2, 2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return x.view(B, nh, T, hs) + + def apply_rotary_pos_emb(self, pos_sin: torch.Tensor, pos_cos: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + if self.config.rope_impl == "interleave": + return ((t * pos_cos) + (self.rotate_every_two(t) * pos_sin)).to(t.dtype) + else: + return ((t * pos_cos) + (self.rotate_half(t) * pos_sin)).to(t.dtype) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.config.rope_full_precision: + q_, k_ = q.float(), k.float() + else: + q_, k_ = q, k + + with torch.autocast(q.device.type, enabled=False): + batch_size = q_.shape[0] + query_len, key_len = q_.shape[-2], k_.shape[-2] # could be different if layer_past not None + if position_ids is not None: + freqs_cis_len = (self.config.max_position_embeddings or self.config.max_sequence_length) + else: + freqs_cis_len = key_len + pos_sin, pos_cos = self.get_rotary_embedding(freqs_cis_len, q_.device) + pos_sin = pos_sin.type_as(q_) + pos_cos = pos_cos.type_as(q_) + if position_ids is not None: + assert query_len == key_len, "Query and key lengths must be equal when using position IDs." + pos_sin = pos_sin[0, 0][position_ids].view( + (batch_size, 1, key_len, pos_sin.shape[-1]) + ) + pos_cos = pos_cos[0, 0][position_ids].view( + (batch_size, 1, key_len, pos_cos.shape[-1]) + ) + q_ = self.apply_rotary_pos_emb( + pos_sin[:, :, key_len - query_len : key_len, :], + pos_cos[:, :, key_len - query_len : key_len, :], + q_, + ) + k_ = self.apply_rotary_pos_emb(pos_sin, pos_cos, k_) + return q_.type_as(q), k_.type_as(k) + + +class MolmoBlock(nn.Module): + """ + A base class for transformer block implementations. + """ + + def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): + super().__init__() + self.layer_id = layer_id + self.config = config + self.hidden_size = ( + config.mlp_hidden_size if config.mlp_hidden_size is not None else config.mlp_ratio * config.d_model + ) + self.__cache = cache + self._activation_checkpoint_fn = None + + # Dropout. + self.dropout = Dropout(config.residual_dropout) + + # Layer norms. + self.k_norm: Optional[LayerNormBase] = None + self.q_norm: Optional[LayerNormBase] = None + if config.attention_layer_norm: + assert config.effective_n_kv_heads is not None + self.k_norm = LayerNormBase.build( + config, + size=(config.d_model // config.n_heads) * config.effective_n_kv_heads, + elementwise_affine=config.attention_layer_norm_with_affine, + ) + self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine) + + # 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 + + # 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 + ) + + # Feed-forward output projection. + self.ff_out = nn.Linear( + int(self.act.output_multiplier * self.hidden_size), + config.d_model, + bias=config.include_bias, + device=config.init_device, + ) + self.ff_out._is_residual = True # type: ignore + + # Rotary embeddings. + if self.config.rope: + self.rotary_emb = RotaryEmbedding(config, self.__cache) + + self.flash_attn_func = None + if config.attention_type == "flash": + try: + from flash_attn import flash_attn_func # type: ignore + + self.flash_attn_func = flash_attn_func + except ModuleNotFoundError: + pass + + 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, + ) + + @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, + 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) + + 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) + + return F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + ) + + def attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + position_ids: 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]]]: + B, T, C = q.size() # batch size, sequence length, d_model + 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: + q = self.q_norm(q).to(dtype=dtype) + k = self.k_norm(k).to(dtype=dtype) + + # Move head forward to be next to the batch dim. + # 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) + # 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) + + if self.config.use_position_ids and self.config.rope: + # Apply rotary embeddings + q, k = self.rotary_emb(q, k, position_ids=position_ids) + + if layer_past is not None: + past_key, past_value = layer_past + k = torch.cat((past_key.to(k.device), k), dim=-2) + v = torch.cat((past_value.to(v.device), v), dim=-2) + + present = (k, v) if use_cache else None + query_len, key_len = q.shape[-2], k.shape[-2] # could be different if layer_past not None + + if not self.config.use_position_ids and self.config.rope: + # Apply rotary embeddings + q, k = self.rotary_emb(q, k) + + if attention_bias is not None: + # Resize and cast attention bias. + # The current dtype of the attention bias might not match the dtype that the SDP attn function will + # run in if AMP is enabled, and this can be a problem if some tokens are masked out due to padding + # as down-casting the attention bias to the autocast precision will result in -infs, which will + # cause the SDP attn function to produce NaNs. + attention_bias = self._cast_attn_bias( + attention_bias[:, :, key_len - query_len : key_len, :key_len], dtype + ) + + # Get the attention scores. + # shape: (B, nh, T, hs) + att = self._scaled_dot_product_attention( + q, + k, + v, + attn_mask=attention_bias, + dropout_p=0.0 if not self.training else self.config.attention_dropout, + response_dropout_p=0.0 if not self.training else self.config.response_attention_dropout, + is_causal=attention_bias is None, + ) + + # Re-assemble all head outputs side-by-side. + att = att.transpose(1, 2).contiguous().view(B, T, C) + + # 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, + 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: MolmoConfig, cache: BufferCache): + return MolmoSequentialBlock(layer_id, config, cache) + + +class MolmoSequentialBlock(MolmoBlock): + """ + This is a typical transformer block where the output is computed as ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). + """ + + def __init__(self, layer_id: int, config: MolmoConfig, cache: BufferCache): + super().__init__(layer_id, config, cache) + # Layer norms. + self.attn_norm = LayerNorm.build(config) + self.ff_norm = LayerNorm.build(config) + # Attention input projection. Projects x -> (q, k, v) + + head_dim = config.d_model // config.n_heads + self.fused_dims = ( + config.d_model, + config.effective_n_kv_heads * head_dim, + config.effective_n_kv_heads * head_dim, + ) + self.att_proj = nn.Linear( + config.d_model, sum(self.fused_dims), + bias=config.include_bias or config.qkv_bias, + device=config.init_device + ) + # Feed-forward input projection. + self.ff_proj = nn.Linear( + config.d_model, self.hidden_size, bias=config.include_bias, device=config.init_device + ) + + def reset_parameters(self): + super().reset_parameters() + self.attn_norm.reset_parameters() + self.ff_norm.reset_parameters() + # NOTE: the standard deviation for these weights does not depend on the layer. + init_weights( + self.config, self.att_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module + ) + init_weights( + self.config, self.ff_proj, d=self.config.d_model, layer_id=None, type_of_module=ModuleType.in_module + ) + + def forward( + self, + x: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + # Get query, key, value projections. + # shape: + # - for regular attn q, k, v: (batch_size, seq_len, d_model) + # - for multi-query attn q: (batch_size, seq_len, d_model) + # k, v: (batch_size, seq_len, d_model // n_heads) + # - for group query attn q: (batch_size, seq_len, d_model) + # k, v: (batch_size, seq_len, d_model // n_kv_heads) + + if not self.config.norm_after: + 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, layer_past=layer_past, use_cache=use_cache + ) + else: + att, cache = self.attention(q, k, v, attention_bias, position_ids=position_ids, 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) + + # Add feed-forward projection. + # shape: (batch_size, seq_len, d_model) + og_x = x + + if not self.config.norm_after: + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.ff_norm, x) # type: ignore + else: + x = self.ff_norm(x) + + x = self.ff_proj(x) + if self._activation_checkpoint_fn is not None: + x = self._activation_checkpoint_fn(self.act, x) # type: ignore + else: + x = self.act(x) + x = self.ff_out(x) + + if self.config.norm_after: + 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) + x = og_x + x + + return x, 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) -> torch.Tensor: + """ + :param input: A tensor of shape `(batch_size, seq_len, embed_dim)` + """ + if self.p == 0.0 and (self.mask_p is None or self.mask_p == 0.0): + return input + else: + if 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_default_input_size: Tuple[int, int] = (336, 336) + image_patch_size: int = 14 + image_pos_patch_size: int = 14 + image_emb_dim: int = 1024 + image_num_heads: int = 16 + image_num_key_value_heads: int = 16 + image_num_layers: int = 24 + image_head_dim: int = 64 + image_mlp_dim: int = 4096 + image_mlp_activations: str = "gelu" + image_dropout_rate: float = 0.0 + image_num_pos: int = 577 + image_norm_eps: float = 1e-5 + attention_dropout: float = 0.0 + residual_dropout: float = 0.0 + initializer_range: float = 0.02 + fsdp_wrap: bool = False + resize_mode: str = "default" + + def __post_init__(self): + self.image_default_input_size = tuple(self.image_default_input_size) # type: ignore[assignment] + + @property + def image_num_patch(self): + h, w = self.image_default_input_size + return h // self.image_patch_size, w // self.image_patch_size + + +@dataclass +class FullMolmoConfig: + d_model: int = 768 + n_heads: int = 12 + n_kv_heads: Optional[int] = None + qkv_bias: bool = False + clip_qkv: Optional[float] = None + n_layers: int = 12 + mlp_ratio: int = 4 + mlp_hidden_size: Optional[int] = None + activation_type: str = "swiglu" + block_group_size: int = 1 + rope: bool = True + rope_full_precision: bool = True + rope_theta: float = 10000. + rope_impl: str = "interleave" + vision_backbone: Optional[VisionBackboneConfig] = None + attention_type: str = "sdpa" + float32_attention: bool = True + attention_dropout: float = 0.1 + response_attention_dropout: float = 0.0 + multi_query_attention: Optional[bool] = None + attention_layer_norm: bool = False + residual_dropout: float = 0.1 + 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 + 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 + initializer_range: float = 0.02 + normalize_input_embeds: bool = False + use_position_ids: bool = True + + @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 MolmoConfigurationError( + "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 ViTMLP(nn.Module): + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + v_cfg = config.vision_backbone + + self.w1 = nn.Linear( + v_cfg.image_emb_dim, + v_cfg.image_mlp_dim, + bias=True, + device=config.init_device, + ) + # Activation function. + cfg = deepcopy(config) + cfg.activation_type = v_cfg.image_mlp_activations + self.act = Activation.build(cfg) + self.w2 = nn.Linear( + v_cfg.image_mlp_dim, + v_cfg.image_emb_dim, + bias=True, + device=config.init_device, + ) + + def reset_parameters(self): + v_cfg = self.config.vision_backbone + nn.init.trunc_normal_(self.w1.weight, std=math.sqrt(1 / v_cfg.image_emb_dim), a=-2.0, b=2.0) + nn.init.trunc_normal_(self.w2.weight, std=math.sqrt(1 / v_cfg.image_mlp_dim), a=-2.0, b=2.0) + nn.init.zeros_(self.w1.bias) + nn.init.zeros_(self.w2.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.w1(x) + x = self.act(x) + x = self.w2(x) + return x + + +class ResidualAttentionBlock(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + v_cfg = config.vision_backbone + self.attention = MultiHeadDotProductAttention(config) + self.feed_forward = ViTMLP(config) + self.attention_norm = nn.LayerNorm( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + device=config.init_device, + ) + self.ffn_norm = nn.LayerNorm( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + device=config.init_device, + ) + + def reset_parameters(self): + self.attention.reset_parameters() + self.feed_forward.reset_parameters() + self.attention_norm.reset_parameters() + self.ffn_norm.reset_parameters() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attention(self.attention_norm(x)) + x = x + self.feed_forward(self.ffn_norm(x)) + return x + + +class BlockCollection(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + self.grad_checkpointing: bool = False + + v_cfg = config.vision_backbone + self.resblocks = nn.ModuleList([ + ResidualAttentionBlock(config) for _ in range(v_cfg.image_num_layers) + ]) + + def reset_parameters(self): + for r in self.resblocks: + r.reset_parameters() + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + hidden_states = [] + for r in self.resblocks: + x = r(x) + hidden_states.append(x) + return hidden_states + + +class LayerNormFp32(nn.LayerNorm): + 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.to(torch.float32), + self.bias.to(torch.float32), self.eps) + return x.to(orig_type) + + +class VisionTransformer(nn.Module): + + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + v_cfg = config.vision_backbone + # class embeddings and positional embeddings + self.scale = v_cfg.image_emb_dim ** -0.5 + self.class_embedding = nn.Parameter( + torch.zeros(v_cfg.image_emb_dim, device=config.init_device), + ) + self.num_prefix_tokens: int = 1 + self.positional_embedding = nn.Parameter( + torch.zeros(v_cfg.image_num_pos, v_cfg.image_emb_dim, device=config.init_device), + ) + + image_patch_size = v_cfg.image_patch_size + self.patch_embedding = nn.Linear( + image_patch_size * image_patch_size * 3, + v_cfg.image_emb_dim, + bias=False, + device=config.init_device, + ) + + self.pre_ln = LayerNormFp32( + v_cfg.image_emb_dim, + eps=v_cfg.image_norm_eps, + ) + + self.transformer = BlockCollection(config) + + @torch.jit.ignore + def set_grad_checkpointing(self, enable=True): + self.transformer.grad_checkpointing = enable + + def reset_parameters(self): + nn.init.normal_(self.class_embedding, std=self.scale) + nn.init.normal_(self.positional_embedding, std=self.scale) + nn.init.normal_(self.patch_embedding.weight, std=0.02) + self.pre_ln.reset_parameters() + self.transformer.reset_parameters() + + def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: + cls_emb = self.positional_embedding[0:1] + pos_emb = self.positional_embedding[1:] + + pos_emb = pos_emb.reshape( + (int(math.sqrt(pos_emb.shape[0])), int(math.sqrt(pos_emb.shape[0])), pos_emb.shape[1]) + ) + + (patch_num_0, patch_num_1) = patch_num + + if pos_emb.shape[0] != patch_num_0 or pos_emb.shape[1] != patch_num_1: + # Dervied from https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + # antialias: default True in jax.image.resize + pos_emb = pos_emb.unsqueeze(0).permute(0, 3, 1, 2) + pos_emb = F.interpolate( + pos_emb, size=(patch_num_0, patch_num_1), mode="bicubic", align_corners=False, antialias=True, + ) + pos_emb = pos_emb.permute(0, 2, 3, 1).squeeze(0) + + pos_emb = pos_emb.reshape(-1, pos_emb.shape[-1]) + x = x + torch.cat([cls_emb[None, :, :], pos_emb[None, :, :]], dim=1).to(x.dtype) + return x + + def forward(self, x: torch.Tensor, patch_num: int = None) -> List[torch.Tensor]: + """ + : param x: (batch_size, num_patch, n_pixels) + """ + if patch_num is None: + patch_num = self.config.vision_backbone.image_num_patch + B, N, D = x.shape + + x = self.patch_embedding(x) + + # class embeddings and positional embeddings + x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1) + x = self.add_pos_emb(x, patch_num) + + x = self.pre_ln(x) + + hidden_states = self.transformer(x) + return hidden_states + + +class MultiHeadDotProductAttention(nn.Module): + def __init__(self, config: FullMolmoConfig, use_bias: bool = True, is_vit_layer: Optional[bool] = True): + super().__init__() + self.config = config + self.use_bias = use_bias + + v_cfg = config.vision_backbone + self.embed_dim = v_cfg.image_emb_dim + self.num_heads = v_cfg.image_num_heads + self.head_dim = v_cfg.image_head_dim + self.num_key_value_heads = v_cfg.image_num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.initializer_range = v_cfg.initializer_range + self.is_vit_layer = is_vit_layer + + nlayers = 1 if (is_vit_layer or config.vit_layers is None) else len(config.vit_layers) + + self.wq = nn.Linear( + nlayers * self.embed_dim, + 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) + nn.init.normal_(self.wv.weight, std=self.initializer_range) + nn.init.normal_(self.wo.weight, std=self.initializer_range) + if self.use_bias: + nn.init.constant_(self.wq.bias, 0) + nn.init.constant_(self.wk.bias, 0) + nn.init.constant_(self.wv.bias, 0) + nn.init.constant_(self.wo.bias, 0) + + def _split_heads(self, hidden_states, num_heads) -> torch.Tensor: + return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim)) + + 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: + + 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) + xk = self._split_heads(xk, self.num_key_value_heads) + xv = self._split_heads(xv, self.num_key_value_heads) + + if self.num_heads != self.num_key_value_heads: + xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads) + xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads) + + og_dtype = xq.dtype + + if self.config.float32_attention: + xq = xq.to(torch.float) + xk = xk.to(torch.float) + + if self.config.attention_type == "direct": + attn_weights = torch.einsum("...qhd,...khd->...hqk", xq / math.sqrt(xq.size(-1)), xk) + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(xq.dtype) + if self.attention_dropout is not None: + attn_weights = self.attention_dropout(attn_weights) + attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights.to(xv.dtype), xv) + + elif self.config.attention_type == "sdpa": + if self.config.float32_attention and not torch.is_autocast_enabled(): + xv = xv.to(torch.float32) + attn_output = F.scaled_dot_product_attention( + xq.transpose(1, 2).contiguous(), + xk.transpose(1, 2).contiguous(), + xv.transpose(1, 2).contiguous(), + is_causal=False, + dropout_p=self.config.vision_backbone.attention_dropout + ).transpose(1, 2) + else: + raise NotImplementedError(self.config.attention_type) + attn_output = attn_output.to(og_dtype) + attn_output = self._merge_heads(attn_output) + attn_output = self.wo(attn_output) + attn_output = self.residual_dropout(attn_output) + + return attn_output + + +class MultiHeadAttentionPool(nn.Module): + def __init__( + self, + config: FullMolmoConfig, + factor: int = 1, + use_bias: bool = True, + dropout: bool = True, + output_layer: bool = True, + mean_residual: bool = False, + query: str = "mean", + is_vit_layer: Optional[bool] = True + ): + super().__init__() + self.config = config + self.factor = factor + self.use_bias = use_bias + self.dropout = dropout + 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 + self.num_heads = v_cfg.image_num_heads + self.head_dim = v_cfg.image_head_dim * factor + self.num_key_value_heads = v_cfg.image_num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.initializer_range = v_cfg.initializer_range + + nlayers = 1 if (is_vit_layer or config.vit_layers is None) else len(config.vit_layers) + + if query != "vector": + self.wq = nn.Linear( + nlayers * input_dim, + 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: + self.wo = nn.Linear( + self.num_heads * self.head_dim, + 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) + + def reset_parameters(self): + if self.query != "vector": + nn.init.normal_(self.wq.weight, std=self.initializer_range) + nn.init.normal_(self.wk.weight, std=self.initializer_range) + nn.init.normal_(self.wv.weight, std=self.initializer_range) + if self.output_layer: + nn.init.normal_(self.wo.weight, std=self.initializer_range) + if self.use_bias: + if self.query != "vector": + nn.init.constant_(self.wq.bias, 0) + nn.init.constant_(self.wk.bias, 0) + nn.init.constant_(self.wv.bias, 0) + if self.output_layer: + nn.init.constant_(self.wo.bias, 0) + if self.query == "vector": + nn.init.normal_(self.attention_query, std=self.initializer_range) + + def _split_heads(self, hidden_states, num_heads): + return hidden_states.reshape(hidden_states.shape[:2] + (num_heads, self.head_dim)) + + def _merge_heads(self, hidden_states): + return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,)) + + def forward(self, inputs_kv: torch.Tensor) -> torch.Tensor: + + xk, xv = self.wk(inputs_kv), self.wv(inputs_kv) + + if self.query == "mean": + inputs_q = inputs_kv.mean(dim=1, keepdim=True) + xq = self.wq(inputs_q) + elif self.query == "first": + inputs_q = inputs_kv[:, :1] + xq = self.wq(inputs_q) + elif self.query == "vector": + xq = self.attention_query.expand(inputs_kv.size(0), -1, -1) + elif self.query == "constant": + inputs_q = torch.ones_like(inputs_kv[:, :1]) / math.sqrt(inputs_kv.shape[-1]) + xq = self.wq(inputs_q) + else: + raise ValueError(f"Unknown query type: {self.query}") + + xq = self._split_heads(xq, self.num_heads) + xk = self._split_heads(xk, self.num_key_value_heads) + xv = self._split_heads(xv, self.num_key_value_heads) + + if self.num_heads != self.num_key_value_heads: + xk = xk.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads) + xv = xv.repeat_interleave(self.num_key_value_groups, dim=2, output_size=self.num_heads) + + xq = xq.to(torch.float) + xk = xk.to(torch.float) + + xq = xq / math.sqrt(xq.size(-1)) + attn_weights = torch.einsum("...qhd,...khd->...hqk", xq, xk) + + attn_weights = F.softmax(attn_weights, dim=-1).to(xq.dtype) + + attn_weights = self.attention_dropout(attn_weights).to(xv.dtype) + + attn_output = torch.einsum("...hqk,...khd->...qhd", attn_weights, xv) + attn_output = self._merge_heads(attn_output) + if self.output_layer: + attn_output = self.wo(attn_output) + if self.dropout: + attn_output = self.residual_dropout(attn_output) + if self.mean_residual: + attn_output += inputs_kv.mean(dim=1, keepdim=True) + + return attn_output + + +class MLP(nn.Module): + def __init__(self, config: FullMolmoConfig, 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 + ) + self.initializer_range = config.initializer_range + + self.w1 = nn.Linear( + input_dim, + 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) + 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 OLMoVisionBackbone(nn.Module): + def __init__(self, config: FullMolmoConfig): + 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}: + self.image_pooling_2d = MultiHeadDotProductAttention(config, is_vit_layer=False) + input_dim = config.vision_backbone.image_emb_dim + elif config.image_pooling_2d == ImagePooling2DType.attention_2wide: + cfg = deepcopy(config) + cfg.vision_backbone.image_emb_dim *= 2 + cfg.vision_backbone.image_head_dim *= 2 + self.image_pooling_2d = MultiHeadDotProductAttention(cfg, is_vit_layer=False) + input_dim = cfg.vision_backbone.image_emb_dim + elif config.image_pooling_2d == ImagePooling2DType.attention_v2: + assert config.vit_layers is not None + use_bias = True + dropout = True + output_layer = True + query = "mean" + mean_residual = False + factor = len(config.vit_layers) + self.image_pooling_2d = MultiHeadAttentionPool( + config, + factor=factor, + use_bias=use_bias, + dropout=dropout, + output_layer=output_layer, + mean_residual=mean_residual, + query=query, + is_vit_layer=False, + ) + input_dim = config.vision_backbone.image_emb_dim * factor + elif config.image_pooling_2d in [ImagePooling2DType.none, ImagePooling2DType.stack]: + self.image_pooling_2d = None + nlayers = 1 if config.vit_layers is None else len(config.vit_layers) + 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: + self.image_projector = MLP(mlp_config, input_dim) + elif config.image_projector == ImageProjectType.linear: + self.image_projector = nn.Linear( + input_dim, + config.d_model, + bias=False, + device=config.init_device, + ) + else: + raise NotImplementedError(f"Unknown image projector: {config.image_projector}") + + self.image_feature_dropout = Dropout(config.image_feature_dropout) + + def reset_parameters(self): + if self.image_pooling_2d is not None: + self.image_pooling_2d.reset_parameters() + if self.config.image_projector == "2mlp": + for module in self.image_projector: + module.reset_parameters() + elif self.config.image_projector == "linear": + nn.init.xavier_uniform_(self.image_projector.weight) + else: + self.image_projector.reset_parameters() + + 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: FullMolmoConfig): + super().__init__(config) + v_cfg = self.config.vision_backbone + self.grad_checkpointing = False + + self.num_prefix_tokens = self.image_vit.num_prefix_tokens + assert self.num_prefix_tokens in {0, 1}, "Only 0 or 1 prefix tokens are supported" + + self.pad_embed = None + if config.image_padding_embed: + image_dim = v_cfg.image_emb_dim*len(self.config.vit_layers) + if config.image_padding_embed in ["pad_embed", "regress"]: + self.pad_embed = nn.Parameter( + torch.zeros((image_dim,), device=config.init_device)) + elif config.image_padding_embed == "pad_and_partial_pad": + self.pad_embed = nn.Parameter( + torch.zeros((2, image_dim), device=config.init_device)) + else: + raise ValueError(config.image_padding_embed) + + def reset_parameters(self): + super().reset_parameters() + self.image_vit.reset_parameters() + + def encode_image(self, images: torch.Tensor) -> torch.Tensor: + """ + : param images: (batch_size, num_crops, num_patch, n_pixels) + """ + cfg = self.config + v_cfg = self.config.vision_backbone + B, T, N, D = images.shape + + mask = ~torch.all(images.view(B * T, N, D) == -1, dim=(1, 2), keepdim=True) + + # Output all hidden states + # n_layers x (batch_num_crops, (1+)n_tokens, image_emb_dim) + images = images.view(B * T, N, D) + image_features = self.image_vit(images) + + if cfg.vit_layers is not None: + 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=image_features.dtype) + all_pad = all_pad.to(dtype=image_features.dtype) + 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) + + # 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: FullMolmoConfig, + module: Union[nn.Linear, nn.Embedding], + d: Optional[int] = None, + layer_id: Optional[int] = None, + std_factor: float = 1.0, + type_of_module: Optional[ModuleType] = None, +) -> None: + d = d if d is not None else config.d_model + std = config.init_std * std_factor + if config.init_cutoff_factor is not None: + cutoff_value = config.init_cutoff_factor * std + nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value) + else: + nn.init.normal_(module.weight, mean=0.0, std=std) + + +class LlamaSwiGLU(nn.Module): + def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: + return F.silu(x1) * x2 + + @property + def output_multiplier(self) -> float: + return 0.5 + + +class SwiGLU(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, gate = x.chunk(2, dim=-1) + return F.silu(gate) * x + + @property + def output_multiplier(self) -> float: + return 0.5 + + +class Activation(nn.Module): + def __init__(self, config: FullMolmoConfig): + super().__init__() + self.config = config + + def forward(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + @property + def output_multiplier(self) -> float: + raise NotImplementedError + + @classmethod + def build(cls, config: FullMolmoConfig) -> 'Activation': + if config.activation_type == "quick_gelu": + return QuickGELU(config) + elif config.activation_type == "gelu": + return cast(Activation, GELU(approximate="none")) + elif config.activation_type == "gelu_tanh": + return cast(Activation, GELU(approximate="tanh")) + elif config.activation_type == "relu": + return cast(Activation, ReLU(inplace=False)) + elif config.activation_type == "silu": + return cast(Activation, SiLU(inplace=False)) + # elif config.activation_type == "llama_geglu": + # return LlamaGEGLU(config) + # elif config.activation_type == "llama_geglu_tanh": + # return LlamaGEGLUTanh(config) + elif config.activation_type == "llama_swiglu": + return LlamaSwiGLU() + elif config.activation_type == "swiglu": + return SwiGLU() + else: + raise NotImplementedError(f"Unknown activation: '{config.activation_type}'") + + +class QuickGELU(Activation): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(1.702 * x) + + @property + def output_multiplier(self) -> float: + return 1.0 + + +class GELU(nn.GELU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +class ReLU(nn.ReLU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +class SiLU(nn.SiLU): + @property + def output_multiplier(self) -> float: + return 1.0 + + +def causal_attention_bias(seq_len: int, device: torch.device) -> torch.FloatTensor: + att_bias = torch.triu( + torch.ones(seq_len, seq_len, device=device, dtype=torch.float), + diagonal=1, + ) + att_bias.masked_fill_(att_bias == 1, torch.finfo(att_bias.dtype).min) + return att_bias.view(1, 1, seq_len, seq_len) # type: ignore + + +def get_causal_attention_bias(cache: BufferCache, seq_len: int, device: torch.device) -> torch.Tensor: + if (causal_bias := cache.get("causal_attention_bias")) is not None and causal_bias.shape[-1] >= seq_len: + if causal_bias.device != device: + causal_bias = causal_bias.to(device) + cache["causal_attention_bias"] = causal_bias + return causal_bias + with torch.autocast(device.type, enabled=False): + causal_bias = causal_attention_bias(seq_len, device) + cache["causal_attention_bias"] = causal_bias + return causal_bias + + +class LayerNormBase(nn.Module): + def __init__( + self, + config: MolmoConfig, + *, + size: Optional[int] = None, + elementwise_affine: Optional[bool] = True, + eps: float = 1e-05, + weight_initializer: Optional[Callable] = torch.ones, + bias_initializer: Optional[Callable] = torch.zeros, + ): + super().__init__() + self.config = config + self.eps = self.config.layer_norm_eps or eps + self.normalized_shape = (size or config.d_model,) + if elementwise_affine or (elementwise_affine is None and self.config.layer_norm_with_affine): + self.weight = nn.Parameter(weight_initializer(self.normalized_shape, device=config.init_device)) + use_bias = self.config.bias_for_layer_norm + if use_bias is None: + use_bias = self.config.include_bias + if use_bias: + self.bias = nn.Parameter(bias_initializer(self.normalized_shape, device=config.init_device)) + else: + self.register_parameter("bias", None) + else: + self.register_parameter("bias", None) + self.register_parameter("weight", None) + + @classmethod + def build(cls, config: FullMolmoConfig, size: Optional[int] = None, **kwargs): + if config.layer_norm_type == "default": + return LayerNorm(config, size=size, low_precision=False, **kwargs) + elif config.layer_norm_type == "low_precision": + return LayerNorm(config, size=size, low_precision=True, **kwargs) + elif config.layer_norm_type == "rms": + return RMSLayerNorm(config, size=size, **kwargs) + else: + raise NotImplementedError(f"Unknown LayerNorm type: '{config.layer_norm_type}'") + + +class RMSLayerNorm(LayerNormBase): + """ + RMS layer norm, a simplified :class:`LayerNorm` implementation + """ + + def __init__( + self, + config: FullMolmoConfig, + size: Optional[int] = None, + elementwise_affine: Optional[bool] = None, + eps: float = 1e-5, + ): + super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + with torch.autocast(enabled=False, device_type=x.device.type): + og_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + x = x.to(og_dtype) + + if self.weight is not None: + if self.bias is not None: + return self.weight * x + self.bias + else: + return self.weight * x + else: + return x + + +class LayerNorm(LayerNormBase): + """ + The default :class:`LayerNorm` implementation which can optionally run in low precision. + """ + + def __init__( + self, + config: FullMolmoConfig, + size: Optional[int] = None, + low_precision: bool = False, + elementwise_affine: Optional[bool] = None, + eps: float = 1e-05, + ): + super().__init__(config, size=size, elementwise_affine=elementwise_affine, eps=eps) + self.low_precision = low_precision + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.low_precision: + module_device = x.device + downcast_x = self._cast_if_autocast_enabled(x) + downcast_weight = ( + self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight + ) + downcast_bias = self._cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias + with torch.autocast(enabled=False, device_type=module_device.type): + return F.layer_norm( + downcast_x, self.normalized_shape, weight=downcast_weight, bias=downcast_bias, eps=self.eps + ) + else: + return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps) + + +class Molmo(nn.Module): + def __init__(self, config: FullMolmoConfig, init_params: bool = True): + super().__init__() + self.config = config + self.__cache = BufferCache() + + # Validate config. + if self.config.embedding_size is not None and self.config.embedding_size != self.config.vocab_size: + if self.config.embedding_size < self.config.vocab_size: + raise MolmoConfigurationError("embedding size should be at least as big as vocab size") + elif self.config.embedding_size % 128 != 0: + import warnings + + warnings.warn( + "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning + ) + 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 + + wte = None + if self.config.additional_vocab_size is not None: + wte = Embedding( + config.embedding_size or config.vocab_size, + config.additional_vocab_size, + config.d_model, + device=config.init_device, + initializer_range=config.initializer_range, + new_embed_initializer_range=config.new_embedding_init_range + ) + else: + wte=nn.Embedding( + config.embedding_size or config.vocab_size, config.d_model, device=config.init_device + ) + + self.transformer = nn.ModuleDict( + dict( + wte=wte, + emb_drop=Dropout(config.embedding_dropout), + ln_f=LayerNorm.build(config), + ) + ) + + blocks = [MolmoBlock.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.rope: + self.transformer.update( + {"wpe": nn.Embedding(config.max_sequence_length, config.d_model, device=config.init_device)} + ) + if not config.weight_tying: + self.transformer.update( + { + "ff_out": nn.Linear( + config.d_model, + config.embedding_size or config.vocab_size, + bias=config.include_bias, + device=config.init_device, + ) + } + ) + + self.vision_backbone: Optional[OLMoVisionBackbone] = None + if config.vision_backbone is not None: + self.vision_backbone = OLMoPretrainedVisionBackbone(config) + + self.__num_fwd_flops: Optional[int] = None + + def reset_parameters(self): + if self.vision_backbone is not None: + self.vision_backbone.reset_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.transformer.ln_f.reset_parameters() # type: ignore + + if hasattr(self.transformer, "ff_out"): + nn.init.normal_(self.transformer.ff_out, mean=0.0, std=0.02) + + if self.config.block_group_size == 1: + for block in self.transformer.blocks: + block.reset_parameters() + else: + for block_group in self.transformer.block_groups: + block_group.reset_parameters() + + + def forward( + self, + input_ids: torch.LongTensor, + input_embeddings: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + response_mask: Optional[torch.Tensor] = None, + images: Optional[torch.Tensor] = None, + image_masks: Optional[torch.Tensor] = None, + image_input_idx: Optional[torch.Tensor] = None, + subsegment_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[Sequence[Tuple[torch.Tensor, torch.Tensor]]] = None, + use_cache: bool = False, + last_logits_only: bool = False, + output_hidden_states: Optional[bool] = None, + append_last_valid_logits: Optional[torch.Tensor] = None, + ) -> ModelOutput: + """ + :param input_ids: A tensor of shape `(batch_size, seq_len)`. + :param input_embeddings: A tensor of shape `(batch_size, seq_len, d_model)` with input + embeddings. When provided, it is treated as the output of the input embedding layer. + :param attention_mask: A tensor of shape `(batch_size, seq_len)` that indicates + 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 + is a response token. A `0` means that the corresponding token is not + a response token. + :param past_key_values: Pre-computed keys and values for each attention block. + Can be used to speed up sequential decoding. The `input_ids` which have + their past given to this model should not be passed as `input_ids` as they have already been computed. + :param use_cache: If `True`, return key and value tensors for each block. + :param last_logits_only: If `True`, only compute the logits for the last token of each sequence. + This can speed up decoding when you only care about the next token. + """ + output_hidden_states = output_hidden_states if output_hidden_states is not None else False + + if past_key_values: + assert len(past_key_values) == self.config.n_layers + + has_image = images is not None + + assert not (has_image and input_embeddings is not None), "Cannot provide both images and input embeddings." + assert not (has_image and past_key_values is not None), "Cached key and values should not be used with images." + + batch_size, seq_len = input_ids.size() if input_embeddings is None else input_embeddings.size()[:2] + if past_key_values is None: + past_length = 0 + else: + past_length = past_key_values[0][0].size(-2) + + 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) + attention_mask = ( + subsegment_mask.to(attention_mask.dtype) * + attention_mask.unsqueeze(2) * + attention_mask.unsqueeze(1)) + if position_ids is None: + raise ValueError(f"Positioned ids must be given if using subsegment_ids") + else: + if self.config.use_position_ids and position_ids is None: + position_ids = torch.clamp( + torch.cumsum(attention_mask.to(torch.int32), dim=-1) - 1, + min=0, + ).broadcast_to((batch_size, attention_mask.shape[-1])) + + # Get embeddings of input. + # shape: (batch_size, seq_len, d_model) + if input_ids is not None: + input_ids = input_ids * (input_ids != -1).to(input_ids.dtype) + x = self.transformer.wte(input_ids) if input_embeddings is None else input_embeddings # type: ignore + + num_image: Optional[int] = None + if images is not None: + # shape: (batch_size, num_image, num_patch, d_model) + # cls_embed: (batch_size, num_image, d_model) + image_features, cls_embed = self.vision_backbone(images, image_masks) + num_image, num_patch = image_features.shape[1:3] + assert image_input_idx.shape == (batch_size, num_image, num_patch) + + # inster the image feature into the embedding. + image_features = image_features.view(batch_size, num_image * num_patch, -1) + image_input_idx = image_input_idx.view(batch_size, num_image * num_patch) + + valid = image_input_idx >= 0 + batch_idx = torch.arange(batch_size, device=x.device) + batch_idx = torch.tile(batch_idx[:, None], [1, image_features.shape[1]]) + + # For hf demo/endpoint + image_features = image_features.to(x.device) + + x[batch_idx[valid], image_input_idx[valid]] += image_features[valid] + + if not self.config.rope: + # Get positional embeddings. + # shape: (1, seq_len) + pos = torch.arange(past_length, past_length + seq_len, dtype=torch.long, device=x.device).unsqueeze(0) + # shape: (1, seq_len, d_model) + pos_emb = self.transformer.wpe(pos) # type: ignore + x = pos_emb + x + + # Add input + positional embeddings and apply dropout. + # shape: (batch_size, seq_len, d_model) + x = self.transformer.emb_drop(x) # type: ignore + + # normalized + if self.config.normalize_input_embeds: + x = x * (self.config.d_model ** 0.5) + + # Transform the attention mask into what the blocks expect. + if attention_mask is not None: + # shape: (batch_size, 1, 1, seq_len) + if len(attention_mask.shape) == 2: + attention_mask = attention_mask[:, :past_length + seq_len] + attention_mask = attention_mask.to(dtype=torch.float).view(batch_size, -1)[:, None, None, :] + else: + attention_mask = attention_mask.unsqueeze(1).to(dtype=torch.float) + attention_mask = (1.0 - attention_mask) * torch.finfo(attention_mask.dtype).min + + # Merge attention mask with attention bias. + if ( + attention_bias is not None + or attention_mask is not None + # 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: + 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) + attention_bias.masked_fill_(attention_bias == 0.0, torch.finfo(attention_bias.dtype).min) + + # Transform to the right shape and data type. + mask_len = seq_len + if attention_mask is not None: + mask_len = attention_mask.shape[-1] + elif past_key_values is not None: + mask_len = past_key_values[0][0].shape[-2] + seq_len + attention_bias = attention_bias[:, :, :mask_len, :mask_len].to(dtype=torch.float) + + # Add in the masking bias. + if attention_mask is not None: + attention_bias = attention_bias + attention_mask + # Might get -infs after adding attention mask, since dtype.min + dtype.min = -inf. + # `F.scaled_dot_product_attention()` doesn't handle -inf like you'd expect, instead + # it can produce NaNs. + ensure_finite_(attention_bias, check_neg_inf=True, check_pos_inf=False) + + attn_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None + + # decoder layers + 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, layer_past=layer_past, use_cache=use_cache) + + if attn_key_values is not None: + assert cache is not None + attn_key_values.append(cache) + else: + for group_idx, block_group in enumerate(self.transformer.block_groups): + if output_hidden_states: + # add hidden states + all_hidden_states.append(x) + + layers_past = ( + None + if past_key_values is None + else past_key_values[ + group_idx * self.config.block_group_size : (group_idx + 1) * self.config.block_group_size + ] + ) + x, cache = block_group( + x, attention_bias=attention_bias, position_ids=position_ids, layers_past=layers_past, use_cache=use_cache + ) + if attn_key_values is not None: + assert cache is not None + attn_key_values.extend(cache) + + if last_logits_only: + # shape: (batch_size, 1, d_model) + if append_last_valid_logits is not None: + last_valid_output = x[ + torch.arange(x.shape[0], device=x.device), append_last_valid_logits.to(x.device)] + x = last_valid_output.unsqueeze(1) + else: + x = x[:, -1, :].unsqueeze(1) + + # Apply final layer norm. + # shape: (batch_size, seq_len or 1, d_model) + x = self.transformer.ln_f(x) # type: ignore + if output_hidden_states: + # add final hidden state post-final-layernorm, following HuggingFace's convention + all_hidden_states.append(x) + + # 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 + else: + logits = self.transformer.ff_out(x) # type: ignore + if self.config.scale_logits: + logits.mul_(1 / math.sqrt(self.config.d_model)) + + if not last_logits_only and append_last_valid_logits is not None: + last_valid_logit = logits[ + torch.arange(logits.shape[0], device=logits.device), append_last_valid_logits] + logits = torch.cat([logits[:, :-1], last_valid_logit[:, None]], dim=1) + + return ModelOutput(logits=logits, attn_key_values=attn_key_values, hidden_states=tuple(all_hidden_states) if output_hidden_states else None) # type: ignore[arg-type] + + +class MolmoForCausalLM(PreTrainedModel): + config_class = MolmoConfig + base_model_prefix = "model" + _no_split_modules = ["MolmoBlock"] + + def __init__(self, config: MolmoConfig, model: Optional[Molmo] = None, init_params: bool = False): + super().__init__(config) + + if not model: + full_config = FullMolmoConfig( + image_padding_embed="pad_and_partial_pad", + image_pooling_2d="attention-meanq", + attention_layer_norm=config.attention_layer_norm, + rope_impl="llama", + vocab_size=config.vocab_size, + max_sequence_length=config.max_position_embeddings, + qkv_bias=config.qkv_bias, + norm_after=config.norm_after, + embedding_size=config.embedding_size, + attention_type="sdpa", + embedding_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=config.rope_theta, + layer_norm_eps=config.layer_norm_eps, + layer_norm_type=config.layer_norm_type, + vit_layers=[-2, -9], + vision_backbone=VisionBackboneConfig( + image_default_input_size=(336, 336), + image_patch_size=14, + image_pos_patch_size=14, + image_emb_dim=1024, + image_num_heads=16, + image_num_key_value_heads=16, + image_num_layers=23, + image_head_dim=64, + image_mlp_dim=4096, + image_mlp_activations="quick_gelu", + image_dropout_rate=0.0, + image_num_pos=577, + image_norm_eps=1e-5, + attention_dropout=0.0, + residual_dropout=0.0, + initializer_range=0.02, + ) + ) + self.model = Molmo(full_config, init_params=init_params) + else: + self.model = model + + + def forward( + self, + input_ids: torch.LongTensor = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + response_mask: Optional[torch.Tensor] = None, + images: Optional[torch.Tensor] = None, + image_masks: Optional[torch.Tensor] = None, + image_input_idx: Optional[torch.Tensor] = None, + subsegment_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + labels: Optional[torch.LongTensor] = None, + loss_masks: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + last_logits_only: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + append_last_valid_logits: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[ + Cache + ] = None, # This is a hack mitigation of an issue in transformers `4.39.x` https://github.com/huggingface/transformers/issues/29426 + ) -> Union[Tuple, CausalLMOutputWithPast]: + if use_cache is None: + use_cache = self.config.use_cache + + if output_attentions: + raise ValueError("output_attentions is not yet supported in Molmo") + + 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( + input_ids=input_ids, + input_embeddings=inputs_embeds, + attention_mask=attention_mask, + attention_bias=attention_bias, + response_mask=response_mask, + images=images, + image_masks=image_masks, + image_input_idx=image_input_idx, + subsegment_ids=subsegment_ids, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + last_logits_only=last_logits_only, + output_hidden_states=output_hidden_states, + append_last_valid_logits=append_last_valid_logits, + ) + + logits = outputs.logits + hidden_states = outputs.hidden_states + + loss = None + if labels is not None: + if loss_masks is not None: + loss_masks = loss_masks * (loss_masks > 0) + batch_size_in_tokens = max(loss_masks.sum().item(), 1) + labels = labels.long() + labels.masked_fill_(~(loss_masks > 0), -100) + labels = labels.view(-1) + logits_for_loss = logits.to(torch.float32).view(-1, logits.size(-1)) + loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none') + loss = loss_fct(logits_for_loss, labels) + loss = loss.view(input_ids.shape[0], -1) + loss = loss * loss_masks + loss = loss.sum() / batch_size_in_tokens + use_zloss = getattr(self.config, "softmax_auxiliary_loss", False) + if use_zloss: + z_squared = logits_for_loss.logsumexp(-1).pow(2) + z_loss = self.config.softmax_auxiliary_loss_scale * z_squared + z_loss = z_loss.view(input_ids.shape[0], -1) + z_loss = z_loss * loss_masks + z_loss = z_loss.sum() / batch_size_in_tokens + loss += z_loss + else: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = torch.nn.CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.embedding_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.attn_key_values, + hidden_states=hidden_states, + ) + + def can_generate(self) -> bool: + return True + + @torch.no_grad() + def generate_from_batch( + self, + batch: Dict[str, Any], + generation_config: Optional[GenerationConfig] = None, + **kwargs, + ): + if generation_config is not None: + assert generation_config.use_cache + + images = batch.get("images") + image_masks = batch.get("image_masks") + image_input_idx = batch.get("image_input_idx") + + # Validate inputs. + input_ids = batch["input_ids"] + batch_size, seq_len = input_ids.shape + attention_mask = batch.get("attention_mask", None) + max_new_tokens = generation_config.max_new_tokens + assert max_new_tokens is not None + mask_len = seq_len + max_new_tokens if self.config.use_position_ids else seq_len + 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"], + 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 + + def prepare_inputs_for_generation( + self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple]] = None, **kwargs + ): + if past_key_values: + # This is because we want the model to only process the last generated token. + input_ids = input_ids[:, -1:] + + if self.config.use_position_ids: + attention_mask = kwargs.get("attention_mask") + images = kwargs.get("images") + image_masks = kwargs.get("image_masks") + image_input_idx = kwargs.get("image_input_idx") + position_ids = kwargs.get("position_ids") + append_last_valid_logits = kwargs.get("append_last_valid_logits") + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": True, + "last_logits_only": True, + } + if past_key_values is None: + model_inputs["images"] = images + 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: + model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values} + + model_inputs.update(kwargs) + model_inputs["use_cache"] = kwargs.pop("use_cache", self.config.use_cache) + return model_inputs + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: Dict[str, Any], + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> Dict[str, Any]: + if self.config.use_position_ids: + model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1 + if "append_last_valid_logits" in model_kwargs: + del model_kwargs["append_last_valid_logits"] + if "images" in model_kwargs: + del model_kwargs["images"] + del model_kwargs["image_masks"] + del model_kwargs["image_input_idx"] + cache_name, cache = super()._extract_past_from_model_output(outputs) + model_kwargs[cache_name] = cache + model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens + return model_kwargs + + def get_input_embeddings(self) -> torch.nn.Module: + return self.model.transformer.wte + + def set_input_embeddings(self, value: torch.nn.Module): + self.model.transformer.wte = value + + def get_output_embeddings(self): + if self.config.weight_tying: + return self.model.transformer.wte + else: + return self.model.transformer.ff_out + + def set_output_embeddings(self, value: torch.nn.Module): + if self.config.weight_tying: + self.model.transformer.wte = value + else: + self.model.transformer.ff_out = value + + def tie_weights(self): + """ + This function is intentionally left as a no-op. + + Weight tying is handled as follows: + - When the model is initialized, the `ff_out` layer is conditionally defined based on the `weight_tying` configuration. + See: `if not config.weight_tying: self.transformer.update(...)` in `olmo/model.py`. + - When computing logits, the `wte` weights are used directly if `weight_tying` is enabled. + See: `if self.config.weight_tying: logits = F.linear(x, self.transformer.wte.weight, None)` in the `forward` method. + + Therefore, there is no need to explicitly tie the weights in this function. + """ + pass + + def resize_token_embeddings( + self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None + ) -> 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 + vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just + returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything. + pad_to_multiple_of (`int`, *optional*): + If set will pad the embedding matrix to a multiple of the provided value. If `new_num_tokens` is set to + `None` will just pad the embedding to a multiple of `pad_to_multiple_of`. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more + details about this, or help on choosing the correct value for resizing, refer to this guide: + https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc + + Return: + `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. + + Note: + This method differs from the base class implementation by resizing the `embedding_size` attribute of the + model configuration instead of the `vocab_size`. It also includes a warning if the resized `embedding_size` + is less than the `vocab_size`. In OLMo, `embedding_size` refers to the dimensionality of the model's token + embeddings, while `vocab_size` refers to the number of unique tokens in the vocabulary. + """ + model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of) + if new_num_tokens is None and pad_to_multiple_of is None: + return model_embeds + + # Update base model and current model config + self.config.embedding_size = model_embeds.weight.shape[0] + self.model.config.embedding_size = model_embeds.weight.shape[0] + + # Check if the embedding size is less than the vocab size + if self.config.embedding_size < self.config.vocab_size: + warning_message = ( + f"Resizing token embeddings to size {self.config.embedding_size}, which is less than the vocab size " + f"{self.config.vocab_size} defined in the model configuration. Make sure your tokenizer's vocabulary " + "size is less than or equal to the new token embedding size." + ) + log.warning(warning_message) + + # Tie weights again if needed + self.tie_weights() + + return model_embeds + + +# Always register for multi-modal features +AutoModelForCausalLM.register(MolmoConfig, MolmoForCausalLM) \ No newline at end of file diff --git a/pdelfin/train/molmo/preprocessing_molmo.py b/pdelfin/train/molmo/preprocessing_molmo.py new file mode 100644 index 0000000..3598399 --- /dev/null +++ b/pdelfin/train/molmo/preprocessing_molmo.py @@ -0,0 +1,192 @@ +""" +Processor class for Molmo. +""" + +from typing import Optional + +import PIL +from PIL import ImageOps +from PIL.Image import Image + +try: + from typing import Unpack +except ImportError: + from typing_extensions import Unpack + +import numpy as np +import torch + +from transformers.image_utils import ImageInput +from transformers.processing_utils import ( + TextKwargs, + ProcessingKwargs, + ProcessorMixin, +) + +from transformers.tokenization_utils_base import TextInput, PreTokenizedInput +from transformers.utils import logging + +from transformers import AutoTokenizer +from .image_preprocessing_molmo import MolmoImagesKwargs, MolmoImageProcessor + + +logger = logging.get_logger(__name__) + + +DEFAULT_IMAGE_PATCH_TOKEN = f"" +DEFAULT_IM_START_TOKEN = f"" +DEFAULT_IM_END_TOKEN = f"" +DEFAULT_IM_COL_TOKEN = f"" +IMAGE_PROMPT = "<|image|>" + +EXTRA_TOKENS = (DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_COL_TOKEN, IMAGE_PROMPT) + + +def get_special_token_ids(tokenizer): + ids = tokenizer.encode("".join(EXTRA_TOKENS), add_special_tokens=False) + assert len(ids) == len(EXTRA_TOKENS) + return {k: i for k, i in zip(EXTRA_TOKENS, ids)} + + +class MolmoTextKwargs(TextKwargs, total=False): + style: Optional[str] + system_prompt: Optional[str] + message_format: Optional[str] + always_start_with_space: Optional[bool] + sequence_length: Optional[int] + + +class MolmoProcessorKwargs(ProcessingKwargs, total=False): + text_kwargs: MolmoTextKwargs + images_kwargs: MolmoImagesKwargs + _defaults = { + "images_kwargs": { + "max_crops": 12, + "overlap_margins": [4, 4], + "base_image_input_size": [336, 336], + "image_token_length_w": 12, + "image_token_length_h": 12, + "image_patch_size": 14, + "image_padding_mask": True, + }, + "text_kwargs": { + "style": "long_caption", + "system_prompt": "none", + "message_format": "role", + "always_start_with_space": True, + "sequence_length": 1536, + "padding": False, + }, + } + + +class MolmoProcessor(ProcessorMixin): + attributes = ["image_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = ("GPT2Tokenizer", "GPT2TokenizerFast") + + def __init__(self, image_processor: MolmoImageProcessor = None, tokenizer : AutoTokenizer = None, **kwargs): + # self.image_processor = image_processor + # self.tokenizer = tokenizer + super().__init__(image_processor, tokenizer) + self._special_tokens = None + + @property + def special_token_ids(self): + if self._special_tokens is None: + self._special_tokens = get_special_token_ids(self.tokenizer) + return self._special_tokens + + def get_tokens_input(self, prompt, message_format, always_start_with_space): + if message_format == "none" or message_format is None: + pass + elif message_format == "role": + prompt = "User: " + prompt + " Assistant:" + else: + raise NotImplementedError(f"Message format {message_format} not implemented") + + if always_start_with_space: + prompt = " " + prompt + + tokens = self.tokenizer.encode(prompt, add_special_tokens=False) + + return tokens + + def process( + self, + text: TextInput = None, + images: ImageInput = None, + *, + tokens: Optional[PreTokenizedInput] = None, + **kwargs: Unpack[MolmoProcessorKwargs], + ): + output_kwargs = self._merge_kwargs( + MolmoProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if tokens is None: + tokens = self.get_tokens_input( + text, + output_kwargs["text_kwargs"]["message_format"], + output_kwargs["text_kwargs"]["always_start_with_space"], + ) + + image_token_id = self.special_token_ids[IMAGE_PROMPT] + + if images is not None: + if not isinstance(images, (list, tuple)): + images = [images] + image_arrays = [] + for image in images: + if isinstance(image, Image): + image = image.convert("RGB") + # Handle images with EXIF orientation tags, which PIL will ignore by default + # https://github.com/python-pillow/Pillow/issues/4703 + img = ImageOps.exif_transpose(image) + image_arrays.append(np.array(image)) + else: + assert len(image.shape) == 3 and image.shape[-1] == 3 + image_arrays.append(image.astype(np.uint8)) + images = image_arrays + # For now only support inserting images at the start + image_idx = [-1]*len(images) + else: + image_idx = None + + sequence_length = output_kwargs["text_kwargs"]["sequence_length"] + + image_patch_token_id = self.special_token_ids[DEFAULT_IMAGE_PATCH_TOKEN] + image_col_token_id = self.special_token_ids[DEFAULT_IM_COL_TOKEN] + image_start_token_id = self.special_token_ids[DEFAULT_IM_START_TOKEN] + image_end_token_id = self.special_token_ids[DEFAULT_IM_END_TOKEN] + out = self.image_processor.multimodal_preprocess( + images=images, + image_idx=image_idx, + tokens=np.asarray(tokens).astype(np.int32), + sequence_length=sequence_length, + image_patch_token_id=image_patch_token_id, + image_col_token_id=image_col_token_id, + image_start_token_id=image_start_token_id, + image_end_token_id=image_end_token_id, + **output_kwargs["images_kwargs"] + ) + + # Prepend BOS + # qwen2 and olmo do not have a BOS, and instead use EOS as a generic seperator token. + bos = self.tokenizer.bos_token_id or self.tokenizer.eos_token_id + decoder_input_tokens = np.pad(out["input_ids"], [[1, 0]], constant_values=bos) + out["input_ids"] = decoder_input_tokens + if "image_input_idx" in out: + # Shift patch mapping up by one since we added BOS + image_input_idx = out["image_input_idx"] + out["image_input_idx"] = np.where(image_input_idx < 0, image_input_idx, image_input_idx + 1) + + for k, v in out.items(): + out[k] = torch.from_numpy(v) + + return out + + +MolmoProcessor.register_for_auto_class() From f42bb02fce67a9dd02008fb9bd9bd7ecf2f9443f Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Thu, 23 Jan 2025 15:18:22 -0800 Subject: [PATCH 2/7] Manually adding gradient checkpointing --- pdelfin/train/molmo/modeling_molmo.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/pdelfin/train/molmo/modeling_molmo.py b/pdelfin/train/molmo/modeling_molmo.py index d68da04..87e687e 100644 --- a/pdelfin/train/molmo/modeling_molmo.py +++ b/pdelfin/train/molmo/modeling_molmo.py @@ -1311,7 +1311,7 @@ class OLMoPretrainedVisionBackbone(OLMoVisionBackbone): def __init__(self, config: FullMolmoConfig): super().__init__(config) v_cfg = self.config.vision_backbone - self.grad_checkpointing = False + self.grad_checkpointing = True 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" @@ -1688,7 +1688,7 @@ class Molmo(nn.Module): "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning ) 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 + torch.backends.cuda.enable_mem_efficient_sdp(True) # this is super slow so make sure torch won't use it wte = None if self.config.additional_vocab_size is not None: @@ -1741,6 +1741,8 @@ class Molmo(nn.Module): self.__num_fwd_flops: Optional[int] = None + self.gradient_checkpointing = False + def reset_parameters(self): if self.vision_backbone is not None: self.vision_backbone.reset_parameters() @@ -1951,7 +1953,11 @@ class Molmo(nn.Module): 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, layer_past=layer_past, use_cache=use_cache) + + if self.gradient_checkpointing and self.training: + x, cache = self._gradient_checkpointing_func(block, x, attention_bias=attention_bias, position_ids=position_ids, layer_past=layer_past, use_cache=use_cache) + else: + x, cache = block(x, attention_bias=attention_bias, position_ids=position_ids, layer_past=layer_past, use_cache=use_cache) if attn_key_values is not None: assert cache is not None @@ -2011,6 +2017,7 @@ class Molmo(nn.Module): class MolmoForCausalLM(PreTrainedModel): config_class = MolmoConfig + supports_gradient_checkpointing = True base_model_prefix = "model" _no_split_modules = ["MolmoBlock"] From 858b49656fdcde418462af0dd3f9d5ff6b24c77d Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Thu, 23 Jan 2025 15:32:04 -0800 Subject: [PATCH 3/7] Getting ready to train molmo 4096 context --- pdelfin/train/config/molmo-o-lora.yaml | 2 +- pdelfin/train/molmo/modeling_molmo.py | 2 +- pdelfin/train/train.py | 7 +++++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/pdelfin/train/config/molmo-o-lora.yaml b/pdelfin/train/config/molmo-o-lora.yaml index 763d01e..1fefbbb 100644 --- a/pdelfin/train/config/molmo-o-lora.yaml +++ b/pdelfin/train/config/molmo-o-lora.yaml @@ -45,7 +45,7 @@ hparams: batch_size: 1 eval_batch_size: 1 gradient_accumulation_steps: 4 - gradient_checkpointing: false + gradient_checkpointing: true clip_grad_norm: 1.0 learning_rate: 1e-4 max_steps: 10000 diff --git a/pdelfin/train/molmo/modeling_molmo.py b/pdelfin/train/molmo/modeling_molmo.py index 87e687e..606cbf8 100644 --- a/pdelfin/train/molmo/modeling_molmo.py +++ b/pdelfin/train/molmo/modeling_molmo.py @@ -1688,7 +1688,7 @@ class Molmo(nn.Module): "Embedding size is not a multiple of 128! This could hurt throughput performance.", UserWarning ) torch.backends.cuda.enable_flash_sdp(True) - torch.backends.cuda.enable_mem_efficient_sdp(True) # this is super slow so make sure torch won't use it + torch.backends.cuda.enable_mem_efficient_sdp(True) # jakep: I found that setting this to true in torch 2.5.1 greatly increased performance (6sec/it from 22sec/it) wte = None if self.config.additional_vocab_size is not None: diff --git a/pdelfin/train/train.py b/pdelfin/train/train.py index 08f97c5..fafe13b 100644 --- a/pdelfin/train/train.py +++ b/pdelfin/train/train.py @@ -122,7 +122,10 @@ def run_train(config: TrainConfig): _attn_implementation="flash_attention_2" if config.model.use_flash_attn else None ) else: - model_config = AutoConfig.from_pretrained(config.model.name_or_path, trust_remote_code=True) + from .molmo.config_molmo import MolmoConfig + from .molmo.modeling_molmo import MolmoForCausalLM + + model_config = MolmoConfig.from_pretrained(config.model.name_or_path, trust_remote_code=True) if model_config.max_position_embeddings < config.generate.max_length: logger.warning(f"ALERT, force adjusting model config max_position_embeddings upwards from {model_config.max_position_embeddings} to {config.generate.max_length}") @@ -131,7 +134,7 @@ def run_train(config: TrainConfig): if config.model.use_flash_attn: model_config.attention_type = "flash" - model = AutoModelForCausalLM.from_pretrained( + model = MolmoForCausalLM.from_pretrained( config.model.name_or_path, torch_dtype=torch.bfloat16, config=model_config, trust_remote_code=True From dabecd9ef0e858b843dddf93690d985252b510b2 Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Thu, 23 Jan 2025 23:39:56 +0000 Subject: [PATCH 4/7] More configs --- pdelfin/train/config/molmo-o-lora-8192.yaml | 88 +++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 pdelfin/train/config/molmo-o-lora-8192.yaml diff --git a/pdelfin/train/config/molmo-o-lora-8192.yaml b/pdelfin/train/config/molmo-o-lora-8192.yaml new file mode 100644 index 0000000..9222d28 --- /dev/null +++ b/pdelfin/train/config/molmo-o-lora-8192.yaml @@ -0,0 +1,88 @@ +model: + name_or_path: allenai/Molmo-7B-O-0924 + arch: causal + use_flash_attn: true + +wandb: + project: pdelfin + entity: ai2-llm + +generate: + max_length: 8192 + +train_data: + seed: 1337 + cache_location: /data/jakep/pdfdata/pdelfin_cache + sources: + - name: openai_batch_data_v5_1_train + response_glob_path: /data/jakep/pdfdata/openai_batch_data_v5_1_train_done/*.json + target_longest_image_dim: [1024] + target_anchor_text_len: [6000] + - name: openai_batch_data_v5_1_iabooks_train + response_glob_path: /data/jakep/pdfdata/openai_batch_data_v5_1_iabooks_train_done/*.json + target_longest_image_dim: [1024] + target_anchor_text_len: [6000] + +valid_data: + cache_location: /data/jakep/pdfdata/pdelfin_cache + metric_for_best_model: openai_batch_data_v5_1_eval_loss + sources: + # These tend to be small, so you can load from s3 it's no big deal + - name: openai_batch_data_v5_1_eval + response_glob_path: s3://ai2-oe-data/jakep/pdfdata/openai_batch_done_v5_1_eval/*.json + target_longest_image_dim: [1024] + target_anchor_text_len: [6000] + - name: openai_batch_data_v5_1_eval + response_glob_path: s3://ai2-oe-data/jakep/pdfdata/openai_batch_done_v5_1_eval/*.json + target_longest_image_dim: [1024] + target_anchor_text_len: [6000] + + + + +# Mostly pulled from https://github.com/QwenLM/Qwen2/blob/main/examples/sft/finetune.sh +hparams: + batch_size: 1 + eval_batch_size: 1 + gradient_accumulation_steps: 4 + gradient_checkpointing: true + clip_grad_norm: 1.0 + learning_rate: 1e-4 + max_steps: 10000 + pad_multiple_of: 16 + log_every_steps: 10 + eval_every_steps: 100 + optim: adamw_torch + lr_scheduler: cosine + weight_decay: 0.01 + warmup_ratio: 0.03 + +# From https://github.com/QwenLM/Qwen2/blob/main/examples/sft/finetune.py +lora: + rank: 32 + alpha: 32 + dropout: 0.05 + task_type: CAUSAL_LM + target_modules: + # attention layers in main transformer + - att_proj + - ff_proj + - attn_out + - ff_out + # vision transformer attention and FF + - attention.wq + - attention.wk + - attention.wv + - attention.wo + - feed_forward.w1 + - feed_forward.w2 + # vision image projector + - vision_backbone.image_projector.w1 + - vision_backbone.image_projector.w2 + - vision_backbone.image_projector.w3 + +save: + path: s3://ai2-oe-data/jakep/experiments/molmo-o-0924/v1/models/ + save_every_steps: 1000 + +max_workers: 10 \ No newline at end of file From d0eea81c00d19d0a236474e0ebcf1deadd04f637 Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Fri, 24 Jan 2025 16:27:42 +0000 Subject: [PATCH 5/7] Dealing with issue with molmo unused params --- pdelfin/train/config/molmo-o-lora-8192.yaml | 1 + pdelfin/train/config/molmo-o-lora.yaml | 1 + scripts/molmo-7b-lora-gantry.sh | 5 ++--- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pdelfin/train/config/molmo-o-lora-8192.yaml b/pdelfin/train/config/molmo-o-lora-8192.yaml index 9222d28..3b8e5c0 100644 --- a/pdelfin/train/config/molmo-o-lora-8192.yaml +++ b/pdelfin/train/config/molmo-o-lora-8192.yaml @@ -46,6 +46,7 @@ hparams: eval_batch_size: 1 gradient_accumulation_steps: 4 gradient_checkpointing: true + find_unused_parameters: true clip_grad_norm: 1.0 learning_rate: 1e-4 max_steps: 10000 diff --git a/pdelfin/train/config/molmo-o-lora.yaml b/pdelfin/train/config/molmo-o-lora.yaml index 1fefbbb..e6b9e70 100644 --- a/pdelfin/train/config/molmo-o-lora.yaml +++ b/pdelfin/train/config/molmo-o-lora.yaml @@ -46,6 +46,7 @@ hparams: eval_batch_size: 1 gradient_accumulation_steps: 4 gradient_checkpointing: true + find_unused_parameters: true clip_grad_norm: 1.0 learning_rate: 1e-4 max_steps: 10000 diff --git a/scripts/molmo-7b-lora-gantry.sh b/scripts/molmo-7b-lora-gantry.sh index 9fc21ef..db71e37 100755 --- a/scripts/molmo-7b-lora-gantry.sh +++ b/scripts/molmo-7b-lora-gantry.sh @@ -22,8 +22,8 @@ run_name=$(basename "$0" .sh) CLUSTER='jupiter' gantry run \ - --description "${run_name}"\ - --task-name "${run_name}"\ + --description "${run_name}-4096"\ + --task-name "${run_name}-4096"\ --allow-dirty \ --host-networking \ --workspace ai2/oe-data-model-based-cleanup \ @@ -32,7 +32,6 @@ gantry run \ --pip gantry-requirements.txt \ --priority high \ --gpus 8 \ - --preemptible \ --cluster "ai2/${CLUSTER}*" \ --budget ai2/oe-data \ --weka "oe-data-default:/data" \ From 5b429ad1002023f169b59b23664594a90cb4b941 Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Fri, 24 Jan 2025 23:15:35 +0000 Subject: [PATCH 6/7] Higher lr for molmo, fixed evals --- pdelfin/train/config/molmo-o-lora-8192.yaml | 8 +++----- scripts/molmo-7b-lora-gantry.sh | 6 +++--- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/pdelfin/train/config/molmo-o-lora-8192.yaml b/pdelfin/train/config/molmo-o-lora-8192.yaml index 3b8e5c0..6c9300c 100644 --- a/pdelfin/train/config/molmo-o-lora-8192.yaml +++ b/pdelfin/train/config/molmo-o-lora-8192.yaml @@ -32,14 +32,12 @@ valid_data: response_glob_path: s3://ai2-oe-data/jakep/pdfdata/openai_batch_done_v5_1_eval/*.json target_longest_image_dim: [1024] target_anchor_text_len: [6000] - - name: openai_batch_data_v5_1_eval - response_glob_path: s3://ai2-oe-data/jakep/pdfdata/openai_batch_done_v5_1_eval/*.json + - name: openai_batch_data_v5_1_iabooks_eval + response_glob_path: s3://ai2-oe-data/jakep/pdfdata/openai_batch_done_v5_1_iabooks_eval/*.json target_longest_image_dim: [1024] target_anchor_text_len: [6000] - - # Mostly pulled from https://github.com/QwenLM/Qwen2/blob/main/examples/sft/finetune.sh hparams: batch_size: 1 @@ -48,7 +46,7 @@ hparams: gradient_checkpointing: true find_unused_parameters: true clip_grad_norm: 1.0 - learning_rate: 1e-4 + learning_rate: 3e-4 max_steps: 10000 pad_multiple_of: 16 log_every_steps: 10 diff --git a/scripts/molmo-7b-lora-gantry.sh b/scripts/molmo-7b-lora-gantry.sh index db71e37..5072c62 100755 --- a/scripts/molmo-7b-lora-gantry.sh +++ b/scripts/molmo-7b-lora-gantry.sh @@ -10,7 +10,7 @@ then fi -EXTRA_ARGS="-c pdelfin/train/config/molmo-o-lora.yaml --num_proc 64 --save.path \"s3://ai2-oe-data/jakep/experiments/molmo-pdf/v1/models/\${BEAKER_USER_ID}\"" +EXTRA_ARGS="-c pdelfin/train/config/molmo-o-lora-8192.yaml --num_proc 64 --save.path \"s3://ai2-oe-data/jakep/experiments/molmo-pdf/v1/models/\${BEAKER_USER_ID}\"" run_name=$(basename "$0" .sh) @@ -22,8 +22,8 @@ run_name=$(basename "$0" .sh) CLUSTER='jupiter' gantry run \ - --description "${run_name}-4096"\ - --task-name "${run_name}-4096"\ + --description "${run_name}-8192"\ + --task-name "${run_name}-8192"\ --allow-dirty \ --host-networking \ --workspace ai2/oe-data-model-based-cleanup \ From ad88a82ee9da3ed1fa346eedab95d58556fb51f3 Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Mon, 27 Jan 2025 17:16:21 +0000 Subject: [PATCH 7/7] more elos --- pdelfin/eval/scoreelo.py | 27 +++++---------------------- 1 file changed, 5 insertions(+), 22 deletions(-) diff --git a/pdelfin/eval/scoreelo.py b/pdelfin/eval/scoreelo.py index 639c9f3..904a00b 100644 --- a/pdelfin/eval/scoreelo.py +++ b/pdelfin/eval/scoreelo.py @@ -279,26 +279,9 @@ def make_report(urls): if __name__ == "__main__": # Example usage - urls = [ - "https://jakep-tinyhost.s3.amazonaws.com/review_page_0-e09ebadf34a7.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=JEQpJxSaMIHuc9DFHyfHuxx0dEU%3D&Expires=1737654586", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_1-c2d267f97a73.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=KMiOTQiFEvgxU94ZrlJRFAgSQZA%3D&Expires=1737654587", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_10-b806c811fb67.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=NaoHNU2ZmEGrgMsxg2JHK%2Fv5zd0%3D&Expires=1737654587", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_11-19c1936b4372.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=BjkVydyKjzzH3uZiZ1GkWAk6cbk%3D&Expires=1737654588", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_12-cd41808a7974.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=jsk8TzJTKJwHi1Ru4%2Bw%2BiHZG638%3D&Expires=1737654589", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_13-8b055079b5eb.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=SE7kkobEBip44O8JY5axoMTV2Bs%3D&Expires=1737654590", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_14-1126e0da563c.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=jLSEWpDUzpmS8P9mNXbBoDYDOwU%3D&Expires=1737654590", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_15-05704e3d000d.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=TaCbyv2%2FDGCnCOgTzUvfEXdO%2Fmo%3D&Expires=1737654591", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_16-e57f795a89da.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=stqm1etAfDIpAQGNvZwe9c%2BYUbA%3D&Expires=1737654592", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_17-041a6d042764.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=rOTroBcSqCh3oM65bOJHEfaeal8%3D&Expires=1737654592", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_18-7a29697cee63.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=abmYM9KtzjicmdacRykPWXCdQr0%3D&Expires=1737654593", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_19-d32f14c067f8.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=iVg3nxrZXVpYybkLJIgOEJ3v37E%3D&Expires=1737654594", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_2-43c553548e69.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=IK27gl7b6NY05YNnnsimMVJc99I%3D&Expires=1737654595", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_3-fb42a458ecd5.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=d1qevJe8ZQONnu7zezYSJe3cbBw%3D&Expires=1737654595", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_4-76a50eed331a.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=qwZu2q1H4Y%2Bf3Kw7DNSYcTxwI7A%3D&Expires=1737654596", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_5-150b4d3583de.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=c%2FeqjnDSIRirgQviFWRLWVowKmA%3D&Expires=1737654597", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_6-6ca285526fd3.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=tkWDDuRinY77BLQCqumtlMiFJU8%3D&Expires=1737654598", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_7-01d711ee8bf7.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=eQtFo6CHJYHGu85wK0YG5khlE5U%3D&Expires=1737654598", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_8-0f36b852f274.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=weI3WB8vhjBYjk6t85DmyLdP97k%3D&Expires=1737654599", - "https://jakep-tinyhost.s3.amazonaws.com/review_page_9-115e33463fd2.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=b4CpkHprCUtZoL0u%2FFYzsu%2BB1yU%3D&Expires=1737654600", - ] + urls = ['https://jakep-tinyhost.s3.amazonaws.com/review_page_0-ff70abb8f517.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=NarEyyCfvusCh%2FHdB47VfHOnnBs%3D&Expires=1738359221', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_1-0800f9af46cf.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=ncTWAu5rSndBJJsU26HRYDaK6i8%3D&Expires=1738359222', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_10-f7081f6ca6f9.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=gYX8yjGyYshRqXGgdsX17%2Fdi9Ig%3D&Expires=1738359223', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_11-355dc69335bc.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=7%2Bc5qoa8Tbk06z0VcvJiIIVAz9M%3D&Expires=1738359224', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_12-95fce9bf0c18.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=fw4PBo0LnxikmLZ8xH%2BGD%2F%2BhXMU%3D&Expires=1738359225', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_13-f88f7d7482bf.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=yXkQp9oFDtroKgiO50EwpYdGLcA%3D&Expires=1738359226', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_14-8ac0b974bfd5.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=EgZTpj1%2FdzMBUgd%2BX4pVZ1Sp%2FrA%3D&Expires=1738359226', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_15-e3136188de5c.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=YKhAv4unNIlRcerQAaHN4kjc4qI%3D&Expires=1738359227', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_16-2c5abde50d49.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=Mj8%2BK5ISKzAYQFeYvmzTgCPcRwA%3D&Expires=1738359228', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_17-f13132a4cdcc.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=%2FHuzw2cjJ4oFm91UXojPnGzYi8Q%3D&Expires=1738359229', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_18-25070f2aa05e.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=ctd%2BUIM%2FxryJm%2FcwA%2BRZ%2FbRzBp8%3D&Expires=1738359230', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_19-d436ee434162.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=jVdFKobIoHlbTQ7zziG%2BXiIQ0Fo%3D&Expires=1738359230', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_2-a5ece743fd31.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=K8hIrjWtvo4SLVQrOB8TiXLgNJk%3D&Expires=1738359231', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_3-9ce03af05f51.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=T0fLGSH%2Bv%2F19veqbxnLxoSf7gVA%3D&Expires=1738359232', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_4-94eec18f8027.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=u2R1LundKpfnAUCcD%2BdGHA6uIR0%3D&Expires=1738359233', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_5-377d0a7d8f5a.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=5R38ZQAR9ew5x%2BRmMVQbTqbfVh0%3D&Expires=1738359234', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_6-537b22646a26.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=PLOELum1qzOXW8Cm5rfZphlFeMw%3D&Expires=1738359235', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_7-a4a7dcb08f20.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=DxPHukGXEpPrEPL6TF9QBKPE1Xg%3D&Expires=1738359236', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_8-48a71c829863.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=TjEINKj69HdmXsKY59k4f3PieeM%3D&Expires=1738359237', 'https://jakep-tinyhost.s3.amazonaws.com/review_page_9-8557438928c3.html?AWSAccessKeyId=AKIASHLPW4FEVZOPGK46&Signature=F7sQxw5A%2FDOcOaa%2FQSeqepH0PQc%3D&Expires=1738359238'] + # import tinyhost + + # print(tinyhost.tinyhost(urls)) + make_report(urls)