Merge branch 'main' into 'main'

Fix: Update OneLogger Instrumentation Points for Optimizer Init

See merge request ADLR/megatron-lm!3698

Author: jaredcasper

megatron/training/one_logger_utils.py

@@ -3,7 +3,7 @@
 
 from .global_vars import get_one_logger, get_args
 
-_one_logger_utils_version = "1.1.0-mlm"
+_one_logger_utils_version = "1.2.0-mlm"
 
 
 def get_timestamp_in_ms():

megatron/training/training.py

@@ -836,14 +836,12 @@ def pretrain(
 
     # Model, optimizer, and learning rate.
     timers('model-and-optimizer-setup', log_level=0).start(barrier=True)
-    app_metrics['app_build_optimizer_start_time'] = one_logger_utils.get_timestamp_in_ms()
     model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
         model_provider, model_type, checkpointing_context=checkpointing_context
     )
 
     timers('model-and-optimizer-setup').stop()
     print_datetime('after model, optimizer, and learning rate ' 'scheduler are built')
-    app_metrics['app_build_optimizer_finish_time'] = one_logger_utils.get_timestamp_in_ms()
     config = get_model_config(model[0])
 
     # Data stuff.
@@ -1234,6 +1232,7 @@ def setup_model_and_optimizer(
     model = get_model(model_provider_func, model_type)
     unwrapped_model = unwrap_model(model)
 
+    one_logger and one_logger.log_metrics({"app_build_optimzer_start_time": one_logger_utils.get_timestamp_in_ms()})
     kwargs = {}
     for f in dataclasses.fields(OptimizerConfig):
         if hasattr(args, f.name):
@@ -1252,6 +1251,7 @@ def setup_model_and_optimizer(
         default_skip_embedding_weight_decay=args.embedding_init_method_std is not None,
     )
     opt_param_scheduler = get_optimizer_param_scheduler(optimizer)
+    one_logger and one_logger.log_metrics({"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()})
 
     if args.moe_use_upcycling:
         torch.distributed.barrier()