fix bug (#123)

* Update word_embedding.py

* Update transformer_encoder.py

Author: hhou435

tencentpretrain/embeddings/word_embedding.py

@@ -9,7 +9,7 @@ class WordEmbedding(nn.Module):
 
     def __init__(self, args, vocab_size):
         super(WordEmbedding, self).__init__()
-        if args.tensor_model_parallel_size > 1:
+        if hasattr(args, "tensor_model_parallel_size") and args.tensor_model_parallel_size > 1:
             self.embedding = mpu.VocabParallelEmbedding(vocab_size, args.emb_size)
         else:
             self.embedding = nn.Embedding(vocab_size, args.emb_size)

tencentpretrain/encoders/transformer_encoder.py

@@ -20,15 +20,18 @@ def __init__(self, args):
         self.relative_position_embedding = args.relative_position_embedding
         self.rotary_position_embedding = args.rotary_position_embedding
         self.has_residual_attention = args.has_residual_attention
-        self.tensor_model_parallel_size = args.tensor_model_parallel_size
+        if hasattr(args, "tensor_model_parallel_size"):
+            self.tensor_model_parallel_size = args.tensor_model_parallel_size
+        else:
+            self.tensor_model_parallel_size = 1
 
         if self.relative_position_embedding:
             args.relative_pos_emb = RelativePositionEmbedding(bidirectional=True, heads_num=args.heads_num,
                                                               num_buckets=args.relative_attention_buckets_num)
         elif self.rotary_position_embedding:
             args.freqs_cis = precompute_freqs_cis(args.hidden_size // args.heads_num, args.max_seq_length * 2)
 
-        if "deepspeed_checkpoint_activations" in args:
+        if hasattr(args, "deepspeed_checkpoint_activations"):
             self.deepspeed_checkpoint_activations = args.deepspeed_checkpoint_activations
             self.deepspeed_checkpoint_layers_num = args.deepspeed_checkpoint_layers_num
         else: