[fix] fix writer in training & engram allreduce mode (flagos-ai#1200)

### PR Category
<!-- One of [ Train | Inference | Compress | Serve | RL | Core |
Hardware | CICD | Tools | Others ] -->
Train
### PR Types
<!-- One of [ User Experience | New Features | Bug Fixes | Improvements
| Performance | Breaking Change| Deprecations | Test Case | Docs |
Others ] -->
Bug Fixes
### PR Description
<!-- Describe what you’ve done -->

- fix writer logic when it is None
- fix engram hang when parallel mode is allreduce

Author: zhaoyinglia

flagscale/train/megatron/training/arguments_fs.py

@@ -422,13 +422,13 @@ def _parse_recompute_refined_config(recom_config, recom_config_name):
                         f"[rank0]: We do not recommend using allreduce for engram embedding, this is deprecated and will be removed in a later version.",
                         DeprecationWarning,
                     )
-                    if self.args.engram_embedding_parallel_size is not None:
-                        warnings.warn(
-                            "[rank0]: If set the embedding_parallel_method to allreduce, "
-                            "the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel."
-                            "So the embedding_parallel_size is useless and set to None."
-                        )
-                        self.args.engram_embedding_parallel_size = None
+                if self.args.engram_embedding_parallel_size is not None:
+                    warnings.warn(
+                        "[rank0]: If set the embedding_parallel_method to allreduce, "
+                        "the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel."
+                        "So the embedding_parallel_size is useless and set to None."
+                    )
+                    self.args.engram_embedding_parallel_size = None
             elif self.args.engram_embedding_parallel_method == "alltoall":
                 assert (
                     self.args.engram_embedding_parallel_size is not None

flagscale/train/megatron/training/global_vars.py

@@ -6,7 +6,6 @@
 import signal
 import sys
 import torch
-import torch.distributed
 
 from datetime import timedelta
 
@@ -15,12 +14,13 @@
 from megatron.core.energy_monitor import EnergyMonitor
 from megatron.core.jit import disable_jit_fuser
 from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator, unset_num_microbatches_calculator
-from megatron.training.tokenizer import build_tokenizer
+from megatron.training.tokenizer import build_tokenizer ########## FlagScale Add ##########
 from megatron.training.dist_signal_handler import DistributedSignalHandler
 
+########## FlagScale Begin ##########
 from megatron.training.spiky_loss import SpikyLossDetector
 from megatron.plugin.utils import get_device_type_for_comm
-
+########## FlagScale End ##########
 
 _GLOBAL_ARGS = None
 _GLOBAL_TOKENIZER = None
@@ -32,8 +32,10 @@
 _GLOBAL_ENERGY_MONITOR = None
 _GLOBAL_SIGNAL_HANDLER = None
 
+########## FlagScale Begin ##########
 _GLOBAL_SPIKY_LOSS_DETECTOR = None
 _GLOBAL_EXTRA_VALID_DATASETS = None
+########## FlagScale End ##########
 
 
 
@@ -165,15 +167,13 @@ def set_global_writers(args):
     Note that this function should be called after calling finish_mpu_init.
     This is because we can know which rank is the last one after the rank mapping in finish_mpu_init.
     """
+    from .utils import is_last_rank
 
     assert args is not None
 
     _ensure_var_is_initialized(_GLOBAL_ARGS, 'args')
-
-    from .utils import is_last_rank
-    if is_last_rank():
-        _set_tensorboard_writer(args)
-        _set_one_logger(args)
+    _set_tensorboard_writer(args)
+    _set_one_logger(args)
 
     # build wandb writers for all processes in the dp group of the last rank
     from megatron.core import mpu
@@ -252,7 +252,7 @@ def _set_tensorboard_writer(args):
                                    'tensorboard writer')
 
     if hasattr(args, 'tensorboard_dir') and \
-       args.tensorboard_dir:
+       args.tensorboard_dir and args.rank == (args.world_size - 1):
         try:
             from torch.utils.tensorboard import SummaryWriter
             print('> setting tensorboard ...')
@@ -274,7 +274,6 @@ def _set_wandb_writer(args):
             raise ValueError("Please specify the wandb experiment name!")
 
         import wandb
-        rank = torch.distributed.get_rank()
         if args.wandb_save_dir:
             save_dir = args.wandb_save_dir
         else:
@@ -286,9 +285,8 @@ def _set_wandb_writer(args):
             # settings were.
             with open(wandb_config['kitchen_config_file'], "r") as f:
                 wandb_config['kitchen_config_file_contents'] = f.read()
+        rank = torch.distributed.get_rank()
         save_dir = os.path.join(save_dir, "rank-{}".format(rank))
-        os.makedirs(save_dir, exist_ok=True)
-
         wandb_id = f"{args.wandb_exp_name}-rank-{rank}"
         name = f'{args.wandb_exp_name}-rank-{rank}'
         group = args.wandb_exp_name
@@ -303,9 +301,7 @@ def _set_wandb_writer(args):
             'config': wandb_config}
         if args.wandb_entity:
             wandb_kwargs['entity'] = args.wandb_entity
-        if args.wandb_entity:
-            wandb_kwargs['entity'] = args.wandb_entity
-
+        os.makedirs(wandb_kwargs['dir'], exist_ok=True)
         if args.wandb_mode == 'online' or args.wandb_api_key:
             assert args.wandb_api_key, 'wandb_api_key is required for online mode'
             wandb.login(key=args.wandb_api_key)
@@ -406,6 +402,7 @@ def destroy_global_vars():
     _GLOBAL_SIGNAL_HANDLER = None
 
 
+########## FlagScale Begin ##########
 def get_spiky_loss_detector():
     """Return spiky loss detector."""
     _ensure_var_is_initialized(_GLOBAL_SPIKY_LOSS_DETECTOR, "spiky loss detector")
@@ -427,4 +424,5 @@ def get_extra_valid_datasets():
 def set_extra_valid_datasets(extra_valid_datasets):
     """Set extra_valid datasets.""" ""
     global _GLOBAL_EXTRA_VALID_DATASETS
-    _GLOBAL_EXTRA_VALID_DATASETS = extra_valid_datasets
+    _GLOBAL_EXTRA_VALID_DATASETS = extra_valid_datasets
+########## FlagScale End ##########