Merge branch 'duncan/fix-mamba-text-gen' into 'main'

Fix model instantiation for text gen server See merge request ADLR/megatron-lm!1906
TrustLLMeu · Aug 15, 2024 · 2d487b1 · 2d487b1
2 parents e8f8e63 + 7b8d43c
commit 2d487b1
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 7 deletions.
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
@@ -190,10 +190,10 @@ def validate_args(args, defaults={}):
     # Checks.
     if args.rank == 0:
         print('using world size: {}, data-parallel size: {}, '
-              'context-parallel size: {} '
+              'context-parallel size: {}, '
               'tensor-model-parallel size: {}, '
-              'encoder-tensor-model-parallel size: {}'
-              'pipeline-model-parallel size: {} '
+              'encoder-tensor-model-parallel size: {}, '
+              'pipeline-model-parallel size: {}, '
               'encoder-pipeline-model-parallel size: {}'.format(
                   args.world_size, args.data_parallel_size,
                   args.context_parallel_size,

diff --git a/pretrain_mamba.py b/pretrain_mamba.py
@@ -75,7 +75,9 @@ def model_provider(pre_process=True, post_process=True) -> MambaModel:
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
-        position_embedding_type=args.position_embedding_type
+        position_embedding_type=args.position_embedding_type,
+        rotary_percent=args.rotary_percent,
+        rotary_base=args.rotary_base
     )
 
     for l in range(model.decoder.num_layers_per_pipeline_rank):

diff --git a/tools/run_mamba_text_generation_server.py b/tools/run_mamba_text_generation_server.py
@@ -63,9 +63,11 @@ def model_provider(pre_process=True, post_process=True) -> MambaModel:
         hybrid_override_pattern=args.hybrid_override_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
-        parallel_output=True,
+        parallel_output=False,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
-        position_embedding_type=args.position_embedding_type
+        position_embedding_type=args.position_embedding_type,
+        rotary_percent=args.rotary_percent,
+        rotary_base=args.rotary_base
     )
 
     for l in range(model.decoder.num_layers_per_pipeline_rank):

diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py
@@ -82,7 +82,8 @@ def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megat
             parallel_output=False,
             share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
             position_embedding_type=args.position_embedding_type,
-            rotary_percent=args.rotary_percent
+            rotary_percent=args.rotary_percent,
+            rotary_base=args.rotary_base
         )
 
     return model