Add files using upload-large-folder tool

InternVL/classification/configs/intern_vit_6b_1k_224.yaml

@@ -0,0 +1,35 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenet-1k'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/configs/intern_vit_6b_1k_224_test_imagenet_a.yaml

@@ -0,0 +1,36 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenet_a'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenet-a'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/configs/intern_vit_6b_1k_224_test_imagenet_r.yaml

@@ -0,0 +1,36 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenet_r'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenet-r'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/configs/intern_vit_6b_1k_224_test_imagenet_real.yaml

@@ -0,0 +1,36 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenet-real'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenet-1k'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/configs/intern_vit_6b_1k_224_test_imagenet_sketch.yaml

@@ -0,0 +1,36 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenet_sketch'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenet-sketch'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/configs/intern_vit_6b_1k_224_test_imagenetv2.yaml

@@ -0,0 +1,36 @@
+DATA:
+  IMG_ON_MEMORY: False
+  BATCH_SIZE: 128
+  DATASET: 'imagenetv2'
+  TRANSFORM: 'build_transform_for_linear_probe'
+  DATA_PATH: './data/imagenetv2'
+MODEL:
+  TYPE: intern_vit_6b
+  DROP_PATH_RATE: 0.0
+  INTERN_VIT_6B:
+    FREEZE_VIT: True
+    PATCH_SIZE: 14
+    PRETRAIN_SIZE: 224
+    QKV_BIAS: False
+    EMBED_DIM: 3200
+    NUM_HEADS: 25
+    MLP_RATIO: 4
+    INIT_VALUES: 0.1
+    QK_NORMALIZATION: True
+    DEPTH: 48
+    USE_FLASH_ATTN: True
+    PRETRAINED: "./pretrained/intern_vit_6b_224px.pth"
+    CLS_TARGET: 'cls_patch_concat'
+TRAIN:
+  EMA:
+    ENABLE: False
+    DECAY: 0.998
+  EPOCHS: 10
+  WARMUP_EPOCHS: 1
+  WEIGHT_DECAY: 0.0
+  BASE_LR: 0.1 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: false
+  OPTIMIZER:
+    NAME: 'sgd'

InternVL/classification/dataset/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .build import build_loader, build_loader2

InternVL/classification/dataset/build.py

@@ -0,0 +1,332 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from timm.data import Mixup, create_transform
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+
+from .cached_image_folder import ImageCephDataset
+from .samplers import NodeDistributedSampler, SubsetRandomSampler
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    def _pil_interp(method):
+        if method == 'bicubic':
+            return InterpolationMode.BICUBIC
+        elif method == 'lanczos':
+            return InterpolationMode.LANCZOS
+        elif method == 'hamming':
+            return InterpolationMode.HAMMING
+        else:
+            return InterpolationMode.BILINEAR
+except:
+    from timm.data.transforms import _pil_interp
+
+
+class TTA(torch.nn.Module):
+
+    def __init__(self, size, scales=[1.0, 1.05, 1.1]):
+        super().__init__()
+        self.size = size
+        self.scales = scales
+
+    def forward(self, img):
+        out = []
+        cc = transforms.CenterCrop(self.size)
+        for scale in self.scales:
+            size_ = int(scale * self.size)
+            rs = transforms.Resize(size_, interpolation=_pil_interp('bicubic'))
+            img_ = rs(img)
+            img_ = cc(img_)
+            out.append(img_)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(size={self.size}, scale={self.scales})'
+
+
+def build_loader(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train', config=config)
+    config.freeze()
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build train dataset')
+
+    dataset_val, _ = build_dataset('val', config=config)
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build val dataset')
+
+    dataset_test, _ = build_dataset('test', config=config)
+    print(f'local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}'
+          'successfully build test dataset')
+
+    num_tasks = dist.get_world_size()
+    global_rank = dist.get_rank()
+
+    if dataset_train is not None:
+        if config.DATA.IMG_ON_MEMORY:
+            sampler_train = NodeDistributedSampler(dataset_train)
+        else:
+            if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part':
+                indices = np.arange(dist.get_rank(), len(dataset_train), dist.get_world_size())
+                sampler_train = SubsetRandomSampler(indices)
+            else:
+                sampler_train = torch.utils.data.DistributedSampler(
+                    dataset_train,
+                    num_replicas=num_tasks,
+                    rank=global_rank,
+                    shuffle=True)
+
+    if dataset_val is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+        else:
+            sampler_val = torch.utils.data.distributed.DistributedSampler(dataset_val, shuffle=False)
+
+    if dataset_test is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+        else:
+            sampler_test = torch.utils.data.distributed.DistributedSampler(dataset_test, shuffle=False)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler_train,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        sampler=sampler_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        sampler=sampler_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+           data_loader_val, data_loader_test, mixup_fn
+
+
+def build_loader2(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train', config=config)
+    config.freeze()
+    dataset_val, _ = build_dataset('val', config=config)
+    dataset_test, _ = build_dataset('test', config=config)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        shuffle=True,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+           data_loader_val, data_loader_test, mixup_fn
+
+
+def build_dataset(split, config):
+    if config.DATA.TRANSFORM == 'build_transform':
+        transform = build_transform(split == 'train', config)
+    elif config.DATA.TRANSFORM == 'build_transform_for_linear_probe':
+        transform = build_transform_for_linear_probe(split == 'train', config)
+    else:
+        raise NotImplementedError
+    print(split, transform)
+    dataset = None
+    nb_classes = None
+    prefix = split
+    if config.DATA.DATASET == 'imagenet' or config.DATA.DATASET == 'imagenet-real':
+        if prefix == 'train' and not config.EVAL_MODE:
+            root = os.path.join(config.DATA.DATA_PATH, 'train')
+            dataset = ImageCephDataset(root, 'train',
+                                       transform=transform,
+                                       on_memory=config.DATA.IMG_ON_MEMORY)
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+        nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet22K':
+        if prefix == 'train':
+            if not config.EVAL_MODE:
+                root = config.DATA.DATA_PATH
+                dataset = ImageCephDataset(root, 'train',
+                                           transform=transform,
+                                           on_memory=config.DATA.IMG_ON_MEMORY)
+            nb_classes = 21841
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenetv2':
+        from .imagenetv2 import ImageNetV2Dataset
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageNetV2Dataset(variant='matched-frequency',
+                                        transform=transform,
+                                        location=config.DATA.DATA_PATH)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet_sketch':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet_a':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000  # actual number of classes is 200
+    elif config.DATA.DATASET == 'imagenet_r':
+        if prefix == 'train' and not config.EVAL_MODE:
+            print(f'Only test split available for {config.DATA.DATASET}')
+        else:
+            dataset = ImageFolder(root=config.DATA.DATA_PATH, transform=transform)
+            nb_classes = 1000  # actual number of classes is 200
+    else:
+        raise NotImplementedError(
+            f'build_dataset does support {config.DATA.DATASET}')
+
+    return dataset, nb_classes
+
+
+def build_transform_for_linear_probe(is_train, config):
+    # linear probe: weak augmentation
+    if is_train:
+        transform = transforms.Compose([
+            transforms.RandomResizedCrop(
+                config.DATA.IMG_SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=config.AUG.MEAN, std=config.AUG.STD)
+        ])
+    else:
+        transform = transforms.Compose([
+            transforms.Resize(
+                config.DATA.IMG_SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.CenterCrop(config.DATA.IMG_SIZE),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=config.AUG.MEAN, std=config.AUG.STD)
+        ])
+    return transform
+
+
+def build_transform(is_train, config):
+    resize_im = config.DATA.IMG_SIZE > 32
+    if is_train:
+        # this should always dispatch to transforms_imagenet_train
+        transform = create_transform(
+            input_size=config.DATA.IMG_SIZE,
+            is_training=True,
+            color_jitter=config.AUG.COLOR_JITTER
+            if config.AUG.COLOR_JITTER > 0 else None,
+            auto_augment=config.AUG.AUTO_AUGMENT
+            if config.AUG.AUTO_AUGMENT != 'none' else None,
+            re_prob=config.AUG.REPROB,
+            re_mode=config.AUG.REMODE,
+            re_count=config.AUG.RECOUNT,
+            interpolation=config.DATA.INTERPOLATION,
+        )
+        if not resize_im:
+            # replace RandomResizedCropAndInterpolation with
+            # RandomCrop
+            transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4)
+
+        return transform
+
+    t = []
+    if resize_im:
+        if config.TEST.CROP:
+            size = int(1.0 * config.DATA.IMG_SIZE)
+            t.append(
+                transforms.Resize(size, interpolation=_pil_interp(config.DATA.INTERPOLATION)),
+                # to maintain same ratio w.r.t. 224 images
+            )
+            t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
+        elif config.AUG.RANDOM_RESIZED_CROP:
+            t.append(
+                transforms.RandomResizedCrop(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+        else:
+            t.append(
+                transforms.Resize(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+    t.append(transforms.ToTensor())
+    t.append(transforms.Normalize(config.AUG.MEAN, config.AUG.STD))
+
+    return transforms.Compose(t)

InternVL/classification/dataset/cached_image_folder.py

@@ -0,0 +1,543 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import io
+import json
+import logging
+import math
+import os
+import os.path as osp
+import re
+import time
+from abc import abstractmethod
+
+import mmcv
+import torch
+import torch.distributed as dist
+import torch.utils.data as data
+from mmcv.fileio import FileClient
+from PIL import Image
+from tqdm import tqdm, trange
+
+from .zipreader import ZipReader, is_zip_path
+
+_logger = logging.getLogger(__name__)
+
+_ERROR_RETRY = 50
+
+
+def has_file_allowed_extension(filename, extensions):
+    """Checks if a file is an allowed extension.
+
+    Args:
+        filename (string): path to a file
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in extensions)
+
+
+def find_classes(dir):
+    classes = [
+        d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
+    ]
+    classes.sort()
+    class_to_idx = {classes[i]: i for i in range(len(classes))}
+    return classes, class_to_idx
+
+
+def make_dataset(dir, class_to_idx, extensions):
+    images = []
+    dir = os.path.expanduser(dir)
+    for target in sorted(os.listdir(dir)):
+        d = os.path.join(dir, target)
+        if not os.path.isdir(d):
+            continue
+        for root, _, fnames in sorted(os.walk(d)):
+            for fname in sorted(fnames):
+                if has_file_allowed_extension(fname, extensions):
+                    path = os.path.join(root, fname)
+                    item = (path, class_to_idx[target])
+                    images.append(item)
+
+    return images
+
+
+def make_dataset_with_ann(ann_file, img_prefix, extensions):
+    images = []
+    with open(ann_file, 'r') as f:
+        contents = f.readlines()
+        for line_str in contents:
+            path_contents = [c for c in line_str.split('\t')]
+            im_file_name = path_contents[0]
+            class_index = int(path_contents[1])
+            assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
+            item = (os.path.join(img_prefix, im_file_name), class_index)
+            images.append(item)
+
+    return images
+
+
+class DatasetFolder(data.Dataset):
+    """A generic data loader where the samples are arranged in this way: ::
+
+    root/class_x/xxx.ext
+        root/class_x/xxy.ext
+        root/class_x/xxz.ext
+        root/class_y/123.ext
+        root/class_y/nsdf3.ext
+        root/class_y/asd932_.ext
+    Args:
+        root (string): Root directory path.
+        loader (callable): A function to load a sample given its path.
+        extensions (list[string]): A list of allowed extensions.
+        transform (callable, optional): A function/transform that takes in
+            a sample and returns a transformed version.
+            E.g, ``transforms.RandomCrop`` for images.
+        target_transform (callable, optional): A function/transform that takes
+            in the target and transforms it.
+     Attributes:
+        samples (list): List of (sample path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 loader,
+                 extensions,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 cache_mode='no'):
+        # image folder mode
+        if ann_file == '':
+            _, class_to_idx = find_classes(root)
+            samples = make_dataset(root, class_to_idx, extensions)
+        # zip mode
+        else:
+            samples = make_dataset_with_ann(os.path.join(root, ann_file),
+                                            os.path.join(root, img_prefix),
+                                            extensions)
+
+        if len(samples) == 0:
+            raise (RuntimeError('Found 0 files in subfolders of: ' + root +
+                                '\n' + 'Supported extensions are: ' +
+                                ','.join(extensions)))
+
+        self.root = root
+        self.loader = loader
+        self.extensions = extensions
+
+        self.samples = samples
+        self.labels = [y_1k for _, y_1k in samples]
+        self.classes = list(set(self.labels))
+
+        self.transform = transform
+        self.target_transform = target_transform
+
+        self.cache_mode = cache_mode
+        if self.cache_mode != 'no':
+            self.init_cache()
+
+    def init_cache(self):
+        assert self.cache_mode in ['part', 'full']
+        n_sample = len(self.samples)
+        global_rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+        samples_bytes = [None for _ in range(n_sample)]
+        start_time = time.time()
+        for index in range(n_sample):
+            if index % (n_sample // 10) == 0:
+                t = time.time() - start_time
+                print(
+                    f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
+                )
+                start_time = time.time()
+            path, target = self.samples[index]
+            if self.cache_mode == 'full':
+                samples_bytes[index] = (ZipReader.read(path), target)
+            elif self.cache_mode == 'part' and index % world_size == global_rank:
+                samples_bytes[index] = (ZipReader.read(path), target)
+            else:
+                samples_bytes[index] = (path, target)
+        self.samples = samples_bytes
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (sample, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        sample = self.loader(path)
+        if self.transform is not None:
+            sample = self.transform(sample)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return sample, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __repr__(self):
+        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
+        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
+        fmt_str += '    Root Location: {}\n'.format(self.root)
+        tmp = '    Transforms (if any): '
+        fmt_str += '{0}{1}\n'.format(
+            tmp,
+            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
+        tmp = '    Target Transforms (if any): '
+        fmt_str += '{0}{1}'.format(
+            tmp,
+            self.target_transform.__repr__().replace('\n',
+                                                     '\n' + ' ' * len(tmp)))
+
+        return fmt_str
+
+
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
+
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    if isinstance(path, bytes):
+        img = Image.open(io.BytesIO(path))
+    elif is_zip_path(path):
+        data = ZipReader.read(path)
+        img = Image.open(io.BytesIO(data))
+    else:
+        with open(path, 'rb') as f:
+            img = Image.open(f)
+            return img.convert('RGB')
+
+    return img.convert('RGB')
+
+
+def accimage_loader(path):
+    import accimage
+    try:
+        return accimage.Image(path)
+    except IOError:
+        # Potentially a decoding problem, fall back to PIL.Image
+        return pil_loader(path)
+
+
+def default_img_loader(path):
+    from torchvision import get_image_backend
+    if get_image_backend() == 'accimage':
+        return accimage_loader(path)
+    else:
+        return pil_loader(path)
+
+
+class CachedImageFolder(DatasetFolder):
+    """A generic data loader where the images are arranged in this way: ::
+
+    root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/xxz.png
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/asd932_.png
+    Args:
+        root (string): Root directory path.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+     Attributes:
+        imgs (list): List of (image path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 loader=default_img_loader,
+                 cache_mode='no'):
+        super(CachedImageFolder,
+              self).__init__(root,
+                             loader,
+                             IMG_EXTENSIONS,
+                             ann_file=ann_file,
+                             img_prefix=img_prefix,
+                             transform=transform,
+                             target_transform=target_transform,
+                             cache_mode=cache_mode)
+        self.imgs = self.samples
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        image = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(image)
+        else:
+            img = image
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+
+class ImageCephDataset(data.Dataset):
+
+    def __init__(self,
+                 root,
+                 split,
+                 parser=None,
+                 transform=None,
+                 target_transform=None,
+                 on_memory=False):
+        if '22k' in root:
+            # Imagenet 22k
+            annotation_root = 'meta_data/'
+        else:
+            # Imagenet
+            annotation_root = 'meta_data/'
+        if parser is None or isinstance(parser, str):
+            parser = ParserCephImage(root=root,
+                                     split=split,
+                                     annotation_root=annotation_root,
+                                     on_memory=on_memory)
+        self.parser = parser
+        self.transform = transform
+        self.target_transform = target_transform
+        self._consecutive_errors = 0
+
+    def __getitem__(self, index):
+        img, target = self.parser[index]
+        self._consecutive_errors = 0
+        if self.transform is not None:
+            img = self.transform(img)
+        if target is None:
+            target = -1
+        elif self.target_transform is not None:
+            target = self.target_transform(target)
+        return img, target
+
+    def __len__(self):
+        return len(self.parser)
+
+    def filename(self, index, basename=False, absolute=False):
+        return self.parser.filename(index, basename, absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return self.parser.filenames(basename, absolute)
+
+
+class Parser:
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def _filename(self, index, basename=False, absolute=False):
+        pass
+
+    def filename(self, index, basename=False, absolute=False):
+        return self._filename(index, basename=basename, absolute=absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return [
+            self._filename(index, basename=basename, absolute=absolute)
+            for index in range(len(self))
+        ]
+
+
+class ParserCephImage(Parser):
+
+    def __init__(self,
+                 root,
+                 split,
+                 annotation_root,
+                 on_memory=False,
+                 **kwargs):
+        super().__init__()
+
+        self.file_client = None
+        self.kwargs = kwargs
+
+        self.root = root  # dataset:s3://imagenet22k
+        if '22k' in root:
+            self.io_backend = 'petrel'
+            with open(osp.join(annotation_root, '22k_class_to_idx.json'),
+                      'r') as f:
+                self.class_to_idx = json.loads(f.read())
+            with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        else:
+            self.io_backend = 'disk'
+            self.class_to_idx = None
+            with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        local_rank = None
+        local_size = None
+        self._consecutive_errors = 0
+        self.on_memory = on_memory
+        if on_memory:
+            self.holder = {}
+            if local_rank is None:
+                local_rank = int(os.environ.get('LOCAL_RANK', 0))
+            if local_size is None:
+                local_size = int(os.environ.get('LOCAL_SIZE', 1))
+            self.local_rank = local_rank
+            self.local_size = local_size
+            self.rank = int(os.environ['RANK'])
+            self.world_size = int(os.environ['WORLD_SIZE'])
+            self.num_replicas = int(os.environ['WORLD_SIZE'])
+            self.num_parts = local_size
+            self.num_samples = int(
+                math.ceil(len(self.samples) * 1.0 / self.num_replicas))
+            self.total_size = self.num_samples * self.num_replicas
+            self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+            self.load_onto_memory_v2()
+
+    def load_onto_memory(self):
+        print('Loading images onto memory...', self.local_rank,
+              self.local_size)
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in trange(len(self.samples)):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+            self.holder[path] = img_bytes
+
+        print('Loading complete!')
+
+    def load_onto_memory_v2(self):
+        # print("Loading images onto memory...", self.local_rank, self.local_size)
+        t = torch.Generator()
+        t.manual_seed(0)
+        indices = torch.randperm(len(self.samples), generator=t).tolist()
+        # indices = range(len(self.samples))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+        assert len(indices) == self.num_samples
+
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in tqdm(indices):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+
+            self.holder[path] = img_bytes
+
+        print('Loading complete!')
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+
+        filepath, target = self.samples[index].split(' ')
+        filepath = osp.join(self.root, filepath)
+
+        try:
+            if self.on_memory:
+                img_bytes = self.holder[filepath]
+            else:
+                # pass
+                img_bytes = self.file_client.get(filepath)
+            img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
+        except Exception as e:
+            _logger.warning(
+                f'Skipped sample (index {index}, file {filepath}). {str(e)}')
+            self._consecutive_errors += 1
+            if self._consecutive_errors < _ERROR_RETRY:
+                return self.__getitem__((index + 1) % len(self))
+            else:
+                raise e
+        self._consecutive_errors = 0
+
+        img = Image.fromarray(img)
+        try:
+            if self.class_to_idx is not None:
+                target = self.class_to_idx[target]
+            else:
+                target = int(target)
+        except:
+            print(filepath, target)
+            exit()
+
+        return img, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _filename(self, index, basename=False, absolute=False):
+        filename, _ = self.samples[index].split(' ')
+        filename = osp.join(self.root, filename)
+
+        return filename
+
+
+def get_temporal_info(date, miss_hour=False):
+    try:
+        if date:
+            if miss_hour:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*)', re.I)
+            else:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*) (\d*):(\d*):(\d*)',
+                                     re.I)
+            m = pattern.match(date.strip())
+
+            if m:
+                year = int(m.group(1))
+                month = int(m.group(2))
+                day = int(m.group(3))
+                x_month = math.sin(2 * math.pi * month / 12)
+                y_month = math.cos(2 * math.pi * month / 12)
+                if miss_hour:
+                    x_hour = 0
+                    y_hour = 0
+                else:
+                    hour = int(m.group(4))
+                    x_hour = math.sin(2 * math.pi * hour / 24)
+                    y_hour = math.cos(2 * math.pi * hour / 24)
+                return [x_month, y_month, x_hour, y_hour]
+            else:
+                return [0, 0, 0, 0]
+        else:
+            return [0, 0, 0, 0]
+    except:
+        return [0, 0, 0, 0]
+
+
+def get_spatial_info(latitude, longitude):
+    if latitude and longitude:
+        latitude = math.radians(latitude)
+        longitude = math.radians(longitude)
+        x = math.cos(latitude) * math.cos(longitude)
+        y = math.cos(latitude) * math.sin(longitude)
+        z = math.sin(latitude)
+        return [x, y, z]
+    else:
+        return [0, 0, 0]

InternVL/classification/dataset/imagenet_a_r_indices.py

@@ -0,0 +1,295 @@
+"""Code from https://github.com/baaivision/EVA/blob/master/EVA-02/asuka/imagenet_a_r_indices.py
+Thanks to the authors of EVA."""
+
+all_wnids = [
+    'n01440764', 'n01443537', 'n01484850', 'n01491361', 'n01494475',
+    'n01496331', 'n01498041', 'n01514668', 'n01514859', 'n01518878',
+    'n01530575', 'n01531178', 'n01532829', 'n01534433', 'n01537544',
+    'n01558993', 'n01560419', 'n01580077', 'n01582220', 'n01592084',
+    'n01601694', 'n01608432', 'n01614925', 'n01616318', 'n01622779',
+    'n01629819', 'n01630670', 'n01631663', 'n01632458', 'n01632777',
+    'n01641577', 'n01644373', 'n01644900', 'n01664065', 'n01665541',
+    'n01667114', 'n01667778', 'n01669191', 'n01675722', 'n01677366',
+    'n01682714', 'n01685808', 'n01687978', 'n01688243', 'n01689811',
+    'n01692333', 'n01693334', 'n01694178', 'n01695060', 'n01697457',
+    'n01698640', 'n01704323', 'n01728572', 'n01728920', 'n01729322',
+    'n01729977', 'n01734418', 'n01735189', 'n01737021', 'n01739381',
+    'n01740131', 'n01742172', 'n01744401', 'n01748264', 'n01749939',
+    'n01751748', 'n01753488', 'n01755581', 'n01756291', 'n01768244',
+    'n01770081', 'n01770393', 'n01773157', 'n01773549', 'n01773797',
+    'n01774384', 'n01774750', 'n01775062', 'n01776313', 'n01784675',
+    'n01795545', 'n01796340', 'n01797886', 'n01798484', 'n01806143',
+    'n01806567', 'n01807496', 'n01817953', 'n01818515', 'n01819313',
+    'n01820546', 'n01824575', 'n01828970', 'n01829413', 'n01833805',
+    'n01843065', 'n01843383', 'n01847000', 'n01855032', 'n01855672',
+    'n01860187', 'n01871265', 'n01872401', 'n01873310', 'n01877812',
+    'n01882714', 'n01883070', 'n01910747', 'n01914609', 'n01917289',
+    'n01924916', 'n01930112', 'n01943899', 'n01944390', 'n01945685',
+    'n01950731', 'n01955084', 'n01968897', 'n01978287', 'n01978455',
+    'n01980166', 'n01981276', 'n01983481', 'n01984695', 'n01985128',
+    'n01986214', 'n01990800', 'n02002556', 'n02002724', 'n02006656',
+    'n02007558', 'n02009229', 'n02009912', 'n02011460', 'n02012849',
+    'n02013706', 'n02017213', 'n02018207', 'n02018795', 'n02025239',
+    'n02027492', 'n02028035', 'n02033041', 'n02037110', 'n02051845',
+    'n02056570', 'n02058221', 'n02066245', 'n02071294', 'n02074367',
+    'n02077923', 'n02085620', 'n02085782', 'n02085936', 'n02086079',
+    'n02086240', 'n02086646', 'n02086910', 'n02087046', 'n02087394',
+    'n02088094', 'n02088238', 'n02088364', 'n02088466', 'n02088632',
+    'n02089078', 'n02089867', 'n02089973', 'n02090379', 'n02090622',
+    'n02090721', 'n02091032', 'n02091134', 'n02091244', 'n02091467',
+    'n02091635', 'n02091831', 'n02092002', 'n02092339', 'n02093256',
+    'n02093428', 'n02093647', 'n02093754', 'n02093859', 'n02093991',
+    'n02094114', 'n02094258', 'n02094433', 'n02095314', 'n02095570',
+    'n02095889', 'n02096051', 'n02096177', 'n02096294', 'n02096437',
+    'n02096585', 'n02097047', 'n02097130', 'n02097209', 'n02097298',
+    'n02097474', 'n02097658', 'n02098105', 'n02098286', 'n02098413',
+    'n02099267', 'n02099429', 'n02099601', 'n02099712', 'n02099849',
+    'n02100236', 'n02100583', 'n02100735', 'n02100877', 'n02101006',
+    'n02101388', 'n02101556', 'n02102040', 'n02102177', 'n02102318',
+    'n02102480', 'n02102973', 'n02104029', 'n02104365', 'n02105056',
+    'n02105162', 'n02105251', 'n02105412', 'n02105505', 'n02105641',
+    'n02105855', 'n02106030', 'n02106166', 'n02106382', 'n02106550',
+    'n02106662', 'n02107142', 'n02107312', 'n02107574', 'n02107683',
+    'n02107908', 'n02108000', 'n02108089', 'n02108422', 'n02108551',
+    'n02108915', 'n02109047', 'n02109525', 'n02109961', 'n02110063',
+    'n02110185', 'n02110341', 'n02110627', 'n02110806', 'n02110958',
+    'n02111129', 'n02111277', 'n02111500', 'n02111889', 'n02112018',
+    'n02112137', 'n02112350', 'n02112706', 'n02113023', 'n02113186',
+    'n02113624', 'n02113712', 'n02113799', 'n02113978', 'n02114367',
+    'n02114548', 'n02114712', 'n02114855', 'n02115641', 'n02115913',
+    'n02116738', 'n02117135', 'n02119022', 'n02119789', 'n02120079',
+    'n02120505', 'n02123045', 'n02123159', 'n02123394', 'n02123597',
+    'n02124075', 'n02125311', 'n02127052', 'n02128385', 'n02128757',
+    'n02128925', 'n02129165', 'n02129604', 'n02130308', 'n02132136',
+    'n02133161', 'n02134084', 'n02134418', 'n02137549', 'n02138441',
+    'n02165105', 'n02165456', 'n02167151', 'n02168699', 'n02169497',
+    'n02172182', 'n02174001', 'n02177972', 'n02190166', 'n02206856',
+    'n02219486', 'n02226429', 'n02229544', 'n02231487', 'n02233338',
+    'n02236044', 'n02256656', 'n02259212', 'n02264363', 'n02268443',
+    'n02268853', 'n02276258', 'n02277742', 'n02279972', 'n02280649',
+    'n02281406', 'n02281787', 'n02317335', 'n02319095', 'n02321529',
+    'n02325366', 'n02326432', 'n02328150', 'n02342885', 'n02346627',
+    'n02356798', 'n02361337', 'n02363005', 'n02364673', 'n02389026',
+    'n02391049', 'n02395406', 'n02396427', 'n02397096', 'n02398521',
+    'n02403003', 'n02408429', 'n02410509', 'n02412080', 'n02415577',
+    'n02417914', 'n02422106', 'n02422699', 'n02423022', 'n02437312',
+    'n02437616', 'n02441942', 'n02442845', 'n02443114', 'n02443484',
+    'n02444819', 'n02445715', 'n02447366', 'n02454379', 'n02457408',
+    'n02480495', 'n02480855', 'n02481823', 'n02483362', 'n02483708',
+    'n02484975', 'n02486261', 'n02486410', 'n02487347', 'n02488291',
+    'n02488702', 'n02489166', 'n02490219', 'n02492035', 'n02492660',
+    'n02493509', 'n02493793', 'n02494079', 'n02497673', 'n02500267',
+    'n02504013', 'n02504458', 'n02509815', 'n02510455', 'n02514041',
+    'n02526121', 'n02536864', 'n02606052', 'n02607072', 'n02640242',
+    'n02641379', 'n02643566', 'n02655020', 'n02666196', 'n02667093',
+    'n02669723', 'n02672831', 'n02676566', 'n02687172', 'n02690373',
+    'n02692877', 'n02699494', 'n02701002', 'n02704792', 'n02708093',
+    'n02727426', 'n02730930', 'n02747177', 'n02749479', 'n02769748',
+    'n02776631', 'n02777292', 'n02782093', 'n02783161', 'n02786058',
+    'n02787622', 'n02788148', 'n02790996', 'n02791124', 'n02791270',
+    'n02793495', 'n02794156', 'n02795169', 'n02797295', 'n02799071',
+    'n02802426', 'n02804414', 'n02804610', 'n02807133', 'n02808304',
+    'n02808440', 'n02814533', 'n02814860', 'n02815834', 'n02817516',
+    'n02823428', 'n02823750', 'n02825657', 'n02834397', 'n02835271',
+    'n02837789', 'n02840245', 'n02841315', 'n02843684', 'n02859443',
+    'n02860847', 'n02865351', 'n02869837', 'n02870880', 'n02871525',
+    'n02877765', 'n02879718', 'n02883205', 'n02892201', 'n02892767',
+    'n02894605', 'n02895154', 'n02906734', 'n02909870', 'n02910353',
+    'n02916936', 'n02917067', 'n02927161', 'n02930766', 'n02939185',
+    'n02948072', 'n02950826', 'n02951358', 'n02951585', 'n02963159',
+    'n02965783', 'n02966193', 'n02966687', 'n02971356', 'n02974003',
+    'n02977058', 'n02978881', 'n02979186', 'n02980441', 'n02981792',
+    'n02988304', 'n02992211', 'n02992529', 'n02999410', 'n03000134',
+    'n03000247', 'n03000684', 'n03014705', 'n03016953', 'n03017168',
+    'n03018349', 'n03026506', 'n03028079', 'n03032252', 'n03041632',
+    'n03042490', 'n03045698', 'n03047690', 'n03062245', 'n03063599',
+    'n03063689', 'n03065424', 'n03075370', 'n03085013', 'n03089624',
+    'n03095699', 'n03100240', 'n03109150', 'n03110669', 'n03124043',
+    'n03124170', 'n03125729', 'n03126707', 'n03127747', 'n03127925',
+    'n03131574', 'n03133878', 'n03134739', 'n03141823', 'n03146219',
+    'n03160309', 'n03179701', 'n03180011', 'n03187595', 'n03188531',
+    'n03196217', 'n03197337', 'n03201208', 'n03207743', 'n03207941',
+    'n03208938', 'n03216828', 'n03218198', 'n03220513', 'n03223299',
+    'n03240683', 'n03249569', 'n03250847', 'n03255030', 'n03259280',
+    'n03271574', 'n03272010', 'n03272562', 'n03290653', 'n03291819',
+    'n03297495', 'n03314780', 'n03325584', 'n03337140', 'n03344393',
+    'n03345487', 'n03347037', 'n03355925', 'n03372029', 'n03376595',
+    'n03379051', 'n03384352', 'n03388043', 'n03388183', 'n03388549',
+    'n03393912', 'n03394916', 'n03400231', 'n03404251', 'n03417042',
+    'n03424325', 'n03425413', 'n03443371', 'n03444034', 'n03445777',
+    'n03445924', 'n03447447', 'n03447721', 'n03450230', 'n03452741',
+    'n03457902', 'n03459775', 'n03461385', 'n03467068', 'n03476684',
+    'n03476991', 'n03478589', 'n03481172', 'n03482405', 'n03483316',
+    'n03485407', 'n03485794', 'n03492542', 'n03494278', 'n03495258',
+    'n03496892', 'n03498962', 'n03527444', 'n03529860', 'n03530642',
+    'n03532672', 'n03534580', 'n03535780', 'n03538406', 'n03544143',
+    'n03584254', 'n03584829', 'n03590841', 'n03594734', 'n03594945',
+    'n03595614', 'n03598930', 'n03599486', 'n03602883', 'n03617480',
+    'n03623198', 'n03627232', 'n03630383', 'n03633091', 'n03637318',
+    'n03642806', 'n03649909', 'n03657121', 'n03658185', 'n03661043',
+    'n03662601', 'n03666591', 'n03670208', 'n03673027', 'n03676483',
+    'n03680355', 'n03690938', 'n03691459', 'n03692522', 'n03697007',
+    'n03706229', 'n03709823', 'n03710193', 'n03710637', 'n03710721',
+    'n03717622', 'n03720891', 'n03721384', 'n03724870', 'n03729826',
+    'n03733131', 'n03733281', 'n03733805', 'n03742115', 'n03743016',
+    'n03759954', 'n03761084', 'n03763968', 'n03764736', 'n03769881',
+    'n03770439', 'n03770679', 'n03773504', 'n03775071', 'n03775546',
+    'n03776460', 'n03777568', 'n03777754', 'n03781244', 'n03782006',
+    'n03785016', 'n03786901', 'n03787032', 'n03788195', 'n03788365',
+    'n03791053', 'n03792782', 'n03792972', 'n03793489', 'n03794056',
+    'n03796401', 'n03803284', 'n03804744', 'n03814639', 'n03814906',
+    'n03825788', 'n03832673', 'n03837869', 'n03838899', 'n03840681',
+    'n03841143', 'n03843555', 'n03854065', 'n03857828', 'n03866082',
+    'n03868242', 'n03868863', 'n03871628', 'n03873416', 'n03874293',
+    'n03874599', 'n03876231', 'n03877472', 'n03877845', 'n03884397',
+    'n03887697', 'n03888257', 'n03888605', 'n03891251', 'n03891332',
+    'n03895866', 'n03899768', 'n03902125', 'n03903868', 'n03908618',
+    'n03908714', 'n03916031', 'n03920288', 'n03924679', 'n03929660',
+    'n03929855', 'n03930313', 'n03930630', 'n03933933', 'n03935335',
+    'n03937543', 'n03938244', 'n03942813', 'n03944341', 'n03947888',
+    'n03950228', 'n03954731', 'n03956157', 'n03958227', 'n03961711',
+    'n03967562', 'n03970156', 'n03976467', 'n03976657', 'n03977966',
+    'n03980874', 'n03982430', 'n03983396', 'n03991062', 'n03992509',
+    'n03995372', 'n03998194', 'n04004767', 'n04005630', 'n04008634',
+    'n04009552', 'n04019541', 'n04023962', 'n04026417', 'n04033901',
+    'n04033995', 'n04037443', 'n04039381', 'n04040759', 'n04041544',
+    'n04044716', 'n04049303', 'n04065272', 'n04067472', 'n04069434',
+    'n04070727', 'n04074963', 'n04081281', 'n04086273', 'n04090263',
+    'n04099969', 'n04111531', 'n04116512', 'n04118538', 'n04118776',
+    'n04120489', 'n04125021', 'n04127249', 'n04131690', 'n04133789',
+    'n04136333', 'n04141076', 'n04141327', 'n04141975', 'n04146614',
+    'n04147183', 'n04149813', 'n04152593', 'n04153751', 'n04154565',
+    'n04162706', 'n04179913', 'n04192698', 'n04200800', 'n04201297',
+    'n04204238', 'n04204347', 'n04208210', 'n04209133', 'n04209239',
+    'n04228054', 'n04229816', 'n04235860', 'n04238763', 'n04239074',
+    'n04243546', 'n04251144', 'n04252077', 'n04252225', 'n04254120',
+    'n04254680', 'n04254777', 'n04258138', 'n04259630', 'n04263257',
+    'n04264628', 'n04265275', 'n04266014', 'n04270147', 'n04273569',
+    'n04275548', 'n04277352', 'n04285008', 'n04286575', 'n04296562',
+    'n04310018', 'n04311004', 'n04311174', 'n04317175', 'n04325704',
+    'n04326547', 'n04328186', 'n04330267', 'n04332243', 'n04335435',
+    'n04336792', 'n04344873', 'n04346328', 'n04347754', 'n04350905',
+    'n04355338', 'n04355933', 'n04356056', 'n04357314', 'n04366367',
+    'n04367480', 'n04370456', 'n04371430', 'n04371774', 'n04372370',
+    'n04376876', 'n04380533', 'n04389033', 'n04392985', 'n04398044',
+    'n04399382', 'n04404412', 'n04409515', 'n04417672', 'n04418357',
+    'n04423845', 'n04428191', 'n04429376', 'n04435653', 'n04442312',
+    'n04443257', 'n04447861', 'n04456115', 'n04458633', 'n04461696',
+    'n04462240', 'n04465501', 'n04467665', 'n04476259', 'n04479046',
+    'n04482393', 'n04483307', 'n04485082', 'n04486054', 'n04487081',
+    'n04487394', 'n04493381', 'n04501370', 'n04505470', 'n04507155',
+    'n04509417', 'n04515003', 'n04517823', 'n04522168', 'n04523525',
+    'n04525038', 'n04525305', 'n04532106', 'n04532670', 'n04536866',
+    'n04540053', 'n04542943', 'n04548280', 'n04548362', 'n04550184',
+    'n04552348', 'n04553703', 'n04554684', 'n04557648', 'n04560804',
+    'n04562935', 'n04579145', 'n04579432', 'n04584207', 'n04589890',
+    'n04590129', 'n04591157', 'n04591713', 'n04592741', 'n04596742',
+    'n04597913', 'n04599235', 'n04604644', 'n04606251', 'n04612504',
+    'n04613696', 'n06359193', 'n06596364', 'n06785654', 'n06794110',
+    'n06874185', 'n07248320', 'n07565083', 'n07579787', 'n07583066',
+    'n07584110', 'n07590611', 'n07613480', 'n07614500', 'n07615774',
+    'n07684084', 'n07693725', 'n07695742', 'n07697313', 'n07697537',
+    'n07711569', 'n07714571', 'n07714990', 'n07715103', 'n07716358',
+    'n07716906', 'n07717410', 'n07717556', 'n07718472', 'n07718747',
+    'n07720875', 'n07730033', 'n07734744', 'n07742313', 'n07745940',
+    'n07747607', 'n07749582', 'n07753113', 'n07753275', 'n07753592',
+    'n07754684', 'n07760859', 'n07768694', 'n07802026', 'n07831146',
+    'n07836838', 'n07860988', 'n07871810', 'n07873807', 'n07875152',
+    'n07880968', 'n07892512', 'n07920052', 'n07930864', 'n07932039',
+    'n09193705', 'n09229709', 'n09246464', 'n09256479', 'n09288635',
+    'n09332890', 'n09399592', 'n09421951', 'n09428293', 'n09468604',
+    'n09472597', 'n09835506', 'n10148035', 'n10565667', 'n11879895',
+    'n11939491', 'n12057211', 'n12144580', 'n12267677', 'n12620546',
+    'n12768682', 'n12985857', 'n12998815', 'n13037406', 'n13040303',
+    'n13044778', 'n13052670', 'n13054560', 'n13133613', 'n15075141'
+]
+
+imagenet_a_wnids = [
+    'n01498041', 'n01531178', 'n01534433', 'n01558993', 'n01580077',
+    'n01614925', 'n01616318', 'n01631663', 'n01641577', 'n01669191',
+    'n01677366', 'n01687978', 'n01694178', 'n01698640', 'n01735189',
+    'n01770081', 'n01770393', 'n01774750', 'n01784675', 'n01819313',
+    'n01820546', 'n01833805', 'n01843383', 'n01847000', 'n01855672',
+    'n01882714', 'n01910747', 'n01914609', 'n01924916', 'n01944390',
+    'n01985128', 'n01986214', 'n02007558', 'n02009912', 'n02037110',
+    'n02051845', 'n02077923', 'n02085620', 'n02099601', 'n02106550',
+    'n02106662', 'n02110958', 'n02119022', 'n02123394', 'n02127052',
+    'n02129165', 'n02133161', 'n02137549', 'n02165456', 'n02174001',
+    'n02177972', 'n02190166', 'n02206856', 'n02219486', 'n02226429',
+    'n02231487', 'n02233338', 'n02236044', 'n02259212', 'n02268443',
+    'n02279972', 'n02280649', 'n02281787', 'n02317335', 'n02325366',
+    'n02346627', 'n02356798', 'n02361337', 'n02410509', 'n02445715',
+    'n02454379', 'n02486410', 'n02492035', 'n02504458', 'n02655020',
+    'n02669723', 'n02672831', 'n02676566', 'n02690373', 'n02701002',
+    'n02730930', 'n02777292', 'n02782093', 'n02787622', 'n02793495',
+    'n02797295', 'n02802426', 'n02814860', 'n02815834', 'n02837789',
+    'n02879718', 'n02883205', 'n02895154', 'n02906734', 'n02948072',
+    'n02951358', 'n02980441', 'n02992211', 'n02999410', 'n03014705',
+    'n03026506', 'n03124043', 'n03125729', 'n03187595', 'n03196217',
+    'n03223299', 'n03250847', 'n03255030', 'n03291819', 'n03325584',
+    'n03355925', 'n03384352', 'n03388043', 'n03417042', 'n03443371',
+    'n03444034', 'n03445924', 'n03452741', 'n03483316', 'n03584829',
+    'n03590841', 'n03594945', 'n03617480', 'n03666591', 'n03670208',
+    'n03717622', 'n03720891', 'n03721384', 'n03724870', 'n03775071',
+    'n03788195', 'n03804744', 'n03837869', 'n03840681', 'n03854065',
+    'n03888257', 'n03891332', 'n03935335', 'n03982430', 'n04019541',
+    'n04033901', 'n04039381', 'n04067472', 'n04086273', 'n04099969',
+    'n04118538', 'n04131690', 'n04133789', 'n04141076', 'n04146614',
+    'n04147183', 'n04179913', 'n04208210', 'n04235860', 'n04252077',
+    'n04252225', 'n04254120', 'n04270147', 'n04275548', 'n04310018',
+    'n04317175', 'n04344873', 'n04347754', 'n04355338', 'n04366367',
+    'n04376876', 'n04389033', 'n04399382', 'n04442312', 'n04456115',
+    'n04482393', 'n04507155', 'n04509417', 'n04532670', 'n04540053',
+    'n04554684', 'n04562935', 'n04591713', 'n04606251', 'n07583066',
+    'n07695742', 'n07697313', 'n07697537', 'n07714990', 'n07718472',
+    'n07720875', 'n07734744', 'n07749582', 'n07753592', 'n07760859',
+    'n07768694', 'n07831146', 'n09229709', 'n09246464', 'n09472597',
+    'n09835506', 'n11879895', 'n12057211', 'n12144580', 'n12267677'
+]
+
+imagenet_a_mask = [wnid in set(imagenet_a_wnids) for wnid in all_wnids]
+
+imagenet_r_wnids = {
+    'n01443537', 'n01484850', 'n01494475', 'n01498041', 'n01514859',
+    'n01518878', 'n01531178', 'n01534433', 'n01614925', 'n01616318',
+    'n01630670', 'n01632777', 'n01644373', 'n01677366', 'n01694178',
+    'n01748264', 'n01770393', 'n01774750', 'n01784675', 'n01806143',
+    'n01820546', 'n01833805', 'n01843383', 'n01847000', 'n01855672',
+    'n01860187', 'n01882714', 'n01910747', 'n01944390', 'n01983481',
+    'n01986214', 'n02007558', 'n02009912', 'n02051845', 'n02056570',
+    'n02066245', 'n02071294', 'n02077923', 'n02085620', 'n02086240',
+    'n02088094', 'n02088238', 'n02088364', 'n02088466', 'n02091032',
+    'n02091134', 'n02092339', 'n02094433', 'n02096585', 'n02097298',
+    'n02098286', 'n02099601', 'n02099712', 'n02102318', 'n02106030',
+    'n02106166', 'n02106550', 'n02106662', 'n02108089', 'n02108915',
+    'n02109525', 'n02110185', 'n02110341', 'n02110958', 'n02112018',
+    'n02112137', 'n02113023', 'n02113624', 'n02113799', 'n02114367',
+    'n02117135', 'n02119022', 'n02123045', 'n02128385', 'n02128757',
+    'n02129165', 'n02129604', 'n02130308', 'n02134084', 'n02138441',
+    'n02165456', 'n02190166', 'n02206856', 'n02219486', 'n02226429',
+    'n02233338', 'n02236044', 'n02268443', 'n02279972', 'n02317335',
+    'n02325366', 'n02346627', 'n02356798', 'n02363005', 'n02364673',
+    'n02391049', 'n02395406', 'n02398521', 'n02410509', 'n02423022',
+    'n02437616', 'n02445715', 'n02447366', 'n02480495', 'n02480855',
+    'n02481823', 'n02483362', 'n02486410', 'n02510455', 'n02526121',
+    'n02607072', 'n02655020', 'n02672831', 'n02701002', 'n02749479',
+    'n02769748', 'n02793495', 'n02797295', 'n02802426', 'n02808440',
+    'n02814860', 'n02823750', 'n02841315', 'n02843684', 'n02883205',
+    'n02906734', 'n02909870', 'n02939185', 'n02948072', 'n02950826',
+    'n02951358', 'n02966193', 'n02980441', 'n02992529', 'n03124170',
+    'n03272010', 'n03345487', 'n03372029', 'n03424325', 'n03452741',
+    'n03467068', 'n03481172', 'n03494278', 'n03495258', 'n03498962',
+    'n03594945', 'n03602883', 'n03630383', 'n03649909', 'n03676483',
+    'n03710193', 'n03773504', 'n03775071', 'n03888257', 'n03930630',
+    'n03947888', 'n04086273', 'n04118538', 'n04133789', 'n04141076',
+    'n04146614', 'n04147183', 'n04192698', 'n04254680', 'n04266014',
+    'n04275548', 'n04310018', 'n04325704', 'n04347754', 'n04389033',
+    'n04409515', 'n04465501', 'n04487394', 'n04522168', 'n04536866',
+    'n04552348', 'n04591713', 'n07614500', 'n07693725', 'n07695742',
+    'n07697313', 'n07697537', 'n07714571', 'n07714990', 'n07718472',
+    'n07720875', 'n07734744', 'n07742313', 'n07745940', 'n07749582',
+    'n07753275', 'n07753592', 'n07768694', 'n07873807', 'n07880968',
+    'n07920052', 'n09472597', 'n09835506', 'n10565667', 'n12267677'
+}
+
+imagenet_r_mask = [wnid in imagenet_r_wnids for wnid in all_wnids]

InternVL/classification/dataset/imagenet_real.py

@@ -0,0 +1,50 @@
+# --------------------------------------------------------
+# EVA: Exploring the Limits of Masked Visual Representation Learning at Scale (https://arxiv.org/abs/2211.07636)
+# Github source: https://github.com/baaivision/EVA
+# Copyright (c) 2022 Beijing Academy of Artificial Intelligence (BAAI)
+# Licensed under The MIT License [see LICENSE for details]
+# By Yuxin Fang
+# Based on timm, DINO, DeiT and BEiT codebases
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm
+# https://github.com/facebookresearch/deit
+# https://github.com/facebookresearch/dino
+# https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------'
+
+import json
+import os
+
+import numpy as np
+
+
+class RealLabelsImagenet:
+
+    def __init__(self, filenames, real_json='real.json', topk=(1, 5)):
+        with open(real_json) as real_labels:
+            real_labels = json.load(real_labels)
+            real_labels = {f'ILSVRC2012_val_{i + 1:08d}.JPEG': labels for i, labels in enumerate(real_labels)}
+        self.real_labels = real_labels
+        self.filenames = filenames
+        assert len(self.filenames) == len(self.real_labels)
+        self.topk = topk
+        self.is_correct = {k: [] for k in topk}
+        self.sample_idx = 0
+
+    def add_result(self, output):
+        maxk = max(self.topk)
+        _, pred_batch = output.topk(maxk, 1, True, True)
+        pred_batch = pred_batch.cpu().numpy()
+        for pred in pred_batch:
+            filename = self.filenames[self.sample_idx]
+            filename = os.path.basename(filename)
+            if self.real_labels[filename]:
+                for k in self.topk:
+                    self.is_correct[k].append(
+                        any([p in self.real_labels[filename] for p in pred[:k]]))
+            self.sample_idx += 1
+
+    def get_accuracy(self, k=None):
+        if k is None:
+            return {k: float(np.mean(self.is_correct[k] for k in self.topk))}
+        else:
+            return float(np.mean(self.is_correct[k])) * 100

InternVL/classification/dataset/imagenetv2.py

@@ -0,0 +1,59 @@
+"""Code from https://github.com/mlfoundations/wise-ft/blob/master/src/datasets/imagenetv2.py
+Thanks to the authors of wise-ft."""
+import pathlib
+import shutil
+import tarfile
+
+import requests
+from PIL import Image
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+URLS = {'matched-frequency': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-matched-frequency.tar.gz',
+        'threshold-0.7': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-threshold0.7.tar.gz',
+        'top-images': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenetv2-top-images.tar.gz',
+        'val': 'https://imagenetv2public.s3-us-west-2.amazonaws.com/imagenet_validation.tar.gz'}
+
+FNAMES = {'matched-frequency': 'imagenetv2-matched-frequency-format-val',
+          'threshold-0.7': 'imagenetv2-threshold0.7-format-val',
+          'top-images': 'imagenetv2-top-images-format-val',
+          'val': 'imagenet_validation'}
+
+V2_DATASET_SIZE = 10000
+VAL_DATASET_SIZE = 50000
+
+
+class ImageNetV2Dataset(Dataset):
+    def __init__(self, variant='matched-frequency', transform=None, location='.'):
+        self.dataset_root = pathlib.Path(f'{location}/ImageNetV2-{variant}/')
+        self.tar_root = pathlib.Path(f'{location}/ImageNetV2-{variant}.tar.gz')
+        self.fnames = list(self.dataset_root.glob('**/*.jpeg'))
+        self.transform = transform
+        assert variant in URLS, f'unknown V2 Variant: {variant}'
+        if not self.dataset_root.exists() or len(self.fnames) != V2_DATASET_SIZE:
+            if not self.tar_root.exists():
+                print(f'Dataset {variant} not found on disk, downloading....')
+                response = requests.get(URLS[variant], stream=True)
+                total_size_in_bytes = int(response.headers.get('content-length', 0))
+                block_size = 1024  # 1 Kibibyte
+                progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+                with open(self.tar_root, 'wb') as f:
+                    for data in response.iter_content(block_size):
+                        progress_bar.update(len(data))
+                        f.write(data)
+                progress_bar.close()
+                if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+                    assert False, f'Downloading from {URLS[variant]} failed'
+            print('Extracting....')
+            tarfile.open(self.tar_root).extractall(f'{location}')
+            shutil.move(f'{location}/{FNAMES[variant]}', self.dataset_root)
+            self.fnames = list(self.dataset_root.glob('**/*.jpeg'))
+
+    def __len__(self):
+        return len(self.fnames)
+
+    def __getitem__(self, i):
+        img, label = Image.open(self.fnames[i]), int(self.fnames[i].parent.name)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label

InternVL/classification/dataset/samplers.py

@@ -0,0 +1,116 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import math
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data.sampler import Sampler
+
+
+class SubsetRandomSampler(torch.utils.data.Sampler):
+    """Samples elements randomly from a given list of indices, without
+    replacement.
+
+    Arguments:
+        indices (sequence): a sequence of indices
+    """
+
+    def __init__(self, indices):
+        self.epoch = 0
+        self.indices = indices
+
+    def __iter__(self):
+        return (self.indices[i] for i in torch.randperm(len(self.indices)))
+
+    def __len__(self):
+        return len(self.indices)
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+
+class NodeDistributedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 local_rank=None,
+                 local_size=None):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    'Requires distributed package to be available')
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    'Requires distributed package to be available')
+            rank = dist.get_rank()
+        if local_rank is None:
+            local_rank = int(os.environ.get('LOCAL_RANK', 0))
+        if local_size is None:
+            local_size = int(os.environ.get('LOCAL_SIZE', 1))
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.num_parts = local_size
+        self.rank = rank
+        self.local_rank = local_rank
+        self.epoch = 0
+        self.num_samples = int(
+            math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+        self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+
+        t = torch.Generator()
+        t.manual_seed(0)
+
+        indices = torch.randperm(len(self.dataset), generator=t).tolist()
+        # indices = range(len(self.dataset))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+
+        index = torch.randperm(len(indices), generator=g).tolist()
+        indices = list(np.array(indices)[index])
+
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

InternVL/classification/dataset/zipreader.py

@@ -0,0 +1,102 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import io
+import os
+import zipfile
+
+import numpy as np
+from PIL import Image, ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def is_zip_path(img_or_path):
+    """judge if this is a zip path."""
+    return '.zip@' in img_or_path
+
+
+class ZipReader(object):
+    """A class to read zipped files."""
+    zip_bank = dict()
+
+    def __init__(self):
+        super(ZipReader, self).__init__()
+
+    @staticmethod
+    def get_zipfile(path):
+        zip_bank = ZipReader.zip_bank
+        if path not in zip_bank:
+            zfile = zipfile.ZipFile(path, 'r')
+            zip_bank[path] = zfile
+        return zip_bank[path]
+
+    @staticmethod
+    def split_zip_style_path(path):
+        pos_at = path.index('@')
+        assert pos_at != -1, "character '@' is not found from the given path '%s'" % path
+
+        zip_path = path[0:pos_at]
+        folder_path = path[pos_at + 1:]
+        folder_path = str.strip(folder_path, '/')
+        return zip_path, folder_path
+
+    @staticmethod
+    def list_folder(path):
+        zip_path, folder_path = ZipReader.split_zip_style_path(path)
+
+        zfile = ZipReader.get_zipfile(zip_path)
+        folder_list = []
+        for file_foler_name in zfile.namelist():
+            file_foler_name = str.strip(file_foler_name, '/')
+            if file_foler_name.startswith(folder_path) and \
+                    len(os.path.splitext(file_foler_name)[-1]) == 0 and \
+                    file_foler_name != folder_path:
+                if len(folder_path) == 0:
+                    folder_list.append(file_foler_name)
+                else:
+                    folder_list.append(file_foler_name[len(folder_path) + 1:])
+
+        return folder_list
+
+    @staticmethod
+    def list_files(path, extension=None):
+        if extension is None:
+            extension = ['.*']
+        zip_path, folder_path = ZipReader.split_zip_style_path(path)
+
+        zfile = ZipReader.get_zipfile(zip_path)
+        file_lists = []
+        for file_foler_name in zfile.namelist():
+            file_foler_name = str.strip(file_foler_name, '/')
+            if file_foler_name.startswith(folder_path) and \
+                    str.lower(os.path.splitext(file_foler_name)[-1]) in extension:
+                if len(folder_path) == 0:
+                    file_lists.append(file_foler_name)
+                else:
+                    file_lists.append(file_foler_name[len(folder_path) + 1:])
+
+        return file_lists
+
+    @staticmethod
+    def read(path):
+        zip_path, path_img = ZipReader.split_zip_style_path(path)
+        zfile = ZipReader.get_zipfile(zip_path)
+        data = zfile.read(path_img)
+        return data
+
+    @staticmethod
+    def imread(path):
+        zip_path, path_img = ZipReader.split_zip_style_path(path)
+        zfile = ZipReader.get_zipfile(zip_path)
+        data = zfile.read(path_img)
+        try:
+            im = Image.open(io.BytesIO(data))
+        except:
+            print('ERROR IMG LOADED: ', path_img)
+            random_img = np.random.rand(224, 224, 3) * 255
+            im = Image.fromarray(np.uint8(random_img))
+        return im

InternVL/classification/meta_data/imagenet_classes.json

@@ -0,0 +1,1002 @@
+{
+    "n01440764": 0,
+    "n01443537": 1,
+    "n01484850": 2,
+    "n01491361": 3,
+    "n01494475": 4,
+    "n01496331": 5,
+    "n01498041": 6,
+    "n01514668": 7,
+    "n01514859": 8,
+    "n01518878": 9,
+    "n01530575": 10,
+    "n01531178": 11,
+    "n01532829": 12,
+    "n01534433": 13,
+    "n01537544": 14,
+    "n01558993": 15,
+    "n01560419": 16,
+    "n01580077": 17,
+    "n01582220": 18,
+    "n01592084": 19,
+    "n01601694": 20,
+    "n01608432": 21,
+    "n01614925": 22,
+    "n01616318": 23,
+    "n01622779": 24,
+    "n01629819": 25,
+    "n01630670": 26,
+    "n01631663": 27,
+    "n01632458": 28,
+    "n01632777": 29,
+    "n01641577": 30,
+    "n01644373": 31,
+    "n01644900": 32,
+    "n01664065": 33,
+    "n01665541": 34,
+    "n01667114": 35,
+    "n01667778": 36,
+    "n01669191": 37,
+    "n01675722": 38,
+    "n01677366": 39,
+    "n01682714": 40,
+    "n01685808": 41,
+    "n01687978": 42,
+    "n01688243": 43,
+    "n01689811": 44,
+    "n01692333": 45,
+    "n01693334": 46,
+    "n01694178": 47,
+    "n01695060": 48,
+    "n01697457": 49,
+    "n01698640": 50,
+    "n01704323": 51,
+    "n01728572": 52,
+    "n01728920": 53,
+    "n01729322": 54,
+    "n01729977": 55,
+    "n01734418": 56,
+    "n01735189": 57,
+    "n01737021": 58,
+    "n01739381": 59,
+    "n01740131": 60,
+    "n01742172": 61,
+    "n01744401": 62,
+    "n01748264": 63,
+    "n01749939": 64,
+    "n01751748": 65,
+    "n01753488": 66,
+    "n01755581": 67,
+    "n01756291": 68,
+    "n01768244": 69,
+    "n01770081": 70,
+    "n01770393": 71,
+    "n01773157": 72,
+    "n01773549": 73,
+    "n01773797": 74,
+    "n01774384": 75,
+    "n01774750": 76,
+    "n01775062": 77,
+    "n01776313": 78,
+    "n01784675": 79,
+    "n01795545": 80,
+    "n01796340": 81,
+    "n01797886": 82,
+    "n01798484": 83,
+    "n01806143": 84,
+    "n01806567": 85,
+    "n01807496": 86,
+    "n01817953": 87,
+    "n01818515": 88,
+    "n01819313": 89,
+    "n01820546": 90,
+    "n01824575": 91,
+    "n01828970": 92,
+    "n01829413": 93,
+    "n01833805": 94,
+    "n01843065": 95,
+    "n01843383": 96,
+    "n01847000": 97,
+    "n01855032": 98,
+    "n01855672": 99,
+    "n01860187": 100,
+    "n01871265": 101,
+    "n01872401": 102,
+    "n01873310": 103,
+    "n01877812": 104,
+    "n01882714": 105,
+    "n01883070": 106,
+    "n01910747": 107,
+    "n01914609": 108,
+    "n01917289": 109,
+    "n01924916": 110,
+    "n01930112": 111,
+    "n01943899": 112,
+    "n01944390": 113,
+    "n01945685": 114,
+    "n01950731": 115,
+    "n01955084": 116,
+    "n01968897": 117,
+    "n01978287": 118,
+    "n01978455": 119,
+    "n01980166": 120,
+    "n01981276": 121,
+    "n01983481": 122,
+    "n01984695": 123,
+    "n01985128": 124,
+    "n01986214": 125,
+    "n01990800": 126,
+    "n02002556": 127,
+    "n02002724": 128,
+    "n02006656": 129,
+    "n02007558": 130,
+    "n02009229": 131,
+    "n02009912": 132,
+    "n02011460": 133,
+    "n02012849": 134,
+    "n02013706": 135,
+    "n02017213": 136,
+    "n02018207": 137,
+    "n02018795": 138,
+    "n02025239": 139,
+    "n02027492": 140,
+    "n02028035": 141,
+    "n02033041": 142,
+    "n02037110": 143,
+    "n02051845": 144,
+    "n02056570": 145,
+    "n02058221": 146,
+    "n02066245": 147,
+    "n02071294": 148,
+    "n02074367": 149,
+    "n02077923": 150,
+    "n02085620": 151,
+    "n02085782": 152,
+    "n02085936": 153,
+    "n02086079": 154,
+    "n02086240": 155,
+    "n02086646": 156,
+    "n02086910": 157,
+    "n02087046": 158,
+    "n02087394": 159,
+    "n02088094": 160,
+    "n02088238": 161,
+    "n02088364": 162,
+    "n02088466": 163,
+    "n02088632": 164,
+    "n02089078": 165,
+    "n02089867": 166,
+    "n02089973": 167,
+    "n02090379": 168,
+    "n02090622": 169,
+    "n02090721": 170,
+    "n02091032": 171,
+    "n02091134": 172,
+    "n02091244": 173,
+    "n02091467": 174,
+    "n02091635": 175,
+    "n02091831": 176,
+    "n02092002": 177,
+    "n02092339": 178,
+    "n02093256": 179,
+    "n02093428": 180,
+    "n02093647": 181,
+    "n02093754": 182,
+    "n02093859": 183,
+    "n02093991": 184,
+    "n02094114": 185,
+    "n02094258": 186,
+    "n02094433": 187,
+    "n02095314": 188,
+    "n02095570": 189,
+    "n02095889": 190,
+    "n02096051": 191,
+    "n02096177": 192,
+    "n02096294": 193,
+    "n02096437": 194,
+    "n02096585": 195,
+    "n02097047": 196,
+    "n02097130": 197,
+    "n02097209": 198,
+    "n02097298": 199,
+    "n02097474": 200,
+    "n02097658": 201,
+    "n02098105": 202,
+    "n02098286": 203,
+    "n02098413": 204,
+    "n02099267": 205,
+    "n02099429": 206,
+    "n02099601": 207,
+    "n02099712": 208,
+    "n02099849": 209,
+    "n02100236": 210,
+    "n02100583": 211,
+    "n02100735": 212,
+    "n02100877": 213,
+    "n02101006": 214,
+    "n02101388": 215,
+    "n02101556": 216,
+    "n02102040": 217,
+    "n02102177": 218,
+    "n02102318": 219,
+    "n02102480": 220,
+    "n02102973": 221,
+    "n02104029": 222,
+    "n02104365": 223,
+    "n02105056": 224,
+    "n02105162": 225,
+    "n02105251": 226,
+    "n02105412": 227,
+    "n02105505": 228,
+    "n02105641": 229,
+    "n02105855": 230,
+    "n02106030": 231,
+    "n02106166": 232,
+    "n02106382": 233,
+    "n02106550": 234,
+    "n02106662": 235,
+    "n02107142": 236,
+    "n02107312": 237,
+    "n02107574": 238,
+    "n02107683": 239,
+    "n02107908": 240,
+    "n02108000": 241,
+    "n02108089": 242,
+    "n02108422": 243,
+    "n02108551": 244,
+    "n02108915": 245,
+    "n02109047": 246,
+    "n02109525": 247,
+    "n02109961": 248,
+    "n02110063": 249,
+    "n02110185": 250,
+    "n02110341": 251,
+    "n02110627": 252,
+    "n02110806": 253,
+    "n02110958": 254,
+    "n02111129": 255,
+    "n02111277": 256,
+    "n02111500": 257,
+    "n02111889": 258,
+    "n02112018": 259,
+    "n02112137": 260,
+    "n02112350": 261,
+    "n02112706": 262,
+    "n02113023": 263,
+    "n02113186": 264,
+    "n02113624": 265,
+    "n02113712": 266,
+    "n02113799": 267,
+    "n02113978": 268,
+    "n02114367": 269,
+    "n02114548": 270,
+    "n02114712": 271,
+    "n02114855": 272,
+    "n02115641": 273,
+    "n02115913": 274,
+    "n02116738": 275,
+    "n02117135": 276,
+    "n02119022": 277,
+    "n02119789": 278,
+    "n02120079": 279,
+    "n02120505": 280,
+    "n02123045": 281,
+    "n02123159": 282,
+    "n02123394": 283,
+    "n02123597": 284,
+    "n02124075": 285,
+    "n02125311": 286,
+    "n02127052": 287,
+    "n02128385": 288,
+    "n02128757": 289,
+    "n02128925": 290,
+    "n02129165": 291,
+    "n02129604": 292,
+    "n02130308": 293,
+    "n02132136": 294,
+    "n02133161": 295,
+    "n02134084": 296,
+    "n02134418": 297,
+    "n02137549": 298,
+    "n02138441": 299,
+    "n02165105": 300,
+    "n02165456": 301,
+    "n02167151": 302,
+    "n02168699": 303,
+    "n02169497": 304,
+    "n02172182": 305,
+    "n02174001": 306,
+    "n02177972": 307,
+    "n02190166": 308,
+    "n02206856": 309,
+    "n02219486": 310,
+    "n02226429": 311,
+    "n02229544": 312,
+    "n02231487": 313,
+    "n02233338": 314,
+    "n02236044": 315,
+    "n02256656": 316,
+    "n02259212": 317,
+    "n02264363": 318,
+    "n02268443": 319,
+    "n02268853": 320,
+    "n02276258": 321,
+    "n02277742": 322,
+    "n02279972": 323,
+    "n02280649": 324,
+    "n02281406": 325,
+    "n02281787": 326,
+    "n02317335": 327,
+    "n02319095": 328,
+    "n02321529": 329,
+    "n02325366": 330,
+    "n02326432": 331,
+    "n02328150": 332,
+    "n02342885": 333,
+    "n02346627": 334,
+    "n02356798": 335,
+    "n02361337": 336,
+    "n02363005": 337,
+    "n02364673": 338,
+    "n02389026": 339,
+    "n02391049": 340,
+    "n02395406": 341,
+    "n02396427": 342,
+    "n02397096": 343,
+    "n02398521": 344,
+    "n02403003": 345,
+    "n02408429": 346,
+    "n02410509": 347,
+    "n02412080": 348,
+    "n02415577": 349,
+    "n02417914": 350,
+    "n02422106": 351,
+    "n02422699": 352,
+    "n02423022": 353,
+    "n02437312": 354,
+    "n02437616": 355,
+    "n02441942": 356,
+    "n02442845": 357,
+    "n02443114": 358,
+    "n02443484": 359,
+    "n02444819": 360,
+    "n02445715": 361,
+    "n02447366": 362,
+    "n02454379": 363,
+    "n02457408": 364,
+    "n02480495": 365,
+    "n02480855": 366,
+    "n02481823": 367,
+    "n02483362": 368,
+    "n02483708": 369,
+    "n02484975": 370,
+    "n02486261": 371,
+    "n02486410": 372,
+    "n02487347": 373,
+    "n02488291": 374,
+    "n02488702": 375,
+    "n02489166": 376,
+    "n02490219": 377,
+    "n02492035": 378,
+    "n02492660": 379,
+    "n02493509": 380,
+    "n02493793": 381,
+    "n02494079": 382,
+    "n02497673": 383,
+    "n02500267": 384,
+    "n02504013": 385,
+    "n02504458": 386,
+    "n02509815": 387,
+    "n02510455": 388,
+    "n02514041": 389,
+    "n02526121": 390,
+    "n02536864": 391,
+    "n02606052": 392,
+    "n02607072": 393,
+    "n02640242": 394,
+    "n02641379": 395,
+    "n02643566": 396,
+    "n02655020": 397,
+    "n02666196": 398,
+    "n02667093": 399,
+    "n02669723": 400,
+    "n02672831": 401,
+    "n02676566": 402,
+    "n02687172": 403,
+    "n02690373": 404,
+    "n02692877": 405,
+    "n02699494": 406,
+    "n02701002": 407,
+    "n02704792": 408,
+    "n02708093": 409,
+    "n02727426": 410,
+    "n02730930": 411,
+    "n02747177": 412,
+    "n02749479": 413,
+    "n02769748": 414,
+    "n02776631": 415,
+    "n02777292": 416,
+    "n02782093": 417,
+    "n02783161": 418,
+    "n02786058": 419,
+    "n02787622": 420,
+    "n02788148": 421,
+    "n02790996": 422,
+    "n02791124": 423,
+    "n02791270": 424,
+    "n02793495": 425,
+    "n02794156": 426,
+    "n02795169": 427,
+    "n02797295": 428,
+    "n02799071": 429,
+    "n02802426": 430,
+    "n02804414": 431,
+    "n02804610": 432,
+    "n02807133": 433,
+    "n02808304": 434,
+    "n02808440": 435,
+    "n02814533": 436,
+    "n02814860": 437,
+    "n02815834": 438,
+    "n02817516": 439,
+    "n02823428": 440,
+    "n02823750": 441,
+    "n02825657": 442,
+    "n02834397": 443,
+    "n02835271": 444,
+    "n02837789": 445,
+    "n02840245": 446,
+    "n02841315": 447,
+    "n02843684": 448,
+    "n02859443": 449,
+    "n02860847": 450,
+    "n02865351": 451,
+    "n02869837": 452,
+    "n02870880": 453,
+    "n02871525": 454,
+    "n02877765": 455,
+    "n02879718": 456,
+    "n02883205": 457,
+    "n02892201": 458,
+    "n02892767": 459,
+    "n02894605": 460,
+    "n02895154": 461,
+    "n02906734": 462,
+    "n02909870": 463,
+    "n02910353": 464,
+    "n02916936": 465,
+    "n02917067": 466,
+    "n02927161": 467,
+    "n02930766": 468,
+    "n02939185": 469,
+    "n02948072": 470,
+    "n02950826": 471,
+    "n02951358": 472,
+    "n02951585": 473,
+    "n02963159": 474,
+    "n02965783": 475,
+    "n02966193": 476,
+    "n02966687": 477,
+    "n02971356": 478,
+    "n02974003": 479,
+    "n02977058": 480,
+    "n02978881": 481,
+    "n02979186": 482,
+    "n02980441": 483,
+    "n02981792": 484,
+    "n02988304": 485,
+    "n02992211": 486,
+    "n02992529": 487,
+    "n02999410": 488,
+    "n03000134": 489,
+    "n03000247": 490,
+    "n03000684": 491,
+    "n03014705": 492,
+    "n03016953": 493,
+    "n03017168": 494,
+    "n03018349": 495,
+    "n03026506": 496,
+    "n03028079": 497,
+    "n03032252": 498,
+    "n03041632": 499,
+    "n03042490": 500,
+    "n03045698": 501,
+    "n03047690": 502,
+    "n03062245": 503,
+    "n03063599": 504,
+    "n03063689": 505,
+    "n03065424": 506,
+    "n03075370": 507,
+    "n03085013": 508,
+    "n03089624": 509,
+    "n03095699": 510,
+    "n03100240": 511,
+    "n03109150": 512,
+    "n03110669": 513,
+    "n03124043": 514,
+    "n03124170": 515,
+    "n03125729": 516,
+    "n03126707": 517,
+    "n03127747": 518,
+    "n03127925": 519,
+    "n03131574": 520,
+    "n03133878": 521,
+    "n03134739": 522,
+    "n03141823": 523,
+    "n03146219": 524,
+    "n03160309": 525,
+    "n03179701": 526,
+    "n03180011": 527,
+    "n03187595": 528,
+    "n03188531": 529,
+    "n03196217": 530,
+    "n03197337": 531,
+    "n03201208": 532,
+    "n03207743": 533,
+    "n03207941": 534,
+    "n03208938": 535,
+    "n03216828": 536,
+    "n03218198": 537,
+    "n03220513": 538,
+    "n03223299": 539,
+    "n03240683": 540,
+    "n03249569": 541,
+    "n03250847": 542,
+    "n03255030": 543,
+    "n03259280": 544,
+    "n03271574": 545,
+    "n03272010": 546,
+    "n03272562": 547,
+    "n03290653": 548,
+    "n03291819": 549,
+    "n03297495": 550,
+    "n03314780": 551,
+    "n03325584": 552,
+    "n03337140": 553,
+    "n03344393": 554,
+    "n03345487": 555,
+    "n03347037": 556,
+    "n03355925": 557,
+    "n03372029": 558,
+    "n03376595": 559,
+    "n03379051": 560,
+    "n03384352": 561,
+    "n03388043": 562,
+    "n03388183": 563,
+    "n03388549": 564,
+    "n03393912": 565,
+    "n03394916": 566,
+    "n03400231": 567,
+    "n03404251": 568,
+    "n03417042": 569,
+    "n03424325": 570,
+    "n03425413": 571,
+    "n03443371": 572,
+    "n03444034": 573,
+    "n03445777": 574,
+    "n03445924": 575,
+    "n03447447": 576,
+    "n03447721": 577,
+    "n03450230": 578,
+    "n03452741": 579,
+    "n03457902": 580,
+    "n03459775": 581,
+    "n03461385": 582,
+    "n03467068": 583,
+    "n03476684": 584,
+    "n03476991": 585,
+    "n03478589": 586,
+    "n03481172": 587,
+    "n03482405": 588,
+    "n03483316": 589,
+    "n03485407": 590,
+    "n03485794": 591,
+    "n03492542": 592,
+    "n03494278": 593,
+    "n03495258": 594,
+    "n03496892": 595,
+    "n03498962": 596,
+    "n03527444": 597,
+    "n03529860": 598,
+    "n03530642": 599,
+    "n03532672": 600,
+    "n03534580": 601,
+    "n03535780": 602,
+    "n03538406": 603,
+    "n03544143": 604,
+    "n03584254": 605,
+    "n03584829": 606,
+    "n03590841": 607,
+    "n03594734": 608,
+    "n03594945": 609,
+    "n03595614": 610,
+    "n03598930": 611,
+    "n03599486": 612,
+    "n03602883": 613,
+    "n03617480": 614,
+    "n03623198": 615,
+    "n03627232": 616,
+    "n03630383": 617,
+    "n03633091": 618,
+    "n03637318": 619,
+    "n03642806": 620,
+    "n03649909": 621,
+    "n03657121": 622,
+    "n03658185": 623,
+    "n03661043": 624,
+    "n03662601": 625,
+    "n03666591": 626,
+    "n03670208": 627,
+    "n03673027": 628,
+    "n03676483": 629,
+    "n03680355": 630,
+    "n03690938": 631,
+    "n03691459": 632,
+    "n03692522": 633,
+    "n03697007": 634,
+    "n03706229": 635,
+    "n03709823": 636,
+    "n03710193": 637,
+    "n03710637": 638,
+    "n03710721": 639,
+    "n03717622": 640,
+    "n03720891": 641,
+    "n03721384": 642,
+    "n03724870": 643,
+    "n03729826": 644,
+    "n03733131": 645,
+    "n03733281": 646,
+    "n03733805": 647,
+    "n03742115": 648,
+    "n03743016": 649,
+    "n03759954": 650,
+    "n03761084": 651,
+    "n03763968": 652,
+    "n03764736": 653,
+    "n03769881": 654,
+    "n03770439": 655,
+    "n03770679": 656,
+    "n03773504": 657,
+    "n03775071": 658,
+    "n03775546": 659,
+    "n03776460": 660,
+    "n03777568": 661,
+    "n03777754": 662,
+    "n03781244": 663,
+    "n03782006": 664,
+    "n03785016": 665,
+    "n03786901": 666,
+    "n03787032": 667,
+    "n03788195": 668,
+    "n03788365": 669,
+    "n03791053": 670,
+    "n03792782": 671,
+    "n03792972": 672,
+    "n03793489": 673,
+    "n03794056": 674,
+    "n03796401": 675,
+    "n03803284": 676,
+    "n03804744": 677,
+    "n03814639": 678,
+    "n03814906": 679,
+    "n03825788": 680,
+    "n03832673": 681,
+    "n03837869": 682,
+    "n03838899": 683,
+    "n03840681": 684,
+    "n03841143": 685,
+    "n03843555": 686,
+    "n03854065": 687,
+    "n03857828": 688,
+    "n03866082": 689,
+    "n03868242": 690,
+    "n03868863": 691,
+    "n03871628": 692,
+    "n03873416": 693,
+    "n03874293": 694,
+    "n03874599": 695,
+    "n03876231": 696,
+    "n03877472": 697,
+    "n03877845": 698,
+    "n03884397": 699,
+    "n03887697": 700,
+    "n03888257": 701,
+    "n03888605": 702,
+    "n03891251": 703,
+    "n03891332": 704,
+    "n03895866": 705,
+    "n03899768": 706,
+    "n03902125": 707,
+    "n03903868": 708,
+    "n03908618": 709,
+    "n03908714": 710,
+    "n03916031": 711,
+    "n03920288": 712,
+    "n03924679": 713,
+    "n03929660": 714,
+    "n03929855": 715,
+    "n03930313": 716,
+    "n03930630": 717,
+    "n03933933": 718,
+    "n03935335": 719,
+    "n03937543": 720,
+    "n03938244": 721,
+    "n03942813": 722,
+    "n03944341": 723,
+    "n03947888": 724,
+    "n03950228": 725,
+    "n03954731": 726,
+    "n03956157": 727,
+    "n03958227": 728,
+    "n03961711": 729,
+    "n03967562": 730,
+    "n03970156": 731,
+    "n03976467": 732,
+    "n03976657": 733,
+    "n03977966": 734,
+    "n03980874": 735,
+    "n03982430": 736,
+    "n03983396": 737,
+    "n03991062": 738,
+    "n03992509": 739,
+    "n03995372": 740,
+    "n03998194": 741,
+    "n04004767": 742,
+    "n04005630": 743,
+    "n04008634": 744,
+    "n04009552": 745,
+    "n04019541": 746,
+    "n04023962": 747,
+    "n04026417": 748,
+    "n04033901": 749,
+    "n04033995": 750,
+    "n04037443": 751,
+    "n04039381": 752,
+    "n04040759": 753,
+    "n04041544": 754,
+    "n04044716": 755,
+    "n04049303": 756,
+    "n04065272": 757,
+    "n04067472": 758,
+    "n04069434": 759,
+    "n04070727": 760,
+    "n04074963": 761,
+    "n04081281": 762,
+    "n04086273": 763,
+    "n04090263": 764,
+    "n04099969": 765,
+    "n04111531": 766,
+    "n04116512": 767,
+    "n04118538": 768,
+    "n04118776": 769,
+    "n04120489": 770,
+    "n04125021": 771,
+    "n04127249": 772,
+    "n04131690": 773,
+    "n04133789": 774,
+    "n04136333": 775,
+    "n04141076": 776,
+    "n04141327": 777,
+    "n04141975": 778,
+    "n04146614": 779,
+    "n04147183": 780,
+    "n04149813": 781,
+    "n04152593": 782,
+    "n04153751": 783,
+    "n04154565": 784,
+    "n04162706": 785,
+    "n04179913": 786,
+    "n04192698": 787,
+    "n04200800": 788,
+    "n04201297": 789,
+    "n04204238": 790,
+    "n04204347": 791,
+    "n04208210": 792,
+    "n04209133": 793,
+    "n04209239": 794,
+    "n04228054": 795,
+    "n04229816": 796,
+    "n04235860": 797,
+    "n04238763": 798,
+    "n04239074": 799,
+    "n04243546": 800,
+    "n04251144": 801,
+    "n04252077": 802,
+    "n04252225": 803,
+    "n04254120": 804,
+    "n04254680": 805,
+    "n04254777": 806,
+    "n04258138": 807,
+    "n04259630": 808,
+    "n04263257": 809,
+    "n04264628": 810,
+    "n04265275": 811,
+    "n04266014": 812,
+    "n04270147": 813,
+    "n04273569": 814,
+    "n04275548": 815,
+    "n04277352": 816,
+    "n04285008": 817,
+    "n04286575": 818,
+    "n04296562": 819,
+    "n04310018": 820,
+    "n04311004": 821,
+    "n04311174": 822,
+    "n04317175": 823,
+    "n04325704": 824,
+    "n04326547": 825,
+    "n04328186": 826,
+    "n04330267": 827,
+    "n04332243": 828,
+    "n04335435": 829,
+    "n04336792": 830,
+    "n04344873": 831,
+    "n04346328": 832,
+    "n04347754": 833,
+    "n04350905": 834,
+    "n04355338": 835,
+    "n04355933": 836,
+    "n04356056": 837,
+    "n04357314": 838,
+    "n04366367": 839,
+    "n04367480": 840,
+    "n04370456": 841,
+    "n04371430": 842,
+    "n04371774": 843,
+    "n04372370": 844,
+    "n04376876": 845,
+    "n04380533": 846,
+    "n04389033": 847,
+    "n04392985": 848,
+    "n04398044": 849,
+    "n04399382": 850,
+    "n04404412": 851,
+    "n04409515": 852,
+    "n04417672": 853,
+    "n04418357": 854,
+    "n04423845": 855,
+    "n04428191": 856,
+    "n04429376": 857,
+    "n04435653": 858,
+    "n04442312": 859,
+    "n04443257": 860,
+    "n04447861": 861,
+    "n04456115": 862,
+    "n04458633": 863,
+    "n04461696": 864,
+    "n04462240": 865,
+    "n04465501": 866,
+    "n04467665": 867,
+    "n04476259": 868,
+    "n04479046": 869,
+    "n04482393": 870,
+    "n04483307": 871,
+    "n04485082": 872,
+    "n04486054": 873,
+    "n04487081": 874,
+    "n04487394": 875,
+    "n04493381": 876,
+    "n04501370": 877,
+    "n04505470": 878,
+    "n04507155": 879,
+    "n04509417": 880,
+    "n04515003": 881,
+    "n04517823": 882,
+    "n04522168": 883,
+    "n04523525": 884,
+    "n04525038": 885,
+    "n04525305": 886,
+    "n04532106": 887,
+    "n04532670": 888,
+    "n04536866": 889,
+    "n04540053": 890,
+    "n04542943": 891,
+    "n04548280": 892,
+    "n04548362": 893,
+    "n04550184": 894,
+    "n04552348": 895,
+    "n04553703": 896,
+    "n04554684": 897,
+    "n04557648": 898,
+    "n04560804": 899,
+    "n04562935": 900,
+    "n04579145": 901,
+    "n04579432": 902,
+    "n04584207": 903,
+    "n04589890": 904,
+    "n04590129": 905,
+    "n04591157": 906,
+    "n04591713": 907,
+    "n04592741": 908,
+    "n04596742": 909,
+    "n04597913": 910,
+    "n04599235": 911,
+    "n04604644": 912,
+    "n04606251": 913,
+    "n04612504": 914,
+    "n04613696": 915,
+    "n06359193": 916,
+    "n06596364": 917,
+    "n06785654": 918,
+    "n06794110": 919,
+    "n06874185": 920,
+    "n07248320": 921,
+    "n07565083": 922,
+    "n07579787": 923,
+    "n07583066": 924,
+    "n07584110": 925,
+    "n07590611": 926,
+    "n07613480": 927,
+    "n07614500": 928,
+    "n07615774": 929,
+    "n07684084": 930,
+    "n07693725": 931,
+    "n07695742": 932,
+    "n07697313": 933,
+    "n07697537": 934,
+    "n07711569": 935,
+    "n07714571": 936,
+    "n07714990": 937,
+    "n07715103": 938,
+    "n07716358": 939,
+    "n07716906": 940,
+    "n07717410": 941,
+    "n07717556": 942,
+    "n07718472": 943,
+    "n07718747": 944,
+    "n07720875": 945,
+    "n07730033": 946,
+    "n07734744": 947,
+    "n07742313": 948,
+    "n07745940": 949,
+    "n07747607": 950,
+    "n07749582": 951,
+    "n07753113": 952,
+    "n07753275": 953,
+    "n07753592": 954,
+    "n07754684": 955,
+    "n07760859": 956,
+    "n07768694": 957,
+    "n07802026": 958,
+    "n07831146": 959,
+    "n07836838": 960,
+    "n07860988": 961,
+    "n07871810": 962,
+    "n07873807": 963,
+    "n07875152": 964,
+    "n07880968": 965,
+    "n07892512": 966,
+    "n07920052": 967,
+    "n07930864": 968,
+    "n07932039": 969,
+    "n09193705": 970,
+    "n09229709": 971,
+    "n09246464": 972,
+    "n09256479": 973,
+    "n09288635": 974,
+    "n09332890": 975,
+    "n09399592": 976,
+    "n09421951": 977,
+    "n09428293": 978,
+    "n09468604": 979,
+    "n09472597": 980,
+    "n09835506": 981,
+    "n10148035": 982,
+    "n10565667": 983,
+    "n11879895": 984,
+    "n11939491": 985,
+    "n12057211": 986,
+    "n12144580": 987,
+    "n12267677": 988,
+    "n12620546": 989,
+    "n12768682": 990,
+    "n12985857": 991,
+    "n12998815": 992,
+    "n13037406": 993,
+    "n13040303": 994,
+    "n13044778": 995,
+    "n13052670": 996,
+    "n13054560": 997,
+    "n13133613": 998,
+    "n15075141": 999
+  }

InternVL/classification/meta_data/map22kto1k.txt

@@ -0,0 +1,1000 @@
+359
+368
+460
+475
+486
+492
+496
+514
+516
+525
+547
+548
+556
+563
+575
+641
+648
+723
+733
+765
+801
+826
+852
+858
+878
+896
+900
+905
+908
+910
+935
+946
+947
+994
+999
+1003
+1005
+1010
+1027
+1029
+1048
+1055
+1064
+1065
+1069
+1075
+1079
+1081
+1085
+1088
+1093
+1106
+1143
+1144
+1145
+1147
+1168
+1171
+1178
+1187
+1190
+1197
+1205
+1216
+1223
+1230
+1236
+1241
+1245
+1257
+1259
+1260
+1267
+1268
+1269
+1271
+1272
+1273
+1277
+1303
+1344
+1349
+1355
+1357
+1384
+1388
+1391
+1427
+1429
+1432
+1437
+1450
+1461
+1462
+1474
+1502
+1503
+1512
+1552
+1555
+1577
+1584
+1587
+1589
+1599
+1615
+1616
+1681
+1692
+1701
+1716
+1729
+1757
+1759
+1764
+1777
+1786
+1822
+1841
+1842
+1848
+1850
+1856
+1860
+1861
+1864
+1876
+1897
+1898
+1910
+1913
+1918
+1922
+1928
+1932
+1935
+1947
+1951
+1953
+1970
+1977
+1979
+2001
+2017
+2067
+2081
+2087
+2112
+2128
+2135
+2147
+2174
+2175
+2176
+2177
+2178
+2181
+2183
+2184
+2187
+2189
+2190
+2191
+2192
+2193
+2197
+2202
+2203
+2206
+2208
+2209
+2211
+2212
+2213
+2214
+2215
+2216
+2217
+2219
+2222
+2223
+2224
+2225
+2226
+2227
+2228
+2229
+2230
+2236
+2238
+2240
+2241
+2242
+2243
+2244
+2245
+2247
+2248
+2249
+2250
+2251
+2252
+2255
+2256
+2257
+2262
+2263
+2264
+2265
+2266
+2268
+2270
+2271
+2272
+2273
+2275
+2276
+2279
+2280
+2281
+2282
+2285
+2289
+2292
+2295
+2296
+2297
+2298
+2299
+2300
+2301
+2302
+2303
+2304
+2305
+2306
+2309
+2310
+2312
+2313
+2314
+2315
+2316
+2318
+2319
+2321
+2322
+2326
+2329
+2330
+2331
+2332
+2334
+2335
+2336
+2337
+2338
+2339
+2341
+2342
+2343
+2344
+2346
+2348
+2349
+2351
+2352
+2353
+2355
+2357
+2358
+2359
+2360
+2364
+2365
+2368
+2369
+2377
+2382
+2383
+2385
+2397
+2398
+2400
+2402
+2405
+2412
+2421
+2428
+2431
+2432
+2433
+2436
+2441
+2445
+2450
+2453
+2454
+2465
+2469
+2532
+2533
+2538
+2544
+2547
+2557
+2565
+2578
+2612
+2658
+2702
+2722
+2731
+2738
+2741
+2747
+2810
+2818
+2833
+2844
+2845
+2867
+2874
+2882
+2884
+2888
+2889
+3008
+3012
+3019
+3029
+3033
+3042
+3091
+3106
+3138
+3159
+3164
+3169
+3280
+3296
+3311
+3318
+3320
+3324
+3330
+3366
+3375
+3381
+3406
+3419
+3432
+3434
+3435
+3493
+3495
+3503
+3509
+3511
+3513
+3517
+3521
+3526
+3546
+3554
+3600
+3601
+3606
+3612
+3613
+3616
+3622
+3623
+3627
+3632
+3634
+3636
+3638
+3644
+3646
+3649
+3650
+3651
+3656
+3663
+3673
+3674
+3689
+3690
+3702
+3733
+3769
+3971
+3974
+4065
+4068
+4073
+4102
+4136
+4140
+4151
+4159
+4165
+4207
+4219
+4226
+4249
+4256
+4263
+4270
+4313
+4321
+4378
+4386
+4478
+4508
+4512
+4536
+4542
+4550
+4560
+4562
+4570
+4571
+4572
+4583
+4588
+4594
+4604
+4608
+4623
+4634
+4636
+4646
+4651
+4652
+4686
+4688
+4691
+4699
+4724
+4727
+4737
+4770
+4774
+4789
+4802
+4807
+4819
+4880
+4886
+4908
+4927
+4931
+4936
+4964
+4976
+4993
+5028
+5033
+5043
+5046
+5096
+5111
+5114
+5131
+5132
+5183
+5199
+5235
+5275
+5291
+5293
+5294
+5343
+5360
+5362
+5364
+5390
+5402
+5418
+5428
+5430
+5437
+5443
+5473
+5484
+5486
+5505
+5507
+5508
+5510
+5567
+5578
+5580
+5584
+5606
+5613
+5629
+5672
+5676
+5692
+5701
+5760
+5769
+5770
+5779
+5814
+5850
+5871
+5893
+5911
+5949
+5954
+6005
+6006
+6012
+6017
+6023
+6024
+6040
+6050
+6054
+6087
+6105
+6157
+6235
+6237
+6256
+6259
+6286
+6291
+6306
+6339
+6341
+6343
+6379
+6383
+6393
+6405
+6479
+6511
+6517
+6541
+6561
+6608
+6611
+6615
+6678
+6682
+6707
+6752
+6798
+6850
+6880
+6885
+6890
+6920
+6981
+7000
+7009
+7038
+7049
+7050
+7052
+7073
+7078
+7098
+7111
+7165
+7198
+7204
+7280
+7283
+7286
+7287
+7293
+7294
+7305
+7318
+7341
+7346
+7354
+7382
+7427
+7428
+7435
+7445
+7450
+7455
+7467
+7469
+7497
+7502
+7506
+7514
+7523
+7651
+7661
+7664
+7672
+7679
+7685
+7696
+7730
+7871
+7873
+7895
+7914
+7915
+7920
+7934
+7935
+7949
+8009
+8036
+8051
+8065
+8074
+8090
+8112
+8140
+8164
+8168
+8178
+8182
+8198
+8212
+8216
+8230
+8242
+8288
+8289
+8295
+8318
+8352
+8368
+8371
+8375
+8376
+8401
+8416
+8419
+8436
+8460
+8477
+8478
+8482
+8498
+8500
+8539
+8543
+8552
+8555
+8580
+8584
+8586
+8594
+8598
+8601
+8606
+8610
+8611
+8622
+8627
+8639
+8649
+8650
+8653
+8654
+8667
+8672
+8673
+8674
+8676
+8684
+8720
+8723
+8750
+8753
+8801
+8815
+8831
+8835
+8842
+8845
+8858
+8897
+8916
+8951
+8954
+8959
+8970
+8976
+8981
+8983
+8989
+8991
+8993
+9019
+9039
+9042
+9043
+9056
+9057
+9070
+9087
+9098
+9106
+9130
+9131
+9155
+9171
+9183
+9198
+9199
+9201
+9204
+9211
+9220
+9224
+9228
+9249
+9259
+9270
+9278
+9294
+9299
+9309
+9321
+9344
+9351
+9375
+9376
+9381
+9391
+9400
+9404
+9440
+9448
+9463
+9474
+9501
+9504
+9513
+9514
+9544
+9566
+9575
+9607
+9608
+9623
+9632
+9638
+9642
+9655
+9673
+9739
+9751
+9759
+9766
+9777
+9801
+9819
+9838
+9878
+9923
+9955
+9960
+9962
+9969
+9996
+10009
+10030
+10039
+10051
+10072
+10074
+10077
+10093
+10096
+10108
+10117
+10120
+10123
+10157
+10225
+10275
+10303
+10306
+10313
+10314
+10331
+10336
+10337
+10412
+10422
+10450
+10462
+10464
+10486
+10518
+10521
+10522
+10531
+10533
+10534
+10550
+10558
+10573
+10582
+10585
+10588
+10611
+10625
+10634
+10637
+10676
+10682
+10725
+10775
+10781
+10782
+10806
+10836
+10839
+10847
+10858
+10870
+10880
+10883
+10907
+10913
+10920
+10935
+10946
+10950
+10951
+10956
+10998
+11002
+11017
+11022
+11024
+11026
+11044
+11054
+11094
+11109
+11136
+11136
+11167
+11185
+11220
+11222
+11241
+11254
+11258
+11278
+11305
+11310
+11330
+11366
+11376
+11388
+11391
+11400
+11406
+11436
+11448
+11465
+11468
+11472
+11477
+11482
+11483
+11506
+11535
+11557
+11565
+11574
+11583
+11593
+11610
+11611
+11618
+11620
+11639
+11642
+11663
+11673
+11688
+11708
+11709
+11715
+11720
+11725
+11728
+11742
+11759
+11770
+11836
+11838
+11855
+11875
+11877
+11883
+11888
+11895
+11916
+11922
+11929
+11943
+11951
+11979
+11983
+12213
+12228
+12238
+12240
+12241
+12246
+12282
+12348
+12368
+12372
+12421
+12559
+12565
+12574
+12687
+12754
+12767
+12777
+12779
+12811
+12831
+12834
+12835
+12842
+12846
+12848
+12849
+12855
+12857
+12872
+12937
+12970
+13016
+13037
+13045
+13058
+13084
+13085
+13087
+13093
+13133
+13181
+13229
+13405
+13443
+13613
+13689
+13697
+13708
+13748
+13803
+13981
+14050
+14058
+14218
+14245
+14255
+14263
+14293
+14323
+14366
+14388
+14393
+14437
+14441
+14964
+15730
+16742
+18035
+18203
+18533
+18790
+19100
+20017
+20460
+21024
+21043
+21161
+21169
+21179
+21194
+21198
+21367
+21815

InternVL/classification/models/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .build import build_model

InternVL/classification/models/build.py

@@ -0,0 +1,36 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .intern_vit_6b import InternViT6B
+
+
+def build_model(config):
+    model_type = config.MODEL.TYPE
+    if model_type == 'intern_vit_6b':
+        model = InternViT6B(
+            num_classes=config.MODEL.NUM_CLASSES,
+            patch_size=config.MODEL.INTERN_VIT_6B.PATCH_SIZE,
+            img_size=config.DATA.IMG_SIZE,
+            pretrain_size=config.MODEL.INTERN_VIT_6B.PRETRAIN_SIZE,
+            qkv_bias=config.MODEL.INTERN_VIT_6B.QKV_BIAS,
+            drop_path_rate=config.MODEL.DROP_PATH_RATE,
+            embed_dim=config.MODEL.INTERN_VIT_6B.EMBED_DIM,
+            num_heads=config.MODEL.INTERN_VIT_6B.NUM_HEADS,
+            mlp_ratio=config.MODEL.INTERN_VIT_6B.MLP_RATIO,
+            init_values=config.MODEL.INTERN_VIT_6B.INIT_VALUES,
+            qk_normalization=config.MODEL.INTERN_VIT_6B.QK_NORMALIZATION,
+            depth=config.MODEL.INTERN_VIT_6B.DEPTH,
+            use_flash_attn=config.MODEL.INTERN_VIT_6B.USE_FLASH_ATTN,
+            with_cp=config.TRAIN.USE_CHECKPOINT,
+            freeze_vit=config.MODEL.INTERN_VIT_6B.FREEZE_VIT,
+            pretrained=config.MODEL.INTERN_VIT_6B.PRETRAINED,
+            cls_target=config.MODEL.INTERN_VIT_6B.CLS_TARGET,
+            head_norm_type=config.MODEL.INTERN_VIT_6B.HEAD_NORM_TYPE,
+        )
+    else:
+        raise NotImplementedError(f'Unkown model: {model_type}')
+
+    return model

InternVL/classification/models/flash_attention.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+try:  # v1
+    from flash_attn.flash_attn_interface import \
+        flash_attn_unpadded_qkvpacked_func
+except:  # v2
+    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+
+from flash_attn.bert_padding import pad_input, unpad_input
+
+
+class FlashAttention(nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
+                max_s=None, need_weights=False):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
+                if unpadded: (nnz, 3, h, d)
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert not need_weights
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
+        assert qkv.is_cuda
+
+        if cu_seqlens is None:
+            batch_size = qkv.shape[0]
+            seqlen = qkv.shape[1]
+            if key_padding_mask is None:
+                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+                max_s = seqlen
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
+                                          device=qkv.device)
+                output = flash_attn_unpadded_qkvpacked_func(
+                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+            else:
+                nheads = qkv.shape[-2]
+                x = rearrange(qkv, 'b s three h d -> b s (three h d)')
+                x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)
+                x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+                output_unpad = flash_attn_unpadded_qkvpacked_func(
+                    x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
+                                             indices, batch_size, seqlen),
+                                   'b s (h d) -> b s h d', h=nheads)
+        else:
+            assert max_s is not None
+            output = flash_attn_unpadded_qkvpacked_func(
+                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+
+        return output, None

InternVL/classification/models/intern_vit_6b.py

@@ -0,0 +1,473 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from einops import rearrange
+from timm.models.layers import DropPath, to_2tuple
+
+try:
+    from .flash_attention import FlashAttention
+    has_flash_attn = True
+except:
+    print('FlashAttention is not installed.')
+    has_flash_attn = False
+
+
+def _freeze_params(module):
+    for param in module.parameters():
+        param.requires_grad = False
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., attn_head_dim=None, out_dim=None):
+        super().__init__()
+        if out_dim is None:
+            out_dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        assert all_head_dim == dim
+
+        self.q = nn.Linear(dim, all_head_dim, bias=False)
+        self.k = nn.Linear(dim, all_head_dim, bias=False)
+        self.v = nn.Linear(dim, all_head_dim, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.k_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, out_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, k=None, v=None):
+        B, N, C = x.shape
+        N_k = k.shape[1]
+        N_v = v.shape[1]
+
+        q_bias, k_bias, v_bias = None, None, None
+        if self.q_bias is not None:
+            q_bias = self.q_bias
+            k_bias = self.k_bias
+            v_bias = self.v_bias
+
+        q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
+        q = q.reshape(B, N, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)  # (B, N_head, N_q, dim)
+
+        k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
+        k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)
+
+        v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
+        v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))  # (B, N_head, N_q, N_k)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class AttentiveBlock(nn.Module):
+
+    def __init__(self, dim, num_heads, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, attn_head_dim=None, out_dim=None):
+        super().__init__()
+
+        self.norm1_q = norm_layer(dim)
+        self.norm1_k = norm_layer(dim)
+        self.norm1_v = norm_layer(dim)
+        self.cross_attn = CrossAttention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+            proj_drop=drop, attn_head_dim=attn_head_dim, out_dim=out_dim)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x_q, x_kv, pos_q, pos_k, bool_masked_pos, rel_pos_bias=None):
+        x_q = self.norm1_q(x_q + pos_q)
+        x_k = self.norm1_k(x_kv + pos_k)
+        x_v = self.norm1_v(x_kv)
+        x = self.cross_attn(x_q, k=x_k, v=x_v)
+
+        return x
+
+
+class AttentionPoolingBlock(AttentiveBlock):
+
+    def forward(self, x):
+        x_q = x.mean(1, keepdim=True)
+        x_kv, pos_q, pos_k = x, 0, 0
+        x = super().forward(x_q, x_kv, pos_q, pos_k, bool_masked_pos=None, rel_pos_bias=None)
+        x = x.squeeze(1)
+        return x
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from apex.normalization import FusedRMSNorm
+
+    RMSNorm = FusedRMSNorm  # noqa
+
+    print('Discovered apex.normalization.FusedRMSNorm - will use it instead of RMSNorm')
+except ImportError:
+    # using the normal RMSNorm
+    pass
+except Exception:
+    print('discovered apex but it failed to load, falling back to RMSNorm')
+    pass
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False, force_fp32=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+        self.force_fp32 = force_fp32
+
+    @torch.cuda.amp.autocast(enabled=False)
+    def forward(self, x):
+        if self.force_fp32:
+            output_type = x.dtype
+            out = x.float().mul_(self.gamma.float()) if self.inplace else x.float() * self.gamma.float()
+            return out.to(dtype=output_type)
+        else:
+            out = x.mul_(self.gamma) if self.inplace else x * self.gamma
+            return out
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_flash_attn=False,
+                 causal=False, norm_layer=nn.LayerNorm, qk_normalization=False):
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.use_flash_attn = use_flash_attn
+        if use_flash_attn:
+            self.causal = causal
+            self.inner_attn = FlashAttention(attention_dropout=attn_drop)
+
+        self.qk_normalization = qk_normalization
+        self.q_norm = norm_layer(dim) if qk_normalization else nn.Identity()
+        self.k_norm = norm_layer(dim) if qk_normalization else nn.Identity()
+
+    def _naive_attn(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        if self.qk_normalization:
+            B_, H_, N_, D_ = q.shape
+            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+
+        attn = ((q * self.scale) @ k.transpose(-2, -1))
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
+        qkv = self.qkv(x)
+        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
+            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        context, _ = self.inner_attn(
+            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=self.causal
+        )
+        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))
+        outs = self.proj_drop(outs)
+        return outs
+
+    def forward(self, x):
+        x = self._naive_attn(x) if not self.use_flash_attn else self._flash_attn(x)
+        return x
+
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU,
+                 bias=True, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        bias = to_2tuple(bias)
+        drop_probs = to_2tuple(drop)
+
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(
+            self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., init_values=None,
+            drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_flash_attn=False, with_cp=False,
+            qk_normalization=False, layerscale_force_fp32=False):
+        super().__init__()
+
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,
+                              use_flash_attn=use_flash_attn, causal=False, norm_layer=norm_layer,
+                              qk_normalization=qk_normalization)
+        self.ls1 = LayerScale(dim, init_values=init_values,
+                              force_fp32=layerscale_force_fp32) if init_values else nn.Identity()
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.ls2 = LayerScale(dim, init_values=init_values,
+                              force_fp32=layerscale_force_fp32) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.with_cp = with_cp
+
+    def forward(self, x):
+
+        def _inner_forward(x):
+            x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+            x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+            return x
+
+        if self.with_cp:
+            return checkpoint.checkpoint(_inner_forward, x)
+        else:
+            return _inner_forward(x)
+
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x, **kwargs):
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x, H, W
+
+
+class InternViT6B(nn.Module):
+
+    def __init__(self, in_chans=3, patch_size=14, img_size=224, pretrain_size=224, qkv_bias=False, drop_path_rate=0.0,
+                 embed_dim=3200, num_heads=25, mlp_ratio=4, init_values=0.1, qk_normalization=True, depth=48,
+                 use_flash_attn=True, with_cp=True, layerscale_force_fp32=False, freeze_vit=True,
+                 cls_target='cls_patch_concat', num_classes=1000, attn_pool_num_heads=16, clip_embed_dim=768,
+                 head_norm_type='bn', pretrained=None):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.pretrain_size = pretrain_size
+        self.drop_path_rate = drop_path_rate
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.cls_target = cls_target
+        self.depth = depth
+
+        use_flash_attn = use_flash_attn and has_flash_attn
+        if use_flash_attn and not has_flash_attn:
+            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')
+        use_flash_attn = [use_flash_attn] * depth if not isinstance(use_flash_attn, list) else use_flash_attn
+
+        norm_layer_for_blocks = partial(RMSNorm, eps=1e-6)
+        self.norm_layer_for_blocks = norm_layer_for_blocks
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+        self.num_patches = num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Identity()
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=qkv_bias,
+                  norm_layer=norm_layer_for_blocks,
+                  drop_path=dpr[i], init_values=init_values, attn_drop=0.,
+                  use_flash_attn=use_flash_attn[i],
+                  with_cp=with_cp,
+                  qk_normalization=qk_normalization,
+                  layerscale_force_fp32=layerscale_force_fp32)
+            for i in range(depth)])
+
+        if cls_target == 'clip_projector':
+            self.clip_projector = AttentionPoolingBlock(
+                dim=embed_dim, num_heads=attn_pool_num_heads, qkv_bias=True, qk_scale=None,
+                drop=0., attn_drop=0., norm_layer=partial(nn.LayerNorm, eps=1e-5), out_dim=clip_embed_dim)
+
+        self.init_weights(pretrained)
+
+        if freeze_vit:
+            _freeze_params(self)
+
+        if cls_target == 'cls_patch_concat':
+            if head_norm_type == 'bn':
+                self.norm = nn.SyncBatchNorm(embed_dim * 2, eps=1e-6)
+            else:
+                self.norm = nn.LayerNorm(embed_dim * 2, eps=1e-6)
+            self.head = nn.Linear(embed_dim * 2, num_classes) if num_classes > 0 else nn.Identity()
+        elif cls_target == 'clip_projector':
+            if head_norm_type == 'bn':
+                self.norm = nn.SyncBatchNorm(clip_embed_dim, eps=1e-6)
+            else:
+                self.norm = nn.LayerNorm(clip_embed_dim, eps=1e-6)
+            self.head = nn.Linear(clip_embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        else:
+            raise NotImplementedError
+
+        if type(self.head) != nn.Identity:
+            self.head.weight.data.normal_(mean=0.0, std=0.01)
+            self.head.bias.data.zero_()
+
+    def init_weights(self, pretrained=None):
+        print(f'pretrained: {pretrained}')
+
+        def resize_pos_embed(pos_embed, H, W):
+            cls = pos_embed[:, :1, :]
+            pos_embed = pos_embed[:, 1:, :].reshape(
+                1, self.pretrain_size // 14, self.pretrain_size // 14, -1).permute(0, 3, 1, 2)
+            pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \
+                reshape(1, -1, H * W).permute(0, 2, 1)
+            pos_embed = torch.cat([cls, pos_embed], dim=1)
+            return pos_embed
+
+        if isinstance(pretrained, str):
+            checkpoint = torch.load(pretrained, map_location='cpu')
+            if 'module' in checkpoint:
+                checkpoint = checkpoint['module']
+
+            # resize pos_embed
+            pos_embed = checkpoint['pos_embed']
+            checkpoint['pos_embed'] = resize_pos_embed(
+                pos_embed, self.img_size // self.patch_size, self.img_size // self.patch_size)
+            # resize patch_embed
+            patch_embed = checkpoint['patch_embed.proj.weight']
+            checkpoint['patch_embed.proj.weight'] = F.interpolate(
+                patch_embed, size=(self.patch_size, self.patch_size),
+                mode='bicubic', align_corners=False)
+            message = self.load_state_dict(checkpoint, strict=False)
+            print(message)
+
+    @property
+    def dtype(self):
+        return self.patch_embed.proj.weight.dtype
+
+    def forward_features(self, x):
+        x, _, _ = self.patch_embed(x.type(self.dtype))
+        batch_size, seq_len, _ = x.size()
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+
+        for idx, blk in enumerate(self.blocks):
+            x = blk(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        if self.cls_target == 'cls_patch_concat':
+            x = torch.cat((x[:, 0, :], x[:, 1:, :].mean(dim=1)), dim=-1)
+        elif self.cls_target == 'clip_projector':
+            x = self.clip_projector(x)
+        else:
+            raise NotImplementedError
+        x = self.norm(x)
+        x = self.head(x)
+        return x
+
+    @torch.jit.ignore
+    def lr_decay_keywords(self, decay_ratio=0.95):
+        lr_ratios = {}
+
+        # blocks
+        for idx in range(self.depth):
+            tag = 'blocks.{}.'.format(idx)
+            decay = 1.0 * (decay_ratio ** (self.depth - idx))
+            lr_ratios[tag] = decay
+
+        # patch_embed
+        lr_ratios['patch_embed'] = 1.0 * (decay_ratio ** (self.depth + 1))
+        lr_ratios['pos_embed'] = 1.0 * (decay_ratio ** (self.depth + 1))
+        lr_ratios['cls_token'] = 1.0 * (decay_ratio ** (self.depth + 1))
+
+        return lr_ratios

InternVL/internvl_chat/internvl/conversation.py

@@ -0,0 +1,406 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Any, Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ''
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep2, self.sep]
+            ret = self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# InternVL-Chat-V1-1 template
+register_conv_template(
+    Conversation(
+        name='internvl_zh',
+        system_template='',
+        roles=('<human>', '<bot>'),
+        sep_style=SeparatorStyle.INTERNVL_ZH,
+        sep='</s>',
+        sep2=' ',
+    )
+)
+
+
+# Both Hermes-2 and internlm2-chat are chatml-format conversation templates. The difference
+# is that during training, the preprocessing function for the Hermes-2 template doesn't add
+# <s> at the beginning of the tokenized sequence, while the internlm2-chat template does.
+# Therefore, they are completely equivalent during inference.
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            6,
+            7,
+            8,
+        ],
+        stop_str='<|endoftext|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internlm2-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            92543,
+            92542
+        ]
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='phi3-chat',
+        system_template='<|system|>\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开���的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|user|>\n', '<|assistant|>\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|end|>',
+        stop_token_ids=[
+            2,
+            32000,
+            32007
+        ]
+    )
+)

InternVL/internvl_chat/internvl/dist_utils.py

@@ -0,0 +1,104 @@
+import os
+import socket
+import subprocess
+from datetime import timedelta
+
+import deepspeed
+import torch
+import torch.multiprocessing as mp
+from torch import distributed as dist
+
+timeout = timedelta(minutes=60)
+
+
+def _find_free_port():
+    # Copied from https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py # noqa: E501
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    # Binding to port 0 will cause the OS to find an available port for us
+    sock.bind(('', 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    # NOTE: there is still a chance the port could be taken by other processes.
+    return port
+
+
+def _is_free_port(port):
+    ips = socket.gethostbyname_ex(socket.gethostname())[-1]
+    ips.append('localhost')
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return all(s.connect_ex((ip, port)) != 0 for ip in ips)
+
+
+def init_dist(launcher, backend='nccl', **kwargs):
+    if mp.get_start_method(allow_none=True) is None:
+        mp.set_start_method('spawn')
+    if launcher == 'pytorch':
+        _init_dist_pytorch(backend, **kwargs)
+    elif launcher == 'mpi':
+        _init_dist_mpi(backend, **kwargs)
+    elif launcher == 'slurm':
+        _init_dist_slurm(backend, **kwargs)
+    else:
+        raise ValueError(f'Invalid launcher type: {launcher}')
+
+
+def _init_dist_pytorch(backend, **kwargs):
+    # TODO: use local_rank instead of rank % num_gpus
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    # dist.init_process_group(backend=backend, **kwargs)
+    deepspeed.init_distributed(dist_backend=backend)
+
+
+def _init_dist_mpi(backend, **kwargs):
+    local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+    torch.cuda.set_device(local_rank)
+    if 'MASTER_PORT' not in os.environ:
+        # 29500 is torch.distributed default port
+        os.environ['MASTER_PORT'] = '29500'
+    if 'MASTER_ADDR' not in os.environ:
+        raise KeyError('The environment variable MASTER_ADDR is not set')
+    os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE']
+    os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK']
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def _init_dist_slurm(backend, port=None):
+    """Initialize slurm distributed training environment.
+
+    If argument ``port`` is not specified, then the master port will be system
+    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
+    environment variable, then a default port ``29500`` will be used.
+
+    Args:
+        backend (str): Backend of torch.distributed.
+        port (int, optional): Master port. Defaults to None.
+    """
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    addr = subprocess.getoutput(
+        f'scontrol show hostname {node_list} | head -n1')
+    # specify master port
+    if port is not None:
+        os.environ['MASTER_PORT'] = str(port)
+    elif 'MASTER_PORT' in os.environ:
+        pass  # use MASTER_PORT in the environment variable
+    else:
+        # if torch.distributed default port(29500) is available
+        # then use it, else find a free port
+        if _is_free_port(29500):
+            os.environ['MASTER_PORT'] = '29500'
+        else:
+            os.environ['MASTER_PORT'] = str(_find_free_port())
+    # use MASTER_ADDR in the environment variable if it already exists
+    if 'MASTER_ADDR' not in os.environ:
+        os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+    os.environ['RANK'] = str(proc_id)
+    # dist.init_process_group(backend=backend, timeout=timeout)
+    deepspeed.init_distributed(dist_backend=backend)

InternVL/internvl_chat/tools/convert_parquet.py

@@ -0,0 +1,83 @@
+import argparse
+import io
+import json
+import os
+import random
+
+import pandas as pd
+from PIL import Image
+from tqdm import tqdm
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('input', type=str, default='')
+argparse.add_argument('output', type=str, default='')
+
+args = argparse.parse_args()
+
+if not os.path.exists(args.output):
+    os.makedirs(args.output)
+
+image_root = os.path.join(args.output, 'images')
+if not os.path.exists(image_root):
+    os.makedirs(image_root)
+
+prompts = [
+    'Please recognize the text in the image.',
+    'Please extract the text from the image.',
+    'Kindly identify and transcribe the text present in the image.',
+    'Could you please perform optical character recognition (OCR) on the image to retrieve the text?',
+    'Please use text recognition techniques to decipher the text within the image.',
+    'Could you extract any readable text contained in the image?',
+    'I need the text within the image recognized and converted into machine-readable format, please.',
+    'Please employ OCR technology to recognize and extract the text from the image.',
+    'Kindly process the image to identify and retrieve any textual content it contains.',
+    'Please analyze the image and retrieve any textual information that is discernible.',
+    'Could you transcribe any visible text from the image, please?',
+    '请从图像中提取文本',
+    '请识别图像中的文本。',
+    '能否使用光学字符识别（OCR）技术在图像上提取文本？',
+    '请使用文本识别技术解读图像中的文字。',
+    '能提取图像中的任何可读文本吗？',
+    '请将图像中的文本识别并转换为机器可读格式。',
+    '请使用OCR技术识别并提取图像中的文本。',
+    '请处理图像以识别并提取其中包含的任何文本内容。',
+    '请分析图像并提取其中可以辨认的任何文本信息。',
+    '你能够将图像中可见的文本转录出来吗？',
+]
+
+cnt = 0
+data = []
+
+
+def process(filename):
+    df_parquet = pd.read_parquet(os.path.join(args.input, filename))
+    for index, row in tqdm(df_parquet.iterrows()):
+        # 在这里对每一行的数据进行处理
+        image = row['image']['bytes']
+        ground_truth = row['ground_truth']
+        ground_truth = json.loads(ground_truth)['gt_parse']['text_sequence']
+        image = Image.open(io.BytesIO(image))
+        global cnt, data
+        image_out_path = os.path.join(args.output, 'images/%08d.jpg' % cnt)
+        image.save(image_out_path)
+        data_item = {'id': cnt, 'image': 'images/%08d.jpg' % cnt}
+        conversations = []
+        conversations.append({'from': 'human', 'value': '<image>\n' + random.choice(prompts)})
+        conversations.append({'from': 'gpt', 'value': ground_truth})
+        data_item['conversations'] = conversations
+        data.append(data_item)
+        cnt += 1
+
+
+process('train-00000-of-00084-26dbc51f3d0903b9.parquet')
+process('train-00001-of-00084-3efa94914043c815.parquet')
+process('train-00002-of-00084-65600b4a95c96e85.parquet')
+process('train-00003-of-00084-45e260eca2cd125f.parquet')
+process('train-00004-of-00084-b1d684c57dc6e3da.parquet')
+
+writer = open(os.path.join(args.output, 'synthdog_en.jsonl'), 'w')
+for item in data:
+    writer.write(json.dumps(item) + '\n')
+writer.close()
+
+print('done')

InternVL/internvl_chat/tools/convert_to_int8.py

@@ -0,0 +1,16 @@
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+path = 'OpenGVLab/InternVL-Chat-V1-5'
+model = AutoModel.from_pretrained(
+    path,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True,
+    load_in_8bit=True).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+
+model.save_pretrained('release/InternVL-Chat-V1-5-Int8')
+tokenizer.save_pretrained('release/InternVL-Chat-V1-5-Int8')
+print('finished')

InternVL/internvl_chat/tools/extract_mlp.py

@@ -0,0 +1,19 @@
+import argparse
+import os.path
+
+import torch
+from internvl.model.internvl_chat import InternVLChatModel
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('model_path', type=str, default='')
+argparse.add_argument('output_path', type=str, default='')
+
+args = argparse.parse_args()
+
+model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+model = model.mlp1.to(torch.bfloat16)
+
+ckpt = model.state_dict()
+output_path = os.path.join(args.output_path, 'mlp_projector.pth')
+torch.save(ckpt, output_path)
+print('finished')

InternVL/internvl_chat/tools/extract_video_frames.py

@@ -0,0 +1,121 @@
+import concurrent.futures
+import json
+import os
+import random
+
+import av
+import numpy as np
+import torch
+from decord import VideoReader, cpu
+from PIL import Image
+from tqdm.auto import tqdm
+
+num_segments = 1
+
+# root directory of evaluation dimension 10
+dimension10_dir = './videos/20bn-something-something-v2'
+# root directory of evaluation dimension 11
+dimension11_dir = './videos/EPIC-KITCHENS'
+# root directory of evaluation dimension 12
+dimension12_dir = './videos/BreakfastII_15fps_qvga_sync'
+
+
+def transform_video(buffer):
+    try:
+        buffer = buffer.numpy()
+    except AttributeError:
+        try:
+            buffer = buffer.asnumpy()
+        except AttributeError:
+            print('Both buffer.numpy() and buffer.asnumpy() failed.')
+            buffer = None
+    images_group = list()
+    for fid in range(len(buffer)):
+        images_group.append(Image.fromarray(buffer[fid]))
+    return images_group
+
+
+def get_index(num_frames, num_segments):
+    if num_segments > num_frames:
+        offsets = np.array([
+            idx for idx in range(num_frames)
+        ])
+    else:
+        # uniform sampling
+        seg_size = float(num_frames - 1) / num_segments
+        start = int(seg_size / 2)
+        offsets = np.array([
+            start + int(np.round(seg_size * idx)) for idx in range(num_segments)
+        ])
+    return offsets
+
+
+def fetch_images(qa_item):
+    use_pyav = False
+    segment = None
+    if qa_item['question_type_id'] == 10:
+        data_path = os.path.join(dimension10_dir, qa_item['data_id'])
+        start = 0.0
+        end = 0.0
+    elif qa_item['question_type_id'] == 11:
+        data_path = os.path.join(dimension11_dir, qa_item['data_id'].split('/')[-1])
+        segment = qa_item['segment']
+        start, end = segment[0], segment[1]
+    elif qa_item['question_type_id'] == 12:
+        data_path = os.path.join(dimension12_dir, qa_item['data_id'])
+        segment = qa_item['segment']
+        start, end = segment[0], segment[1]
+        use_pyav = True
+
+    if use_pyav:
+        # using pyav for decoding videos in evaluation dimension 12
+        reader = av.open(data_path)
+        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
+        video_len = len(frames)
+        start_frame, end_frame = start, end
+        end_frame = min(end_frame, video_len)
+        offset = get_index(end_frame - start_frame, num_segments)
+        frame_indices = offset + start_frame
+        buffer = torch.stack([frames[idx] for idx in frame_indices])
+    else:
+        # using decord for decoding videos in evaluation dimension 10-11
+        vr = VideoReader(data_path, num_threads=1, ctx=cpu(0))
+        video_len = len(vr)
+        fps = vr.get_avg_fps()
+        if segment is not None:
+            # obtain start and end frame for the video segment in evaluation dimension 11
+            start_frame = int(min(max(start * fps, 0), video_len - 1))
+            end_frame = int(min(max(end * fps, 0), video_len - 1))
+            tot_frames = int(end_frame - start_frame)
+            offset = get_index(tot_frames, num_segments)
+            frame_indices = offset + start_frame
+        else:
+            # sample frames of the video in evaluation dimension 10
+            frame_indices = get_index(video_len - 1, num_segments)
+        vr.seek(0)
+        buffer = vr.get_batch(frame_indices)
+    return transform_video(buffer)
+
+
+def fetch_images_parallel(qa_item):
+    return qa_item, fetch_images(qa_item)
+
+
+if __name__ == '__main__':
+    data = json.load(open('SEED-Bench.json'))
+    video_img_dir = 'SEED-Bench-video-image'
+    ques_type_id_to_name = {id:n for n,id in data['question_type'].items()}
+
+    video_data = [x for x in data['questions'] if x['data_type'] == 'video']
+
+    with open(output, 'w') as f, concurrent.futures.ThreadPoolExecutor() as executor:
+        future_to_images = {executor.submit(fetch_images_parallel, qa_item): qa_item for qa_item in video_data}
+        for future in tqdm(concurrent.futures.as_completed(future_to_images), total=len(future_to_images)):
+            qa_item = future_to_images[future]
+            try:
+                qa_item, images = future.result()
+            except Exception as exc:
+                print(f'{qa_item} generated an exception: {exc}')
+            else:
+                img_file = f"{qa_item['question_type_id']}_{qa_item['question_id']}.png"
+                images[0].save(os.path.join(video_img_dir, img_file))

InternVL/internvl_chat/tools/extract_vit.py

@@ -0,0 +1,16 @@
+import argparse
+
+import torch
+from internvl.model.internvl_chat import InternVLChatModel
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('model_path', type=str, default='')
+argparse.add_argument('output_path', type=str, default='')
+
+args = argparse.parse_args()
+
+model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+model = model.vision_model.to(torch.bfloat16)
+
+model.save_pretrained(args.output_path)
+print('finished')

InternVL/internvl_chat/tools/json2jsonl.py

@@ -0,0 +1,21 @@
+
+import argparse
+import json
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('path', type=str)
+
+args = argparse.parse_args()
+
+assert args.path.endswith('.json')
+
+data = json.load(open(args.path))
+writer = open(args.path.replace('.json', '.jsonl'), 'w')
+for idx, item in enumerate(data):
+    conversations = item['conversations']
+    if conversations[0]['from'] == 'system':
+        item['conversations'] = item['conversations'][1:]
+    item['id'] = idx
+    writer.write(json.dumps(item, ensure_ascii=False) + '\n')
+
+writer.close()

InternVL/internvl_chat/tools/jsonl2jsonl.py

@@ -0,0 +1,23 @@
+import argparse
+import json
+import os
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('path', type=str)
+
+args = argparse.parse_args()
+
+assert args.path.endswith('.jsonl')
+
+f = open(args.path)
+data = [json.loads(line) for line in f.readlines()]
+writer = open(args.path.replace('.jsonl', '_new.jsonl'), 'w')
+for idx, item in enumerate(data):
+    item['id'] = idx
+    # item['image'] = os.path.join('train/documents/', item['image'])
+    conversations = item['conversations']
+    if conversations[0]['from'] == 'system':
+        item['conversations'] = item['conversations'][1:]
+    writer.write(json.dumps(item, ensure_ascii=False) + '\n')
+
+writer.close()

InternVL/internvl_chat/tools/merge_lora.py

@@ -0,0 +1,31 @@
+import argparse
+
+import torch
+from internvl.model.internvl_chat import InternVLChatModel
+from transformers import AutoTokenizer
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('input_path', type=str, help='Path to the input model')
+argparse.add_argument('output_path', type=str, help='Path to the output model')
+args = argparse.parse_args()
+
+print('Loading model...')
+model = InternVLChatModel.from_pretrained(
+    args.input_path, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).eval()
+print('Loading tokenizer...')
+tokenizer = AutoTokenizer.from_pretrained(args.input_path, trust_remote_code=True)
+
+if model.config.use_backbone_lora:
+    model.vision_model.merge_and_unload()
+    model.vision_model = model.vision_model.model
+    model.config.use_backbone_lora = 0
+if model.config.use_llm_lora:
+    model.language_model.merge_and_unload()
+    model.language_model = model.language_model.model
+    model.config.use_llm_lora = 0
+
+print('Saving model...')
+model.save_pretrained(args.output_path)
+print('Saving tokenizer...')
+tokenizer.save_pretrained(args.output_path)
+print('Done!')

InternVL/internvl_chat/tools/replace_llm.py

@@ -0,0 +1,29 @@
+import argparse
+
+import torch
+from internvl.model.internvl_chat import InternVLChatModel
+from transformers import AutoModel, AutoTokenizer
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('model_path', type=str, default='')
+argparse.add_argument('llm_path', type=str, default='')
+
+args = argparse.parse_args()
+
+if args.model_path[-1] == '/':
+    args.model_path = args.model_path[:-1]
+
+model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+
+llm = AutoModel.from_pretrained(
+    args.llm_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained(
+    args.llm_path, trust_remote_code=True)
+model.language_model = llm
+model.config.llm_config = llm.config
+model.to(torch.bfloat16)
+
+output_path = args.model_path + '_replace_llm'
+model.save_pretrained(output_path)
+tokenizer.save_pretrained(output_path)
+print('finished')

InternVL/internvl_chat/tools/resize_pos_embed.py

@@ -0,0 +1,25 @@
+import argparse
+
+import torch
+from internvl.model.internvl_chat import InternVLChatModel
+from transformers import AutoTokenizer
+
+argparse = argparse.ArgumentParser()
+argparse.add_argument('model_path', type=str, default='')
+argparse.add_argument('output_path', type=str, default='')
+argparse.add_argument('force_image_size', type=int, default=448)
+
+args = argparse.parse_args()
+
+model = InternVLChatModel.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+model.vision_model.resize_pos_embeddings(old_size=model.config.vision_config.image_size,
+                                         new_size=args.force_image_size,
+                                         patch_size=14)
+model.config.vision_config.image_size = args.force_image_size
+model.config.force_image_size = args.force_image_size
+
+model.save_pretrained(args.output_path)
+
+tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+tokenizer.save_pretrained(args.output_path)
+print('finished')

InternVL/internvl_chat_llava/llava/model/language_model/llava_llama.py

@@ -0,0 +1,140 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+
+
+class LlavaConfig(LlamaConfig):
+    model_type = "llava_llama"
+
+
+class LlavaLlamaModel(LlavaMetaModel, LlamaModel):
+    config_class = LlavaConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(LlavaLlamaModel, self).__init__(config)
+
+
+class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
+    config_class = LlavaConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = LlavaLlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model/pipeline 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.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": kwargs.get("images", None),
+            }
+        )
+        return model_inputs
+
+AutoConfig.register("llava_llama", LlavaConfig)
+AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)

InternVL/internvl_chat_llava/llava/model/language_model/llava_mpt.py

@@ -0,0 +1,97 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import Optional, Tuple
+
+import torch
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+    MptConfig, MptForCausalLM, MptModel
+from llava.model.llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
+
+
+class LlavaMptConfig(MptConfig):
+    model_type = "llava_mpt"
+
+
+class LlavaMptModel(LlavaMetaModel, MptModel):
+    config_class = LlavaMptConfig
+
+    def __init__(self, config: MptConfig):
+        config.hidden_size = config.d_model
+        super(LlavaMptModel, self).__init__(config)
+
+    def embed_tokens(self, x):
+        return self.wte(x)
+
+
+class LlavaMptForCausalLM(MptForCausalLM, LlavaMetaForCausalLM):
+    config_class = LlavaMptConfig
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config):
+        super(MptForCausalLM, self).__init__(config)
+
+        self.transformer = LlavaMptModel(config)
+        self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.transformer
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, LlavaMptModel):
+            module.gradient_checkpointing = value
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            labels: Optional[torch.Tensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            images=None):
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(
+            input_ids, attention_mask, past_key_values, labels, images)
+
+        return super().forward(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        images = kwargs.pop("images", None)
+        _inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
+        )
+        _inputs['images'] = images
+        return _inputs
+
+
+AutoConfig.register("llava_mpt", LlavaMptConfig)
+AutoModelForCausalLM.register(LlavaMptConfig, LlavaMptForCausalLM)

InternVL/internvl_chat_llava/llava/model/language_model/mpt/adapt_tokenizer.py

@@ -0,0 +1,41 @@
+from typing import Union
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+NUM_SENTINEL_TOKENS: int = 100
+
+def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
+    """Adds sentinel tokens and padding token (if missing).
+
+    Expands the tokenizer vocabulary to include sentinel tokens
+    used in mixture-of-denoiser tasks as well as a padding token.
+
+    All added tokens are added as special tokens. No tokens are
+    added if sentinel tokens and padding token already exist.
+    """
+    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
+    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
+    if tokenizer.pad_token is None:
+        tokenizer.add_tokens('<pad>', special_tokens=True)
+        tokenizer.pad_token = '<pad>'
+        assert tokenizer.pad_token_id is not None
+    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
+    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
+    tokenizer.sentinel_token_ids = _sentinel_token_ids
+
+class AutoTokenizerForMOD(AutoTokenizer):
+    """AutoTokenizer + Adaptation for MOD.
+
+    A simple wrapper around AutoTokenizer to make instantiating
+    an MOD-adapted tokenizer a bit easier.
+
+    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),
+    a padding token, and a property to get the token ids of the
+    sentinel tokens.
+    """
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """See `AutoTokenizer.from_pretrained` docstring."""
+        tokenizer = super().from_pretrained(*args, **kwargs)
+        adapt_tokenizer_for_denoising(tokenizer)
+        return tokenizer

InternVL/internvl_chat_llava/llava/model/language_model/mpt/attention.py

@@ -0,0 +1,300 @@
+"""Attention layers."""
+import math
+import warnings
+from typing import Optional
+import torch
+import torch.nn as nn
+from einops import rearrange
+from packaging import version
+from torch import nn
+from .norm import LPLayerNorm
+
+def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
+    if original_is_causal and num_query_tokens != num_key_tokens:
+        if num_query_tokens != 1:
+            raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
+        else:
+            return False
+    return original_is_causal
+
+def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
+    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
+    kv_n_heads = 1 if multiquery else n_heads
+    k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
+    v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
+    if past_key_value is not None:
+        if len(past_key_value) != 0:
+            k = torch.cat([past_key_value[0], k], dim=3)
+            v = torch.cat([past_key_value[1], v], dim=2)
+        past_key_value = (k, v)
+    (b, _, s_q, d) = q.shape
+    s_k = k.size(-1)
+    if softmax_scale is None:
+        softmax_scale = 1 / math.sqrt(d)
+    attn_weight = q.matmul(k) * softmax_scale
+    if attn_bias is not None:
+        _s_q = max(0, attn_bias.size(2) - s_q)
+        _s_k = max(0, attn_bias.size(3) - s_k)
+        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
+        if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q):
+            raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
+        attn_weight = attn_weight + attn_bias
+    min_val = torch.finfo(q.dtype).min
+    if key_padding_mask is not None:
+        if attn_bias is not None:
+            warnings.warn('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unneccessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
+        attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
+    if is_causal and (not q.size(2) == 1):
+        s = max(s_q, s_k)
+        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
+        causal_mask = causal_mask.tril()
+        causal_mask = causal_mask.to(torch.bool)
+        causal_mask = ~causal_mask
+        causal_mask = causal_mask[-s_q:, -s_k:]
+        attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    if dropout_p:
+        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
+    out = attn_weight.to(v.dtype).matmul(v)
+    out = rearrange(out, 'b h s d -> b s (h d)')
+    if needs_weights:
+        return (out, attn_weight, past_key_value)
+    return (out, None, past_key_value)
+
+def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
+    for tensor in tensors:
+        if tensor.dtype not in valid_dtypes:
+            raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
+        if not tensor.is_cuda:
+            raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
+
+def flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
+    try:
+        from flash_attn import bert_padding, flash_attn_interface
+    except:
+        raise RuntimeError('Please install flash-attn==1.0.3.post0')
+    check_valid_inputs(query, key, value)
+    if past_key_value is not None:
+        if len(past_key_value) != 0:
+            key = torch.cat([past_key_value[0], key], dim=1)
+            value = torch.cat([past_key_value[1], value], dim=1)
+        past_key_value = (key, value)
+    if attn_bias is not None:
+        _s_q = max(0, attn_bias.size(2) - query.size(1))
+        _s_k = max(0, attn_bias.size(3) - key.size(1))
+        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
+    if attn_bias is not None:
+        raise NotImplementedError(f'attn_bias not implemented for flash attn.')
+    (batch_size, seqlen) = query.shape[:2]
+    if key_padding_mask is None:
+        key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
+    query_padding_mask = key_padding_mask[:, -query.size(1):]
+    (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
+    query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
+    (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
+    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
+    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    if multiquery:
+        key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
+        value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
+    dropout_p = dropout_p if training else 0.0
+    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
+    output_unpad = flash_attn_interface.flash_attn_unpadded_func(query_unpad, key_unpad, value_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
+    output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
+    return (output, None, past_key_value)
+
+def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
+    try:
+        from .flash_attn_triton import flash_attn_func
+    except:
+        _installed = False
+        if version.parse(torch.__version__) < version.parse('2.0.0'):
+            _installed = True
+            try:
+                from flash_attn.flash_attn_triton import flash_attn_func
+            except:
+                _installed = False
+        if not _installed:
+            raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU and `pip install .[gpu]` if installing from llm-foundry source or `pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). Note: (1) requires you have CMake and PyTorch already installed.')
+    check_valid_inputs(query, key, value)
+    if past_key_value is not None:
+        if len(past_key_value) != 0:
+            key = torch.cat([past_key_value[0], key], dim=1)
+            value = torch.cat([past_key_value[1], value], dim=1)
+        past_key_value = (key, value)
+    if attn_bias is not None:
+        _s_q = max(0, attn_bias.size(2) - query.size(1))
+        _s_k = max(0, attn_bias.size(3) - key.size(1))
+        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
+    if dropout_p:
+        raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
+    if needs_weights:
+        raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
+    if key_padding_mask is not None:
+        warnings.warn('Propagating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unnecessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
+        (b_size, s_k) = key_padding_mask.shape[:2]
+        if attn_bias is None:
+            attn_bias = query.new_zeros(b_size, 1, 1, s_k)
+        attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
+    query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
+    key = rearrange(key, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
+    value = rearrange(value, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
+    if multiquery:
+        key = key.expand(*key.shape[:2], n_heads, key.size(-1))
+        value = value.expand(*value.shape[:2], n_heads, value.size(-1))
+    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
+    attn_output = flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
+    output = attn_output.view(*attn_output.shape[:2], -1)
+    return (output, None, past_key_value)
+
+class MultiheadAttention(nn.Module):
+    """Multi-head self attention.
+
+    Using torch or triton attention implementation enables user to also use
+    additive bias.
+    """
+
+    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
+        super().__init__()
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.qk_ln = qk_ln
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.softmax_scale = softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
+        self.attn_dropout_p = attn_pdrop
+        self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)
+        fuse_splits = (d_model, 2 * d_model)
+        self.Wqkv._fused = (0, fuse_splits)
+        if self.qk_ln:
+            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
+            self.q_ln = layernorm_class(self.d_model, device=device)
+            self.k_ln = layernorm_class(self.d_model, device=device)
+        if self.attn_impl == 'flash':
+            self.attn_fn = flash_attn_fn
+        elif self.attn_impl == 'triton':
+            self.attn_fn = triton_flash_attn_fn
+            if verbose:
+                warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
+        elif self.attn_impl == 'torch':
+            self.attn_fn = scaled_multihead_dot_product_attention
+            if torch.cuda.is_available() and verbose:
+                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
+        else:
+            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
+        self.out_proj._is_residual = True
+
+    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+        qkv = self.Wqkv(x)
+        if self.clip_qkv:
+            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
+        (query, key, value) = qkv.chunk(3, dim=2)
+        key_padding_mask = attention_mask
+        if self.qk_ln:
+            dtype = query.dtype
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)
+        return (self.out_proj(context), attn_weights, past_key_value)
+
+class MultiQueryAttention(nn.Module):
+    """Multi-Query self attention.
+
+    Using torch or triton attention implementation enables user to also use
+    additive bias.
+    """
+
+    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, verbose: int=0, device: Optional[str]=None):
+        super().__init__()
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.qk_ln = qk_ln
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.softmax_scale = softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.head_dim)
+        self.attn_dropout_p = attn_pdrop
+        self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
+        fuse_splits = (d_model, d_model + self.head_dim)
+        self.Wqkv._fused = (0, fuse_splits)
+        if self.qk_ln:
+            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
+            self.q_ln = layernorm_class(d_model, device=device)
+            self.k_ln = layernorm_class(self.head_dim, device=device)
+        if self.attn_impl == 'flash':
+            self.attn_fn = flash_attn_fn
+        elif self.attn_impl == 'triton':
+            self.attn_fn = triton_flash_attn_fn
+            if verbose:
+                warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
+        elif self.attn_impl == 'torch':
+            self.attn_fn = scaled_multihead_dot_product_attention
+            if torch.cuda.is_available() and verbose:
+                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
+        else:
+            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
+        self.out_proj._is_residual = True
+
+    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+        qkv = self.Wqkv(x)
+        if self.clip_qkv:
+            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
+        (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
+        key_padding_mask = attention_mask
+        if self.qk_ln:
+            dtype = query.dtype
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights, multiquery=True)
+        return (self.out_proj(context), attn_weights, past_key_value)
+
+def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
+    if attn_impl == 'flash':
+        return None
+    elif attn_impl in ['torch', 'triton']:
+        if alibi:
+            if (prefix_lm or not causal) or use_sequence_id:
+                return (1, n_heads, seq_len, seq_len)
+            return (1, n_heads, 1, seq_len)
+        elif prefix_lm or use_sequence_id:
+            return (1, 1, seq_len, seq_len)
+        return None
+    else:
+        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+
+def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
+    if attn_impl == 'flash':
+        return None
+    elif attn_impl in ['torch', 'triton']:
+        if alibi:
+            (device, dtype) = (attn_bias.device, attn_bias.dtype)
+            attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=not causal, alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))
+        return attn_bias
+    else:
+        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+
+def gen_slopes(n_heads, alibi_bias_max=8, device=None):
+    _n_heads = 2 ** math.ceil(math.log2(n_heads))
+    m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
+    m = m.mul(alibi_bias_max / _n_heads)
+    slopes = 1.0 / torch.pow(2, m)
+    if _n_heads != n_heads:
+        slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
+    return slopes.view(1, n_heads, 1, 1)
+
+def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
+    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
+    if full:
+        alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
+        alibi_bias = alibi_bias.abs().mul(-1)
+    slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
+    alibi_bias = alibi_bias * slopes
+    return alibi_bias.to(dtype=dtype)
+ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}

InternVL/internvl_chat_llava/llava/model/language_model/mpt/blocks.py

@@ -0,0 +1,41 @@
+"""GPT Blocks used for the GPT Model."""
+from typing import Dict, Optional, Tuple
+import torch
+import torch.nn as nn
+from .attention import ATTN_CLASS_REGISTRY
+from .norm import NORM_CLASS_REGISTRY
+
+class MPTMLP(nn.Module):
+
+    def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
+        super().__init__()
+        self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)
+        self.act = nn.GELU(approximate='none')
+        self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)
+        self.down_proj._is_residual = True
+
+    def forward(self, x):
+        return self.down_proj(self.act(self.up_proj(x)))
+
+class MPTBlock(nn.Module):
+
+    def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', verbose: int=0, device: Optional[str]=None, **kwargs):
+        del kwargs
+        super().__init__()
+        norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
+        attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
+        self.norm_1 = norm_class(d_model, device=device)
+        self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, verbose=verbose, device=device)
+        self.norm_2 = norm_class(d_model, device=device)
+        self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
+        self.resid_attn_dropout = nn.Dropout(resid_pdrop)
+        self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
+
+    def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
+        a = self.norm_1(x)
+        (b, attn_weights, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)
+        x = x + self.resid_attn_dropout(b)
+        m = self.norm_2(x)
+        n = self.ffn(m)
+        x = x + self.resid_ffn_dropout(n)
+        return (x, attn_weights, past_key_value)

InternVL/internvl_chat_llava/llava/model/language_model/mpt/configuration_mpt.py

@@ -0,0 +1,118 @@
+"""A HuggingFace-style model configuration."""
+from typing import Dict, Optional, Union
+from transformers import PretrainedConfig
+attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
+init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
+
+class MPTConfig(PretrainedConfig):
+    model_type = 'mpt'
+
+    def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, verbose: int=0, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, **kwargs):
+        """The MPT configuration class.
+
+        Args:
+            d_model (int): The size of the embedding dimension of the model.
+            n_heads (int): The number of attention heads.
+            n_layers (int): The number of layers in the model.
+            expansion_ratio (int): The ratio of the up/down scale in the MLP.
+            max_seq_len (int): The maximum sequence length of the model.
+            vocab_size (int): The size of the vocabulary.
+            resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
+            emb_pdrop (float): The dropout probability for the embedding layer.
+            learned_pos_emb (bool): Whether to use learned positional embeddings
+            attn_config (Dict):  A dictionary used to configure the model's attention module:
+                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
+                attn_pdrop (float): The dropout probability for the attention layers.
+                attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
+                qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
+                clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
+                    this value.
+                softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
+                    use the default scale of ``1/sqrt(d_keys)``.
+                prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
+                    extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
+                    can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
+                attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
+                    When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
+                    which sub-sequence each token belongs to.
+                    Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
+                alibi (bool): Whether to use the alibi bias instead of position embeddings.
+                alibi_bias_max (int): The maximum value of the alibi bias.
+            init_device (str): The device to use for parameter initialization.
+            logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
+            no_bias (bool): Whether to use bias in all layers.
+            verbose (int): The verbosity level. 0 is silent.
+            embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
+            norm_type (str): choose type of norm to use
+            multiquery_attention (bool): Whether to use multiquery attention implementation.
+            use_cache (bool): Whether or not the model should return the last key/values attentions
+            init_config (Dict): A dictionary used to configure the model initialization:
+                init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
+                    'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or
+                    'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.
+                init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
+                emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
+                emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
+                    used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
+                init_std (float): The standard deviation of the normal distribution used to initialize the model,
+                    if using the baseline_ parameter initialization scheme.
+                init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
+                fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
+                init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
+                ---
+                See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
+        """
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.expansion_ratio = expansion_ratio
+        self.max_seq_len = max_seq_len
+        self.vocab_size = vocab_size
+        self.resid_pdrop = resid_pdrop
+        self.emb_pdrop = emb_pdrop
+        self.learned_pos_emb = learned_pos_emb
+        self.attn_config = attn_config
+        self.init_device = init_device
+        self.logit_scale = logit_scale
+        self.no_bias = no_bias
+        self.verbose = verbose
+        self.embedding_fraction = embedding_fraction
+        self.norm_type = norm_type
+        self.use_cache = use_cache
+        self.init_config = init_config
+        if 'name' in kwargs:
+            del kwargs['name']
+        if 'loss_fn' in kwargs:
+            del kwargs['loss_fn']
+        super().__init__(**kwargs)
+        self._validate_config()
+
+    def _set_config_defaults(self, config, config_defaults):
+        for (k, v) in config_defaults.items():
+            if k not in config:
+                config[k] = v
+        return config
+
+    def _validate_config(self):
+        self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
+        self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
+        if self.d_model % self.n_heads != 0:
+            raise ValueError('d_model must be divisible by n_heads')
+        if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
+            raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
+        if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
+            raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
+        if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
+        if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('alibi only implemented with torch and triton attention.')
+        if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
+        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
+            raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
+        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
+            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+        if self.init_config.get('name', None) is None:
+            raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
+        if not self.learned_pos_emb and (not self.attn_config['alibi']):
+            raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')

InternVL/internvl_chat_llava/llava/model/language_model/mpt/custom_embedding.py

@@ -0,0 +1,11 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+class SharedEmbedding(nn.Embedding):
+
+    def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
+        if unembed:
+            return F.linear(input, self.weight)
+        return super().forward(input)

InternVL/internvl_chat_llava/llava/model/language_model/mpt/flash_attn_triton.py

@@ -0,0 +1,484 @@
+"""
+Copied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py
+update imports to use 'triton_pre_mlir'
+
+*Experimental* implementation of FlashAttention in Triton.
+Tested with triton==2.0.0.dev20221202.
+Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions
+other than 64:
+https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207
+We'll update this implementation with the new Triton backend once this is fixed.
+
+We use the FlashAttention implementation from Phil Tillet a starting point.
+https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
+
+Changes:
+- Implement both causal and non-causal attention.
+- Implement both self-attention and cross-attention.
+- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.
+- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.
+- Support attention bias.
+- Speed up the forward pass a bit, and only store the LSE instead of m and l.
+- Make the backward for d=128 much faster by reducing register spilling.
+- Optionally parallelize the backward pass across seqlen_k, to deal with the case of
+small batch size * nheads.
+
+Caution:
+- This is an *experimental* implementation. The forward pass should be quite robust but
+I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler).
+- This implementation has only been tested on A100.
+- If you plan to use headdim other than 64 and 128, you should test for race conditions
+(due to the Triton compiler), as done in tests/test_flash_attn.py
+"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
+for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
+that there are none left for other head dimensions.
+
+Differences between this Triton version and the CUDA version:
+- Triton version doesn't support dropout.
+- Triton forward is generally faster than CUDA forward, while Triton backward is
+generally slower than CUDA backward. Overall Triton forward + backward is slightly slower
+than CUDA forward + backward.
+- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
+- Triton version supports attention bias, while CUDA version doesn't.
+"""
+import math
+import torch
+import triton_pre_mlir as triton
+import triton_pre_mlir.language as tl
+
+@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0, 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0, 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM']})
+@triton.jit
+def _fwd_kernel(Q, K, V, Bias, Out, Lse, TMP, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_ob, stride_oh, stride_om, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
+    start_m = tl.program_id(0)
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :])
+    k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :])
+    v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :])
+    if BIAS_TYPE == 'vector':
+        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
+    elif BIAS_TYPE == 'matrix':
+        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :])
+    t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
+    lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
+    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
+    if EVEN_M & EVEN_N:
+        if EVEN_HEADDIM:
+            q = tl.load(q_ptrs)
+        else:
+            q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+    elif EVEN_HEADDIM:
+        q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
+    else:
+        q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
+    end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
+    for start_n in range(0, end_n, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        if EVEN_N & EVEN_M:
+            if EVEN_HEADDIM:
+                k = tl.load(k_ptrs + start_n * stride_kn)
+            else:
+                k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0)
+        elif EVEN_HEADDIM:
+            k = tl.load(k_ptrs + start_n * stride_kn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0)
+        else:
+            k = tl.load(k_ptrs + start_n * stride_kn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k, trans_b=True)
+        if not EVEN_N:
+            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float('-inf'))
+        if IS_CAUSAL:
+            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float('-inf'))
+        if BIAS_TYPE != 'none':
+            if BIAS_TYPE == 'vector':
+                if EVEN_N:
+                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs + start_n, mask=start_n + offs_n < seqlen_k, other=0.0).to(tl.float32)
+                bias = bias[None, :]
+            elif BIAS_TYPE == 'matrix':
+                if EVEN_M & EVEN_N:
+                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs + start_n, mask=(offs_m[:, None] < seqlen_q) & ((start_n + offs_n)[None, :] < seqlen_k), other=0.0).to(tl.float32)
+            qk = qk * softmax_scale + bias
+            m_ij = tl.maximum(tl.max(qk, 1), lse_i)
+            p = tl.exp(qk - m_ij[:, None])
+        else:
+            m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
+            p = tl.exp(qk * softmax_scale - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+        acc_o_scale = tl.exp(m_i - m_ij)
+        tl.store(t_ptrs, acc_o_scale)
+        acc_o_scale = tl.load(t_ptrs)
+        acc_o = acc_o * acc_o_scale[:, None]
+        if EVEN_N & EVEN_M:
+            if EVEN_HEADDIM:
+                v = tl.load(v_ptrs + start_n * stride_vn)
+            else:
+                v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0)
+        elif EVEN_HEADDIM:
+            v = tl.load(v_ptrs + start_n * stride_vn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0)
+        else:
+            v = tl.load(v_ptrs + start_n * stride_vn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
+        p = p.to(v.dtype)
+        acc_o += tl.dot(p, v)
+        m_i = m_ij
+        l_i_new = tl.exp(lse_i - m_ij) + l_ij
+        lse_i = m_ij + tl.log(l_i_new)
+    o_scale = tl.exp(m_i - lse_i)
+    tl.store(t_ptrs, o_scale)
+    o_scale = tl.load(t_ptrs)
+    acc_o = acc_o * o_scale[:, None]
+    start_m = tl.program_id(0)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
+    tl.store(lse_ptrs, lse_i)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:, None] * stride_om + offs_d[None, :])
+    if EVEN_M:
+        if EVEN_HEADDIM:
+            tl.store(out_ptrs, acc_o)
+        else:
+            tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)
+    elif EVEN_HEADDIM:
+        tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
+    else:
+        tl.store(out_ptrs, acc_o, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
+
+@triton.jit
+def _bwd_preprocess_do_o_dot(Out, DO, Delta, stride_ob, stride_oh, stride_om, stride_dob, stride_doh, stride_dom, nheads, seqlen_q, seqlen_q_rounded, headdim, BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr):
+    start_m = tl.program_id(0)
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    o = tl.load(Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
+    do = tl.load(DO + off_b * stride_dob + off_h * stride_doh + offs_m[:, None] * stride_dom + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
+
+@triton.jit
+def _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr):
+    if EVEN_N & EVEN_M:
+        if EVEN_HEADDIM:
+            tl.store(dv_ptrs, dv)
+            tl.store(dk_ptrs, dk)
+        else:
+            tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
+            tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
+    elif EVEN_HEADDIM:
+        tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
+        tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
+    else:
+        tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
+        tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
+
+@triton.jit
+def _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD: tl.constexpr, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
+    begin_m = 0 if not IS_CAUSAL else start_n * BLOCK_N // BLOCK_M * BLOCK_M
+    offs_qm = begin_m + tl.arange(0, BLOCK_M)
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
+    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
+    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
+    do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
+    dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
+    if BIAS_TYPE == 'vector':
+        b_ptrs = Bias + offs_n
+    elif BIAS_TYPE == 'matrix':
+        b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
+    dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
+    if begin_m >= seqlen_q:
+        dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
+        dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
+        _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
+        return
+    if EVEN_N & EVEN_M:
+        if EVEN_HEADDIM:
+            k = tl.load(k_ptrs)
+            v = tl.load(v_ptrs)
+        else:
+            k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+            v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+    elif EVEN_HEADDIM:
+        k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
+        v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
+    else:
+        k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
+        v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0)
+    num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
+    for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
+        start_m = tl.multiple_of(start_m, BLOCK_M)
+        offs_m_curr = start_m + offs_m
+        if EVEN_M & EVEN_HEADDIM:
+            q = tl.load(q_ptrs)
+        elif EVEN_HEADDIM:
+            q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
+        else:
+            q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
+        qk = tl.dot(q, k, trans_b=True)
+        if not EVEN_N:
+            qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf'))
+        if IS_CAUSAL:
+            qk = tl.where(offs_m_curr[:, None] >= offs_n[None, :], qk, float('-inf'))
+        if BIAS_TYPE != 'none':
+            tl.debug_barrier()
+            if BIAS_TYPE == 'vector':
+                if EVEN_N:
+                    bias = tl.load(b_ptrs).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32)
+                bias = bias[None, :]
+            elif BIAS_TYPE == 'matrix':
+                if EVEN_M & EVEN_N:
+                    bias = tl.load(b_ptrs).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k), other=0.0).to(tl.float32)
+            qk = qk * softmax_scale + bias
+        if not EVEN_M & EVEN_HEADDIM:
+            tl.debug_barrier()
+        lse_i = tl.load(LSE + offs_m_curr)
+        if BIAS_TYPE == 'none':
+            p = tl.exp(qk * softmax_scale - lse_i[:, None])
+        else:
+            p = tl.exp(qk - lse_i[:, None])
+        if EVEN_M & EVEN_HEADDIM:
+            do = tl.load(do_ptrs)
+        else:
+            do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0)
+        dv += tl.dot(p.to(do.dtype), do, trans_a=True)
+        if not EVEN_M & EVEN_HEADDIM:
+            tl.debug_barrier()
+        dp = tl.dot(do, v, trans_b=True)
+        if not EVEN_HEADDIM:
+            tl.debug_barrier()
+        Di = tl.load(D + offs_m_curr)
+        ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
+        dk += tl.dot(ds, q, trans_a=True)
+        if not EVEN_M & EVEN_HEADDIM:
+            tl.debug_barrier()
+        if not ATOMIC_ADD:
+            if EVEN_M & EVEN_HEADDIM:
+                dq = tl.load(dq_ptrs, eviction_policy='evict_last')
+                dq += tl.dot(ds, k)
+                tl.store(dq_ptrs, dq, eviction_policy='evict_last')
+            elif EVEN_HEADDIM:
+                dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0, eviction_policy='evict_last')
+                dq += tl.dot(ds, k)
+                tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q, eviction_policy='evict_last')
+            else:
+                dq = tl.load(dq_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0, eviction_policy='evict_last')
+                dq += tl.dot(ds, k)
+                tl.store(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), eviction_policy='evict_last')
+        else:
+            dq = tl.dot(ds, k)
+            if EVEN_M & EVEN_HEADDIM:
+                tl.atomic_add(dq_ptrs, dq)
+            elif EVEN_HEADDIM:
+                tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q)
+            else:
+                tl.atomic_add(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
+        dq_ptrs += BLOCK_M * stride_dqm
+        q_ptrs += BLOCK_M * stride_qm
+        do_ptrs += BLOCK_M * stride_dom
+        if BIAS_TYPE == 'matrix':
+            b_ptrs += BLOCK_M * stride_bm
+    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
+    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
+    _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
+
+def init_to_zero(name):
+    return lambda nargs: nargs[name].zero_()
+
+@triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ'))], key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'])
+@triton.heuristics({'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0, 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0, 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM']})
+@triton.jit
+def _bwd_kernel(Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_dob, stride_doh, stride_dom, stride_dqb, stride_dqh, stride_dqm, stride_dkb, stride_dkh, stride_dkn, stride_dvb, stride_dvh, stride_dvn, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, SEQUENCE_PARALLEL: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    Q += off_b * stride_qb + off_h * stride_qh
+    K += off_b * stride_kb + off_h * stride_kh
+    V += off_b * stride_vb + off_h * stride_vh
+    DO += off_b * stride_dob + off_h * stride_doh
+    DQ += off_b * stride_dqb + off_h * stride_dqh
+    DK += off_b * stride_dkb + off_h * stride_dkh
+    DV += off_b * stride_dvb + off_h * stride_dvh
+    if BIAS_TYPE != 'none':
+        Bias += off_b * stride_bb + off_h * stride_bh
+    D += off_hb * seqlen_q_rounded
+    LSE += off_hb * seqlen_q_rounded
+    if not SEQUENCE_PARALLEL:
+        num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
+        for start_n in range(0, num_block_n):
+            _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=False, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
+    else:
+        start_n = tl.program_id(0)
+        _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=True, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N)
+
+def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
+    (batch, seqlen_q, nheads, d) = q.shape
+    (_, seqlen_k, _, _) = k.shape
+    assert k.shape == (batch, seqlen_k, nheads, d)
+    assert v.shape == (batch, seqlen_k, nheads, d)
+    assert d <= 128, 'FlashAttention only support head dimensions up to 128'
+    assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
+    assert q.dtype in [torch.float16, torch.bfloat16], 'Only support fp16 and bf16'
+    assert q.is_cuda and k.is_cuda and v.is_cuda
+    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
+    has_bias = bias is not None
+    bias_type = 'none'
+    if has_bias:
+        assert bias.dtype in [q.dtype, torch.float]
+        assert bias.is_cuda
+        assert bias.dim() == 4
+        if bias.stride(-1) != 1:
+            bias = bias.contiguous()
+        if bias.shape[2:] == (1, seqlen_k):
+            bias_type = 'vector'
+        elif bias.shape[2:] == (seqlen_q, seqlen_k):
+            bias_type = 'matrix'
+        else:
+            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
+        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
+    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
+    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
+    lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
+    tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
+    o = torch.empty_like(q)
+    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
+    BLOCK = 128
+    num_warps = 4 if d <= 64 else 8
+    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
+    _fwd_kernel[grid](q, k, v, bias, o, lse, tmp, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, o.stride(0), o.stride(2), o.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM, BLOCK_M=BLOCK, BLOCK_N=BLOCK, num_warps=num_warps, num_stages=1)
+    return (o, lse, softmax_scale)
+
+def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None):
+    if do.stride(-1) != 1:
+        do = do.contiguous()
+    (batch, seqlen_q, nheads, d) = q.shape
+    (_, seqlen_k, _, _) = k.shape
+    assert d <= 128
+    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
+    assert lse.shape == (batch, nheads, seqlen_q_rounded)
+    assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1
+    assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
+    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
+    dq_accum = torch.empty_like(q, dtype=torch.float32)
+    delta = torch.empty_like(lse)
+    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
+    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
+    _bwd_preprocess_do_o_dot[grid](o, do, delta, o.stride(0), o.stride(2), o.stride(1), do.stride(0), do.stride(2), do.stride(1), nheads, seqlen_q, seqlen_q_rounded, d, BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM)
+    has_bias = bias is not None
+    bias_type = 'none'
+    if has_bias:
+        assert bias.dtype in [q.dtype, torch.float]
+        assert bias.is_cuda
+        assert bias.dim() == 4
+        assert bias.stride(-1) == 1
+        if bias.shape[2:] == (1, seqlen_k):
+            bias_type = 'vector'
+        elif bias.shape[2:] == (seqlen_q, seqlen_k):
+            bias_type = 'matrix'
+        else:
+            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)')
+        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
+    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
+    grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N']) if META['SEQUENCE_PARALLEL'] else 1, batch * nheads)
+    _bwd_kernel[grid](q, k, v, bias, do, dq_accum, dk, dv, lse, delta, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, do.stride(0), do.stride(2), do.stride(1), dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1), dk.stride(0), dk.stride(2), dk.stride(1), dv.stride(0), dv.stride(2), dv.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM)
+    dq.copy_(dq_accum)
+
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
+        """
+            qkv: (batch, seqlen, 3, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
+        """
+        if qkv.stride(-1) != 1:
+            qkv = qkv.contiguous()
+        (o, lse, ctx.softmax_scale) = _flash_attn_forward(qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal, softmax_scale=softmax_scale)
+        ctx.save_for_backward(qkv, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        (qkv, o, lse, bias) = ctx.saved_tensors
+        assert not ctx.needs_input_grad[1], 'FlashAttention does not support bias gradient yet'
+        with torch.inference_mode():
+            dqkv = torch.empty_like(qkv)
+            _flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return (dqkv, None, None, None)
+flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply
+
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
+        """
+            q: (batch, seqlen_q, nheads, headdim)
+            kv: (batch, seqlen_k, 2, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
+        """
+        (q, kv) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]]
+        (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale)
+        ctx.save_for_backward(q, kv, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        (q, kv, o, lse, bias) = ctx.saved_tensors
+        if len(ctx.needs_input_grad) >= 3:
+            assert not ctx.needs_input_grad[2], 'FlashAttention does not support bias gradient yet'
+        with torch.inference_mode():
+            dq = torch.empty_like(q)
+            dkv = torch.empty_like(kv)
+            _flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse, dq, dkv[:, :, 0], dkv[:, :, 1], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return (dq, dkv, None, None, None)
+flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply
+
+class FlashAttnFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
+        """
+            q: (batch_size, seqlen_q, nheads, headdim)
+            k, v: (batch_size, seqlen_k, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
+        """
+        (q, k, v) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]]
+        (o, lse, ctx.softmax_scale) = _flash_attn_forward(q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
+        ctx.save_for_backward(q, k, v, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        (q, k, v, o, lse, bias) = ctx.saved_tensors
+        assert not ctx.needs_input_grad[3], 'FlashAttention does not support bias gradient yet'
+        with torch.inference_mode():
+            dq = torch.empty_like(q)
+            dk = torch.empty_like(k)
+            dv = torch.empty_like(v)
+            _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return (dq, dk, dv, None, None, None)
+flash_attn_func = FlashAttnFunc.apply

InternVL/internvl_chat_llava/llava/model/language_model/mpt/hf_prefixlm_converter.py

@@ -0,0 +1,415 @@
+"""Converts Huggingface Causal LM to Prefix LM.
+
+Conversion does lightweight surgery on a HuggingFace
+Causal LM to convert it to a Prefix LM.
+
+Prefix LMs accepts a `bidirectional_mask` input in `forward`
+and treat the input prompt as the prefix in `generate`.
+"""
+import math
+import warnings
+from types import MethodType
+from typing import Any, Dict, List, Optional, Tuple, Union
+import torch
+from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
+from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
+from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
+from transformers.models.bloom.modeling_bloom import logging
+from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
+from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
+from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
+from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
+from transformers.models.opt.modeling_opt import OPTForCausalLM
+from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
+from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
+logger = logging.get_logger(__name__)
+_SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
+CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
+
+def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
+    """Converts a GPT-style Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `GPT2LMHeadModel`
+        - `GPTNeoForCausalLM`
+        - `GPTNeoXForCausalLM`
+        - `GPTJForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, _SUPPORTED_GPT_MODELS)
+    assert model.config.add_cross_attention == False, 'Only supports GPT-style decoder-only models'
+
+    def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
+        """Helper that gets a list of the model's attention modules.
+
+        Each module has a `bias` buffer used for causal masking. The Prefix LM
+        conversion adds logic to dynamically manipulate these biases to support
+        Prefix LM attention masking.
+        """
+        attn_modules = []
+        if isinstance(model, GPTNeoXForCausalLM):
+            blocks = model.gpt_neox.layers
+        else:
+            blocks = model.transformer.h
+        for block in blocks:
+            if isinstance(model, GPTNeoForCausalLM):
+                if block.attn.attention_type != 'global':
+                    continue
+                attn_module = block.attn.attention
+            elif isinstance(model, GPTNeoXForCausalLM):
+                attn_module = block.attention
+            else:
+                attn_module = block.attn
+            attn_modules.append(attn_module)
+        return attn_modules
+    setattr(model, '_original_forward', getattr(model, 'forward'))
+    setattr(model, '_original_generate', getattr(model, 'generate'))
+
+    def forward(self: CAUSAL_GPT_TYPES, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.FloatTensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
+        """Wraps original forward to enable PrefixLM attention."""
+
+        def call_og_forward():
+            if isinstance(self, GPTNeoXForCausalLM):
+                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+            else:
+                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+        if bidirectional_mask is None:
+            return call_og_forward()
+        assert isinstance(bidirectional_mask, torch.Tensor)
+        attn_modules = _get_attn_modules(model)
+        (b, s) = bidirectional_mask.shape
+        max_length = attn_modules[0].bias.shape[-1]
+        if s > max_length:
+            raise ValueError(f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).')
+        assert s <= max_length
+        if s < max_length:
+            pad = torch.zeros((int(b), int(max_length - s)), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
+            bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
+        bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
+        output = call_og_forward()
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
+        return output
+
+    def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):
+        """Wraps original generate to enable PrefixLM attention."""
+        attn_modules = _get_attn_modules(model)
+        for attn_module in attn_modules:
+            attn_module.bias.data[:] = 1
+        output = self._original_generate(*args, **kwargs)
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
+        return output
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'generate', MethodType(generate, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+
+def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
+    """Converts a BLOOM Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `BloomForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, BloomForCausalLM)
+    assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
+
+    def _prepare_attn_mask(self: BloomModel, attention_mask: torch.Tensor, bidirectional_mask: Optional[torch.Tensor], input_shape: Tuple[int, int], past_key_values_length: int) -> torch.BoolTensor:
+        combined_attention_mask = None
+        device = attention_mask.device
+        (_, src_length) = input_shape
+        if src_length > 1:
+            combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
+            if bidirectional_mask is not None:
+                assert attention_mask.shape == bidirectional_mask.shape
+                expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
+                combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
+        expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
+        combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+        return combined_attention_mask
+
+    def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+        num_heads = self.config.n_head
+        closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+        base = torch.tensor(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)
+        powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)
+        slopes = torch.pow(base, powers)
+        if closest_power_of_2 != num_heads:
+            extra_base = torch.tensor(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32)
+            num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+            extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)
+            slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+        qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)
+        ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)
+        diffs = qa - ka + key_length - query_length
+        diffs = -diffs.abs()
+        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length)
+        alibi = alibi.expand(batch_size, -1, -1, -1).reshape(-1, query_length, key_length)
+        return alibi.to(dtype)
+    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
+
+    def forward(self: BloomModel, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+        if deprecated_arguments.pop('position_ids', False) is not False:
+            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' + 'You can safely ignore passing `position_ids`.', FutureWarning)
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
+        elif input_ids is not None:
+            (batch_size, seq_length) = input_ids.shape
+        elif inputs_embeds is not None:
+            (batch_size, seq_length, _) = inputs_embeds.shape
+        else:
+            raise ValueError('You have to specify either input_ids or inputs_embeds')
+        if past_key_values is None:
+            past_key_values = tuple([None] * len(self.h))
+        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            tmp = past_key_values[0][0]
+            past_key_values_length = tmp.shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        if attention_mask is None:
+            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+        else:
+            attention_mask = attention_mask.to(hidden_states.device)
+        alibi = self._build_alibi_tensor(batch_size=batch_size, query_length=seq_length, key_length=seq_length_with_past, dtype=hidden_states.dtype, device=hidden_states.device)
+        causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
+        for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
+            if output_hidden_states:
+                hst = (hidden_states,)
+                all_hidden_states = all_hidden_states + hst
+            if self.gradient_checkpointing and self.training:
+                if use_cache:
+                    logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
+                    use_cache = False
+
+                def create_custom_forward(module):
+
+                    def custom_forward(*inputs):
+                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
+                    return custom_forward
+                outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
+            else:
+                outputs = block(hidden_states, layer_past=layer_past, attention_mask=causal_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, alibi=alibi)
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+            if output_attentions:
+                oa = (outputs[2 if use_cache else 1],)
+                all_self_attentions = all_self_attentions + oa
+        hidden_states = self.ln_f(hidden_states)
+        if output_hidden_states:
+            hst = (hidden_states,)
+            all_hidden_states = all_hidden_states + hst
+        if not return_dict:
+            return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None))
+        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
+    setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
+    setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
+    setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
+    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
+
+    def forward(self: BloomForCausalLM, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
+        """Replacement forward method for BloomCausalLM."""
+        if deprecated_arguments.pop('position_ids', False) is not False:
+            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed ' + 'in v5.0.0. You can safely ignore passing `position_ids`.', FutureWarning)
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, bidirectional_mask=bidirectional_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            (batch_size, seq_length, vocab_size) = shift_logits.shape
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length))
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithCrossAttentions(loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)
+
+    def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            bidirectional_mask = None
+            if past[0][0].shape[0] == input_ids.shape[0]:
+                past = self._convert_to_bloom_cache(past)
+        else:
+            bidirectional_mask = torch.ones_like(input_ids)
+        return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+
+def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
+    """Converts an OPT Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `OPTForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, OPTForCausalLM)
+    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
+    setattr(model, '_original_forward', getattr(model, 'forward'))
+    setattr(model, '_original_generate', getattr(model, 'generate'))
+    model.model.decoder.bidirectional_mask = None
+
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            if self.bidirectional_mask == 'g':
+                (bsz, src_length) = input_shape
+                combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+            else:
+                combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
+                if self.bidirectional_mask is not None:
+                    assert attention_mask.shape == self.bidirectional_mask.shape
+                    expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
+                    combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
+        if attention_mask is not None:
+            expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
+            combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+        return combined_attention_mask
+    setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
+
+    def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
+
+        def call_og_forward():
+            return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+        if bidirectional_mask is None:
+            return call_og_forward()
+        self.model.decoder.bidirectional_mask = bidirectional_mask
+        try:
+            outputs = call_og_forward()
+        except:
+            self.model.decoder.bidirectional_mask = None
+            raise
+        self.model.decoder.bidirectional_mask = None
+        return outputs
+
+    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
+        """Wraps original generate to enable PrefixLM-style attention."""
+        self.model.decoder.bidirectional_mask = 'g'
+        try:
+            output = self._original_generate(*args, **kwargs)
+        except:
+            self.model.decoder.bidirectional_mask = None
+            raise
+        self.model.decoder.bidirectional_mask = None
+        return output
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'generate', MethodType(generate, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)
+CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
+
+def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
+    """Converts a HuggingFace Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `GPT2LMHeadModel`
+        - `GPTNeoForCausalLM`
+        - `GPTNeoXForCausalLM`
+        - `GPTJForCausalLM`
+        - `BloomForCausalLM`
+        - `OPTForCausalLM`
+
+    Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
+    `generate` method and/or select underlying methods depending on the model class.
+
+    These changes preserve the model API, but add a new input to `forward`: "bidirectional_mask".
+
+    Notes on training:
+        To actually train the converted model as a Prefix LM, training batches will need to indicate
+        the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.
+
+        **This is not a standard input and requires custom layers either within or after your dataloader.**
+
+        In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`
+        such that `batch['labels'][batch['bidirectional_mask'] == 1] == -100`.
+        That is, the prefix portion of the sequence should not generate any loss. Loss should only be
+        generated by the target portion of the sequence.
+
+    Notes on `GPTNeoForCausalLM`:
+        To simplify the implementation, "global" and "local" attention layers are handled differently.
+        For "global" layers, we handle conversion as described above. For "local" layers, which use a
+        causal attention mask within a restricted local window, we do not alter the masking.
+
+    Notes on `forward` method conversion:
+        After conversion, the `forward` method will handle a new input, `bidirectional_mask`,
+        which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions
+        belonging to the prefix (prefix tokens can attend to one another bidirectionally), and
+        0 indicates token positions belonging to the target.
+
+        The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing
+        causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset
+        the causal masks before returning the result.
+
+    Notes on `generate` method conversion:
+        After conversion, the `generate` method will have the same signature but will internally
+        convert all causal masks to be purely bidirectional, call the original `generate` method, and
+        (where appropriate) reset the causal masks before returning the result.
+
+        This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token
+        "prompt" passed to `generate` (which is treated as the prefix) and then sequentially generates
+        each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one
+        another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and
+        previously-generated tokens (also as expected in a Prefix LM).
+
+    To preserve the API, the original methods are renamed to `_original_forward` and
+    `_original_generate`, and replaced with new `forward` and `generate` methods that wrap
+    them, respectively. Although implementation details vary by model class.
+    """
+    if isinstance(model, _SUPPORTED_GPT_MODELS):
+        return _convert_gpt_causal_lm_to_prefix_lm(model)
+    elif isinstance(model, BloomForCausalLM):
+        return _convert_bloom_causal_lm_to_prefix_lm(model)
+    elif isinstance(model, OPTForCausalLM):
+        return _convert_opt_causal_lm_to_prefix_lm(model)
+    else:
+        raise TypeError(f'Cannot convert model to Prefix LM. ' + f'Model does not belong to set of supported HF models:' + f'\n{_SUPPORTED_HF_MODELS}')
+
+def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
+    """Attempts to add bidirectional_mask to batch if missing.
+
+    Raises:
+        KeyError if bidirectional_mask is missing and can't be inferred
+    """
+    if 'bidirectional_mask' not in batch:
+        if batch.get('mode', None) == 'icl_task':
+            batch['bidirectional_mask'] = batch['attention_mask'].clone()
+            for (i, continuation_indices) in enumerate(batch['continuation_indices']):
+                batch['bidirectional_mask'][i, continuation_indices] = 0
+        elif 'labels' in batch and 'attention_mask' in batch:
+            batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], -100)).type_as(batch['attention_mask'])
+        else:
+            raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')

InternVL/internvl_chat_llava/llava/model/language_model/mpt/meta_init_context.py

@@ -0,0 +1,94 @@
+from contextlib import contextmanager
+import torch
+import torch.nn as nn
+
+@contextmanager
+def init_empty_weights(include_buffers: bool=False):
+    """Meta initialization context manager.
+
+    A context manager under which models are initialized with all parameters
+    on the meta device, therefore creating an empty model. Useful when just
+    initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
+            not to also put all buffers on the meta device while initializing.
+
+    Example:
+    ```python
+    import torch.nn as nn
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+
+    <Tip warning={true}>
+
+    Any model created under this context manager has no weights. As such you can't do something like
+    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+    """
+    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
+        yield f
+
+@contextmanager
+def init_on_device(device: torch.device, include_buffers: bool=False):
+    """Device initialization context manager.
+
+    A context manager under which models are initialized with all parameters
+    on the specified device.
+
+    Args:
+        device (`torch.device`): Device to initialize all parameters on.
+        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
+            not to also put all buffers on the meta device while initializing.
+
+    Example:
+    ```python
+    import torch.nn as nn
+
+    with init_on_device(device=torch.device("cuda")):
+        tst = nn.Liner(100, 100)  # on `cuda` device
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+    if include_buffers:
+        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
+    else:
+        tensor_constructors_to_patch = {}
+
+    def patch_tensor_constructor(fn):
+
+        def wrapper(*args, **kwargs):
+            kwargs['device'] = device
+            return fn(*args, **kwargs)
+        return wrapper
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)

InternVL/internvl_chat_llava/llava/model/language_model/mpt/modeling_mpt.py

@@ -0,0 +1,331 @@
+"""A simple, flexible implementation of a GPT model.
+
+Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
+"""
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from .attention import attn_bias_shape, build_attn_bias
+from .blocks import MPTBlock
+from .custom_embedding import SharedEmbedding
+from .norm import NORM_CLASS_REGISTRY
+from .configuration_mpt import MPTConfig
+from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
+from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
+from .meta_init_context import init_empty_weights
+from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
+try:
+    from .flash_attn_triton import flash_attn_func
+except:
+    pass
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+
+class MPTPreTrainedModel(PreTrainedModel):
+    config_class = MPTConfig
+    base_model_prefix = 'model'
+    _no_split_modules = ['MPTBlock']
+
+class MPTModel(MPTPreTrainedModel):
+
+    def __init__(self, config: MPTConfig):
+        config._validate_config()
+        super().__init__(config)
+        self.attn_impl = config.attn_config['attn_impl']
+        self.prefix_lm = config.attn_config['prefix_lm']
+        self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
+        self.alibi = config.attn_config['alibi']
+        self.alibi_bias_max = config.attn_config['alibi_bias_max']
+        if config.init_device == 'mixed':
+            if dist.get_local_rank() == 0:
+                config.init_device = 'cpu'
+            else:
+                config.init_device = 'meta'
+        if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
+            norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
+            raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
+        norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
+        self.embedding_fraction = config.embedding_fraction
+        self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
+        if not self.alibi:
+            self.wpe = torch.nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
+        self.emb_drop = nn.Dropout(config.emb_pdrop)
+        self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
+        self.norm_f = norm_class(config.d_model, device=config.init_device)
+        if config.init_device != 'meta':
+            print(f'You are using config.init_device={config.init_device!r}, but you can also use config.init_device="meta" with Composer + FSDP for fast initialization.')
+            self.apply(self.param_init_fn)
+        self.is_causal = not self.prefix_lm
+        self._attn_bias_initialized = False
+        self.attn_bias = None
+        self.attn_bias_shape = attn_bias_shape(self.attn_impl, config.n_heads, config.max_seq_len, self.alibi, prefix_lm=self.prefix_lm, causal=self.is_causal, use_sequence_id=self.attn_uses_sequence_id)
+        if config.no_bias:
+            for module in self.modules():
+                if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
+                    if config.verbose:
+                        warnings.warn(f'Removing bias ({module.bias}) from {module}.')
+                    module.register_parameter('bias', None)
+        if config.verbose and config.verbose > 2:
+            print(self)
+        if 'verbose' not in self.config.init_config:
+            self.config.init_config['verbose'] = self.config.verbose
+        if self.config.init_config['verbose'] > 1:
+            init_fn_name = self.config.init_config['name']
+            warnings.warn(f'Using {init_fn_name} initialization.')
+        self.gradient_checkpointing = False
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, value):
+        self.wte = value
+
+    @torch.no_grad()
+    def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
+        if not self._attn_bias_initialized:
+            if self.attn_bias_shape:
+                self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
+                self.attn_bias = build_attn_bias(self.attn_impl, self.attn_bias, self.config.n_heads, self.config.max_seq_len, causal=self.is_causal, alibi=self.alibi, alibi_bias_max=self.alibi_bias_max)
+            self._attn_bias_initialized = True
+        if self.attn_impl == 'flash':
+            return (self.attn_bias, attention_mask)
+        if self.attn_bias is not None:
+            self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)
+        attn_bias = self.attn_bias
+        if self.prefix_lm:
+            assert isinstance(attn_bias, torch.Tensor)
+            assert isinstance(prefix_mask, torch.Tensor)
+            attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
+        if self.attn_uses_sequence_id and sequence_id is not None:
+            assert isinstance(attn_bias, torch.Tensor)
+            attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
+        if attention_mask is not None:
+            s_k = attention_mask.shape[-1]
+            if attn_bias is None:
+                attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)
+            else:
+                _s_k = max(0, attn_bias.size(-1) - s_k)
+                attn_bias = attn_bias[:, :, :, _s_k:]
+            if prefix_mask is not None and attention_mask.shape != prefix_mask.shape:
+                raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
+            min_val = torch.finfo(attn_bias.dtype).min
+            attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
+        return (attn_bias, None)
+
+    def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
+        (s_k, s_q) = attn_bias.shape[-2:]
+        if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:
+            raise ValueError('attn_bias does not match the expected shape. ' + f'The last two dimensions should both be {self.config.max_length} ' + f'but are {s_k} and {s_q}.')
+        seq_len = prefix_mask.shape[-1]
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f'prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
+        attn_bias = attn_bias[..., :seq_len, :seq_len]
+        causal = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)).view(1, 1, seq_len, seq_len)
+        prefix = prefix_mask.view(-1, 1, 1, seq_len)
+        cannot_attend = ~torch.logical_or(causal, prefix.bool())
+        min_val = torch.finfo(attn_bias.dtype).min
+        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
+        return attn_bias
+
+    def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
+        seq_len = sequence_id.shape[-1]
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
+        attn_bias = attn_bias[..., :seq_len, :seq_len]
+        cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
+        min_val = torch.finfo(attn_bias.dtype).min
+        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
+        return attn_bias
+
+    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.Tensor]=None):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if attention_mask is not None:
+            attention_mask = attention_mask.bool()
+        if prefix_mask is not None:
+            prefix_mask = prefix_mask.bool()
+        if not return_dict:
+            raise NotImplementedError('return_dict False is not implemented yet for MPT')
+        if output_attentions:
+            if self.attn_impl != 'torch':
+                raise NotImplementedError('output_attentions is not implemented for MPT when using attn_impl `flash` or `triton`.')
+        if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
+            raise NotImplementedError('MPT does not support training with left padding.')
+        if self.prefix_lm and prefix_mask is None:
+            raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
+        if self.training:
+            if self.attn_uses_sequence_id and sequence_id is None:
+                raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
+            elif self.attn_uses_sequence_id is False and sequence_id is not None:
+                warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
+        if input_ids is not None:
+            S = input_ids.size(1)
+            assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
+            tok_emb = self.wte(input_ids)
+        else:
+            assert inputs_embeds is not None
+            assert self.alibi, 'inputs_embeds is not implemented for MPT unless for alibi.'
+            S = inputs_embeds.size(1)
+            tok_emb = inputs_embeds
+        if self.alibi:
+            x = tok_emb
+        else:
+            past_position = 0
+            if past_key_values is not None:
+                if len(past_key_values) != self.config.n_layers:
+                    raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
+                past_position = past_key_values[0][0].size(1)
+                if self.attn_impl == 'torch':
+                    past_position = past_key_values[0][0].size(3)
+            if S + past_position > self.config.max_seq_len:
+                raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length {S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
+            pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
+            if attention_mask is not None:
+                pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
+            pos_emb = self.wpe(pos)
+            x = tok_emb + pos_emb
+        if self.embedding_fraction == 1:
+            x = self.emb_drop(x)
+        else:
+            x_shrunk = x * self.embedding_fraction + x.detach() * (1 - self.embedding_fraction)
+            assert isinstance(self.emb_drop, nn.Module)
+            x = self.emb_drop(x_shrunk)
+        (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=torch.float32, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
+        if use_cache and past_key_values is None:
+            past_key_values = [() for _ in range(self.config.n_layers)]
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        for (b_idx, block) in enumerate(self.blocks):
+            if output_hidden_states:
+                assert all_hidden_states is not None
+                all_hidden_states = all_hidden_states + (x,)
+            past_key_value = past_key_values[b_idx] if past_key_values is not None else None
+            if self.gradient_checkpointing and self.training:
+                (x, attn_weights, past_key_value) = torch.utils.checkpoint.checkpoint(block, x, past_key_value, attn_bias, attention_mask, self.is_causal)
+            else:
+                (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
+            if past_key_values is not None:
+                past_key_values[b_idx] = past_key_value
+            if output_attentions:
+                assert all_self_attns is not None
+                all_self_attns = all_self_attns + (attn_weights,)
+        x = self.norm_f(x)
+        if output_hidden_states:
+            assert all_hidden_states is not None
+            all_hidden_states = all_hidden_states + (x,)
+        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns)
+
+    def param_init_fn(self, module):
+        init_fn_name = self.config.init_config['name']
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def activation_checkpointing_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+class MPTForCausalLM(MPTPreTrainedModel):
+
+    def __init__(self, config: MPTConfig):
+        super().__init__(config)
+        if not config.tie_word_embeddings:
+            raise ValueError('MPTForCausalLM only supports tied word embeddings')
+        print(f'Instantiating an MPTForCausalLM model from {__file__}')
+        self.transformer = MPTModel(config)
+        for child in self.transformer.children():
+            if isinstance(child, torch.nn.ModuleList):
+                continue
+            if isinstance(child, torch.nn.Module):
+                child._fsdp_wrap = True
+        self.logit_scale = None
+        if config.logit_scale is not None:
+            logit_scale = config.logit_scale
+            if isinstance(logit_scale, str):
+                if logit_scale == 'inv_sqrt_d_model':
+                    logit_scale = 1 / math.sqrt(config.d_model)
+                else:
+                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+            self.logit_scale = logit_scale
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value):
+        self.transformer.wte = value
+
+    def get_output_embeddings(self):
+        return self.transformer.wte
+
+    def set_output_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.transformer = decoder
+
+    def get_decoder(self):
+        return self.transformer
+
+    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.FloatTensor]=None):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if inputs_embeds is not None:
+            raise NotImplementedError('inputs_embeds has to be None (for hf/peft support).')
+        outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
+        logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
+        if self.logit_scale is not None:
+            if self.logit_scale == 0:
+                warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
+            logits *= self.logit_scale
+        loss = None
+        if labels is not None:
+            labels = torch.roll(labels, shifts=-1)
+            labels[:, -1] = -100
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
+
+    def param_init_fn(self, module):
+        init_fn_name = self.config.init_config['name']
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def activation_checkpointing_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is not None:
+            raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
+        attention_mask = kwargs['attention_mask'].bool()
+        if attention_mask[:, -1].sum() != attention_mask.shape[0]:
+            raise NotImplementedError('MPT does not support generation with right padding.')
+        if self.transformer.attn_uses_sequence_id and self.training:
+            sequence_id = torch.zeros_like(input_ids[:1])
+        else:
+            sequence_id = None
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+        if self.transformer.prefix_lm:
+            prefix_mask = torch.ones_like(attention_mask)
+            if kwargs.get('use_cache') == False:
+                raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
+        else:
+            prefix_mask = None
+        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        """Used by HuggingFace generate when using beam search with kv-caching.
+
+        See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
+        for an example in transformers.
+        """
+        reordered_past = []
+        for layer_past in past_key_values:
+            reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
+        return reordered_past

InternVL/internvl_chat_llava/llava/model/language_model/mpt/norm.py

@@ -0,0 +1,56 @@
+import torch
+
+def _cast_if_autocast_enabled(tensor):
+    if torch.is_autocast_enabled():
+        if tensor.device.type == 'cuda':
+            dtype = torch.get_autocast_gpu_dtype()
+        elif tensor.device.type == 'cpu':
+            dtype = torch.get_autocast_cpu_dtype()
+        else:
+            raise NotImplementedError()
+        return tensor.to(dtype=dtype)
+    return tensor
+
+class LPLayerNorm(torch.nn.LayerNorm):
+
+    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
+        super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
+
+    def forward(self, x):
+        module_device = x.device
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+        downcast_bias = _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 torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
+
+def rms_norm(x, weight=None, eps=1e-05):
+    output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+    if weight is not None:
+        return output * weight
+    return output
+
+class RMSNorm(torch.nn.Module):
+
+    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+        super().__init__()
+        self.eps = eps
+        if weight:
+            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
+        else:
+            self.register_parameter('weight', None)
+
+    def forward(self, x):
+        return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
+
+class LPRMSNorm(RMSNorm):
+
+    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+        super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
+
+    def forward(self, x):
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+        with torch.autocast(enabled=False, device_type=x.device.type):
+            return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
+NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}

InternVL/internvl_chat_llava/llava/model/language_model/mpt/param_init_fns.py

@@ -0,0 +1,181 @@
+import math
+import warnings
+from collections.abc import Sequence
+from functools import partial
+from typing import Optional, Tuple, Union
+import torch
+from torch import nn
+from .norm import NORM_CLASS_REGISTRY
+
+def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f"Initializing network using module's reset_parameters attribute")
+    if hasattr(module, 'reset_parameters'):
+        module.reset_parameters()
+
+def fused_init_helper_(module: nn.Module, init_fn_):
+    _fused = getattr(module, '_fused', None)
+    if _fused is None:
+        raise RuntimeError(f'Internal logic error')
+    (dim, splits) = _fused
+    splits = (0, *splits, module.weight.size(dim))
+    for (s, e) in zip(splits[:-1], splits[1:]):
+        slice_indices = [slice(None)] * module.weight.ndim
+        slice_indices[dim] = slice(s, e)
+        init_fn_(module.weight[slice_indices])
+
+def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'If model has bias parameters they are initialized to 0.')
+    init_div_is_residual = init_div_is_residual
+    if init_div_is_residual is False:
+        div_is_residual = 1.0
+    elif init_div_is_residual is True:
+        div_is_residual = math.sqrt(2 * n_layers)
+    elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
+        div_is_residual = init_div_is_residual
+    elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():
+        div_is_residual = float(init_div_is_residual)
+    else:
+        div_is_residual = 1.0
+        raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
+    if init_div_is_residual is not False:
+        if verbose > 1:
+            warnings.warn(f'Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. ' + f'Set `init_div_is_residual: false` in init config to disable this.')
+    if isinstance(module, nn.Linear):
+        if hasattr(module, '_fused'):
+            fused_init_helper_(module, init_fn_)
+        else:
+            init_fn_(module.weight)
+        if module.bias is not None:
+            torch.nn.init.zeros_(module.bias)
+        if init_div_is_residual is not False and getattr(module, '_is_residual', False):
+            with torch.no_grad():
+                module.weight.div_(div_is_residual)
+    elif isinstance(module, nn.Embedding):
+        if emb_init_std is not None:
+            std = emb_init_std
+            if std == 0:
+                warnings.warn(f'Embedding layer initialized to 0.')
+            emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
+            if verbose > 1:
+                warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
+        elif emb_init_uniform_lim is not None:
+            lim = emb_init_uniform_lim
+            if isinstance(lim, Sequence):
+                if len(lim) > 2:
+                    raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
+                if lim[0] == lim[1]:
+                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')
+            else:
+                if lim == 0:
+                    warnings.warn(f'Embedding layer initialized to 0.')
+                lim = [-lim, lim]
+            (a, b) = lim
+            emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
+            if verbose > 1:
+                warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
+        else:
+            emb_init_fn_ = init_fn_
+        emb_init_fn_(module.weight)
+    elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
+        if verbose > 1:
+            warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
+        if hasattr(module, 'weight') and module.weight is not None:
+            torch.nn.init.ones_(module.weight)
+        if hasattr(module, 'bias') and module.bias is not None:
+            torch.nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.MultiheadAttention):
+        if module._qkv_same_embed_dim:
+            assert module.in_proj_weight is not None
+            assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None)
+            assert d_model is not None
+            _d = d_model
+            splits = (0, _d, 2 * _d, 3 * _d)
+            for (s, e) in zip(splits[:-1], splits[1:]):
+                init_fn_(module.in_proj_weight[s:e])
+        else:
+            assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None)
+            assert module.in_proj_weight is None
+            init_fn_(module.q_proj_weight)
+            init_fn_(module.k_proj_weight)
+            init_fn_(module.v_proj_weight)
+        if module.in_proj_bias is not None:
+            torch.nn.init.zeros_(module.in_proj_bias)
+        if module.bias_k is not None:
+            torch.nn.init.zeros_(module.bias_k)
+        if module.bias_v is not None:
+            torch.nn.init.zeros_(module.bias_v)
+        init_fn_(module.out_proj.weight)
+        if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False):
+            with torch.no_grad():
+                module.out_proj.weight.div_(div_is_residual)
+        if module.out_proj.bias is not None:
+            torch.nn.init.zeros_(module.out_proj.bias)
+    else:
+        for _ in module.parameters(recurse=False):
+            raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
+
+def _normal_init_(std, mean=0.0):
+    return partial(torch.nn.init.normal_, mean=mean, std=std)
+
+def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+    del kwargs
+    init_fn_ = _normal_init_(std=std)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
+    generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def baseline_param_init_fn_(module: nn.Module, init_std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+    del kwargs
+    if init_std is None:
+        raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
+    _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+    del kwargs
+    std = math.sqrt(2 / (5 * d_model))
+    _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+    """From section 2.3.1 of GPT-NeoX-20B:
+
+    An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)
+    see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151
+    and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py
+    """
+    del kwargs
+    residual_div = n_layers / math.sqrt(10)
+    if verbose > 1:
+        warnings.warn(f'setting init_div_is_residual to {residual_div}')
+    small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
+    kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
+    generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
+    kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
+    generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
+    del kwargs
+    xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}')
+    generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+
+def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
+    xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}')
+    generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_}

InternVL/internvl_chat_llava/llava/model/multimodal_encoder/builder.py

@@ -0,0 +1,12 @@
+import os
+from .clip_encoder import CLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, **kwargs):
+    vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
+    is_absolute_path_exists = os.path.exists(vision_tower)
+    if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion") \
+            or "intern" in vision_tower.lower():
+        return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
+
+    raise ValueError(f'Unknown vision tower: {vision_tower}')