WIP

Author: thomasw21

megatron/model/__init__.py

@@ -18,6 +18,7 @@
 from .distributed import DistributedDataParallel
 from .bert_model import BertModel
 from .gpt_model import GPTModel, GPTModelPipe
+from .shared_t5_model import SharedT5ModelPipe
 from .t5_model import T5Model
 from .language_model import get_language_model
 from .module import Float16Module

megatron/model/gpt_model.py

@@ -21,7 +21,7 @@
 from megatron import get_args
 from megatron import mpu
 from megatron.enums import AttnMaskType
-from .module import MegatronModule, fp32_to_float16
+from .module import MegatronModule, fp32_to_16bit
 
 from .language_model import parallel_lm_logits
 from .language_model import get_language_model
@@ -213,9 +213,9 @@ def __init__(
 
         def _to_float16(inputs):
             if args.fp16:
-                return fp32_to_float16(inputs, lambda v: v.half())
+                return fp32_to_16bit(inputs, lambda v: v.half())
             elif args.bf16:
-                return fp32_to_float16(inputs, lambda v: v.bfloat16())
+                return fp32_to_16bit(inputs, lambda v: v.bfloat16())
             else:
                 return inputs
 

megatron/model/module.py

@@ -122,7 +122,7 @@ def conversion_helper(val, conversion):
     return rtn
 
 
-def fp32_to_float16(val, float16_convertor):
+def fp32_to_16bit(val, float16_convertor):
     """Convert fp32 `val` to fp16/bf16"""
     def half_conversion(val):
         val_typecheck = val
@@ -168,7 +168,7 @@ def float16_convertor(val):
 
     def forward(self, *inputs, **kwargs):
         if mpu.is_pipeline_first_stage():
-            inputs = fp32_to_float16(inputs, self.float16_convertor)
+            inputs = fp32_to_16bit(inputs, self.float16_convertor)
         outputs = self.module(*inputs, **kwargs)
         if mpu.is_pipeline_last_stage():
             outputs = float16_to_fp32(outputs)

megatron/model/shared_t5_model.py

@@ -0,0 +1,189 @@
+import torch
+from deepspeed import PipelineModule
+from deepspeed.runtime.pipe import TiedLayerSpec, LayerSpec
+from torch.nn import LayerNorm
+
+from megatron.enums import AttnMaskType, LayerType
+
+from megatron.model.transformer import ParallelTransformerLayerPipe
+
+from megatron.model.language_model import EmbeddingPipe, parallel_lm_logits
+
+from megatron.model.utils import init_method_normal, scaled_init_method_normal
+
+from megatron import get_args, mpu
+
+from megatron.model.module import MegatronModule, fp32_to_16bit, float16_to_fp32
+
+def cross_entropy(output, labels):
+    labels, loss_mask = labels[0], labels[1]
+
+    losses = mpu.vocab_parallel_cross_entropy(output.contiguous().float(), labels)
+
+    expected_number_of_tokens = loss_mask.sum()
+
+    loss_mask = loss_mask.view(-1)
+    loss = torch.sum(losses.view(-1) * loss_mask) / expected_number_of_tokens
+    return loss
+
+class SharedT5ModelPipe(PipelineModule, MegatronModule):
+    """Share encoder decoder language model."""
+
+    def __init__(
+        self,
+        num_tokentypes=0,
+        parallel_output=True,
+    ):
+        args = get_args()
+        self.parallel_output = parallel_output
+
+        init_method = init_method_normal(args.init_method_std)
+
+        self.specs = []
+
+        def _to_16bit(inputs):
+            if args.fp16:
+                return fp32_to_16bit(inputs, lambda v: v.half())
+            elif args.bf16:
+                return fp32_to_16bit(inputs, lambda v: v.bfloat16())
+            else:
+                return inputs
+
+        self.specs.append(lambda inputss: tuple(_to_16bit(inputs) for inputs in inputss))
+
+        # Embedding layer
+        self.specs.append(TiedLayerSpec('embed',
+                                        EmbeddingPipe,
+                                        args.hidden_size,
+                                        args.padded_vocab_size,
+                                        args.hidden_dropout,
+                                        init_method=init_method,
+                                        num_tokentypes=num_tokentypes,
+                                        tied_weight_attr='word_embeddings_weight'))
+
+        assert hasattr(args, 'attn_mask'), "Deepspeed integration should have attention mask s"
+        if args.fp32_residual_connection:
+            self.specs.append(lambda x: x.transpose(0, 1).contiguous().float())
+        else:
+            self.specs.append(lambda x: x.transpose(0, 1).contiguous())
+
+        ### -----  Encoder -----
+        for layer_idx in range(args.num_layers):
+            self.specs.append(
+                TiedLayerSpec(
+                    f"block_{layer_idx}",
+                    ParallelTransformerLayerPipe,
+                    init_method=init_method,
+                    # Inputs: (input_tokens, target_tokens,
+                    forward_fn=lambda module, *inputs: ,
+                    output_layer_init_method=scaled_init_method_normal(args.init_method_std,
+                                                                       args.num_layers),
+                    layer_type=LayerType.encoder,
+                    layer_number=layer_idx,
+                    self_attn_mask_type=AttnMaskType.causal,
+                    tied_weight_attrs=["input_layernorm", "self_attention", "post_attention_layernorm", "mlp"]
+                ))
+
+        # Final layernorm after encoder layers
+        self.specs.append(
+            TiedLayerSpec(
+                "final_layer_norm",
+                LayerNorm,
+                args.hidden_size,
+                eps=args.layernorm_epsilon
+            ))
+
+        # Decoder
+        for layer_idx in range(args.num_layers):
+            self.specs.append(
+                TiedLayerSpec(
+                    f"block_{layer_idx}",
+                    ParallelTransformerLayerPipe,
+                    init_method=init_method,
+                    output_layer_init_method=scaled_init_method_normal(args.init_method_std,
+                                                                       args.num_layers),
+                    layer_number=layer_idx,
+                    layer_type=LayerType.decoder,
+                    self_attn_mask_type=AttnMaskType.padding,
+                    tied_weight_attrs=["input_layernorm", "self_attention", "post_attention_layernorm", "mlp"]
+                )
+            )
+
+        # Final layernorm after decoder layers
+        self.specs.append(
+            TiedLayerSpec(
+                "final_layer_norm",
+                LayerNorm,
+                args.hidden_size,
+                eps=args.layernorm_epsilon
+            ))
+
+        # Undo data format change
+        self.specs.append(lambda x: x.transpose(0, 1).contiguous())
+
+        def _logits_helper(embedding, lm_output):
+            """A wrapper to massage inputs/outputs from pipeline. """
+            return parallel_lm_logits(
+                lm_output,
+                embedding.word_embeddings_weight,
+                self.parallel_output)
+
+        self.specs.append(
+            TiedLayerSpec('embed',
+                          EmbeddingPipe,
+                          args.hidden_size,
+                          args.padded_vocab_size,
+                          args.hidden_dropout,
+                          init_method=init_method,
+                          num_tokentypes=num_tokentypes,
+                          forward_fn=_logits_helper,
+                          tied_weight_attr='word_embeddings_weight')
+        )
+
+        if not hasattr(args, 'attn_mask'):
+            # We drop attention mask from the pipeline
+            self.specs.append(lambda x: x[0])
+
+        # Final layernorm after transformer layers
+        self.specs.append(
+            TiedLayerSpec(
+                "final_layer_norm",
+                LayerNorm,
+                args.hidden_size,
+                eps=args.layernorm_epsilon
+            ))
+
+        # Undo data format change
+        self.specs.append(lambda x: x.transpose(0, 1).contiguous())
+
+        # Convert to fp32 if needed
+        if args.fp16 or args.bf16:
+            self.specs.append(float16_to_fp32)
+
+        if args.checkpoint_activations:
+            interval = args.checkpoint_num_layers
+        else:
+            interval = 0
+
+        from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
+        topo = PipeModelDataParallelTopology(num_pp=mpu.get_pipeline_model_parallel_world_size(),
+                                             num_mp=mpu.get_tensor_model_parallel_world_size(),
+                                             num_dp=mpu.get_data_parallel_world_size())
+
+        # here one can extend the regex to include more layers to be counted towards partitioning,
+        # e.g. 'type:transformer|embedding' will add up all the transformer blocks and also the first
+        # and last embedding layers and then partition that transformers+2 layers - so to get a good
+        # balance you may want to use less transformer layers
+        #
+        # caveat emptor: the current implementation of PP fails unless each stage has at least one
+        # transformer layer
+        if args.pp_partition_method is not None:
+            partition_method = args.pp_partition_method
+        else:
+            partition_method = 'type:transformer'
+
+        super().__init__(layers=self.specs,
+                         loss_fn=cross_entropy,
+                         topology=topo,
+                         activation_checkpoint_interval=interval,
+                         partition_method=partition_method)

megatron/text_generation_utils.py

@@ -26,7 +26,7 @@
 from megatron import get_args
 from megatron import get_tokenizer
 from megatron import mpu
-from megatron.utils import get_ltor_masks_and_position_ids, unwrap_model
+from megatron.utils import get_attention_masks_and_position_ids, unwrap_model
 from megatron.p2p_communication import recv_forward, send_forward
 
 # These are needed to unwrap the model, would be nice to put these in megatron.utils if possible?
@@ -42,7 +42,7 @@ def get_batch(context_tokens):
     # Move to GPU.
     tokens = context_tokens.view(args.micro_batch_size, -1).contiguous().cuda()
     # Get the attention mask and position ids.
-    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
+    attention_mask, _, position_ids = get_attention_masks_and_position_ids(
         tokens,
         tokenizer.eod,
         args.reset_position_ids,

megatron/utils.py

@@ -151,14 +151,16 @@ def check_adlr_autoresume_termination(iteration, model,
         sys.exit(0)
 
 
-def get_ltor_masks_and_position_ids(
+
+def get_attention_masks_and_position_ids(
         data,
         eod_token,
         reset_position_ids,
         reset_attention_mask,
         eod_mask_loss,
         prefix_indices,
         loss_on_targets_only,
+        ltor=True,
     ):
     """
     Build masks and position id for left to right model.
@@ -177,9 +179,10 @@ def get_ltor_masks_and_position_ids(
         att_mask_batch = micro_batch_size
     else:
         att_mask_batch = 1
-    attention_mask = torch.tril(torch.ones(
-        (att_mask_batch, seq_length, seq_length), device=data.device)).view(
-            att_mask_batch, 1, seq_length, seq_length)
+    attention_mask = torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)
+    if ltor:
+        attention_mask = torch.tril(attention_mask)
+    attention_mask = attention_mask.view(att_mask_batch, 1, seq_length, seq_length)
 
     # Loss mask.
     loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)

pretrain_gpt.py

@@ -25,7 +25,7 @@
 from megatron.data.gpt_dataset import build_train_valid_test_datasets, build_dataset_group
 from megatron.model import GPTModel, GPTModelPipe
 from megatron.training import pretrain
-from megatron.utils import get_ltor_masks_and_position_ids, get_prefix_indices
+from megatron.utils import get_attention_masks_and_position_ids, get_prefix_indices
 from megatron.utils import average_losses_across_data_parallel_group
 
 import deepspeed
@@ -110,7 +110,7 @@ def get_batch(data_iterator):
     tokens = tokens_[:, :-1].contiguous()
 
     # Get the masks and postition ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+    attention_mask, loss_mask, position_ids = get_attention_masks_and_position_ids(
         tokens,
         tokenizer.eod,
         args.reset_position_ids,
@@ -141,7 +141,7 @@ def get_batch_pipe(data):
     tokens = tokens_[:, :-1].contiguous()
 
     # Get the masks and position ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+    attention_mask, loss_mask, position_ids = get_attention_masks_and_position_ids(
         tokens,
         tokenizer.eod,
         args.reset_position_ids,

pretrain_prefix_lm.py

@@ -25,7 +25,7 @@
 from megatron.data.gpt_dataset import build_train_valid_test_datasets, build_dataset_group
 from megatron.model import GPTModel, GPTModelPipe
 from megatron.training import pretrain
-from megatron.utils import get_ltor_masks_and_position_ids, get_prefix_indices, reweight_loss_mask_
+from megatron.utils import get_attention_masks_and_position_ids, get_prefix_indices, reweight_loss_mask_
 from megatron.utils import average_losses_across_data_parallel_group
 
 import deepspeed
@@ -97,7 +97,7 @@ def get_batch(data_iterator):
     )
 
     # Get the masks and postition ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+    attention_mask, loss_mask, position_ids = get_attention_masks_and_position_ids(
         tokens,
         tokenizer.eod,
         args.reset_position_ids,
@@ -131,6 +131,7 @@ def get_batch_pipe(data):
     tokens = tokens_[:, :-1].contiguous()
 
     # Prefix
+    # TODO @thomasw21 actually since this step is random, we need to make sure that random state are synchronized. Otherwise we need to broadcast after this step.
     prefix_indices = get_prefix_indices(
         tokens,
         tokenizer.eod,
@@ -139,7 +140,7 @@ def get_batch_pipe(data):
     )
 
     # Get the masks and position ids.
-    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+    attention_mask, loss_mask, position_ids = get_attention_masks_and_position_ids(
         tokens,
         tokenizer.eod,
         args.reset_position_ids,