diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 0f016473b6a..6211197b9b7 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -5,11 +5,11 @@ import torch from megatron.core import parallel_state -from megatron.core.pipeline_parallel.hybrid_cp_schedule import BalancedCPScheduler +from megatron.core.pipeline_parallel.dynamic_cp_schedule import BalancedCPScheduler from megatron.core.process_groups_config import ProcessGroupCollection -class HybridCPDataLoaderWrapper: +class DynamicCPDataLoaderWrapper: """ A wrapper class that wraps around an existing data_iterator. For every __next__ call, @@ -40,7 +40,7 @@ def __init__( self.tp_group = pg_collection.tp assert ( self.dp_cp_group is not None and self.dp_group is not None and self.tp_group is not None - ), "dp_cp_group, dp_group, tp_group must not be None when using hybrid context parallel" + ), "dp_cp_group, dp_group, tp_group must not be None when using dynamic context parallel" self.cp_balancing_scheduler = BalancedCPScheduler( max_seq_len_per_rank=self.config.max_seqlen_per_dp_cp_rank, dp_cp_group=self.dp_cp_group diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py index 42146d1acd2..82ef605aa6d 100644 --- a/megatron/core/datasets/gpt_dataset.py +++ b/megatron/core/datasets/gpt_dataset.py @@ -58,8 +58,8 @@ class GPTDatasetConfig(BlendedMegatronDatasetConfig): Set to 0 if sequence parallel is not enabled regardless of TP size. """ - hybrid_context_parallel: bool = False - """Option to enable hybrid context parallelism. When setting this to True, + dynamic_context_parallel: bool = False + """Option to enable dynamic context parallelism. When setting this to True, each sample should be divisible by the data parallel size * context parallel size * 2. If sequence parallel is enabled, it should be divisible by the data parallel size * context parallel size * sequence parallel size * 2. diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index bb913d97446..db7aee710a8 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1363,21 +1363,17 @@ def forward( """Forward.""" if packed_seq_params is not None: # If Dynamic CP group is provided, update TE DPA CP group - if packed_seq_params.cp_group is not None: - self.cp_group = packed_seq_params.cp_group - super().set_context_parallel_group( - self.cp_group, - torch.distributed.get_process_group_ranks(self.cp_group), - TEDotProductAttention.cp_stream, - self.cp_comm_type, - ) - # If cp_group is None but local_cp_size is provided, - # Indicates to turn off CP dynamically - elif packed_seq_params.local_cp_size is not None: - assert ( - packed_seq_params.local_cp_size == 1 - ), "local_cp_size must be == 1 if provided without cp_group" - super().set_context_parallel_group(None, None, None, self.cp_comm_type) + if packed_seq_params.local_cp_size is not None: + if packed_seq_params.local_cp_size == 1: + super().set_context_parallel_group(None, None, None, self.cp_comm_type) + else: + self.cp_group = packed_seq_params.cp_group + super().set_context_parallel_group( + self.cp_group, + torch.distributed.get_process_group_ranks(self.cp_group), + TEDotProductAttention.cp_stream, + self.cp_comm_type, + ) self.kept_packed_seq_params.discard("cp_group") self.kept_packed_seq_params.discard("local_cp_size") diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 5da94b802ce..a70846da18f 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -62,11 +62,11 @@ class ModelParallelConfig: each rank when using hybrid_context_parallel. """ - hybrid_context_parallel: bool = False + dynamic_context_parallel: bool = False """ - If true, enables hybrid context parallel. This is used to balance the workload of + If true, enables dynamic context parallel. This is used to balance the workload of each CP rank when we use packed samples with variable sequence lengths. - Please set max_seqlen_per_dp_cp_rank when using hybrid_context_parallel. + Please set max_seqlen_per_dp_cp_rank when using dynamic_context_parallel. """ expert_model_parallel_size: int = 1 diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 47a1a59eea7..34f1bc7f476 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -113,8 +113,8 @@ _CONTEXT_PARALLEL_GLOBAL_RANKS = None # Hierarchical context parallel groups _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS = None -# Hybrid context parallel groups -_HYBRID_DP_CP_GROUPS = {} +# Dynamic context parallel groups +_DYNAMIC_DP_CP_GROUPS = {} # Data parallel group information with context parallel combined. _DATA_PARALLEL_GROUP_WITH_CP = None @@ -420,29 +420,29 @@ def create_hierarchical_groups( return hierarchical_groups, hierarchical_groups_gloo -def create_hybrid_dp_cp_groups(rank, ranks, pg_options): +def create_dynamic_dp_cp_groups(rank, ranks, pg_options): """ - Creates groups required for hybrid DPxCP. + Creates groups required for dynamic DPxCP. Creates a new group for every power of 2 up to the number of DPxCP ranks. Returns a dictionary indexed by group size. """ - hybrid_dp_cp_groups = {} + dynamic_dp_cp_groups = {} # Generate group for every power of 2 up to the number of CP ranks # We limit the allowed group sizes in order to avoid excessive overhead. - group_sizes = [2**i for i in range(int(log2(len(ranks))))][1:] + group_sizes = [2**i for i in range(int(log2(len(ranks))))] for group_size in group_sizes: for i in range(0, len(ranks), group_size): group = create_group( ranks[i : i + group_size], pg_options=pg_options, - group_desc=f"HYBRID_DP_CP_GROUP_{group_size}", + group_desc=f"DYNAMIC_DP_CP_GROUP_{group_size}", ) if rank in ranks[i : i + group_size]: assert ( - group_size not in hybrid_dp_cp_groups - ), f"Rank {rank} appears in multiple Hybrid DP CP groups of size {group_size}" - hybrid_dp_cp_groups[group_size] = group - return hybrid_dp_cp_groups + group_size not in dynamic_dp_cp_groups + ), f"Rank {rank} appears in multiple Dynamic DP CP groups of size {group_size}" + dynamic_dp_cp_groups[group_size] = group + return dynamic_dp_cp_groups class RankGenerator(object): @@ -554,7 +554,7 @@ def initialize_model_parallel( use_sharp: bool = False, context_parallel_size: int = 1, hierarchical_context_parallel_sizes: Optional[List[int]] = None, - hybrid_context_parallel: bool = False, + dynamic_context_parallel: bool = False, expert_model_parallel_size: int = 1, num_distributed_optimizer_instances: int = 1, expert_tensor_parallel_size: Optional[int] = None, @@ -940,18 +940,34 @@ def initialize_model_parallel( if "NCCL_COLLNET_ENABLE" in os.environ: del os.environ["NCCL_COLLNET_ENABLE"] - if hybrid_context_parallel: - global _HYBRID_DP_CP_GROUPS + if dynamic_context_parallel: + # TODO: Are gloo groups needed for Dynamic CP? + global _DYNAMIC_DP_CP_GROUPS for ranks_with_cp in decoder_rank_generator.get_ranks('dp-cp'): assert ( len(ranks_with_cp) % 2 == 0 - ), "Hybrid context parallel requires an even number of ranks" - _HYBRID_DP_CP_GROUPS.update( - create_hybrid_dp_cp_groups( + ), "Dynamic context parallel requires an even number of ranks" + _DYNAMIC_DP_CP_GROUPS.update( + create_dynamic_dp_cp_groups( rank, ranks_with_cp, get_nccl_options("dp_cp", nccl_comm_cfgs) ) ) - # TODO: Are gloo groups needed for hybrid cp? + + # 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( + 0, 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: + group = get_dynamic_data_context_parallel_groups(group_size=group_size) + torch.distributed.barrier(group=group, device_ids=[torch.cuda.current_device()]) + torch.cuda.synchronize() + for ranks in decoder_rank_generator.get_ranks('dp'): group = create_group( @@ -1475,16 +1491,16 @@ def get_hierarchical_context_parallel_groups(check_initialized=True): return _HIERARCHICAL_CONTEXT_PARALLEL_GROUPS -def get_hybrid_data_context_parallel_groups(check_initialized=True, group_size=None): - """Get the hybrid context parallel groups the caller rank belongs to.""" +def get_dynamic_data_context_parallel_groups(check_initialized=True, group_size=None): + """Get the dynamic context parallel groups the caller rank belongs to.""" # If the group size is the same as the entire DPxCP group, return the original group if get_data_parallel_world_size(with_context_parallel=True) == group_size: if check_initialized: assert _DATA_PARALLEL_GROUP_WITH_CP is not None return _DATA_PARALLEL_GROUP_WITH_CP if check_initialized: - assert _HYBRID_DP_CP_GROUPS is not None - return _HYBRID_DP_CP_GROUPS[group_size] + assert _DYNAMIC_DP_CP_GROUPS is not None + return _DYNAMIC_DP_CP_GROUPS[group_size] def get_embedding_group(check_initialized=True): diff --git a/megatron/core/pipeline_parallel/hybrid_cp_schedule.py b/megatron/core/pipeline_parallel/dynamic_cp_schedule.py similarity index 99% rename from megatron/core/pipeline_parallel/hybrid_cp_schedule.py rename to megatron/core/pipeline_parallel/dynamic_cp_schedule.py index 27b5fc87945..48dd633aeba 100644 --- a/megatron/core/pipeline_parallel/hybrid_cp_schedule.py +++ b/megatron/core/pipeline_parallel/dynamic_cp_schedule.py @@ -48,7 +48,7 @@ def gpus_needed(self, seq_len: int) -> int: This is used to determine the CP size of a sub-sample. The number is rounded up to the next power of 2 to match the available - hybrid context parallel process group sizes. + dynamic context parallel process group sizes. """ return max(1, 2 ** ceil(log2((seq_len / self.max_seq_len_per_rank)))) @@ -370,7 +370,7 @@ def fill_empty_gpus( "try to increase 'max-seqlen-per-cp-rank'." min_group_size = min(existing_group_sizes) - # We have Hybrid DPxCP groups for every power of 2 of GPUs or the entire DPxCP group. + # We have Dynamic DPxCP groups for every power of 2 of GPUs or the entire DPxCP group. next_power = min(min_group_size * 2, total_gpus) # Find the first group of min_group_size that can be expanded @@ -474,7 +474,7 @@ def get_groups_and_subsamples(self, sample_id_seqlens, config): return groups, sample_id_groups -def hybrid_context_parallel_forward_backward( +def dynamic_context_parallel_forward_backward( forward_step_func, data_iterator, model, @@ -492,7 +492,7 @@ def hybrid_context_parallel_forward_backward( model_type, ): """ - Scheduler for Hybrid Context Parallel. + Scheduler for Dynamic Context Parallel. This function performs the packed sample scheduling and determines 1. The number of microbatches to schedule for each CP rank diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 15c5adfc7a2..62616b8732a 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -36,7 +36,7 @@ combined_1f1b_schedule_for_interleaved_pipelining, combined_1f1b_schedule_for_no_pipelining, ) -from .hybrid_cp_schedule import hybrid_context_parallel_forward_backward +from .dynamic_cp_schedule import dynamic_context_parallel_forward_backward # Types Shape = Union[List[int], torch.Size] @@ -615,8 +615,8 @@ def forward_backward_no_pipelining( total_num_tokens, partial(check_first_val_step, first_val_step, forward_only), ) - elif config.hybrid_context_parallel: - forward_data_store, total_num_tokens = hybrid_context_parallel_forward_backward( + elif config.dynamic_context_parallel: + forward_data_store, total_num_tokens = dynamic_context_parallel_forward_backward( forward_step_func, data_iterator, model, diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 28e3dde01c4..39aa3a79166 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -924,6 +924,11 @@ def forward( (Tuple[Tensor, Tensor]) Attention output and bias. """ + # here we need to set the right cp group for dynamic-cp + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" + self.pg_collection.cp = packed_seq_params.cp_group + # Check if we need to skip RoPE # no_rope is 0-indexed array and self.layer_number is 1-indexed no_rope = ( diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index a9cdc697cc8..4ff59535cb6 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -530,8 +530,8 @@ def get_query_key_value_tensors( if packed_seq_params is not None: assert ( packed_seq_params.local_cp_size is None - ), "hybrid_context_parallel is not supported with MLA yet and is planned for future. \ - Please disable hybrid_context_parallel." + ), "dynamic_context_parallel is not supported with MLA yet and is planned for future. \ + Please disable dynamic_context_parallel." inference_context = deprecate_inference_params(inference_context, inference_params) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 42b82f8d00f..60c5678a3a2 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2168,11 +2168,11 @@ def get_thd_batch_on_this_cp_rank( ################################ -### hybrid context parallel ### +### dynamic context parallel ### ################################ -def get_batch_on_this_hybrid_cp_rank( +def get_batch_on_this_dynamic_cp_rank( batch: Dict[str, Any], local_cp_size: int, cp_group: Optional[torch.distributed.ProcessGroup] = None, @@ -2182,18 +2182,17 @@ def get_batch_on_this_hybrid_cp_rank( """ assert local_cp_size is not None if cp_group is None: - # Get the local cp group required for as defined by the HybridCPDataLoaderWrapper - if local_cp_size > 1: - cp_group = parallel_state.get_hybrid_data_context_parallel_groups( - group_size=local_cp_size - ) + # Get the local cp group required for as defined by the DynamicCPDataLoaderWrapper + cp_group = parallel_state.get_dynamic_data_context_parallel_groups( + group_size=local_cp_size + ) else: # If cp group is provided, it must match the local cp size - # as defined by the HybridCPDataLoaderWrapper + # as defined by the DynamicCPDataLoaderWrapper assert cp_group.size() == local_cp_size # Convert [seqlen] to [1, seqlen] similar to default collate_fn - # as hybrid_context_parallel dataloader wrapper does not go through default collate_fn + # as dynamic_context_parallel dataloader wrapper does not go through default collate_fn for key, data in batch.items(): if key in ['attention_mask']: continue @@ -2213,8 +2212,8 @@ def get_batch_on_this_hybrid_cp_rank( cp_group=cp_group, ) - if cp_group is not None and cp_group.size() > 1: - # When using hybrid_context_parallel, each sub-sample of a packed sample is + if cp_group.size() > 1: + # When using dynamic_context_parallel, each sub-sample of a packed sample is # required to be divisible by CP*DP*2 or CP*DP*TP*2 (if using sequence parallel) batch = get_batch_on_this_cp_rank(batch, cp_group=cp_group) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 91e26af99c6..f9e84636309 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1150,12 +1150,12 @@ def validate_args(args, defaults={}): if args.tp_comm_overlap: assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' - if args.hybrid_context_parallel: - assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' - assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' - assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' - assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' - assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' + if args.dynamic_context_parallel: + assert not args.pipeline_model_parallel_size > 1, 'Dynamic context parallelism not supported with pipeline parallelism' + assert not args.enable_cuda_graph, 'Dynamic context parallelism not supported with CUDA Graph' + assert not args.use_megatron_fsdp, 'Dynamic context parallelism not supported with Megatron FSDP' + assert args.dataloader_type == 'single', 'Dynamic context parallelism only supported with single dataloader type' + assert args.calculate_per_token_loss, 'Dynamic context parallelism must be used with --calculate-per-token-loss' # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index ca4cc1b36a3..166d4597a97 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -39,8 +39,8 @@ def build_pretraining_data_loader(dataset, consumed_samples): data_parallel_size=mpu.get_data_parallel_world_size(), ) elif args.dataloader_type == 'single': - if args.hybrid_context_parallel: - batch_sampler = HybridCPMegatronPretrainingSampler( + if args.dynamic_context_parallel: + batch_sampler = DynamicCPMegatronPretrainingSampler( total_samples=len(dataset), consumed_samples=consumed_samples, micro_batch_size=args.micro_batch_size, @@ -79,7 +79,7 @@ def worker_init_fn(_): worker_init_fn if args.exit_signal_handler and args.num_workers > 0 else None ) # Torch dataloader. - if args.hybrid_context_parallel: + if args.dynamic_context_parallel: extra_kwargs = {"collate_fn": lambda x: x,} else: extra_kwargs = {} @@ -162,11 +162,11 @@ def __iter__(self): start_idx, end_idx = self.get_start_end_idx() yield batch[start_idx:end_idx] -class HybridCPMegatronPretrainingSampler(MegatronPretrainingSampler): +class DynamicCPMegatronPretrainingSampler(MegatronPretrainingSampler): """ - Data sampler for hybrid context parallel (Hybrid CP) format. + Data sampler for dynamic context parallel (Dynamic CP) format. This data sampler pulls in the entire global batch at once across all data parallel ranks. - This helps provide the Hybrid CP Dataloader Wrapper to schedule and load balance sub-samples + This helps provide the Dynamic CP Dataloader Wrapper to schedule and load balance sub-samples of the entire global batch. """ diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index c150ac3d5ca..b364cc77c8d 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -369,7 +369,7 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s use_sharp=args.use_sharp, context_parallel_size=args.context_parallel_size, hierarchical_context_parallel_sizes=args.hierarchical_context_parallel_sizes, - hybrid_context_parallel=args.hybrid_context_parallel, + dynamic_context_parallel=args.dynamic_context_parallel, expert_model_parallel_size=args.expert_model_parallel_size, num_distributed_optimizer_instances=args.num_distributed_optimizer_instances, expert_tensor_parallel_size=args.expert_tensor_parallel_size, diff --git a/megatron/training/training.py b/megatron/training/training.py index b508da02ef1..74e7799c71b 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -139,7 +139,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.training.initialize import set_jit_fusion_options from megatron.training.utils import get_batch_on_this_cp_rank, get_batch_on_this_tp_rank, is_hybrid_model from megatron.training.datasets.data_samplers import build_pretraining_data_loader -from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper +from megatron.core.datasets.data_schedule import DynamicCPDataLoaderWrapper from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler from megatron.core.transformer.moe import upcycling_utils from megatron.core.transformer.moe.moe_utils import track_moe_metrics, clear_aux_losses_tracker @@ -2568,8 +2568,8 @@ def train( energy_monitor = get_energy_monitor() one_logger = get_one_logger() - if args.hybrid_context_parallel: - train_data_iterator = iter(HybridCPDataLoaderWrapper(train_data_iterator, config)) + if args.dynamic_context_parallel: + train_data_iterator = iter(DynamicCPDataLoaderWrapper(train_data_iterator, config)) if args.run_workload_inspector_server: try: diff --git a/megatron/training/utils.py b/megatron/training/utils.py index f0a5cac3176..98d6614c30d 100644 --- a/megatron/training/utils.py +++ b/megatron/training/utils.py @@ -577,7 +577,7 @@ def _broadcast_cu_seqlens(cu_seqlens): buf = cu_seqlens.to(device=dev, non_blocking=True).contiguous() _broadcast(buf) - if args.hybrid_context_parallel: + if args.dynamic_context_parallel: seq_len = torch.tensor(batch['tokens'].shape[0], dtype=torch.int32, device=torch.cuda.current_device()) _broadcast(seq_len) @@ -607,7 +607,7 @@ def _broadcast_cu_seqlens(cu_seqlens): _broadcast(batch['attention_mask']) else: - if args.hybrid_context_parallel: + if args.dynamic_context_parallel: seq_len = torch.tensor(0, dtype=torch.int32, device=torch.cuda.current_device()) _broadcast(seq_len) shape = (seq_len.item()) @@ -630,7 +630,7 @@ def _broadcast_cu_seqlens(cu_seqlens): device=torch.cuda.current_device(), ) if args.create_attention_mask_in_dataloader: - shape_attention_mask = (args.micro_batch_size, 1, args.seq_length, args.seq_length) if not args.hybrid_context_parallel else (1, 1, shape[0], shape[0]) + shape_attention_mask = (args.micro_batch_size, 1, args.seq_length, args.seq_length) if not args.dynamic_context_parallel else (1, 1, shape[0], shape[0]) attention_mask = torch.empty( shape_attention_mask, dtype=torch.bool, @@ -644,7 +644,7 @@ def _broadcast_cu_seqlens(cu_seqlens): device=torch.cuda.current_device(), ) cu_seqlens = None - if args.hybrid_context_parallel or args.sft: + if args.dynamic_context_parallel or args.sft: max_seqlen = torch.empty( 1, dtype=torch.int32, @@ -657,7 +657,7 @@ def _broadcast_cu_seqlens(cu_seqlens): 1, dtype=torch.int32, device=torch.cuda.current_device(), - ) if args.hybrid_context_parallel else None + ) if args.dynamic_context_parallel else None def _broadcast_cu_seqlens(): dev = torch.cuda.current_device() diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 81768944623..cd27d195435 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -29,7 +29,7 @@ from megatron.core.models.gpt import GPTModel from megatron.core.rerun_state_machine import get_rerun_state_machine from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer -from megatron.core.utils import get_attr_wrapped_model, get_thd_batch_on_this_cp_rank, get_batch_on_this_hybrid_cp_rank, StragglerDetector +from megatron.core.utils import get_attr_wrapped_model, get_thd_batch_on_this_cp_rank, get_batch_on_this_dynamic_cp_rank, StragglerDetector from megatron.training import ( get_args, get_timers, @@ -90,8 +90,8 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): elif local_cp_size is None: # Packed THD format assert max_seqlen.dim() == 1 batch, packed_seq_params = get_thd_batch_on_this_cp_rank(batch, cu_seqlens, cu_seqlens_padded, max_seqlen) - else: # Hybrid CP format - batch, packed_seq_params = get_batch_on_this_hybrid_cp_rank(batch, local_cp_size) + else: # Dynamic CP format + batch, packed_seq_params = get_batch_on_this_dynamic_cp_rank(batch, local_cp_size) return (*batch.values(), packed_seq_params) @@ -248,7 +248,7 @@ def core_gpt_dataset_config_from_args(args): "context_parallel_size": args.context_parallel_size, "data_parallel_size": args.data_parallel_size, "sequence_parallel_size": args.tensor_model_parallel_size*args.sequence_parallel, - "hybrid_context_parallel": args.hybrid_context_parallel, + "dynamic_context_parallel": args.dynamic_context_parallel, } # add FIM args to the config diff --git a/pretrain_mamba.py b/pretrain_mamba.py index 0fecbef2c71..037f1817d99 100644 --- a/pretrain_mamba.py +++ b/pretrain_mamba.py @@ -94,7 +94,7 @@ def get_batch(data_iterator, vp_stage=None): cu_seqlens = batch['cu_seqlens'] # Unused at the moment cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) - # Support for Hybrid Context Parallel (Unused in this script) + # Support for Dynamic Context Parallel (Unused in this script) local_cp_size = batch.pop('local_cp_size', None) if cu_seqlens is not None: diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py index f8d1cde7028..b1b699e1c28 100644 --- a/tests/unit_tests/models/test_mamba_moe_model.py +++ b/tests/unit_tests/models/test_mamba_moe_model.py @@ -279,7 +279,7 @@ "fine_grained_activation_offloading": False, "min_offloaded_tensor_size": 1024 * 1024, "offload_modules": [], - "hybrid_context_parallel": False, + "dynamic_context_parallel": False, "max_seqlen_per_dp_cp_rank": None, "inference_disable_torch_grouped_mm": False, "inference_disable_triton_nvls_kernels": False, diff --git a/tests/unit_tests/test_parallel_state.py b/tests/unit_tests/test_parallel_state.py index 21dc740cdf4..19f55525fb1 100644 --- a/tests/unit_tests/test_parallel_state.py +++ b/tests/unit_tests/test_parallel_state.py @@ -507,9 +507,9 @@ def golden_rank_result_from_past_code( "world_size, tp_size, cp_size, dp_size", [(8, 1, 2, 4), (8, 1, 1, 8)], # 8 GPUs, 1 TP, 2 CP, 4 DP # 8 GPUs, 1 TP, 1 CP, 8 DP ) -def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size): +def test_dynamic_dp_cp_groups(world_size, tp_size, cp_size, dp_size): """ - Test that hybrid DPxCP groups are created correctly. + Test that dynamic DPxCP groups are created correctly. """ Utils.destroy_model_parallel() @@ -520,13 +520,13 @@ def test_hybrid_dp_cp_groups(world_size, tp_size, cp_size, dp_size): Utils.initialize_model_parallel( tensor_model_parallel_size=tp_size, context_parallel_size=cp_size, - hybrid_context_parallel=True, + dynamic_context_parallel=True, ) dp_cp_size = ps.get_data_parallel_world_size(with_context_parallel=True) group_sizes = [2**i for i in range(int(log2(dp_cp_size)))][1:] for group_size in group_sizes: - group = ps.get_hybrid_data_context_parallel_groups(group_size=group_size) + group = ps.get_dynamic_data_context_parallel_groups(group_size=group_size) assert group.size() == group_size Utils.destroy_model_parallel()