add only_packing_no_scheduling for hybrid-cp

Signed-off-by: tailaim <[email protected]>

Author: xiaoyao0115

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -3342,9 +3342,6 @@ def wait_bucket_ready(self, bucket_id, empty_ok=False):
         # Wait for asynchronous / overlapped NCCL operations to complete.
         param_gather_event, mark_bucket_ready_to_use = self.param_gather_event_map.pop(bucket_id)
         param_gather_event.wait()
-        # debugmtl
-        if self.ag_stream is not None:
-            torch.cuda.current_stream().wait_stream(self.ag_stream)
         mark_bucket_ready_to_use()
 
     @torch.no_grad()

megatron/core/model_parallel_config.py

@@ -69,6 +69,14 @@ class ModelParallelConfig:
     When enabling hybrid_context_parallel, sft_sequence_packing must be true.
     """
 
+    hybrid_context_parallel_scheduler: str = 'balanced'
+    """
+    Scheduler for hybrid context parallel.
+    balanced: balanced scheduler for hybrid context parallel.
+    only_packing_no_scheduling: scheduling is already handled by the data sampler,
+    this scheduler only performs packing.
+    """
+
     sft_sequence_packing: bool = False
     """
     If true, enables sft sequence packing.

megatron/core/models/gpt/gpt_model.py

@@ -37,65 +37,6 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
 
-# #debugmtl
-# _ACT_GRAD_DUMP_COUNTERS = {}
-
-# def _sanitize_name(name: str) -> str:
-#     return str(name).replace('/', '_').replace('\\', '_').replace('.', '_').replace(' ', '_')
-
-# def _next_act_dump_index(rank: int, layer_name: str) -> int:
-#     key = (rank, layer_name)
-#     cnt = _ACT_GRAD_DUMP_COUNTERS.get(key, 0) + 1
-#     _ACT_GRAD_DUMP_COUNTERS[key] = cnt
-#     return cnt
-
-# def get_debug_hook(layer_name: str):
-#     """
-#     Tensor-level grad hook: save activation grad by (rank, layer_name, index).
-#     """
-#     import os
-#     def hook(grad: torch.Tensor):
-#         if grad is None:
-#             return
-
-#         rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
-
-#         # 基础目录自行改成你想要的
-#         base_dir = "/home/tailaim/act_grad_dump"
-#         if not base_dir:
-#             return
-
-#         try:
-#             idx = _next_act_dump_index(rank, layer_name)
-#             layer_dir = os.path.join(
-#                 base_dir,
-#                 f"rank_{rank}",
-#                 _sanitize_name(layer_name),
-#             )
-#             os.makedirs(layer_dir, exist_ok=True)
-#             file_path = os.path.join(layer_dir, f"grad_{idx:06d}.pt")
-
-#             # 只前几次写盘，避免太多文件
-#             if idx <= 16:
-#                 torch.save(grad.detach().cpu(), file_path)
-
-#             # 只在第一次写时打印一行日志
-#             if idx == 1:
-#                 try:
-#                     g_shape = tuple(grad.shape)
-#                     g_dtype = str(grad.dtype)
-#                 except Exception:
-#                     g_shape = "unknown"
-#                     g_dtype = "unknown"
-#                 print(
-#                     f"[Rank {rank}] Saved act grad: layer={layer_name}, "
-#                     f"idx={idx:06d}, shape={g_shape}, dtype={g_dtype}, path={file_path}"
-#                 )
-#         except Exception as e:
-#             print(f"[Rank {rank}] act grad dump failed for {layer_name}: {e}")
-
-#     return hook
-
 
 class GPTModel(LanguageModule):
     """GPT Transformer language model.

megatron/core/parallel_state.py

@@ -559,6 +559,7 @@ def initialize_model_parallel(
     create_gloo_process_groups: bool = True,
     high_priority_stream_groups: Optional[List[str]] = None,
     sharp_enabled_group: Optional[str] = None,
+    min_hybrid_context_parallel_size: int = 1,
 ) -> None:
     """Initialize model data parallel groups.
 
@@ -973,7 +974,12 @@ def initialize_model_parallel(
     if hybrid_context_parallel:
         # PyTorch is performing lazy initialization of the communicator group.
         # Therefore, we need to perform a nccl call to ensure that the communicator group is created.
-        group_sizes = [2**i for i in range(int(log2(data_parallel_size)))]
+        group_sizes = [
+            2**i
+            for i in range(
+                int(log2(min_hybrid_context_parallel_size)), int(log2(data_parallel_size))
+            )
+        ]
         if group_sizes[-1] * 2 == data_parallel_size:
             group_sizes.append(data_parallel_size)
         for group_size in group_sizes: