[megatron] add logging jsonl (#4908)

Author: Jintao-Huang

swift/megatron/init.py

@@ -11,7 +11,7 @@
 from packaging import version
 
 from swift.llm import git_clone_github
-from swift.utils import get_logger, is_megatron_available, safe_ddp_context, subprocess_run
+from swift.utils import JsonlWriter, get_logger, is_master, is_megatron_available, safe_ddp_context, subprocess_run
 
 logger = get_logger()
 
@@ -60,12 +60,18 @@ def _patch_training_log():
     from megatron.training.training import num_floating_point_operations
     from megatron.core.num_microbatches_calculator import get_num_microbatches
     from megatron.training.utils import reduce_max_stat_across_model_parallel_group, report_memory
+    jsonl_writer = None
 
     # Code borrowed from NVIDIA/Megatron-LM
     def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_rate, iteration, loss_scale,
                      report_memory_flag, skipped_iter, grad_norm, params_norm, num_zeros_in_grad):
         """Log training information such as losses, timing, ...."""
+        nonlocal jsonl_writer
         args = get_args()
+        if is_master() and jsonl_writer is None:
+            logging_path = os.path.join(args.save, 'logging.jsonl')
+            logger.info(f'logging_path: {logging_path}')
+            jsonl_writer = JsonlWriter(logging_path, enable_async=True)
         timers = get_timers()
         writer = get_tensorboard_writer()
         wandb_writer = get_wandb_writer()
@@ -209,6 +215,8 @@ def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_r
             mtp_loss_scale = 1 / get_num_microbatches()
             MTPLossLoggingHelper.track_mtp_metrics(mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict)
         if iteration % args.log_interval == 0 or iteration == 1:
+            origin_total_loss_dict = total_loss_dict.copy()
+
             if args.record_memory_history and is_last_rank():
                 snapshot = torch.cuda.memory._snapshot()
                 from pickle import dump
@@ -277,6 +285,26 @@ def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_r
                 report_memory_flag = False
             timers.log(timers_to_log, normalizer=args.log_interval)
 
+            if is_master():
+                logs = {}
+                for key in origin_total_loss_dict:
+                    if key not in [advanced_iters_key, skipped_iters_key, nan_iters_key]:
+                        avg = origin_total_loss_dict[key].item() / float(
+                            max(1, origin_total_loss_dict[advanced_iters_key]))
+                        logs[key] = round(avg, 8)
+                if grad_norm is not None:
+                    logs['grad_norm'] = round(grad_norm, 8)
+                if params_norm is not None:
+                    logs['params_norm'] = round(params_norm, 8)
+                logs['learning_rate'] = round(learning_rate, 8)
+                logs['elapsed_time_per_iteration'] = round(elapsed_time_per_iteration, 8)
+                if args.log_throughput:
+                    logs['throughput'] = round(throughput, 8)
+                logs['loss_scale'] = round(loss_scale, 8)
+                logs['consumed_samples'] = args.consumed_train_samples
+                logs['global_step/max_steps'] = f'{iteration}/{args.train_iters}'
+                jsonl_writer.append(logs)
+
         return report_memory_flag
 
     training.training_log = training_log

swift/megatron/trainers/base.py

@@ -1,5 +1,6 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
+import os
 import time
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
@@ -18,7 +19,7 @@
 from megatron.training.checkpointing import load_checkpoint
 from packaging import version
 
-from swift.utils import get_logger
+from swift.utils import JsonlWriter, get_logger, is_master
 from ..utils import adapter_state_dict_context, prepare_mcore_model
 from .utils import get_swift_datasets_provider
 
@@ -30,6 +31,9 @@ class BaseMegatronTrainer(ABC):
     def __init__(self, args):
         self.args = args
         self.stimer = StragglerDetector()
+        logging_path = os.path.join(args.save, 'logging.jsonl')
+        logger.info(f'logging_path: {logging_path}')
+        self.jsonl_writer = JsonlWriter(logging_path, enable_async=True)
         self._patch_megatron()
 
     @contextmanager
@@ -305,9 +309,11 @@ def evaluate(self,
         timers.log(['evaluate'])
 
         rerun_state_machine.set_mode(rerun_mode)
-
-        rerun_state_machine.set_mode(rerun_mode)
-
+        if is_master():
+            logs = {}
+            for key, val in total_loss_dict.items():
+                logs[f'eval_{key}'] = round(val.item(), 8)
+            self.jsonl_writer.append(logs)
         return total_loss_dict, collected_non_loss_data, False
 
     def save_checkpoint(self, *args, **kwargs):