Merge branch 'fix-ep-tp' into 'main'

Fix Issue with EP and TP

See merge request ADLR/megatron-lm!3699

Author: ko3n1g

megatron/core/inference/contexts/dynamic_context.py

@@ -155,7 +155,6 @@ def __init__(
             tp_size = tensor_model_parallel_size
         hidden_size_per_attention_head = core_divide(projection_size, num_attention_heads)
         num_attention_heads_per_partition = core_divide(num_attention_heads, tp_size)
-
         # Chunk size tokens, bytes.
         dtype_size_bytes = params_dtype.itemsize
         self.chunk_size_tokens = chunk_size_tokens
@@ -177,23 +176,24 @@ def __init__(
         def bytes_to_max_requests_and_tokens(n_bytes):
             n_tokens = n_bytes / self.chunk_size_bytes * self.chunk_size_tokens
             n_requests = n_tokens / max_sequence_length
-            return int(n_requests), int(n_tokens)
+            return self.round_up_requests(int(n_requests), tp_size=tp_size), self.round_up_tokens(
+                int(n_tokens), tp_size=tp_size
+            )
 
         self.max_requests, self.max_tokens = bytes_to_max_requests_and_tokens(buffer_size_bytes)
-
         if buffer_overflow_factor is not None:
             self.max_requests = self.round_up_requests(
-                int(self.max_requests * buffer_overflow_factor)
+                int(self.max_requests * buffer_overflow_factor), tp_size=tp_size
             )
             self.max_tokens = self.round_up_tokens(
-                int(self.max_tokens * buffer_overflow_factor / 50.0)
+                int(self.max_tokens * buffer_overflow_factor / 50.0), tp_size=tp_size
             )
 
         if max_requests_override is not None:
-            self.max_requests = self.round_up_requests(max_requests_override)
+            self.max_requests = self.round_up_requests(max_requests_override, tp_size=tp_size)
 
         if max_tokens_override is not None:
-            self.max_tokens = self.round_up_tokens(max_tokens_override)
+            self.max_tokens = self.round_up_tokens(max_tokens_override, tp_size=tp_size)
 
         self.max_requests = min(self.max_requests, self.max_tokens)  # e.g., decode only.
 
@@ -277,7 +277,8 @@ def bytes_to_max_requests_and_tokens(n_bytes):
             self.cuda_graph_step_size = cuda_graph_rounder * int(
                 math.ceil(int(self.cuda_graph_step_size) / cuda_graph_rounder)
             )
-
+            # Make sure divisble by TP size
+            self.cuda_graph_step_size = math.ceil(self.cuda_graph_step_size / tp_size) * tp_size
             # Cuda graph request counts.
             if num_cuda_graphs == 1:
                 self.cuda_graph_request_counts = [self.max_requests]
@@ -355,26 +356,46 @@ def bytes_to_max_requests_and_tokens(n_bytes):
     REQUEST_ROUNDER = 4
 
     @classmethod
-    def round_up_tokens(cls, value):
-        """Round up to nearest multiple of `TOKEN_ROUNDER` (above)."""
+    def round_up_tokens(cls, value, tp_size=None):
+        """Round up to nearest multiple of `TOKEN_ROUNDER` (above) that is also divisible by tensor model parallel size."""
         if not HAVE_PACKAGING:
             raise ImportError(
                 "`packaging` is required for this functionality, please install it with `pip install packaging`"
             )
         if PkgVersion(mcore_version) < PkgVersion("0.13"):
             return cls.round_up(value)
-        return cls.TOKEN_ROUNDER * int(math.ceil(int(value) / cls.TOKEN_ROUNDER))
+
+        # Make sure divisible by TP size
+        if tp_size is None:
+            # Check if parallel state is initialized before trying to get TP size
+            if parallel_state.is_initialized():
+                tp_size = parallel_state.get_tensor_model_parallel_world_size()
+            else:
+                tp_size = 1
+        token_rounder = math.ceil(cls.TOKEN_ROUNDER / tp_size) * tp_size
+
+        return token_rounder * int(math.ceil(int(value) / token_rounder))
 
     @classmethod
-    def round_up_requests(cls, value):
-        """Round up to nearest multiple of `REQUEST_ROUNDER` (above)."""
+    def round_up_requests(cls, value, tp_size=None):
+        """Round up to nearest multiple of `REQUEST_ROUNDER` (above) that is also divisible by tensor model parallel size."""
         if not HAVE_PACKAGING:
             raise ImportError(
                 "`packaging` is required for this functionality, please install it with `pip install packaging`"
             )
         if PkgVersion(mcore_version) < PkgVersion("0.13"):
             return cls.round_up(value)
-        return cls.REQUEST_ROUNDER * int(math.ceil(int(value) / cls.REQUEST_ROUNDER))
+
+        # Make sure divisible by TP size
+        if tp_size is None:
+            # Check if parallel state is initialized before trying to get TP size
+            if parallel_state.is_initialized():
+                tp_size = parallel_state.get_tensor_model_parallel_world_size()
+            else:
+                tp_size = 1
+        request_rounder = math.ceil(cls.REQUEST_ROUNDER / tp_size) * tp_size
+
+        return request_rounder * int(math.ceil(int(value) / request_rounder))
 
     @classmethod
     def round_up(cls, value):

megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -26,6 +26,8 @@
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.utils import get_attention_mask
 from megatron.core.transformer.cuda_graphs import create_cudagraphs
+from megatron.core.transformer.moe.moe_layer import BaseMoELayer
+from megatron.core.transformer.utils import set_model_to_sequence_parallel
 from megatron.core.utils import get_model_config
 
 try:
@@ -429,9 +431,11 @@ def generate_output_tokens_dynamic_batch(
         # Get flat tokens, position ids.
         input_ids, position_ids = context.current_input_and_position_ids()
 
+        model_config = get_model_config(self.inference_wrapped_model.model)
+
         # If using symmetric kernels and we are using using nccl
         # for prefill turn off symmetric kernels
-        symmetric_ar_type = get_model_config(self.inference_wrapped_model.model).symmetric_ar_type
+        symmetric_ar_type = model_config.symmetric_ar_type
         nccl_all_reduce_for_prefill = (
             self.inference_wrapped_model.inference_wrapper_config.nccl_all_reduce_for_prefill
         )
@@ -681,6 +685,21 @@ def generate_all_output_tokens_static_batch(
                 not self.inference_wrapped_model.inference_context.is_decode_only()
             ), f"Generation must start in prefill mode"
 
+            # Sequence parallelism is required for MoE layers when using expert parallelism (EP)
+            # becausethe expert routing mechanism relies on sequence parallelism's communication
+            # infrastructure to distribute tokens across expert ranks. However, sequence parallelism
+            # is not currently supported for non-MoE layers during inference,so we selectively
+            # disable it for all other layer types. This is safe because MoE layers perform an
+            # all-gather operation on sequences before passing data to subsequent layers, ensuring
+            # that each rank has the complete sequence data needed for the next non-MoE layer.
+            tp_size = model_config.tensor_model_parallel_size
+            ep_size = model_config.expert_model_parallel_size
+            model_is_tp_ep = tp_size > 1 and ep_size > 1
+            if model_is_tp_ep:
+                set_model_to_sequence_parallel(
+                    self.inference_wrapped_model.model.module, False, exclude_modules=[BaseMoELayer]
+                )
+
             # If using symmetric kernels and we are using using nccl
             # for prefill turn off symmetric kernels
             symmetric_ar_type = model_config.symmetric_ar_type

megatron/core/transformer/utils.py

@@ -199,10 +199,14 @@ def sharded_state_dict_default(
 _sequence_parallel_attr_cache = None
 
 
-def _init_sequence_parallel_cache(model):
+def _init_sequence_parallel_cache(model, exclude_modules):
     """
     Initialize the cache of modules with sequence parallel attributes.
     Only needs to be called once, subsequent calls have no effect.
+
+    Args:
+        model: model to change sequence parallelism attributes
+        exclude_modules: Modules to exclude from changing sequence parallelism
     """
     global _sequence_parallel_attr_cache
     model_id = id(model)
@@ -229,33 +233,35 @@ def _init_sequence_parallel_cache(model):
 
     # Recursive function to find all modules with our target attributes
     def find_modules_with_attrs(module):
-        # Check if this module has any of our target attributes
-        for attr in sequence_parallel_attrs:
-            if hasattr(module, attr):
-                _sequence_parallel_attr_cache[model_id][attr].append(module)
+        if exclude_modules is None or module not in exclude_modules:
+            # Check if this module has any of our target attributes
+            for attr in sequence_parallel_attrs:
+                if hasattr(module, attr):
+                    _sequence_parallel_attr_cache[model_id][attr].append(module)
 
-        # Check all children modules recursively
-        for child in module._modules.values():
-            if child is not None:
-                find_modules_with_attrs(child)
+            # Check all children modules recursively
+            for child in module._modules.values():
+                if child is not None:
+                    find_modules_with_attrs(child)
 
     # Start the search from each major component
     find_modules_with_attrs(model_modules)
 
 
-def set_model_to_sequence_parallel(model, set_to=False):
+def set_model_to_sequence_parallel(model, set_to=False, exclude_modules=None):
     """
     Set sequence parallel attributes for the model.
 
     Args:
         set_to: Value to set for sequence_parallel attributes
+        exclude_modules: Modules to exclude from changing sequence parallelism
     """
     global _sequence_parallel_attr_cache
     model_id = id(model)
 
     # Initialize cache if needed
     if _sequence_parallel_attr_cache is None or model_id not in _sequence_parallel_attr_cache:
-        _init_sequence_parallel_cache(model)
+        _init_sequence_parallel_cache(model, exclude_modules)
 
     model.config.sequence_parallel = set_to
 

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies.", "generated_text": " Wait for the moment when the music stops, and the lights come up, and the DJ says, \"I'm going to play a song for you", "generated_tokens": [32844, 1394, 1278, 4735, 2200, 1278, 7146, 30774, 1044, 1321, 1278, 26466, 3930, 2015, 1044, 1321, 1278, 30245, 8223, 1044, 1429, 1073, 4525, 4670, 1317, 3354, 1261, 6947, 1394, 1636], "latency": 19.686351776123047, "logprobs": [-5.292269229888916, -7.716421127319336, -9.068008422851562, -12.118106842041016, -3.741875648498535, -1.8551081418991089, -1.8765699863433838, -9.52701187133789, -15.140599250793457, -9.51123046875, -10.250877380371094, -8.108641624450684, -9.099360466003418, -9.529533386230469, -10.495244979858398, -9.094446182250977, -9.802777290344238, -8.999850273132324, -9.103967666625977, -9.895696640014648, -8.00230884552002, -7.3570966720581055, -7.892587184906006, -12.32270622253418, -20.62922477722168, -9.672601699829102, -8.485877990722656, -10.270708084106445, -11.473578453063965, -15.617767333984375, -7.8881988525390625, -12.872822761535645, -8.940616607666016, -7.3508501052856445, -10.157344818115234, -12.235904693603516, -9.32239818572998, -6.516319751739502, -8.389573097229004, -8.860508918762207, -16.462238311767578, -7.349939346313477, -11.075740814208984, -14.01966667175293, -9.536352157592773, -9.535323143005371, -11.839295387268066, -12.564751625061035, -9.356565475463867, -9.240681648254395, -9.669130325317383, -8.965482711791992, -11.053977012634277, -14.045623779296875, -13.299668312072754, -14.695284843444824, -13.231292724609375, -9.543293952941895, -11.672986030578613, -10.587867736816406, -8.400468826293945, -10.324536323547363, -13.930037498474121, -15.27256965637207, -10.176668167114258, -13.777766227722168, -8.423280715942383, -7.511598110198975, -14.0129976272583, -5.561246871948242, -9.51725959777832, -10.10839557647705, -8.918962478637695, -8.14908218383789, -7.653857707977295, -11.743547439575195, -10.011963844299316, -12.899750709533691, -15.401609420776367, -6.838616847991943, -9.010682106018066, -10.37846565246582, -6.819251537322998, -13.074575424194336, -10.851410865783691, -8.874515533447266, -10.204574584960938, -16.298084259033203, -13.584976196289062, -10.295950889587402, -8.796205520629883, -12.162117004394531, -9.572405815124512, -8.92280101776123, -10.94050407409668, -15.27184772491455, -13.962615966796875, -9.328908920288086, -9.781393051147461, -12.07744026184082, -11.402749061584473, -11.740723609924316, -17.354206085205078, -9.84351634979248, -9.201858520507812, -8.702098846435547, -12.5997314453125, -14.244935989379883, -14.273555755615234, -16.253263473510742, -13.604464530944824, -11.363554000854492, -9.675899505615234, -12.930312156677246, -10.388641357421875, -11.593982696533203, -10.904473304748535]}}

tests/functional_tests/test_cases/moe/gpt_dynamic_inference_tp4_pp1_ep4_16B_logitsmatch/model_config.yaml

@@ -0,0 +1,79 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+TEST_TYPE: frozen-start
+MODE: inference
+MODEL_ARGS:
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --log-memory-to-tensorboard: true
+  --timing-log-level: 2
+  --load: ${CHECKPOINT_LOAD_PATH}/deepseek_16b_pyt/model/checkpoints
+  --tokenizer-model: ${DATA_PATH}/deepseek_16b_pyt/tokenizer/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json
+  --tokenizer-type: TikTokenizer
+  --tiktoken-pattern: v2
+  --distributed-backend: nccl
+  --log-interval: 1
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 4
+  --pipeline-model-parallel-size: 1
+  --expert-model-parallel-size: 4
+  --expert-tensor-parallel-size: 1
+  --sequence-parallel: true
+  --use-mcore-models: true
+  --moe-token-dispatcher-type: alltoall
+  --moe-grouped-gemm: true
+  --num-experts: 64
+  --moe-router-topk: 6
+  --moe-z-loss-coeff: 0
+  --moe-router-load-balancing-type: seq_aux_loss
+  --moe-aux-loss-coeff: 1e-3
+  --moe-router-score-function: sigmoid
+  --untie-embeddings-and-output-weights: true
+  --disable-bias-linear: true
+  --init-method-std: 0.014
+  --position-embedding-type: rope
+  --rotary-base: 1000000
+  --rotary-percent: 1.0
+  --num-layers: 27
+  --hidden-size: 2048
+  --moe-ffn-hidden-size: 1408
+  --moe-shared-expert-intermediate-size: 2816
+  --ffn-hidden-size: 10944
+  --num-attention-heads: 16
+  --kv-channels: 128
+  --normalization: RMSNorm
+  --swiglu: true
+  --attention-dropout: 0.0
+  --hidden-dropout: 0.0
+  --seq-length: 4096
+  --max-position-embeddings: 4096
+  --micro-batch-size: 1
+  --ckpt-format: torch_dist
+  --ckpt-fully-parallel-save: true
+  --ckpt-fully-parallel-load: true
+  --ckpt-assume-constant-structure: true
+  --dist-ckpt-strictness: log_unexpected
+  --bf16: true
+  --attention-backend: flash
+  --no-create-attention-mask-in-dataloader: true
+  --num-workers: 8
+  --use-checkpoint-args: true
+  --no-use-tokenizer-model-from-checkpoint-args: true
+  --no-load-optim: true
+  --deterministic-mode: true
+  --save-interval: 2000
+  --temperature: 1.0
+  --top_k: 1
+  --return-log-probs: true
+  --num-tokens-to-generate: 30
+  --max-tokens-to-oom: 3600000
+  --inference-max-seq-length: 4096
+  --output-path: ${TENSORBOARD_PATH}
+  --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
+METRICS:
+  - "generated_tokens"
+  - "logprobs"

tests/functional_tests/test_cases/moe/gpt_static_inference_tp4_pp1_ep4_16B_logitsmatch/golden_values_dev_dgx_h100.json

@@ -0,0 +1 @@
+{"0": {"input_prompt": "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies.", "generated_text": " Wait for the moment when the music stops, and the lights come up, and the DJ says, \"I'm going to play a song for you", "generated_tokens": [32844, 1394, 1278, 4735, 2200, 1278, 7146, 30774, 1044, 1321, 1278, 26466, 3930, 2015, 1044, 1321, 1278, 30245, 8223, 1044, 1429, 1073, 4525, 4670, 1317, 3354, 1261, 6947, 1394, 1636], "tpot": [11.527735710144043, 0.591648280620575, 0.07851152122020721, 0.08397766202688217, 0.07938633859157562, 0.07749997079372406, 0.078935906291008, 0.08028851449489594, 0.0792686715722084, 0.07622275501489639, 0.07548975944519043, 0.07680464535951614, 0.07543014734983444, 0.07560738921165466, 0.07399769872426987, 0.07543785870075226, 0.07579366117715836, 0.0751631036400795, 0.07548335939645767, 0.07564015686511993, 0.07764911651611328, 0.07656403630971909, 0.07500188797712326, 0.07636498659849167, 0.07543619722127914, 0.07694652676582336, 0.07432099431753159, 0.0751761943101883, 0.07691458612680435, 0.07628953456878662], "latency": 14.289155829232186, "logprobs": [-10.448518753051758, -3.693941593170166, -2.833103656768799, -1.2445695400238037, -0.23799529671669006, -1.7522815465927124, -2.378152370452881, -1.9484899044036865, -2.108924388885498, -6.127920150756836, -0.8197959661483765, -2.477976083755493, -3.492497444152832, -4.170319557189941, -1.9918553829193115, -1.8618279695510864, -2.2335567474365234, -7.071791172027588, -0.039936937391757965, -1.9948835372924805, -5.008172512054443, -8.708097457885742, -9.903486251831055, -0.851460337638855, -4.765171051025391, -0.8707393407821655, -2.219733238220215, -0.01853257417678833, -0.035978663712739944, -3.387631416320801, -8.754067420959473, -1.2686023712158203, -6.662981986999512, -3.7872395515441895, -3.6667354106903076, -4.171259880065918, -2.2128500938415527, -1.091404914855957, -0.22139909863471985, -0.8265669941902161, -4.746159553527832, -9.04170036315918, -0.013459297828376293, -3.17301607131958, -1.3139652013778687, -3.9821701049804688, -0.7707944512367249, -0.002040567807853222, -2.9162371158599854, -10.677328109741211, -3.1504364013671875, -1.1485933065414429, -4.871399402618408, -0.20786719024181366, -0.06325722485780716, -1.3587590456008911, -2.207646369934082, -4.407937049865723, -0.36253970861434937, -4.0189995765686035, -0.3988611698150635, -0.13855230808258057, -2.7199528217315674, -10.558171272277832, -0.04671315476298332, -3.5006980895996094, -0.9756439328193665, -4.673828125, -0.2634696066379547, -2.5747756958007812, -0.8531911969184875, -1.6041897535324097, -5.738401412963867, -16.978456497192383, -2.6206722259521484, -0.14098073542118073, -7.450814247131348, -1.076573371887207, -2.129807472229004, -1.5724716186523438, -0.29326727986335754, -5.609436511993408, -0.0065282415598630905, -7.79502010345459, -2.715085744857788, -3.0889575481414795, -3.0355961322784424, -2.4395439624786377, -0.3983170986175537, -1.5089631080627441, -2.276723861694336, -0.6004312038421631, -1.3054823875427246, -1.9454480409622192, -1.7226327657699585, -0.7742734551429749, -0.49186939001083374, -1.2962923049926758, -1.567298173904419, -1.0149078369140625, -0.40288272500038147, -0.4789682626724243, -0.04533138871192932, -1.2695876359939575, -2.223480224609375, -2.6703481674194336, -0.7677091956138611, -0.42749911546707153, -2.8563802242279053, -1.5350499153137207, -1.6456167697906494, -0.05149398744106293, -1.3739523887634277, -1.3543274402618408, -1.2655469179153442, -1.307403326034546, -0.497008740901947]}}