diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 9a1a11a8c2b..93e30f6aa25 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -74,6 +74,7 @@ def __init__( data_parallel_size: int, tokenizer, inference_coordinator_port: int | None = None, + deterministic_mode: bool = False, ): """ Initializes the inference coordinator. @@ -145,6 +146,12 @@ def __init__( assert identity not in self.identities_of_data_parallel_ranks self.identities_of_data_parallel_ranks.append(identity) logging.info("Inference Coordinator: Connected with data parallel ranks...") + + # In deterministic mode, sort identities for consistent scheduling order. + if deterministic_mode: + self.identities_of_data_parallel_ranks = deque( + sorted(self.identities_of_data_parallel_ranks) + ) self.data_parallel_rank_iterator = cycle(self.identities_of_data_parallel_ranks) self.data_parallel_pause_acks = set() self.data_parallel_stop_acks = set() @@ -343,6 +350,7 @@ def entrypoint( data_parallel_size: int, tokenizer, inference_coordinator_port: int | None = None, + deterministic_mode: bool = False, ): """ Class method to instantiate and run the coordinator, for use in a separate process. @@ -356,9 +364,14 @@ def entrypoint( once the coordinator is ready to accept connections. inference_coordinator_port (int): The port to bind to. data_parallel_size (int): The number of expected TP-coordinators. + deterministic_mode (bool): Whether to enable deterministic scheduling. """ coordinator = cls( - pipe_connection, data_parallel_size, tokenizer, inference_coordinator_port + pipe_connection, + data_parallel_size, + tokenizer, + inference_coordinator_port, + deterministic_mode=deterministic_mode, ) ready_event.set() try: diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 0a95e8f4a53..a09d15ae20b 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -424,6 +424,7 @@ async def start_listening_to_data_parallel_coordinator( # Spawn a DP coordinator process and get the connection info. if launch_inference_coordinator and self.is_dp_coordinator: spawn_context = multiprocessing.get_context('spawn') + deterministic_mode = torch.are_deterministic_algorithms_enabled() dp_pipe, dp_process_pipe = spawn_context.Pipe() coordinator_ready_event = spawn_context.Event() self.inference_coordinator_process = spawn_context.Process( @@ -434,6 +435,7 @@ async def start_listening_to_data_parallel_coordinator( get_pg_size(self.pg_collection.dp), self.controller.tokenizer, inference_coordinator_port, + deterministic_mode, ), ) self.inference_coordinator_process.start() diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 364a80db81e..3058db78f41 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -518,6 +518,10 @@ def get_environment_rollouts( rollouts = [ loop.run_until_complete(anext(rollout_generator)) for _ in range(n_prompts) ] + # In deterministic mode, sort rollouts by problem_id for consistent ordering + # regardless of completion order due to system timing jitter. + if torch.are_deterministic_algorithms_enabled(): + rollouts.sort(key=lambda group: group[0].problem_id if group and group[0].problem_id else "") if not args.rl_partial_rollouts: while True: try: diff --git a/tests/functional_tests/python_test_utils/test_grpo_training_loop.py b/tests/functional_tests/python_test_utils/test_grpo_training_loop.py index 6faca9b11b3..b4447f5f761 100644 --- a/tests/functional_tests/python_test_utils/test_grpo_training_loop.py +++ b/tests/functional_tests/python_test_utils/test_grpo_training_loop.py @@ -84,9 +84,8 @@ def test_grpo_training_loop( with open(model_config_path, 'r') as f: model_config = yaml.safe_load(f) metrics = model_config["METRICS"] - if "THROUGHPUT_TEST_PARAMS" in model_config: - throughput_test_params = model_config["THROUGHPUT_TEST_PARAMS"] - start_step = throughput_test_params["--start_step"] + if "ENV_VARS" in model_config and "THROUGHPUT_START_STEP" in model_config["ENV_VARS"]: + start_step = model_config["ENV_VARS"]["THROUGHPUT_START_STEP"] else: start_step = 1 diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json index d985f671cab..05bc35e362f 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json @@ -1,173 +1,187 @@ { + "lm loss": { + "start_step": 1, + "end_step": 30, + "step_interval": 1, + "values": { + "1": 0.0, + "2": 0.0, + "3": 0.0, + "4": 0.04559, + "5": 0.0, + "6": 0.0523, + "7": 0.0, + "8": 0.0, + "9": 0.04887, + "10": 0.0, + "11": 0.0, + "12": 0.0, + "13": 0.0, + "14": 0.0, + "15": 0.0, + "16": 0.0, + "17": 0.04299, + "18": 0.0, + "19": 0.03797, + "20": 0.0, + "21": 0.0, + "22": 0.0, + "23": 0.0, + "24": 0.0, + "25": 0.0, + "26": 0.0, + "27": 0.0, + "28": 0.0, + "29": 0.0, + "30": 0.0 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 30, + "step_interval": 1, + "values": { + "1": 583687296.0, + "2": 583687296.0, + "3": 583687296.0, + "4": 31.0, + "5": 583687296.0, + "6": 12.0, + "7": 583687296.0, + "8": 583687296.0, + "9": 16.0, + "10": 583687296.0, + "11": 583687296.0, + "12": 583687296.0, + "13": 583687296.0, + "14": 583687296.0, + "15": 583687296.0, + "16": 583687296.0, + "17": 47.0, + "18": 583687296.0, + "19": 43.0, + "20": 583687296.0, + "21": 583687296.0, + "22": 583687296.0, + "23": 583687296.0, + "24": 583687296.0, + "25": 583687296.0, + "26": 583687296.0, + "27": 583687296.0, + "28": 583687296.0, + "29": 583687296.0, + "30": 583687296.0 + } + }, "mem-allocated-bytes": { "start_step": 1, - "end_step": 50, + "end_step": 30, "step_interval": 1, "values": { - "1": 48967716864.0, - "2": 48973631488.0, - "3": 48974528512.0, - "4": 48971538432.0, - "5": 48974340096.0, - "6": 48974143488.0, - "7": 48977002496.0, - "8": 48975851520.0, - "9": 48974036992.0, - "10": 48973709312.0, - "11": 48973262848.0, - "12": 48973705216.0, - "13": 48973598720.0, - "14": 48976703488.0, - "15": 48975118336.0, - "16": 48977072128.0, - "17": 48976465920.0, - "18": 48976470016.0, - "19": 48976478208.0, - "20": 48976654336.0, - "21": 48976793600.0, - "22": 48976052224.0, - "23": 48976277504.0, - "24": 48974708736.0, - "25": 48973062144.0, - "26": 48976236544.0, - "27": 48975970304.0, - "28": 48976711680.0, - "29": 48975593472.0, - "30": 48977321984.0, - "31": 48977506304.0, - "32": 48976646144.0, - "33": 48976072704.0, - "34": 48973631488.0, - "35": 48976650240.0, - "36": 48975650816.0, - "37": 48974950400.0, - "38": 48972750848.0, - "39": 48976617472.0, - "40": 48979308544.0, - "41": 48978587648.0, - "42": 48975626240.0, - "43": 48975089664.0, - "44": 48973688832.0, - "45": 48975327232.0, - "46": 48975159296.0, - "47": 48975372288.0, - "48": 48973856768.0, - "49": 48973377536.0, - "50": 48975568896.0 + "1": 48985034752.0, + "2": 48991363072.0, + "3": 48993005568.0, + "4": 48991928320.0, + "5": 48992874496.0, + "6": 48991891456.0, + "7": 48991338496.0, + "8": 48993873920.0, + "9": 48993124352.0, + "10": 48994115584.0, + "11": 48994050048.0, + "12": 48993181696.0, + "13": 48993918976.0, + "14": 48992014336.0, + "15": 48992256000.0, + "16": 48989933568.0, + "17": 48992645120.0, + "18": 48992890880.0, + "19": 48992821248.0, + "20": 48992821248.0, + "21": 48991612928.0, + "22": 48993181696.0, + "23": 48992821248.0, + "24": 48992821248.0, + "25": 48993931264.0, + "26": 48992022528.0, + "27": 48993173504.0, + "28": 48992821248.0, + "29": 48993935360.0, + "30": 48994017280.0 } }, "mem-max-allocated-bytes": { "start_step": 1, - "end_step": 50, + "end_step": 30, "step_interval": 1, "values": { - "1": 49090379776.0, - "2": 49937022976.0, - "3": 49938366464.0, - "4": 49938366464.0, - "5": 49938366464.0, - "6": 49938698240.0, - "7": 49939156992.0, - "8": 49939156992.0, - "9": 49939156992.0, - "10": 49939156992.0, - "11": 49939156992.0, - "12": 49939156992.0, - "13": 49939156992.0, - "14": 49940287488.0, - "15": 49940287488.0, - "16": 49940287488.0, - "17": 49941729280.0, - "18": 49941733376.0, - "19": 49941741568.0, - "20": 49941778432.0, - "21": 49941778432.0, - "22": 49941778432.0, - "23": 49941778432.0, - "24": 49941778432.0, - "25": 49941778432.0, - "26": 49941778432.0, - "27": 49941934080.0, - "28": 49941934080.0, - "29": 49941934080.0, - "30": 49941934080.0, - "31": 49942675456.0, - "32": 49942675456.0, - "33": 49942675456.0, - "34": 49942675456.0, - "35": 49942675456.0, - "36": 49942675456.0, - "37": 49942675456.0, - "38": 49942675456.0, - "39": 49942675456.0, - "40": 49944379392.0, - "41": 49944379392.0, - "42": 49944379392.0, - "43": 49944379392.0, - "44": 49944379392.0, - "45": 49944379392.0, - "46": 49944379392.0, - "47": 49944379392.0, - "48": 49944379392.0, - "49": 49944379392.0, - "50": 49944379392.0 + "1": 49104257024.0, + "2": 49953497088.0, + "3": 49955368960.0, + "4": 49955368960.0, + "5": 49955368960.0, + "6": 49955368960.0, + "7": 49955368960.0, + "8": 49955745792.0, + "9": 49955745792.0, + "10": 49957498880.0, + "11": 49957838848.0, + "12": 49957838848.0, + "13": 49957838848.0, + "14": 49957838848.0, + "15": 49957838848.0, + "16": 49957838848.0, + "17": 49957838848.0, + "18": 49957838848.0, + "19": 49957838848.0, + "20": 49957838848.0, + "21": 49957838848.0, + "22": 49957838848.0, + "23": 49957838848.0, + "24": 49957838848.0, + "25": 49957838848.0, + "26": 49957838848.0, + "27": 49957838848.0, + "28": 49957838848.0, + "29": 49957838848.0, + "30": 49957838848.0 } }, "iteration-time": { "start_step": 1, - "end_step": 50, + "end_step": 30, "step_interval": 1, "values": { - "1": 63.07516, - "2": 4.36236, - "3": 3.83222, - "4": 3.85784, - "5": 3.74494, - "6": 3.82661, - "7": 4.05458, - "8": 3.76622, - "9": 3.90518, - "10": 4.09283, - "11": 3.96358, - "12": 3.85778, - "13": 3.84546, - "14": 3.85497, - "15": 4.35749, - "16": 3.7861, - "17": 3.8896, - "18": 3.6267, - "19": 3.76463, - "20": 3.6953, - "21": 3.63427, - "22": 3.66652, - "23": 3.60379, - "24": 3.57701, - "25": 3.57327, - "26": 3.71371, - "27": 3.69626, - "28": 3.89285, - "29": 3.62405, - "30": 3.58297, - "31": 3.56993, - "32": 3.75257, - "33": 3.72279, - "34": 3.48095, - "35": 3.60831, - "36": 3.74971, - "37": 3.72155, - "38": 3.51054, - "39": 3.64562, - "40": 3.66038, - "41": 3.86018, - "42": 3.58341, - "43": 3.82647, - "44": 3.85728, - "45": 3.62416, - "46": 3.59141, - "47": 3.74512, - "48": 3.61762, - "49": 3.57079, - "50": 3.66209 + "1": "nan", + "2": 54.85374, + "3": 4.04314, + "4": 3.83505, + "5": 4.00853, + "6": 3.71939, + "7": 3.66436, + "8": 4.07479, + "9": 3.90049, + "10": 4.34491, + "11": 3.98659, + "12": 3.90765, + "13": 4.12679, + "14": 3.75558, + "15": 3.72381, + "16": 3.45749, + "17": 3.73387, + "18": 3.71406, + "19": 3.75517, + "20": 3.94287, + "21": 3.88534, + "22": 3.86744, + "23": 3.87809, + "24": 3.86352, + "25": 3.87829, + "26": 3.76391, + "27": 3.76762, + "28": 3.96514, + "29": 3.92952, + "30": 3.87378 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/model_config.yaml index b74417a898b..b5788d64049 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + THROUGHPUT_START_STEP: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 @@ -70,7 +71,7 @@ MODEL_ARGS: --dist-ckpt-strictness: log_unexpected --perform-rl-step: true --train-samples: 48828125 - --exit-interval: 50 + --exit-interval: 30 --tensorboard-dir: ${TENSORBOARD_PATH} --save-interval: 1000000 --eval-interval: 1000000 @@ -79,6 +80,11 @@ MODEL_ARGS: --rl-inference-tensor-model-parallel-size: 1 --rl-inference-pipeline-model-parallel-size: 2 --refit-method: gloo + --deterministic-mode: true METRICS: + - "iteration-time" + - "lm loss" + - "num-zeros" - "mem-allocated-bytes" - "mem-max-allocated-bytes" + diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml index 3037e2e0803..ada0350b876 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + THROUGHPUT_START_STEP: 10 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 @@ -99,5 +100,3 @@ METRICS: - "mem-allocated-bytes" - "mem-max-allocated-bytes" - "iteration-time" -THROUGHPUT_TEST_PARAMS: - --start_step: 10 diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml index 456280fdb04..4490ced3988 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp4_pp1_dp2_8b_throughput_github/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + THROUGHPUT_START_STEP: 10 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 @@ -99,5 +100,3 @@ METRICS: - "mem-allocated-bytes" - "mem-max-allocated-bytes" - "iteration-time" -THROUGHPUT_TEST_PARAMS: - --start_step: 10 diff --git a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml index ed5d123892e..c7dcfa594d8 100644 --- a/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt_grpo_tp8tp4_pp1_ep8ep2_dp8_throughputtest/model_config.yaml @@ -1,5 +1,6 @@ ENV_VARS: CUDA_DEVICE_MAX_CONNECTIONS: 1 + THROUGHPUT_START_STEP: 1 NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 NCCL_ALGO: Ring CUBLAS_WORKSPACE_CONFIG: :4096:8 @@ -137,5 +138,3 @@ METRICS: - "mem-allocated-bytes" - "mem-max-allocated-bytes" -THROUGHPUT_TEST_PARAMS: - --start_step: 1 diff --git a/tests/test_utils/recipes/h100/gpt-grpo.yaml b/tests/test_utils/recipes/h100/gpt-grpo.yaml index e707c1c2431..faaccee73dd 100644 --- a/tests/test_utils/recipes/h100/gpt-grpo.yaml +++ b/tests/test_utils/recipes/h100/gpt-grpo.yaml @@ -77,5 +77,5 @@ products: - test_case: [gpt_grpo_tp2tp1_pp4pp2_dp8_583m_throughputtest] products: - environment: [dev] - scope: [mr-broken] + scope: [mr] platforms: [dgx_h100]