add megatron log

Author: yangbofun

swift/megatron/arguments/megatron_args.py

@@ -453,6 +453,8 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     overlap_grad_reduce: bool = False
     overlap_param_gather: bool = False
     overlap_param_gather_with_optimizer_step: bool = False
+    """DDP bucket size for grad reduce. Default: max(40M, 1M*dp_size). Only used when overlap_grad_reduce=True."""
+    bucket_size: Optional[int] = None
     align_grad_reduce: bool = True
     virtual_pipeline_model_parallel_size: Optional[int] = None
     microbatch_group_size_per_vp_stage: Optional[int] = None

swift/megatron/trainers/base.py

@@ -38,6 +38,8 @@
 from swift.trainers import dynamic_gradient_checkpointing
 from swift.trainers.utils import patch_modelscope_hub_timeout
 from swift.utils import deep_getattr, get_last_valid_indices, get_logger, is_last_rank, is_master, ms_logger_context
+from swift.utils import setup_megatron_logging
+from swift.utils.logger import add_file_handler_if_needed
 from .batch_sampler import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
 from .utils import (TrainerState, build_streaming_dataloader, get_batch_on_this_cp_rank, get_batch_on_this_pp_rank,
                     get_packed_seq_params)
@@ -97,6 +99,10 @@ def __init__(self, args, template: Template):
         if args.async_save and args.use_persistent_ckpt_worker:
             init_persistent_async_worker()
 
+        # Configure Megatron logging after main logger is ready so Megatron logs
+        # (e.g. DDP bucket info) appear in both console and log file.
+        setup_megatron_logging()
+
     def _load_checkpoint(self):
         args = self.args
         if not args.finetune:

swift/utils/__init__.py

@@ -9,7 +9,7 @@
                            is_lmdeploy_available, is_megatron_available, is_swanlab_available, is_trl_available,
                            is_unsloth_available, is_vllm_ascend_available, is_vllm_available, is_wandb_available)
 from .io_utils import JsonlWriter, append_to_jsonl, get_file_mm_type, read_from_jsonl, write_to_jsonl
-from .logger import get_logger, ms_logger_context
+from .logger import get_logger, ms_logger_context, setup_megatron_logging
 from .np_utils import get_seed, stat_array, transform_jsonl_to_df
 from .processor_utils import Processor, ProcessorMixin
 from .safetensors import LazyTensor, SafetensorLazyLoader, StreamingSafetensorSaver

swift/utils/logger.py

@@ -154,3 +154,43 @@ def add_file_handler_if_needed(logger, log_file, file_mode, log_level):
         file_handler.setFormatter(logger_format)
         file_handler.setLevel(log_level)
         logger.addHandler(file_handler)
+
+
+def setup_megatron_logging(log_file: Optional[str] = None):
+    """Configure Megatron loggers to output to the same stream/file as swift logs.
+
+    Call this after init_megatron_env() so Megatron logs (e.g. DDP bucket info,
+    param_and_grad_buffer) appear in ms-swift log files.
+
+    Reuses Swift's handlers to avoid log interleaving/corruption when both
+    Swift and Megatron write concurrently (shared handler lock serializes writes).
+
+    Args:
+        log_file: Optional log file path. If None, only uses Swift's StreamHandler.
+            When using tee (e.g. 2>&1 | tee node_rank_0.log), stderr is captured.
+    """
+    if importlib.util.find_spec('torch') is not None:
+        is_worker0 = int(os.getenv('LOCAL_RANK', -1)) in {-1, 0}
+    else:
+        is_worker0 = True
+    if not is_worker0:
+        return
+    log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
+    log_level = getattr(logging, log_level, logging.INFO)
+    megatron_logger = logging.getLogger('megatron.core')
+    megatron_logger.setLevel(log_level)
+    megatron_logger.propagate = False  # Don't propagate to root (filters at ERROR)
+    # Reuse Swift logger's handlers to avoid interleaved/corrupted output
+    swift_logger = logging.getLogger('swift')
+    for handler in swift_logger.handlers:
+        if handler not in megatron_logger.handlers:
+            megatron_logger.addHandler(handler)
+    if log_file is not None:
+        for h in megatron_logger.handlers:
+            if isinstance(h, logging.FileHandler) and getattr(h, 'baseFilename', '') == log_file:
+                break
+        else:
+            file_handler = logging.FileHandler(log_file, 'a')
+            file_handler.setFormatter(logger_format)
+            file_handler.setLevel(log_level)
+            megatron_logger.addHandler(file_handler)