gpt2

libai/inference/text_classification.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import oneflow as flow
+
+from libai.data.structures import DistTensorData, Instance
+from libai.inference.basic import BasePipeline
+
+
+class TextClassificationPipeline(BasePipeline):
+    def __init__(
+        self,
+        config_file,
+        data_parallel=None,
+        tensor_parallel=None,
+        pipeline_parallel=None,
+        pipeline_stage_id=None,
+        pipeline_num_layers=None,
+        model_path=None,
+        mode="libai",
+        **kwargs,
+    ):
+        super().__init__(
+            config_file,
+            data_parallel,
+            tensor_parallel,
+            pipeline_parallel,
+            pipeline_stage_id,
+            model_path,
+            pipeline_num_layers,
+            mode,
+            **kwargs,
+        )
+
+    def update_cfg(
+        self,
+        data_parallel=1,
+        tensor_parallel=1,
+        pipeline_parallel=1,
+        pipeline_stage_id=None,
+        pipeline_num_layers=None,
+    ):
+        super().update_cfg(
+            data_parallel,
+            tensor_parallel,
+            pipeline_parallel,
+            pipeline_stage_id,
+            pipeline_num_layers,
+        )
+        self.cfg.model.cfg.hidden_dropout_prob = 0.0
+        self.cfg.model.cfg.attention_probs_dropout_prob = 0.0
+
+        assert "num_labels" in self.cfg.model.cfg, "The model's config must contain num_labels"
+        if "label2id" not in self.cfg.model.cfg:
+            label2id = {"Label_" + str(i): i for i in range(self.cfg.model.cfg.num_labels)}
+            id2label = {ind: label for label, ind in label2id.items()}
+            self.cfg.model.cfg["label2id"] = label2id
+            self.cfg.model.cfg["id2label"] = id2label
+
+    def _parse_parameters(self, **pipeline_parameters):
+        preprocess_params = {}
+        forward_params = {}
+        postprocess_params = {**pipeline_parameters}
+        return preprocess_params, forward_params, postprocess_params
+
+    def preprocess(
+        self,
+        inputs,
+        pad: bool = False,
+        **kwargs,
+    ) -> dict:
+        # tokenizer encoder
+        input_ids = flow.tensor(np.array(self.tokenizer.encode(inputs)))
+        padding_mask = flow.tensor(np.ones(input_ids.shape), dtype=flow.bool)
+        # set batch size = 1
+        input_ids = input_ids.unsqueeze(0)
+        padding_mask = padding_mask.unsqueeze(0)
+
+        # to global tensor
+        model_input = Instance(
+            input_ids=DistTensorData(input_ids),
+            attention_mask=DistTensorData(padding_mask),
+        )
+        mdoel_input_dict = {}
+        for key, value in model_input.get_fields().items():
+            value.to_global()
+            mdoel_input_dict[key] = value.tensor
+        return mdoel_input_dict
+
+    def forward(self, mdoel_input_dict) -> dict:
+        model_outputs_dict = self.model(**mdoel_input_dict)
+        return model_outputs_dict
+
+    def postprocess(
+        self, model_outputs_dict, function_to_apply=None, return_all_scores=False, **kwargs
+    ) -> dict:
+        # prepare
+        num_labels = self.cfg.model.cfg.num_labels
+        if function_to_apply is not None:
+            function_to_apply = function_to_apply.lower()
+            assert function_to_apply in [
+                "sigmoid",
+                "softmax",
+                "none",
+            ], f"Unrecognized `function_to_apply` argument: {function_to_apply}"
+        else:
+            if num_labels == 1:
+                function_to_apply = "sigmoid"
+            elif num_labels > 1:
+                function_to_apply = "softmax"
+
+        # process, logits: [num_labels]
+        logits = model_outputs_dict["logits"][0]
+
+        if function_to_apply == "sigmoid":
+            scores = flow.sigmoid(logits)
+        elif function_to_apply == "softmax":
+            scores = flow.softmax(logits)
+        else:
+            scores = logits
+        scores = scores.detach().numpy()
+        if return_all_scores:
+            return [
+                {"label": self.cfg.model.cfg.id2label[i], "score": score.item()}
+                for i, score in enumerate(scores)
+            ]
+        else:
+            return {
+                "label": self.cfg.model.cfg.id2label[scores.argmax().item()],
+                "score": scores.max().item(),
+            }

libai/inference/text_generation.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from libai.inference.basic import BasePipeline
+from libai.utils import distributed as dist
+
+
+class TextGenerationPipeline(BasePipeline):
+    def load_pretrain_weight(self, libai_cfg_model, model_path, mode="huggingface"):
+        """load pretrained model.
+
+        Args:
+            libai_cfg_model (libai.models): Lazy config Model in Libai, you can import it
+                by `from libai.config.configs.common.models.bert
+                    import pretrain_model as libai_cfg_model`
+            model_path (str): The directory path of pretrained model,
+        """
+        if mode == "huggingface":
+            from projects.MT5.utils.mt5_loader import T5LoaderHuggerFace
+
+            model_loader = T5LoaderHuggerFace(
+                libai_cfg_model,
+                libai_cfg_model.cfg,
+                model_path,
+                hidden_dropout_prob=0.0,
+                attention_probs_dropout_prob=0.0,
+                embedding_dropout_prob=0.0,
+            )
+            return model_loader.load()
+        elif mode == "libai":
+            from projects.MT5.utils.mt5_loader import T5LoaderLibai
+
+            model_loader = T5LoaderLibai(
+                libai_cfg_model,
+                libai_cfg_model.cfg,
+                model_path,
+            )
+            return model_loader.load()
+        elif mode == "random":
+            from libai.engine import DefaultTrainer
+
+            return DefaultTrainer.build_model(self.cfg)
+        else:
+            raise NotImplementedError
+
+    def _parse_parameters(self, **pipeline_parameters):
+        preprocess_params = {}
+        forward_params = {**pipeline_parameters}
+        postprocess_params = {}
+
+        return preprocess_params, forward_params, postprocess_params
+
+    def preprocess(
+        self,
+        inputs,
+        pad: bool = False,
+        **kwargs,
+    ) -> dict:
+        # tokenizer encoder
+        encoder_ids = self.tokenizer.encode(inputs, return_tensors="of", is_global=True)
+
+        encoder_input_dict = {
+            "encoder_ids": encoder_ids,
+        }
+
+        return encoder_input_dict
+
+    def forward(self, encoder_input_dict, **kwargs) -> dict:
+        outputs = self.model.generate(encoder_input_dict["encoder_ids"], **kwargs)
+        return {"return_ids": outputs}
+
+    def postprocess(self, model_output_dict, **kwargs) -> dict:
+        return_ids = model_output_dict["return_ids"]
+        records = [
+            {"generated_text": self.tokenizer.decode(return_ids[i], skip_special_tokens=True)}
+            for i in range(return_ids.size(0))
+        ]
+        return records
+
+
+if __name__ == "__main__":
+    pipeline = TextGenerationPipeline(
+        "/path/to/libai/projects/MT5/configs/t5_inference.py",
+        data_parallel=1,
+        tensor_parallel=2,
+        pipeline_parallel=2,
+        pipeline_stage_id=[0] * 12 + [1] * 12,
+        pipeline_num_layers=12 * 2,
+        model_path="/path/to/t5-base",
+        mode="huggingface",
+    )
+
+    text = ["summarize: She is a student, She is tall, She loves study"]
+    dict1 = pipeline(text)
+    if dist.is_main_process():
+        print(dict1)

libai/inference/utils/imagenet_class.py

+IMAGENET_LABELS = [
+    "tench, Tinca tinca",
+    "goldfish, Carassius auratus",
+    "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",  # noqa: E501
+    "tiger shark, Galeocerdo cuvieri",
+    "hammerhead, hammerhead shark",
+    "electric ray, crampfish, numbfish, torpedo",
+    "stingray",
+    "cock",
+    "hen",
+    "ostrich, Struthio camelus",
+    "brambling, Fringilla montifringilla",
+    "goldfinch, Carduelis carduelis",
+    "house finch, linnet, Carpodacus mexicanus",
+    "junco, snowbird",
+    "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
+    "robin, American robin, Turdus migratorius",
+    "bulbul",
+    "jay",
+    "magpie",
+    "chickadee",
+    "water ouzel, dipper",
+    "kite",
+    "bald eagle, American eagle, Haliaeetus leucocephalus",
+    "vulture",
+    "great grey owl, great gray owl, Strix nebulosa",
+    "European fire salamander, Salamandra salamandra",
+    "common newt, Triturus vulgaris",
+    "eft",
+    "spotted salamander, Ambystoma maculatum",
+    "axolotl, mud puppy, Ambystoma mexicanum",
+    "bullfrog, Rana catesbeiana",
+    "tree frog, tree-frog",
+    "tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui",
+    "loggerhead, loggerhead turtle, Caretta caretta",
+    "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",  # noqa: E501
+    "mud turtle",
+    "terrapin",
+    "box turtle, box tortoise",
+    "banded gecko",
+    "common iguana, iguana, Iguana iguana",
+    "American chameleon, anole, Anolis carolinensis",
+    "whiptail, whiptail lizard",
+    "agama",
+    "frilled lizard, Chlamydosaurus kingi",
+    "alligator lizard",
+    "Gila monster, Heloderma suspectum",
+    "green lizard, Lacerta viridis",
+    "African chameleon, Chamaeleo chamaeleon",
+    "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",  # noqa: E501
+    "African crocodile, Nile crocodile, Crocodylus niloticus",
+    "American alligator, Alligator mississipiensis",
+    "triceratops",
+    "thunder snake, worm snake, Carphophis amoenus",
+    "ringneck snake, ring-necked snake, ring snake",
+    "hognose snake, puff adder, sand viper",
+    "green snake, grass snake",
+    "king snake, kingsnake",
+    "garter snake, grass snake",
+    "water snake",
+    "vine snake",
+    "night snake, Hypsiglena torquata",
+    "boa constrictor, Constrictor constrictor",
+    "rock python, rock snake, Python sebae",
+    "Indian cobra, Naja naja",
+    "green mamba",
+    "sea snake",
+    "horned viper, cerastes, sand viper, horned asp, Cerastes cornutus",
+    "diamondback, diamondback rattlesnake, Crotalus adamanteus",
+    "sidewinder, horned rattlesnake, Crotalus cerastes",
+    "trilobite",
+    "harvestman, daddy longlegs, Phalangium opilio",
+    "scorpion",
+    "black and gold garden spider, Argiope aurantia",
+    "barn spider, Araneus cavaticus",
+    "garden spider, Aranea diademata",
+    "black widow, Latrodectus mactans",
+    "tarantula",
+    "wolf spider, hunting spider",
+    "tick",
+    "centipede",
+    "black grouse",
+    "ptarmigan",
+    "ruffed grouse, partridge, Bonasa umbellus",
+    "prairie chicken, prairie grouse, prairie fowl",
+    "peacock",
+    "quail",
+    "partridge",
+    "African grey, African gray, Psittacus erithacus",
+    "macaw",
+    "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
+    "lorikeet",
+    "coucal",
+    "bee eater",
+    "hornbill",
+    "hummingbird",
+    "jacamar",
+    "toucan",
+    "drake",
+    "red-breasted merganser, Mergus serrator",
+    "goose",
+    "black swan, Cygnus atratus",
+    "tusker",
+    "echidna, spiny anteater, anteater",
+    "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",  # noqa: E501
+    "wallaby, brush kangaroo",
+    "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",  # noqa: E501
+    "wombat",
+    "jellyfish",
+    "sea anemone, anemone",
+    "brain coral",
+    "flatworm, platyhelminth",
+    "nematode, nematode worm, roundworm",
+    "conch",
+    "snail",
+    "slug",
+    "sea slug, nudibranch",
+    "chiton, coat-of-mail shell, sea cradle, polyplacophore",
+    "chambered nautilus, pearly nautilus, nautilus",
+    "Dungeness crab, Cancer magister",
+    "rock crab, Cancer irroratus",
+    "fiddler crab",
+    "king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica",  # noqa: E501
+    "American lobster, Northern lobster, Maine lobster, Homarus americanus",  # noqa: E501
+    "spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish",  # noqa: E501
+    "crayfish, crawfish, crawdad, crawdaddy",
+    "hermit crab",
+    "isopod",
+    "white stork, Ciconia ciconia",
+    "black stork, Ciconia nigra",
+    "spoonbill",
+    "flamingo",
+    "little blue heron, Egretta caerulea",
+    "American egret, great white heron, Egretta albus",
+    "bittern",
+    "crane",
+    "limpkin, Aramus pictus",
+    "European gallinule, Porphyrio porphyrio",
+    "American coot, marsh hen, mud hen, water hen, Fulica americana",
+    "bustard",
+    "ruddy turnstone, Arenaria interpres",
+    "red-backed sandpiper, dunlin, Erolia alpina",
+    "redshank, Tringa totanus",
+    "dowitcher",
+    "oystercatcher, oyster catcher",
+    "pelican",
+    "king penguin, Aptenodytes patagonica",
+    "albatross, mollymawk",
+    "grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus",  # noqa: E501
+    "killer whale, killer, orca, grampus, sea wolf, Orcinus orca",
+    "dugong, Dugong dugon",
+    "sea lion",
+    "Chihuahua",
+    "Japanese spaniel",
+    "Maltese dog, Maltese terrier, Maltese",
+    "Pekinese, Pekingese, Peke",
+    "Shih-Tzu",
+    "Blenheim spaniel",
+    "papillon",
+    "toy terrier",
+    "Rhodesian ridgeback",
+    "Afghan hound, Afghan",
+    "basset, basset hound",
+    "beagle",
+    "bloodhound, sleuthhound",
+    "bluetick",
+    "black-and-tan coonhound",
+    "Walker hound, Walker foxhound",
+    "English foxhound",
+    "redbone",
+    "borzoi, Russian wolfhound",
+    "Irish wolfhound",
+    "Italian greyhound",
+    "whippet",
+    "Ibizan hound, Ibizan Podenco",
+    "Norwegian elkhound, elkhound",
+    "otterhound, otter hound",
+    "Saluki, gazelle hound",
+    "Scottish deerhound, deerhound",
+    "Weimaraner",
+    "Staffordshire bullterrier, Staffordshire bull terrier",
+    "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",  # noqa: E501
+    "Bedlington terrier",
+    "Border terrier",
+    "Kerry blue terrier",
+    "Irish terrier",
+    "Norfolk terrier",
+    "Norwich terrier",
+    "Yorkshire terrier",
+    "wire-haired fox terrier",
+    "Lakeland terrier",
+    "Sealyham terrier, Sealyham",
+    "Airedale, Airedale terrier",
+    "cairn, cairn terrier",
+    "Australian terrier",
+    "Dandie Dinmont, Dandie Dinmont terrier",
+    "Boston bull, Boston terrier",
+    "miniature schnauzer",
+    "giant schnauzer",
+    "standard schnauzer",
+    "Scotch terrier, Scottish terrier, Scottie",
+    "Tibetan terrier, chrysanthemum dog",
+    "silky terrier, Sydney silky",
+    "soft-coated wheaten terrier",
+    "West Highland white terrier",
+    "Lhasa, Lhasa apso",
+    "flat-coated retriever",
+    "curly-coated retriever",
+    "golden retriever",
+    "Labrador retriever",
+    "Chesapeake Bay retriever",
+    "German short-haired pointer",
+    "vizsla, Hungarian pointer",
+    "English setter",
+    "Irish setter, red setter",
+    "Gordon setter",
+    "Brittany spaniel",
+    "clumber, clumber spaniel",
+    "English springer, English springer spaniel",
+    "Welsh springer spaniel",
+    "cocker spaniel, English cocker spaniel, cocker",
+    "Sussex spaniel",
+    "Irish water spaniel",
+    "kuvasz",
+    "schipperke",
+    "groenendael",
+    "malinois",
+    "briard",
+    "kelpie",
+    "komondor",
+    "Old English sheepdog, bobtail",
+    "Shetland sheepdog, Shetland sheep dog, Shetland",
+    "collie",
+    "Border collie",
+    "Bouvier des Flandres, Bouviers des Flandres",
+    "Rottweiler",
+    "German shepherd, German shepherd dog, German police dog, alsatian",
+    "Doberman, Doberman pinscher",
+    "miniature pinscher",
+    "Greater Swiss Mountain dog",
+    "Bernese mountain dog",
+    "Appenzeller",
+    "EntleBucher",
+    "boxer",
+    "bull mastiff",
+    "Tibetan mastiff",
+    "French bulldog",
+    "Great Dane",
+    "Saint Bernard, St Bernard",
+    "Eskimo dog, husky",
+    "malamute, malemute, Alaskan malamute",
+    "Siberian husky",
+    "dalmatian, coach dog, carriage dog",
+    "affenpinscher, monkey pinscher, monkey dog",
+    "basenji",
+    "pug, pug-dog",
+    "Leonberg",
+    "Newfoundland, Newfoundland dog",
+    "Great Pyrenees",
+    "Samoyed, Samoyede",
+    "Pomeranian",
+    "chow, chow chow",
+    "keeshond",
+    "Brabancon griffon",
+    "Pembroke, Pembroke Welsh corgi",
+    "Cardigan, Cardigan Welsh corgi",
+    "toy poodle",
+    "miniature poodle",
+    "standard poodle",
+    "Mexican hairless",
+    "timber wolf, grey wolf, gray wolf, Canis lupus",
+    "white wolf, Arctic wolf, Canis lupus tundrarum",
+    "red wolf, maned wolf, Canis rufus, Canis niger",
+    "coyote, prairie wolf, brush wolf, Canis latrans",
+    "dingo, warrigal, warragal, Canis dingo",
+    "dhole, Cuon alpinus",
+    "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
+    "hyena, hyaena",
+    "red fox, Vulpes vulpes",
+    "kit fox, Vulpes macrotis",
+    "Arctic fox, white fox, Alopex lagopus",
+    "grey fox, gray fox, Urocyon cinereoargenteus",
+    "tabby, tabby cat",
+    "tiger cat",
+    "Persian cat",
+    "Siamese cat, Siamese",
+    "Egyptian cat",
+    "cougar, puma, catamount, mountain lion, painter, panther, Felis concolor",  # noqa: E501
+    "lynx, catamount",
+    "leopard, Panthera pardus",
+    "snow leopard, ounce, Panthera uncia",
+    "jaguar, panther, Panthera onca, Felis onca",
+    "lion, king of beasts, Panthera leo",
+    "tiger, Panthera tigris",
+    "cheetah, chetah, Acinonyx jubatus",
+    "brown bear, bruin, Ursus arctos",
+    "American black bear, black bear, Ursus americanus, Euarctos americanus",  # noqa: E501
+    "ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus",
+    "sloth bear, Melursus ursinus, Ursus ursinus",
+    "mongoose",
+    "meerkat, mierkat",
+    "tiger beetle",
+    "ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle",
+    "ground beetle, carabid beetle",
+    "long-horned beetle, longicorn, longicorn beetle",
+    "leaf beetle, chrysomelid",
+    "dung beetle",
+    "rhinoceros beetle",
+    "weevil",
+    "fly",
+    "bee",
+    "ant, emmet, pismire",
+    "grasshopper, hopper",
+    "cricket",
+    "walking stick, walkingstick, stick insect",
+    "cockroach, roach",
+    "mantis, mantid",
+    "cicada, cicala",
+    "leafhopper",
+    "lacewing, lacewing fly",
+    "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",  # noqa: E501
+    "damselfly",
+    "admiral",
+    "ringlet, ringlet butterfly",
+    "monarch, monarch butterfly, milkweed butterfly, Danaus plexippus",
+    "cabbage butterfly",
+    "sulphur butterfly, sulfur butterfly",
+    "lycaenid, lycaenid butterfly",
+    "starfish, sea star",
+    "sea urchin",
+    "sea cucumber, holothurian",
+    "wood rabbit, cottontail, cottontail rabbit",
+    "hare",
+    "Angora, Angora rabbit",
+    "hamster",
+    "porcupine, hedgehog",
+    "fox squirrel, eastern fox squirrel, Sciurus niger",
+    "marmot",
+    "beaver",
+    "guinea pig, Cavia cobaya",
+    "sorrel",
+    "zebra",
+    "hog, pig, grunter, squealer, Sus scrofa",
+    "wild boar, boar, Sus scrofa",
+    "warthog",
+    "hippopotamus, hippo, river horse, Hippopotamus amphibius",
+    "ox",
+    "water buffalo, water ox, Asiatic buffalo, Bubalus bubalis",
+    "bison",
+    "ram, tup",
+    "bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis",  # noqa: E501
+    "ibex, Capra ibex",
+    "hartebeest",
+    "impala, Aepyceros melampus",
+    "gazelle",
+    "Arabian camel, dromedary, Camelus dromedarius",
+    "llama",
+    "weasel",
+    "mink",
+    "polecat, fitch, foulmart, foumart, Mustela putorius",
+    "black-footed ferret, ferret, Mustela nigripes",
+    "otter",
+    "skunk, polecat, wood pussy",
+    "badger",
+    "armadillo",
+    "three-toed sloth, ai, Bradypus tridactylus",
+    "orangutan, orang, orangutang, Pongo pygmaeus",
+    "gorilla, Gorilla gorilla",
+    "chimpanzee, chimp, Pan troglodytes",
+    "gibbon, Hylobates lar",
+    "siamang, Hylobates syndactylus, Symphalangus syndactylus",
+    "guenon, guenon monkey",
+    "patas, hussar monkey, Erythrocebus patas",
+    "baboon",
+    "macaque",
+    "langur",
+    "colobus, colobus monkey",
+    "proboscis monkey, Nasalis larvatus",
+    "marmoset",
+    "capuchin, ringtail, Cebus capucinus",
+    "howler monkey, howler",
+    "titi, titi monkey",
+    "spider monkey, Ateles geoffroyi",
+    "squirrel monkey, Saimiri sciureus",
+    "Madagascar cat, ring-tailed lemur, Lemur catta",
+    "indri, indris, Indri indri, Indri brevicaudatus",
+    "Indian elephant, Elephas maximus",
+    "African elephant, Loxodonta africana",
+    "lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens",
+    "giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca",
+    "barracouta, snoek",
+    "eel",
+    "coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch",  # noqa: E501
+    "rock beauty, Holocanthus tricolor",
+    "anemone fish",
+    "sturgeon",
+    "gar, garfish, garpike, billfish, Lepisosteus osseus",
+    "lionfish",
+    "puffer, pufferfish, blowfish, globefish",
+    "abacus",
+    "abaya",
+    "academic gown, academic robe, judge's robe",
+    "accordion, piano accordion, squeeze box",
+    "acoustic guitar",
+    "aircraft carrier, carrier, flattop, attack aircraft carrier",
+    "airliner",
+    "airship, dirigible",
+    "altar",
+    "ambulance",
+    "amphibian, amphibious vehicle",
+    "analog clock",
+    "apiary, bee house",
+    "apron",
+    "ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin",  # noqa: E501
+    "assault rifle, assault gun",
+    "backpack, back pack, knapsack, packsack, rucksack, haversack",
+    "bakery, bakeshop, bakehouse",
+    "balance beam, beam",
+    "balloon",
+    "ballpoint, ballpoint pen, ballpen, Biro",
+    "Band Aid",
+    "banjo",
+    "bannister, banister, balustrade, balusters, handrail",
+    "barbell",
+    "barber chair",
+    "barbershop",
+    "barn",
+    "barometer",
+    "barrel, cask",
+    "barrow, garden cart, lawn cart, wheelbarrow",
+    "baseball",
+    "basketball",
+    "bassinet",
+    "bassoon",
+    "bathing cap, swimming cap",
+    "bath towel",
+    "bathtub, bathing tub, bath, tub",
+    "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",  # noqa: E501
+    "beacon, lighthouse, beacon light, pharos",
+    "beaker",
+    "bearskin, busby, shako",
+    "beer bottle",
+    "beer glass",
+    "bell cote, bell cot",
+    "bib",
+    "bicycle-built-for-two, tandem bicycle, tandem",
+    "bikini, two-piece",
+    "binder, ring-binder",
+    "binoculars, field glasses, opera glasses",
+    "birdhouse",
+    "boathouse",
+    "bobsled, bobsleigh, bob",
+    "bolo tie, bolo, bola tie, bola",
+    "bonnet, poke bonnet",
+    "bookcase",
+    "bookshop, bookstore, bookstall",
+    "bottlecap",
+    "bow",
+    "bow tie, bow-tie, bowtie",
+    "brass, memorial tablet, plaque",
+    "brassiere, bra, bandeau",
+    "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
+    "breastplate, aegis, egis",
+    "broom",
+    "bucket, pail",
+    "buckle",
+    "bulletproof vest",
+    "bullet train, bullet",
+    "butcher shop, meat market",
+    "cab, hack, taxi, taxicab",
+    "caldron, cauldron",
+    "candle, taper, wax light",
+    "cannon",
+    "canoe",
+    "can opener, tin opener",
+    "cardigan",
+    "car mirror",
+    "carousel, carrousel, merry-go-round, roundabout, whirligig",
+    "carpenter's kit, tool kit",
+    "carton",
+    "car wheel",
+    "cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM",  # noqa: E501
+    "cassette",
+    "cassette player",
+    "castle",
+    "catamaran",
+    "CD player",
+    "cello, violoncello",
+    "cellular telephone, cellular phone, cellphone, cell, mobile phone",
+    "chain",
+    "chainlink fence",
+    "chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour",  # noqa: E501
+    "chain saw, chainsaw",
+    "chest",
+    "chiffonier, commode",
+    "chime, bell, gong",
+    "china cabinet, china closet",
+    "Christmas stocking",
+    "church, church building",
+    "cinema, movie theater, movie theatre, movie house, picture palace",
+    "cleaver, meat cleaver, chopper",
+    "cliff dwelling",
+    "cloak",
+    "clog, geta, patten, sabot",
+    "cocktail shaker",
+    "coffee mug",
+    "coffeepot",
+    "coil, spiral, volute, whorl, helix",
+    "combination lock",
+    "computer keyboard, keypad",
+    "confectionery, confectionary, candy store",
+    "container ship, containership, container vessel",
+    "convertible",
+    "corkscrew, bottle screw",
+    "cornet, horn, trumpet, trump",
+    "cowboy boot",
+    "cowboy hat, ten-gallon hat",
+    "cradle",
+    "crane",
+    "crash helmet",
+    "crate",
+    "crib, cot",
+    "Crock Pot",
+    "croquet ball",
+    "crutch",
+    "cuirass",
+    "dam, dike, dyke",
+    "desk",
+    "desktop computer",
+    "dial telephone, dial phone",
+    "diaper, nappy, napkin",
+    "digital clock",
+    "digital watch",
+    "dining table, board",
+    "dishrag, dishcloth",
+    "dishwasher, dish washer, dishwashing machine",
+    "disk brake, disc brake",
+    "dock, dockage, docking facility",
+    "dogsled, dog sled, dog sleigh",
+    "dome",
+    "doormat, welcome mat",
+    "drilling platform, offshore rig",
+    "drum, membranophone, tympan",
+    "drumstick",
+    "dumbbell",
+    "Dutch oven",
+    "electric fan, blower",
+    "electric guitar",
+    "electric locomotive",
+    "entertainment center",
+    "envelope",
+    "espresso maker",
+    "face powder",
+    "feather boa, boa",
+    "file, file cabinet, filing cabinet",
+    "fireboat",
+    "fire engine, fire truck",
+    "fire screen, fireguard",
+    "flagpole, flagstaff",
+    "flute, transverse flute",
+    "folding chair",
+    "football helmet",
+    "forklift",
+    "fountain",
+    "fountain pen",
+    "four-poster",
+    "freight car",
+    "French horn, horn",
+    "frying pan, frypan, skillet",
+    "fur coat",
+    "garbage truck, dustcart",
+    "gasmask, respirator, gas helmet",
+    "gas pump, gasoline pump, petrol pump, island dispenser",
+    "goblet",
+    "go-kart",
+    "golf ball",
+    "golfcart, golf cart",
+    "gondola",
+    "gong, tam-tam",
+    "gown",
+    "grand piano, grand",
+    "greenhouse, nursery, glasshouse",
+    "grille, radiator grille",
+    "grocery store, grocery, food market, market",
+    "guillotine",
+    "hair slide",
+    "hair spray",
+    "half track",
+    "hammer",
+    "hamper",
+    "hand blower, blow dryer, blow drier, hair dryer, hair drier",
+    "hand-held computer, hand-held microcomputer",
+    "handkerchief, hankie, hanky, hankey",
+    "hard disc, hard disk, fixed disk",
+    "harmonica, mouth organ, harp, mouth harp",
+    "harp",
+    "harvester, reaper",
+    "hatchet",
+    "holster",
+    "home theater, home theatre",
+    "honeycomb",
+    "hook, claw",
+    "hoopskirt, crinoline",
+    "horizontal bar, high bar",
+    "horse cart, horse-cart",
+    "hourglass",
+    "iPod",
+    "iron, smoothing iron",
+    "jack-o'-lantern",
+    "jean, blue jean, denim",
+    "jeep, landrover",
+    "jersey, T-shirt, tee shirt",
+    "jigsaw puzzle",
+    "jinrikisha, ricksha, rickshaw",
+    "joystick",
+    "kimono",
+    "knee pad",
+    "knot",
+    "lab coat, laboratory coat",
+    "ladle",
+    "lampshade, lamp shade",
+    "laptop, laptop computer",
+    "lawn mower, mower",
+    "lens cap, lens cover",
+    "letter opener, paper knife, paperknife",
+    "library",
+    "lifeboat",
+    "lighter, light, igniter, ignitor",
+    "limousine, limo",
+    "liner, ocean liner",
+    "lipstick, lip rouge",
+    "Loafer",
+    "lotion",
+    "loudspeaker, speaker, speaker unit, loudspeaker system, speaker system",  # noqa: E501
+    "loupe, jeweler's loupe",
+    "lumbermill, sawmill",
+    "magnetic compass",
+    "mailbag, postbag",
+    "mailbox, letter box",
+    "maillot",
+    "maillot, tank suit",
+    "manhole cover",
+    "maraca",
+    "marimba, xylophone",
+    "mask",
+    "matchstick",
+    "maypole",
+    "maze, labyrinth",
+    "measuring cup",
+    "medicine chest, medicine cabinet",
+    "megalith, megalithic structure",
+    "microphone, mike",
+    "microwave, microwave oven",
+    "military uniform",
+    "milk can",
+    "minibus",
+    "miniskirt, mini",
+    "minivan",
+    "missile",
+    "mitten",
+    "mixing bowl",
+    "mobile home, manufactured home",
+    "Model T",
+    "modem",
+    "monastery",
+    "monitor",
+    "moped",
+    "mortar",
+    "mortarboard",
+    "mosque",
+    "mosquito net",
+    "motor scooter, scooter",
+    "mountain bike, all-terrain bike, off-roader",
+    "mountain tent",
+    "mouse, computer mouse",
+    "mousetrap",
+    "moving van",
+    "muzzle",
+    "nail",
+    "neck brace",
+    "necklace",
+    "nipple",
+    "notebook, notebook computer",
+    "obelisk",
+    "oboe, hautboy, hautbois",
+    "ocarina, sweet potato",
+    "odometer, hodometer, mileometer, milometer",
+    "oil filter",
+    "organ, pipe organ",
+    "oscilloscope, scope, cathode-ray oscilloscope, CRO",
+    "overskirt",
+    "oxcart",
+    "oxygen mask",
+    "packet",
+    "paddle, boat paddle",
+    "paddlewheel, paddle wheel",
+    "padlock",
+    "paintbrush",
+    "pajama, pyjama, pj's, jammies",
+    "palace",
+    "panpipe, pandean pipe, syrinx",
+    "paper towel",
+    "parachute, chute",
+    "parallel bars, bars",
+    "park bench",
+    "parking meter",
+    "passenger car, coach, carriage",
+    "patio, terrace",
+    "pay-phone, pay-station",
+    "pedestal, plinth, footstall",
+    "pencil box, pencil case",
+    "pencil sharpener",
+    "perfume, essence",
+    "Petri dish",
+    "photocopier",
+    "pick, plectrum, plectron",
+    "pickelhaube",
+    "picket fence, paling",
+    "pickup, pickup truck",
+    "pier",
+    "piggy bank, penny bank",
+    "pill bottle",
+    "pillow",
+    "ping-pong ball",
+    "pinwheel",
+    "pirate, pirate ship",
+    "pitcher, ewer",
+    "plane, carpenter's plane, woodworking plane",
+    "planetarium",
+    "plastic bag",
+    "plate rack",
+    "plow, plough",
+    "plunger, plumber's helper",
+    "Polaroid camera, Polaroid Land camera",
+    "pole",
+    "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",  # noqa: E501
+    "poncho",
+    "pool table, billiard table, snooker table",
+    "pop bottle, soda bottle",
+    "pot, flowerpot",
+    "potter's wheel",
+    "power drill",
+    "prayer rug, prayer mat",
+    "printer",
+    "prison, prison house",
+    "projectile, missile",
+    "projector",
+    "puck, hockey puck",
+    "punching bag, punch bag, punching ball, punchball",
+    "purse",
+    "quill, quill pen",
+    "quilt, comforter, comfort, puff",
+    "racer, race car, racing car",
+    "racket, racquet",
+    "radiator",
+    "radio, wireless",
+    "radio telescope, radio reflector",
+    "rain barrel",
+    "recreational vehicle, RV, R.V.",
+    "reel",
+    "reflex camera",
+    "refrigerator, icebox",
+    "remote control, remote",
+    "restaurant, eating house, eating place, eatery",
+    "revolver, six-gun, six-shooter",
+    "rifle",
+    "rocking chair, rocker",
+    "rotisserie",
+    "rubber eraser, rubber, pencil eraser",
+    "rugby ball",
+    "rule, ruler",
+    "running shoe",
+    "safe",
+    "safety pin",
+    "saltshaker, salt shaker",
+    "sandal",
+    "sarong",
+    "sax, saxophone",
+    "scabbard",
+    "scale, weighing machine",
+    "school bus",
+    "schooner",
+    "scoreboard",
+    "screen, CRT screen",
+    "screw",
+    "screwdriver",
+    "seat belt, seatbelt",
+    "sewing machine",
+    "shield, buckler",
+    "shoe shop, shoe-shop, shoe store",
+    "shoji",
+    "shopping basket",
+    "shopping cart",
+    "shovel",
+    "shower cap",
+    "shower curtain",
+    "ski",
+    "ski mask",
+    "sleeping bag",
+    "slide rule, slipstick",
+    "sliding door",
+    "slot, one-armed bandit",
+    "snorkel",
+    "snowmobile",
+    "snowplow, snowplough",
+    "soap dispenser",
+    "soccer ball",
+    "sock",
+    "solar dish, solar collector, solar furnace",
+    "sombrero",
+    "soup bowl",
+    "space bar",
+    "space heater",
+    "space shuttle",
+    "spatula",
+    "speedboat",
+    "spider web, spider's web",
+    "spindle",
+    "sports car, sport car",
+    "spotlight, spot",
+    "stage",
+    "steam locomotive",
+    "steel arch bridge",
+    "steel drum",
+    "stethoscope",
+    "stole",
+    "stone wall",
+    "stopwatch, stop watch",
+    "stove",
+    "strainer",
+    "streetcar, tram, tramcar, trolley, trolley car",
+    "stretcher",
+    "studio couch, day bed",
+    "stupa, tope",
+    "submarine, pigboat, sub, U-boat",
+    "suit, suit of clothes",
+    "sundial",
+    "sunglass",
+    "sunglasses, dark glasses, shades",
+    "sunscreen, sunblock, sun blocker",
+    "suspension bridge",
+    "swab, swob, mop",
+    "sweatshirt",
+    "swimming trunks, bathing trunks",
+    "swing",
+    "switch, electric switch, electrical switch",
+    "syringe",
+    "table lamp",
+    "tank, army tank, armored combat vehicle, armoured combat vehicle",
+    "tape player",
+    "teapot",
+    "teddy, teddy bear",
+    "television, television system",
+    "tennis ball",
+    "thatch, thatched roof",
+    "theater curtain, theatre curtain",
+    "thimble",
+    "thresher, thrasher, threshing machine",
+    "throne",
+    "tile roof",
+    "toaster",
+    "tobacco shop, tobacconist shop, tobacconist",
+    "toilet seat",
+    "torch",
+    "totem pole",
+    "tow truck, tow car, wrecker",
+    "toyshop",
+    "tractor",
+    "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",  # noqa: E501
+    "tray",
+    "trench coat",
+    "tricycle, trike, velocipede",
+    "trimaran",
+    "tripod",
+    "triumphal arch",
+    "trolleybus, trolley coach, trackless trolley",
+    "trombone",
+    "tub, vat",
+    "turnstile",
+    "typewriter keyboard",
+    "umbrella",
+    "unicycle, monocycle",
+    "upright, upright piano",
+    "vacuum, vacuum cleaner",
+    "vase",
+    "vault",
+    "velvet",
+    "vending machine",
+    "vestment",
+    "viaduct",
+    "violin, fiddle",
+    "volleyball",
+    "waffle iron",
+    "wall clock",
+    "wallet, billfold, notecase, pocketbook",
+    "wardrobe, closet, press",
+    "warplane, military plane",
+    "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
+    "washer, automatic washer, washing machine",
+    "water bottle",
+    "water jug",
+    "water tower",
+    "whiskey jug",
+    "whistle",
+    "wig",
+    "window screen",
+    "window shade",
+    "Windsor tie",
+    "wine bottle",
+    "wing",
+    "wok",
+    "wooden spoon",
+    "wool, woolen, woollen",
+    "worm fence, snake fence, snake-rail fence, Virginia fence",
+    "wreck",
+    "yawl",
+    "yurt",
+    "web site, website, internet site, site",
+    "comic book",
+    "crossword puzzle, crossword",
+    "street sign",
+    "traffic light, traffic signal, stoplight",
+    "book jacket, dust cover, dust jacket, dust wrapper",
+    "menu",
+    "plate",
+    "guacamole",
+    "consomme",
+    "hot pot, hotpot",
+    "trifle",
+    "ice cream, icecream",
+    "ice lolly, lolly, lollipop, popsicle",
+    "French loaf",
+    "bagel, beigel",
+    "pretzel",
+    "cheeseburger",
+    "hotdog, hot dog, red hot",
+    "mashed potato",
+    "head cabbage",
+    "broccoli",
+    "cauliflower",
+    "zucchini, courgette",
+    "spaghetti squash",
+    "acorn squash",
+    "butternut squash",
+    "cucumber, cuke",
+    "artichoke, globe artichoke",
+    "bell pepper",
+    "cardoon",
+    "mushroom",
+    "Granny Smith",
+    "strawberry",
+    "orange",
+    "lemon",
+    "fig",
+    "pineapple, ananas",
+    "banana",
+    "jackfruit, jak, jack",
+    "custard apple",
+    "pomegranate",
+    "hay",
+    "carbonara",
+    "chocolate sauce, chocolate syrup",
+    "dough",
+    "meat loaf, meatloaf",
+    "pizza, pizza pie",
+    "potpie",
+    "burrito",
+    "red wine",
+    "espresso",
+    "cup",
+    "eggnog",
+    "alp",
+    "bubble",
+    "cliff, drop, drop-off",
+    "coral reef",
+    "geyser",
+    "lakeside, lakeshore",
+    "promontory, headland, head, foreland",
+    "sandbar, sand bar",
+    "seashore, coast, seacoast, sea-coast",
+    "valley, vale",
+    "volcano",
+    "ballplayer, baseball player",
+    "groom, bridegroom",
+    "scuba diver",
+    "rapeseed",
+    "daisy",
+    "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",  # noqa: E501
+    "corn",
+    "acorn",
+    "hip, rose hip, rosehip",
+    "buckeye, horse chestnut, conker",
+    "coral fungus",
+    "agaric",
+    "gyromitra",
+    "stinkhorn, carrion fungus",
+    "earthstar",
+    "hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa",  # noqa: E501
+    "bolete",
+    "ear, spike, capitulum",
+    "toilet tissue, toilet paper, bathroom tissue",
+]

libai/layers/__init__.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .activation import build_activation
+from .cross_entropy import ParallelCrossEntropyLoss
+from .embedding import Embedding, SinePositionalEmbedding, VocabEmbedding, PatchEmbedding
+from .layer_norm import LayerNorm, RMSLayerNorm
+from .linear import Linear, Linear1D
+from .lm_logits import LMLogits
+from .mlp import MLP
+from .transformer_layer import TransformerLayer
+from .attention import MultiheadAttention
+from .droppath import DropPath, drop_path
+
+__all__ = [
+    "Embedding",
+    "VocabEmbedding",
+    "SinePositionalEmbedding",
+    "PatchEmbedding",
+    "build_activation",
+    "Linear",
+    "Linear1D",
+    "MLP",
+    "LayerNorm",
+    "RMSLayerNorm",
+    "TransformerLayer",
+    "MultiheadAttention",
+    "ParallelCrossEntropyLoss",
+    "LMLogits",
+    "drop_path",
+    "DropPath",
+]

libai/layers/activation.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from enum import Enum
+from typing import Optional
+
+import oneflow as flow
+from oneflow import nn
+
+
+class Activation(str, Enum):
+    SquaredReLU = "squared_relu"
+    GeLU = "gelu"
+    GeLUTanh = "gelu_tanh"
+    LeakyReLU = "leaky_relu"
+    ReLU = "relu"
+    Tanh = "tanh"
+    QuickGELU = "quick_gelu"
+
+
+# For unit testing / parity comparisons, probably not the fastest way
+class SquaredReLU(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        x_ = flow._C.relu(x)
+        return x_ * x_
+
+
+class Passthrough(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        return x
+
+
+class GeLUTanh(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        """When the approximate argument is 'tanh', Gelu is estimated with:
+        0.5 * x * (1.0 + flow.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * flow.pow(x, 3.0))))
+        """
+        return flow.nn.functional.gelu(x, approximate="tanh")
+
+
+class QuickGELU(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        return x * flow.sigmoid(1.702 * x)
+
+
+def build_activation(activation: Optional[Activation]):
+    """
+    Fetching activation layers by name, e.g.,
+    ``build_activation("gelu")`` returns ``nn.GELU()`` module.
+    """
+    if not activation:
+        return Passthrough()
+
+    return {
+        Activation.ReLU: nn.ReLU,
+        Activation.GeLU: nn.GELU,
+        Activation.GeLUTanh: GeLUTanh,
+        Activation.LeakyReLU: nn.LeakyReLU,
+        Activation.SquaredReLU: SquaredReLU,
+        Activation.Tanh: nn.Tanh,
+        Activation.QuickGELU: QuickGELU,
+    }[activation]()

libai/layers/attention.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+import math
+from typing import Tuple
+
+import oneflow as flow
+from oneflow import nn
+
+from .linear import Linear
+
+
+class AttnMaskType(enum.Enum):
+    padding = 1
+    causal = 2
+
+
+class MultiheadAttention(nn.Module):
+    """Multi-head attention layer, support self attention and cross attention.
+
+    Args:
+        hidden_size: size of hidden state.
+        num_attention_heads: number of attention heads.
+        is_cross_attention: used to specify whether it is self attention or cross attention.
+            Defaults to False.
+        attention_dropout_prob: dropout probability of attention weights.
+            Defaults to 0.0.
+        output_dropout_prob: dropout probability of output. Defaults to 0.0.
+        init_method: method to initialize the input layer weights.
+            Defaults to ``init.xavier_normal_``.
+        output_layer_init_method: method to initialize the output layer weights.
+            If None, use ``init_method``.
+        bias_dropout_fusion: whether to fuse add bias and dropout.
+            Defaults to False.
+        scale_mask_softmax_fusion: whether to fuse scale, mask and softmax.
+            Defaults to False.
+        apply_query_key_layer_scaling: if `True`, scaling the attention score by layer index.
+            Defaults to False.
+        layer_idx: a layer_idx sign which determines the placements.
+            It will be used in pipeline parallelism. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        is_cross_attention=False,
+        attention_dropout_prob=0.0,
+        output_dropout_prob=0.0,
+        init_method=nn.init.xavier_normal_,
+        output_layer_init_method=None,
+        bias_dropout_fusion=False,
+        scale_mask_softmax_fusion=False,
+        apply_query_key_layer_scaling=False,
+        attn_mask_type=AttnMaskType.padding,
+        *,
+        layer_idx=0
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        if output_layer_init_method is None:
+            output_layer_init_method = init_method
+
+        assert (
+            hidden_size % num_attention_heads == 0
+        ), "hidden_size must be divisible by num_attention_heads."
+
+        self.num_heads = num_attention_heads
+        self.head_size = hidden_size // num_attention_heads
+        self.attn_mask_type = attn_mask_type
+
+        self.attention_dropout_prob = attention_dropout_prob
+        self.dropout = nn.Dropout(p=attention_dropout_prob)
+        self.norm_factor = 1.0 / math.sqrt(float(self.head_size))
+        self.coeff = None
+        if apply_query_key_layer_scaling:
+            self.coeff = layer_idx + 1
+            self.norm_factor /= self.coeff
+
+        self.is_cross_attention = is_cross_attention
+        self.scale_mask_softmax_fusion = scale_mask_softmax_fusion
+        self.bias_dropout_fusion = bias_dropout_fusion
+
+        if self.bias_dropout_fusion:
+            self.output_dropout_prob = output_dropout_prob
+        else:
+            self.output_dropout = nn.Dropout(p=output_dropout_prob)
+
+        if self.is_cross_attention:
+            self.query = Linear(
+                self.hidden_size,
+                self.hidden_size,
+                parallel="col",
+                init_method=init_method,
+                layer_idx=layer_idx,
+            )
+            self.key_value = Linear(
+                self.hidden_size,
+                self.hidden_size * 2,
+                parallel="col",
+                init_method=init_method,
+                layer_idx=layer_idx,
+            )
+        else:
+            self.query_key_value = Linear(
+                self.hidden_size,
+                self.hidden_size * 3,
+                parallel="col",
+                init_method=init_method,
+                layer_idx=layer_idx,
+            )
+
+        self.dense = Linear(
+            self.hidden_size,
+            self.hidden_size,
+            parallel="row",
+            init_method=output_layer_init_method,
+            skip_bias_add=self.bias_dropout_fusion,
+            layer_idx=layer_idx,
+        )
+
+    def forward(
+        self,
+        hidden_states: flow.Tensor,
+        encoder_states: flow.Tensor = None,
+        attention_mask: flow.Tensor = None,
+        past_key_value: Tuple[flow.Tensor, flow.Tensor] = None,
+        use_cache: bool = False,
+    ):
+        """
+
+        Args:
+            hidden_states (flow.Tensor): shape is [bsz, tgt_len, hidden_size].
+            encoder_states (flow.Tensor, optional): shape is [bsz, src_len, hidden_size].
+                Defaults to None.
+            attention_mask (flow.Tensor, optional): shape is [bsz, 1, tgt_len, src_len].
+                It should be the combination of padding mask and casual mask.
+                It is the padding mask of source input when used with self-attention in encoder.
+                And it is the combination of padding mask of target input and casual mask when
+                used with self-attention in decoder. It is the padding mask of source input when
+                used with cross-attention in decoder.
+                Defaults to None.
+            past_key_value (Tuple[flow.Tensor, flow.Tensor], optional): tuple of key and value,
+                each shape is [bsz, num_heads, src_len, head_size]. Defaults to None.
+            use_cache (bool, optional): it will be set to True, when the model is in the inference
+                phase and used for incremental decoding. Defaults to False.
+        """
+
+        # hidden_states, encoder_states: [S(0), B]
+        # attention_mask: [S(0), B]
+
+        if encoder_states is not None:
+            encoder_states = encoder_states.to_global(placement=hidden_states.placement)
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.to_global(placement=hidden_states.placement)
+
+        bsz, tgt_len = hidden_states.size()[:2]
+
+        if self.is_cross_attention:
+            # if it is cross attention, key and value should be calculated only once, and the
+            # result can be reused.
+            query = self.query(hidden_states)
+            query = query.view(bsz, -1, self.num_heads, self.head_size)
+            query = query.permute(0, 2, 1, 3)
+            if past_key_value is not None:
+                key, value = past_key_value
+            elif encoder_states is not None:
+                key_value = self.key_value(encoder_states)
+                key_value = key_value.view(bsz, -1, self.num_heads, 2 * self.head_size)
+                key_value = key_value.permute(0, 2, 1, 3)
+                key, value = flow.chunk(key_value, chunks=2, dim=-1)
+            else:
+                raise ValueError(
+                    "past_key_value and encoder_states cannot be None at the same time."
+                )
+        else:
+            # if it is self attention, query, key, and value are all obtained from hidden_states.
+            # when in the inference phase of an incremental decoder,
+            # hidden_states is the last-added state,
+            # the full key and value could be obtained by concatenating with past_key_value.
+            query_key_value = self.query_key_value(hidden_states)
+            query_key_value = query_key_value.view(bsz, -1, self.num_heads, 3 * self.head_size)
+            query_key_value = query_key_value.permute(
+                0, 2, 1, 3
+            )  # [bsz, num_heads, src_len, 3 * head_size]
+            query, key, value = flow.chunk(query_key_value, chunks=3, dim=-1)
+            if past_key_value is not None:
+                past_key, past_value = past_key_value
+                key = flow.cat((past_key.type_as(key), key), dim=2)
+                value = flow.cat((past_value.type_as(value), value), dim=2)
+
+        # query, key, value: [S(0), S(1)], shape: [bsz, num_heads, seq_length, head_size]
+        if use_cache:
+            past_key_value = (key, value)
+
+        # [bsz, num_heads, tgt_len, src_len] with [S(0), S(1)]
+        attention_scores = flow.matmul(query, key, transpose_b=True, alpha=self.norm_factor)
+
+        # [S(0), S(1)] x [S(0), B] = [S(0), S(1)]
+        if attention_mask is not None:
+            if self.scale_mask_softmax_fusion:
+                if self.attn_mask_type == AttnMaskType.padding:
+                    attention_mask = (
+                        attention_mask.expand_as(attention_scores) if use_cache else attention_mask
+                    )
+                    attention_weights = flow._C.fused_scale_mask_softmax_dropout(
+                        attention_scores,
+                        attention_mask,
+                        fill_value=-10000.0,
+                        scale=self.coeff,
+                        p=self.attention_dropout_prob,
+                    )[0]
+            else:
+                if self.coeff is not None:
+                    attention_scores *= self.coeff
+                attention_scores = flow.mul(attention_scores, attention_mask)
+                attention_scores = attention_scores - 10000.0 * (1 - attention_mask)
+                # TODO(xingyu.liao): graph will occur `where_scalar` errors
+                # when using `masked_fill`
+                # attention_scores = attention_scores.masked_fill(1 - attention_mask, -10000.0)
+                attention_weights = flow.softmax(attention_scores, dim=-1)
+                # [bsz, num_heads, tgt_len, src_len]
+                attention_weights = self.dropout(attention_weights)
+        else:
+            if self.scale_mask_softmax_fusion and self.attn_mask_type == AttnMaskType.causal:
+                attention_weights = flow._C.fused_scale_tril_softmax_mask_scale(
+                    attention_scores,
+                    p=self.attention_dropout_prob,
+                    diagonal=0,
+                    tril_scale_value=self.coeff,
+                    tril_fill_value=-10000.0,
+                )[0]
+            else:
+                attention_weights = flow.softmax(attention_scores, dim=-1)
+                # [bsz, num_heads, tgt_len, src_len]
+                attention_weights = self.dropout(attention_weights)
+
+        # Context shape: [bsz, num_heads, tgt_len, head_size] with [S(0), S(1)]
+        context = flow.matmul(attention_weights, value)
+        # Change shape: [bsz, num_heads, tgt_len, head_size] -> [bsz, tgt_len, num_heads, head_size]
+        context = context.transpose(1, 2)
+
+        # Concat multi-head results from
+        # [bsz, tgt_len, num_heads, head_size] -> [bsz, tgt_len, num_heads * head_size]
+        # SBP sign: [S(0), S(2)]
+        # [S(0), S(2)] x [B, S(0)] = [S(0), P] -> [S(0), B]
+        output = self.dense(context.flatten(2))
+
+        if self.bias_dropout_fusion:
+            output, bias = output
+            output = flow._C.fused_bias_add_dropout(
+                output, bias, p=self.output_dropout_prob, axis=output.ndim - 1
+            )
+        else:
+            output = self.output_dropout(output)
+
+        if use_cache:
+            output = (output, past_key_value)
+
+        return output
+
+    def extra_repr(self) -> str:
+        return "hidden_size={}, num_heads={}, is_cross_attention={}".format(
+            self.hidden_size,
+            self.num_heads,
+            self.is_cross_attention,
+        )

libai/layers/cross_entropy.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import oneflow as flow
+from oneflow import nn
+
+
+class ParallelCrossEntropyLoss(nn.Module):
+    """This criterion acts like :class:`~flow.nn.CrossEntropyLoss` except it will
+    execute distributed cross entropy loss computation cross different GPUs.
+    """
+
+    def forward(self, logits: flow.Tensor, target: flow.Tensor):
+        """Function for the distributed cross entropy.
+
+        Args:
+            logits (flow.Tensor): vocab_parallel_logits with shape
+                (batch_size, seq_length, vocab_size) and sbp signature is [S(0), S(2)].
+            target (flow.Tensor): target with shape (batch_size, seq_length) and
+                sbp signature is [S(0), B].
+        """
+        assert logits.ndim == 3
+        assert target.ndim == 2
+        assert logits.shape[0:2] == target.shape
+
+        target = target.to_global(placement=logits.placement)
+
+        # Change -1 in target to 0 because sparse_softmax_cross_entropy don't accept -1
+        target = target * (target >= 0)
+
+        lm_loss = flow._C.sparse_softmax_cross_entropy(
+            logits.view(-1, logits.shape[-1]),
+            target.view(-1),
+        )
+        return lm_loss

libai/layers/droppath.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import oneflow as flow
+import oneflow.nn as nn
+
+
+def drop_path(x, drop_prob: float = 0.5, training: bool = False, scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+
+    # similar opeartion to new_tensor(shape).bernoulli_(keep_prob)
+    random_tensor = flow.rand(*shape, dtype=x.dtype, sbp=x.sbp, placement=x.placement)
+    random_tensor = (random_tensor < keep_prob).to(flow.float32)
+
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor = random_tensor / keep_prob
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

libai/layers/embedding.py

+# coding=utf-8
+# Copyright 2021 The OneFlow Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import oneflow as flow
+from oneflow import nn
+from oneflow.nn import init
+
+from libai.utils import distributed as dist
+
+
+class Embedding(nn.Module):
+    """Construct the trainable embedding module, which does not support parallelization.
+    This can be used for positional embedding and token type embedding.
+
+    Arguments:
+        num_embeddings: size of vocabulary.
+        embedding_dim: dimension of embeddings.
+        padding_idx: pad index. Defaults to None.
+        init_method: method to initialize weights. Defaults to ``flow.nn.init.xavier_normal_``.
+        amp_enabled: fp16 option for embedding weight. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        num_embeddings,
+        embedding_dim,
+        padding_idx=None,
+        init_method=init.xavier_normal_,
+        amp_enabled=False,
+        layer_idx=0,
+    ):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        if padding_idx is not None:
+            if padding_idx > 0:
+                assert (
+                    padding_idx < self.num_embeddings
+                ), "Padding_idx must be within num_embeddings"
+            elif padding_idx < 0:
+                assert (
+                    padding_idx >= -self.num_embeddings
+                ), "Padding_idx must be within num_embeddings"
+                padding_idx = self.num_embeddings + padding_idx
+        self.padding_idx = padding_idx
+        self.init_method = init_method
+        self.amp_enabled = amp_enabled
+
+        assert num_embeddings > 0
+        self.weight = nn.Parameter(
+            flow.empty(
+                (num_embeddings, embedding_dim),
+                dtype=flow.float32,
+                placement=dist.get_layer_placement(layer_idx),
+                sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            )
+        )
+        self.init_method(self.weight)
+        # FIXME(lxy): Fill padding_idx is not supported in nd_sbp right now.
+        # self._fill_padding_idx_with_zero()
+
+    def forward(self, input_ids):
+        weight = flow._C.amp_white_identity(self.weight) if self.amp_enabled else self.weight
+        # embeddings with sbp sign: [B, B]
+        #   [B, B] x [S(0), B] --> [S(0), B]
+        #     ↑         ↑              ↑
+        #   embed    pos_ids       pos_embed
+        input_embeds = flow._C.gather(weight, input_ids, axis=0)
+        return input_embeds
+
+    def _fill_padding_idx_with_zero(self) -> None:
+        if self.padding_idx is not None:
+            with flow.no_grad():
+                self.weight[self.padding_idx] = flow.zeros(
+                    self.embedding_dim,
+                    placement=dist.get_layer_placement(0),
+                    sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                )
+
+    def extra_repr(self) -> str:
+        s = "num_embeddings={num_embeddings}, embedding_dim={embedding_dim}"
+        if self.padding_idx is not None:
+            s += ", padding_idx={padding_idx}"
+        return s.format(**self.__dict__)
+
+
+class VocabEmbedding(nn.Module):
+    """Construct the word embeddings, which may be split along vocabulary dimension.
+
+    Arguments:
+        num_embeddings: size of vocabulary.
+        embedding_dim: dimension of embeddings.
+        padding_idx: pad index. Defaults to None.
+        init_method: method to initialize weights. Defaults to ``flow.nn.init.xavier_normal_``.
+        amp_enabled: fp16 option for embedding weight. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        num_embeddings,
+        embedding_dim,
+        padding_idx=None,
+        init_method=init.xavier_normal_,
+        amp_enabled=False,
+    ):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        if padding_idx is not None:
+            if padding_idx > 0:
+                assert (
+                    padding_idx < self.num_embeddings
+                ), "Padding_idx must be within num_embeddings"
+            elif padding_idx < 0:
+                assert (
+                    padding_idx >= -self.num_embeddings
+                ), "Padding_idx must be within num_embeddings"
+                padding_idx = self.num_embeddings + padding_idx
+        self.padding_idx = padding_idx
+        self.init_method = init_method
+        self.amp_enabled = amp_enabled
+
+        # Word token embedding shape with (vocab_size, hidden_size)
+        # sbp: [B, S(0)]
+        self.weight = nn.Parameter(
+            flow.empty(
+                (num_embeddings, embedding_dim),
+                dtype=flow.float32,
+                placement=dist.get_layer_placement(0),
+                sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.split(0)]),
+            )
+        )
+        # Initialize the word embedding
+        self.init_method(self.weight)
+        # FIXME(Lxy): Fill padding_idx is not supported in nd_sbp right now.
+        # self._fill_padding_idx_with_zero()
+
+    def forward(self, input_ids):
+        weight = flow._C.amp_white_identity(self.weight) if self.amp_enabled else self.weight
+        # input_ids with shape (batch_size, seq_len), and sbp sign: [S(0), B]
+
+        # Gather forward sbp sign
+        # [B, S(0)] x [S(0), B] --> [S(0), P]
+        #     ↑           ↑            ↑
+        #   embed  input_ids    input_embeds
+        input_embeds = flow._C.gather(weight, input_ids, axis=0)
+        # Set the embeds sbp from [S(0), P] --> [S(0), B] to get complete embedding results.
+        input_embeds = input_embeds.to_global(sbp=dist.get_hidden_sbp())
+
+        return input_embeds
+
+    def _fill_padding_idx_with_zero(self) -> None:
+        if self.padding_idx is not None:
+            with flow.no_grad():
+                self.weight[self.padding_idx] = flow.zeros(
+                    self.embedding_dim,
+                    placement=dist.get_layer_placement(0),
+                    sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                )
+
+    def extra_repr(self) -> str:
+        s = "num_embeddings={num_embeddings}, embedding_dim={embedding_dim}"
+        if self.padding_idx is not None:
+            s += ", padding_idx={padding_idx}"
+        return s.format(**self.__dict__)
+
+
+class SinePositionalEmbedding(nn.Module):
+    """Construct the sinusoidal positional embeddings.
+
+    Arguments:
+        num_embeddings: size of vocabulary.
+        embedding_dim: dimension of embeddings.
+    """
+
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+
+        self.embedding_dim = embedding_dim
+        self.num_embeddings = num_embeddings
+
+        position_embedding = flow.zeros(
+            num_embeddings,
+            embedding_dim,
+            dtype=flow.float32,
+            placement=dist.get_layer_placement(0),
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+        )
+        position = flow._C.global_arange(
+            start=0,
+            end=num_embeddings,
+            placement=dist.get_layer_placement(0),
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            dtype=flow.float32,
+        ).unsqueeze(1)
+        position_range = flow._C.global_arange(
+            start=0,
+            end=embedding_dim,
+            step=2,
+            placement=dist.get_layer_placement(0),
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            dtype=flow.float32,
+        )
+        div_term = flow.exp(position_range * (-math.log(10000.0) / embedding_dim))
+
+        position_embedding[:, 0::2] = flow.sin(position * div_term)
+        position_embedding[:, 1::2] = flow.cos(position * div_term)
+        self.register_buffer("position_embedding", position_embedding)
+
+    def forward(self, position_ids):
+        position_embeds = flow._C.gather(self.position_embedding, position_ids, axis=0)
+        return position_embeds
+
+    def extra_repr(self) -> str:
+        s = "num_embeddings={num_embeddings}, embedding_dim={embedding_dim}"
+        return s.format(**self.__dict__)
+
+
+class PatchEmbedding(nn.Module):
+    """2D Image to Patch Embedding
+
+    Arguments:
+        img_size: size of input image. Default to 224.
+        patch_size: embedded patch size. Default to 16.
+        in_chans: input channel's size. Default to 3.
+        embed_dim: dimension of embedded patch. Default to 768.
+        norm_layer: normalization patch embedding or not. Default to None.
+        flatten: flatten patch embedding or keep the 2-D shape. Default to True.
+        layer_idx: A layer_idx sign which determines the placement. It will be used in pipeline
+        parallelism. Default to 0.
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+        flatten=True,
+        *,
+        layer_idx=0,
+    ):
+        super().__init__()
+        img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
+        patch_size = patch_size if isinstance(patch_size, tuple) else (patch_size, patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        ).to_global(
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            placement=dist.get_layer_placement(layer_idx),
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0]
+        ), f"Input image height ({H}) doesn't match model ({self.img_size[0]})."
+        assert (
+            W == self.img_size[1]
+        ), f"Input image width ({W}) doesn't match model ({self.img_size[1]})."
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x