Convert t5 to use config object.

Author: jaredcasper

megatron/model/t5_model.py

@@ -11,9 +11,7 @@
 from megatron.model import LayerNorm
 from megatron.model.utils import (
     openai_gelu,
-    get_linear_layer,
-    init_method_normal,
-    scaled_init_method_normal
+    get_linear_layer
 )
 from .module import MegatronModule
 
@@ -43,17 +41,12 @@ class T5LMHead(MegatronModule):
 
     Arguments:
         mpu_vocab_size: model parallel size of vocabulary.
-        hidden_size: hidden size
-        init_method: init method for weight initialization
-        layernorm_epsilon: tolerance for layer norm divisions
         parallel_output: wether output logits being distributed or not.
     """
 
     def __init__(self, mpu_vocab_size, parallel_output):
         super(T5LMHead, self).__init__()
 
-        args = get_args()
-
         self.bias = torch.nn.Parameter(torch.zeros(mpu_vocab_size))
         self.bias.model_parallel = True
         self.bias.partition_dim = 0
@@ -72,37 +65,34 @@ class T5Model(MegatronModule):
     """T5 Language model."""
 
     def __init__(self,
+                 config,
                  num_tokentypes=0,
                  parallel_output=True,
                  pre_process=True,
                  post_process=True,
                  add_encoder=True,
                  add_decoder=True):
-        super(T5Model, self).__init__()
+        super().__init__(config=config)
         args = get_args()
 
         self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy
         self.parallel_output = parallel_output
-        init_method = init_method_normal(args.init_method_std)
-        scaled_init_method = scaled_init_method_normal(args.init_method_std,
-                                                       args.num_layers)
         self.pre_process = pre_process
         self.post_process = post_process
         self.add_encoder = add_encoder
         self.add_decoder = add_decoder
 
         self.language_model, self._language_model_key = get_language_model(
+            config=config,
             num_tokentypes=num_tokentypes,
             add_pooler=False,
             add_encoder=add_encoder,
             add_decoder=add_decoder,
             encoder_attn_mask_type=AttnMaskType.padding,
-            init_method=init_method,
-            scaled_init_method=scaled_init_method,
             pre_process=self.pre_process,
             post_process=self.post_process)
 
-        self.initialize_word_embeddings(init_method_normal)
+        self.initialize_word_embeddings()
 
         if self.post_process and self.add_decoder:
             self.lm_head = T5LMHead(

megatron/model/transformer.py

@@ -747,15 +747,14 @@ def __init__(self, config,
                                LayerType.retro_decoder_with_retriever,
                                LayerType.retro_encoder):
             self.inter_attention = ParallelAttention(
-                config.init_method,
-                config.output_layer_init_method,
+                config,
                 layer_number,
                 attention_type=AttnType.cross_attn)
             # Layernorm on the attention output.
             self.post_inter_attention_layernorm = LayerNorm(
                 config.hidden_size,
                 eps=config.layernorm_epsilon,
-                no_persist_layer_norm=config.no_persist_layer_norm,
+                no_persist_layer_norm=not config.persist_layer_norm,
                 sequence_parallel=config.sequence_parallel,
                 apply_layernorm_1p=args.apply_layernorm_1p)
 

pretrain_t5.py

@@ -17,6 +17,7 @@
 from megatron.model import T5Model
 from megatron.training import pretrain
 from megatron.utils import average_losses_across_data_parallel_group
+from megatron.arguments import core_transformer_config_from_args
 
 
 """
@@ -60,7 +61,9 @@ def model_provider(pre_process=True, post_process=True,
     """Build the model."""
 
     print_rank_0('building T5 model ...')
-    model = T5Model(num_tokentypes=0,
+    config = core_transformer_config_from_args(get_args())
+    model = T5Model(config=config,
+                    num_tokentypes=0,
                     parallel_output=True,
                     pre_process=pre_process,
                     post_process=post_process,