glm

CONTRIBUTING.md

+## Contributing to LiBai
\ No newline at end of file

LICENSE

+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README.md

+# GLM(General Language Model Pretraining with Autoregressive Blank Infilling )
+## 模型介绍
+2017 年, Google 提出了 Transformer 架构, 随后 BERT 、GPT、T5等预训练模型不断涌现, 并在各项任务中都不断刷新 SOTA 纪录。2022年, 清华提出了 GLM 模型(https://github.com/THUDM/GLM), 不同于上述预训练模型架构，它采用了一种自回归的空白填充方法, 在 NLP 领域三种主要的任务（自然语言理解、无条件生成、有条件生成）上都取得了不错的结果。
+
+在LiBai中主要实现了GLM推理部分的工作。
+
+## GLM-Inference
+当模型规模过于庞大，单个 GPU 设备无法容纳大规模模型参数时，便捷好用的分布式训练和推理需求就相继出现，业内也随之推出相应的工具。
+
+基于 OneFlow 构建的 LiBai 模型库让分布式上手难度降到最低，用户不需要关注模型如何分配在不同的显卡设备，只需要修改几个配置数据就可以设置不同的分布式策略。当然，加速性能更是出众。
+
+用 LiBai 搭建的 GLM 可以便捷地实现model parallel + pipeline parallel推理, 很好地解决单卡放不下大规模模型的问题。
+
+### 分布式推理具有天然优势
+
+要知道，模型的参数其实就是许多 tensor，也就是以矩阵的形式出现，大模型的参数也就是大矩阵，并行策略就是把大矩阵分为多个小矩阵，并分配到不同的显卡或不同的设备上，基础的 LinearLayer 在LiBai中的实现代码如下：
+
+```python
+class Linear1D(nn.Module):
+    def __init__(self, in_features, out_features, parallel="data", layer_idx=0, ...):
+        super().__init__()
+
+        if parallel == "col":
+            weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.split(0)])
+        elif parallel == "row":
+            weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.split(1)])
+        elif parallel == "data":
+            weight_sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])
+        else:
+            raise KeyError(f"{parallel} is not supported! Only support ('data', 'row' and 'col')")
+
+        self.weight = flow.nn.Parameter(
+            flow.empty(
+                (out_features, in_features),
+                dtype=flow.float32,
+                placement=dist.get_layer_placement(layer_idx),  # for pipeline parallelism placement
+                sbp=weight_sbp,
+            )
+        )
+        init_method(self.weight)
+        ...
+    
+    def forward(self, x):
+        ...
+```
+
+在这里，用户可选择去如何切分 Linear 层的矩阵，如何切分数据矩阵，而OneFlow 中的 SBP 控制竖着切、横着切以及其他拆分矩阵的方案（模型并行、数据并行），以及通过设置 Placement 来控制这个 LinearLayer 是放在第几张显卡上（流水并行）。
+
+所以，根据 LiBai 中各种 layer 的设计原理以及基于 OneFlow 中 tensor 自带的 SBP 和 Placement 属性的天然优势，使得用户搭建的模型能够很简单地就实现数据并行、模型并行以及流水并行操作。
+
+## GLM-10B-chinese推理
+### 环境配置
+提供[光源](https://www.sourcefind.cn/#/service-details)拉取的训练以及推理的docker镜像：image.sourcefind.cn:5000/dcu/admin/base/oneflow:0.9.1-centos7.6-dtk-22.10.1-py39-latest
+
+    cd libai
+    pip3 install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple
+    pip3 install pybind11 -i https://mirrors.aliyun.com/pypi/simple
+    pip3 install -e . -i https://mirrors.aliyun.com/pypi/simple
+需要先准备好模型权重：https://huggingface.co/THUDM/glm-10b-chinese/tree/main
+
+### Glm-10b-chinese的文件结构
+
+```python
+$ tree data
+path/to/glm-10b-chinese
+├── added_tokens.json
+├── cog-pretrain.model
+├── config.json
+└── pytorch_model.bin
+```
+
+### 推理
+
+采用1节点，4张DCU-Z100-16G，采用tp=2，pp=2的并行配置。
+
+运行以下代码：
+
+    cd projects/GLM
+    # 运行前修改 configs/glm_inference.py 中 `pad_token_id=50000, eos_token_id=50007, bos_token_id=None`
+    python3 -m oneflow.distributed.launch --nproc_per_node 4 demo.py
+
+demo.py如下：
+
+    # model parallel + pipeline parallel demo
+    
+    import oneflow as flow
+    from projects.GLM.tokenizer.glm_tokenizer import GLMChineseTokenzier
+    from libai.utils import distributed as dist
+    from projects.GLM.configs.glm_inference import cfg
+    from projects.GLM.modeling_glm import GLMForConditionalGeneration
+    from projects.GLM.utils.glm_loader import GLMLoaderHuggerFace
+    from omegaconf import DictConfig
+    import time
+    
+    # 只需简单配置并行方案
+    parallel_config = DictConfig(
+        dict(
+            data_parallel_size=1,
+            tensor_parallel_size=2,
+            pipeline_parallel_size=2,
+            pipeline_num_layers=2 * 24
+        )
+    )
+    dist.setup_dist_util(parallel_config)
+    
+    tokenizer = GLMChineseTokenzier.from_pretrained("glm-10b-chinese")
+    input_ids = tokenizer.encode(
+        [
+            "冬天，中国哪座城市最适合避寒？问题描述：能推荐一些国内适合冬天避寒的城市吗？回答用户：旅游爱好者 回答： [gMASK]"
+        ],
+        return_tensors="of",
+    )
+    inputs = {"input_ids": input_ids, "attention_mask": flow.ones(input_ids.size())}
+    inputs = tokenizer.build_inputs_for_generation(inputs, max_gen_length=128)
+    
+    sbp = dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])
+    placement = dist.get_layer_placement(0)
+    
+    loader = GLMLoaderHuggerFace(
+        GLMForConditionalGeneration, 
+        cfg, 
+        "glm-10b-chinese",
+        embedding_dropout_prob=0,
+        attention_dropout_prob=0,
+        output_dropout_prob=0,
+    )
+    
+    T1 = time.time()
+    model = loader.load()
+    T2 = time.time()
+    if dist.is_main_process():
+        print('模型加载时间:%s秒' % (T2 - T1))
+    
+    T3 = time.time()
+    outputs = model.generate(
+        inputs=inputs['input_ids'].to_global(sbp=sbp, placement=placement), 
+        position_ids=inputs['position_ids'].to_global(sbp=sbp, placement=placement), 
+        generation_attention_mask=inputs['generation_attention_mask'].to_global(sbp=sbp, placement=placement), 
+        max_length=128
+    )
+    T4 = time.time()
+    if dist.is_main_process():
+        print('model.generate: %s秒' % (T4 - T3))
+    
+    T5 = time.time()
+    res = tokenizer.decode(outputs[0])
+    T6 = time.time()
+    if dist.is_main_process():
+        print('tokenizer.decode: %s秒' % (T6 - T5))
+    
+    if dist.is_main_process():
+        print(res)
+
+输出：
+
+```
+>>>Total number of model parameters: 9,879,633,920
+模型加载时间:59.47162699699402秒
+model.generate: 72.28496813774109秒
+tokenizer.decode: 0.0698804759979248秒
+[CLS] 冬天,中国哪座城市最适合避寒?问题描述:能推荐一些国内适合冬天避寒的城市吗?回答用户:旅游爱好者 回答: [gMASK] <|endoftext|> <|startofpiece|> 避寒,当然是去海南呀!<n><n>海南的冬天,阳光明媚,温度适宜,而且空气清新,没有雾霾,没有沙尘暴,没有雾霾,没有雾霾!<n><n>海南的冬天,阳光明媚,温度适宜,而且空气清新,没有雾霾,没有沙尘暴,没有雾霾!<n><n>海南的冬天,阳光明媚,温度适宜,而且空气清新,没有雾霾,没有沙尘暴,没有雾霾!
+```
+
+## 性能和准确率数据
+
+使用的加速卡：4张DCU-Z100-16G：
+
+| bs | max_input_length | max_gen_length | model.generate耗时/s |
+| :------: | :------: | :------: | :------: |
+| 1 | 128 | 128 | 72.2 |
+| 1 | 512 | 512 | 201.3 |
+## 参考
+* https://github.com/Oneflow-Inc/libai
+* https://github.com/Oneflow-Inc/one-glm

changelog.md

+## Changelog
+
+### Beta 0.1.0 (22/03/2022)
+
+**New Features:**
+- Support Data Parallelism
+- Support 1D Tensor Parallelism
+- Support Pipeline Parallelism
+- Unified distributed Layers for both single-GPU and multi-GPU training
+- `LazyConfig` system for more flexible syntax and no predefined structures
+- Easy-to-use trainer and engine
+- Support both CV and NLP data processing
+- Mixed Precision Training
+- Activation Checkpointing
+- Gradient Accumulation
+- Gradient Clipping
+- Zero Redundancy Optimizer (ZeRO)
+
+**Supported Models:**
+- Support 3D parallel [BERT](https://arxiv.org/abs/1810.04805) model
+- Support 3D parallel [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) model
+- Support 3D parallel [T5](https://arxiv.org/abs/1910.10683) model
+- Support 3D parallel [Vision Transformer](https://arxiv.org/abs/2010.11929)
+- Support Data parallel [Swin Transformer](https://arxiv.org/abs/2103.14030) model
+- Support finetune task in [QQP project](/projects/QQP/)
+- Support text classification task in [text classification project](/projects/text_classification/)

configs/bert_classification.py

+from libai.config import LazyCall
+from libai.models.bert_model import BertForClassification
+from .common.models.bert import cfg as bert_cfg
+from .common.models.graph import graph
+from .common.train import train
+from .common.optim import optim
+from .common.data.bert_dataset import tokenization, dataloader
+
+vocab_file = "./data_test/bert_data/bert-base-chinese-vocab.txt"
+data_prefix = "./data_test/bert_data/loss_compara_content_sentence"
+
+dataloader.train.dataset[0].data_prefix = data_prefix
+dataloader.train.dataset[0].indexed_dataset.data_prefix = data_prefix
+
+bert_cfg.num_labels = 2
+bert_cfg.classifier_dropout = 0.1
+
+model = LazyCall(BertForClassification)(cfg=bert_cfg)
+tokenization.tokenizer.vocab_file = vocab_file
+
+model.cfg.vocab_size = 21128
+model.cfg.intermediate_size = 3072
+model.cfg.num_attention_heads = 12
+model.cfg.hidden_layers = 12
+model.cfg.hidden_size = 768
+
+train.amp.enabled = True
+train.activation_checkpoint.enabled = True
+train.dist.pipeline_num_layers = model.cfg.hidden_layers
+
+train.output_dir = "output/bert_classification_output"

configs/bert_large_pretrain.py

+from libai.config import LazyCall
+from libai.evaluation import PPLEvaluator
+from .common.models.bert import pretrain_model as model
+from .common.models.graph import graph
+from .common.train import train
+from .common.optim import optim
+from .common.data.bert_dataset import dataloader, tokenization
+
+vocab_file = "./data_test/bert_data/bert-base-chinese-vocab.txt"
+data_prefix = "./data_test/bert_data/loss_compara_content_sentence"
+
+tokenization.tokenizer.vocab_file = vocab_file
+dataloader.train.dataset[0].data_prefix = data_prefix
+dataloader.train.dataset[0].indexed_dataset.data_prefix = data_prefix
+dataloader.test[0].dataset.data_prefix = data_prefix
+dataloader.test[0].dataset.indexed_dataset.data_prefix = data_prefix
+
+# Bert-large model config
+model.cfg.num_attention_heads = 16
+model.cfg.hidden_size = 768
+model.cfg.hidden_layers = 8
+
+train.input_placement_device = "cpu"
+
+train.dist.pipeline_num_layers = model.cfg.hidden_layers
+
+train.train_micro_batch_size = 16
+
+train.amp.enabled = True
+
+for ds in dataloader.train.dataset:
+    ds.max_seq_length = model.cfg.max_position_embeddings
+
+train.evaluation.evaluator = LazyCall(PPLEvaluator)()
+
+train.output_dir = "output/bert_output"

configs/common/data/bert_dataset.py

+from libai.config import LazyCall
+from omegaconf import OmegaConf
+from libai.data import build_nlp_test_loader, build_nlp_train_val_test_loader
+from libai.data.datasets import BertDataset
+from libai.data.data_utils import get_indexed_dataset
+
+from libai.tokenizer import BertTokenizer
+
+
+tokenization = OmegaConf.create()
+
+tokenization.tokenizer = LazyCall(BertTokenizer)(
+    vocab_file="bert-base-chinese-vocab.txt",
+    do_lower_case=True,
+    do_chinese_wwm=True,
+)
+tokenization.append_eod = False
+tokenization.make_vocab_size_divisible_by = 128
+
+dataloader = OmegaConf.create()
+
+dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
+    dataset=[
+        LazyCall(BertDataset)(
+            name="bert",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_seq_length=512,
+            mask_lm_prob=0.15,
+            short_seq_prob=0.1,
+            binary_head=True,
+            seed=1234,
+            masking_style="bert-cn-wwm",
+        ),
+    ],
+    train_val_test_num_samples=None,  # a hint for deferred assignment
+    splits=[[949.0, 50.0, 1.0]],
+    weights=[1.0],
+    num_workers=4,
+)
+
+dataloader.test = [
+    LazyCall(build_nlp_test_loader)(
+        dataset=LazyCall(BertDataset)(
+            name="bert",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_num_samples=10,
+            max_seq_length=512,
+            mask_lm_prob=0.15,
+            short_seq_prob=0.1,
+            binary_head=True,
+            seed=1234,
+            masking_style="bert-cn-wwm",
+        ),
+        test_batch_size=4,
+    )
+]

configs/common/data/cifar100.py

+from omegaconf import OmegaConf
+from flowvision import transforms
+from flowvision.data.mixup import Mixup
+from flowvision.transforms import InterpolationMode
+from flowvision.transforms.functional import str_to_interp_mode
+
+from libai.data.datasets import CIFAR100Dataset
+from libai.data.build import build_image_train_loader, build_image_test_loader
+from libai.config import LazyCall
+
+# mean and std of cifar100 dataset
+CIFAR100_TRAIN_MEAN = (0.5070751592371323, 0.48654887331495095, 0.4409178433670343)
+CIFAR100_TRAIN_STD = (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)
+
+train_aug = LazyCall(transforms.Compose)(
+    transforms=[
+        LazyCall(transforms.RandomResizedCrop)(
+            size=(224, 224),
+            scale=(0.08, 1.0),
+            ratio=(3.0 / 4.0, 4.0 / 3.0),
+            interpolation=str_to_interp_mode("bicubic"),
+        ),
+        LazyCall(transforms.RandomHorizontalFlip)(),
+        LazyCall(transforms.ToTensor)(),
+        LazyCall(transforms.Normalize)(mean=CIFAR100_TRAIN_MEAN, std=CIFAR100_TRAIN_STD),
+    ]
+)
+
+test_aug = LazyCall(transforms.Compose)(
+    transforms=[
+        LazyCall(transforms.Resize)(
+            size=256,
+            interpolation=InterpolationMode.BICUBIC,
+        ),
+        LazyCall(transforms.CenterCrop)(
+            size=224,
+        ),
+        LazyCall(transforms.ToTensor)(),
+        LazyCall(transforms.Normalize)(
+            mean=CIFAR100_TRAIN_MEAN,
+            std=CIFAR100_TRAIN_STD,
+        ),
+    ]
+)
+
+
+# Dataloader config
+dataloader = OmegaConf.create()
+dataloader.train = LazyCall(build_image_train_loader)(
+    dataset=[
+        LazyCall(CIFAR100Dataset)(
+            root="./",
+            train=True,
+            download=True,
+            transform=train_aug,
+        ),
+    ],
+    num_workers=4,
+    mixup_func=LazyCall(Mixup)(
+        mixup_alpha=0.8,
+        cutmix_alpha=1.0,
+        prob=1.0,
+        switch_prob=0.5,
+        mode="batch",
+        num_classes=100,
+    ),
+)
+
+dataloader.test = [
+    LazyCall(build_image_test_loader)(
+        dataset=LazyCall(CIFAR100Dataset)(
+            root="./",
+            train=False,
+            download=True,
+            transform=test_aug,
+        ),
+        num_workers=4,
+    )
+]

configs/common/data/gpt_dataset.py

+from libai.config import LazyCall
+from omegaconf import OmegaConf
+from libai.data import build_nlp_test_loader, build_nlp_train_val_test_loader
+from libai.data.datasets import GPT2Dataset
+from libai.data.data_utils import get_indexed_dataset
+
+from libai.tokenizer import GPT2Tokenizer
+
+
+tokenization = OmegaConf.create()
+
+tokenization.tokenizer = LazyCall(GPT2Tokenizer)(
+    vocab_file="/workspace/data/gpt_dataset/gpt2-vocab.json",
+    merges_file="/workspace/data/gpt_dataset/gpt2-merges.txt",
+    do_lower_case=True,
+    do_chinese_wwm=True,
+)
+tokenization.append_eod = False
+tokenization.make_vocab_size_divisible_by = 128
+
+dataloader = OmegaConf.create()
+
+dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
+    dataset=[
+        LazyCall(GPT2Dataset)(
+            name="gpt-2",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_seq_length=1024,
+            seed=1234,
+        ),
+    ],
+    train_val_test_num_samples=None,  # a hint for deferred assignment
+    splits=[[949.0, 50.0, 1.0]],
+    weights=[1.0],
+    num_workers=4,
+)
+
+dataloader.test = [
+    LazyCall(build_nlp_test_loader)(
+        dataset=LazyCall(GPT2Dataset)(
+            name="gpt-2",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_seq_length=1024,
+            max_num_samples=10,
+            seed=1234,
+        ),
+        test_batch_size=4,
+    )
+]

configs/common/data/imagenet.py

+from omegaconf import OmegaConf
+from flowvision import transforms
+from flowvision.transforms import InterpolationMode
+from flowvision.transforms.functional import str_to_interp_mode
+from flowvision.data.constants import (
+    IMAGENET_DEFAULT_MEAN,
+    IMAGENET_DEFAULT_STD,
+)
+from flowvision.data.auto_augment import rand_augment_transform
+from flowvision.data.random_erasing import RandomErasing
+
+from libai.config import LazyCall
+from libai.data.datasets import ImageNetDataset
+from libai.data.build import build_image_train_loader, build_image_test_loader
+
+train_aug = LazyCall(transforms.Compose)(
+    transforms=[
+        LazyCall(transforms.RandomResizedCrop)(
+            size=224,
+            scale=(0.08, 1.0),
+            ratio=(3.0 / 4.0, 4.0 / 3.0),
+            interpolation=InterpolationMode.BICUBIC,
+        ),
+        LazyCall(transforms.RandomHorizontalFlip)(p=0.5),
+        LazyCall(rand_augment_transform)(
+            config_str="rand-m9-mstd0.5-inc1",
+            hparams=dict(
+                translate_const=int(224 * 0.45),
+                img_mean=tuple([min(255, round(255 * x)) for x in IMAGENET_DEFAULT_MEAN]),
+                interpolation=str_to_interp_mode("bicubic"),
+            ),
+        ),
+        LazyCall(transforms.ToTensor)(),
+        LazyCall(transforms.Normalize)(
+            mean=IMAGENET_DEFAULT_MEAN,
+            std=IMAGENET_DEFAULT_STD,
+        ),
+        LazyCall(RandomErasing)(
+            probability=0.25,
+            mode="pixel",
+            max_count=1,
+            num_splits=0,
+            device="cpu",
+        ),
+    ]
+)
+
+
+test_aug = LazyCall(transforms.Compose)(
+    transforms=[
+        LazyCall(transforms.Resize)(
+            size=256,
+            interpolation=InterpolationMode.BICUBIC,
+        ),
+        LazyCall(transforms.CenterCrop)(
+            size=224,
+        ),
+        LazyCall(transforms.ToTensor)(),
+        LazyCall(transforms.Normalize)(
+            mean=IMAGENET_DEFAULT_MEAN,
+            std=IMAGENET_DEFAULT_STD,
+        ),
+    ]
+)
+
+
+dataloader = OmegaConf.create()
+dataloader.train = LazyCall(build_image_train_loader)(
+    dataset=[
+        LazyCall(ImageNetDataset)(
+            root="./dataset",
+            train=True,
+            transform=train_aug,
+        ),
+    ],
+    num_workers=4,
+    mixup_func=None,
+)
+
+
+dataloader.test = [
+    LazyCall(build_image_test_loader)(
+        dataset=LazyCall(ImageNetDataset)(
+            root="./dataset",
+            train=False,
+            transform=test_aug,
+        ),
+        num_workers=4,
+    )
+]

configs/common/data/roberta_dataset.py

+from libai.config import LazyCall
+from omegaconf import OmegaConf
+from libai.data import build_nlp_test_loader, build_nlp_train_val_test_loader
+from libai.data.datasets import RobertaDataset
+from libai.data.data_utils import get_indexed_dataset
+
+from libai.tokenizer import RobertaTokenizer
+
+
+tokenization = OmegaConf.create()
+
+tokenization.tokenizer = LazyCall(RobertaTokenizer)(
+    vocab_file="roberta-vocab.json", merges_file="roberta-merges.txt"
+)
+tokenization.append_eod = False
+tokenization.make_vocab_size_divisible_by = 128
+
+dataloader = OmegaConf.create()
+
+dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
+    dataset=[
+        LazyCall(RobertaDataset)(
+            name="roberta",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_seq_length=514,
+            mask_lm_prob=0.15,
+            short_seq_prob=0.0,
+            seed=1234,
+            masking_style="bert",
+        ),
+    ],
+    train_val_test_num_samples=None,  # a hint for deferred assignment
+    splits=[[949.0, 50.0, 1.0]],
+    weights=[1.0],
+    num_workers=4,
+)
+
+dataloader.test = [
+    LazyCall(build_nlp_test_loader)(
+        dataset=LazyCall(RobertaDataset)(
+            name="roberta",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_num_samples=10,
+            max_seq_length=514,
+            mask_lm_prob=0.15,
+            short_seq_prob=0.1,
+            seed=1234,
+            masking_style="bert",
+        ),
+        test_batch_size=4,
+    )
+]

configs/common/data/t5_dataset.py

+from libai.config import LazyCall
+from omegaconf import OmegaConf
+from libai.data import build_nlp_test_loader, build_nlp_train_val_test_loader
+from libai.data.datasets import T5Dataset
+from libai.data.data_utils import get_indexed_dataset
+
+from libai.tokenizer import BertTokenizer
+
+
+tokenization = OmegaConf.create()
+
+tokenization.setup = True
+
+special_tokens = []
+for i in range(100):
+    special_tokens.append(f"<extra_id_{i}>")
+tokenization.tokenizer = LazyCall(BertTokenizer)(
+    vocab_file="/workspace/data/libai_dataset/bert-base-chinese-vocab.txt",
+    do_lower_case=True,
+    do_chinese_wwm=True,
+    bos_token="[BOS]",
+    eos_token="[EOS]",
+    additional_special_tokens=special_tokens,
+)
+tokenization.append_eod = False
+tokenization.make_vocab_size_divisible_by = 128
+
+dataloader = OmegaConf.create()
+
+dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
+    dataset=[
+        LazyCall(T5Dataset)(
+            name="t5",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_seq_length=512,
+            max_seq_length_dec=128,
+            masked_lm_prob=0.15,
+            short_seq_prob=0.1,
+            seed=1234,
+        ),
+    ],
+    train_val_test_num_samples=None,  # a hint for deferred assignment
+    splits=[[949.0, 50.0, 1.0]],
+    weights=[1.0],
+    num_workers=4,
+)
+
+dataloader.test = [
+    LazyCall(build_nlp_test_loader)(
+        dataset=LazyCall(T5Dataset)(
+            name="t5",
+            data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix="/workspace/data/libai_dataset/loss_compara_content_sentence",
+                data_impl="mmap",
+                skip_warmup=False,
+            ),
+            max_num_samples=10,
+            max_seq_length=512,
+            max_seq_length_dec=128,
+            masked_lm_prob=0.15,
+            short_seq_prob=0.1,
+            seed=1234,
+        ),
+        test_batch_size=4,
+    )
+]

configs/common/models/bert.py

+from omegaconf import DictConfig
+from libai.config import LazyCall
+from libai.models import BertModel, BertForPreTraining
+
+
+cfg = dict(
+    vocab_size=30522,
+    hidden_size=768,
+    hidden_layers=24,
+    num_attention_heads=12,
+    intermediate_size=4096,
+    hidden_dropout_prob=0.1,
+    attention_probs_dropout_prob=0.1,
+    max_position_embeddings=512,
+    num_tokentypes=2,
+    add_pooling_layer=True,
+    initializer_range=0.02,
+    layernorm_eps=1e-5,
+    bias_gelu_fusion=True,
+    bias_dropout_fusion=True,
+    scale_mask_softmax_fusion=True,
+    apply_query_key_layer_scaling=True,
+    apply_residual_post_layernorm=False,
+    add_binary_head=True,
+    amp_enabled=False,
+)
+
+cfg = DictConfig(cfg)
+
+bert_model = LazyCall(BertModel)(cfg=cfg)
+
+pretrain_model = LazyCall(BertForPreTraining)(cfg=cfg)

configs/common/models/gpt.py

+from omegaconf import DictConfig
+from libai.config import LazyCall
+from libai.models import GPTModel, GPTForPreTraining
+
+
+cfg = dict(
+    hidden_layers=6,
+    vocab_size=30522,
+    hidden_size=384,
+    ffn_hidden_size=1536,
+    num_attention_heads=12,
+    max_seq_length=1024,
+    embedding_dropout_prob=0,
+    attention_dropout_prob=0,
+    output_dropout_prob=0,
+    layernorm_epsilon=1e-5,
+    initializer_range=0.02,
+    use_scaled_init_for_output_weights=True,
+    bias_gelu_fusion=True,
+    bias_dropout_fusion=True,
+    scale_mask_softmax_fusion=True,
+    apply_query_key_layer_scaling=True,
+    apply_residual_post_layernorm=False,
+    amp_enabled=False,
+)
+
+cfg = DictConfig(cfg)
+
+gpt_model = LazyCall(GPTModel)(cfg=cfg)
+
+pretrain_model = LazyCall(GPTForPreTraining)(cfg=cfg)

configs/common/models/graph.py

+from omegaconf import DictConfig
+from libai.config import LazyCall
+from libai.models.utils import GraphBase
+
+
+graph = dict(
+    # options for graph or eager mode
+    enabled=True,
+    debug=-1,  # debug mode for graph
+    auto_parallel=dict(
+        enabled=False,
+        enable_auto_parallel_ignore_user_sbp_config=False,  # ignore all .to_global() in graph
+        trunk_algo=True,  # consider overlapping calculate time and transfer time
+        sbp_collector=False,  # use proxy node when one node transfer to many nodes
+    ),
+    train_graph=LazyCall(GraphBase)(
+        is_train=True,
+    ),
+    eval_graph=LazyCall(GraphBase)(is_train=False),
+)
+
+graph = DictConfig(graph)

configs/common/models/resmlp/resmlpB_24.py

+from libai.config import LazyCall
+from libai.models import ResMLP
+
+from .resmlp_12 import cfg
+
+
+cfg.patch_size = 8
+cfg.embed_dim = 768
+cfg.depth = 24
+cfg.init_scale = 1e-6
+
+model = LazyCall(ResMLP)(cfg=cfg)

configs/common/models/resmlp/resmlp_12.py

+from omegaconf import DictConfig
+from libai.config import LazyCall
+from libai.models import ResMLP
+
+
+cfg = dict(
+    img_size=224,
+    patch_size=16,
+    in_chans=3,
+    embed_dim=384,
+    depth=12,
+    drop_rate=0.0,
+    drop_path_rate=0.05,
+    init_scale=0.1,
+    num_classes=1000,
+    loss_func=None,
+)
+
+cfg = DictConfig(cfg)
+
+model = LazyCall(ResMLP)(cfg=cfg)

configs/common/models/resmlp/resmlp_24.py

+from libai.config import LazyCall
+from libai.models import ResMLP
+
+from .resmlp_12 import cfg
+
+
+cfg.depth = 24
+cfg.init_scale = 1e-5
+
+model = LazyCall(ResMLP)(cfg=cfg)

configs/common/models/resmlp/resmlp_36.py

+from libai.config import LazyCall
+from libai.models import ResMLP
+
+from .resmlp_12 import cfg
+
+
+cfg.depth = 36
+cfg.init_scale = 1e-6
+
+model = LazyCall(ResMLP)(cfg=cfg)

configs/common/models/roberta.py

+from omegaconf import DictConfig
+from libai.config import LazyCall
+from libai.models import RobertaModel, RobertaForPreTraining, RobertaForCausalLM
+
+
+cfg = dict(
+    vocab_size=50265,
+    hidden_size=768,
+    hidden_layers=12,
+    num_attention_heads=12,
+    intermediate_size=3072,
+    hidden_dropout_prob=0.1,
+    attention_probs_dropout_prob=0.1,
+    max_position_embeddings=514,
+    num_tokentypes=1,
+    add_pooling_layer=True,
+    initializer_range=0.02,
+    layernorm_eps=1e-5,
+    pad_token_id=1,
+    bias_gelu_fusion=True,
+    bias_dropout_fusion=True,
+    scale_mask_softmax_fusion=True,
+    apply_query_key_layer_scaling=True,
+    apply_residual_post_layernorm=False,
+    amp_enabled=False,
+)
+
+cfg = DictConfig(cfg)
+
+roberta_model = LazyCall(RobertaModel)(cfg=cfg)
+
+roberta_causal_lm = LazyCall(RobertaForCausalLM)(cfg=cfg)
+
+pretrain_model = LazyCall(RobertaForPreTraining)(cfg=cfg)