bloom

projects/MT5/layers/loss_layer.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 libai.layers import ParallelCrossEntropyLoss
+from libai.utils import distributed as dist
+
+
+class MT5Loss(flow.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.lm_loss = ParallelCrossEntropyLoss()
+
+    def forward(self, logits, lm_labels, loss_mask):
+        lm_labels = lm_labels.to_global(placement=logits.placement)
+        lm_loss = self.lm_loss(logits, lm_labels)
+        loss_mask = loss_mask.to_global(placement=lm_loss.placement)
+        loss_mask = loss_mask.float()
+        denominator = loss_mask.sum().to_global(
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])
+        )
+        lm_loss = flow._C.amp_white_identity(lm_loss)
+        lm_loss = flow._C.amp_black_identity(lm_loss)
+        masked_lm_loss = flow.sum(lm_loss.view(-1) * loss_mask.view(-1)) / denominator
+        masked_lm_loss = masked_lm_loss.to_global(
+            sbp=dist.get_nd_sbp([flow.sbp.partial_sum, flow.sbp.broadcast])
+        )
+
+        if self.training:
+            # token throughput
+            done_tokens = (
+                flow.zeros(
+                    1,
+                    sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                    placement=lm_labels.placement,
+                )
+                + logits.shape[0] * logits.shape[1]
+            )
+
+            # correct token
+            correct_tokens = flow.sum(
+                (
+                    logits.to_global(
+                        sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+                        placement=lm_labels.placement,
+                    )
+                    .argmax(dim=-1)
+                    .eq(lm_labels)
+                ).float()
+            )
+
+            return {
+                "mlm_loss": masked_lm_loss,
+                "done_tokens": done_tokens,
+                "correct_tokens": correct_tokens,
+                "denominator": denominator,
+            }
+
+        else:
+            return {
+                "mlm_loss": masked_lm_loss,
+            }

projects/MT5/layers/mask_layer.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 libai.utils import distributed as dist
+
+
+class ExtendedMask(flow.nn.Module):
+    def forward(self, x, input_tensor=None, is_decoder=False):
+        if x.dim() == 3:
+            extended_mask = x[:, None, :, :]
+        elif x.dim() == 2:
+            if is_decoder:
+                extended_mask = self.create_extended_mask_for_decoder(x, input_tensor)
+            else:
+                extended_mask = x[:, None, None, :]
+
+        return extended_mask
+
+    def create_extended_mask_for_decoder(self, x, input_tensor):
+        batch_size, seq_len = input_tensor.size()
+        seq_ids = flow.arange(
+            seq_len,
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            placement=x.placement,
+        )
+        causal_mask = (
+            seq_ids[None, None, :].repeat(batch_size, seq_len, 1) <= seq_ids[None, :, None]
+        )
+
+        causal_mask = causal_mask.to(x.dtype)
+        causal_mask = causal_mask.to_global(sbp=x.sbp)
+        if causal_mask.shape[1] < x.shape[1]:
+            prefix_seq_len = x.shape[1] - causal_mask.shape[1]
+            ones = flow.ones(
+                (batch_size, seq_len, prefix_seq_len),
+                dtype=causal_mask.dtype,
+                sbp=causal_mask.sbp,
+                placement=causal_mask.placement,
+            )
+            causal_mask = flow.cat(
+                [
+                    ones,
+                    causal_mask,
+                ],
+                dim=-1,
+            )
+
+        extended_mask = causal_mask[:, None, :, :] * x[:, None, None, :]
+        return extended_mask

projects/MT5/layers/mlp_layer.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
+
+from libai.layers import Linear, build_activation
+
+
+class T5MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        ffn_hidden_size,
+        output_dropout_prob=0.0,
+        init_method=nn.init.xavier_normal_,
+        output_layer_init_method=None,
+        *,
+        layer_idx=0,
+    ):
+        super().__init__()
+        self.output_dropout_prob = output_dropout_prob
+
+        if output_layer_init_method is None:
+            output_layer_init_method = init_method
+
+        self.dense_h_to_4h = Linear(
+            hidden_size,
+            ffn_hidden_size,
+            bias=False,
+            parallel="col",
+            skip_bias_add=False,
+            init_method=init_method,
+            layer_idx=layer_idx,
+        )
+
+        self.activation_func = build_activation("relu")
+
+        self.dense_4h_to_h = Linear(
+            ffn_hidden_size,
+            hidden_size,
+            bias=False,
+            parallel="row",
+            skip_bias_add=False,
+            init_method=output_layer_init_method,
+            layer_idx=layer_idx,
+        )
+
+        self.dropout = nn.Dropout(self.output_dropout_prob)
+
+    def forward(self, hidden_states):
+        intermediate = self.dense_h_to_4h(hidden_states)
+        intermediate = self.activation_func(intermediate)
+        output = self.dense_4h_to_h(intermediate)
+        output = self.dropout(output)
+        return output
+
+
+class MT5MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        ffn_hidden_size,
+        output_dropout_prob=0.0,
+        init_method=nn.init.xavier_normal_,
+        output_layer_init_method=None,
+        *,
+        layer_idx=0,
+    ):
+        super().__init__()
+        self.output_dropout_prob = output_dropout_prob
+
+        if output_layer_init_method is None:
+            output_layer_init_method = init_method
+
+        self.wi_0 = Linear(
+            hidden_size,
+            ffn_hidden_size,
+            bias=False,
+            parallel="col",
+            skip_bias_add=False,
+            init_method=init_method,
+            layer_idx=layer_idx,
+        )
+
+        self.wi_1 = Linear(
+            hidden_size,
+            ffn_hidden_size,
+            bias=False,
+            parallel="col",
+            skip_bias_add=False,
+            init_method=init_method,
+            layer_idx=layer_idx,
+        )
+
+        self.wo = Linear(
+            ffn_hidden_size,
+            hidden_size,
+            bias=False,
+            parallel="row",
+            skip_bias_add=False,
+            init_method=output_layer_init_method,
+            layer_idx=layer_idx,
+        )
+
+        self.dropout = nn.Dropout(self.output_dropout_prob)
+
+    def forward(self, hidden_states):
+        wi_0_out = self.wi_0(hidden_states)
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = flow._C.fused_fast_gelu_mul(wi_0_out, hidden_linear)
+        output = self.wo(hidden_states)
+        output = self.dropout(output)
+        return output

projects/MT5/layers/transformer_layer.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.nn as nn
+
+from libai.layers.droppath import DropPath
+from libai.layers.layer_norm import RMSLayerNorm as LayerNorm
+from libai.utils import distributed as dist
+from projects.MT5.layers.attention_layer import MultiheadAttention
+from projects.MT5.layers.mlp_layer import MT5MLP, T5MLP
+
+
+class TransformerLayer(nn.Module):
+    """A single transformer layer.
+
+    Transformer layer takes input with size [bsz, seq_length, hidden size] and returns an
+    output of the same size.
+    The input and output has same sbp sign, (S(0), B).
+
+    Arguments:
+        hidden_size: size of hidden state.
+        ffn_hidden_size: size of feed forword neural network.
+        num_attention_heads: number of attention heads.
+        is_decoder: used to specify whether this is transformer encoder layer or transformer
+            decoder layer. Default: ``False``.
+        attention_dropout_prob: dropout probability of attention weights.
+        output_dropout_prob: dropout probability of output.
+        layernorm_epsilon: epsilon used in layernorm layer. Default: `1e-5`.
+        init_method: method to initialize the input layer weights.
+        output_layer_init_method: method to initialize the output layer weights.
+            If None, use `init_method`.
+        layer_idx: the layer index, which determines the placement.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        ffn_hidden_size,
+        num_attention_heads,
+        head_size,
+        relative_attention_num_buckets,
+        is_decoder=False,
+        attention_dropout_prob=0.0,
+        output_dropout_prob=0.0,
+        drop_path_prob=0.0,
+        layernorm_epsilon=1e-5,
+        init_method=nn.init.xavier_normal_,
+        output_layer_init_method=None,
+        padding_idx=None,
+        *,
+        layer_idx=0,
+        model_type="t5",
+        has_relative_attention_bias=False
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.ffn_hidden_size = ffn_hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.head_size = head_size
+        self.attention_dropout_prob = attention_dropout_prob
+        self.output_dropout_prob = output_dropout_prob
+        self.layernorm_epsilon = layernorm_epsilon
+        self.layer_idx = layer_idx
+        self.is_decoder = is_decoder
+
+        self.init_method = init_method
+        if output_layer_init_method is None:
+            output_layer_init_method = init_method
+        self.output_layer_init_method = output_layer_init_method
+
+        self.drop_path = DropPath(drop_path_prob) if drop_path_prob > 0.0 else nn.Identity()
+
+        self.input_layernorm = LayerNorm(
+            self.hidden_size, eps=self.layernorm_epsilon, layer_idx=self.layer_idx
+        )
+
+        self.self_attention = self.build_attention(
+            is_cross_attention=False,
+            relative_attention_num_buckets=relative_attention_num_buckets,
+            padding_idx=padding_idx,
+            has_relative_attention_bias=has_relative_attention_bias,
+            is_decoder=self.is_decoder,
+        )
+        self.post_attention_layernorm = LayerNorm(
+            self.hidden_size, eps=self.layernorm_epsilon, layer_idx=self.layer_idx
+        )
+
+        if self.is_decoder:
+            self.cross_attention = self.build_attention(
+                is_cross_attention=True,
+                relative_attention_num_buckets=relative_attention_num_buckets,
+                padding_idx=padding_idx,
+                is_decoder=self.is_decoder,
+            )
+            self.post_cross_attention_layernorm = LayerNorm(
+                self.hidden_size, eps=self.layernorm_epsilon, layer_idx=self.layer_idx
+            )
+        if model_type == "mt5":
+            self.mlp = MT5MLP(
+                self.hidden_size,
+                self.ffn_hidden_size,
+                self.output_dropout_prob,
+                self.init_method,
+                output_layer_init_method=self.output_layer_init_method,
+                layer_idx=self.layer_idx,
+            )
+        elif model_type == "t5":
+            self.mlp = T5MLP(
+                self.hidden_size,
+                self.ffn_hidden_size,
+                self.output_dropout_prob,
+                self.init_method,
+                output_layer_init_method=self.output_layer_init_method,
+                layer_idx=self.layer_idx,
+            )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        encoder_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        use_cache=False,
+        position_bias=None,
+        encoder_decoder_position_bias=None,
+    ):
+        """
+        Args:
+            hidden_states: shape is (batch_size, seq_length, hidden_size),
+                sbp signature is (S(0), B).
+            attention_mask: the combination of key padding mask and casual mask of hidden states
+                with shape (batch_size, 1, seq_length, seq_length) and the sbp
+                signature is (S(0), B),
+            encoder_states: encoder output with shape (batch_size, seq_length, hidden_size)
+                and the sbp signature is (S(0), B), which will be used in cross attention.
+            encoder_attention_mask: key padding mask of encoder states with shape
+                (batch_size, 1, seq_length, seq_length) and the sbp signature is (S(0), B).
+            past_key_value: tuple of key and value, each shape is
+                (seq_length, bsz, num_heads, head_size), For decoder layer,
+                the past_key_value contains the states both from self attention
+                and cross attention.
+            use_cache: it will be set to `True` when the model is in the inference phase and
+                used for incremental decoding.
+        """
+        hidden_states = hidden_states.to_global(placement=dist.get_layer_placement(self.layer_idx))
+
+        if attention_mask is not None:
+            attention_mask = attention_mask.to_global(
+                placement=dist.get_layer_placement(self.layer_idx)
+            )
+
+        if past_key_value is not None:
+            if self.is_decoder:
+                assert len(past_key_value) == 4
+                self_attn_past_key_value = past_key_value[:2]
+                cross_attn_past_key_value = past_key_value[2:]
+            else:
+                self_attn_past_key_value = past_key_value
+                cross_attn_past_key_value = None
+        else:
+            self_attn_past_key_value, cross_attn_past_key_value = None, None
+
+        layernorm_output = self.input_layernorm(hidden_states)
+
+        attention_output, position_bias = self.self_attention(
+            layernorm_output,
+            attention_mask=attention_mask,
+            past_key_value=self_attn_past_key_value,
+            position_bias=position_bias,
+            use_cache=use_cache,
+        )
+
+        attention_output = self.drop_path(attention_output)
+
+        if use_cache:
+            attention_output, presents = attention_output
+        else:
+            presents = None
+
+        hidden_states = hidden_states + attention_output
+
+        layernorm_output = self.post_attention_layernorm(hidden_states)
+
+        if self.is_decoder:
+            if presents is not None:
+                query_length = presents[0].shape[2]
+            else:
+                query_length = None
+
+            attention_output, encoder_decoder_position_bias = self.cross_attention(
+                layernorm_output,
+                encoder_states,
+                attention_mask=encoder_attention_mask,
+                past_key_value=cross_attn_past_key_value,
+                position_bias=encoder_decoder_position_bias,
+                use_cache=use_cache,
+                query_length=query_length,
+            )
+            if use_cache:
+                attention_output, decoder_presents = attention_output
+                presents = presents + decoder_presents
+
+            attention_output = self.drop_path(attention_output)
+
+            hidden_states = hidden_states + attention_output
+            layernorm_output = self.post_cross_attention_layernorm(hidden_states)
+
+        mlp_output = self.mlp(layernorm_output)
+        mlp_output = self.drop_path(mlp_output)
+
+        output = hidden_states + mlp_output
+
+        if use_cache:
+            output = (output, presents)
+        output = (output,) + (position_bias,)
+        if self.is_decoder:
+            output = output + (encoder_decoder_position_bias,)
+        return output
+
+    def build_attention(
+        self,
+        is_cross_attention=False,
+        relative_attention_num_buckets=None,
+        padding_idx=None,
+        has_relative_attention_bias=False,
+        is_decoder=False,
+    ):
+        return MultiheadAttention(
+            self.hidden_size,
+            self.num_attention_heads,
+            head_size=self.head_size,
+            relative_attention_num_buckets=relative_attention_num_buckets,
+            is_cross_attention=is_cross_attention,
+            attention_dropout_prob=self.attention_dropout_prob,
+            output_dropout_prob=self.output_dropout_prob,
+            init_method=self.init_method,
+            output_layer_init_method=self.output_layer_init_method,
+            padding_idx=padding_idx,
+            layer_idx=self.layer_idx,
+            has_relative_attention_bias=has_relative_attention_bias,
+            is_decoder=is_decoder,
+        )

projects/MT5/mt5_model.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
+
+from libai.config import configurable
+from libai.inference.generator.generation_utils import Generator
+from libai.layers import Linear, LMLogits, RMSLayerNorm
+from libai.models.utils import init_method_normal, scaled_init_method_normal
+from libai.utils import distributed as dist
+from projects.MT5.layers.embed_layer import MT5Embedding
+from projects.MT5.layers.loss_layer import MT5Loss
+from projects.MT5.layers.mask_layer import ExtendedMask
+from projects.MT5.layers.transformer_layer import TransformerLayer
+from projects.MT5.utils.mt5_loader import T5LoaderHuggerFace
+
+
+class MT5Model(flow.nn.Module, Generator):
+    @configurable
+    def __init__(
+        self,
+        vocab_size,
+        hidden_size,
+        hidden_layers,
+        num_attention_heads,
+        head_size,
+        intermediate_size,
+        embedding_dropout_prob,
+        hidden_dropout_prob,
+        attention_probs_dropout_prob,
+        relative_attention_num_buckets,
+        padding_idx=None,
+        initializer_range=0.02,
+        layernorm_eps=1e-12,
+        amp_enabled=False,
+        model_type="mt5",
+        cfg=None,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.model_type = model_type
+        init_method = init_method_normal(initializer_range)
+        scaled_init_method = scaled_init_method_normal(initializer_range, hidden_layers)
+        self.embedding = MT5Embedding(
+            hidden_size=hidden_size,
+            vocab_size=vocab_size,
+            embedding_dropout_prob=embedding_dropout_prob,
+            init_method=init_method,
+            amp_enabled=amp_enabled,
+        )
+        self.extended_attn_mask = ExtendedMask()
+
+        encoder_layers = flow.nn.ModuleList(
+            [
+                TransformerLayer(
+                    hidden_size=hidden_size,
+                    ffn_hidden_size=intermediate_size,
+                    num_attention_heads=num_attention_heads,
+                    head_size=head_size,
+                    relative_attention_num_buckets=relative_attention_num_buckets,
+                    is_decoder=False,
+                    attention_dropout_prob=attention_probs_dropout_prob,
+                    output_dropout_prob=hidden_dropout_prob,
+                    layernorm_epsilon=layernorm_eps,
+                    init_method=init_method,
+                    output_layer_init_method=scaled_init_method,
+                    padding_idx=padding_idx,
+                    layer_idx=i,
+                    model_type=model_type,
+                    has_relative_attention_bias=bool(i == 0),
+                )
+                for i in range(hidden_layers)
+            ]
+        )
+
+        encoder_final_layernorm = RMSLayerNorm(
+            (hidden_size,),
+            eps=layernorm_eps,
+            layer_idx=hidden_layers - 1,
+        )
+
+        self.encoder = flow.nn.Sequential()
+        self.encoder.add_module("layers", encoder_layers)
+        self.encoder.add_module("final_layernorm", encoder_final_layernorm)
+
+        decoder_layers = flow.nn.ModuleList(
+            [
+                TransformerLayer(
+                    hidden_size=hidden_size,
+                    ffn_hidden_size=intermediate_size,
+                    num_attention_heads=num_attention_heads,
+                    head_size=head_size,
+                    relative_attention_num_buckets=relative_attention_num_buckets,
+                    is_decoder=True,
+                    attention_dropout_prob=attention_probs_dropout_prob,
+                    output_dropout_prob=hidden_dropout_prob,
+                    layernorm_epsilon=layernorm_eps,
+                    init_method=init_method,
+                    output_layer_init_method=scaled_init_method,
+                    padding_idx=padding_idx,
+                    layer_idx=i,
+                    model_type=model_type,
+                    has_relative_attention_bias=bool(i - hidden_layers == 0),
+                )
+                for i in range(hidden_layers, 2 * hidden_layers)
+            ]
+        )
+
+        decoder_final_layernorm = RMSLayerNorm(
+            (hidden_size,),
+            eps=layernorm_eps,
+            layer_idx=2 * hidden_layers - 1,
+        )
+
+        self.decoder = flow.nn.Sequential()
+        self.decoder.add_module("layers", decoder_layers)
+        self.decoder.add_module("final_layernorm", decoder_final_layernorm)
+        self.past_key_values = [None] * len(self.decoder.layers)
+        self.encoder_states = None
+        self.past_length = 0
+
+        if model_type == "mt5":
+            self.lm_head = Linear(
+                hidden_size, vocab_size, bias=False, layer_idx=2 * hidden_layers - 1
+            )
+        else:
+            self.lm_head = LMLogits(vocab_size, bias=False)
+
+    @classmethod
+    def from_config(cls, cfg):
+        return {
+            "vocab_size": cfg.vocab_size,
+            "hidden_size": cfg.hidden_size,
+            "hidden_layers": cfg.hidden_layers,
+            "num_attention_heads": cfg.num_attention_heads,
+            "head_size": cfg.head_size,
+            "intermediate_size": cfg.intermediate_size,
+            "embedding_dropout_prob": cfg.embedding_dropout_prob,
+            "hidden_dropout_prob": cfg.hidden_dropout_prob,
+            "attention_probs_dropout_prob": cfg.attention_probs_dropout_prob,
+            "relative_attention_num_buckets": cfg.relative_attention_num_buckets,
+            "padding_idx": cfg.padding_idx,
+            "initializer_range": cfg.initializer_range,
+            "layernorm_eps": cfg.layernorm_eps,
+            "amp_enabled": cfg.amp_enabled,
+            "model_type": cfg.model_type,
+            "cfg": cfg,
+        }
+
+    def forward(
+        self,
+        encoder_input_ids=None,
+        decoder_input_ids=None,
+        encoder_attn_mask=None,
+        decoder_attn_mask=None,
+        encoder_decoder_attn_mask=None,
+        use_cache=False,
+        only_encoder=False,
+    ):
+
+        encoder_input_ids = (
+            encoder_input_ids.to_global(placement=dist.get_layer_placement(0))
+            if encoder_input_ids is not None
+            else encoder_input_ids
+        )
+        decoder_input_ids = (
+            decoder_input_ids.to_global(placement=dist.get_layer_placement(0))
+            if decoder_input_ids is not None
+            else decoder_input_ids
+        )
+        encoder_attn_mask = (
+            encoder_attn_mask.to_global(placement=dist.get_layer_placement(0))
+            if encoder_attn_mask is not None
+            else encoder_attn_mask
+        )
+        decoder_attn_mask = (
+            decoder_attn_mask.to_global(placement=dist.get_layer_placement(0))
+            if decoder_attn_mask is not None
+            else decoder_attn_mask
+        )
+        encoder_decoder_attn_mask = (
+            encoder_decoder_attn_mask.to_global(placement=dist.get_layer_placement(0))
+            if encoder_decoder_attn_mask is not None
+            else encoder_decoder_attn_mask
+        )
+
+        if use_cache and self.encoder_states is not None:
+            encoder_states = self.encoder_states
+        else:
+            position_bias = None
+            encoder_decoder_position_bias = None
+            self.set_cache(encoder_states=None, past_key_values=None)
+            encoder_attn_mask = self.extended_attn_mask(encoder_attn_mask)
+            enc_embedding_output = self.embedding(encoder_input_ids)
+            # transpose [batch_size, seq_len, embed_size] to [seq_len, batch_size, embed_size]
+            enc_hidden_states = enc_embedding_output.transpose(0, 1)
+
+            for layer in self.encoder.layers:
+                enc_hidden_states, position_bias = layer(
+                    enc_hidden_states,
+                    encoder_attn_mask,
+                    position_bias=position_bias,
+                )
+            encoder_states = self.encoder.final_layernorm(enc_hidden_states)
+
+        if only_encoder:
+            return encoder_states
+
+        decoder_attn_mask = self.extended_attn_mask(
+            decoder_attn_mask, decoder_input_ids, is_decoder=True
+        )
+        encoder_decoder_attn_mask = self.extended_attn_mask(encoder_decoder_attn_mask)
+
+        dec_embedding_output = self.embedding(decoder_input_ids)
+        # transpose [batch_size, seq_len, embed_size] to [seq_len, batch_size, embed_size]
+        dec_hidden_states = dec_embedding_output.transpose(0, 1)
+        if use_cache:
+            presents = []
+
+        position_bias = None
+        encoder_decoder_position_bias = None
+        for layer, past_key_value in zip(self.decoder.layers, self.past_key_values):
+            dec_hidden_states, position_bias, encoder_decoder_position_bias = layer(
+                dec_hidden_states,
+                decoder_attn_mask,
+                encoder_states,
+                encoder_decoder_attn_mask,
+                past_key_value=past_key_value,
+                position_bias=position_bias,
+                encoder_decoder_position_bias=encoder_decoder_position_bias,
+                use_cache=use_cache,
+            )
+            if use_cache:
+                dec_hidden_states, present = dec_hidden_states
+                presents.append(present)
+        if use_cache:
+            self.set_cache(encoder_states, past_key_values=presents)
+
+        decoder_states = self.decoder.final_layernorm(dec_hidden_states)
+
+        if self.cfg.tie_word_embeddings:
+            decoder_states = decoder_states * (self.cfg.hidden_size ** -0.5)
+
+        if self.model_type == "mt5":
+            logits = self.lm_head(decoder_states)
+        else:
+            logits = self.lm_head(decoder_states, self.embedding.word_embeddings.weight)
+
+        return {"logits": logits}
+
+    def set_cache(self, encoder_states, past_key_values):
+        self.encoder_states = encoder_states
+        self.past_length = 0 if past_key_values is None else past_key_values[0][0].shape[2]
+
+        if past_key_values is None:
+            past_key_values = [None] * len(self.decoder.layers)
+        assert len(past_key_values) == len(self.decoder.layers), (
+            f"past_key_values's length {len(past_key_values)} doesn't match "
+            f"decoder num_layers' length {self.decoder.layers}"
+        )
+        self.past_key_values = past_key_values
+
+    def _reorder_cache(self, beam_idx):
+        past_key_values = self.past_key_values
+        reordered_decoder_past = ()
+        for layer_past_states in past_key_values:
+            # get the correct batch idx from layer past batch dim
+            # batch dim of `past` is at 2nd position
+            reordered_layer_past_states = ()
+            for layer_past_state in layer_past_states:
+                # need to set correct `past` for each of the four key / value states
+                beam_idx = beam_idx.to_global(placement=layer_past_state.placement)
+                reordered_layer_past_states = reordered_layer_past_states + (
+                    layer_past_state.index_select(0, beam_idx),
+                )
+
+            assert reordered_layer_past_states[0].shape == layer_past_states[0].shape
+            assert len(reordered_layer_past_states) == len(layer_past_states)
+
+            reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,)
+        return reordered_decoder_past
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past=None,
+        encoder_attn_mask=None,
+        encoder_decoder_attn_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+    ):
+        # cut decoder_input_ids if past is used
+        if past is not None:
+            input_ids = input_ids[:, -1:]
+            self.past_key_values = past
+
+        self.encoder_states = encoder_outputs
+        decoder_attn_maks = flow.ones(
+            input_ids.size(),
+            dtype=flow.bool,
+            sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast]),
+            placement=flow.placement("cuda", list(range(dist.get_world_size()))),
+        )
+        return {
+            "decoder_input_ids": input_ids,
+            "decoder_attn_mask": decoder_attn_maks,
+            "encoder_attn_mask": encoder_attn_mask,
+            "encoder_decoder_attn_mask": encoder_decoder_attn_mask,
+            "use_cache": use_cache,
+        }
+
+
+class MT5ForPreTraining(flow.nn.Module):
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        if cfg.pretrained_model_path is not None:
+            loader = T5LoaderHuggerFace(MT5Model, cfg, cfg.pretrained_model_path)
+            self.mt5_model = loader.load()
+        else:
+            self.mt5_model = MT5Model(cfg)
+        self.loss_func = MT5Loss()
+
+    def set_cache(self, encoder_states, past_key_values):
+        self.mt5_model.set_cache(encoder_states, past_key_values)
+
+    def forward(
+        self,
+        encoder_input_ids,
+        decoder_input_ids,
+        encoder_attn_mask,
+        decoder_attn_mask,
+        encoder_decoder_attn_mask,
+        lm_labels=None,
+        loss_mask=None,
+        use_cache=False,
+    ):
+        logits = self.mt5_model(
+            encoder_input_ids,
+            decoder_input_ids,
+            encoder_attn_mask,
+            decoder_attn_mask,
+            encoder_decoder_attn_mask,
+            use_cache=use_cache,
+        )["logits"]
+        # transpose [seq_len, batch_size, vocab_size] to [batch_size, seq_len, vocab_size]
+        logits = logits.transpose(0, 1)
+        if lm_labels is not None:
+            lm_loss = self.loss_func(logits, lm_labels, loss_mask)
+            return lm_loss
+        else:
+            return {
+                "prediction_scores": logits,
+            }
+
+    @staticmethod
+    def set_pipeline_stage_id(model):
+        dist_utils = dist.get_dist_util()
+
+        # Set pipeline parallelism stage_id
+        if hasattr(model.mt5_model.encoder.final_layernorm, "config"):
+            # Old API in OneFlow 0.8
+            for module_block in model.modules():
+                if isinstance(module_block.origin, MT5Embedding):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.origin, ExtendedMask):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.origin, TransformerLayer):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(module_block.layer_idx),
+                        dist.get_layer_placement(module_block.layer_idx),
+                    )
+                elif isinstance(module_block.origin, MT5Loss):
+                    module_block.config.set_stage(
+                        dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+                    )
+
+            model.mt5_model.encoder.final_layernorm.config.set_stage(
+                dist_utils.get_layer_stage_id(model.mt5_model.encoder.final_layernorm.layer_idx),
+                dist.get_layer_placement(model.mt5_model.encoder.final_layernorm.layer_idx),
+            )
+            model.mt5_model.decoder.final_layernorm.config.set_stage(
+                dist_utils.get_layer_stage_id(model.mt5_model.decoder.final_layernorm.layer_idx),
+                dist.get_layer_placement(model.mt5_model.decoder.final_layernorm.layer_idx),
+            )
+        else:
+            for module_block in model.modules():
+                if isinstance(module_block.to(nn.Module), MT5Embedding):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.to(nn.Module), ExtendedMask):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(0), dist.get_layer_placement(0)
+                    )
+                elif isinstance(module_block.to(nn.Module), TransformerLayer):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(module_block.layer_idx),
+                        dist.get_layer_placement(module_block.layer_idx),
+                    )
+                elif isinstance(module_block.to(nn.Module), MT5Loss):
+                    module_block.to(flow.nn.graph.GraphModule).set_stage(
+                        dist_utils.get_layer_stage_id(-1), dist.get_layer_placement(-1)
+                    )
+
+            model.mt5_model.encoder.final_layernorm.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(model.mt5_model.encoder.final_layernorm.layer_idx),
+                dist.get_layer_placement(model.mt5_model.encoder.final_layernorm.layer_idx),
+            )
+            model.mt5_model.decoder.final_layernorm.to(flow.nn.graph.GraphModule).set_stage(
+                dist_utils.get_layer_stage_id(model.mt5_model.decoder.final_layernorm.layer_idx),
+                dist.get_layer_placement(model.mt5_model.decoder.final_layernorm.layer_idx),
+            )
+
+    @staticmethod
+    def set_activation_checkpoint(model):
+        for module_block in model.modules():
+            # Old API in OneFlow 0.8
+            if hasattr(module_block, "origin"):
+                if isinstance(module_block.origin, TransformerLayer):
+                    module_block.config.activation_checkpointing = True
+            else:
+                if isinstance(module_block.to(nn.Module), TransformerLayer):
+                    module_block.to(nn.graph.GraphModule).activation_checkpointing = True

projects/MT5/readme.md

+# MT5
+
+Reproduce T5Model and MT5Model with OneFlow, which effect are equivalent to HuggingFace's [T5](https://huggingface.co/docs/transformers/v4.19.4/en/model_doc/t5#overview) and [T5v1.1](https://github.com/google-research/text-to-text-transfer-transformer/blob/main/released_checkpoints.md#t511).
+
+## Introduce
+The t5 and mt5 pretraining project can support 3D parallel and [ZERO](https://arxiv.org/abs/2202.10435).
+
+## Training MT5
+Training MT5 on 8 GPUs using 3D parallelism and ZERO.
+
+### 1. Prepare your training config file
+
+> set the pretrain parameters in `MT5/configs/mt5_pretrain.py`, such as `vocab_file` and `data_prefix`.
+
+> If you would like to use the t5 model, please set `model_type`="t5".
+
+### 2. Prepare the demo training data
+
+Prepare the demo training data by running:
+```bash
+# path/to/libai
+wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/ci-files/dataset/libai/bert_dataset/bert-base-chinese-vocab.txt -P ./data_test/bert_data/
+wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/ci-files/dataset/libai/bert_dataset/loss_compara_content_sentence.bin -P ./data_test/bert_data/
+wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/ci-files/dataset/libai/bert_dataset/loss_compara_content_sentence.idx -P ./data_test/bert_data/
+```
+
+### 3. Prepare your own training data
+
+If you want to use your own training data, please skip the step2, and refer [Preprocessing Dataset](https://libai.readthedocs.io/en/latest/tutorials/basics/Preprocessing_Dataset.html#).
+
+```bash
+IMPL=mmap
+KEYS=text
+
+python tools/preprocess_data.py \
+        --input /path/to/libai/projects/MT5/data/test.json \
+        --json-keys ${KEYS} \
+        --vocab-file /path/to/libai/projects/MT5/data/vocab.txt \
+        --dataset-impl ${IMPL} \
+        --tokenizer-name BertTokenizer \
+        --do-lower-case \
+        --do-chinese-wwm \
+        --split-sentences \
+        --output-prefix magic_prompt_${IMPL} \
+        --workers 4 \
+        --log-interval 2
+
+```
+
+### 4. Run the following code to start training
+```bash
+# cd /path/to/libai
+bash tools/train.sh projects/MT5/train_net.py projects/MT5/configs/mt5_pretrain.py 8
+```

projects/MT5/train_net.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 logging
+import os
+import random
+import sys
+
+import numpy as np
+import oneflow as flow
+
+from libai.config import LazyConfig, default_argument_parser, try_get_key
+from libai.engine import DefaultTrainer, default_setup
+from libai.utils.checkpoint import Checkpointer
+from libai.utils.events import JSONWriter, TensorboardXWriter
+from projects.MT5.utils.mt5_metrc_printer import MT5MetricPrinter
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
+
+
+logger = logging.getLogger("libai." + __name__)
+
+
+class Mt5Trainer(DefaultTrainer):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def build_writers(self):
+        """
+        Build a list of writers to be used. By default it contains
+        writers that write metrics to the screen,
+        a json file, and a tensorboard event file respectively.
+        If you'd like a different list of writers, you can overwrite it in
+        your trainer.
+
+        Returns:
+            list[EventWriter]: a list of :class:`EventWriter` objects.
+
+        It is now implemented by:
+
+        .. code-block:: python
+
+            return [
+                MT5MetricPrinter(self.global_batch_size, self.max_iter),
+                JSONWriter(os.path.join(self.cfg.train.output_dir, "metrics.json")),
+                TensorboardXWriter(self.cfg.train.output_dir),
+            ]
+        """
+        # Assume the default print/log frequency.
+        return [
+            # It may not always print what you want to see, since it prints "common" metrics only.
+            MT5MetricPrinter(self.global_batch_size, self.max_iter, self.cfg.train.log_period),
+            JSONWriter(os.path.join(self.cfg.train.output_dir, "metrics.json")),
+            TensorboardXWriter(self.cfg.train.output_dir),
+        ]
+
+
+def main(args):
+    cfg = LazyConfig.load(args.config_file)
+    cfg = LazyConfig.apply_overrides(cfg, args.opts)
+    default_setup(cfg, args)
+
+    seed_for_rank = cfg.train.seed + flow.env.get_rank()
+    flow.manual_seed(seed_for_rank)
+    flow.cuda.manual_seed(seed_for_rank)
+    np.random.seed(seed_for_rank)
+    random.seed(seed_for_rank)
+
+    if args.fast_dev_run:
+        cfg.train.train_epoch = 0
+        cfg.train.train_iter = 20
+        cfg.train.evaluation.eval_period = 10
+        cfg.train.log_period = 1
+
+    if args.eval_only:
+        tokenizer = None
+        if try_get_key(cfg, "tokenization") is not None:
+            tokenizer = Mt5Trainer.build_tokenizer(cfg)
+        model = Mt5Trainer.build_model(cfg)
+        Checkpointer(model, save_dir=cfg.train.output_dir).resume_or_load(
+            cfg.train.load_weight, resume=args.resume
+        )
+        if try_get_key(cfg, "train.graph.enabled", default=False):
+            model = Mt5Trainer.build_graph(cfg, model, is_train=False)
+        test_loader = Mt5Trainer.build_test_loader(cfg, tokenizer)
+        if len(test_loader) == 0:
+            logger.info("No dataset in dataloader.test, please set dataset for dataloader.test")
+        _ = Mt5Trainer.test(cfg, test_loader, model)
+        return
+
+    trainer = Mt5Trainer(cfg)
+    return trainer.train()
+
+
+if __name__ == "__main__":
+    args = default_argument_parser().parse_args()
+    main(args)

projects/MT5/utils/mt5_loader.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 json
+
+import oneflow as flow
+
+from libai.models.utils import ModelLoaderHuggerFace, ModelLoaderLiBai
+
+
+class T5LoaderHuggerFace(ModelLoaderHuggerFace):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+
+        """NOTE: base_model_prefix_1 is T5's prefix in Transformers.
+        base_model_prefix_2 is T5's prefix in LiBai."""
+        self.base_model_prefix_1 = "transformer"
+        self.base_model_prefix_2 = "mt5_model"
+
+    def _convert_state_dict(self, flow_state_dict, cfg):
+        """Convert state_dict's keys to match model.
+
+        Args:
+            flow_state_dict (OrderedDict): model state dict.
+            cfg (dict): model's default config dict in LiBai.
+
+        Returns:
+            OrderedDict: flow state dict.
+        """
+        # The converted checkpoint.
+        oneflow_state_dict = flow_state_dict.copy()
+        old_keys = list(oneflow_state_dict.keys())
+        # Get configs
+        num_heads = cfg.get("num_attention_heads")
+        hidden_size = cfg.get("hidden_size")
+        head_size = cfg.get("head_size", None)
+        if head_size is None:
+            head_size = int(hidden_size / num_heads)
+
+        has_prefix = any(s.startswith(self.base_model_prefix_1) for s in oneflow_state_dict)
+        prefix1 = self.base_model_prefix_1 + "." if has_prefix else ""
+        prefix2 = self.base_model_prefix_2 + "." if has_prefix else ""
+        encoder_decoder_idx = 1 if has_prefix else 0
+        layer_idx1 = 3 if has_prefix else 2
+        layer_idx2 = 5 if has_prefix else 4
+        op_idx = 6 if has_prefix else 5
+
+        # Convert T5's Embedding layers.
+        # NOTE: Transformers' T5 has no position embedding layer.
+        new_key = prefix2 + "embedding.word_embeddings.weight"
+        old_keys.remove(prefix1 + "shared.weight")
+        oneflow_state_dict[new_key] = oneflow_state_dict.pop(prefix1 + "shared.weight")
+
+        # Convert T5's final_layer_norm
+        new_key = prefix2 + "encoder.final_layernorm.weight"
+        old_keys.remove(prefix1 + "encoder.final_layer_norm.weight")
+        oneflow_state_dict[new_key] = oneflow_state_dict.pop(
+            prefix1 + "encoder.final_layer_norm.weight"
+        )
+        new_key = prefix2 + "decoder.final_layernorm.weight"
+        old_keys.remove(prefix1 + "decoder.final_layer_norm.weight")
+        oneflow_state_dict[new_key] = oneflow_state_dict.pop(
+            prefix1 + "decoder.final_layer_norm.weight"
+        )
+
+        # Convert MT5's lm_head
+        if cfg.model_type == "mt5" and "lm_head.weight" in oneflow_state_dict:
+            new_key = prefix2 + "lm_head.weight"
+            old_keys.remove("lm_head.weight")
+            oneflow_state_dict[new_key] = oneflow_state_dict.pop("lm_head.weight")
+
+        # NOTE: Each layers has no bias in Transformer's T5.
+        for key in old_keys:
+            keys = key.split(".")
+            if layer_idx1 > len(keys) or layer_idx2 > len(keys):
+                continue
+            layer1 = keys[layer_idx1]
+            layer2 = keys[layer_idx2]
+            op_name = keys[op_idx]
+
+            if keys[op_idx + 1] == "relative_attention_bias" and keys[op_idx] == "SelfAttention":
+                new_key = (
+                    prefix2
+                    + keys[encoder_decoder_idx]
+                    + ".layers.0.self_attention.relative_attention_bias.weight"
+                )
+                oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert T5's Encoder layers.
+            if keys[encoder_decoder_idx] == "encoder":
+                if op_name == "SelfAttention":
+                    new_key = (
+                        prefix2
+                        + "encoder.layers."
+                        + layer1
+                        + ".self_attention.query_key_value.weight"
+                    )
+                    if new_key in oneflow_state_dict.keys():
+                        continue
+                    q_w = ".".join(keys[: op_idx + 1]) + ".q." + "weight"
+                    k_w = ".".join(keys[: op_idx + 1]) + ".k." + "weight"
+                    v_w = ".".join(keys[: op_idx + 1]) + ".v." + "weight"
+                    qkv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_w),
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    qkv_w = self._fix_qkv_ordering(qkv_w, head_size, num_heads, hidden_size)
+                    oneflow_state_dict[new_key] = qkv_w
+
+                    o_w = ".".join(keys[: op_idx + 1]) + ".o." + "weight"
+                    new_key = prefix2 + "encoder.layers." + layer1 + ".self_attention.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(o_w)
+                elif op_name == "layer_norm":
+                    if layer2 == "0":
+                        new_key = prefix2 + "encoder.layers." + layer1 + ".input_layernorm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif layer2 == "1":
+                        new_key = (
+                            prefix2
+                            + "encoder.layers."
+                            + layer1
+                            + ".post_attention_layernorm.weight"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif op_name == "DenseReluDense":
+                    if cfg.get("model_type") == "t5":
+                        if keys[op_idx + 1] == "wi":
+                            new_key = (
+                                prefix2 + "encoder.layers." + layer1 + ".mlp.dense_h_to_4h.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wo":
+                            new_key = (
+                                prefix2 + "encoder.layers." + layer1 + ".mlp.dense_4h_to_h.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif cfg.get("model_type") == "mt5":
+                        if keys[op_idx + 1] == "wi_0":
+                            new_key = prefix2 + "encoder.layers." + layer1 + ".mlp.wi_0.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wi_1":
+                            new_key = prefix2 + "encoder.layers." + layer1 + ".mlp.wi_1.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wo":
+                            new_key = prefix2 + "encoder.layers." + layer1 + ".mlp.wo.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+
+            # Convert T5's decoder Layers.
+            elif keys[encoder_decoder_idx] == "decoder":
+                if op_name == "SelfAttention":
+                    new_key = (
+                        prefix2
+                        + "decoder.layers."
+                        + layer1
+                        + ".self_attention.query_key_value.weight"
+                    )
+                    if new_key in oneflow_state_dict.keys():
+                        continue
+                    q_w = ".".join(keys[: op_idx + 1]) + ".q." + "weight"
+                    k_w = ".".join(keys[: op_idx + 1]) + ".k." + "weight"
+                    v_w = ".".join(keys[: op_idx + 1]) + ".v." + "weight"
+                    qkv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(q_w),
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    qkv_w = self._fix_qkv_ordering(qkv_w, head_size, num_heads, hidden_size)
+
+                    oneflow_state_dict[new_key] = qkv_w
+
+                    o_w = ".".join(keys[: op_idx + 1]) + ".o." + "weight"
+                    new_key = prefix2 + "decoder.layers." + layer1 + ".self_attention.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(o_w)
+                elif op_name == "layer_norm":
+                    if layer2 == "0":
+                        new_key = prefix2 + "decoder.layers." + layer1 + ".input_layernorm.weight"
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif layer2 == "1":
+                        new_key = (
+                            prefix2
+                            + "decoder.layers."
+                            + layer1
+                            + ".post_attention_layernorm.weight"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif layer2 == "2":
+                        new_key = (
+                            prefix2
+                            + "decoder.layers."
+                            + layer1
+                            + ".post_cross_attention_layernorm.weight"
+                        )
+                        oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                elif op_name == "EncDecAttention":
+                    new_key = prefix2 + "decoder.layers." + layer1 + ".cross_attention.query.weight"
+                    if new_key in oneflow_state_dict.keys():
+                        continue
+                    q_w = ".".join(keys[: op_idx + 1]) + ".q." + "weight"
+                    k_w = ".".join(keys[: op_idx + 1]) + ".k." + "weight"
+                    v_w = ".".join(keys[: op_idx + 1]) + ".v." + "weight"
+
+                    q_w = oneflow_state_dict.pop(q_w)
+                    kv_w = flow.cat(
+                        (
+                            oneflow_state_dict.pop(k_w),
+                            oneflow_state_dict.pop(v_w),
+                        ),
+                        dim=0,
+                    )
+                    q_w = self._fix_qkv_ordering(q_w, head_size, num_heads, hidden_size)
+                    kv_w = self._fix_qkv_ordering(kv_w, head_size, num_heads, hidden_size)
+
+                    oneflow_state_dict[new_key] = q_w
+                    new_key = (
+                        prefix2 + "decoder.layers." + layer1 + ".cross_attention.key_value.weight"
+                    )
+                    oneflow_state_dict[new_key] = kv_w
+
+                    o_w = ".".join(keys[: op_idx + 1]) + ".o." + "weight"
+                    new_key = prefix2 + "decoder.layers." + layer1 + ".cross_attention.dense.weight"
+                    oneflow_state_dict[new_key] = oneflow_state_dict.pop(o_w)
+                elif op_name == "DenseReluDense":
+                    if cfg.get("model_type") == "t5":
+                        if keys[op_idx + 1] == "wi":
+                            new_key = (
+                                prefix2 + "decoder.layers." + layer1 + ".mlp.dense_h_to_4h.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wo":
+                            new_key = (
+                                prefix2 + "decoder.layers." + layer1 + ".mlp.dense_4h_to_h.weight"
+                            )
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                    elif cfg.get("model_type") == "mt5":
+                        if keys[op_idx + 1] == "wi_0":
+                            new_key = prefix2 + "decoder.layers." + layer1 + ".mlp.wi_0.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wi_1":
+                            new_key = prefix2 + "decoder.layers." + layer1 + ".mlp.wi_1.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+                        elif keys[op_idx + 1] == "wo":
+                            new_key = prefix2 + "decoder.layers." + layer1 + ".mlp.wo.weight"
+                            oneflow_state_dict[new_key] = oneflow_state_dict.pop(key)
+        return oneflow_state_dict
+
+    def _load_config_from_json(self, config_file):
+        """load config from `config.json`, and update default config.
+
+        Args:
+            config_file (str): Path of config file.
+        """
+        with open(config_file, mode="r", encoding="utf-8") as f:
+            cfg_dict = json.load(f)
+
+        self._update_cfg("vocab_size", cfg_dict["vocab_size"])
+        self._update_cfg("hidden_size", cfg_dict["d_model"])
+        self._update_cfg("hidden_layers", cfg_dict["num_layers"])
+        self._update_cfg("num_attention_heads", cfg_dict["num_heads"])
+        self._update_cfg("intermediate_size", cfg_dict["d_ff"])
+        self._update_cfg("hidden_dropout_prob", cfg_dict["dropout_rate"])
+        self._update_cfg("attention_probs_dropout_prob", cfg_dict["dropout_rate"])
+        self._update_cfg(
+            "relative_attention_num_buckets", cfg_dict["relative_attention_num_buckets"]
+        )
+        self._update_cfg("embedding_dropout_prob", cfg_dict["dropout_rate"])
+        self._update_cfg("initializer_range", cfg_dict["initializer_factor"])
+        self._update_cfg("layernorm_eps", cfg_dict["layer_norm_epsilon"])
+        self._update_cfg("head_size", cfg_dict["d_kv"])
+        if "tie_word_embeddings" in self.libai_cfg:
+            self._update_cfg("tie_word_embeddings", cfg_dict.get("tie_word_embeddings", True))
+
+        # update libai_cfg by kwargs
+        for k, v in self.kwargs.items():
+            self._update_cfg(k, v)
+
+        self._update_cfg_log()
+
+
+class T5LoaderLibai(ModelLoaderLiBai):
+    def __init__(self, model, libai_cfg, pretrained_model_path, **kwargs):
+        super().__init__(model, libai_cfg, pretrained_model_path, **kwargs)
+        self.base_model_prefix_2 = "mt5_model"

projects/MT5/utils/mt5_metrc_printer.py

+import datetime
+import logging
+import time
+
+from libai.utils.events import EventWriter, get_event_storage
+
+
+class MT5MetricPrinter(EventWriter):
+    """
+    Print **MT5** metrics to the terminal, including
+    iteration time, ETA, memory, all losses, and the learning rate.
+    It also applies smoothing using a window of 20 elements.
+    It's meant to print MT5 metrics in MT5 ways.
+    To print something in more customized ways, please implement a similar printer by yourself.
+    """
+
+    def __init__(self, batch_size, max_iter, log_period):
+        """
+        Args:
+            max_iter (int): the maximum number of iterations to train.
+                Used to compute ETA.
+        """
+        self.logger = logging.getLogger("libai." + __name__)
+        self._batch_size = batch_size
+        self._max_iter = max_iter
+        self._last_write = None
+        self._log_period = log_period
+
+    def write(self):
+        storage = get_event_storage()
+        iteration = storage.iter
+        consumed_samples = storage.samples
+
+        try:
+            done_tokens = storage.history("done_tokens").avg(self._log_period)
+            token_time = storage.history("time").avg(self._log_period)
+        except KeyError:
+            done_tokens = None
+
+        try:
+            correct_tokens = storage.history("correct_tokens").avg(self._log_period)
+            denominator = storage.history("denominator").avg(self._log_period)
+            acc_mlm = correct_tokens / denominator
+        except KeyError:
+            acc_mlm = None
+
+        if iteration == self._max_iter:
+            # This hook only reports training progress (loss, ETA, etc) but not other data,
+            # therefore do not write anything after training succeeds, even if this method
+            # is called.
+            return
+
+        try:
+            data_time = storage.history("data_time").avg(self._log_period)
+        except KeyError:
+            # they may not exist in the first few iterations (due to warmup)
+            # or when SimpleTrainer is not used
+            data_time = None
+
+        eta_string = None
+        try:
+            iter_time = storage.history("time").global_avg()
+            eta_seconds = storage.history("time").median(1000) * (self._max_iter - iteration - 1)
+            storage.put_scalar("eta_seconds", eta_seconds, smoothing_hint=False)
+            eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+        except KeyError:
+            iter_time = None
+            # estimate eta on our own - more noisy
+            if self._last_write is not None:
+                estimate_iter_time = (time.perf_counter() - self._last_write[1]) / (
+                    iteration - self._last_write[0]
+                )
+                eta_seconds = estimate_iter_time * (self._max_iter - iteration - 1)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+            self._last_write = (iteration, time.perf_counter())
+
+        try:
+            lr = "{:.2e}".format(storage.history("lr").latest())
+        except KeyError:
+            lr = "N/A"
+
+        max_mem_mb = None
+
+        # NOTE: max_mem is parsed by grep in "dev/parse_results.sh"
+        self.logger.info(
+            " {eta}  {iter}  {sample}  {losses}  {time}  {data_time}  {tpt}  lr: {lr}  {memory} "
+            " {tokens_speed}  {acc_mlm}".format(
+                eta=f"eta: {eta_string}" if eta_string else "",
+                iter=f"iteration: {iteration}/{self._max_iter}",
+                sample=f"consumed_samples: {consumed_samples}",
+                losses="  ".join(
+                    [
+                        "{}: {:.4g}".format(k, v.median(200))
+                        for k, v in storage.histories().items()
+                        if "loss" in k
+                    ]
+                ),
+                time="time: {:.4f} s/iter  ".format(iter_time) if iter_time is not None else "",
+                data_time="data_time: {:.4f} s/iter".format(data_time)
+                if data_time is not None
+                else "",
+                tpt="total_throughput: {:.2f} samples/s".format(self._batch_size / iter_time)
+                if iter_time is not None
+                else "",
+                lr=lr,
+                memory="max_mem: {:.0f}M".format(max_mem_mb) if max_mem_mb is not None else "",
+                tokens_speed="tokens_throughput: {:.4f} tokens/s".format(done_tokens / token_time)
+                if done_tokens is not None
+                else "",
+                acc_mlm="acc_mlm: {:.4f}".format(acc_mlm) if acc_mlm is not None else "",
+            )
+        )

projects/MagicPrompt/configs/gpt2_dataset.py

+from libai.config import LazyCall
+from libai.data import build_nlp_train_val_test_loader
+from configs.common.data.gpt_dataset import tokenization, dataloader
+from libai.tokenizer import GPT2Tokenizer
+from libai.data.datasets import GPT2Dataset
+from libai.data.data_utils import get_indexed_dataset
+
+
+data_prefix = "/data/home/magicprompt/train/en_train_mmap_text_sentence"
+
+tokenization.tokenizer = LazyCall(GPT2Tokenizer)(
+    vocab_file="/data/home/magicprompt/vocab.json",
+    merges_file="/data/home/magicprompt/merges.txt",
+    do_lower_case=True,
+    do_chinese_wwm=True,
+)
+tokenization.append_eod = False
+
+dataloader.train = LazyCall(build_nlp_train_val_test_loader)(
+    dataset=[
+        LazyCall(GPT2Dataset)(
+            name="gpt-2",
+            data_prefix=data_prefix,
+            indexed_dataset=LazyCall(get_indexed_dataset)(
+                data_prefix=data_prefix,
+                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,
+)

projects/MagicPrompt/configs/gpt2_inference.py

+from configs.common.models.gpt import cfg
+from libai.config import LazyCall
+from libai.tokenizer.tokenization_gpt2 import GPT2Tokenizer
+from projects.MagicPrompt.gpt2 import GPTModel, GPTForPreTraining
+from configs.common.data.gpt_dataset import tokenization
+from configs.common.train import train
+
+
+cfg.update(
+    # Model
+    embedding_dropout_prob=0,
+    attention_dropout_prob=0,
+    output_dropout_prob=0,
+    bias_gelu_fusion=False,
+    bias_dropout_fusion=False,
+    scale_mask_softmax_fusion=False,
+    apply_query_key_layer_scaling=False,
+    apply_residual_post_layernorm=False,
+    amp_enabled=False,
+    num_attention_heads=12,
+    hidden_size=768,
+    ffn_hidden_size=4 * 768,
+    hidden_layers=12,
+    max_seq_length=1024,
+    initializer_range=0.02,
+    vocab_size=50304,
+    # Inference
+    is_encoder_decoder=False,
+    max_length=20,
+    min_length=0,
+    do_sample=False,
+    early_stopping=False,
+    num_beams=1,
+    num_beam_groups=1,
+    diversity_penalty=0.0,
+    temperature=1.0,
+    top_k=50,
+    top_p=1.0,
+    typical_p=1.0,
+    repetition_penalty=1.0,
+    length_penalty=1.0,
+    no_repeat_ngram_size=0,
+    encoder_no_repeat_ngram_size=0,
+    num_return_sequences=1,
+    chunk_size_feed_forward=0,
+    output_scores=False,
+    forced_bos_token_id=None,
+    forced_eos_token_id=None,
+    remove_invalid_values=False,
+    exponential_decay_length_penalty=None,
+    use_cache=True,
+    # Tokenizer
+    pad_token_id=0,
+    eos_token_id=50256,
+    bos_token_id=50256,
+    sep_token_id=None,
+    decoder_start_token_id=None,
+    # train
+    pretrained_model_path="/data/home/magicprompt",
+)
+
+
+model = LazyCall(GPTModel)(cfg=cfg)
+pretrain_model = LazyCall(GPTForPreTraining)(cfg=cfg)
+tokenization.tokenizer = LazyCall(GPT2Tokenizer)(
+    vocab_file="/data/home/magicprompt/vocab.json",
+    merges_file="/data/home/magicprompt/merges.txt",
+    add_bos_token=True,
+)

projects/MagicPrompt/configs/gpt2_training.py

+from libai.config import LazyCall
+from libai.evaluation import PPLEvaluator
+from projects.MagicPrompt.configs.gpt2_inference import pretrain_model as model
+from projects.MagicPrompt.configs.gpt2_dataset import dataloader, tokenization
+from configs.common.optim import optim
+
+from libai.scheduler import WarmupExponentialLR
+
+from configs.common.train import train
+from configs.common.models.graph import graph
+
+
+vocab_file = "/data/home/magicprompt/vocab.json"
+merge_files = "/data/home/magicprompt/merges.txt"
+train_data_prefix = "/data/home/magicprompt/train/en_train_mmap_text_sentence"
+
+tokenization.tokenizer.vocab_file = vocab_file
+tokenization.tokenizer.merges_file = merge_files
+dataloader.train.dataset[0].data_prefix = train_data_prefix
+dataloader.train.dataset[0].indexed_dataset.data_prefix = train_data_prefix
+
+# gpt2 model config
+model.cfg.pretrained_model_path = None
+model.cfg.embedding_dropout_prob = 0.1
+model.cfg.attention_dropout_prob = 0.1
+model.cfg.output_dropout_prob = 0.1
+model.cfg.num_attention_heads = 12
+model.cfg.hidden_size = 768
+model.cfg.ffn_hidden_size = 4 * 768
+model.cfg.hidden_layers = 12
+model.cfg.max_seq_length = 1024
+model.cfg.initializer_range = 0.02
+model.cfg.vocab_size = 50257
+model.cfg.layernorm_epsilon = 1e-5
+model.cfg.use_scaled_init_for_output_weights = True
+model.cfg.bias_gelu_fusion = True
+model.cfg.bias_dropout_fusion = True
+model.cfg.scale_mask_softmax_fusion = True
+model.cfg.apply_query_key_layer_scaling = True
+model.cfg.apply_residual_post_layernorm = False
+model.cfg.amp_enabled = True
+
+train.input_placement_device = "cpu"
+
+train.dist.pipeline_num_layers = model.cfg.hidden_layers
+
+for ds in dataloader.train.dataset:
+    ds.max_seq_length = model.cfg.max_seq_length
+
+optim.lr = 5.0e-05
+
+train.update(
+    dict(
+        output_dir="projects/MagicPrompt/oneflow_magicprompt",
+        train_micro_batch_size=4,
+        test_micro_batch_size=4,
+        train_epoch=33,
+        train_iter=10000,
+        log_period=50,
+        amp=dict(enabled=True),
+        warmup_ratio=0,
+        checkpointer=dict(period=8000, max_to_keep=20),
+        dist=dict(
+            data_parallel_size=1,
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            # pipeline_num_layers=model.cfg.hidden_layers,
+        ),
+        scheduler=LazyCall(WarmupExponentialLR)(
+            warmup_factor=0.0,
+            gamma=1.0,
+            warmup_method="linear",
+            warmup_iter=0.0,
+        ),
+        evaluation=dict(
+            enabled=True,
+            evaluator=LazyCall(PPLEvaluator)(),
+            eval_iter=250,
+            eval_period=4000,
+        ),
+        rdma_enabled=False,
+    )
+)