Upload safetensors export

.gitattributes

@@ -1,35 +1,7 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+sp.model filter=lfs diff=lfs merge=lfs -text
+vlmo/tokenizer/sp.model filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,91 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: zero-shot-image-classification
+tags:
+- multimodal
+- image-text-retrieval
+- bilingual
+- chinese
+- english
+- vision-language
+- custom-code
+---
+
+# M2-Encoder-0.4B Hugging Face Export
+
+This folder is generated from `Ant-Multi-Modal-Framework/prj/M2_Encoder` and is structured for direct upload to Hugging Face Hub.
+
+## What This Repo Supports
+
+- `AutoConfig.from_pretrained(..., trust_remote_code=True)`
+- `AutoProcessor.from_pretrained(..., trust_remote_code=True)`
+- `AutoModel.from_pretrained(..., trust_remote_code=True)`
+- Zero-shot image-text retrieval and zero-shot image classification
+
+## Required Weight File
+
+Put the model weight file in the repo root with this exact filename:
+
+`m2_encoder_0.4B.safetensors`
+
+Large files should be tracked by Git LFS. A `.gitattributes` file is included for that.
+
+## Usage
+
+### ModelScope-equivalent scoring
+
+The original ModelScope sample computes probabilities from the raw normalized embeddings:
+
+```python
+from transformers import AutoModel, AutoProcessor
+
+repo_id = "your-name/your-m2-encoder-repo"
+
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
+processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
+
+text_inputs = processor(
+    text=["杰尼龟", "妙蛙种子", "小火龙", "皮卡丘"],
+    return_tensors="pt",
+)
+image_inputs = processor(images="pokemon.jpeg", return_tensors="pt")
+
+text_outputs = model(**text_inputs)
+image_outputs = model(**image_inputs)
+
+probs = (image_outputs.image_embeds @ text_outputs.text_embeds.t()).softmax(dim=-1)
+print(probs)
+```
+
+### CLIP-style logits
+
+`model(**inputs)` also returns `logits_per_image` and `logits_per_text`, which use the model's learned `logit_scale`.
+Those logits are useful, but they are not the same computation as the raw dot product in the original ModelScope demo.
+
+## Upload
+
+Option 1:
+
+```bash
+python upload_to_hub.py --repo-id your-name/your-m2-encoder-repo
+```
+
+Option 2:
+
+```bash
+huggingface-cli login
+git init
+git lfs install
+git remote add origin https://huggingface.co/your-name/your-m2-encoder-repo
+git add .
+git commit -m "Upload M2-Encoder HF export"
+git push origin main
+```
+
+## Notes
+
+- This is a Hugging Face remote-code adapter, not a native `transformers` implementation.
+- The underlying model code still comes from the official M2-Encoder repo.
+- You need `trust_remote_code=True`.
+- The weights are not bundled by default when exporting unless you pass `--checkpoint`.

config.json

@@ -0,0 +1,37 @@
+{
+  "beit_version": "base",
+  "encoder_embed_dim": 768,
+  "out_embed_dim": 768,
+  "image_size": 224,
+  "visual_mask_size": 14,
+  "loss_names": {
+    "itc": 1
+  },
+  "encoder_layers": 9,
+  "beit3_vl_layers": 3,
+  "tokenizer_type": "GLMChineseTokenizer",
+  "tokenizer": ".",
+  "vocab_size": 115244,
+  "whole_word_masking": true,
+  "precision": 32,
+  "test_only": true,
+  "flash_attn": false,
+  "modelscope": {
+    "model_id": "M2Cognition/M2-Encoder"
+  },
+  "model_file": "m2_encoder_0.4B.safetensors",
+  "model_type": "m2_encoder",
+  "architectures": [
+    "M2EncoderModel"
+  ],
+  "processor_class": "M2EncoderProcessor",
+  "auto_map": {
+    "AutoConfig": "configuration_m2_encoder.M2EncoderConfig",
+    "AutoModel": "modeling_m2_encoder.M2EncoderModel",
+    "AutoProcessor": "processing_m2_encoder.M2EncoderProcessor",
+    "AutoTokenizer": [
+      "tokenization_glm.GLMChineseTokenizer",
+      null
+    ]
+  }
+}

configuration_m2_encoder.py

@@ -0,0 +1,90 @@
+import json
+import os
+from typing import Any, Dict
+
+from transformers import PretrainedConfig
+
+
+class M2EncoderConfig(PretrainedConfig):
+    model_type = "m2_encoder"
+
+    def __init__(
+        self,
+        loss_names=None,
+        beit_version="large",
+        encoder_embed_dim=1024,
+        out_embed_dim=1024,
+        encoder_layers=21,
+        beit3_vl_layers=3,
+        image_size=224,
+        visual_mask_size=14,
+        tokenizer_type="GLMChineseTokenizer",
+        tokenizer=".",
+        vocab_size=115244,
+        whole_word_masking=False,
+        precision=32,
+        test_only=True,
+        flash_attn=False,
+        model_file="m2_encoder_1B.ckpt",
+        architectures=None,
+        auto_map=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.loss_names = loss_names or {"itc": 1}
+        self.beit_version = beit_version
+        self.encoder_embed_dim = encoder_embed_dim
+        self.out_embed_dim = out_embed_dim
+        self.encoder_layers = encoder_layers
+        self.beit3_vl_layers = beit3_vl_layers
+        self.image_size = image_size
+        self.visual_mask_size = visual_mask_size
+        self.tokenizer_type = tokenizer_type
+        self.tokenizer = tokenizer
+        self.vocab_size = vocab_size
+        self.whole_word_masking = whole_word_masking
+        self.precision = precision
+        self.test_only = test_only
+        self.flash_attn = flash_attn
+        self.model_file = model_file
+        self.architectures = architectures or ["M2EncoderModel"]
+        self.auto_map = auto_map or {
+            "AutoConfig": "configuration_m2_encoder.M2EncoderConfig",
+            "AutoModel": "modeling_m2_encoder.M2EncoderModel",
+            "AutoProcessor": "processing_m2_encoder.M2EncoderProcessor",
+            "AutoTokenizer": ["tokenization_glm.GLMChineseTokenizer", None],
+        }
+
+    @classmethod
+    def from_encoder_json(cls, config_path: str, **kwargs) -> "M2EncoderConfig":
+        with open(config_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        data.update(kwargs)
+        return cls(**data)
+
+    def to_vlmo_overrides(self, model_dir: str) -> Dict[str, Any]:
+        return {
+            "loss_names": self.loss_names,
+            "beit_version": self.beit_version,
+            "encoder_embed_dim": self.encoder_embed_dim,
+            "out_embed_dim": self.out_embed_dim,
+            "encoder_layers": self.encoder_layers,
+            "beit3_vl_layers": self.beit3_vl_layers,
+            "image_size": self.image_size,
+            "visual_mask_size": self.visual_mask_size,
+            "tokenizer_type": self.tokenizer_type,
+            "tokenizer": self._resolve_tokenizer_dir(model_dir),
+            "vocab_size": self.vocab_size,
+            "whole_word_masking": self.whole_word_masking,
+            "precision": self.precision,
+            "test_only": self.test_only,
+            "flash_attn": self.flash_attn,
+            "load_path": os.path.join(model_dir, self.model_file),
+        }
+
+    def _resolve_tokenizer_dir(self, model_dir: str) -> str:
+        if os.path.isabs(self.tokenizer):
+            return self.tokenizer
+        if self.tokenizer in (".", "./", ""):
+            return model_dir
+        return os.path.join(model_dir, self.tokenizer)

image_processing_m2_encoder.py

@@ -0,0 +1,42 @@
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
+from transformers.image_utils import ImageFeatureExtractionMixin
+
+
+class M2EncoderImageProcessor(FeatureExtractionMixin, ImageFeatureExtractionMixin):
+    model_input_names = ["pixel_values"]
+
+    def __init__(self, size: int = 224, resample: int = Image.BICUBIC, **kwargs):
+        super().__init__(**kwargs)
+        if isinstance(size, dict):
+            size = int(size.get("height") or size.get("width"))
+        self.size = size
+        self.resample = resample
+
+    def __call__(
+        self,
+        images,
+        return_tensors: Optional[Union[str, torch.Tensor]] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        if not isinstance(images, (list, tuple)):
+            images = [images]
+
+        pixel_values: List[np.ndarray] = []
+        for image in images:
+            if not isinstance(image, Image.Image):
+                image = Image.fromarray(np.asarray(image))
+            image = image.convert("RGB")
+            image = image.resize((self.size, self.size), resample=self.resample)
+            array = np.asarray(image, dtype=np.float32) / 255.0
+            array = np.transpose(array, (2, 0, 1))
+            pixel_values.append(array)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values},
+            tensor_type=return_tensors,
+        )

m2_encoder_0.4B.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a75931d296881382b0c8282dafc63cbe2d1cacbb315c6d32572acbdaa2d9203e
+size 1053218384

modeling_m2_encoder.py

@@ -0,0 +1,150 @@
+import os
+import sys
+import importlib
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_m2_encoder import M2EncoderConfig
+
+
+@dataclass
+class M2EncoderOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    text_embeds: Optional[torch.FloatTensor] = None
+    image_embeds: Optional[torch.FloatTensor] = None
+    logits_per_image: Optional[torch.FloatTensor] = None
+    logits_per_text: Optional[torch.FloatTensor] = None
+
+
+class M2EncoderModel(PreTrainedModel):
+    config_class = M2EncoderConfig
+    base_model_prefix = "m2_encoder"
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: M2EncoderConfig):
+        super().__init__(config)
+        model_dir = getattr(config, "_model_dir", None)
+        if model_dir is None:
+            raise ValueError(
+                "M2EncoderConfig is missing `_model_dir`. Use "
+                "`M2EncoderModel.from_pretrained(...)` so the checkpoint path can be resolved."
+            )
+        if model_dir not in sys.path:
+            sys.path.insert(0, model_dir)
+
+        vlmo_default_config = importlib.import_module("vlmo.config").config
+        VLMo = importlib.import_module("vlmo.modules").VLMo
+
+        vlmo_config = vlmo_default_config()
+        vlmo_config.update(config.to_vlmo_overrides(model_dir))
+        load_path = vlmo_config["load_path"]
+        use_safetensors = load_path.endswith(".safetensors")
+        if use_safetensors:
+            vlmo_config["load_path"] = ""
+
+        if vlmo_config["flash_attn"]:
+            patch_torch_scale_with_flash_attn = importlib.import_module(
+                "vlmo.utils.patch_utils"
+            ).patch_torch_scale_with_flash_attn
+            patch_torch_scale_with_flash_attn()
+
+        self.model = VLMo(vlmo_config)
+        if use_safetensors:
+            state_dict = load_file(load_path)
+            self.model.load_state_dict(state_dict, strict=False)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *model_args,
+        config: Optional[M2EncoderConfig] = None,
+        **kwargs,
+    ):
+        model_dir = pretrained_model_name_or_path
+        if not os.path.isdir(model_dir):
+            model_dir = snapshot_download(repo_id=pretrained_model_name_or_path)
+
+        if config is None:
+            config = M2EncoderConfig.from_pretrained(model_dir, **kwargs)
+        checkpoint_path = os.path.join(
+            model_dir,
+            kwargs.pop("m2_checkpoint_name", config.model_file),
+        )
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(
+                f"Missing M2-Encoder checkpoint: {checkpoint_path}"
+            )
+        config._model_dir = model_dir
+        return cls(config, *model_args)
+
+    def get_text_features(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+    ) -> torch.FloatTensor:
+        outputs = self.model.infer_text(
+            {
+                "text_ids": input_ids,
+                "text_masks": attention_mask,
+                "text_labels": None,
+            }
+        )
+        return outputs["cls_vlffn_feats"]
+
+    def get_image_features(self, pixel_values: torch.FloatTensor) -> torch.FloatTensor:
+        outputs = self.model.infer_image({"image": [pixel_values]})
+        return outputs["cls_vlffn_feats"]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = True,
+        **kwargs,
+    ) -> Union[M2EncoderOutput, Tuple[torch.FloatTensor, ...]]:
+        text_embeds = None
+        image_embeds = None
+
+        if input_ids is not None:
+            if attention_mask is None:
+                attention_mask = torch.ones_like(input_ids)
+            text_embeds = self.get_text_features(
+                input_ids=input_ids, attention_mask=attention_mask
+            )
+
+        if pixel_values is not None:
+            image_embeds = self.get_image_features(pixel_values=pixel_values)
+
+        logits_per_image = None
+        logits_per_text = None
+        if image_embeds is not None and text_embeds is not None:
+            logit_scale = self.model.logit_scale.exp()
+            logits_per_image = logit_scale * image_embeds @ text_embeds.t()
+            logits_per_text = logits_per_image.t()
+
+        if not return_dict:
+            return tuple(
+                value
+                for value in (
+                    text_embeds,
+                    image_embeds,
+                    logits_per_image,
+                    logits_per_text,
+                )
+                if value is not None
+            )
+
+        return M2EncoderOutput(
+            text_embeds=text_embeds,
+            image_embeds=image_embeds,
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+        )

preprocessor_config.json

@@ -0,0 +1,11 @@
+{
+  "processor_class": "M2EncoderProcessor",
+  "image_processor_type": "M2EncoderImageProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_m2_encoder.M2EncoderProcessor"
+  },
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}

processing_m2_encoder.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from transformers import AutoTokenizer
+from transformers.processing_utils import ProcessorMixin
+
+from .image_processing_m2_encoder import M2EncoderImageProcessor
+
+
+class M2EncoderProcessor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "M2EncoderImageProcessor"
+    tokenizer_class = ("GLMChineseTokenizer", None)
+
+    def __init__(self, image_processor, tokenizer):
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        image_processor = M2EncoderImageProcessor.from_pretrained(
+            pretrained_model_name_or_path, **kwargs
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        return cls(image_processor=image_processor, tokenizer=tokenizer)
+
+    def __call__(
+        self,
+        text=None,
+        images=None,
+        padding="max_length",
+        truncation=True,
+        max_length: Optional[int] = 52,
+        return_tensors=None,
+        **kwargs,
+    ):
+        encoding = {}
+        if text is not None:
+            encoding.update(
+                self.tokenizer(
+                    text,
+                    padding=padding,
+                    truncation=truncation,
+                    max_length=max_length,
+                    return_special_tokens_mask=True,
+                    return_tensors=return_tensors,
+                    **kwargs,
+                )
+            )
+        if images is not None:
+            encoding.update(
+                self.image_processor(images, return_tensors=return_tensors, **kwargs)
+            )
+        return encoding

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "processor_class": "M2EncoderProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_m2_encoder.M2EncoderProcessor"
+  }
+}

requirements.txt

@@ -0,0 +1,15 @@
+torch
+pytorch_lightning<=2.0.8
+transformers
+safetensors
+Pillow
+tqdm
+einops
+sacred
+timm
+torchvision
+fairscale
+numpy
+opencv-python
+sentencepiece
+huggingface_hub

sp.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7fe3bcc8d284fcb782691411e8b6fd4f45d7245565b094de6ab795e66bcd32f
+size 2270960

tokenization_glm.py

@@ -0,0 +1,307 @@
+import os
+from shutil import copyfile
+from typing import Optional, Tuple, List, Union
+
+import sentencepiece as spm
+import torch
+from transformers import PreTrainedTokenizer
+from transformers.models.auto.tokenization_auto import get_tokenizer_config
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class GLMBatchEncoding(BatchEncoding):
+    def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
+        """
+        Send all values to device by calling `v.to(device)` (PyTorch only).
+
+        Args:
+            device (`str` or `torch.device`): The device to put the tensors on.
+
+        Returns:
+            [`BatchEncoding`]: The same instance after modification.
+        """
+
+        # This check catches things like APEX blindly calling "to" on all inputs to a module
+        # Otherwise it passes the casts down and casts the LongTensor containing the token idxs
+        # into a HalfTensor
+        if isinstance(device, str) or isinstance(device, int):
+        #if isinstance(device, str) or _is_torch_device(device) or isinstance(device, int):
+            self.data = {k: v.to(device=device) if torch.is_tensor(v) else v for k, v in self.data.items()}
+        else:
+            logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")
+        return self
+
+
+class GLMTokenizerMixin:
+    @property
+    def sop_token(self) -> Optional[str]:
+        return "<|startofpiece|>"
+
+    @property
+    def sop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the start token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.sop_token)
+
+    @property
+    def eop_token(self) -> Optional[str]:
+        return "<|endofpiece|>"
+
+    @property
+    def eop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the end token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.eop_token)
+
+    @property
+    def gmask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[gMASK]")
+
+    @property
+    def smask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[sMASK]")
+
+    @property
+    def mask_token_ids(self):
+        return [self.mask_token_id, self.smask_token_id, self.gmask_token_id]
+
+    def _build_input_for_multiple_choice(self, context, choices):
+        context_id = context["input_ids"]
+        if torch.is_tensor(context_id):
+            context_id = context_id.tolist()
+
+        division = len(context_id)
+        mask_position = context_id.index(self.mask_token_id)
+
+        token = torch.tensor(context_id, dtype=torch.long)
+        attention_mask = [context["attention_mask"].expand(division, -1)]
+        position_id = torch.arange(division, dtype=torch.long)
+        block_position_id = torch.zeros(division, dtype=torch.long)
+
+        choice_ids, choice_indices = [], []
+
+        for choice_str in choices:
+            choice = torch.tensor(self(choice_str, add_special_tokens=False, padding=False)['input_ids'],
+                                  dtype=torch.long)
+            choice_ids.append(choice)
+            choice_indices.append(torch.arange(len(token), len(token) + len(choice), dtype=torch.long))
+            attention_mask.append(torch.tril(torch.ones((len(choice), len(choice)), dtype=torch.long)))
+
+            token = torch.cat((token, torch.tensor([self.sop_token_id], dtype=torch.long), choice[:-1]))
+            position_id = torch.cat((position_id, torch.tensor([mask_position] * len(choice), dtype=torch.long)))
+            block_position_id = torch.cat((block_position_id, torch.arange(1, 1 + len(choice), dtype=torch.long)))
+
+        attention_mask = torch.block_diag(*attention_mask)
+        attention_mask[division:, :division] = context["attention_mask"].unsqueeze(0)
+
+        return {
+            "input_ids": token,
+            "position_ids": torch.stack((position_id, block_position_id)),
+            "attention_mask": attention_mask,
+            "choice_ids": choice_ids,
+            "choice_indices": choice_indices
+        }
+
+    def _pad_batch(self, tokens, position_ids, attention_mask, max_seq_length):
+        pad_length = max_seq_length - len(tokens)
+        attention_mask = torch.nn.functional.pad(
+            attention_mask,
+            (0, pad_length, 0, pad_length),
+            mode="constant",
+            value=0,
+        )
+        tokens = torch.cat((tokens, torch.zeros(pad_length, dtype=torch.long)))
+        position_ids = torch.cat((position_ids, position_ids[..., -1:].expand(-1, pad_length)), dim=-1)
+        return tokens, position_ids, attention_mask
+
+    def _collate(self, samples):
+        TILE = 1
+        length_to_pad = (max(map(lambda spl: len(spl["input_ids"]), samples)) + TILE - 1) // TILE * TILE
+
+        token_batch, position_id_batch, attention_mask_batch = [], [], []
+        choices_batch, choice_target_ids_batch = [], []
+
+        for sample in samples:
+            token, position_id, attention_mask = self._pad_batch(
+                sample["input_ids"], sample["position_ids"], sample["attention_mask"], length_to_pad
+            )
+            token_batch.append(token)
+            position_id_batch.append(position_id)
+            attention_mask_batch.append(attention_mask)
+            choices_batch.append(sample["choice_ids"])
+            choice_target_ids_batch.append(sample["choice_indices"])
+        return {
+            "input_ids": torch.stack(token_batch),
+            "position_ids": torch.stack(position_id_batch),
+            "attention_mask": torch.stack(attention_mask_batch).unsqueeze(1),
+            "choice_ids": choices_batch,
+            "choice_indices": choice_target_ids_batch,
+        }
+
+    def build_inputs_for_multiple_choice(self, model_input: BatchEncoding, choices, max_length=None):
+        samples = [{key: value[i] for key, value in model_input.items()} for i in range(len(model_input["input_ids"]))]
+        samples = [self._build_input_for_multiple_choice(sample, choice) for sample, choice in
+                   zip(samples, choices)]
+        inputs = self._collate(samples)
+        return GLMBatchEncoding(inputs)
+
+    def build_inputs_for_generation(self, model_input: BatchEncoding, max_gen_length=512, targets=None, padding=False):
+        mask_ids = self.mask_token_ids
+        input_ids = model_input.input_ids
+        batch_size, seq_length = input_ids.shape[:2]
+        position_id, block_position_id = list(range(seq_length)), [0 for _ in range(seq_length)]
+        position_ids, block_position_ids = [], []
+        labels = None
+        if targets is not None:
+            is_batched = isinstance(targets, (list, tuple))
+            targets = self(targets, add_special_tokens=False, padding=False).input_ids
+            if not is_batched:
+                targets = [targets]
+            assert len(targets) == len(input_ids)
+            targets = [(target + [self.eop_token_id])[:max_gen_length] for target in targets]
+            if not padding:
+                max_gen_length = max(map(len, targets))
+            targets = [[self.sop_token_id] + target for target in targets]
+            labels = [target[1:] for target in targets]
+            targets = [target + [self.pad_token_id] * (max_gen_length + 1 - len(target)) for target in targets]
+            labels = [label + [-100] * (max_gen_length - len(label)) for label in labels]
+            targets = torch.tensor(targets, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.tensor(labels, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.cat((input_ids.new_full((batch_size, seq_length), -100), labels), dim=1)
+        for i in range(batch_size):
+            mask_positions = []
+            for mask_id in mask_ids:
+                mask_positions += (input_ids[i] == mask_id).nonzero(as_tuple=True)[0].tolist()
+            if not mask_positions:
+                raise ValueError("Cannot find mask token in the input")
+            mask_positions.sort()
+            mask_pos = mask_positions[0]
+            position_ids.append(position_id + [mask_pos] * max_gen_length)
+            block_position_ids.append(block_position_id + list(range(1, max_gen_length + 1)))
+        position_ids = torch.tensor(position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        block_position_ids = torch.tensor(block_position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        position_ids = torch.stack((position_ids, block_position_ids), dim=1)
+        attention_mask = model_input.attention_mask
+        attention_mask = attention_mask.unsqueeze(1).expand(-1, seq_length + max_gen_length, -1)
+        generation_attention_mask = torch.cat([attention_mask.new_zeros((seq_length, max_gen_length)),
+                                               torch.tril(attention_mask.new_ones((max_gen_length, max_gen_length)))],
+                                              dim=0).unsqueeze(0).expand(batch_size, -1, -1)
+        attention_mask = torch.cat((attention_mask, generation_attention_mask), dim=2)
+        attention_mask = attention_mask.unsqueeze(1)
+        if targets is None:
+            input_ids = torch.cat((input_ids, input_ids.new_full((batch_size, 1), self.sop_token_id)), dim=-1)
+        else:
+            input_ids = torch.cat((input_ids, targets[:, :-1]), dim=1)
+        batch = {"input_ids": input_ids, "position_ids": position_ids}
+        if labels is None:
+            batch["generation_attention_mask"] = attention_mask
+        else:
+            batch["attention_mask"] = attention_mask
+            batch["labels"] = labels
+        return BatchEncoding(batch)
+
+def encode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(' '*i, f'<|blank_{i}|>')
+    return content
+
+def decode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(f'<|blank_{i}|>', ' '*i)
+    return content
+
+
+class GLMChineseTokenizer(PreTrainedTokenizer, GLMTokenizerMixin):
+    vocab_files_names = {"vocab_file": "sp.model"}
+    truncation_side: str = "left"
+
+    def __init__(self, vocab_file, **kwargs):
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return len(self.sp_model)
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text, **kwargs):
+        text = encode_whitespaces(text)
+        return self.sp_model.EncodeAsPieces(text)
+        #return self.sp_model.EncodeAsPieces(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.PieceToId(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.sp_model.IdToPiece(index)
+
+    def convert_tokens_to_string(self, tokens):
+        res = self.sp_model.DecodeIds(tokens)
+        return decode_whitespaces(res)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: ``[CLS] X [SEP]``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        assert token_ids_1 is None
+        cls = [self.cls_token_id]
+        eos = [self.eos_token_id]
+        return cls + token_ids_0 + eos
+
+
+class GLMTokenizer:
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+        tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
+        config_tokenizer_class = tokenizer_config.get("tokenizer_class")
+
+        if config_tokenizer_class == "GLMChineseTokenizer":
+            tokenizer_class = GLMChineseTokenizer
+        else:
+            raise NotImplementedError("Not implemented tokenizer type:", config_tokenizer_class)
+        return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "name_or_path": "THUDM/glm-10b-chinese",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]",
+  "add_prefix_space": false,
+  "tokenizer_class": "GLMChineseTokenizer",
+  "use_fast": false,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_glm.GLMChineseTokenizer",
+      null
+      ]
+  }
+}

upload_to_hub.py

@@ -0,0 +1,31 @@
+import argparse
+from pathlib import Path
+
+from huggingface_hub import HfApi, create_repo, upload_folder
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-id", required=True, help="Hugging Face repo id, e.g. user/M2-Encoder-Large")
+    parser.add_argument(
+        "--folder",
+        default=str(Path(__file__).resolve().parent),
+        help="Folder to upload. Defaults to this script's directory.",
+    )
+    parser.add_argument("--private", action="store_true", help="Create the repo as private.")
+    parser.add_argument("--commit-message", default="Upload M2-Encoder HF export")
+    args = parser.parse_args()
+
+    folder = Path(args.folder).resolve()
+    api = HfApi()
+    create_repo(repo_id=args.repo_id, private=args.private, exist_ok=True)
+    upload_folder(
+        repo_id=args.repo_id,
+        folder_path=str(folder),
+        commit_message=args.commit_message,
+    )
+    print(f"Uploaded {folder} -> {args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()

vlmo/config.py

@@ -0,0 +1,165 @@
+from sacred import Experiment
+
+ex = Experiment("VLMo")
+
+
+def _loss_names(d):
+    ret = {
+        "itm": 0,  # image-text matching loss
+        "itc": 0,  # image-text contrastive loss
+        "caption": 0,  # image captioning loss
+        "mvlm": 0,  # masked language modeling loss
+        "textmlm": 0,  # text-only masked language modeling
+        "imagemlm": 0,  # image-only masked language modeling
+        "vqa": 0,
+        "nlvr2": 0,
+        "irtr": 0,  # retrieval task ft
+    }
+    ret.update(d)
+    return ret
+
+
+@ex.config
+def config():
+    exp_name = "vlmo"
+    seed = 1
+    datasets = ["coco", "vg", "sbu", "gcc"]  # dataset name, the definition can refer to: vlmo/datamodules/__init__.py  # noqa
+    loss_names = _loss_names({"itm": 0, "itc": 0, "mvlm": 0})  # training loss
+    batch_size = 1024  # this is a desired batch size; pl trainer will accumulate gradients.
+
+    # BEiT-v3 setting
+    encoder_layers = 12  # the layer number of backbone
+    encoder_embed_dim = 768  # the hidden size of tokenizer
+    out_embed_dim = 768  # the hidden size of output embedding
+    beit_version = "base"  # model size: base(0.4B)|large(1B)|huge(10B)
+    beit3_vl_layers = 3  # the layer number of vl_backbone
+    deepnorm_init = True  # init method
+    share_layer = False  # if share the weight between layer within backbone
+    share_attn = False  # if share the attention weight of different layer
+    one_attn = False  # if share the attention weight of vision and language
+
+    # Image setting
+    train_transform_keys = ["square_transform_randaug"]  # train transform: refer to vlmo/transforms/__init__.py
+    val_transform_keys = ["square_transform"]  # test transform: refer to refer to vlmo/transforms/__init__.py
+    image_size = 224  # image size
+    reclip_image_size = None  # reclip image size
+    patch_size = 16  # patch size
+    draw_false_image = 0  # if get negative image
+    image_only = False  # only input image
+    text_only = False  # # only input text
+
+    # Video setting, video_num_frm is not None means video input
+    video_num_frm = None
+
+    # Visual tokenizer setting based on beit2
+    tokenizer_model = "beit2_visual_tokenizer"
+    codebook_size = 8192
+    codebook_dim = 32
+    visual_mask_size = 14
+    visual_mask_num = 80
+
+    # Text Setting
+    lang = 'cn'  # language for zero-shot imagenet testing: cn|en
+    vqav2_label_size = 3129
+    max_text_len = 52  # the number of characters
+    max_text_len_of_initckpt = 196
+    tokenizer_type = "BertTokenizer"  # Chinese text
+    vocab_size = 21128
+    tokenizer = "./vocab.txt"
+    whole_word_masking = True
+    mlm_prob = 0.15  # language mask ratio
+    draw_false_text = 0
+    mvlm_prob = 0.50  # vision-langurage mlm task
+    mask_ratio = 0  # flip: mask ratio for image
+
+    # cap setting
+    cap_onlytext = False  # default caption image to text
+
+    # imagemlm setting
+    split_data_for_imagemlm = False  # if True, split a batch data to two parts, and the first part for imagemlm.
+
+    # itc setting
+    itc_mask = False  # itc use masked token
+    aggregate_nodes = -1  # aggregate nodes num for compute_itc, default -1 is for all nodes
+
+    # Transformer Setting
+    model_arch = "vlmo_base_patch16"
+    drop_path_rate = 0.1
+
+    # Downstream Setting
+    get_recall_metric = False
+    get_recall_rerank_metric = False
+    get_zeroshot_metric = False
+    get_muge_feat = False
+    get_f30k_feat = False
+    k_test = 32
+
+    # PL Trainer Setting
+    resume_from = None
+    fast_dev_run = False
+    val_check_interval = 1.0
+    test_only = False
+    use_sharded_training = False
+    resume_during_training = False
+    save_top_k = 10
+    every_n_train_steps = 2000  # the step to save checkpoint
+    log_metric_steps = 100  # the step to log metric
+
+    # below params varies with the environment
+    use_pcache = False  # data storage method: pcache or nas
+    pcache_root = ""
+    # main_site: pcache://multimodalproxyi-pool.cz50c.alipay.com:39999/mnt/
+    # public_cloud: pcache://pcache_public_cloud.pcache.local:39999/mnt/abc7c88079a60b45ddfce7afa40720b7/
+    gpu_env = "main_site"  # public_cloud or main_site
+    data_root = ""  # data root for data list
+
+
+    log_dir = "result"
+    per_gpu_batchsize = 4  # you should define this manually with per_gpu_batch_size=#
+    num_gpus = 1
+    num_nodes = 1
+    load_path = ""
+    num_workers = 8
+    precision = 16
+    local_run = True
+    flash_attn = False
+    deepspeed_config = None  # "ds_config.json"
+    coalesce_backbone = False
+    mask_data = "v+l"  # 'v+l':choose input of imagemlm+textmlm task, 'vl': choose input of mvlm task.
+    communication_benchmark = False
+    checkpoint_activations = False
+
+    # dataset setting
+    single_cap = True  # if have only one caption
+    random_one = False  # if choose one caption from caption list
+
+    # ITC setting
+    itc_feats_name = "cls_vlffn_feats"  # feat for itc loss
+    itc_distill = ""
+    itc_distill_dim = 1024
+    itc_teacher_weights = ""
+
+    # mup training setting
+    mup = False
+    base_encoder_embed_dim = 1
+    delta_encoder_embed_dim = 2
+    mup_encoder_attention_heads = 1
+    base_encoder_ffn_embed_dim = 1
+    delta_encoder_ffn_embed_dim = 2
+
+    # atorch
+    atorch_config = None
+    compile_op = False
+    optimizer_state_shard_save = False
+    model_state_shard_save = False
+
+    # itc loss
+    local_loss = False
+    use_dual_softmax = False
+
+    num_frames = 1
+# ----------------------- LMM pretraining config -----------------------
+
+    # norm setting
+    deepnorm = False
+

vlmo/modules/__init__.py

@@ -0,0 +1 @@
+from .vlmo_module import VLMo

vlmo/modules/heads.py

@@ -0,0 +1,24 @@
+import torch.nn as nn
+
+
+class Pooler(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ITCHead(nn.Module):
+    def __init__(self, hidden_size, out_size):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, out_size, bias=False)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x

vlmo/modules/modeling_utils.py

@@ -0,0 +1,179 @@
+# --------------------------------------------------------
+# Image as a Foreign Language: BEiT Pretraining for Vision and Vision-Language Tasks (https://arxiv.org/abs/2208.10442)
+# Github source: https://github.com/microsoft/unilm/tree/master/beit3
+# Copyright (c) 2023 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------'
+
+import math
+import torch
+import torch.nn as nn
+from timm.models.layers import trunc_normal_ as __call_trunc_normal_
+
+from vlmo.torchscale.model.BEiT3 import BEiT3
+from vlmo.torchscale.architecture.config import EncoderConfig
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0):
+    __call_trunc_normal_(tensor, mean=mean, std=std, a=-std, b=std)
+
+
+def _get_base_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    encoder_layers=12,
+    encoder_embed_dim=768,
+    encoder_attention_heads=12,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+def _get_large_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=64010,
+    encoder_layers=24,
+    encoder_embed_dim=1024,
+    encoder_attention_heads=16,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+def _get_huge_config(
+    img_size=224,
+    patch_size=16,
+    drop_path_rate=0,
+    checkpoint_activations=None,
+    mlp_ratio=4,
+    vocab_size=30522,
+    encoder_layers=32,
+    encoder_embed_dim=4096,
+    encoder_attention_heads=32,
+    share_layer=False,
+    share_attn=False,
+    deepnorm=False,
+    mask_ratio=0,
+    max_text_len=52,
+    one_attn=False,
+    **kwargs
+):
+    return EncoderConfig(
+        img_size=img_size,
+        patch_size=patch_size,
+        vocab_size=vocab_size,
+        multiway=True,
+        layernorm_embedding=False,
+        normalize_output=True,
+        no_output_layer=True,
+        drop_path_rate=drop_path_rate,
+        encoder_embed_dim=encoder_embed_dim,
+        encoder_attention_heads=encoder_attention_heads,
+        encoder_layers=encoder_layers,
+        encoder_ffn_embed_dim=int(encoder_embed_dim * mlp_ratio),
+        checkpoint_activations=checkpoint_activations,
+        share_layer=share_layer,
+        share_attn=share_attn,
+        deepnorm=deepnorm,
+        mask_ratio=mask_ratio,
+        max_text_len=max_text_len,
+        one_attn=one_attn,
+    )
+
+
+class BEiT3Wrapper(nn.Module):
+    def __init__(self, args, **kwargs):
+        super().__init__()
+        self.args = args
+        self.beit3 = BEiT3(args)
+        self.apply(self._init_weights)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_num_layers(self):
+        return self.beit3.encoder.num_layers
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {
+            "pos_embed",
+            "cls_token",
+            "beit3.encoder.embed_positions.A.weight",
+            "beit3.vision_embed.cls_token",
+            "logit_scale",
+        }
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)

vlmo/modules/multiway_transformer.py

@@ -0,0 +1,396 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+DeiT model defs and weights from https://github.com/facebookresearch/deit,
+paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+from pytorch_lightning.utilities.distributed import rank_zero_info
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(dim))
+            self.v_bias = nn.Parameter(torch.zeros(dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, mask=None, relative_position_bias=None):
+        B, N, C = x.shape
+
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = q.float() @ k.float().transpose(-2, -1)
+
+        if relative_position_bias is not None:
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            mask = mask.bool()
+            attn = attn.masked_fill(~mask[:, None, None, :], float("-inf"))
+        attn = attn.softmax(dim=-1).type_as(x)
+        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
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        with_vlffn=False,
+        layer_scale_init_values=0.1,
+        max_text_len=40,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2_text = norm_layer(dim)
+        self.norm2_imag = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_text = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_imag = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_vl = None
+        if with_vlffn:
+            self.mlp_vl = Mlp(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop=drop,
+            )
+            self.norm2_vl = norm_layer(dim)
+
+        self.gamma_1 = (
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_values is not None
+            else 1.0
+        )
+        self.gamma_2 = (
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_values is not None
+            else 1.0
+        )
+
+        self.max_text_len = max_text_len
+
+    def forward(self, x, mask=None, modality_type=None, relative_position_bias=None):
+        x = x + self.drop_path(
+            self.gamma_1 * self.attn(self.norm1(x), mask=mask, relative_position_bias=relative_position_bias)
+        )
+
+        if modality_type == "image":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x)))
+        elif modality_type == "text":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x)))
+        else:
+            if self.mlp_vl is None:
+                x_text = x[:, : self.max_text_len]
+                x_imag = x[:, self.max_text_len :]
+                x_text = x_text + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x_text)))
+                x_imag = x_imag + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x_imag)))
+                x = torch.cat([x_text, x_imag], dim=1)
+            else:
+                x = x + self.drop_path(self.gamma_2 * self.mlp_vl(self.norm2_vl(x)))
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        no_patch_embed_bias=False,
+    ):
+        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.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False if no_patch_embed_bias else True,
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # FIXME look at relaxing size constraints
+        x = self.proj(x)
+        return x
+
+
+class MultiWayTransformer(nn.Module):
+    """Vision Transformer
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
+        https://arxiv.org/abs/2010.11929
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=None,
+        need_relative_position_embed=True,
+        use_abs_pos_emb=False,
+        layer_scale_init_values=0.1,
+        vlffn_start_layer_index=10,
+        config=None,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            norm_layer: (nn.Module): normalization layer
+            need_relative_position_embed (bool): enable relative position bias on self-attention
+            use_abs_pos_emb (bool): enable abs pos emb
+            layer_scale_init_values (float or None): layer scale init values, set None to disable
+            vlffn_start_layer_index (int): vl-ffn start index
+            config: (dict): other hyper from pytorch-lighting
+        """
+        super().__init__()
+        drop_path_rate = drop_path_rate if config is None else config["drop_path_rate"]
+        rank_zero_info("drop path rate: {}".format(drop_path_rate))
+        self.use_abs_pos_emb = use_abs_pos_emb
+        self.need_relative_position_embed = need_relative_position_embed
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.vlffn_start_layer_index = vlffn_start_layer_index
+        if config["loss_names"]["textmlm"] > 0:
+            self.vlffn_start_layer_index = depth
+            rank_zero_info(
+                "Set vlffn_start_layer_index={} for text-only pretraining".format(self.vlffn_start_layer_index)
+            )
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) if self.use_abs_pos_emb else None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    with_vlffn=(i >= self.vlffn_start_layer_index),
+                    layer_scale_init_values=layer_scale_init_values,
+                    max_text_len=config["max_text_len"],
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=0.02)
+        trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def visual_embed(self, _x):
+        x = self.patch_embed(_x)
+        x = x.flatten(2).transpose(1, 2)
+        B, L, _ = x.shape
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        x_mask = torch.ones(x.shape[0], x.shape[1])
+
+        return x, x_mask
+
+
+# VLMo base/p16
+@register_model
+def vlmo_base_patch16(pretrained=False, **kwargs):
+    img_size = kwargs.pop("img_size", 224)
+    model = MultiWayTransformer(
+        img_size=img_size,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        vlffn_start_layer_index=10,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    return model

vlmo/modules/objectives.py

@@ -0,0 +1,12 @@
+import torch.nn as nn
+
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+    if isinstance(module, nn.Linear) and module.bias is not None:
+        module.bias.data.zero_()

vlmo/modules/vlmo_module.py

@@ -0,0 +1,405 @@
+import math
+import os
+import time
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from pytorch_lightning.utilities.distributed import rank_zero_info
+from timm.models import create_model
+from transformers import AutoTokenizer, BertTokenizer, XLMRobertaTokenizer  # noqa
+from vlmo.modules import heads, objectives, vlmo_utils
+from vlmo.tokenizer.tokenization_glm import GLMChineseTokenizer  # noqa
+from vlmo.torchscale.architecture.encoder import Encoder
+from vlmo.torchscale.model.BEiT3 import BEiT3 as ts_backbone
+from vlmo.transforms.utils import inception_normalize as img_norm
+
+from .modeling_utils import _get_base_config, _get_large_config, _get_huge_config, trunc_normal_  # noqa
+
+
+def convert_pl_ckpt(state_dict, num_visual_token=197):
+    print("start convert_pl_ckpt!!!")
+    new_state_dict = {}
+    for key in state_dict:
+        value = state_dict[key]
+        if "visual_tokenizer" in key:
+            continue
+        elif "backbone.encoder.embed_positions.A.weight" in key:
+            if value.shape[0] < num_visual_token + 2:
+                N = value.shape[0] - 3
+                dim = value.shape[-1]
+                class_pos_embed = value[:3, ]
+                patch_pos_embed = value[3:, ]
+                w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                patch_pos_embed = patch_pos_embed.float()
+                patch_pos_embed = nn.functional.interpolate(
+                    patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                    size=(w0, h0),
+                    mode="area",
+                )
+                patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(-1, dim)
+                new_value = torch.cat((class_pos_embed, patch_pos_embed), dim=0)
+                new_state_dict[key] = new_value
+                print("reshape ", key, "raw shape: ", value.shape, "new shape: ", new_value.shape, num_visual_token)
+            elif value.shape[0] > num_visual_token + 2:
+                new_state_dict[key] = value[: num_visual_token + 2, :]
+                print("first ", key, "raw shape: ", value.shape, new_state_dict[key].shape, num_visual_token)
+            else:
+                new_state_dict[key] = value
+                print("raw shape")
+        else:
+            new_state_dict[key] = state_dict[key]
+
+    return new_state_dict
+
+
+def convert_deepspeed_ckpt(state_dict, num_visual_token=197):
+    new_state_dict = {}
+    for key in state_dict:
+        if key.startswith("_forward_module."):
+            new_key = key[len("_forward_module."):]
+            value = state_dict[key]
+            new_state_dict[new_key] = value
+            if "visual_tokenizer.encoder.pos_embed" in new_key or "visual_tokenizer.decoder.pos_embed" in new_key:
+                if value.shape[1] != num_visual_token:
+                    N = value.shape[1] - 1
+                    dim = value.shape[-1]
+                    class_pos_embed = value[:, 0]
+                    patch_pos_embed = value[:, 1:]
+                    w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                    patch_pos_embed = patch_pos_embed.float()
+                    patch_pos_embed = nn.functional.interpolate(
+                        patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                        size=(w0, h0),
+                        mode="area",
+                    )
+                    patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+                    new_value = torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+                    new_state_dict[new_key] = new_value
+                    print("reshape ", new_key, "raw shape: ", value.shape, "new_shape: ", new_value.shape)
+            if "backbone.encoder.embed_positions.A.weight" in new_key:
+                if value.shape[1] != num_visual_token + 2:
+                    N = value.shape[0] - 3
+                    dim = value.shape[-1]
+                    class_pos_embed = value[:3, ]
+                    patch_pos_embed = value[3:, ]
+                    w0, h0 = int(math.sqrt(num_visual_token - 1)), int(math.sqrt(num_visual_token - 1))
+                    patch_pos_embed = patch_pos_embed.float()
+                    patch_pos_embed = nn.functional.interpolate(
+                        patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+                        size=(w0, h0),
+                        mode="area",
+                    )
+                    patch_pos_embed = patch_pos_embed.to(class_pos_embed.dtype)
+                    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(-1, dim)
+                    new_value = torch.cat((class_pos_embed, patch_pos_embed), dim=0)
+                    new_state_dict[new_key] = new_value
+                    print("reshape ", new_key, "raw shape: ", value.shape, "new_shape: ", new_value.shape)
+
+        else:
+            new_state_dict[key] = state_dict[key]
+
+    return new_state_dict
+
+
+def get_visual_tokenizer(config):
+    tokenizer_name = config["tokenizer_model"]
+    print(f"Creating visual tokenizer: {tokenizer_name}")
+    model = create_model(
+        config["tokenizer_model"],
+        img_size=config["image_size"],
+        n_code=config["codebook_size"],
+        code_dim=config["codebook_dim"],
+    ).eval()
+    return model
+
+
+def get_pretrained_tokenizer(tokenizer_type, from_pretrained):
+    _Tokenizer = eval(f"{tokenizer_type}")
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            _Tokenizer.from_pretrained(from_pretrained)
+        torch.distributed.barrier()
+    return _Tokenizer.from_pretrained(from_pretrained)
+
+
+class VLMo(pl.LightningModule):
+    def __init__(self, config):
+        super().__init__()
+        self.save_hyperparameters()
+        s_t = time.time()
+
+        # tokenizer & backbone
+        self.img_size = config["image_size"]
+        if not config["test_only"]:
+            self.visual_tokenizer = get_visual_tokenizer(config)
+        kwargs = {}
+        if "encoder_attention_heads" in config:
+            kwargs["encoder_attention_heads"] = config["encoder_attention_heads"]
+        if "atorch_config" in config and config["atorch_config"]:
+            checkpoint_activations = False  # ?
+        else:
+            checkpoint_activations = config["checkpoint_activations"]
+        args = eval(f'_get_{config["beit_version"]}_config')(
+            img_size=config["image_size"],
+            patch_size=config["patch_size"],
+            vocab_size=config["vocab_size"],
+            encoder_layers=config["encoder_layers"],
+            encoder_embed_dim=config["encoder_embed_dim"],
+            checkpoint_activations=checkpoint_activations,
+            share_layer=config["share_layer"],
+            share_attn=config["share_attn"],
+            deepnorm=config["deepnorm"],
+            mask_ratio=config["mask_ratio"],
+            max_text_len=config["max_text_len"],
+            one_attn=config["one_attn"],
+            **kwargs,
+        )
+        self.num_features = args.encoder_embed_dim
+        self.out_features = config["out_embed_dim"]
+        self.cap_onlytext = config["cap_onlytext"]
+        self.lang = config["lang"]
+        self.num_frames = config["num_frames"]
+        self.tokenizer_type = config["tokenizer_type"]
+        self.text_tokenizer = get_pretrained_tokenizer(self.tokenizer_type, from_pretrained=config["tokenizer"])  # noqa
+        print("BEiT args", args.__dict__)
+        self.backbone = ts_backbone(args)
+
+        self.use_vl = config["beit3_vl_layers"] > 0
+        if self.use_vl:
+            args.encoder_layers = config["beit3_vl_layers"]
+            self.backbone_vl = Encoder(args)
+
+        self.norm = nn.LayerNorm(self.num_features, eps=1e-6)
+
+        # task layers
+        self.pooler = heads.Pooler(self.num_features)
+        self.pooler.apply(objectives.init_weights)
+
+        # contrastive loss (or sampling for global hard negative)
+        if config["loss_names"]["itc"] > 0:
+            self.itc_text_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_image_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_text_proj.apply(objectives.init_weights)
+            self.itc_image_proj.apply(objectives.init_weights)
+
+            self.itc_vl_text_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_vl_image_proj = heads.ITCHead(self.num_features, self.out_features)
+            self.itc_vl_text_proj.apply(objectives.init_weights)
+            self.itc_vl_image_proj.apply(objectives.init_weights)
+
+            self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+            self.logit_vl_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        lp_s_t = time.time()
+
+        self.load_pretrained_weight()
+        load_pretrain_time = time.time() - lp_s_t
+
+        self.current_tasks = list()
+
+        # ===================== load downstream (test_only) ======================
+
+        if self.hparams.config["load_path"] != "" and self.hparams.config["test_only"]:
+            rank_zero_info("Load ckpt from: {}".format(self.hparams.config["load_path"]))
+            ckpt = torch.load(self.hparams.config["load_path"], map_location="cpu")
+
+            state_dict = None
+
+            for state_dict_key in ("state_dict", "module", "model"):
+                if state_dict_key in ckpt:
+                    rank_zero_info("Read state dict from ckpt[%s]. " % state_dict_key)
+                    state_dict = ckpt[state_dict_key]
+                    break
+            if state_dict_key == "module":
+                state_dict = convert_deepspeed_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict_key == "state_dict":
+                state_dict = convert_pl_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict is None:
+                if list(ckpt.keys())[0].startswith('_forward_module.'):
+                    rank_zero_info("Read state dict from ckpt with _forward_module prefix. ")
+                    state_dict = convert_deepspeed_ckpt(ckpt, self.backbone.vision_embed.num_position_embeddings())
+                else:
+                    rank_zero_info("Read state dict from ckpt. ")
+                    state_dict = ckpt
+
+            missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
+            rank_zero_info("missing_keys: {}".format(missing_keys))
+            rank_zero_info("unexpected_keys: {}".format(unexpected_keys))
+
+        construct_time = time.time() - s_t
+        print(
+            f"Process {os.getpid()}. VLMo Constructor time: {construct_time}s;",
+            f"load_pretrain_time: {load_pretrain_time}s",
+            flush=True,
+        )
+        # coalesce backbone calls
+        self._coalesce_backbone = config["coalesce_backbone"]
+        self._mask_data = config["mask_data"]
+        self._backbone_inputs = {}
+        self._backbone_inputs_current_size = 0
+        self._backbone_inputs_keys = {}
+        self._backbone_outputs = None
+        self._default_attn_masks = {}
+        self._itc_group = None
+        self._itc_aggregate_dict = None
+        self._itc_mask = config["itc_mask"]
+        self._local_loss = config["local_loss"]
+        self._aggregate_nodes = config["aggregate_nodes"]
+        self.accumulated_batches_reached = False
+        vlmo_utils.set_task(self)
+        self._only_itc_single_machine = (
+                self._aggregate_nodes > 0 and len(self.current_tasks) == 1 and "itc" in self.current_tasks
+        )
+        self._split_data_for_imagemlm = config["split_data_for_imagemlm"]
+        self.log_metric_steps = config["log_metric_steps"]
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.backbone.encoder.layers):
+            rescale(layer.self_attn.v_proj.A.weight.data, layer_id + 1)
+            rescale(layer.self_attn.v_proj.B.weight.data, layer_id + 1)
+            rescale(layer.self_attn.out_proj.A.weight.data, layer_id + 1)
+            rescale(layer.self_attn.out_proj.B.weight.data, layer_id + 1)
+            rescale(layer.ffn.A.fc2.weight.data, layer_id + 1)
+            rescale(layer.ffn.B.fc2.weight.data, layer_id + 1)
+
+        if self.use_vl:
+            pre_layers = len(self.backbone.encoder.layers) + 1
+            for layer_id, layer in enumerate(self.backbone_vl.layers):
+                rescale(layer.self_attn.v_proj.A.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.v_proj.B.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.out_proj.A.weight.data, layer_id + pre_layers)
+                rescale(layer.self_attn.out_proj.B.weight.data, layer_id + pre_layers)
+                rescale(layer.ffn.A.fc2.weight.data, layer_id + pre_layers)
+                rescale(layer.ffn.B.fc2.weight.data, layer_id + pre_layers)
+
+    def load_pretrained_weight(self):
+        if self.hparams.config["load_path"] != "" and not self.hparams.config["test_only"]:
+            config = self.hparams.config
+            ckpt = torch.load(self.hparams.config["load_path"], map_location="cpu")
+            rank_zero_info("Load ckpt from: {}".format(self.hparams.config["load_path"]))
+
+            state_dict = None
+
+            for state_dict_key in ("state_dict", "module", "model"):
+                if state_dict_key in ckpt:
+                    rank_zero_info("Read state dict from ckpt[%s]. " % state_dict_key)
+                    state_dict = ckpt[state_dict_key]
+                    break
+            if state_dict_key == "module":
+                state_dict = convert_deepspeed_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict_key == "state_dict":
+                state_dict = convert_pl_ckpt(state_dict, self.backbone.vision_embed.num_position_embeddings())
+            if state_dict is None:
+                if list(ckpt.keys())[0].startswith('_forward_module.'):
+                    rank_zero_info("Read state dict from ckpt with _forward_module prefix. ")
+                    state_dict = convert_deepspeed_ckpt(ckpt,
+                                                        self.backbone.vision_embed.num_position_embeddings())
+                else:
+                    rank_zero_info("Read state dict from ckpt. ")
+                    state_dict = ckpt
+
+            missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
+            missing_keys = [k for k in missing_keys if "itc_teacher" not in k]
+            rank_zero_info("missing_keys: {}".format(missing_keys))
+            rank_zero_info("unexpected_keys: {}".format(unexpected_keys))
+
+    def infer_text(
+            self,
+            batch,
+            mask_text=False,
+    ):
+        do_mlm = "_mlm" if mask_text else ""
+        text_ids = batch[f"text_ids{do_mlm}"]
+        text_labels = batch[f"text_labels{do_mlm}"]
+        text_masks = batch[f"text_masks"]
+        text_embed = self.backbone.text_embed(text_ids)
+        text_padding_position = 1 - text_masks
+        lffn_hiddens = self.backbone(
+            textual_tokens=text_ids,
+            text_padding_position=text_padding_position,
+        )["encoder_out"]
+        vlffn_hiddens = self.backbone_vl(
+            src_tokens=None,
+            token_embeddings=lffn_hiddens,
+            encoder_padding_mask=text_padding_position,
+            multiway_split_position=-1,
+        )["encoder_out"]
+
+        cls_feats = self.itc_text_proj(lffn_hiddens[:, 0])
+        cls_feats = cls_feats / cls_feats.norm(dim=-1, keepdim=True)
+
+        cls_vlffn_feats = self.itc_vl_text_proj(vlffn_hiddens[:, 0])
+        cls_vlffn_feats = cls_vlffn_feats / cls_vlffn_feats.norm(dim=-1, keepdim=True)
+
+        ret = {
+            "cls_feats": cls_feats,
+            "cls_vlffn_feats": cls_vlffn_feats,
+            "text_embed": text_embed,
+        }
+
+        return ret
+
+    def infer_image(
+            self,
+            batch,
+            mask_image=False,
+            image_token_type_idx=1,
+            image_embeds=None,
+            image_masks=None,
+    ):
+        if f"image_{image_token_type_idx - 1}" in batch:
+            imgkey = f"image_{image_token_type_idx - 1}"
+        else:
+            imgkey = "image"
+
+        img = batch[imgkey][0]
+        if mask_image:
+            image_masks = batch[f"{imgkey}_masks"][0].flatten(1)
+
+            with torch.no_grad():
+                img = self.visual_tokenizer.pre_process(img)
+                quantize, embed_ind, _ = self.visual_tokenizer.encode(img)
+                image_ids = embed_ind.view(img.shape[0], -1)
+
+                image_labels = torch.full_like(image_ids, -100)
+                bool_masked_pos = image_masks.to(torch.bool)
+                image_labels[bool_masked_pos] = image_ids[bool_masked_pos]
+
+        img_tensor = img_norm(img)
+        vffn_hiddens = self.backbone(visual_tokens=img_tensor)["encoder_out"]
+        vlffn_hiddens = self.backbone_vl(
+            src_tokens=None,
+            token_embeddings=vffn_hiddens,
+            multiway_split_position=-1,
+        )["encoder_out"]
+
+        cls_feats = self.itc_image_proj(vffn_hiddens[:, 0])
+        cls_feats = cls_feats / cls_feats.norm(dim=-1, keepdim=True)
+
+        cls_vlffn_feats = self.itc_vl_image_proj(vlffn_hiddens[:, 0])
+        cls_vlffn_feats = cls_vlffn_feats / cls_vlffn_feats.norm(dim=-1, keepdim=True)
+
+        ret = {
+            "image_feats": vffn_hiddens,
+            "cls_feats": cls_feats,
+            "cls_vlffn_feats": cls_vlffn_feats,
+        }
+
+        return ret

vlmo/modules/vlmo_utils.py

@@ -0,0 +1,12 @@
+def set_task(pl_module):
+    pl_module.current_tasks = [k for k, v in pl_module.hparams.config["loss_names"].items() if v >= 1]
+    return
+
+
+def no_sync_module_apply(module, fn):
+    """FSDP module .apply will use _unshard_params_recurse which will sync params across ranks.
+    using this function when apply fn is unnecessary to sync params across ranks.
+    """
+    for child in module.children():
+        fn(child)
+        no_sync_module_apply(child, fn)

vlmo/tokenizer/__init__.py

@@ -0,0 +1,6 @@
+# coding: utf-8
+# Copyright (c) Antfin, Inc. All rights reserved.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function

vlmo/tokenizer/sp.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7fe3bcc8d284fcb782691411e8b6fd4f45d7245565b094de6ab795e66bcd32f
+size 2270960

vlmo/tokenizer/tokenization_glm.py

@@ -0,0 +1,307 @@
+import os
+from shutil import copyfile
+from typing import Optional, Tuple, List, Union
+
+import sentencepiece as spm
+import torch
+from transformers import PreTrainedTokenizer
+from transformers.models.auto.tokenization_auto import get_tokenizer_config
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class GLMBatchEncoding(BatchEncoding):
+    def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":
+        """
+        Send all values to device by calling `v.to(device)` (PyTorch only).
+
+        Args:
+            device (`str` or `torch.device`): The device to put the tensors on.
+
+        Returns:
+            [`BatchEncoding`]: The same instance after modification.
+        """
+
+        # This check catches things like APEX blindly calling "to" on all inputs to a module
+        # Otherwise it passes the casts down and casts the LongTensor containing the token idxs
+        # into a HalfTensor
+        if isinstance(device, str) or isinstance(device, int):
+        #if isinstance(device, str) or _is_torch_device(device) or isinstance(device, int):
+            self.data = {k: v.to(device=device) if torch.is_tensor(v) else v for k, v in self.data.items()}
+        else:
+            logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")
+        return self
+
+
+class GLMTokenizerMixin:
+    @property
+    def sop_token(self) -> Optional[str]:
+        return "<|startofpiece|>"
+
+    @property
+    def sop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the start token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.sop_token)
+
+    @property
+    def eop_token(self) -> Optional[str]:
+        return "<|endofpiece|>"
+
+    @property
+    def eop_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the end token in the vocabulary, used when training a model with autoregressive blank filling.
+        """
+        return self.convert_tokens_to_ids(self.eop_token)
+
+    @property
+    def gmask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[gMASK]")
+
+    @property
+    def smask_token_id(self) -> int:
+        return self.convert_tokens_to_ids("[sMASK]")
+
+    @property
+    def mask_token_ids(self):
+        return [self.mask_token_id, self.smask_token_id, self.gmask_token_id]
+
+    def _build_input_for_multiple_choice(self, context, choices):
+        context_id = context["input_ids"]
+        if torch.is_tensor(context_id):
+            context_id = context_id.tolist()
+
+        division = len(context_id)
+        mask_position = context_id.index(self.mask_token_id)
+
+        token = torch.tensor(context_id, dtype=torch.long)
+        attention_mask = [context["attention_mask"].expand(division, -1)]
+        position_id = torch.arange(division, dtype=torch.long)
+        block_position_id = torch.zeros(division, dtype=torch.long)
+
+        choice_ids, choice_indices = [], []
+
+        for choice_str in choices:
+            choice = torch.tensor(self(choice_str, add_special_tokens=False, padding=False)['input_ids'],
+                                  dtype=torch.long)
+            choice_ids.append(choice)
+            choice_indices.append(torch.arange(len(token), len(token) + len(choice), dtype=torch.long))
+            attention_mask.append(torch.tril(torch.ones((len(choice), len(choice)), dtype=torch.long)))
+
+            token = torch.cat((token, torch.tensor([self.sop_token_id], dtype=torch.long), choice[:-1]))
+            position_id = torch.cat((position_id, torch.tensor([mask_position] * len(choice), dtype=torch.long)))
+            block_position_id = torch.cat((block_position_id, torch.arange(1, 1 + len(choice), dtype=torch.long)))
+
+        attention_mask = torch.block_diag(*attention_mask)
+        attention_mask[division:, :division] = context["attention_mask"].unsqueeze(0)
+
+        return {
+            "input_ids": token,
+            "position_ids": torch.stack((position_id, block_position_id)),
+            "attention_mask": attention_mask,
+            "choice_ids": choice_ids,
+            "choice_indices": choice_indices
+        }
+
+    def _pad_batch(self, tokens, position_ids, attention_mask, max_seq_length):
+        pad_length = max_seq_length - len(tokens)
+        attention_mask = torch.nn.functional.pad(
+            attention_mask,
+            (0, pad_length, 0, pad_length),
+            mode="constant",
+            value=0,
+        )
+        tokens = torch.cat((tokens, torch.zeros(pad_length, dtype=torch.long)))
+        position_ids = torch.cat((position_ids, position_ids[..., -1:].expand(-1, pad_length)), dim=-1)
+        return tokens, position_ids, attention_mask
+
+    def _collate(self, samples):
+        TILE = 1
+        length_to_pad = (max(map(lambda spl: len(spl["input_ids"]), samples)) + TILE - 1) // TILE * TILE
+
+        token_batch, position_id_batch, attention_mask_batch = [], [], []
+        choices_batch, choice_target_ids_batch = [], []
+
+        for sample in samples:
+            token, position_id, attention_mask = self._pad_batch(
+                sample["input_ids"], sample["position_ids"], sample["attention_mask"], length_to_pad
+            )
+            token_batch.append(token)
+            position_id_batch.append(position_id)
+            attention_mask_batch.append(attention_mask)
+            choices_batch.append(sample["choice_ids"])
+            choice_target_ids_batch.append(sample["choice_indices"])
+        return {
+            "input_ids": torch.stack(token_batch),
+            "position_ids": torch.stack(position_id_batch),
+            "attention_mask": torch.stack(attention_mask_batch).unsqueeze(1),
+            "choice_ids": choices_batch,
+            "choice_indices": choice_target_ids_batch,
+        }
+
+    def build_inputs_for_multiple_choice(self, model_input: BatchEncoding, choices, max_length=None):
+        samples = [{key: value[i] for key, value in model_input.items()} for i in range(len(model_input["input_ids"]))]
+        samples = [self._build_input_for_multiple_choice(sample, choice) for sample, choice in
+                   zip(samples, choices)]
+        inputs = self._collate(samples)
+        return GLMBatchEncoding(inputs)
+
+    def build_inputs_for_generation(self, model_input: BatchEncoding, max_gen_length=512, targets=None, padding=False):
+        mask_ids = self.mask_token_ids
+        input_ids = model_input.input_ids
+        batch_size, seq_length = input_ids.shape[:2]
+        position_id, block_position_id = list(range(seq_length)), [0 for _ in range(seq_length)]
+        position_ids, block_position_ids = [], []
+        labels = None
+        if targets is not None:
+            is_batched = isinstance(targets, (list, tuple))
+            targets = self(targets, add_special_tokens=False, padding=False).input_ids
+            if not is_batched:
+                targets = [targets]
+            assert len(targets) == len(input_ids)
+            targets = [(target + [self.eop_token_id])[:max_gen_length] for target in targets]
+            if not padding:
+                max_gen_length = max(map(len, targets))
+            targets = [[self.sop_token_id] + target for target in targets]
+            labels = [target[1:] for target in targets]
+            targets = [target + [self.pad_token_id] * (max_gen_length + 1 - len(target)) for target in targets]
+            labels = [label + [-100] * (max_gen_length - len(label)) for label in labels]
+            targets = torch.tensor(targets, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.tensor(labels, dtype=input_ids.dtype, device=input_ids.device)
+            labels = torch.cat((input_ids.new_full((batch_size, seq_length), -100), labels), dim=1)
+        for i in range(batch_size):
+            mask_positions = []
+            for mask_id in mask_ids:
+                mask_positions += (input_ids[i] == mask_id).nonzero(as_tuple=True)[0].tolist()
+            if not mask_positions:
+                raise ValueError("Cannot find mask token in the input")
+            mask_positions.sort()
+            mask_pos = mask_positions[0]
+            position_ids.append(position_id + [mask_pos] * max_gen_length)
+            block_position_ids.append(block_position_id + list(range(1, max_gen_length + 1)))
+        position_ids = torch.tensor(position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        block_position_ids = torch.tensor(block_position_ids, dtype=input_ids.dtype, device=input_ids.device)
+        position_ids = torch.stack((position_ids, block_position_ids), dim=1)
+        attention_mask = model_input.attention_mask
+        attention_mask = attention_mask.unsqueeze(1).expand(-1, seq_length + max_gen_length, -1)
+        generation_attention_mask = torch.cat([attention_mask.new_zeros((seq_length, max_gen_length)),
+                                               torch.tril(attention_mask.new_ones((max_gen_length, max_gen_length)))],
+                                              dim=0).unsqueeze(0).expand(batch_size, -1, -1)
+        attention_mask = torch.cat((attention_mask, generation_attention_mask), dim=2)
+        attention_mask = attention_mask.unsqueeze(1)
+        if targets is None:
+            input_ids = torch.cat((input_ids, input_ids.new_full((batch_size, 1), self.sop_token_id)), dim=-1)
+        else:
+            input_ids = torch.cat((input_ids, targets[:, :-1]), dim=1)
+        batch = {"input_ids": input_ids, "position_ids": position_ids}
+        if labels is None:
+            batch["generation_attention_mask"] = attention_mask
+        else:
+            batch["attention_mask"] = attention_mask
+            batch["labels"] = labels
+        return BatchEncoding(batch)
+
+def encode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(' '*i, f'<|blank_{i}|>')
+    return content
+
+def decode_whitespaces(content):
+    for i in range(10, 1, -1):
+        content = content.replace(f'<|blank_{i}|>', ' '*i)
+    return content
+
+
+class GLMChineseTokenizer(PreTrainedTokenizer, GLMTokenizerMixin):
+    vocab_files_names = {"vocab_file": "sp.model"}
+    truncation_side: str = "left"
+
+    def __init__(self, vocab_file, **kwargs):
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return len(self.sp_model)
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text, **kwargs):
+        text = encode_whitespaces(text)
+        return self.sp_model.EncodeAsPieces(text)
+        #return self.sp_model.EncodeAsPieces(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.PieceToId(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.sp_model.IdToPiece(index)
+
+    def convert_tokens_to_string(self, tokens):
+        res = self.sp_model.DecodeIds(tokens)
+        return decode_whitespaces(res)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: ``[CLS] X [SEP]``
+        - pair of sequences: ``[CLS] A [SEP] B [SEP]``
+
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (:obj:`List[int]`, `optional`):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        assert token_ids_1 is None
+        cls = [self.cls_token_id]
+        eos = [self.eos_token_id]
+        return cls + token_ids_0 + eos
+
+
+class GLMTokenizer:
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+        tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
+        config_tokenizer_class = tokenizer_config.get("tokenizer_class")
+
+        if config_tokenizer_class == "GLMChineseTokenizer":
+            tokenizer_class = GLMChineseTokenizer
+        else:
+            raise NotImplementedError("Not implemented tokenizer type:", config_tokenizer_class)
+        return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)

vlmo/tokenizer/tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "name_or_path": "THUDM/glm-10b-chinese",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "unk_token": "[UNK]",
+  "add_prefix_space": false,
+  "tokenizer_class": "GLMChineseTokenizer",
+  "use_fast": false,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_glm.GLMChineseTokenizer",
+      null
+      ]
+  }
+}

vlmo/torchscale/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/torchscale/architecture/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]

vlmo/torchscale/architecture/config.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+
+class EncoderConfig(object):
+    def __init__(self, **kwargs):
+        self.encoder_embed_dim = kwargs.pop("encoder_embed_dim", 768)
+        self.encoder_attention_heads = kwargs.pop("encoder_attention_heads", 12)
+        self.encoder_ffn_embed_dim = kwargs.pop("encoder_ffn_embed_dim", 3072)
+        self.encoder_layers = kwargs.pop("encoder_layers", 12)
+        self.encoder_normalize_before = kwargs.pop("encoder_normalize_before", True)
+        self.normalize_output = kwargs.pop("normalize_output", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_encoder_input_output_embed = kwargs.pop("share_encoder_input_output_embed", False)
+        self.max_source_positions = kwargs.pop("max_source_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        self.share_layer = kwargs.pop("share_layer", False)
+        self.share_attn = kwargs.pop("share_attn", False)
+        self.mask_ratio = kwargs.pop("mask_ratio", 0)
+        self.max_text_len = kwargs.pop("max_text_len", 52)
+        self.one_attn = kwargs.pop('one_attn', False)
+
+
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Vision
+        self.img_size = kwargs.pop("img_size", 224)
+        self.patch_size = kwargs.pop("patch_size", 16)
+        self.in_chans = kwargs.pop("in_chans", 3)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.encoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.encoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)
+
+
+class DecoderConfig(object):
+    def __init__(self, **kwargs):
+        self.decoder_embed_dim = kwargs.pop("decoder_embed_dim", 768)
+        self.decoder_attention_heads = kwargs.pop("decoder_attention_heads", 12)
+        self.decoder_ffn_embed_dim = kwargs.pop("decoder_ffn_embed_dim", 3072)
+        self.decoder_layers = kwargs.pop("decoder_layers", 12)
+        self.decoder_normalize_before = kwargs.pop("decoder_normalize_before", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_decoder_input_output_embed = kwargs.pop("share_decoder_input_output_embed", False)
+        self.max_target_positions = kwargs.pop("max_target_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.decoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.decoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)
+
+
+class EncoderDecoderConfig(object):
+    def __init__(self, **kwargs):
+        self.encoder_embed_dim = kwargs.pop("encoder_embed_dim", 768)
+        self.encoder_attention_heads = kwargs.pop("encoder_attention_heads", 12)
+        self.encoder_ffn_embed_dim = kwargs.pop("encoder_ffn_embed_dim", 3072)
+        self.encoder_layers = kwargs.pop("encoder_layers", 12)
+        self.encoder_normalize_before = kwargs.pop("encoder_normalize_before", True)
+        self.decoder_embed_dim = kwargs.pop("decoder_embed_dim", 768)
+        self.decoder_attention_heads = kwargs.pop("decoder_attention_heads", 12)
+        self.decoder_ffn_embed_dim = kwargs.pop("decoder_ffn_embed_dim", 3072)
+        self.decoder_layers = kwargs.pop("decoder_layers", 12)
+        self.decoder_normalize_before = kwargs.pop("decoder_normalize_before", True)
+        self.activation_fn = kwargs.pop("activation_fn", "gelu")
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.drop_path_rate = kwargs.pop("drop_path_rate", 0.0)
+        self.attention_dropout = kwargs.pop("attention_dropout", 0.0)
+        self.activation_dropout = kwargs.pop("activation_dropout", 0.0)
+        self.no_scale_embedding = kwargs.pop("no_scale_embedding", True)
+        self.layernorm_embedding = kwargs.pop("layernorm_embedding", False)
+        self.moe_freq = kwargs.pop("moe_freq", 0)
+        self.moe_top1_expert = kwargs.pop("moe_top1_expert", False)
+        self.moe_expert_count = kwargs.pop("moe_expert_count", 0)
+        self.moe_gating_use_fp32 = kwargs.pop("moe_gating_use_fp32", True)
+        self.moe_eval_capacity_token_fraction = kwargs.pop("moe_eval_capacity_token_fraction", 0.25)
+        self.moe_second_expert_policy = kwargs.pop("moe_second_expert_policy", "random")
+        self.moe_normalize_gate_prob_before_dropping = kwargs.pop("moe_normalize_gate_prob_before_dropping", False)
+        self.use_xmoe = kwargs.pop("use_xmoe", False)
+        self.rel_pos_buckets = kwargs.pop("rel_pos_buckets", 0)
+        self.max_rel_pos = kwargs.pop("max_rel_pos", 0)
+        self.deepnorm = kwargs.pop("deepnorm", False)
+        self.subln = kwargs.pop("subln", True)
+        self.bert_init = kwargs.pop("bert_init", False)
+        self.multiway = kwargs.pop("multiway", False)
+        self.share_all_embeddings = kwargs.pop("share_all_embeddings", False)
+        self.share_decoder_input_output_embed = kwargs.pop("share_decoder_input_output_embed", False)
+        self.max_source_positions = kwargs.pop("max_source_positions", 1024)
+        self.max_target_positions = kwargs.pop("max_target_positions", 1024)
+        self.no_output_layer = kwargs.pop("no_output_layer", False)
+        self.layernorm_eps = kwargs.pop("layernorm_eps", 1e-5)
+        # Text
+        self.vocab_size = kwargs.pop("vocab_size", -1)
+        # Fairscale
+        self.checkpoint_activations = kwargs.pop("checkpoint_activations", False)
+        self.fsdp = kwargs.pop("fsdp", False)
+        self.ddp_rank = kwargs.pop("ddp_rank", 0)
+        self.xpos_rel_pos = kwargs.pop("xpos_rel_pos", False)
+        self.xpos_scale_base = kwargs.pop("xpos_scale_base", 512)
+
+        if self.deepnorm:
+            self.encoder_normalize_before = False
+            self.decoder_normalize_before = False
+            self.subln = False
+        if self.subln:
+            self.encoder_normalize_before = True
+            self.decoder_normalize_before = True
+            self.deepnorm = False
+        if self.use_xmoe:
+            self.moe_normalize_gate_prob_before_dropping = True
+            self.moe_second_expert_policy = "random"
+            assert self.moe_freq > 0 and self.moe_expert_count > 0
+
+    def override(self, args):
+        for hp in self.__dict__.keys():
+            if getattr(args, hp, None) is not None:
+                self.__dict__[hp] = getattr(args, hp, None)

vlmo/torchscale/architecture/decoder.py

@@ -0,0 +1,428 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+from fairscale.nn import checkpoint_wrapper, wrap
+
+from vlmo.torchscale.architecture.utils import init_bert_params
+from vlmo.torchscale.component.droppath import DropPath
+from vlmo.torchscale.component.feedforward_network import FeedForwardNetwork, make_experts
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.relative_position_bias import RelativePositionBias
+from vlmo.torchscale.component.xmoe.moe_layer import MOELayer
+from vlmo.torchscale.component.xmoe.routing import Top1Gate, Top2Gate
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+
+class DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        args,
+        depth,
+        is_moe_layer=False,
+        is_encoder_decoder=False,
+    ):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.decoder_embed_dim
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        if args.drop_path_rate > 0:
+            drop_path_prob = np.linspace(0, args.drop_path_rate, args.decoder_layers)[depth]
+            self.drop_path = DropPath(drop_path_prob)
+        else:
+            self.drop_path = None
+
+        self.self_attn = self.build_self_attention(self.embed_dim, args)
+
+        self.normalize_before = args.decoder_normalize_before
+
+        self.self_attn_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        if not is_encoder_decoder:
+            self.encoder_attn = None
+            self.encoder_attn_layer_norm = None
+        else:
+            self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
+            self.encoder_attn_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        self.is_moe_layer = is_moe_layer
+        self.ffn_dim = args.decoder_ffn_embed_dim
+
+        if not self.is_moe_layer:
+            self.ffn = self.build_ffn(
+                self.embed_dim,
+                self.args,
+            )
+        else:
+            if args.moe_top1_expert:
+                gate = Top1Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    use_fp32=args.moe_gating_use_fp32,
+                    moe_eval_capacity_token_fraction=args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            else:
+                gate = Top2Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    args.moe_gating_use_fp32,
+                    args.moe_second_expert_policy,
+                    args.moe_normalize_gate_prob_before_dropping,
+                    args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            experts = make_experts(args, self.embed_dim, self.ffn_dim)
+            self.moe_layer = MOELayer(gate, experts, args)
+
+        self.final_layer_norm = LayerNorm(self.embed_dim, eps=args.layernorm_eps)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                self.alpha = math.pow(3.0 * args.decoder_layers, 0.25)
+            else:
+                self.alpha = math.pow(2.0 * args.decoder_layers, 0.25)
+        else:
+            self.alpha = 1.0
+
+    def build_ffn(self, embed_dim, args):
+        return FeedForwardNetwork(
+            embed_dim,
+            self.ffn_dim,
+            args.activation_fn,
+            args.dropout,
+            args.activation_dropout,
+            args.layernorm_eps,
+            args.subln,
+        )
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+        )
+
+    def build_encoder_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=False,
+            encoder_decoder_attention=True,
+            subln=args.subln,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual * self.alpha + x
+
+    def forward(
+        self,
+        x,
+        encoder_out=None,
+        encoder_padding_mask=None,
+        incremental_state=None,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+        self_attn_rel_pos=None,
+        cross_attn_rel_pos=None,
+    ):
+        residual = x
+        if self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        x, attn = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=self_attn_padding_mask,
+            incremental_state=incremental_state,
+            attn_mask=self_attn_mask,
+            rel_pos=self_attn_rel_pos,
+        )
+        x = self.dropout_module(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        if self.encoder_attn is not None and encoder_out is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.encoder_attn_layer_norm(x)
+
+            x, attn = self.encoder_attn(
+                query=x,
+                key=encoder_out,
+                value=encoder_out,
+                key_padding_mask=encoder_padding_mask,
+                incremental_state=None,
+                rel_pos=cross_attn_rel_pos,
+            )
+            x = self.dropout_module(x)
+
+            if self.drop_path is not None:
+                x = self.drop_path(x)
+
+            x = self.residual_connection(x, residual)
+            if not self.normalize_before:
+                x = self.encoder_attn_layer_norm(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.final_layer_norm(x)
+        if not self.is_moe_layer:
+            x = self.ffn(x)
+            l_aux = None
+        else:
+            x, l_aux = self.moe_layer(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.final_layer_norm(x)
+
+        return x, attn, None, l_aux
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self, args, embed_tokens=None, embed_positions=None, output_projection=None, is_encoder_decoder=False, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.args = args
+
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        embed_dim = args.decoder_embed_dim
+        self.embed_dim = embed_dim
+        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
+
+        self.embed_tokens = embed_tokens
+        self.embed_positions = embed_positions
+
+        if output_projection is None and not args.no_output_layer and args.vocab_size > 0:
+            self.output_projection = self.build_output_projection(args)
+        else:
+            self.output_projection = output_projection
+
+        if args.layernorm_embedding:
+            self.layernorm_embedding = LayerNorm(embed_dim, eps=args.layernorm_eps)
+        else:
+            self.layernorm_embedding = None
+
+        self.layers = nn.ModuleList([])
+
+        moe_freq = args.moe_freq
+        for i in range(args.decoder_layers):
+            is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+            self.layers.append(
+                self.build_decoder_layer(
+                    args,
+                    depth=i,
+                    is_moe_layer=is_moe_layer,
+                    is_encoder_decoder=is_encoder_decoder,
+                )
+            )
+
+        self.num_layers = len(self.layers)
+
+        if args.decoder_normalize_before:
+            self.layer_norm = LayerNorm(embed_dim, eps=args.layernorm_eps)
+        else:
+            self.layer_norm = None
+
+        self.self_attn_relative_position = None
+        self.cross_attn_relative_position = None
+
+        if args.rel_pos_buckets > 0 and args.max_rel_pos > 0:
+            self.self_attn_relative_position = RelativePositionBias(
+                num_buckets=args.rel_pos_buckets,
+                max_distance=args.max_rel_pos,
+                n_heads=args.decoder_attention_heads,
+            )
+            if is_encoder_decoder:
+                self.cross_attn_relative_position = RelativePositionBias(
+                    num_buckets=args.rel_pos_buckets,
+                    max_distance=args.max_rel_pos,
+                    n_heads=args.decoder_attention_heads,
+                )
+
+        if args.bert_init:
+            self.apply(init_bert_params)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                init_scale = math.pow(12.0 * args.decoder_layers, 0.25)
+            else:
+                init_scale = math.pow(8.0 * args.decoder_layers, 0.25)
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.div_(init_scale)
+
+        if args.subln:
+            if is_encoder_decoder:
+                init_scale = math.sqrt(math.log(args.decoder_layers * 3))
+            else:
+                init_scale = math.sqrt(math.log(args.decoder_layers * 2))
+            for name, p in self.named_parameters():
+                if "encoder_attn" in name:
+                    continue
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.mul_(init_scale)
+
+    def build_output_projection(
+        self,
+        args,
+    ):
+        if args.share_decoder_input_output_embed:
+            output_projection = torch.nn.Linear(
+                self.embed_tokens.weight.shape[1],
+                self.embed_tokens.weight.shape[0],
+                bias=False,
+            )
+            output_projection.weight = self.embed_tokens.weight
+        else:
+            output_projection = torch.nn.Linear(args.decoder_embed_dim, args.vocab_size, bias=False)
+            torch.nn.init.normal_(output_projection.weight, mean=0, std=args.decoder_embed_dim**-0.5)
+        return output_projection
+
+    def build_decoder_layer(self, args, depth, is_moe_layer=False, is_encoder_decoder=False):
+        layer = DecoderLayer(
+            args,
+            depth,
+            is_moe_layer=is_moe_layer,
+            is_encoder_decoder=is_encoder_decoder,
+        )
+        if args.checkpoint_activations:
+            layer = checkpoint_wrapper(layer)
+        if args.fsdp:
+            layer = wrap(layer)
+        return layer
+
+    def forward_embedding(
+        self,
+        tokens,
+        token_embedding=None,
+        incremental_state=None,
+    ):
+        positions = None
+        if self.embed_positions is not None:
+            positions = self.embed_positions(tokens, incremental_state=incremental_state)
+
+        if incremental_state is not None:
+            tokens = tokens[:, -1:]
+            if positions is not None:
+                positions = positions[:, -1:]
+
+        if token_embedding is None:
+            token_embedding = self.embed_tokens(tokens)
+
+        x = embed = self.embed_scale * token_embedding
+
+        if positions is not None:
+            x += positions
+
+        if self.layernorm_embedding is not None:
+            x = self.layernorm_embedding(x)
+
+        x = self.dropout_module(x)
+
+        return x, embed
+
+    def forward(
+        self,
+        prev_output_tokens,
+        self_attn_padding_mask=None,
+        encoder_out=None,
+        incremental_state=None,
+        features_only=False,
+        return_all_hiddens=False,
+        token_embeddings=None,
+        **kwargs
+    ):
+        # embed tokens and positions
+        x, _ = self.forward_embedding(prev_output_tokens, token_embeddings, incremental_state)
+
+        # relative position
+        self_attn_rel_pos_bias = None
+        slen = prev_output_tokens.size(1)
+        if self.self_attn_relative_position is not None:
+            self_attn_rel_pos_bias = self.self_attn_relative_position(batch_size=x.size(0), qlen=slen, klen=slen)
+            if incremental_state is not None:
+                self_attn_rel_pos_bias = self_attn_rel_pos_bias[-1:, :, :]
+        cross_attn_rel_pos_bias = None
+        if self.cross_attn_relative_position is not None:
+            cross_attn_rel_pos_bias = self.cross_attn_relative_position(
+                batch_size=x.size(0),
+                qlen=slen,
+                klen=encoder_out["encoder_out"].size(1),
+            )
+            if incremental_state is not None:
+                cross_attn_rel_pos_bias = cross_attn_rel_pos_bias[-1:, :, :]
+
+        # decoder layers
+        inner_states = [x]
+
+        if encoder_out is None:
+            l_aux = []
+        else:
+            l_aux = encoder_out["l_aux"] if "l_aux" in encoder_out else []
+
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = torch.triu(
+                    torch.zeros([x.size(1), x.size(1)]).float().fill_(float("-inf")).type_as(x),
+                    1,
+                )
+            else:
+                self_attn_mask = None
+                if idx not in incremental_state:
+                    incremental_state[idx] = {}
+
+            x, layer_attn, _, l_aux_i = layer(
+                x,
+                encoder_out["encoder_out"] if encoder_out is not None else None,
+                encoder_out["encoder_padding_mask"] if encoder_out is not None else None,
+                incremental_state[idx] if incremental_state is not None else None,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                self_attn_rel_pos=self_attn_rel_pos_bias,
+                cross_attn_rel_pos=cross_attn_rel_pos_bias,
+            )
+            l_aux.append(l_aux_i)
+            inner_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        if not features_only:
+            x = self.output_layer(x)
+
+        return x, {
+            "inner_states": inner_states,
+            "l_aux": l_aux,
+            "attn": None,
+        }
+
+    def output_layer(self, features):
+        return self.output_projection(features)

vlmo/torchscale/architecture/encoder.py

@@ -0,0 +1,482 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+from fairscale.nn import checkpoint_wrapper, wrap
+
+try:
+    from apex.normalization import FusedLayerNorm as LayerNorm
+except ModuleNotFoundError:
+    from torch.nn import LayerNorm
+
+from vlmo.torchscale.architecture.utils import init_bert_params
+from vlmo.torchscale.component.droppath import DropPath
+from vlmo.torchscale.component.feedforward_network import FeedForwardNetwork, make_experts
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.multiway_network import MultiwayWrapper, set_split_position
+from vlmo.torchscale.component.relative_position_bias import RelativePositionBias
+from vlmo.torchscale.component.xmoe.moe_layer import MOELayer
+from vlmo.torchscale.component.xmoe.routing import Top1Gate, Top2Gate
+from vlmo.modules.vlmo_utils import no_sync_module_apply
+from pytorch_lightning.utilities.distributed import rank_zero_info
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, args, depth, attn=None, is_moe_layer=False, is_encoder_decoder=False):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.encoder_embed_dim
+        self.self_attn = self.build_self_attention(self.embed_dim, args) if attn is None else attn
+        self.self_attn_layer_norm = MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps))
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        if args.drop_path_rate > 0:
+            drop_path_prob = np.linspace(0, args.drop_path_rate, args.encoder_layers)[depth]
+            self.drop_path = DropPath(drop_path_prob)
+        else:
+            self.drop_path = None
+
+        self.normalize_before = args.encoder_normalize_before
+        self.is_moe_layer = is_moe_layer
+        self.ffn_dim = args.encoder_ffn_embed_dim
+
+        if not self.is_moe_layer:
+            self.ffn = MultiwayWrapper(
+                args,
+                self.build_ffn(
+                    self.embed_dim,
+                    self.args,
+                ),
+            )
+        else:
+            assert not self.args.multiway
+            if args.moe_top1_expert:
+                gate = Top1Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    use_fp32=args.moe_gating_use_fp32,
+                    moe_eval_capacity_token_fraction=args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            else:
+                gate = Top2Gate(
+                    self.embed_dim,
+                    args.moe_expert_count,
+                    args.moe_gating_use_fp32,
+                    args.moe_second_expert_policy,
+                    args.moe_normalize_gate_prob_before_dropping,
+                    args.moe_eval_capacity_token_fraction,
+                    use_xmoe=args.use_xmoe,
+                )
+            experts = make_experts(args, self.embed_dim, self.ffn_dim)
+            self.moe_layer = MOELayer(gate, experts, args)
+        self.final_layer_norm = MultiwayWrapper(args, LayerNorm(self.embed_dim, eps=args.layernorm_eps))
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                self.alpha = math.pow(math.pow(args.encoder_layers, 4) * args.decoder_layers, 0.0625) * 0.81
+            else:
+                self.alpha = math.pow(2.0 * args.encoder_layers, 0.25)
+        else:
+            self.alpha = 1.0
+
+    def build_ffn(self, embed_dim, args):
+        return FeedForwardNetwork(
+            embed_dim,
+            self.ffn_dim,
+            args.activation_fn,
+            args.dropout,
+            args.activation_dropout,
+            args.layernorm_eps,
+            args.subln,
+        )
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+            one_attn=args.one_attn,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual * self.alpha + x
+
+    def forward(
+        self,
+        x,
+        encoder_padding_mask,
+        attn_mask=None,
+        rel_pos=None,
+        multiway_split_position=None,
+        incremental_state=None,
+    ):
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+
+        if attn_mask is not None:
+            # float16: -1e8 equal 0
+            attn_mask = attn_mask.masked_fill(attn_mask.to(torch.bool), -1e8)
+
+        residual = x
+        if self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+        x, _ = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=encoder_padding_mask,
+            attn_mask=attn_mask,
+            rel_pos=rel_pos,
+            incremental_state=incremental_state,
+        )
+        x = self.dropout_module(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.self_attn_layer_norm(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.final_layer_norm(x)
+        if not self.is_moe_layer:
+            x = self.ffn(x)
+            l_aux = None
+        else:
+            x = x.transpose(0, 1)
+            x, l_aux = self.moe_layer(x)
+            x = x.transpose(0, 1)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x = self.residual_connection(x, residual)
+        if not self.normalize_before:
+            x = self.final_layer_norm(x)
+        return x, l_aux
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self, args, embed_tokens=None, embed_positions=None, output_projection=None, is_encoder_decoder=False, **kwargs
+    ):
+        self.args = args
+        super().__init__(**kwargs)
+
+        self.dropout_module = torch.nn.Dropout(args.dropout)
+
+        embed_dim = args.encoder_embed_dim
+        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
+        self.mask_ratio = args.mask_ratio
+        self.max_text_len = args.max_text_len
+        self.vision_len = (args.img_size // args.patch_size) * (args.img_size // args.patch_size)
+
+        self.embed_tokens = embed_tokens
+        self.embed_positions = embed_positions
+
+        if output_projection is None and not is_encoder_decoder and not args.no_output_layer and args.vocab_size > 0:
+            self.output_projection = self.build_output_projection(args)
+        else:
+            self.output_projection = output_projection
+
+        if args.layernorm_embedding:
+            self.layernorm_embedding = MultiwayWrapper(args, LayerNorm(embed_dim, eps=args.layernorm_eps), dim=1)
+        else:
+            self.layernorm_embedding = None
+
+        self.layers = nn.ModuleList([])
+        if self.args.share_layer:
+            single_layer = self.build_encoder_layer(
+                args, depth=0, is_moe_layer=False, is_encoder_decoder=is_encoder_decoder
+            )
+            for i in range(args.encoder_layers):
+                self.layers.append(single_layer)
+        elif self.args.share_attn:
+            moe_freq = args.moe_freq
+            embed_dim = args.encoder_embed_dim
+            shared_attn = self.build_self_attention(embed_dim, self.args)
+            for i in range(args.encoder_layers):
+                is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+                self.layers.append(
+                    self.build_encoder_layer(
+                        args,
+                        depth=i,
+                        attn=shared_attn,
+                        is_moe_layer=is_moe_layer,
+                        is_encoder_decoder=is_encoder_decoder,
+                    )
+                )
+
+        else:
+            moe_freq = args.moe_freq
+            for i in range(args.encoder_layers):
+                is_moe_layer = moe_freq != 0 and (i + 1) % moe_freq == 0
+                self.layers.append(
+                    self.build_encoder_layer(
+                        args,
+                        depth=i,
+                        is_moe_layer=is_moe_layer,
+                        is_encoder_decoder=is_encoder_decoder,
+                    )
+                )
+        self.num_layers = len(self.layers)
+
+        if args.encoder_normalize_before and args.normalize_output:
+            self.layer_norm = MultiwayWrapper(args, LayerNorm(embed_dim, eps=args.layernorm_eps))
+        else:
+            self.layer_norm = None
+
+        if args.rel_pos_buckets > 0 and args.max_rel_pos > 0:
+            self.relative_position = RelativePositionBias(
+                num_buckets=args.rel_pos_buckets,
+                max_distance=args.max_rel_pos,
+                n_heads=args.encoder_attention_heads,
+            )
+        else:
+            self.relative_position = None
+
+        if args.bert_init:
+            self.apply(init_bert_params)
+
+        if args.deepnorm:
+            if is_encoder_decoder:
+                init_scale = math.pow(math.pow(args.encoder_layers, 4) * args.decoder_layers, 0.0625) / 1.15
+            else:
+                init_scale = math.pow(8.0 * args.encoder_layers, 0.25)
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.div_(init_scale)
+
+        if args.subln:
+            if is_encoder_decoder:
+                init_scale = math.sqrt(math.log(3 * args.decoder_layers) * math.log(2 * args.encoder_layers) / 3)
+            else:
+                init_scale = math.sqrt(math.log(args.encoder_layers * 2))
+            for name, p in self.named_parameters():
+                if "fc1" in name or "fc2" in name or "out_proj" in name or "v_proj" in name:
+                    p.data.mul_(init_scale)
+
+    def random_masking(self, x, mask_ratio):
+        N, L, D = x.shape  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L - 1, device=x.device)
+        ids_shuffle = torch.argsort(noise, dim=1) + torch.ones(N, L - 1, device=x.device, dtype=int)
+        ids_keep = ids_shuffle[:, :len_keep]
+
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        x0 = x[:, 0, :]
+        x0 = x0.reshape(N, 1, D)
+        x_masked_add = torch.cat([x0, x_masked], axis=1)
+        return x_masked_add, ids_keep
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            args,
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+            encoder_decoder_attention=False,
+            subln=args.subln,
+            one_attn=args.one_attn,
+        )
+
+    def build_output_projection(
+        self,
+        args,
+    ):
+        if args.share_encoder_input_output_embed:
+            assert args.encoder_embedding_type == "language"
+            output_projection = torch.nn.Linear(
+                self.embed_tokens.weight.shape[1],
+                self.embed_tokens.weight.shape[0],
+                bias=False,
+            )
+            output_projection.weight = self.embed_tokens.weight
+        else:
+            output_projection = torch.nn.Linear(args.encoder_embed_dim, args.vocab_size, bias=False)
+            torch.nn.init.normal_(output_projection.weight, mean=0, std=args.encoder_embed_dim**-0.5)
+        return output_projection
+
+    def checkpointing_and_params_allgather(
+        self,
+        origin_layer,
+    ):
+        origin_forward = origin_layer.forward
+
+        from deepspeed import checkpointing
+        def forward(*args, **kwargs):
+            # deepspeed checkpoint not support kwargs
+            ret = checkpointing.checkpoint(origin_forward, *args, **kwargs)
+            return ret
+
+        return forward
+
+    def build_encoder_layer(self, args, depth, attn=None, is_moe_layer=False, is_encoder_decoder=False):
+        layer = EncoderLayer(
+            args,
+            depth,
+            attn,
+            is_moe_layer=is_moe_layer,
+            is_encoder_decoder=is_encoder_decoder,
+        )
+        if args.checkpoint_activations:
+            rank_zero_info("EncoderLayer params: %s", sum(p.numel() for p in layer.parameters() if p.requires_grad))
+            layer = checkpoint_wrapper(layer)
+            # layer.ffn = checkpoint_wrapper(layer.ffn,)
+        if args.fsdp:
+            layer = wrap(layer)
+        return layer
+
+    def checkpointing_layers(self):
+        for i, layer in enumerate(self.layers):
+            rank_zero_info(f"Checkpointing wrapper EncoderLayers: {i}")
+            self.layers[i] = checkpoint_wrapper(layer)
+
+    def forward_embedding(
+        self,
+        src_tokens,
+        token_embedding=None,
+        positions=None,
+    ):
+        if token_embedding is None:
+            token_embedding = self.embed_tokens(src_tokens)
+        x = embed = self.embed_scale * token_embedding
+        if self.embed_positions is not None:
+            if src_tokens is not None:
+                x = embed + self.embed_positions(src_tokens, positions=positions)
+            else:
+                x = embed + self.embed_positions(x, positions=positions)
+        is_flip, ids_keep = 0, None
+        if self.mask_ratio > 0:
+            if x.shape[1] == self.vision_len + 1:
+                x, ids_keep = self.random_masking(x, self.mask_ratio)
+                is_flip = 1
+            elif x.shape[1] == self.vision_len + self.max_text_len + 1:
+                vision_tokens = x[:, : self.vision_len + 1, :]
+                vision_tokens, ids_keep = self.random_masking(vision_tokens, self.mask_ratio)
+                x = torch.cat(
+                    [
+                        vision_tokens,
+                        x[
+                            :,
+                            self.vision_len + 1 :,
+                        ],
+                    ],
+                    dim=1,
+                )
+                is_flip = 2
+        if self.layernorm_embedding is not None:
+            x = self.layernorm_embedding(x)
+        x = self.dropout_module(x)
+        return x, embed, ids_keep, is_flip
+
+    def forward(
+        self,
+        src_tokens,
+        encoder_padding_mask=None,
+        attn_mask=None,
+        return_all_hiddens=False,
+        token_embeddings=None,
+        multiway_split_position=None,
+        features_only=False,
+        incremental_state=None,
+        positions=None,
+        **kwargs
+    ):
+        assert src_tokens is not None or token_embeddings is not None
+
+        if encoder_padding_mask is None:
+            if src_tokens is not None:
+                encoder_padding_mask = torch.zeros_like(src_tokens, device=src_tokens.device).bool()
+            else:
+                encoder_padding_mask = torch.zeros(
+                    [token_embeddings.size(0), token_embeddings.size(1)],
+                    device=token_embeddings.device,
+                ).bool()
+
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+
+        x, encoder_embedding, ids_keep, is_flip = self.forward_embedding(src_tokens, token_embeddings, positions)
+        if is_flip > 0:
+            if is_flip == 2:
+                text_ids = (
+                    torch.arange(
+                        self.vision_len + 1, self.vision_len + 1 + self.max_text_len, device=x.device, dtype=torch.int64
+                    )
+                    .unsqueeze(0)
+                    .repeat(ids_keep.shape[0], 1)
+                )
+                cls_ids = torch.zeros(ids_keep.shape[0], 1, device=x.device, dtype=torch.int64)
+                ids_keep = torch.cat([cls_ids, ids_keep, text_ids], dim=1)
+            elif is_flip == 1:
+                cls_ids = torch.zeros(ids_keep.shape[0], 1, device=x.device, dtype=torch.int64)
+                ids_keep = torch.cat([cls_ids, ids_keep], dim=1)
+            if encoder_padding_mask is not None:
+                encoder_padding_mask = torch.gather(encoder_padding_mask, dim=1, index=ids_keep)
+            if attn_mask is not None:
+                attn_mask = torch.gather(
+                    attn_mask, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, attn_mask.shape[-1])
+                )
+                attn_mask = torch.gather(attn_mask, dim=2, index=ids_keep.unsqueeze(1).repeat(1, attn_mask.shape[1], 1))
+            if multiway_split_position > 0:
+                multiway_split_position = ids_keep.shape[1] - self.max_text_len
+        x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
+
+        encoder_states = []
+
+        if return_all_hiddens:
+            encoder_states.append(x)
+
+        rel_pos_bias = None
+        if self.relative_position is not None:
+            rel_pos_bias = self.relative_position(batch_size=x.size(0), qlen=x.size(1), klen=x.size(1))
+
+        l_aux = []
+        for idx, layer in enumerate(self.layers):
+            x, l_aux_i = layer(
+                x,
+                encoder_padding_mask=encoder_padding_mask if incremental_state is None else None,
+                attn_mask=attn_mask,
+                rel_pos=rel_pos_bias,
+                multiway_split_position=multiway_split_position,
+                incremental_state=incremental_state[idx] if incremental_state is not None else None,
+            )
+            if return_all_hiddens:
+                assert encoder_states is not None
+                encoder_states.append(x)
+            l_aux.append(l_aux_i)
+
+        if multiway_split_position is not None:
+            assert self.args.multiway
+            no_sync_module_apply(self, set_split_position(multiway_split_position))
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        if not features_only and self.output_projection is not None:
+            x = self.output_projection(x)
+
+        return {
+            "encoder_out": x,
+            "encoder_embedding": encoder_embedding,
+            "encoder_padding_mask": encoder_padding_mask,
+            "encoder_states": encoder_states,
+            "l_aux": l_aux,
+            "multiway_split_position": multiway_split_position,
+        }

vlmo/torchscale/architecture/encoder_decoder.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch.nn as nn
+
+from vlmo.torchscale.architecture.decoder import Decoder
+from vlmo.torchscale.architecture.encoder import Encoder
+
+
+class EncoderDecoder(nn.Module):
+    def __init__(
+        self,
+        args,
+        encoder_embed_tokens=None,
+        encoder_embed_positions=None,
+        decoder_embed_tokens=None,
+        decoder_embed_positions=None,
+        output_projection=None,
+        **kwargs
+    ):
+        super().__init__()
+        self.args = args
+        if args.share_all_embeddings:
+            args.share_decoder_input_output_embed = True
+
+        self.encoder = Encoder(args, encoder_embed_tokens, encoder_embed_positions, is_encoder_decoder=True, **kwargs)
+
+        if args.share_all_embeddings and decoder_embed_tokens is None:
+            decoder_embed_tokens = self.encoder.embed_tokens
+
+        self.decoder = Decoder(
+            args, decoder_embed_tokens, decoder_embed_positions, output_projection, is_encoder_decoder=True, **kwargs
+        )
+
+    def forward(self, src_tokens, prev_output_tokens, return_all_hiddens=False, features_only=False, **kwargs):
+        encoder_out = self.encoder(src_tokens, return_all_hiddens=return_all_hiddens)
+        decoder_out = self.decoder(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            features_only=features_only,
+            return_all_hiddens=return_all_hiddens,
+        )
+        return decoder_out

vlmo/torchscale/architecture/utils.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
+import torch.nn as nn
+
+from vlmo.torchscale.component.multihead_attention import MultiheadAttention
+from vlmo.torchscale.component.multiway_network import MultiwayNetwork
+
+
+def init_bert_params(module):
+    def normal_(data):
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        if isinstance(module.q_proj, MultiwayNetwork):
+            normal_(module.q_proj.A.weight.data)
+            normal_(module.q_proj.B.weight.data)
+            normal_(module.k_proj.A.weight.data)
+            normal_(module.k_proj.B.weight.data)
+            normal_(module.v_proj.A.weight.data)
+            normal_(module.v_proj.B.weight.data)
+        else:
+            normal_(module.q_proj.weight.data)
+            normal_(module.k_proj.weight.data)
+            normal_(module.v_proj.weight.data)