Merge pull request #87 from epfLLM/evalonly_and_wbresume

Added eval only and wandb resume options
epfLLM · Nov 6, 2023 · 820f102 · 820f102
2 parents 402f7e8 + 47e6f9c
commit 820f102
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/megatron/arguments.py b/megatron/arguments.py
@@ -536,10 +536,12 @@ def _add_logging_args(parser):
                        help='Project name for Weights & Biases.')
     group.add_argument('--wandb_entity', type=str, default="meditron",
                        help='Entity/team name for Weights & Biases.')
+    group.add_argument('--wandb_name', type=str, default=None,
+                        help='Name for this run, alternatively can set `WANDB_NAME`.')
     group.add_argument('--wandb_id',type=str,default=None,
                        help="Unique ID to identify this run, alternatively can set `WANDB_RUN_ID`.")
-    group.add_argument('--wandb_resume',action="store_true",
-                       help="If set, we resume logging for the id given instead of launching a new run (errors if id given and resume=False).")
+    group.add_argument('--wandb_resume',type=str,default="allow",
+                       help="If set, we resume logging for the id given instead of launching a new run (errors if id given and resume=None).")
     group.add_argument("--wandb_api_key",type=str,default=None,
                        help="API key for Weights & Biases, needs to be set if not set in environment variable `WANDB_API_KEY`.")
     group.add_argument("--metrics", default=[], nargs="+", choices=list(METRICS) + ["all"],
@@ -878,6 +880,8 @@ def _add_distributed_args(parser):
 
 def _add_validation_args(parser):
     group = parser.add_argument_group(title='validation')
+    group.add_argument('--eval_only', action='store_true',
+                        help='Run evaluation only.')
     group.add_argument('--eval_iters', type=int, default=100,
                        help='Number of iterations to run for evaluation'
                        'validation/test for.')

diff --git a/megatron/training.py b/megatron/training.py
@@ -936,6 +936,11 @@ def build_train_valid_test_data_iterators(build_train_valid_test_datasets_provid
     args.do_valid = flags[1].item()
     args.do_test = flags[2].item()
 
+    if args.eval_only:
+        args.do_train = False
+        args.do_valid = False
+        args.do_test = True
+
     # Build iterators.
     dl_type = args.dataloader_type
     assert dl_type in ['single', 'cyclic']

diff --git a/megatron/wandb_logger.py b/megatron/wandb_logger.py
@@ -60,6 +60,7 @@ def from_args(args)->'WandBConfig':
                            log_interval=args.log_interval,
                            config=args,entity=args.wandb_entity,
                            project=args.wandb_project,
+                           name=args.wandb_name,
                            run_id=args.wandb_id,
                            resume=args.wandb_resume,
                            api_key=args.wandb_api_key,