Skip to content

Commit

Permalink
fix bug (#123)
Browse files Browse the repository at this point in the history
* Update word_embedding.py

* Update transformer_encoder.py
  • Loading branch information
hhou435 authored Jan 24, 2024
1 parent f5ce4e5 commit 6feab42
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 3 deletions.
2 changes: 1 addition & 1 deletion tencentpretrain/embeddings/word_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 5 additions & 2 deletions tencentpretrain/encoders/transformer_encoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 6feab42

Please sign in to comment.