From d586ddbf47791236ea035d50b9dab672165bdc36 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Tue, 17 Feb 2026 18:11:02 -0800 Subject: [PATCH 01/23] Remove do_not_average_loss Signed-off-by: Yi-Fu Wu --- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- nemo_rl/models/megatron/train.py | 21 ++++++++++++++----- .../policy/workers/megatron_policy_worker.py | 2 +- tests/unit/models/megatron/test_train.py | 10 +++++---- 4 files changed, 24 insertions(+), 11 deletions(-) diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index 193463c4f8..b12071b947 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit 193463c4f8414e6906a40dd527a450bca50706b1 +Subproject commit b12071b947f9ee3c6616306662069fc4ca77be4c diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 95ccc3761d..8459eada93 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -220,7 +220,6 @@ def megatron_forward_backward( defer_fp32_logits: Optional[bool] = False, global_valid_seqs: Optional[torch.Tensor] = None, global_valid_toks: Optional[torch.Tensor] = None, - do_not_average_loss: bool = False, straggler_timer: Optional[StragglerDetector] = None, ) -> Any: """Execute forward and backward passes using Megatron's utilities. @@ -241,7 +240,6 @@ def megatron_forward_backward( defer_fp32_logits: Whether to skip the conversion of logits to fp32 global_valid_seqs: Global valid sequence count for loss normalization global_valid_toks: Global valid token count for loss normalization - do_not_average_loss: If True, do not average loss across microbatches straggler_timer: Straggler detector for profiling the forward pass Returns: @@ -266,7 +264,6 @@ def megatron_forward_backward( micro_batch_size=mbs, decoder_seq_length=seq_length, forward_only=forward_only, - do_not_average_loss=do_not_average_loss, ) @@ -275,10 +272,12 @@ def __init__( self, loss_fn: LossFunction, cfg: PolicyConfig, + num_microbatches: int = 1, cp_normalize: bool = True, ): self.loss_fn = loss_fn self.cfg = cfg + self.num_microbatches = num_microbatches self.cp_normalize = cp_normalize def __call__( @@ -325,14 +324,26 @@ def __call__( if self.cp_normalize: cp_size = get_context_parallel_world_size() - orig_loss_fn_wrapped = loss_fn_wrapped + prev_loss_fn = loss_fn_wrapped def _div_by_cp_size(*args, **kwargs): - loss, metrics = orig_loss_fn_wrapped(*args, **kwargs) + loss, metrics = prev_loss_fn(*args, **kwargs) return loss / cp_size, metrics loss_fn_wrapped = _div_by_cp_size + # Counteract Megatron's default loss averaging in schedules.py, + # which applies (* cp_size / num_microbatches) to the loss. + cp_size = get_context_parallel_world_size() + num_microbatches = self.num_microbatches + loss_fn_before_mcore_scaling = loss_fn_wrapped + + def _counteract_mcore_loss_averaging(*args, **kwargs): + loss, metrics = loss_fn_before_mcore_scaling(*args, **kwargs) + return loss * num_microbatches / cp_size, metrics + + loss_fn_wrapped = _counteract_mcore_loss_averaging + return loss_fn_wrapped diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index d9a1c3d8a3..5f1483ed9a 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -318,6 +318,7 @@ def train( loss_post_processor = LossPostProcessor( loss_fn=loss_fn, cfg=self.cfg, + num_microbatches=num_microbatches, ) rerun_state_machine = get_rerun_state_machine() @@ -339,7 +340,6 @@ def train( defer_fp32_logits=self.defer_fp32_logits, global_valid_seqs=global_valid_seqs, global_valid_toks=global_valid_toks, - do_not_average_loss=True, straggler_timer=self.mcore_state.straggler_timer, ) diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index cf261c3d75..24dda67eec 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -719,13 +719,15 @@ def test_loss_post_processor_no_packing( def test_loss_post_processor_with_cp_normalize( self, mock_cp_size, mock_cp_grp, mock_tp_grp, mock_tp_rank ): - """Test LossPostProcessor with CP normalization.""" + """Test LossPostProcessor with CP normalization and microbatch pre-scaling.""" from nemo_rl.models.megatron.train import LossPostProcessor mock_loss_fn = MagicMock(return_value=(torch.tensor(1.0), {})) cfg = {"sequence_packing": {"enabled": False}} - processor = LossPostProcessor(loss_fn=mock_loss_fn, cfg=cfg, cp_normalize=True) + processor = LossPostProcessor( + loss_fn=mock_loss_fn, cfg=cfg, num_microbatches=4, cp_normalize=True + ) # Set up mock return values for process groups mock_tp_grp.return_value = MagicMock() @@ -736,8 +738,8 @@ def test_loss_post_processor_with_cp_normalize( output_tensor = torch.randn(2, 10, 100) loss, _ = wrapped_fn(output_tensor) - # Loss should be divided by CP size (2) - assert torch.isclose(loss, torch.tensor(0.5)) + # Loss should be scaled by num_microbatches / (cp_size * cp_size) = 4 / (2 * 2) = 1.0 + assert torch.isclose(loss, torch.tensor(1.0)) @patch( "nemo_rl.models.megatron.train.get_tensor_model_parallel_rank", return_value=0 From 96d4a1181d64b5e65aeaaff56ac9954b3c202a46 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Wed, 18 Feb 2026 15:13:40 -0800 Subject: [PATCH 02/23] Update gitmodules for mcore branch Signed-off-by: Yi-Fu Wu --- .gitmodules | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index 81d066b8b0..c1b0c5a56f 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM url = https://github.com/yaoyu-33/Megatron-LM.git - branch = main + branch = yifu/remove_do_not_average_loss shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge From e5d1ae9a32f72b923316bf7bf8111883f4dc3fba Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Thu, 19 Feb 2026 11:44:03 +0200 Subject: [PATCH 03/23] Switching mcore to upstream main Signed-off-by: Ahmad Kiswani --- .gitmodules | 4 +- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- 3rdparty/Megatron-LM-workspace/setup.py | 7 +- uv.lock | 76 ++++++++++++++++++---- 4 files changed, 73 insertions(+), 16 deletions(-) diff --git a/.gitmodules b/.gitmodules index c1b0c5a56f..8d7c7be7e5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM - url = https://github.com/yaoyu-33/Megatron-LM.git - branch = yifu/remove_do_not_average_loss + url = https://github.com/NVIDIA/Megatron-LM.git + branch = main shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index b12071b947..0d0943c6bf 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit b12071b947f9ee3c6616306662069fc4ca77be4c +Subproject commit 0d0943c6bfa9cbb30fcd62d40ce1792c4cb201e8 diff --git a/3rdparty/Megatron-LM-workspace/setup.py b/3rdparty/Megatron-LM-workspace/setup.py index fb0a7cf92e..380864cc25 100644 --- a/3rdparty/Megatron-LM-workspace/setup.py +++ b/3rdparty/Megatron-LM-workspace/setup.py @@ -43,7 +43,7 @@ # VCS dependencies use full "pkg @ git+URL@rev" format matching pyproject.toml [tool.uv.sources] CACHED_DEPENDENCIES = [ # Default dependencies from pyproject.toml - "torch", + "torch>=2.6.0", "numpy", "packaging>=24.2", # Dev dependencies from pyproject.toml @@ -58,7 +58,7 @@ "opentelemetry-api~=1.33.1", "mamba-ssm~=2.2", "causal-conv1d~=1.5", - "flash-linear-attention~=0.3.2", + "flash-linear-attention~=0.4.0", "nv-grouped-gemm~=1.1", "megatron-energon[av_decode]~=6.0", "av", @@ -69,6 +69,9 @@ "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0", "datasets", "fastapi~=0.50", + "flask[async]", + "hypercorn", + "openai", ] diff --git a/uv.lock b/uv.lock index e0c3cda97f..3a379f1303 100644 --- a/uv.lock +++ b/uv.lock @@ -555,6 +555,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ae/69/fe387b0f70ed608a363a90036e08ef8c1e844e5c98145502160661012dc0/apache_tvm_ffi-0.1.4-cp314-cp314t-win_amd64.whl", hash = "sha256:1bceda57240d03a3cf026334521c0595d097ab92b6d0df7485cbb37b2c056c27", size = 1794928, upload-time = "2025-11-30T07:21:25.234Z" }, ] +[[package]] +name = "asgiref" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/40/f03da1264ae8f7cfdbf9146542e5e7e8100a4c66ab48e791df9a03d3f6c0/asgiref-3.11.1.tar.gz", hash = "sha256:5f184dc43b7e763efe848065441eac62229c9f7b0475f41f80e207a114eda4ce", size = 38550, upload-time = "2026-02-03T13:30:14.33Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" }, +] + [[package]] name = "astor" version = "0.8.1" @@ -1924,16 +1933,16 @@ wheels = [ [[package]] name = "fla-core" -version = "0.3.2" +version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "einops" }, { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.9.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/67/c6/10a1149b07e6bab45b2cb2d07f6b827716c2baf5f3404161753f25c6389b/fla_core-0.3.2.tar.gz", hash = "sha256:d38db16bc4e1c6fa8c04df442f246da1e6926a209426bc6ef703d41bfbc37c92", size = 296725, upload-time = "2025-09-10T07:43:40.155Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/de/0d6bd5664ba2e711cabdde11ccb41ddcdd866c531e40900af3601bd7b8c6/fla_core-0.4.1.tar.gz", hash = "sha256:38ab28966eeadc2141b29e87c2bf72a8a4851e00af9d25bbbc3596b1fb53450d", size = 319608, upload-time = "2025-12-24T18:07:37.669Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/f5/74947b33c07682280e65adbdf17c4ee94b30232df2f728bafecf13d1d820/fla_core-0.3.2-py3-none-any.whl", hash = "sha256:e751d5a41e33eee721a6fb6588bd857f6f36e0d14719a23b1ebdbd617d307209", size = 413594, upload-time = "2025-09-10T07:43:37.786Z" }, + { url = "https://files.pythonhosted.org/packages/f6/43/945ef69eb48a14c30fd7323d3e0b560c821ae71e6d3ef979e06a901bc3b9/fla_core-0.4.1-py3-none-any.whl", hash = "sha256:93c6afe4c80fc7bc705fa8aeea6a46d2cf2d77383f9619a41863c7114c801bab", size = 437282, upload-time = "2025-12-24T18:07:34.41Z" }, ] [[package]] @@ -1952,17 +1961,15 @@ sdist = { url = "https://files.pythonhosted.org/packages/e8/6d/7066d160bdffa2f9d [[package]] name = "flash-linear-attention" -version = "0.3.2" +version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "datasets" }, { name = "fla-core" }, - { name = "pytest" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/84/f6/e62c1e562a288557eba7f06f168a7615813d1a227327b8beb8ba426da2c5/flash_linear_attention-0.3.2.tar.gz", hash = "sha256:9147747316c2951fed4ebeb4fa87977c05d807dc70c93b46250b68a6eb1183e2", size = 150880, upload-time = "2025-09-10T07:43:41.37Z" } +sdist = { url = "https://files.pythonhosted.org/packages/46/83/7d8ec7ffb5229080b1c9b772338ff588cbd63282ac355ede2a12a6e174a8/flash_linear_attention-0.4.1.tar.gz", hash = "sha256:127ee7273ed15ac17f72bcf4c75e1051719d8fbe0a2d1d047e59406f36d81ee2", size = 158280, upload-time = "2025-12-24T18:07:38.812Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/d0/35ce9eac5f52c72005095aaa12a393d2656ed7ffedf925b2381a6b76d10c/flash_linear_attention-0.3.2-py3-none-any.whl", hash = "sha256:604e73361437ba786420ab195e2caa3fd19280503761e703fa353c5ce5c65376", size = 274592, upload-time = "2025-09-10T07:43:39.107Z" }, + { url = "https://files.pythonhosted.org/packages/63/d5/6327559a9d5b9243b10c3984f1bcef256ed2ad06d105a3bb8f7b2979659c/flash_linear_attention-0.4.1-py3-none-any.whl", hash = "sha256:d18bdfe9d1f4b424676444eac9d50fb8433b70e5d4e0e0878b20bcbcdbea57ce", size = 287415, upload-time = "2025-12-24T18:07:35.815Z" }, ] [[package]] @@ -2141,6 +2148,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/f9/7f9263c5695f4bd0023734af91bedb2ff8209e8de6ead162f35d8dc762fd/flask-3.1.2-py3-none-any.whl", hash = "sha256:ca1d8112ec8a6158cc29ea4858963350011b5c846a414cdb7a954aa9e967d03c", size = 103308, upload-time = "2025-08-19T21:03:19.499Z" }, ] +[package.optional-dependencies] +async = [ + { name = "asgiref" }, +] + [[package]] name = "flask-cors" version = "6.0.1" @@ -2851,6 +2863,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, ] +[[package]] +name = "hypercorn" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "h2" }, + { name = "priority" }, + { name = "wsproto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/44/01/39f41a014b83dd5c795217362f2ca9071cf243e6a75bdcd6cd5b944658cc/hypercorn-0.18.0.tar.gz", hash = "sha256:d63267548939c46b0247dc8e5b45a9947590e35e64ee73a23c074aa3cf88e9da", size = 68420, upload-time = "2025-11-08T13:54:04.78Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/35/850277d1b17b206bd10874c8a9a3f52e059452fb49bb0d22cbb908f6038b/hypercorn-0.18.0-py3-none-any.whl", hash = "sha256:225e268f2c1c2f28f6d8f6db8f40cb8c992963610c5725e13ccfcddccb24b1cd", size = 61640, upload-time = "2025-11-08T13:54:03.202Z" }, +] + [[package]] name = "hyperframe" version = "6.1.0" @@ -3652,6 +3679,8 @@ dependencies = [ { name = "flash-linear-attention" }, { name = "flashinfer-python", version = "0.5.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-vllm' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang')" }, { name = "flashinfer-python", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-sglang' or extra != 'extra-7-nemo-rl-vllm'" }, + { name = "flask", extra = ["async"] }, + { name = "hypercorn" }, { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"] }, { name = "multi-storage-client" }, @@ -3661,6 +3690,7 @@ dependencies = [ { name = "nvidia-resiliency-ext" }, { name = "nvtx" }, { name = "onnxscript" }, + { name = "openai" }, { name = "opentelemetry-api" }, { name = "packaging" }, { name = "tensorstore", version = "0.1.74", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -3680,8 +3710,10 @@ requires-dist = [ { name = "einops", specifier = "~=0.8" }, { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.1.0" }, { name = "fastapi", specifier = "~=0.50" }, - { name = "flash-linear-attention", specifier = "~=0.3.2" }, + { name = "flash-linear-attention", specifier = "~=0.4.0" }, { name = "flashinfer-python", specifier = "~=0.5.0" }, + { name = "flask", extras = ["async"] }, + { name = "hypercorn" }, { name = "mamba-ssm", git = "https://github.com/state-spaces/mamba.git?rev=d68d16ed7d5d5164eb5a57c0285f3b7eb8394ec1" }, { name = "megatron-energon", extras = ["av-decode"], specifier = "~=6.0" }, { name = "multi-storage-client", specifier = "~=0.27" }, @@ -3691,11 +3723,12 @@ requires-dist = [ { name = "nvidia-resiliency-ext" }, { name = "nvtx", specifier = "~=0.2" }, { name = "onnxscript" }, + { name = "openai" }, { name = "opentelemetry-api", specifier = "~=1.33.1" }, { name = "packaging", specifier = ">=24.2" }, { name = "tensorstore", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, - { name = "torch", marker = "sys_platform != 'darwin'", index = "https://download.pytorch.org/whl/cu129" }, - { name = "torch", marker = "sys_platform == 'darwin'", index = "https://pypi.org/simple" }, + { name = "torch", marker = "sys_platform != 'darwin'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cu129" }, + { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.6.0", index = "https://pypi.org/simple" }, { name = "tqdm" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], specifier = ">=2.9.0a0,<2.12.0" }, { name = "wget" }, @@ -6132,6 +6165,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/c7/5613524e606ea1688b3bdbf48aa64bafb6d0a4ac3750274c43b6158a390f/prettytable-3.16.0-py3-none-any.whl", hash = "sha256:b5eccfabb82222f5aa46b798ff02a8452cf530a352c31bddfa29be41242863aa", size = 33863, upload-time = "2025-03-24T19:39:02.359Z" }, ] +[[package]] +name = "priority" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/3c/eb7c35f4dcede96fca1842dac5f4f5d15511aa4b52f3a961219e68ae9204/priority-2.0.0.tar.gz", hash = "sha256:c965d54f1b8d0d0b19479db3924c7c36cf672dbf2aec92d43fbdaf4492ba18c0", size = 24792, upload-time = "2021-06-27T10:15:05.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/5f/82c8074f7e84978129347c2c6ec8b6c59f3584ff1a20bc3c940a3e061790/priority-2.0.0-py3-none-any.whl", hash = "sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa", size = 8946, upload-time = "2021-06-27T10:15:03.856Z" }, +] + [[package]] name = "prometheus-client" version = "0.22.1" @@ -9861,6 +9903,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" }, ] +[[package]] +name = "wsproto" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" }, +] + [[package]] name = "xattr" version = "1.3.0" From 50aa2de969f81b6f8a3b25e09872a18d87f37381 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Thu, 19 Feb 2026 14:52:00 -0800 Subject: [PATCH 04/23] refit latest update based on ahmads refactor --- examples/configs/grpo_math_1B.yaml | 14 +- examples/configs/grpo_math_1B_megatron.yaml | 12 +- migration_notes_delete.md | 675 ++++++++++++++++ nemo_rl/algorithms/grpo.py | 51 +- .../models/generation/megatron/__init__.py | 19 + .../megatron/megatron_generation.py | 205 +++++ nemo_rl/models/megatron/setup.py | 19 + nemo_rl/models/policy/lm_policy.py | 42 +- .../policy/workers/megatron_policy_worker.py | 726 ++++++++++++++---- 9 files changed, 1591 insertions(+), 172 deletions(-) create mode 100644 migration_notes_delete.md create mode 100644 nemo_rl/models/generation/megatron/__init__.py create mode 100644 nemo_rl/models/generation/megatron/megatron_generation.py diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 11aebb9e84..dfe684095d 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -146,6 +146,11 @@ policy: moe_enable_deepep: false moe_token_dispatcher_type: "allgather" moe_shared_expert_overlap: false + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: "local" + cuda_graph_scope: null + use_te_rng_tracker: true + inference_rng_tracker: true optimizer: optimizer: "adam" @@ -252,13 +257,16 @@ policy: stop_strings: null mcore_generation_config: buffer_size_gb: 20 # Total GPU memory (in GB) allocated for KV cache buffers - buffer_guaranteed_fraction: 0.1 # Fraction of buffer reserved for guaranteed active requests num_cuda_graphs: 16 # Number of CUDA graphs to pre-compile for different batch sizes block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity) use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing - enable_chunked_prefill: true # Split long prefills into chunks for better memory management - unified_memory_level: 0 # Unified memory usage level (0=disabled, higher values enable more aggressive paging) + unified_memory_level: 0 # Unified memory usage level (0=disabled, 1+=enables unified memory ) max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens + enable_chunked_prefill: false + kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute" + static_kv_memory_pointers: false # Relevant only for offload and recompute modes + materialize_only_last_token_logits: false + vllm_cfg: async_engine: false precision: ${policy.precision} diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index b240c6519c..13d9634eb3 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -100,6 +100,11 @@ policy: moe_shared_expert_overlap: false #gives ~20% training perf speedup with sequence packing apply_rope_fusion: True + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: "local" + cuda_graph_scope: null + use_te_rng_tracker: true + inference_rng_tracker: true optimizer: optimizer: "adam" @@ -151,9 +156,12 @@ policy: num_cuda_graphs: 16 # Number of CUDA graphs to pre-compile for different batch sizes block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity) use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing - enable_chunked_prefill: false # Split long prefills into chunks for better memory management - unified_memory_level: 0 # Unified memory usage level (0=disabled, higher values enable more aggressive paging) + unified_memory_level: 1 # Unified memory usage level (0=disabled, 1+=enables unified memory ) max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens + enable_chunked_prefill: false + kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute" + static_kv_memory_pointers: false # Relevant only for offload and recompute modes + materialize_only_last_token_logits: false vllm_cfg: tensor_parallel_size: 1 diff --git a/migration_notes_delete.md b/migration_notes_delete.md new file mode 100644 index 0000000000..161fcb3593 --- /dev/null +++ b/migration_notes_delete.md @@ -0,0 +1,675 @@ +# NOTES +### MegatronTokenizer Issue. +``` + from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer as MegatronTokenizer +ModuleNotFoundError: No module named 'megatron.core.datasets.megatron_tokenizer' +``` +Fix : +Create this file in megatron/core/datasets/megatron_tokenizer.py +``` +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Backward-compatibility shim. The legacy tokenizer base class was moved out of +# megatron.core.datasets in newer Megatron-LM versions. Megatron-Bridge still +# imports from this path, so we keep a thin re-export here. + +import json +import logging +from abc import ABC, abstractmethod +from collections import OrderedDict +from typing import Any + +import numpy + +logger = logging.getLogger(__name__) + + +class MegatronLegacyTokenizer(ABC): + """Abstract class for tokenizer + + Absent a config or class-specific tracking of which objects are uniquely identifying, we must + include all key word arguments as unique identifiers + + Args: + tokenizer_paths (Tuple[str]): All tokenizer source paths or prefixes + + tokenizer_options (Dict[str, Any]): All tokenizer options + """ + + def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any): + logger.warning( + "You're using the legacy tokenizer system, which is deprecated " + "and will be removed in a future release. Please migrate to the new tokenizer system " + "(`megatron.core.tokenizers.MegatronTokenizer`)." + ) + self.unique_identifiers = OrderedDict() + self.unique_identifiers["class"] = type(self).__name__ + self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths) + for option in tokenizer_options: + self.unique_identifiers[option] = str(tokenizer_options[option]) + + self.unique_description = json.dumps(self.unique_identifiers, indent=4) + + super().__init__() + + @abstractmethod + def tokenize(self, text: str) -> numpy.ndarray: + pass + + def detokenize(self, ids: numpy.ndarray) -> str: + raise NotImplementedError("{} has no method 'detokenize'".format(type(self).__name__)) + + def offsets(self, ids: list[int], text: str) -> list[int]: + raise NotImplementedError("{} has no method 'offsets'".format(type(self).__name__)) + + @property + @abstractmethod + def vocab(self): + pass + + @property + @abstractmethod + def inv_vocab(self): + pass + + @property + @abstractmethod + def vocab_size(self): + pass + + @property + def cls(self): + raise NotImplementedError("{} has no attribute 'cls'".format(type(self).__name__)) + + @property + def sep(self): + raise NotImplementedError("{} has no attribute 'sep'".format(type(self).__name__)) + + @property + def pad(self): + raise NotImplementedError("{} has no attribute 'pad'".format(type(self).__name__)) + + @property + def eod(self): + raise NotImplementedError("{} has no attribute 'eod'".format(type(self).__name__)) + + @property + def bos(self): + raise NotImplementedError("{} has no attribute 'bos'".format(type(self).__name__)) + + @property + def eos(self): + raise NotImplementedError("{} has no attribute 'eos'".format(type(self).__name__)) + + @property + def mask(self): + raise NotImplementedError("{} has no attribute 'mask'".format(type(self).__name__)) + + +# Older code imported this class under the name ``MegatronTokenizer``. +MegatronTokenizer = MegatronLegacyTokenizer + +``` + + +### DDP ISSUE +``` +@@ -369,12 +369,15 @@ class DistributedDataParallel(_BaseDataParallel): + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook ++ for module, handle in list(self.remove_forward_pre_hook_handles.items()): ++ handle.remove() ++ self.remove_forward_pre_hook_handles.clear() + +- for module in self.module.modules(): +- assert self.remove_forward_pre_hook_handles[module] is not None +- self.remove_forward_pre_hook_handles[module].remove() +- del self.remove_forward_pre_hook_handles[module] +- assert len(self.remove_forward_pre_hook_handles) == 0 + + # Force synchronize parameters. + if param_sync: +``` + +### EXPLANATION +## Why `disable_forward_pre_hook` is Called + +### The Context: `use_reference_model` + +The `use_reference_model` context manager (lines 1437-1486) temporarily **swaps the model weights** with the reference model weights: + +1. **On entry**: Copies the current model's state_dict to CPU, then loads the reference model's state_dict into the model +2. **On exit**: Restores the original model weights + +This allows running inference with the reference model's weights without having two full models in GPU memory. + +### What is the Forward Pre-Hook? + +Looking at the DDP code you attached, when **overlap_param_gather** is enabled with distributed optimizer: + +```376:386:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def enable_forward_pre_hook(self): + """ + Enable forward pre-hooks needed for param all-gather overlap with forward compute. + """ + assert self.use_forward_hook + assert len(self.remove_forward_pre_hook_handles) == 0 + # Register forward pre-hook for all sub-modules. + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +The forward pre-hook is used to **overlap parameter all-gather with forward compute**. Here's how it works: + +1. With **distributed optimizer**, model parameters are **sharded across data-parallel ranks** (each rank only holds a portion of the parameters) +2. Before forward pass, parameters need to be **all-gathered** to reconstruct full parameters +3. The forward pre-hook intercepts each module's forward call to **wait for the all-gather to complete** for that module's parameters before executing + +```411:437:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def hook(module, *unused): + // ... + # Make sure all parameters in this module have been all-gathered as necessary. + for param in module.parameters(recurse=False): + # Skip parameters without an associated buffer + if param not in self.param_to_bucket_group: + continue + // ... + self.param_to_bucket_group[param].finish_param_sync( + skip_next_bucket_dispatch=skip_next_bucket_dispatch + ) +``` + +### Why Disable It During Weight Swap? + +When swapping weights in `use_reference_model`: + +```1459:1459:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/nemo_rl/models/policy/workers/megatron_policy_worker.py + self.model.load_state_dict(self.reference_state_dict, strict=True) +``` + +**The forward pre-hook would interfere because:** + +1. The hook maintains state about **which parameters have been all-gathered** via `param_to_bucket_group` +2. When you swap weights using `load_state_dict`, the parameter tensors are replaced, but the **bucket mappings become stale/invalid** +3. If the hook tries to call `finish_param_sync` on parameters that weren't properly set up for the new weights, it could cause errors or incorrect behavior +4. Additionally, `disable_forward_pre_hook` calls `start_param_sync(force_sync=True)` which **synchronously gathers all parameters** first, ensuring a clean state before the weight swap + +After the weight swap is complete (in `finally` block), the hooks are re-enabled: + +```1484:1486:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/nemo_rl/models/policy/workers/megatron_policy_worker.py + ## re-enable overlap param gather after weight swap + if self.should_disable_forward_pre_hook: + self.enable_forward_pre_hook() +``` + +### Summary + +| Step | What Happens | +|------|-------------| +| 1. `disable_forward_pre_hook()` | Removes all forward hooks, forces a synchronous param all-gather to ensure all params are complete | +| 2. `load_state_dict()` | Safely swap weights without hooks interfering | +| 3. Forward passes | Run without overlap optimization (simpler, safer) | +| 4. `enable_forward_pre_hook()` | Re-register hooks after restoring original weights | + +This is essentially a safety measure: disable the async overlap optimization during the weight-swapping operation to avoid race conditions or stale state in the DDP infrastructure. + +```python + def enable_forward_pre_hook(self): + """ + Enable forward pre-hooks needed for param all-gather overlap with forward compute. + """ + assert self.use_forward_hook + assert len(self.remove_forward_pre_hook_handles) == 0 + # Register forward pre-hook for all sub-modules. + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +```python + def hook(module, *unused): + // ... + # Make sure all parameters in this module have been all-gathered as necessary. + for param in module.parameters(recurse=False): + # Skip parameters without an associated buffer + if param not in self.param_to_bucket_group: + continue + // ... + self.param_to_bucket_group[param].finish_param_sync( + skip_next_bucket_dispatch=skip_next_bucket_dispatch + ) +``` + +```python + self.model.load_state_dict(self.reference_state_dict, strict=True) +``` + +```python + ## re-enable overlap param gather after weight swap + if self.should_disable_forward_pre_hook: + self.enable_forward_pre_hook() +``` + +--- + +**User** + +When I run the application I get an error like this +``` + File "/opt/nemo-rl/examples/run_grpo_math.py", line 260, in + main() + File "/opt/nemo-rl/examples/run_grpo_math.py", line 243, in main + grpo_train( + File "/opt/nemo-rl/nemo_rl/algorithms/grpo.py", line 1339, in grpo_train + reference_logprobs = policy.get_reference_policy_logprobs( + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 390, in get_reference_policy_logprobs + self.worker_group.get_all_worker_results(futures) + File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 953, in get_all_worker_results + return future_bundle.get_results( + ^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 103, in get_results + all_results = ray.get(object_refs) + ^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper + return fn(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper + return func(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, in get + values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 968, in get_objects + raise value.as_instanceof_cause() +ray.exceptions.RayTaskError(AssertionError): ray::MegatronPolicyWorker.get_reference_policy_logprobs() (pid=710294, ip=10.65.29.29, actor_id=65a88a2a253f3b406b563ea701000000, repr=MegatronPolicyWorker[rank=0]) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/utils/nsys.py", line 88, in wrapper + ret = func(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/base_policy_worker.py", line 143, in get_reference_policy_logprobs + with self.use_reference_model(): + ^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/root/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/contextlib.py", line 137, in __enter__ + return next(self.gen) + ^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 1446, in use_reference_model + self.disable_forward_pre_hook() + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 944, in disable_forward_pre_hook + self.model.disable_forward_pre_hook(param_sync=param_sync) + File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py", line 399, in disable_forward_pre_hook + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: remove_forward_pre_hook_handles is not empty, it is: {CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): } +``` +What does this error mean ? + +--- + +**Cursor** + +```plaintext +AssertionError: remove_forward_pre_hook_handles is not empty, it is: {CudaGraphManager(): ...} +``` + +```python +assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) + ... +``` + +This error is revealing a **mismatch between hook registration and deregistration** due to `CudaGraphManager` modules. Let me explain: + +## What the Error Means + +The assertion failure happens in `disable_forward_pre_hook`: + +```388:399:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +### The Problem + +1. **When `enable_forward_pre_hook` was called**: It iterated through `self.module.modules()` and registered hooks for **ALL** modules, including `CudaGraphManager` objects + +2. **When `disable_forward_pre_hook` is called**: It again iterates through `self.module.modules()` to remove hooks + +3. **The mismatch**: The `CudaGraphManager` modules that were present during registration are **no longer being returned** by `self.module.modules()` during deregistration. This leaves their handles in the dictionary, causing the assertion to fail. + +### Why This Happens + +Looking at the leftover keys in the error: +``` +{CudaGraphManager(): , + CudaGraphManager(): , + ...} +``` + +These `CudaGraphManager` objects were registered as submodules when `enable_forward_pre_hook` ran, but they're either: +- Dynamically removed from the module tree after registration +- Not yielded by `modules()` during the second iteration for some reason +- Created/destroyed dynamically as part of CUDA graph management + +### Root Cause + +This is likely a **bug or incompatibility** between: +1. **CUDA Graphs** being enabled in your configuration (`cuda_graph_impl = "local"` as I saw earlier in the code) +2. **Overlap param gather** being enabled (`overlap_param_gather = True`) + +When CUDA graphs are used, the model may dynamically add/remove `CudaGraphManager` wrapper modules, which breaks the assumption that `modules()` returns the same set both times. + +### Potential Fixes + +**Option 1**: Disable CUDA graphs when using reference model weight swapping (if that's acceptable for your use case) + +**Option 2**: Fix the `disable_forward_pre_hook` to be more robust by clearing ALL handles, not just those from current `modules()`: + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove all registered handles + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` + +**Option 3**: Skip calling `disable_forward_pre_hook` when reference model doesn't need it (if `init_reference_model=False` was used or reference model was never initialized) + +Would you like me to help you implement one of these fixes? + +```python + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +```plaintext +{CudaGraphManager(): , + CudaGraphManager(): , + ...} +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove all registered handles + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` + +```python +def toggle_cuda_graphs(model, set_to="none", reset_cuda_graphs=True): + ... + elif attribute == "cudagraph_manager": + for module in modules: + if set_to == "local": + if reset_cuda_graphs: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + # If we are resetting cuda graphs we create a new cuda graph manager + setattr(module[0], attribute, CudaGraphManager(model.config)) + else: + # If we are not resetting cuda graphs we set it to its cached cuda graph + setattr(module[0], attribute, module[1]) + else: + for module in modules: + # If we are deleting the cuda graph, we delete its attribute + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") + ... +``` + +```python +setattr(module[0], attribute, CudaGraphManager(model.config)) +``` + +```python +delattr(module[0], "cudagraph_manager") +``` + +```python +toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) + + +``` + +```plaintext +Two potential solutions emerge: +1. Disable hooks before structural module changes +2. Modify hook removal logic to track references independently of module hierarchy +``` + +Let me check where `toggle_cuda_graphs` is being called in relation to the hook operations: + +**Yes, this is exactly the cause!** Let me explain the problem: + +## Root Cause Analysis + +### The Sequence of Events + +1. **During `generate()`** (line 1875): + ```python + toggle_cuda_graphs(lang_module, set_to="local", reset_cuda_graphs=True) + ``` + This **creates new `CudaGraphManager` objects** and adds them as submodules: + ```python + # From toggle_cuda_graphs when set_to="local": + setattr(module[0], attribute, CudaGraphManager(model.config)) + ``` + +2. **At some point**, `enable_forward_pre_hook()` is called, which registers hooks for **ALL** current modules, including the newly created `CudaGraphManager` objects: + + ```376:386:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def enable_forward_pre_hook(self): + # ... + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) + ``` + +3. **At end of `generate()`** (line 2045): + ```python + toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) + ``` + This **DELETES the `CudaGraphManager` attributes** from their parent modules: + ```python + # From toggle_cuda_graphs when set_to="none": + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") # ← REMOVES from module tree! + ``` + +4. **When `get_reference_policy_logprobs` → `use_reference_model` → `disable_forward_pre_hook()` is called**: + It iterates through `self.module.modules()`, but the `CudaGraphManager` objects are **no longer in the module tree** (they were deleted in step 3). + +5. **Result**: The hook handles for `CudaGraphManager` objects remain in `remove_forward_pre_hook_handles` but can't be removed → **Assertion fails!** + +## Visual Timeline + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ generate() starts │ +│ └─ toggle_cuda_graphs(set_to="local") │ +│ └─ Creates CudaGraphManager as submodule │ +├─────────────────────────────────────────────────────────────────────┤ +│ enable_forward_pre_hook() called somewhere │ +│ └─ Registers hooks for ALL modules (including CudaGraphManager) │ +│ └─ remove_forward_pre_hook_handles[CudaGraphManager] = handle │ +├─────────────────────────────────────────────────────────────────────┤ +│ generate() ends │ +│ └─ toggle_cuda_graphs(set_to="none") │ +│ └─ delattr(module, "cudagraph_manager") ← DELETES IT! │ +├─────────────────────────────────────────────────────────────────────┤ +│ get_reference_policy_logprobs() │ +│ └─ use_reference_model() │ +│ └─ disable_forward_pre_hook() │ +│ └─ for module in self.module.modules(): ← CudaGraphManager│ +│ # is NOT here anymore! │ +│ └─ ASSERTION FAILS: handles dict not empty! │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +## Solution Options + +**Option 1**: Fix `disable_forward_pre_hook` in DDP to be robust (best fix): + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove ALL registered handles (not just current modules) + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + if param_sync: + self.start_param_sync(force_sync=True) +``` + +**Option 2**: Call `disable_forward_pre_hook` BEFORE `toggle_cuda_graphs(set_to="none")` in your code. + +**Option 3**: Don't delete cuda graph managers, just disable them (modify `toggle_cuda_graphs`). + +Would you like me to implement Option 1 as a fix in the Megatron DDP code? + +```python + toggle_cuda_graphs(lang_module, set_to="local", reset_cuda_graphs=True) +``` + +```python + # From toggle_cuda_graphs when set_to="local": + setattr(module[0], attribute, CudaGraphManager(model.config)) +``` + +```python + def enable_forward_pre_hook(self): + # ... + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +```python + toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) +``` + +```python + # From toggle_cuda_graphs when set_to="none": + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") # ← REMOVES from module tree! +``` + +```plaintext +┌─────────────────────────────────────────────────────────────────────┐ +│ generate() starts │ +│ └─ toggle_cuda_graphs(set_to="local") │ +│ └─ Creates CudaGraphManager as submodule │ +├─────────────────────────────────────────────────────────────────────┤ +│ enable_forward_pre_hook() called somewhere │ +│ └─ Registers hooks for ALL modules (including CudaGraphManager) │ +│ └─ remove_forward_pre_hook_handles[CudaGraphManager] = handle │ +├─────────────────────────────────────────────────────────────────────┤ +│ generate() ends │ +│ └─ toggle_cuda_graphs(set_to="none") │ +│ └─ delattr(module, "cudagraph_manager") ← DELETES IT! │ +├─────────────────────────────────────────────────────────────────────┤ +│ get_reference_policy_logprobs() │ +│ └─ use_reference_model() │ +│ └─ disable_forward_pre_hook() │ +│ └─ for module in self.module.modules(): ← CudaGraphManager│ +│ # is NOT here anymore! │ +│ └─ ASSERTION FAILS: handles dict not empty! │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove ALL registered handles (not just current modules) + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + if param_sync: + self.start_param_sync(force_sync=True) +``` + +AFTER FIX 1 +### What Changed + +| Before | After | +|--------|-------| +| Iterated over `self.module.modules()` | Iterates over `self.remove_forward_pre_hook_handles.items()` directly | +| Asserted each module exists in dict | No assertion needed - we're iterating over what we have | +| Deleted entries one by one | Uses `clear()` after removing all handles | +| Would fail if modules were removed from tree | Handles dynamically removed modules (like `CudaGraphManager`) | + +This fix ensures that all registered hook handles are properly removed regardless of whether the modules are still part of the module tree. You can now run your application again to test it. + +```python + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all registered modules. + # We iterate over the dict keys directly (instead of self.module.modules()) to handle + # cases where modules may have been dynamically removed from the module tree + # (e.g., CudaGraphManager objects removed by toggle_cuda_graphs). + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` \ No newline at end of file diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d349500516..43aa43ce77 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -60,6 +60,7 @@ from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.models.generation.megatron import MegatronGeneration from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, run_async_nemo_gym_rollout, @@ -539,6 +540,18 @@ def init_sglang(): pg.finish_generation() return pg, time.perf_counter() - t0 + def init_megatron_generation(): + """Initialize Megatron generation workers for non-colocated inference.""" + t0 = time.perf_counter() + mg = MegatronGeneration( + cluster=inference_cluster, + config=policy_config, + tokenizer=tokenizer, + processor=processor, + weights_path=weights_path, + ) + return mg, time.perf_counter() - t0 + def initialize_generation_with_policy( init_generation_fn, generation_name: str, @@ -603,14 +616,38 @@ def initialize_generation_with_policy( # Handle generation-specific setup if backend == "megatron": # Megatron generation: policy_generation is None, only initialize policy - policy_generation = None - print( - f" ✓ Using {backend} backend for generation with {policy_config['model_name']}", - flush=True, - ) + if colocated_inference: + policy_generation = None + print( + f" ✓ Using {backend} backend for generation with {policy_config['model_name']}", + flush=True, + ) - policy, policy_time = init_policy() - worker_init_timing_metrics["policy_init_time_s"] = policy_time + policy, policy_time = init_policy() + worker_init_timing_metrics["policy_init_time_s"] = policy_time + else: + # Non-colocated Megatron backend: separate inference workers + print( + " ⚡ Using parallel worker initialization (non-colocated Megatron mode)", + flush=True, + ) + + # Execute both initializations in parallel + parallel_start_time = time.perf_counter() + with ThreadPoolExecutor(max_workers=2) as executor: + megatron_gen_future = executor.submit(init_megatron_generation) + policy_future = executor.submit(init_policy) + policy_generation, megatron_gen_time = megatron_gen_future.result() + policy, policy_time = policy_future.result() + parallel_wall_time = time.perf_counter() - parallel_start_time + + # Store timing metrics + worker_init_timing_metrics["megatron_generation_init_time_s"] = ( + megatron_gen_time + ) + worker_init_timing_metrics["policy_init_time_s"] = policy_time + worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time + worker_init_timing_metrics["parallel_init_enabled"] = True elif backend == "vllm": # vLLM generation: setup config, then initialize with policy diff --git a/nemo_rl/models/generation/megatron/__init__.py b/nemo_rl/models/generation/megatron/__init__.py new file mode 100644 index 0000000000..cf9d2aa8e3 --- /dev/null +++ b/nemo_rl/models/generation/megatron/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo_rl.models.generation.megatron.megatron_generation import ( + MegatronGeneration, +) + +__all__ = ["MegatronGeneration"] \ No newline at end of file diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py new file mode 100644 index 0000000000..24aab7de7d --- /dev/null +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -0,0 +1,205 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MegatronGeneration: A GenerationInterface implementation for non-colocated +Megatron-based inference. + +This module wraps a Policy object (configured for inference only, without +optimizer or reference model) and exposes it through the GenerationInterface. +It enables non-colocated inference where training and generation run on +separate GPU clusters, with weights synchronized via NCCL collective +communication. + +The init_collective and update_weights_from_collective methods are currently +placeholders that will be implemented in a future PR. +""" + +from typing import Any, Optional + +import ray +from transformers import AutoProcessor +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.interfaces import ( + GenerationDatumSpec, + GenerationInterface, + GenerationOutputSpec, +) +from nemo_rl.models.policy import PolicyConfig + + +class MegatronGeneration(GenerationInterface): + """Generation interface backed by Megatron for non-colocated inference. + + This class creates a Policy instance configured for inference only + (no optimizer, no reference model) on a dedicated inference cluster. + It implements the GenerationInterface so it can be used as a drop-in + replacement for VllmGeneration in the non-colocated inference flow. + + The weight synchronization methods (init_collective, update_weights_from_collective) + are placeholders that will be implemented in a future PR. + """ + + def __init__( + self, + cluster: RayVirtualCluster, + config: PolicyConfig, + tokenizer: PreTrainedTokenizerBase, + name_prefix: str = "megatron_generation", + processor: Optional[AutoProcessor] = None, + weights_path: Optional[str] = None, + ): + """Initialize a MegatronGeneration instance. + + Args: + cluster: The RayVirtualCluster to deploy inference workers on. + config: PolicyConfig for the Megatron model. + tokenizer: The tokenizer for the model. + name_prefix: Prefix for naming the worker group. + processor: Optional processor for VLMs. + weights_path: Optional path to model weights for initialization. + """ + # Import here to avoid circular imports + from nemo_rl.models.policy.lm_policy import Policy + + self.cfg = config + + # Create a Policy object configured for inference only: + # - No optimizer (not training on this cluster) + # - No reference model (not needed for generation) + self._policy = Policy( + cluster=cluster, + config=config, + tokenizer=tokenizer, + name_prefix=name_prefix, + processor=processor, + init_optimizer=False, + init_reference_model=False, + weights_path=weights_path, + ) + + def init_collective( + self, ip: str, port: int, world_size: int, *, train_world_size: int + ) -> list[ray.ObjectRef]: + """Initialize the collective communication for weight synchronization. + + This sets up NCCL communication between training workers and these + inference workers so that updated model weights can be broadcast + from the training cluster to the inference cluster. + + Uses init_collective_as_inference on the workers, which offsets each + worker's rank by train_world_size to avoid colliding with training + workers' ranks (rank = train_world_size + worker_rank). + + Args: + ip: IP address for the process group rendezvous. + port: Port for the process group rendezvous. + world_size: Total world size (train + inference workers). + train_world_size: Number of training workers (used to offset ranks). + + Returns: + List of Ray ObjectRefs for the collective init futures. + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "init_collective_as_inference", + ip=ip, + port=port, + world_size=world_size, + train_world_size=train_world_size, + ) + return futures + + def update_weights_from_collective(self) -> list[ray.ObjectRef]: + """Receive updated weights from the training cluster via collective communication. + + This method is called after the training side calls + policy.broadcast_weights_for_collective(). It receives the broadcast + weights and updates the local model parameters. + + TODO: This is a placeholder. The actual implementation will: + 1. Iterate over the model's state_dict info + 2. Use packed_broadcast_consumer to receive weights from the training side + 3. Update the local model parameters with the received weights + + Returns: + List of Ray ObjectRefs for the weight update futures. + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "update_weights_from_collective", + ) + return futures + + def generate( + self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False + ) -> BatchedDataDict[GenerationOutputSpec]: + """Generate a batch of data using the Megatron generation backend. + + Delegates to the internal Policy's generate method. + + Args: + data: BatchedDataDict containing input_ids and input_lengths. + greedy: Whether to use greedy decoding. + + Returns: + BatchedDataDict conforming to GenerationOutputSpec. + """ + return self._policy.generate(data, greedy=greedy) + + def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: + """Prepare the inference workers for generation. + + For Megatron generation, this is a no-op since the workers + are always ready for inference. + """ + return self._policy.prepare_for_generation(*args, **kwargs) + + def finish_generation(self, *args: Any, **kwargs: Any) -> bool: + """Clean up after generation. + + For Megatron generation, this is a no-op. + """ + return self._policy.finish_generation(*args, **kwargs) + + def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + """Prepare state dict metadata for weight refitting. + + This stores the state dict info (tensor names, shapes, dtypes) on each + inference worker so that update_weights_from_collective knows what + tensors to expect during the weight broadcast. + + Note: This calls store_refit_info on workers (not prepare_refit_info), + because prepare_refit_info on MegatronPolicyWorker calculates and + returns metadata (training-side), while store_refit_info accepts and + stores metadata (inference-side). + + Args: + state_dict_info: Dictionary mapping tensor names to (shape, dtype) tuples, + as returned by the training-side prepare_refit_info(). + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "store_refit_info", + state_dict_info=state_dict_info, + ) + ray.get(futures) + + def shutdown(self) -> bool: + """Shut down all inference workers and clean up resources.""" + return self._policy.shutdown() + + def __del__(self) -> None: + """Safety net to ensure workers are shut down.""" + if hasattr(self, "_policy"): + self._policy.shutdown() diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index e9fc2da9e1..e811f711a5 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -318,6 +318,10 @@ def setup_model_config( # Apply performance settings _apply_performance_config(model_cfg, config) + # Apply generation settings + if config["generation"][f"backend"] == "megatron": + _apply_cuda_graph_and_rng_tracker_config(model_cfg, config) + # Validate optimizer configuration _validate_optimizer_config(config) @@ -367,6 +371,14 @@ def _apply_parallelism_config(model_cfg: Any, config: PolicyConfig) -> None: ) +def _apply_cuda_graph_and_rng_tracker_config(model_cfg: Any, config: PolicyConfig) -> None: + """Apply CUDA GRAPH and RNG TRACKER configuration.""" + model_cfg.cuda_graph_impl = config["megatron_cfg"]["cuda_graph_impl"] + model_cfg.cuda_graph_scope = config["megatron_cfg"]["cuda_graph_scope"] + model_cfg.use_te_rng_tracker = config["megatron_cfg"]["use_te_rng_tracker"] + model_cfg.inference_rng_tracker = config["megatron_cfg"]["inference_rng_tracker"] + + def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: """Apply Mixture of Experts configuration.""" model_cfg.expert_tensor_parallel_size = config["megatron_cfg"][ @@ -399,6 +411,9 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.moe_token_dispatcher_type = config["megatron_cfg"][ "moe_token_dispatcher_type" ] + model_cfg.moe_pad_experts_for_cuda_graph_inference = config["megatron_cfg"][ + "moe_pad_experts_for_cuda_graph_inference" + ] model_cfg.moe_shared_expert_overlap = config["megatron_cfg"][ "moe_shared_expert_overlap" ] @@ -850,6 +865,10 @@ def setup_reference_model_state( ref_ckpt_context = init_checkpointing_context(ref_checkpoint_config) + megatron_cfg.model.cuda_graph_impl = "none" + megatron_cfg.model.use_te_rng_tracker = False + megatron_cfg.model.inference_rng_tracker = False + # Create a separate megatron config for the reference model ref_megatron_cfg = ConfigContainer( model=megatron_cfg.model, diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 29f034b065..c3787692c1 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -620,7 +620,13 @@ def train( def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using the policy.""" + """Generate a batch of data using the policy. + + For coordinator-based inference (Megatron backend), all data is sent to DP rank 0 + only, which submits requests to the coordinator. The coordinator then distributes + work across all DP engines. Other DP ranks participate in the engine loop but + don't receive input data directly. + """ # Verify input data is right-padded assert isinstance(data, BatchedDataDict), ( f"data must be a BatchedDataDict, got type: {type(data)}" @@ -629,14 +635,38 @@ def generate( "Missing required input fields" ) - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data = data.shard_by_batch_size(dp_size, batch_size=None) + if self.cfg["generation"]['backend'] == "vllm": + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + data = data.shard_by_batch_size(dp_size, batch_size=None) + in_sharded_axes = ["data_parallel"] + output_is_replicated = [ + "tensor_parallel", + "pipeline_parallel", + ] + elif self.cfg["generation"]['backend'] == "megatron": + # For coordinator-based inference: send ALL data to DP rank 0 only. + # Other DP ranks are called with data=None but still participate in the + # inference engine loop. The coordinator handles load balancing across DP ranks. + # + # With in_sharded_axes=[] and data_parallel not in replicate_on_axes, + # data_parallel becomes a "free axis". Only workers at DP coord 0 receive data, + # while workers at other DP coords get None (via make_dummy_calls_to_free_axes). + in_sharded_axes = [] + output_is_replicated = [ + "data_parallel", + "tensor_parallel", + "pipeline_parallel", + ] + else: + raise ValueError(f"Invalid generation backend: {self.cfg['generation']['backend']}, expected 'vllm' or 'megatron'") + futures = self.worker_group.run_all_workers_sharded_data( "generate", - data=sharded_data, - in_sharded_axes=["data_parallel"], + data=data, # Full data goes to DP=0 only (free axis behavior) + in_sharded_axes=in_sharded_axes, replicate_on_axes=["tensor_parallel", "pipeline_parallel"], - output_is_replicated=["tensor_parallel", "pipeline_parallel"], + output_is_replicated=output_is_replicated, + make_dummy_calls_to_free_axes=True, # Call all DP ranks, but only DP=0 gets data common_kwargs={"greedy": greedy}, ) assert self.cfg["generation"] is not None, "Generation config is not set" diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 5f1483ed9a..c47beb19d1 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -11,9 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import gc import os import re +import time import warnings from collections import defaultdict from contextlib import AbstractContextManager, contextmanager, nullcontext @@ -31,17 +33,13 @@ reduce_max_stat_across_model_parallel_group, ) from megatron.bridge.utils.common_utils import get_rank_safe +from megatron.core.transformer.utils import toggle_cuda_graphs from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( FullyShardedDataParallel as custom_FSDP, ) -from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( - InferenceWrapperConfig, -) -from megatron.core.inference.text_generation_controllers.text_generation_controller import ( - TextGenerationController, -) +from megatron.core.inference.config import InferenceConfig, KVCacheManagementMode from megatron.core.optimizer import ChainedOptimizer from megatron.core.parallel_state import ( get_pipeline_model_parallel_group, @@ -664,53 +662,23 @@ def get_topk_logits( [{"topk_logits": topk_logits.cpu(), "topk_indices": topk_indices.cpu()}] ) - @wrap_with_nvtx_name("megatron_policy_worker/generate") - def generate( - self, *, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False - ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using huggingface framework generation. - - Args: - data: BatchedDataDict containing input_ids and input_lengths tensors - Returns: - BatchedDataDict conforming to GenerationOutputSpec: - - output_ids: input + generated token IDs - - logprobs: Log probabilities for each token - - generation_lengths: Lengths of each response - """ - # 512 bATCH SIZE (200 tokens) - no_grad = torch.no_grad() - no_grad.__enter__() - self.model.config.flash_decode = False - if self.should_disable_forward_pre_hook: - self.model = self.move_model( - self.model, "cuda", move_params=True, move_grads=False - ) - # Verify input is right padded - assert isinstance(data, BatchedDataDict), ( - f"data must be a BatchedDataDict, got type: {type(data)}" - ) - assert "input_ids" in data and "input_lengths" in data, ( - f"input_ids and input_lengths must be present in the BatchedDataDict, got keys: {data.keys()}" + def _get_lang_module(self): + """Get the underlying language module from the wrapped model.""" + return ( + self.model.module.module + if hasattr(self.model.module, "module") + else self.model.module ) - is_right_padded, error_msg = verify_right_padding( - data, pad_value=self.tokenizer.pad_token_id - ) - if not is_right_padded: - warnings.warn( - f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" - ) - model_cfg = self.megatron_cfg.model - inference_wrapper_config = InferenceWrapperConfig( - hidden_size=model_cfg.hidden_size, - inference_batch_times_seqlen_threshold=1000000, - fp32_residual_connection=model_cfg.fp32_residual_connection, - params_dtype=model_cfg.params_dtype, - padded_vocab_size=self.final_padded_vocab_size, # Use the potentially updated value - inference_max_seq_length=self.cfg["generation"]["max_new_tokens"], # type: ignore - inference_max_requests=self.cfg["generation_batch_size"], - ) + def _initialize_inference_engine(self, mcore_generation_config: dict): + """Initialize the persistent inference engine and client. + + This method sets up the DynamicInferenceEngine, DynamicInferenceContext, + and InferenceClient for coordinator-based inference. The engine is created + once and reused across multiple generate() calls. + """ + if self._inference_engine_initialized: + return from megatron.core.inference.contexts.dynamic_context import ( DynamicInferenceContext, @@ -721,139 +689,385 @@ def generate( from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( GPTInferenceWrapper, ) - from megatron.core.inference.sampling_params import SamplingParams - - mcore_generation_config = cast( - MegatronGenerationConfig, self.cfg["generation"]["mcore_generation_config"] + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, ) - buffer_size_gb = mcore_generation_config["buffer_size_gb"] + model_cfg = self.megatron_cfg.model + + + from megatron.core.utils import get_attr_wrapped_model + pg_collection = get_attr_wrapped_model(self.model, "pg_collection") + + buffer_size_gb = mcore_generation_config["buffer_size_gb"] num_cuda_graphs = mcore_generation_config["num_cuda_graphs"] block_size_tokens = mcore_generation_config["block_size_tokens"] + enable_chunked_prefill = mcore_generation_config["enable_chunked_prefill"] use_cuda_graphs_for_non_decode_steps = mcore_generation_config[ "use_cuda_graphs_for_non_decode_steps" ] - enable_chunked_prefill = mcore_generation_config["enable_chunked_prefill"] - unified_memory_level = mcore_generation_config["unified_memory_level"] max_tokens = mcore_generation_config["max_tokens"] + # Level 0: No unified memory, CUDA graphs are deleted/recreated on pause/resume + # Level 1: Unified memory enabled, tensors maintain static addresses + unified_memory_level = mcore_generation_config["unified_memory_level"] + kv_cache_management_mode = mcore_generation_config["kv_cache_management_mode"] + static_kv_memory_pointers = mcore_generation_config["static_kv_memory_pointers"] + materialize_only_last_token_logits = mcore_generation_config["materialize_only_last_token_logits"] + model_config = self.model.config - model_config.cuda_graph_impl = "local" - dynamic_context = DynamicInferenceContext( - params_dtype=inference_wrapper_config.params_dtype, - num_layers=model_config.num_layers, - kv_channels=model_config.kv_channels, - num_attention_heads=model_config.num_query_groups, - max_sequence_length=self.cfg["generation"]["max_new_tokens"], + inference_config = InferenceConfig( + block_size_tokens=block_size_tokens, buffer_size_gb=buffer_size_gb, - materialize_only_last_token_logits=False, num_cuda_graphs=num_cuda_graphs, - block_size_tokens=block_size_tokens, - tensor_model_parallel_size=self.cfg["megatron_cfg"][ - "tensor_model_parallel_size" - ], - use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, - use_flashinfer_fused_rope=False, - unified_memory_level=unified_memory_level, max_tokens=max_tokens, - ) - inference_wrapped_model = GPTInferenceWrapper( - self.model, inference_wrapper_config, dynamic_context - ) + max_sequence_length=self.cfg["generation"]["max_new_tokens"], + unified_memory_level=unified_memory_level, + kv_cache_management_mode=KVCacheManagementMode(kv_cache_management_mode), + static_kv_memory_pointers=static_kv_memory_pointers, + use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, + materialize_only_last_token_logits=materialize_only_last_token_logits, + enable_chunked_prefill=enable_chunked_prefill, + pg_collection=pg_collection, + ) - inference_wrapped_model.prep_model_for_inference() - # Set pipeline parallel flag - inference_wrapped_model.model_is_pipeline_parallel = ( - self.cfg["megatron_cfg"]["pipeline_model_parallel_size"] > 1 - ) + # Create inference context + self.inference_context = DynamicInferenceContext(model_config, inference_config) + # Create inference wrapper + self.inference_wrapped_model = GPTInferenceWrapper( + self.model, self.inference_context + ) + # Create text generation controller text_generation_controller = TextGenerationController( - inference_wrapped_model=inference_wrapped_model, + inference_wrapped_model=self.inference_wrapped_model, tokenizer=self.megatron_tokenizer, ) - # Calculate seed based on node and rank to ensure reproducibility across workers - local_rank = torch.cuda.current_device() # Local GPU index on the node - num_gpus_per_node = torch.cuda.device_count() - node_idx = self.rank // num_gpus_per_node if num_gpus_per_node > 0 else 0 - seed = (node_idx * 1024) + local_rank - - # New API: DynamicInferenceEngine has additional parameters - dynamic_engine = DynamicInferenceEngine( + # Create the inference engine + self.dynamic_inference_engine = DynamicInferenceEngine( text_generation_controller, - dynamic_context, - enable_cuda_graph=True, - random_seed=seed, - track_paused_request_events=False, - enable_chunked_prefill=enable_chunked_prefill, - inference_logging_step_interval=0, + self.inference_context ) - # Handle None values for top_k - convert to integer as required by Megatron - top_k_cfg = self.cfg["generation"]["top_k"] - top_k_val = 1 if greedy else (int(top_k_cfg) if top_k_cfg is not None else 0) + self._inference_engine_initialized = True + self._inference_engine_alseep = True # Engine starts in paused state + print(f"[Rank {self.rank}] Initialized persistent inference engine") - top_p_cfg = self.cfg["generation"]["top_p"] - top_p_val = ( - 0.0 if greedy else (float(top_p_cfg) if top_p_cfg is not None else 0.0) + async def _start_inference_coordinator(self, coordinator_port: int): + """Start the inference coordinator and engine loop. + + This is called once when the inference infrastructure is first needed. + The engine's start_listening_to_data_parallel_coordinator returns the + actual coordinator address (dp_addr) which is used to create the client. + """ + dp_addr = await self.dynamic_inference_engine.start_listening_to_data_parallel_coordinator( + inference_coordinator_port=coordinator_port, + launch_inference_coordinator=True, ) - # New API: SamplingParams now includes termination_id and uses num_tokens_total - sampling_params = SamplingParams( - temperature=self.cfg["generation"]["temperature"] if not greedy else 0, - top_k=top_k_val, - top_p=top_p_val, - skip_prompt_log_probs=False, - return_log_probs=True, - num_tokens_total=self.cfg["generation"]["max_new_tokens"], - num_tokens_to_generate=None, - termination_id=self.megatron_tokenizer.eod, + dist_rank = torch.distributed.get_rank() + if dist_rank == 0: + from megatron.core.inference.inference_client import InferenceClient + self.inference_client = InferenceClient(inference_coordinator_address=dp_addr) + await self.inference_client.start() + + self._inference_engine_alseep = False + + def _sleep(self): + """pause the inference engine to free GPU memory for training. + + This method should be called before training to: + 1. Deallocate KV cache and other inference-specific GPU memory + 2. Disable CUDA graphs for inference + 3. Toggle model configuration for training mode + + Uses the coordinator's pause mechanism to properly pause the engine loop + and then pause the engine (deallocate tensors, etc.). + + For coordinator-based inference: + - Only rank 0 sends pause signals via the coordinator + - The coordinator broadcasts to all DP engines + - Non-rank-0 workers wait for their engine to be paused via the event loop + """ + + future = asyncio.run_coroutine_threadsafe( + self._sleep_engine(), + self._inference_loop ) + future.result() + # Synchronize all ranks + torch.distributed.barrier() + + self._inference_engine_alseep = True + print(f"[Rank {self.rank}] paused inference engine") + + async def _sleep_engine(self): + """Send suspend signals via the coordinator and wait for acknowledgment. + + Mirrors MegatronLocal.suspend() from megatron/rl/inference/megatron.py: + 1. Rank 0 sends suspend (PAUSE + SUSPEND) to coordinator + 2. All ranks wait for engine to be paused + 3. All ranks call engine.suspend() to deallocate GPU state + """ + if torch.distributed.get_rank() == 0: + # Send PAUSE signals + self.inference_client.suspend_engines() + # Wait for the engine to acknowledge the pause + await self.dynamic_inference_engine.paused.wait() + self.dynamic_inference_engine.suspend() + + def _wake(self): + """Resume the inference engine after training. + + This method should be called before generation to: + 1. Reallocate KV cache and inference-specific GPU memory + 2. Enable CUDA graphs for inference + 3. Toggle model configuration for inference mode + + Uses the coordinator's resume mechanism to properly resume the engine loop. + + For coordinator-based inference: + - Only rank 0 sends resume signals via the coordinator + - The coordinator broadcasts to all DP engines + - Non-rank-0 workers wait for their engine to be running via the event loop + """ + + # Use the coordinator-based resume mechanism + # Only rank 0 sends the signal - coordinator broadcasts to all DP engines + future = asyncio.run_coroutine_threadsafe( + self._wake_engine(), + self._inference_loop + ) + future.result() + # Synchronize all ranks + torch.distributed.barrier() + + self._inference_engine_alseep = False + print(f"[Rank {self.rank}] Resumed inference engine") + + async def _wake_engine(self): + """Send resume signals via the coordinator and wait for acknowledgment. + + Mirrors MegatronLocal.resume() from megatron/rl/inference/megatron.py: + 1. Rank 0 sends resume (RESUME + UNPAUSE) to coordinator + 2. All ranks wait for engine to be running + 3. All ranks call engine.resume() to reallocate GPU state + """ + if torch.distributed.get_rank() == 0: + self.inference_client.resume_engines() + await self.dynamic_inference_engine.running.wait() + self.dynamic_inference_engine.resume() - input_ids = data["input_ids"] - prompt_tokens_tensor = input_ids.cuda() - prompt_lengths_tensor = data["input_lengths"] - request_id = 0 + @contextmanager + def inference_mode(self, mcore_generation_config: dict): + """Context manager for inference mode, following Megatron RL's pattern. + + This mirrors megatron_rl_inference_mode from megatron/rl/rl_utils.py + + ENTER order: + 1. Put model in eval mode + 2. Clear rotary cache + 3. Toggle CUDA graphs ON + 4. Initialize inference engine (first time only) + 5. Resume engine (reallocates KV cache, recreates CUDA graphs as needed) + + EXIT order: + 1. Suspend engine (deallocates KV cache and GPU state) + 2. Toggle CUDA graphs OFF + 3. Clear rotary cache + 4. Put model back in train mode + + KV cache lifecycle is managed by the engine's suspend/resume mechanism + via KVCacheManagementMode in InferenceConfig. + + Yields: + The dynamic inference engine for use during inference. + """ + # Get the language module (unwrap from precision wrappers if needed) + lang_module = self._get_lang_module() + + # Get config settings + cuda_graph_impl = mcore_generation_config.get("cuda_graph_impl", "local") + + # Save training state + was_training = lang_module.training + + # === ENTER INFERENCE MODE === + + # 1. Put model in eval mode + lang_module.eval() + + # 2. Clear rotary position embedding caches (Megatron RL does this) + rotary_module = getattr(lang_module, "rotary_pos_emb", None) + has_lru_cache = rotary_module is not None and hasattr(rotary_module.forward, "cache_parameters") + if has_lru_cache: + rotary_module.forward.cache_clear() + + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) + + # 4. Initialize inference engine if not already done + if not self._inference_engine_initialized: + self._initialize_inference_engine(mcore_generation_config) + # Start the coordinator and engine loop (first time only) + coordinator_port = self.cfg["generation"].get( + "inference_coordinator_port", 5995 + ) + self._run_async_coordinator_start(coordinator_port) - # New API: add_request now takes sampling_params as a parameter - for p, prompt_len in zip( - prompt_tokens_tensor, prompt_lengths_tensor, strict=True - ): - dynamic_engine.add_request( - request_id, - p[:prompt_len], - sampling_params=sampling_params, + if self._inference_engine_alseep: + self._wake() + + try: + # Yield the inference engine for use + yield self.dynamic_inference_engine + + finally: + + # 1. pause the inference engine + if self._inference_engine_initialized and not self._inference_engine_alseep: + self._sleep() + + # 2. Toggle CUDA graphs OFF + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to="none") + + # 4. Clear rotary embedding cache again (Megatron RL does this on exit too) + if has_lru_cache: + rotary_module.forward.cache_clear() + + # 5. Restore training state + if was_training: + lang_module.train() + + # 6. Force garbage collection and CUDA memory cleanup + gc.collect() + torch.cuda.empty_cache() + + + @wrap_with_nvtx_name("megatron_policy_worker/generate") + def generate( + self, *, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False + ) -> BatchedDataDict[GenerationOutputSpec]: + """Generate a batch of data using Megatron Core inference with coordinator. + + This method uses the coordinator-based inference pattern from Megatron Core, + which enables better parallelism across data-parallel ranks through a central + coordinator that routes requests to available engines. + + The inference engine is created once and reused across generate() calls. + The engine is paused between generate() calls to free GPU memory for training. + + For coordinator-based inference: + - Only DP rank 0 receives actual data and submits requests to the coordinator + - Other DP ranks receive data=None but still participate in the inference engine loop + - The coordinator distributes work across all DP engines + - Results are broadcast from rank 0 to all ranks + + Args: + data: BatchedDataDict containing input_ids and input_lengths tensors, + or None for non-DP-0 workers (they participate in engine loop only) + BatchedDataDict conforming to GenerationOutputSpec: + - output_ids: input + generated token IDs + - logprobs: Log probabilities for each token + - generation_lengths: Lengths of each response + """ + no_grad = torch.no_grad() + no_grad.__enter__() + from megatron.core.inference.sampling_params import SamplingParams + + self.model.config.flash_decode = False + if self.should_disable_forward_pre_hook: + self.model = self.move_model( + self.model, "cuda", move_params=True, move_grads=False ) - request_id += 1 - - result = [] - while dynamic_engine.has_unfinished_requests(): - result_step = dynamic_engine.step_modern() - result.extend(result_step["finished_request_records"]) - - # Sort results by request_id to maintain original batch order - result.sort(key=lambda x: x.request_id) - - out = { - "tokens": [ - x.requests[0].prompt_tokens.tolist() + x.requests[0].generated_tokens - for x in result - ], - "logprobs": [ - x.requests[0].prompt_log_probs + x.requests[0].generated_log_probs - for x in result - ], - } + + dist_rank = torch.distributed.get_rank() + is_request_submitter = (dist_rank == 0) + + # For non-rank-0 workers, data may be None (they participate in engine loop only) + if data is not None: + # Verify input is right padded + assert isinstance(data, BatchedDataDict), ( + f"data must be a BatchedDataDict, got type: {type(data)}" + ) + is_right_padded, error_msg = verify_right_padding( + data, pad_value=self.tokenizer.pad_token_id + ) + if not is_right_padded: + warnings.warn( + f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" + ) + + + mcore_generation_config = self.cfg["generation"]["mcore_generation_config"] + # Use inference_mode context manager (mirrors megatron_rl_inference_mode from Megatron RL) + # This handles: eval mode, CUDA graph toggle, engine init/resume, and cleanup + with torch.no_grad(), self.inference_mode(mcore_generation_config) as inference_engine: + # Handle None values for top_k - convert to integer as required by Megatron + top_k_cfg = self.cfg["generation"]["top_k"] + top_k_val = 1 if greedy else (int(top_k_cfg) if top_k_cfg is not None else 0) + + top_p_cfg = self.cfg["generation"]["top_p"] + top_p_val = ( + 0.0 if greedy else (float(top_p_cfg) if top_p_cfg is not None else 0.0) + ) + + sampling_params = SamplingParams( + temperature=self.cfg["generation"]["temperature"] if not greedy else 0, + top_k=top_k_val, + top_p=top_p_val, + skip_prompt_log_probs=False, + return_log_probs=True, + num_tokens_total=self.cfg["generation"]["max_new_tokens"], + num_tokens_to_generate=None, + termination_id=self.megatron_tokenizer.eod, + ) + + # Only rank 0 has actual data to submit + if is_request_submitter: + input_ids = data["input_ids"] + print(f"[Rank {dist_rank}] input_ids: {input_ids.shape}") + prompt_tokens_tensor = input_ids.cuda() + prompt_lengths_tensor = data["input_lengths"] + else: + print(f"[Rank {dist_rank}] Participating in engine loop (no data to submit)") + prompt_tokens_tensor = torch.empty(0, dtype=torch.long, device="cuda") + prompt_lengths_tensor = torch.empty(0, dtype=torch.long, device="cuda") + + # Run the coordinator-based generation using the persistent engine + # Rank 0 submits requests, other ranks participate in engine loop + # Results are broadcast to all ranks inside this method + result = self._run_async_generation_with_persistent_engine( + prompt_tokens_tensor, + prompt_lengths_tensor, + sampling_params, + ) + + self.model.config.flash_decode = False + + # Context manager has exited - CUDA graphs are now disabled, model is back in train mode + + # Only rank 0 needs to format and return results + # Other ranks return None (their results are ignored due to output_is_replicated) + if not is_request_submitter: + # Return empty result for non-submitter ranks + # Use BatchedDataDict directly instead of from_batches to avoid padding issues with empty tensors + return BatchedDataDict({ + "output_ids": torch.empty(0, 0, dtype=torch.long), + "logprobs": torch.empty(0, 0, dtype=torch.float), + "generation_lengths": torch.empty(0, dtype=torch.long), + "unpadded_sequence_lengths": torch.empty(0, dtype=torch.long), + }).to("cpu") input_lengths = data["input_lengths"] - # pad the out "tokens" and "logprobs" and make them into tensors from lists batch_size = data["input_ids"].size(0) - max_gen_seq_len = max([len(x.requests[0].generated_tokens) for x in result]) + max_gen_seq_len = max([len(x.generated_tokens) for x in result]) padded_input_length = input_ids.size(1) max_seq_len = padded_input_length + max_gen_seq_len - # Create padded tensors for tokens and logprobs output_ids_padded = torch.full( (batch_size, max_seq_len), self.tokenizer.pad_token_id, @@ -867,7 +1081,6 @@ def generate( device=data["input_ids"].device, ) - # Fill in the padded tensors with actual values generation_lengths = torch.zeros( batch_size, dtype=torch.long, device=data["input_ids"].device ) @@ -875,15 +1088,17 @@ def generate( batch_size, dtype=torch.long, device=data["input_ids"].device ) for i in range(batch_size): - seq_len = len(out["tokens"][i]) + tokens = result[i].prompt_tokens.tolist() + result[i].generated_tokens + logprobs = result[i].prompt_log_probs + result[i].generated_log_probs + seq_len = len(tokens) output_ids_padded[i, :seq_len] = torch.tensor( - out["tokens"][i], dtype=torch.long, device=data["input_ids"].device + tokens, dtype=torch.long, device=data["input_ids"].device ) generation_lengths[i] = seq_len - input_lengths[i].item() unpadded_sequence_lengths[i] = seq_len - logprob_len = len(out["logprobs"][i]) + logprob_len = len(logprobs) logprobs_padded[i, 1 : logprob_len + 1] = torch.tensor( - out["logprobs"][i], + logprobs, dtype=torch.float, device=data["input_ids"].device, ) @@ -895,11 +1110,149 @@ def generate( "unpadded_sequence_lengths": unpadded_sequence_lengths, } - self.model.config.flash_decode = False no_grad.__exit__(None, None, None) return BatchedDataDict.from_batches([out_dict]).to("cpu") + def _start_inference_loop_thread(self): + """Start a background thread with a persistent event loop for inference. + + This thread runs the event loop that hosts the engine loop task. + The loop runs forever until explicitly stopped. + """ + import threading + + def run_loop(): + asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) + self._inference_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._inference_loop) + # Run forever - the engine loop task will run in this loop + self._inference_loop.run_forever() + + self._inference_thread = threading.Thread(target=run_loop, daemon=True) + self._inference_thread.start() + + # Wait for the loop to be created + while self._inference_loop is None: + time.sleep(0.001) + + def _run_async_coordinator_start(self, coordinator_port: int): + """Start the coordinator and engine loop in the background thread. + + This is called once during the first generate() call to initialize + the persistent inference infrastructure. + """ + import concurrent.futures + + # Start the background thread with the event loop if not already running + if self._inference_loop is None: + self._start_inference_loop_thread() + + # Schedule the coordinator start in the inference loop + future = asyncio.run_coroutine_threadsafe( + self._start_inference_coordinator(coordinator_port), + self._inference_loop + ) + # Wait for completion + return future.result() + + def _run_async_generation_with_persistent_engine( + self, + prompt_tokens_tensor: torch.Tensor, + prompt_lengths_tensor: torch.Tensor, + sampling_params: "SamplingParams", + ) -> list: + """Run generation using the persistent inference engine. + + This method uses the pre-initialized engine and client to run generation. + Unlike the original method, it doesn't start/stop the coordinator each time. + The async operation runs in the persistent inference loop. + """ + if self._inference_loop is None: + raise RuntimeError("Inference loop not initialized. Call generate() first.") + + # Schedule the generation in the inference loop + future = asyncio.run_coroutine_threadsafe( + self._generate_with_persistent_engine( + prompt_tokens_tensor, + prompt_lengths_tensor, + sampling_params, + ), + self._inference_loop + ) + # Wait for completion and return the result + return future.result() + + async def _generate_with_persistent_engine( + self, + prompt_tokens_tensor: torch.Tensor, + prompt_lengths_tensor: torch.Tensor, + sampling_params: "SamplingParams", + ) -> list: + """Run generation using the persistent coordinator-based inference. + + This method uses the already-running engine and submits requests through + the persistent client. The engine loop continues running between calls. + + For coordinator-based inference with centralized request submission: + - Only rank 0 (the request submitter) submits requests and collects results + - Other ranks return early but their engine loops continue running in the + background, processing requests distributed by the coordinator + - No broadcast is needed since only rank 0's results are used by the caller + + Args: + prompt_tokens_tensor: Tensor of prompt token IDs [batch_size, seq_len] + prompt_lengths_tensor: Tensor of prompt lengths [batch_size] + sampling_params: Sampling parameters for generation + + Returns: + List of completed request records sorted by request_id (rank 0), + or empty list (other ranks) + """ + from megatron.core.inference.inference_request import DynamicInferenceRequestRecord + + dist_rank = torch.distributed.get_rank() + + if dist_rank == 0: + assert self.inference_client is not None, "Inference client not initialized" + + # Non-rank-0 workers: return immediately with empty results + # Their engine loops will continue processing requests from the coordinator + # in the background (the engine loop runs as a separate task in _inference_loop) + if dist_rank != 0: + print(f"[Rank {dist_rank}] Participating in engine loop only (not submitting requests)") + # Return empty results - the caller only uses rank 0's results + return [] + + # Rank 0: submit ALL requests and collect results + print(f"[Rank {dist_rank}] Submitting {prompt_tokens_tensor.size(0)} requests to coordinator") + + futures = [] + for request_id, (prompt_tokens, prompt_len) in enumerate( + zip(prompt_tokens_tensor, prompt_lengths_tensor, strict=True) + ): + # Extract the actual prompt tokens (without padding) and convert to list + prompt = prompt_tokens[: prompt_len.item()].tolist() + future = self.inference_client.add_request(prompt, sampling_params) + futures.append(future) + + # Wait for all requests to complete + # The coordinator distributes work to all DP engines, including this one + completed_records: list[DynamicInferenceRequestRecord] = await asyncio.gather( + *futures + ) + + # Extract the merged request from each record + results = [record.merge() for record in completed_records] + + # Sort by request_id to maintain original batch order + results.sort(key=lambda x: x.request_id) + + print(f"[Rank {dist_rank}] Completed {len(results)} requests") + + return results + + @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/prepare_refit_info") def prepare_refit_info(self) -> None: @@ -1056,6 +1409,71 @@ def broadcast_weights_for_collective( post_iter_func=lambda x: x[1], ) + @torch.no_grad() + def update_weights_from_collective(self) -> bool: + """Receive updated weights from collective communication (inference side). + + This method is the consumer counterpart of broadcast_weights_for_collective. + It receives weights broadcast by the training workers and updates the local + model parameters. + + TODO: Implement the actual weight update logic using packed_broadcast_consumer. + The implementation should: + 1. Iterate over the stored state_dict_info + 2. Use packed_broadcast_consumer to receive weights from the training side + 3. Update the local Megatron model parameters with the received weights + + Returns: + bool: True if weights were successfully updated. + """ + raise NotImplementedError( + "update_weights_from_collective for MegatronPolicyWorker is not yet implemented. " + "This placeholder will be replaced with actual NCCL collective weight reception logic." + ) + + def init_collective_as_inference( + self, ip: str, port: int, world_size: int, *, train_world_size: int + ) -> None: + """Initialize collective communication for inference-side workers. + + Unlike the base init_collective (used by training workers which use + self.rank directly), this method offsets the rank by train_world_size + so inference workers get globally unique ranks that don't collide + with training workers. + + Args: + ip: IP address for the process group rendezvous. + port: Port for the process group rendezvous. + world_size: Total world size (train + inference workers). + train_world_size: Number of training workers (used to offset ranks). + """ + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + # Offset rank by train_world_size so inference workers get unique global ranks + rank = train_world_size + self.rank + pg = StatelessProcessGroup.create( + host=ip, port=port, rank=rank, world_size=world_size + ) + device = torch.cuda.current_device() + self.model_update_group = PyNcclCommunicator(pg, device=device) + + def store_refit_info(self, state_dict_info: dict[str, Any]) -> None: + """Store state dict metadata for weight refitting on the inference side. + + This is the inference-side counterpart of prepare_refit_info(). Instead of + calculating the metadata from the model, it accepts pre-computed metadata + from the training side. + + TODO: Implement proper storage and use of state_dict_info for + update_weights_from_collective. + + Args: + state_dict_info: Dictionary mapping tensor names to (shape, dtype) tuples, + as returned by the training-side prepare_refit_info(). + """ + self.state_dict_info = state_dict_info + def prepare_for_lp_inference(self): self.model = self.move_model(self.model, "cuda", move_grads=False) self.model.eval() From 13a5c0aa434e15c15272e33f3a5dbbdc203d0d24 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Thu, 19 Feb 2026 15:35:00 -0800 Subject: [PATCH 05/23] Some fixes: --- nemo_rl/algorithms/grpo.py | 5 ---- .../policy/workers/megatron_policy_worker.py | 27 ++++++++++++++----- tests/functional/grpo_non_colocated.sh | 4 ++- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 43aa43ce77..25641805e8 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -396,11 +396,6 @@ def setup( ) else: - assert generation_config["backend"] != "megatron", ( - "Non-colocated inference is not supported for Megatron generation backends. " - "Please use vLLM backend for generation." - ) - # train resources will be updated through overall and inference resources below train_gpus_per_node = cluster_config["gpus_per_node"] train_nodes = policy_nodes diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c47beb19d1..5d491c973b 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -230,6 +230,15 @@ def __init__( ## used for streaming update inference engine weights self._held_gather_buffer = None + + self.dynamic_inference_engine = None + self.inference_client = None + self.inference_context = None + self.inference_wrapped_model = None + self._inference_engine_initialized = False + self._inference_engine_alseep = True # Start paused since we begin with training + self._inference_loop = None # Event loop for inference operations + self._inference_thread = None # Thread running the event loop def enable_forward_pre_hook(self): assert isinstance(self.model, DistributedDataParallel) @@ -237,7 +246,14 @@ def enable_forward_pre_hook(self): def disable_forward_pre_hook(self, param_sync=True): assert isinstance(self.model, DistributedDataParallel) - self.model.disable_forward_pre_hook(param_sync=param_sync) + for module, handle in list(self.model.remove_forward_pre_hook_handles.items()): + handle.remove() + self.model.remove_forward_pre_hook_handles.clear() + if param_sync: + self.model.start_param_sync(force_sync=True) + + # TODO : Check why this doesnt work. + #self.model.disable_forward_pre_hook(param_sync=param_sync) @wrap_with_nvtx_name("megatron_policy_worker/train") def train( @@ -1447,16 +1463,15 @@ def init_collective_as_inference( world_size: Total world size (train + inference workers). train_world_size: Number of training workers (used to offset ranks). """ - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator - from vllm.distributed.utils import StatelessProcessGroup + from nemo_rl.distributed.stateless_process_group import StatelessProcessGroup # Offset rank by train_world_size so inference workers get unique global ranks rank = train_world_size + self.rank - pg = StatelessProcessGroup.create( - host=ip, port=port, rank=rank, world_size=world_size + self.model_update_group = StatelessProcessGroup( + master_address=ip, port=port, rank=rank, world_size=world_size ) device = torch.cuda.current_device() - self.model_update_group = PyNcclCommunicator(pg, device=device) + self.model_update_group.init_nccl_communicator(device=device) def store_refit_info(self, state_dict_info: dict[str, Any]) -> None: """Store state dict metadata for weight refitting on the inference side. diff --git a/tests/functional/grpo_non_colocated.sh b/tests/functional/grpo_non_colocated.sh index 8c65aedda2..4570e50279 100755 --- a/tests/functional/grpo_non_colocated.sh +++ b/tests/functional/grpo_non_colocated.sh @@ -20,6 +20,7 @@ mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ + --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ grpo.num_generations_per_prompt=4 \ @@ -27,7 +28,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.train_micro_batch_size=1 \ policy.generation.colocated.enabled=false \ policy.generation.colocated.resources.gpus_per_node=1 \ - policy.generation.vllm_cfg.async_engine=true \ + policy.generation.backend=vllm \ + policy.generation.vllm_cfg.async_engine=false \ cluster.gpus_per_node=2 \ grpo.max_num_steps=2 \ logger.tensorboard_enabled=true \ From af270543fdb2b1cf9fc7a84063aa974cf02b81c8 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Thu, 19 Feb 2026 17:27:06 -0800 Subject: [PATCH 06/23] Fixes for offload --- examples/configs/grpo_math_1B_megatron.yaml | 3 ++- examples/configs/grpo_math_8B_megatron.yaml | 3 ++- pyproject.toml | 1 + uv.lock | 2 ++ 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 13d9634eb3..f9eebea735 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -130,6 +130,7 @@ policy: clip_grad: ${policy.max_grad_norm} scheduler: + override_opt_param_scheduler: true start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} weight_decay_incr_style: "constant" @@ -156,7 +157,7 @@ policy: num_cuda_graphs: 16 # Number of CUDA graphs to pre-compile for different batch sizes block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity) use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing - unified_memory_level: 1 # Unified memory usage level (0=disabled, 1+=enables unified memory ) + unified_memory_level: 0 # Unified memory usage level (0=disabled, 1+=enables unified memory ) max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens enable_chunked_prefill: false kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute" diff --git a/examples/configs/grpo_math_8B_megatron.yaml b/examples/configs/grpo_math_8B_megatron.yaml index 977ab394b5..94f14aaf2d 100644 --- a/examples/configs/grpo_math_8B_megatron.yaml +++ b/examples/configs/grpo_math_8B_megatron.yaml @@ -17,7 +17,7 @@ policy: train_global_batch_size: 512 train_micro_batch_size: 1 generation_batch_size: 32 # Only used when generating using HF backend - logprob_batch_size: 4 + logprob_batch_size: ${policy.train_micro_batch_size} max_total_sequence_length: 4096 precision: "bfloat16" @@ -48,6 +48,7 @@ policy: params_dtype: "float32" scheduler: + override_opt_param_scheduler: true start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} weight_decay_incr_style: "constant" diff --git a/pyproject.toml b/pyproject.toml index 7b702d2662..878657c725 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,7 @@ mcore = [ # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 "flash-attn==2.8.1", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", + "torch-memory-saver", ] nemo_gym = ["nemo_gym"] diff --git a/uv.lock b/uv.lock index 3a379f1303..5ce78f1a4c 100644 --- a/uv.lock +++ b/uv.lock @@ -4746,6 +4746,7 @@ mcore = [ { name = "flash-attn" }, { name = "megatron-bridge" }, { name = "megatron-core" }, + { name = "torch-memory-saver" }, { name = "transformer-engine", extra = ["pytorch"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] nemo-gym = [ @@ -4883,6 +4884,7 @@ requires-dist = [ { name = "tiktoken" }, { name = "torch", marker = "sys_platform != 'darwin'", specifier = "==2.9.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.9.0", index = "https://pypi.org/simple" }, + { name = "torch-memory-saver", marker = "extra == 'mcore'" }, { name = "torch-memory-saver", marker = "extra == 'sglang'" }, { name = "torchao", marker = "extra == 'sglang'" }, { name = "torchdata" }, From 6b4ad6f654fe01701bb6996e2e7eabe374d971b1 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Thu, 19 Feb 2026 22:20:36 -0800 Subject: [PATCH 07/23] Fixes for offload --- .../models/policy/workers/megatron_policy_worker.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 5d491c973b..c7d3e2feec 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -945,16 +945,16 @@ def inference_mode(self, mcore_generation_config: dict): # 1. pause the inference engine if self._inference_engine_initialized and not self._inference_engine_alseep: self._sleep() - + # 2. Toggle CUDA graphs OFF if cuda_graph_impl != "none": toggle_cuda_graphs(lang_module, set_to="none") - # 4. Clear rotary embedding cache again (Megatron RL does this on exit too) + # 3. Clear rotary embedding cache again (Megatron RL does this on exit too) if has_lru_cache: rotary_module.forward.cache_clear() - # 5. Restore training state + # 4. Restore training state if was_training: lang_module.train() @@ -990,8 +990,6 @@ def generate( - logprobs: Log probabilities for each token - generation_lengths: Lengths of each response """ - no_grad = torch.no_grad() - no_grad.__enter__() from megatron.core.inference.sampling_params import SamplingParams self.model.config.flash_decode = False @@ -1126,8 +1124,6 @@ def generate( "unpadded_sequence_lengths": unpadded_sequence_lengths, } - no_grad.__exit__(None, None, None) - return BatchedDataDict.from_batches([out_dict]).to("cpu") def _start_inference_loop_thread(self): From c19ee75d65a9b9b450d7c13eb9e78abecbd732cf Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Thu, 19 Feb 2026 11:44:03 +0200 Subject: [PATCH 08/23] chore: Switching mcore to upstream main Signed-off-by: Ahmad Kiswani --- .gitmodules | 4 +- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- 3rdparty/Megatron-LM-workspace/setup.py | 7 +- uv.lock | 76 ++++++++++++++++++---- 4 files changed, 73 insertions(+), 16 deletions(-) diff --git a/.gitmodules b/.gitmodules index c1b0c5a56f..8d7c7be7e5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM - url = https://github.com/yaoyu-33/Megatron-LM.git - branch = yifu/remove_do_not_average_loss + url = https://github.com/NVIDIA/Megatron-LM.git + branch = main shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index b12071b947..0d0943c6bf 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit b12071b947f9ee3c6616306662069fc4ca77be4c +Subproject commit 0d0943c6bfa9cbb30fcd62d40ce1792c4cb201e8 diff --git a/3rdparty/Megatron-LM-workspace/setup.py b/3rdparty/Megatron-LM-workspace/setup.py index fb0a7cf92e..380864cc25 100644 --- a/3rdparty/Megatron-LM-workspace/setup.py +++ b/3rdparty/Megatron-LM-workspace/setup.py @@ -43,7 +43,7 @@ # VCS dependencies use full "pkg @ git+URL@rev" format matching pyproject.toml [tool.uv.sources] CACHED_DEPENDENCIES = [ # Default dependencies from pyproject.toml - "torch", + "torch>=2.6.0", "numpy", "packaging>=24.2", # Dev dependencies from pyproject.toml @@ -58,7 +58,7 @@ "opentelemetry-api~=1.33.1", "mamba-ssm~=2.2", "causal-conv1d~=1.5", - "flash-linear-attention~=0.3.2", + "flash-linear-attention~=0.4.0", "nv-grouped-gemm~=1.1", "megatron-energon[av_decode]~=6.0", "av", @@ -69,6 +69,9 @@ "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0", "datasets", "fastapi~=0.50", + "flask[async]", + "hypercorn", + "openai", ] diff --git a/uv.lock b/uv.lock index e0c3cda97f..3a379f1303 100644 --- a/uv.lock +++ b/uv.lock @@ -555,6 +555,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ae/69/fe387b0f70ed608a363a90036e08ef8c1e844e5c98145502160661012dc0/apache_tvm_ffi-0.1.4-cp314-cp314t-win_amd64.whl", hash = "sha256:1bceda57240d03a3cf026334521c0595d097ab92b6d0df7485cbb37b2c056c27", size = 1794928, upload-time = "2025-11-30T07:21:25.234Z" }, ] +[[package]] +name = "asgiref" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/40/f03da1264ae8f7cfdbf9146542e5e7e8100a4c66ab48e791df9a03d3f6c0/asgiref-3.11.1.tar.gz", hash = "sha256:5f184dc43b7e763efe848065441eac62229c9f7b0475f41f80e207a114eda4ce", size = 38550, upload-time = "2026-02-03T13:30:14.33Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" }, +] + [[package]] name = "astor" version = "0.8.1" @@ -1924,16 +1933,16 @@ wheels = [ [[package]] name = "fla-core" -version = "0.3.2" +version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "einops" }, { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.9.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/67/c6/10a1149b07e6bab45b2cb2d07f6b827716c2baf5f3404161753f25c6389b/fla_core-0.3.2.tar.gz", hash = "sha256:d38db16bc4e1c6fa8c04df442f246da1e6926a209426bc6ef703d41bfbc37c92", size = 296725, upload-time = "2025-09-10T07:43:40.155Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/de/0d6bd5664ba2e711cabdde11ccb41ddcdd866c531e40900af3601bd7b8c6/fla_core-0.4.1.tar.gz", hash = "sha256:38ab28966eeadc2141b29e87c2bf72a8a4851e00af9d25bbbc3596b1fb53450d", size = 319608, upload-time = "2025-12-24T18:07:37.669Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/f5/74947b33c07682280e65adbdf17c4ee94b30232df2f728bafecf13d1d820/fla_core-0.3.2-py3-none-any.whl", hash = "sha256:e751d5a41e33eee721a6fb6588bd857f6f36e0d14719a23b1ebdbd617d307209", size = 413594, upload-time = "2025-09-10T07:43:37.786Z" }, + { url = "https://files.pythonhosted.org/packages/f6/43/945ef69eb48a14c30fd7323d3e0b560c821ae71e6d3ef979e06a901bc3b9/fla_core-0.4.1-py3-none-any.whl", hash = "sha256:93c6afe4c80fc7bc705fa8aeea6a46d2cf2d77383f9619a41863c7114c801bab", size = 437282, upload-time = "2025-12-24T18:07:34.41Z" }, ] [[package]] @@ -1952,17 +1961,15 @@ sdist = { url = "https://files.pythonhosted.org/packages/e8/6d/7066d160bdffa2f9d [[package]] name = "flash-linear-attention" -version = "0.3.2" +version = "0.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "datasets" }, { name = "fla-core" }, - { name = "pytest" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/84/f6/e62c1e562a288557eba7f06f168a7615813d1a227327b8beb8ba426da2c5/flash_linear_attention-0.3.2.tar.gz", hash = "sha256:9147747316c2951fed4ebeb4fa87977c05d807dc70c93b46250b68a6eb1183e2", size = 150880, upload-time = "2025-09-10T07:43:41.37Z" } +sdist = { url = "https://files.pythonhosted.org/packages/46/83/7d8ec7ffb5229080b1c9b772338ff588cbd63282ac355ede2a12a6e174a8/flash_linear_attention-0.4.1.tar.gz", hash = "sha256:127ee7273ed15ac17f72bcf4c75e1051719d8fbe0a2d1d047e59406f36d81ee2", size = 158280, upload-time = "2025-12-24T18:07:38.812Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/d0/35ce9eac5f52c72005095aaa12a393d2656ed7ffedf925b2381a6b76d10c/flash_linear_attention-0.3.2-py3-none-any.whl", hash = "sha256:604e73361437ba786420ab195e2caa3fd19280503761e703fa353c5ce5c65376", size = 274592, upload-time = "2025-09-10T07:43:39.107Z" }, + { url = "https://files.pythonhosted.org/packages/63/d5/6327559a9d5b9243b10c3984f1bcef256ed2ad06d105a3bb8f7b2979659c/flash_linear_attention-0.4.1-py3-none-any.whl", hash = "sha256:d18bdfe9d1f4b424676444eac9d50fb8433b70e5d4e0e0878b20bcbcdbea57ce", size = 287415, upload-time = "2025-12-24T18:07:35.815Z" }, ] [[package]] @@ -2141,6 +2148,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/f9/7f9263c5695f4bd0023734af91bedb2ff8209e8de6ead162f35d8dc762fd/flask-3.1.2-py3-none-any.whl", hash = "sha256:ca1d8112ec8a6158cc29ea4858963350011b5c846a414cdb7a954aa9e967d03c", size = 103308, upload-time = "2025-08-19T21:03:19.499Z" }, ] +[package.optional-dependencies] +async = [ + { name = "asgiref" }, +] + [[package]] name = "flask-cors" version = "6.0.1" @@ -2851,6 +2863,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, ] +[[package]] +name = "hypercorn" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, + { name = "h2" }, + { name = "priority" }, + { name = "wsproto" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/44/01/39f41a014b83dd5c795217362f2ca9071cf243e6a75bdcd6cd5b944658cc/hypercorn-0.18.0.tar.gz", hash = "sha256:d63267548939c46b0247dc8e5b45a9947590e35e64ee73a23c074aa3cf88e9da", size = 68420, upload-time = "2025-11-08T13:54:04.78Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/35/850277d1b17b206bd10874c8a9a3f52e059452fb49bb0d22cbb908f6038b/hypercorn-0.18.0-py3-none-any.whl", hash = "sha256:225e268f2c1c2f28f6d8f6db8f40cb8c992963610c5725e13ccfcddccb24b1cd", size = 61640, upload-time = "2025-11-08T13:54:03.202Z" }, +] + [[package]] name = "hyperframe" version = "6.1.0" @@ -3652,6 +3679,8 @@ dependencies = [ { name = "flash-linear-attention" }, { name = "flashinfer-python", version = "0.5.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-vllm' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang')" }, { name = "flashinfer-python", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-sglang' or extra != 'extra-7-nemo-rl-vllm'" }, + { name = "flask", extra = ["async"] }, + { name = "hypercorn" }, { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"] }, { name = "multi-storage-client" }, @@ -3661,6 +3690,7 @@ dependencies = [ { name = "nvidia-resiliency-ext" }, { name = "nvtx" }, { name = "onnxscript" }, + { name = "openai" }, { name = "opentelemetry-api" }, { name = "packaging" }, { name = "tensorstore", version = "0.1.74", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, @@ -3680,8 +3710,10 @@ requires-dist = [ { name = "einops", specifier = "~=0.8" }, { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.1.0" }, { name = "fastapi", specifier = "~=0.50" }, - { name = "flash-linear-attention", specifier = "~=0.3.2" }, + { name = "flash-linear-attention", specifier = "~=0.4.0" }, { name = "flashinfer-python", specifier = "~=0.5.0" }, + { name = "flask", extras = ["async"] }, + { name = "hypercorn" }, { name = "mamba-ssm", git = "https://github.com/state-spaces/mamba.git?rev=d68d16ed7d5d5164eb5a57c0285f3b7eb8394ec1" }, { name = "megatron-energon", extras = ["av-decode"], specifier = "~=6.0" }, { name = "multi-storage-client", specifier = "~=0.27" }, @@ -3691,11 +3723,12 @@ requires-dist = [ { name = "nvidia-resiliency-ext" }, { name = "nvtx", specifier = "~=0.2" }, { name = "onnxscript" }, + { name = "openai" }, { name = "opentelemetry-api", specifier = "~=1.33.1" }, { name = "packaging", specifier = ">=24.2" }, { name = "tensorstore", specifier = "~=0.1,!=0.1.46,!=0.1.72" }, - { name = "torch", marker = "sys_platform != 'darwin'", index = "https://download.pytorch.org/whl/cu129" }, - { name = "torch", marker = "sys_platform == 'darwin'", index = "https://pypi.org/simple" }, + { name = "torch", marker = "sys_platform != 'darwin'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cu129" }, + { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.6.0", index = "https://pypi.org/simple" }, { name = "tqdm" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], specifier = ">=2.9.0a0,<2.12.0" }, { name = "wget" }, @@ -6132,6 +6165,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/c7/5613524e606ea1688b3bdbf48aa64bafb6d0a4ac3750274c43b6158a390f/prettytable-3.16.0-py3-none-any.whl", hash = "sha256:b5eccfabb82222f5aa46b798ff02a8452cf530a352c31bddfa29be41242863aa", size = 33863, upload-time = "2025-03-24T19:39:02.359Z" }, ] +[[package]] +name = "priority" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/3c/eb7c35f4dcede96fca1842dac5f4f5d15511aa4b52f3a961219e68ae9204/priority-2.0.0.tar.gz", hash = "sha256:c965d54f1b8d0d0b19479db3924c7c36cf672dbf2aec92d43fbdaf4492ba18c0", size = 24792, upload-time = "2021-06-27T10:15:05.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/5f/82c8074f7e84978129347c2c6ec8b6c59f3584ff1a20bc3c940a3e061790/priority-2.0.0-py3-none-any.whl", hash = "sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa", size = 8946, upload-time = "2021-06-27T10:15:03.856Z" }, +] + [[package]] name = "prometheus-client" version = "0.22.1" @@ -9861,6 +9903,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" }, ] +[[package]] +name = "wsproto" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" }, +] + [[package]] name = "xattr" version = "1.3.0" From b0bd9f956c3196a2bc3d8b57049cb7d4d5c7c7f6 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Fri, 20 Feb 2026 08:01:57 -0800 Subject: [PATCH 09/23] Point Megatron-LM submodule to fork with fixes --- .gitmodules | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitmodules b/.gitmodules index 8d7c7be7e5..724b5b4f31 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM - url = https://github.com/NVIDIA/Megatron-LM.git - branch = main + url = https://github.com/shanmugamr1992/Megatron-LM.git + branch = fixes shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge From fb3905c708a50700c40828772d34e8296bf996f3 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Fri, 20 Feb 2026 08:12:02 -0800 Subject: [PATCH 10/23] Update Megatron-LM submodule --- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index 0d0943c6bf..de895f8448 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit 0d0943c6bfa9cbb30fcd62d40ce1792c4cb201e8 +Subproject commit de895f84482eeda48cc3c0bbbb58d82da0a700fc From 924e32a696409fb095a9eea7eef5fc97cdea55c6 Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 02:30:36 +0200 Subject: [PATCH 11/23] TE version bump to match mcore Signed-off-by: Ahmad Kiswani --- pyproject.toml | 8 ++++---- uv.lock | 26 +++++++++++++------------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7b702d2662..7b1ced085c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,7 @@ automodel = [ "mamba-ssm", "causal-conv1d", "nv-grouped-gemm", - "transformer-engine[pytorch]==2.8.0", + "transformer-engine[pytorch]>=2.9.0a0,<2.12.0", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", ] vllm = [ @@ -108,7 +108,7 @@ mcore = [ # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml. # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml - "transformer-engine[pytorch]==2.8.0", + "transformer-engine[pytorch]>=2.9.0a0,<2.12.0", "megatron-core", "megatron-bridge", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) @@ -235,12 +235,12 @@ default-groups = ["dev", "build"] # --link-mode=copy (slower but more reliable; supresses warning) # --link-mode=symlink (fastest option when uv cache and venv on different file-system; caveat: venv is brittle since it depends on the environment/container) link-mode = "copy" -# The TE override is needed because automodel/mbridge we are on is still on 2.5.0 +# The TE override is needed because automodel/mbridge we are on is still on an older version # The opencv-python-headless override is needed because automodel pins it to 4.10.0.84, whereas vllm>=0.11.0 needs >= 4.11.0 # The timm override is needed because current automodel pins to 1.0.16. This can be removed once we move ToT automodel # The nvidia-modelopt override is needed because mcore is still on 0.33 override-dependencies = [ - "transformer-engine[pytorch]==2.8.0", + "transformer-engine[pytorch]>=2.9.0a0,<2.12.0", "opencv-python-headless>=4.11.0", "timm<=1.0.22", "nvidia-modelopt[torch]>=0.39.0", diff --git a/uv.lock b/uv.lock index 3a379f1303..8c04ceaac7 100644 --- a/uv.lock +++ b/uv.lock @@ -211,7 +211,7 @@ overrides = [ { name = "torch", marker = "sys_platform != 'darwin'", specifier = "==2.9.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.9.0", index = "https://pypi.org/simple" }, { name = "torchaudio", specifier = "==2.9.0" }, - { name = "transformer-engine", extras = ["pytorch"], specifier = "==2.8.0" }, + { name = "transformer-engine", extras = ["pytorch"], specifier = ">=2.9.0a0,<2.12.0" }, ] [[manifest.dependency-metadata]] @@ -4888,8 +4888,8 @@ requires-dist = [ { name = "torchdata" }, { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = ">=0.22.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = ">=0.22.0", index = "https://pypi.org/simple" }, - { name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'automodel'", specifier = "==2.8.0" }, - { name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'mcore'", specifier = "==2.8.0" }, + { name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'automodel'", specifier = ">=2.9.0a0,<2.12.0" }, + { name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'mcore'", specifier = ">=2.9.0a0,<2.12.0" }, { name = "transformers", specifier = "==4.57.1" }, { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", index = "https://download.pytorch.org/whl/cu129" }, { name = "uvloop", marker = "extra == 'sglang'" }, @@ -9220,13 +9220,10 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.8.0" +version = "2.11.0" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "transformer-engine-cu12" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/72/be/a7cf5f28b7abbe966956217b18208fc34cc9bfaf62fa9472c0603db74899/transformer_engine-2.8.0-py3-none-any.whl", hash = "sha256:795b056d31b0f67f5d7432725177782dd5084090c1e3c52532577070947fe9a7", size = 638319, upload-time = "2025-10-07T04:55:34.115Z" }, + { url = "https://files.pythonhosted.org/packages/00/33/44571ec584c88e1715f4c2afefc0ddd45064c7065ac1c6ffc8e832bc3ba3/transformer_engine-2.11.0-py3-none-any.whl", hash = "sha256:7ee1eae8fa6b0cb471c6066aa3555304fda8537174e5019929dc0c8655071df3", size = 723110, upload-time = "2026-01-02T09:58:23.245Z" }, ] [package.optional-dependencies] @@ -9236,7 +9233,7 @@ pytorch = [ [[package]] name = "transformer-engine-cu12" -version = "2.8.0" +version = "2.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "importlib-metadata" }, @@ -9244,22 +9241,25 @@ dependencies = [ { name = "pydantic" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/53/db/cde3e772cf5cd7e941b64d37e4a61e2762f36ecc2e6508525af536076f8d/transformer_engine_cu12-2.8.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:3a92c781fc3c1a3a6a0009871a36903fa364b2d51ce06b06641d29aeefd59310", size = 480373707, upload-time = "2025-10-07T05:03:05.392Z" }, - { url = "https://files.pythonhosted.org/packages/b9/14/67860f2f1f9d0eca4a8e5e0cef5a0de5c4fc26340625051f032d16913d8c/transformer_engine_cu12-2.8.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:8ee5f9df586491a35fd1a01cb95b6970a9f01a5e8f935ecdacd56173c44d6a67", size = 480875025, upload-time = "2025-10-07T04:54:43.762Z" }, + { url = "https://files.pythonhosted.org/packages/05/27/5c4c27cb245a3513e5ad7ccef50e2e9688996e2cc558edbbb575dfcca276/transformer_engine_cu12-2.11.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ed5fda0925cb304d6864b451d8d012c579d5bd097bfefefca769b2704b06381a", size = 287630565, upload-time = "2026-01-02T09:56:43.645Z" }, + { url = "https://files.pythonhosted.org/packages/fa/a2/1439bbb6bc7d4d6045bad7d213884f7be92301c0982f009e3bbafa40e4ff/transformer_engine_cu12-2.11.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:6e5c0707583b2a90b2570da6f57409c6802653e069dfec38cf07a3b77ba9b12d", size = 288159349, upload-time = "2026-01-02T09:57:56.435Z" }, ] [[package]] name = "transformer-engine-torch" -version = "2.8.0" +version = "2.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "einops" }, { name = "onnx" }, { name = "onnxscript" }, + { name = "packaging" }, + { name = "pydantic" }, { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "torch", version = "2.9.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "transformer-engine-cu12" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/38/63/1e3953244ed4f318f87889309a56cdd664759f007967eb850ee415a5584d/transformer_engine_torch-2.8.0.tar.gz", hash = "sha256:ce09f1bd9b8e532a5c347b9e9b3a3a771722095daddca673ae82ccce8e68d759", size = 209805, upload-time = "2025-10-07T04:54:11.134Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/42/068a40f5b213a3a8899e3885eb178776662897abed03cd725953d1106c39/transformer_engine_torch-2.11.0.tar.gz", hash = "sha256:b58d6322bdf885dfab0646da572aff9cf090b332ad470559aa58883c231e1816", size = 242065, upload-time = "2026-01-02T09:58:58.423Z" } [[package]] name = "transformers" From 4ba7f0959097eee1ec8e2a82bef86d2aeae27bc2 Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 05:09:17 +0200 Subject: [PATCH 12/23] bumped mbridge Signed-off-by: Ahmad Kiswani --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 3rdparty/Megatron-Bridge-workspace/setup.py | 6 ++++-- uv.lock | 10 ++++++++-- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 15398e08fc..f91542b909 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 15398e08fc86be3de084c7382116527246ab1852 +Subproject commit f91542b90908ad08b7e13672feea03e27bedee27 diff --git a/3rdparty/Megatron-Bridge-workspace/setup.py b/3rdparty/Megatron-Bridge-workspace/setup.py index a0beea9449..d0bdd8711f 100644 --- a/3rdparty/Megatron-Bridge-workspace/setup.py +++ b/3rdparty/Megatron-Bridge-workspace/setup.py @@ -27,7 +27,7 @@ CACHED_DEPENDENCIES = [ "transformers<5.0.0", - "datasets", + "datasets>=2.20.0", "accelerate", "omegaconf>=2.3.0", "tensorboard>=2.19.0", @@ -41,13 +41,15 @@ "hydra-core>1.3,<=1.3.2", "megatron-core[dev,mlm]>=0.15.0a0,<0.17.0", "qwen-vl-utils", - "transformer-engine[pytorch]>=2.10.0a0,<2.12.0", + "transformer-engine[pytorch,core_cu13]>=2.10.0a0,<2.13.0", "mamba-ssm", "nvidia-resiliency-ext", "causal-conv1d", "flash-linear-attention", "timm", "open-clip-torch>=3.2.0", + "mlflow>=3.5.0", + "torch>=2.6.0", ] # If the bridge source exists, compare cached dependencies with the submodule's pyproject diff --git a/uv.lock b/uv.lock index 8c04ceaac7..8c95e36ee9 100644 --- a/uv.lock +++ b/uv.lock @@ -3623,6 +3623,7 @@ dependencies = [ { name = "hydra-core" }, { name = "mamba-ssm" }, { name = "megatron-core" }, + { name = "mlflow" }, { name = "nvidia-resiliency-ext" }, { name = "omegaconf" }, { name = "open-clip-torch" }, @@ -3633,6 +3634,8 @@ dependencies = [ { name = "six" }, { name = "tensorboard" }, { name = "timm" }, + { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "torch", version = "2.9.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform != 'darwin' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "tqdm" }, { name = "transformer-engine", extra = ["pytorch"] }, { name = "transformers" }, @@ -3644,11 +3647,12 @@ dependencies = [ requires-dist = [ { name = "accelerate" }, { name = "causal-conv1d", git = "https://github.com/Dao-AILab/causal-conv1d?rev=67e0a9dfe1518fc0036444e9ab5fe06ab78299e0" }, - { name = "datasets" }, + { name = "datasets", specifier = ">=2.20.0" }, { name = "flash-linear-attention" }, { name = "hydra-core", specifier = ">1.3,<=1.3.2" }, { name = "mamba-ssm", git = "https://github.com/state-spaces/mamba.git?rev=d68d16ed7d5d5164eb5a57c0285f3b7eb8394ec1" }, { name = "megatron-core", extras = ["dev", "mlm"], editable = "3rdparty/Megatron-LM-workspace" }, + { name = "mlflow", specifier = ">=3.5.0" }, { name = "nvidia-resiliency-ext" }, { name = "omegaconf", specifier = ">=2.3.0" }, { name = "open-clip-torch", specifier = ">=3.2.0" }, @@ -3659,8 +3663,10 @@ requires-dist = [ { name = "six", specifier = ">=1.17.0" }, { name = "tensorboard", specifier = ">=2.19.0" }, { name = "timm" }, + { name = "torch", marker = "sys_platform != 'darwin'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cu129" }, + { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.6.0", index = "https://pypi.org/simple" }, { name = "tqdm", specifier = ">=4.67.1" }, - { name = "transformer-engine", extras = ["pytorch"], specifier = ">=2.10.0a0,<2.12.0" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], specifier = ">=2.10.0a0,<2.13.0" }, { name = "transformers", specifier = "<5.0.0" }, { name = "typing-extensions" }, { name = "wandb", specifier = ">=0.19.10" }, From 07f73ce10f4b7a142188652adb9d313e36252bf7 Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 05:29:49 +0200 Subject: [PATCH 13/23] Fixed mcore inference module Signed-off-by: Ahmad Kiswani --- .../policy/workers/megatron_policy_worker.py | 75 ++++++------------- 1 file changed, 21 insertions(+), 54 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 5f1483ed9a..5a6a683765 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -36,9 +36,7 @@ from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( FullyShardedDataParallel as custom_FSDP, ) -from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( - InferenceWrapperConfig, -) +from megatron.core.inference.config import InferenceConfig from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) @@ -702,14 +700,8 @@ def generate( ) model_cfg = self.megatron_cfg.model - inference_wrapper_config = InferenceWrapperConfig( - hidden_size=model_cfg.hidden_size, - inference_batch_times_seqlen_threshold=1000000, - fp32_residual_connection=model_cfg.fp32_residual_connection, - params_dtype=model_cfg.params_dtype, - padded_vocab_size=self.final_padded_vocab_size, # Use the potentially updated value - inference_max_seq_length=self.cfg["generation"]["max_new_tokens"], # type: ignore - inference_max_requests=self.cfg["generation_batch_size"], + mcore_generation_config = cast( + MegatronGenerationConfig, self.cfg["generation"]["mcore_generation_config"] ) from megatron.core.inference.contexts.dynamic_context import ( @@ -723,45 +715,32 @@ def generate( ) from megatron.core.inference.sampling_params import SamplingParams - mcore_generation_config = cast( - MegatronGenerationConfig, self.cfg["generation"]["mcore_generation_config"] - ) - buffer_size_gb = mcore_generation_config["buffer_size_gb"] - - num_cuda_graphs = mcore_generation_config["num_cuda_graphs"] - block_size_tokens = mcore_generation_config["block_size_tokens"] - use_cuda_graphs_for_non_decode_steps = mcore_generation_config[ - "use_cuda_graphs_for_non_decode_steps" - ] - enable_chunked_prefill = mcore_generation_config["enable_chunked_prefill"] - unified_memory_level = mcore_generation_config["unified_memory_level"] - max_tokens = mcore_generation_config["max_tokens"] - model_config = self.model.config model_config.cuda_graph_impl = "local" - dynamic_context = DynamicInferenceContext( - params_dtype=inference_wrapper_config.params_dtype, - num_layers=model_config.num_layers, - kv_channels=model_config.kv_channels, - num_attention_heads=model_config.num_query_groups, + local_rank = torch.cuda.current_device() + num_gpus_per_node = torch.cuda.device_count() + node_idx = self.rank // num_gpus_per_node if num_gpus_per_node > 0 else 0 + model_config.inference_sampling_seed = (node_idx * 1024) + local_rank + + inference_config = InferenceConfig( max_sequence_length=self.cfg["generation"]["max_new_tokens"], - buffer_size_gb=buffer_size_gb, - materialize_only_last_token_logits=False, - num_cuda_graphs=num_cuda_graphs, - block_size_tokens=block_size_tokens, - tensor_model_parallel_size=self.cfg["megatron_cfg"][ - "tensor_model_parallel_size" + buffer_size_gb=mcore_generation_config["buffer_size_gb"], + num_cuda_graphs=mcore_generation_config["num_cuda_graphs"], + block_size_tokens=mcore_generation_config["block_size_tokens"], + use_cuda_graphs_for_non_decode_steps=mcore_generation_config[ + "use_cuda_graphs_for_non_decode_steps" ], - use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, + enable_chunked_prefill=mcore_generation_config["enable_chunked_prefill"], + unified_memory_level=mcore_generation_config["unified_memory_level"], + max_tokens=mcore_generation_config["max_tokens"], + materialize_only_last_token_logits=False, use_flashinfer_fused_rope=False, - unified_memory_level=unified_memory_level, - max_tokens=max_tokens, - ) - inference_wrapped_model = GPTInferenceWrapper( - self.model, inference_wrapper_config, dynamic_context ) + dynamic_context = DynamicInferenceContext(model_config, inference_config) + inference_wrapped_model = GPTInferenceWrapper(self.model, dynamic_context) + inference_wrapped_model.prep_model_for_inference() # Set pipeline parallel flag inference_wrapped_model.model_is_pipeline_parallel = ( @@ -773,21 +752,9 @@ def generate( tokenizer=self.megatron_tokenizer, ) - # Calculate seed based on node and rank to ensure reproducibility across workers - local_rank = torch.cuda.current_device() # Local GPU index on the node - num_gpus_per_node = torch.cuda.device_count() - node_idx = self.rank // num_gpus_per_node if num_gpus_per_node > 0 else 0 - seed = (node_idx * 1024) + local_rank - - # New API: DynamicInferenceEngine has additional parameters dynamic_engine = DynamicInferenceEngine( text_generation_controller, dynamic_context, - enable_cuda_graph=True, - random_seed=seed, - track_paused_request_events=False, - enable_chunked_prefill=enable_chunked_prefill, - inference_logging_step_interval=0, ) # Handle None values for top_k - convert to integer as required by Megatron From 3b9f96b63d7f9d702361c8037754a6ad9398d054 Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 05:40:09 +0200 Subject: [PATCH 14/23] bumped megatron-lm again Signed-off-by: Ahmad Kiswani --- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index 0d0943c6bf..a6d6dc6a85 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit 0d0943c6bfa9cbb30fcd62d40ce1792c4cb201e8 +Subproject commit a6d6dc6a853e94cf222881f9d67084383ddf5b65 From f6f06c9cf0fa535c7f58182feceee70862e531ad Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 05:56:12 +0200 Subject: [PATCH 15/23] fix: move SequencePackingGradientTestActor to separate module to avoid pytest import in mcore worker venv Signed-off-by: Ahmad Kiswani --- .../sequence_packing_gradient_actor.py | 380 ++++++++++++++++++ .../test_sequence_packing_gradients.py | 363 +---------------- 2 files changed, 383 insertions(+), 360 deletions(-) create mode 100644 tests/unit/algorithms/sequence_packing_gradient_actor.py diff --git a/tests/unit/algorithms/sequence_packing_gradient_actor.py b/tests/unit/algorithms/sequence_packing_gradient_actor.py new file mode 100644 index 0000000000..20564d77af --- /dev/null +++ b/tests/unit/algorithms/sequence_packing_gradient_actor.py @@ -0,0 +1,380 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ray actor for sequence packing gradient tests. + +Separated from test_sequence_packing_gradients.py to avoid importing pytest +in Ray worker environments that use PY_EXECUTABLES.MCORE. +""" + +import os +from unittest.mock import MagicMock + +import ray +import torch + +from nemo_rl.algorithms.loss_functions import ( + ClippedPGLossFn, + SequencePackingLossWrapper, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +@ray.remote(num_gpus=1) +class SequencePackingGradientTestActor: + def __init__(self, cp_size): + self.cp_size = cp_size + self.env_vars = dict(os.environ) + + def test_sequence_packing_gradients(self): + from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank + from nemo_rl.models.megatron.data import ( + _pack_sequences_for_megatron, + make_processed_microbatch_iterator, + ) + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + forward_with_post_processing_fn, + ) + + # Initialize process group + torch.distributed.init_process_group(backend="nccl") + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + # Create CP group - all ranks participate in CP + cp_group = torch.distributed.new_group(ranks=list(range(world_size))) + + # Patch get_context_parallel_group to always return cp_group + # (Assume it's imported from nemo_rl.models.megatron.common) + import megatron.core.parallel_state as parallel_state + + parallel_state._CONTEXT_PARALLEL_GROUP = cp_group + parallel_state._TENSOR_MODEL_PARALLEL_GROUP = torch.distributed.new_group( + ranks=[rank] + ) + + # Test parameters + batch_size = 4 + max_seq_len = 512 + vocab_size = 1000 + cp_size = self.cp_size + + # Ensure sequence length is compatible with CP load balancing + if max_seq_len % (2 * cp_size) != 0: + max_seq_len = (max_seq_len // (2 * cp_size) + 1) * (2 * cp_size) + + # Create test data with varying sequence lengths + torch.manual_seed(42) # For reproducibility + seq_lengths = torch.tensor( + [ + max_seq_len // 4, + max_seq_len * 1 // 4, + max_seq_len // 4, + max_seq_len * 3 // 4, + ] + ) + + # Create input data + input_ids = torch.zeros( + batch_size, max_seq_len, dtype=torch.long, device="cuda" + ) + token_mask = torch.zeros( + batch_size, max_seq_len, dtype=torch.float, device="cuda" + ) + + # Fill with random tokens up to seq_length + for i in range(batch_size): + length = seq_lengths[i] + input_ids[i, :length] = torch.randint( + 0, vocab_size, (length,), device="cuda" + ) + token_mask[i, :length] = 1.0 + + # Create other required tensors + sample_mask = torch.ones(batch_size, dtype=torch.float, device="cuda") + advantages = torch.randn(batch_size, max_seq_len, device="cuda") + prev_logprobs = torch.randn(batch_size, max_seq_len, device="cuda") + generation_logprobs = torch.randn(batch_size, max_seq_len, device="cuda") + reference_policy_logprobs = generation_logprobs.clone() + + original_data = { + "input_ids": input_ids, + "input_lengths": seq_lengths, + "token_mask": token_mask, + "sample_mask": sample_mask, + "advantages": advantages, + "prev_logprobs": prev_logprobs, + "generation_logprobs": generation_logprobs, + "reference_policy_logprobs": reference_policy_logprobs, + } + + # ===== TEST 1: Baseline (no sequence packing) ===== + print(f"Rank {rank}: Testing baseline (no sequence packing)") + + baseline_logits = torch.randn( + batch_size, max_seq_len, vocab_size, requires_grad=True, device="cuda" + ) + + loss_config = { + "reference_policy_kl_penalty": 0.1, + "reference_policy_kl_type": "k3", + "kl_input_clamp_value": 20.0, + "kl_output_clamp_value": 10.0, + "ratio_clip_min": 0.2, + "ratio_clip_max": 0.2, + "ratio_clip_c": 3.0, + "use_on_policy_kl_approximation": False, + "use_importance_sampling_correction": False, + "truncated_importance_sampling_ratio": None, + "sequence_level_importance_ratios": False, + "token_level_loss": True, + "force_on_policy_ratio": False, + } + + base_loss_fn = ClippedPGLossFn(loss_config) + data_dict = BatchedDataDict(original_data) + + global_valid_toks = torch.tensor( + sum(seq_lengths).item(), dtype=torch.float, device="cuda" + ) + global_valid_seqs = torch.tensor(batch_size, dtype=torch.float, device="cuda") + + # Forward pass + baseline_loss, baseline_metrics = base_loss_fn( + baseline_logits, + data_dict, + global_valid_seqs, + global_valid_toks, + ) + + # Backward pass + baseline_loss.backward() + + # Check baseline gradients + baseline_grad_norm = torch.norm(baseline_logits.grad).item() + baseline_grad_max = torch.max(torch.abs(baseline_logits.grad)).item() + baseline_grad_mean = torch.mean(torch.abs(baseline_logits.grad)).item() + baseline_grad_store = baseline_logits.grad.clone() + baseline_logits.grad.zero_() + + print( + f"Rank {rank}: Baseline gradient stats - norm: {baseline_grad_norm:.4f}, max: {baseline_grad_max:.4f}, mean: {baseline_grad_mean:.4f}" + ) + + # ===== TEST 2: Sequence packing with context parallelism ===== + print(f"Rank {rank}: Testing with sequence packing + CP") + + # Pack sequences + pad_to_multiple = cp_size * 2 # Common requirement for CP + ( + packed_input_ids, + packed_input_ids_cp, + packed_seq_params, + cu_seqlens, + cu_seqlens_padded, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=pad_to_multiple, + pad_packed_seq_to=max_seq_len * batch_size if cp_size > 1 else None, + cp_rank=rank, + cp_size=cp_size, + ) + + # For CP, logits are sharded across context parallel ranks + def make_packed_logits(logits): + packed_logits = torch.zeros( + 1, packed_input_ids_cp.shape[1], vocab_size, device="cuda" + ) + run_seq = 0 + for i, seq_len in enumerate(seq_lengths): + padded_seqlen = cu_seqlens_padded[i + 1] - cu_seqlens_padded[i] + if padded_seqlen > baseline_logits.shape[1]: + # pad the logits with zeros + tmp_logits = torch.zeros( + 1, padded_seqlen, vocab_size, device="cuda" + ) + tmp_logits[:, :seq_len] = baseline_logits[i : i + 1, :seq_len] + else: + tmp_logits = baseline_logits[i : i + 1, :padded_seqlen] + packed_logits[ + :, run_seq // cp_size : (run_seq + padded_seqlen) // cp_size, : + ] = _get_tokens_on_this_cp_rank(tmp_logits, rank, cp_size) + run_seq += padded_seqlen + return packed_logits + + packed_logits = make_packed_logits(baseline_logits) + + # Create sequence packing wrapper + wrapper = SequencePackingLossWrapper( + loss_fn=base_loss_fn, + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + ) + + # Create data dict for packed sequences + packed_data_dict = BatchedDataDict(original_data) + + tp_group = torch.distributed.new_group(ranks=[rank]) + + # Forward pass + packed_loss, packed_metrics = wrapper( + packed_logits, + packed_data_dict, + global_valid_seqs, + global_valid_toks, + vocab_parallel_rank=0, + vocab_parallel_group=tp_group, + context_parallel_group=cp_group, + ) + + # Backward pass + packed_loss /= cp_size + packed_loss.backward() + + # Check packed gradients + packed_grad = baseline_logits.grad.clone() + # all-reduce across cp ranks + torch.distributed.all_reduce(packed_grad, op=torch.distributed.ReduceOp.SUM) + + packed_grad_norm = torch.norm(packed_grad).item() + packed_grad_max = torch.max(torch.abs(packed_grad)).item() + packed_grad_mean = torch.mean(torch.abs(packed_grad)).item() + + print( + f"Rank {rank}: Packed gradient stats - norm: {packed_grad_norm:.4f}, max: {packed_grad_max:.4f}, mean: {packed_grad_mean:.4f}" + ) + + # ===== ANALYSIS ===== + gradient_ratio_norm = ( + packed_grad_norm / baseline_grad_norm + if baseline_grad_norm > 0 + else float("inf") + ) + gradient_ratio_max = ( + packed_grad_max / baseline_grad_max + if baseline_grad_max > 0 + else float("inf") + ) + gradient_ratio_mean = ( + packed_grad_mean / baseline_grad_mean + if baseline_grad_mean > 0 + else float("inf") + ) + + print( + f"Rank {rank}: Gradient ratios - norm: {gradient_ratio_norm:.4f}, max: {gradient_ratio_max:.4f}, mean: {gradient_ratio_mean:.4f}" + ) + + print( + f"differences by token: {torch.sum(torch.abs(packed_grad - baseline_grad_store), dim=-1)}" + ) + + torch.testing.assert_close( + packed_grad, baseline_grad_store, atol=1e-5, rtol=1e-5 + ) + + # test 3: with forward_with_post_processing_fn + # reset grad + baseline_logits.grad.zero_() + packed_logits = make_packed_logits(baseline_logits) + + # mock straggler detector with dummy context manager + mock_straggler_timer = MagicMock() + mock_straggler_timer.return_value = MagicMock( + __enter__=MagicMock(return_value=None), + __exit__=MagicMock(return_value=False), + ) + + # mock model forward + class MockModel: + def __init__(self): + self.logits = packed_logits + + def __call__(self, *args, **kwargs): + return self.logits + + def forward( + self, input_ids, position_ids, attention_mask, packed_seq_params=None + ): + return self.logits + + cfg = { + "sequence_packing": {"enabled": True}, + "dynamic_batching": {"enabled": False}, + "megatron_cfg": { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": cp_size, + }, + } + + post_processor = LossPostProcessor( + loss_fn=base_loss_fn, + cfg=cfg, + cp_normalize=True, + ) + + output_tensor, wrapped_loss_fn = forward_with_post_processing_fn( + data_iterator=make_processed_microbatch_iterator( + iter([packed_data_dict]), + cfg=cfg, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=pad_to_multiple, + pad_packed_seq_to_multiple_of=1, + straggler_timer=mock_straggler_timer, + pad_full_seq_to=max_seq_len * batch_size if cp_size > 1 else None, + ), + model=MockModel(), + cfg=cfg, + post_processing_fn=post_processor, + global_valid_seqs=global_valid_seqs, + global_valid_toks=global_valid_toks, + straggler_timer=mock_straggler_timer, + ) + loss, metrics = wrapped_loss_fn(output_tensor) + + loss.backward() + + # Check packed gradients + packed_grad = baseline_logits.grad.clone() + # all-reduce across cp ranks + torch.distributed.all_reduce(packed_grad, op=torch.distributed.ReduceOp.SUM) + + packed_grad_norm = torch.norm(packed_grad).item() + packed_grad_max = torch.max(torch.abs(packed_grad)).item() + packed_grad_mean = torch.mean(torch.abs(packed_grad)).item() + print( + f"Rank {rank}: Packed gradient stats - norm: {packed_grad_norm:.4f}, max: {packed_grad_max:.4f}, mean: {packed_grad_mean:.4f}" + ) + + gradient_ratio_norm = ( + packed_grad_norm / baseline_grad_norm + if baseline_grad_norm > 0 + else float("inf") + ) + gradient_ratio_max = ( + packed_grad_max / baseline_grad_max + if baseline_grad_max > 0 + else float("inf") + ) + + print( + f"Rank {rank}: Gradient ratios - norm: {gradient_ratio_norm:.4f}, max: {gradient_ratio_max:.4f}" + ) + print( + f"differences by token: {torch.sum(torch.abs(packed_grad - baseline_grad_store), dim=-1)}" + ) diff --git a/tests/unit/algorithms/test_sequence_packing_gradients.py b/tests/unit/algorithms/test_sequence_packing_gradients.py index f0ce832eb0..88ec9ce1b4 100644 --- a/tests/unit/algorithms/test_sequence_packing_gradients.py +++ b/tests/unit/algorithms/test_sequence_packing_gradients.py @@ -13,18 +13,10 @@ # limitations under the License. """Test script to debug high gradients with sequence packing + context parallelism.""" -import os -from unittest.mock import MagicMock - import pytest import ray import torch -from nemo_rl.algorithms.loss_functions import ( - ClippedPGLossFn, - SequencePackingLossWrapper, -) -from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.distributed.ray_actor_environment_registry import ( ACTOR_ENVIRONMENT_REGISTRY, @@ -32,358 +24,9 @@ ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup - - -@ray.remote(num_gpus=1) -class SequencePackingGradientTestActor: - def __init__(self, cp_size): - self.cp_size = cp_size - self.env_vars = dict(os.environ) - - def test_sequence_packing_gradients(self): - from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank - from nemo_rl.models.megatron.data import ( - _pack_sequences_for_megatron, - make_processed_microbatch_iterator, - ) - from nemo_rl.models.megatron.train import ( - LossPostProcessor, - forward_with_post_processing_fn, - ) - - # Initialize process group - torch.distributed.init_process_group(backend="nccl") - - rank = int(os.environ["RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - - # Create CP group - all ranks participate in CP - cp_group = torch.distributed.new_group(ranks=list(range(world_size))) - - # Patch get_context_parallel_group to always return cp_group - # (Assume it's imported from nemo_rl.models.megatron.common) - import megatron.core.parallel_state as parallel_state - - parallel_state._CONTEXT_PARALLEL_GROUP = cp_group - parallel_state._TENSOR_MODEL_PARALLEL_GROUP = torch.distributed.new_group( - ranks=[rank] - ) - - # Test parameters - batch_size = 4 - max_seq_len = 512 - vocab_size = 1000 - cp_size = self.cp_size - - # Ensure sequence length is compatible with CP load balancing - if max_seq_len % (2 * cp_size) != 0: - max_seq_len = (max_seq_len // (2 * cp_size) + 1) * (2 * cp_size) - - # Create test data with varying sequence lengths - torch.manual_seed(42) # For reproducibility - seq_lengths = torch.tensor( - [ - max_seq_len // 4, - max_seq_len * 1 // 4, - max_seq_len // 4, - max_seq_len * 3 // 4, - ] - ) - - # Create input data - input_ids = torch.zeros( - batch_size, max_seq_len, dtype=torch.long, device="cuda" - ) - token_mask = torch.zeros( - batch_size, max_seq_len, dtype=torch.float, device="cuda" - ) - - # Fill with random tokens up to seq_length - for i in range(batch_size): - length = seq_lengths[i] - input_ids[i, :length] = torch.randint( - 0, vocab_size, (length,), device="cuda" - ) - token_mask[i, :length] = 1.0 - - # Create other required tensors - sample_mask = torch.ones(batch_size, dtype=torch.float, device="cuda") - advantages = torch.randn(batch_size, max_seq_len, device="cuda") - prev_logprobs = torch.randn(batch_size, max_seq_len, device="cuda") - generation_logprobs = torch.randn(batch_size, max_seq_len, device="cuda") - reference_policy_logprobs = generation_logprobs.clone() - - original_data = { - "input_ids": input_ids, - "input_lengths": seq_lengths, - "token_mask": token_mask, - "sample_mask": sample_mask, - "advantages": advantages, - "prev_logprobs": prev_logprobs, - "generation_logprobs": generation_logprobs, - "reference_policy_logprobs": reference_policy_logprobs, - } - - # ===== TEST 1: Baseline (no sequence packing) ===== - print(f"Rank {rank}: Testing baseline (no sequence packing)") - - baseline_logits = torch.randn( - batch_size, max_seq_len, vocab_size, requires_grad=True, device="cuda" - ) - - loss_config = { - "reference_policy_kl_penalty": 0.1, - "reference_policy_kl_type": "k3", - "kl_input_clamp_value": 20.0, - "kl_output_clamp_value": 10.0, - "ratio_clip_min": 0.2, - "ratio_clip_max": 0.2, - "ratio_clip_c": 3.0, - "use_on_policy_kl_approximation": False, - "use_importance_sampling_correction": False, - "truncated_importance_sampling_ratio": None, - "sequence_level_importance_ratios": False, - "token_level_loss": True, - "force_on_policy_ratio": False, - } - - base_loss_fn = ClippedPGLossFn(loss_config) - data_dict = BatchedDataDict(original_data) - - global_valid_toks = torch.tensor( - sum(seq_lengths).item(), dtype=torch.float, device="cuda" - ) - global_valid_seqs = torch.tensor(batch_size, dtype=torch.float, device="cuda") - - # Forward pass - baseline_loss, baseline_metrics = base_loss_fn( - baseline_logits, - data_dict, - global_valid_seqs, - global_valid_toks, - ) - - # Backward pass - baseline_loss.backward() - - # Check baseline gradients - baseline_grad_norm = torch.norm(baseline_logits.grad).item() - baseline_grad_max = torch.max(torch.abs(baseline_logits.grad)).item() - baseline_grad_mean = torch.mean(torch.abs(baseline_logits.grad)).item() - baseline_grad_store = baseline_logits.grad.clone() - baseline_logits.grad.zero_() - - print( - f"Rank {rank}: Baseline gradient stats - norm: {baseline_grad_norm:.4f}, max: {baseline_grad_max:.4f}, mean: {baseline_grad_mean:.4f}" - ) - - # ===== TEST 2: Sequence packing with context parallelism ===== - print(f"Rank {rank}: Testing with sequence packing + CP") - - # Pack sequences - pad_to_multiple = cp_size * 2 # Common requirement for CP - ( - packed_input_ids, - packed_input_ids_cp, - packed_seq_params, - cu_seqlens, - cu_seqlens_padded, - ) = _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=pad_to_multiple, - pad_packed_seq_to=max_seq_len * batch_size if cp_size > 1 else None, - cp_rank=rank, - cp_size=cp_size, - ) - - # For CP, logits are sharded across context parallel ranks - def make_packed_logits(logits): - packed_logits = torch.zeros( - 1, packed_input_ids_cp.shape[1], vocab_size, device="cuda" - ) - run_seq = 0 - for i, seq_len in enumerate(seq_lengths): - padded_seqlen = cu_seqlens_padded[i + 1] - cu_seqlens_padded[i] - if padded_seqlen > baseline_logits.shape[1]: - # pad the logits with zeros - tmp_logits = torch.zeros( - 1, padded_seqlen, vocab_size, device="cuda" - ) - tmp_logits[:, :seq_len] = baseline_logits[i : i + 1, :seq_len] - else: - tmp_logits = baseline_logits[i : i + 1, :padded_seqlen] - packed_logits[ - :, run_seq // cp_size : (run_seq + padded_seqlen) // cp_size, : - ] = _get_tokens_on_this_cp_rank(tmp_logits, rank, cp_size) - run_seq += padded_seqlen - return packed_logits - - packed_logits = make_packed_logits(baseline_logits) - - # Create sequence packing wrapper - wrapper = SequencePackingLossWrapper( - loss_fn=base_loss_fn, - cu_seqlens_q=cu_seqlens, - cu_seqlens_q_padded=cu_seqlens_padded, - ) - - # Create data dict for packed sequences - packed_data_dict = BatchedDataDict(original_data) - - tp_group = torch.distributed.new_group(ranks=[rank]) - - # Forward pass - packed_loss, packed_metrics = wrapper( - packed_logits, - packed_data_dict, - global_valid_seqs, - global_valid_toks, - vocab_parallel_rank=0, - vocab_parallel_group=tp_group, - context_parallel_group=cp_group, - ) - - # Backward pass - packed_loss /= cp_size - packed_loss.backward() - - # Check packed gradients - packed_grad = baseline_logits.grad.clone() - # all-reduce across cp ranks - torch.distributed.all_reduce(packed_grad, op=torch.distributed.ReduceOp.SUM) - - packed_grad_norm = torch.norm(packed_grad).item() - packed_grad_max = torch.max(torch.abs(packed_grad)).item() - packed_grad_mean = torch.mean(torch.abs(packed_grad)).item() - # print(f"max grad on dims {torch.max(torch.abs(packed_grad), dim=0)}, {torch.max(torch.abs(packed_grad), dim=1)}, {torch.max(torch.abs(packed_grad), dim=2)}") - - print( - f"Rank {rank}: Packed gradient stats - norm: {packed_grad_norm:.4f}, max: {packed_grad_max:.4f}, mean: {packed_grad_mean:.4f}" - ) - - # ===== ANALYSIS ===== - gradient_ratio_norm = ( - packed_grad_norm / baseline_grad_norm - if baseline_grad_norm > 0 - else float("inf") - ) - gradient_ratio_max = ( - packed_grad_max / baseline_grad_max - if baseline_grad_max > 0 - else float("inf") - ) - gradient_ratio_mean = ( - packed_grad_mean / baseline_grad_mean - if baseline_grad_mean > 0 - else float("inf") - ) - - print( - f"Rank {rank}: Gradient ratios - norm: {gradient_ratio_norm:.4f}, max: {gradient_ratio_max:.4f}, mean: {gradient_ratio_mean:.4f}" - ) - - print( - f"differences by token: {torch.sum(torch.abs(packed_grad - baseline_grad_store), dim=-1)}" - ) - - torch.testing.assert_close( - packed_grad, baseline_grad_store, atol=1e-5, rtol=1e-5 - ) - - # test 3: with forward_with_post_processing_fn - # reset grad - baseline_logits.grad.zero_() - packed_logits = make_packed_logits(baseline_logits) - - # mock straggler detector with dummy context manager - mock_straggler_timer = MagicMock() - mock_straggler_timer.return_value = MagicMock( - __enter__=MagicMock(return_value=None), - __exit__=MagicMock(return_value=False), - ) - - # mock model forward - class MockModel: - def __init__(self): - self.logits = packed_logits - - def __call__(self, *args, **kwargs): - return self.logits - - def forward( - self, input_ids, position_ids, attention_mask, packed_seq_params=None - ): - return self.logits - - cfg = { - "sequence_packing": {"enabled": True}, - "dynamic_batching": {"enabled": False}, - "megatron_cfg": { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": cp_size, - }, - } - - post_processor = LossPostProcessor( - loss_fn=base_loss_fn, - cfg=cfg, - cp_normalize=True, - ) - - output_tensor, wrapped_loss_fn = forward_with_post_processing_fn( - data_iterator=make_processed_microbatch_iterator( - iter([packed_data_dict]), - cfg=cfg, - seq_length_key="input_lengths", - pad_individual_seqs_to_multiple_of=pad_to_multiple, - pad_packed_seq_to_multiple_of=1, - straggler_timer=mock_straggler_timer, - pad_full_seq_to=max_seq_len * batch_size if cp_size > 1 else None, - ), - model=MockModel(), - cfg=cfg, - post_processing_fn=post_processor, - global_valid_seqs=global_valid_seqs, - global_valid_toks=global_valid_toks, - straggler_timer=mock_straggler_timer, - ) - loss, metrics = wrapped_loss_fn(output_tensor) - - loss.backward() - - # Check packed gradients - packed_grad = baseline_logits.grad.clone() - # all-reduce across cp ranks - torch.distributed.all_reduce(packed_grad, op=torch.distributed.ReduceOp.SUM) - - packed_grad_norm = torch.norm(packed_grad).item() - packed_grad_max = torch.max(torch.abs(packed_grad)).item() - packed_grad_mean = torch.mean(torch.abs(packed_grad)).item() - print( - f"Rank {rank}: Packed gradient stats - norm: {packed_grad_norm:.4f}, max: {packed_grad_max:.4f}, mean: {packed_grad_mean:.4f}" - ) - - gradient_ratio_norm = ( - packed_grad_norm / baseline_grad_norm - if baseline_grad_norm > 0 - else float("inf") - ) - gradient_ratio_max = ( - packed_grad_max / baseline_grad_max - if baseline_grad_max > 0 - else float("inf") - ) - - print( - f"Rank {rank}: Gradient ratios - norm: {gradient_ratio_norm:.4f}, max: {gradient_ratio_max:.4f}" - ) - print( - f"differences by token: {torch.sum(torch.abs(packed_grad - baseline_grad_store), dim=-1)}" - ) - +from tests.unit.algorithms.sequence_packing_gradient_actor import ( + SequencePackingGradientTestActor, +) SEQUENCE_PACKING_GRADIENT_TEST_ACTOR_FQN = ( f"{SequencePackingGradientTestActor.__module__}.SequencePackingGradientTestActor" From 04d300ff4173bf41c9bd9a9e9f886dd4e100022d Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 13:32:20 +0200 Subject: [PATCH 16/23] fix: moved PackSequencesTestActor to a to separate module to avoid pytest import in mcore worker venv Signed-off-by: Ahmad Kiswani --- .../models/megatron/megatron_data_actors.py | 982 ++++++++++++++++++ .../models/megatron/test_megatron_data.py | 964 +---------------- 2 files changed, 986 insertions(+), 960 deletions(-) create mode 100644 tests/unit/models/megatron/megatron_data_actors.py diff --git a/tests/unit/models/megatron/megatron_data_actors.py b/tests/unit/models/megatron/megatron_data_actors.py new file mode 100644 index 0000000000..0687cf076b --- /dev/null +++ b/tests/unit/models/megatron/megatron_data_actors.py @@ -0,0 +1,982 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ray actors for megatron data tests. + +Separated from test_megatron_data.py to avoid importing pytest +in Ray worker environments that use PY_EXECUTABLES.MCORE. +""" + +import os + +import ray +import torch + + +@ray.remote(num_gpus=1) +class PackSequencesTestActor: + def __init__(self, cp_size): + self.cp_size = cp_size + self.env_vars = dict(os.environ) + + def run_all_pack_sequences_tests(self): + """Run all sequence packing tests in a single call to avoid expensive reinitializations.""" + from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank + from nemo_rl.models.megatron.data import _pack_sequences_for_megatron + + # Initialize process group if CP > 1 + if self.cp_size > 1: + torch.distributed.init_process_group(backend="nccl") + rank = int(os.environ["RANK"]) + else: + rank = 0 + + results = {} + + # Test 1: Basic packing functionality + results["basic"] = self._test_basic_packing(_pack_sequences_for_megatron) + if not results["basic"]["success"]: + return results["basic"] + + # Test 2: Variable sequence lengths + results["variable_lengths"] = self._test_variable_lengths( + _pack_sequences_for_megatron + ) + if not results["variable_lengths"]["success"]: + return results["variable_lengths"] + + # Test 3: Content preservation and consistency + results["consistency"] = self._test_consistency(_pack_sequences_for_megatron) + if not results["consistency"]["success"]: + return results["consistency"] + + # Test 4: Edge cases + results["edge_cases"] = self._test_edge_cases(_pack_sequences_for_megatron) + if not results["edge_cases"]["success"]: + return results["edge_cases"] + + # Test 5: Context parallelism (only if CP > 1) + if self.cp_size > 1: + results["context_parallel"] = self._test_context_parallel( + _pack_sequences_for_megatron, _get_tokens_on_this_cp_rank, rank + ) + if not results["context_parallel"]["success"]: + return results["context_parallel"] + else: + results["context_parallel"] = { + "success": True, + "error": None, + "skipped": "CP=1", + } + + return {"success": True, "error": None, "detailed_results": results} + + def _test_basic_packing(self, _pack_sequences_for_megatron): + """Test basic sequence packing without context parallelism.""" + try: + # Test parameters + batch_size = 3 + max_seq_len = 10 + vocab_size = 100 + + # Create test data with variable sequence lengths + input_ids = torch.randint( + 0, vocab_size, (batch_size, max_seq_len), device="cuda" + ) + seq_lengths = torch.tensor([8, 5, 7], device="cuda") + + # Test 1: Basic packing without CP + packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( + _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + ) + + # Verify shapes + expected_total_tokens = seq_lengths.sum().item() + if packed_input_ids.shape != (1, expected_total_tokens): + return { + "success": False, + "error": f"Basic packing shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", + } + + # Verify cu_seqlens + expected_cu_seqlens = torch.tensor( + [0, 8, 13, 20], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens, expected_cu_seqlens): + return { + "success": False, + "error": f"cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", + } + + # Verify PackedSeqParams + if packed_seq_params.qkv_format != "thd": + return { + "success": False, + "error": f"Wrong qkv_format: expected 'thd', got {packed_seq_params.qkv_format}", + } + + if packed_seq_params.max_seqlen_q != 8: + return { + "success": False, + "error": f"Wrong max_seqlen_q: expected 8, got {packed_seq_params.max_seqlen_q}", + } + + # Test 2: Packing with individual sequence padding + ( + packed_input_ids_pad, + _, + packed_seq_params_pad, + cu_seqlens_pad, + cu_seqlens_padded_pad, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=4, + cp_rank=0, + cp_size=1, + ) + + # With padding to multiple of 4: [8, 5, 7] -> [8, 8, 8] = 24 tokens + expected_total_tokens_pad = 24 + if packed_input_ids_pad.shape != (1, expected_total_tokens_pad): + return { + "success": False, + "error": f"Padded packing shape mismatch: expected (1, {expected_total_tokens_pad}), got {packed_input_ids_pad.shape}", + } + + # Verify padded cu_seqlens + expected_cu_seqlens_padded = torch.tensor( + [0, 8, 16, 24], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens_padded_pad, expected_cu_seqlens_padded): + return { + "success": False, + "error": f"Padded cu_seqlens mismatch: expected {expected_cu_seqlens_padded}, got {cu_seqlens_padded_pad}", + } + + return {"success": True, "error": None} + + except Exception as e: + return {"success": False, "error": f"Basic packing test failed: {str(e)}"} + + def _test_variable_lengths(self, _pack_sequences_for_megatron): + """Test sequence packing with variable sequence lengths.""" + try: + # Test parameters + batch_size = 4 + max_seq_len = 12 + vocab_size = 50 + + # Create test data with highly variable sequence lengths + input_ids = torch.randint( + 0, vocab_size, (batch_size, max_seq_len), device="cuda" + ) + seq_lengths = torch.tensor([12, 3, 8, 1], device="cuda") + + # Test 1: Variable lengths without padding + packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( + _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + ) + + # Verify total tokens + expected_total_tokens = seq_lengths.sum().item() # 12 + 3 + 8 + 1 = 24 + if packed_input_ids.shape != (1, expected_total_tokens): + return { + "success": False, + "error": f"Variable lengths shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", + } + + # Verify cu_seqlens + expected_cu_seqlens = torch.tensor( + [0, 12, 15, 23, 24], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens, expected_cu_seqlens): + return { + "success": False, + "error": f"Variable lengths cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", + } + + # Test 2: Variable lengths with padding + ( + packed_input_ids_pad, + _, + packed_seq_params_pad, + cu_seqlens_pad, + cu_seqlens_padded_pad, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=4, + cp_rank=0, + cp_size=1, + ) + + # With padding to multiple of 4: [12, 3, 8, 1] -> [12, 4, 8, 4] = 28 tokens + expected_total_tokens_pad = 28 + if packed_input_ids_pad.shape != (1, expected_total_tokens_pad): + return { + "success": False, + "error": f"Variable lengths padded shape mismatch: expected (1, {expected_total_tokens_pad}), got {packed_input_ids_pad.shape}", + } + + # Verify padded cu_seqlens + expected_cu_seqlens_padded = torch.tensor( + [0, 12, 16, 24, 28], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens_padded_pad, expected_cu_seqlens_padded): + return { + "success": False, + "error": f"Variable lengths padded cu_seqlens mismatch: expected {expected_cu_seqlens_padded}, got {cu_seqlens_padded_pad}", + } + + # Verify max_seqlen + if packed_seq_params.max_seqlen_q != 12: + return { + "success": False, + "error": f"Variable lengths wrong max_seqlen_q: expected 12, got {packed_seq_params.max_seqlen_q}", + } + + if packed_seq_params_pad.max_seqlen_q != 12: + return { + "success": False, + "error": f"Variable lengths padded wrong max_seqlen_q: expected 12, got {packed_seq_params_pad.max_seqlen_q}", + } + + return {"success": True, "error": None} + + except Exception as e: + return { + "success": False, + "error": f"Variable lengths test failed: {str(e)}", + } + + def _test_consistency(self, _pack_sequences_for_megatron): + """Test that packing produces consistent results and that content is preserved.""" + try: + # Test parameters + batch_size = 2 + seq_len = 8 + vocab_size = 20 + + # Create deterministic test data + torch.manual_seed(123) + input_ids = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda" + ) + seq_lengths = torch.tensor([6, 4], device="cuda") + + # Test consistency between multiple calls + ( + packed_input_ids_1, + _, + packed_seq_params_1, + cu_seqlens_1, + cu_seqlens_padded_1, + ) = _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + + ( + packed_input_ids_2, + _, + packed_seq_params_2, + cu_seqlens_2, + cu_seqlens_padded_2, + ) = _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + + # Verify consistency + if not torch.equal(packed_input_ids_1, packed_input_ids_2): + return { + "success": False, + "error": "Inconsistent packed_input_ids between calls", + } + + if not torch.equal(cu_seqlens_1, cu_seqlens_2): + return { + "success": False, + "error": "Inconsistent cu_seqlens between calls", + } + + # Verify content preservation + # Extract the first sequence (length 6) and compare with original + first_seq_packed = packed_input_ids_1[0, :6] + first_seq_original = input_ids[0, :6] + + if not torch.equal(first_seq_packed, first_seq_original): + return { + "success": False, + "error": "Content not preserved in first sequence", + } + + # Extract the second sequence (length 4) and compare with original + second_seq_packed = packed_input_ids_1[0, 6:10] + second_seq_original = input_ids[1, :4] + + if not torch.equal(second_seq_packed, second_seq_original): + return { + "success": False, + "error": "Content not preserved in second sequence", + } + + return {"success": True, "error": None} + + except Exception as e: + return {"success": False, "error": f"Consistency test failed: {str(e)}"} + + def _test_edge_cases(self, _pack_sequences_for_megatron): + """Test edge cases and error conditions.""" + try: + # Test 1: Single sequence + batch_size = 1 + seq_len = 10 + vocab_size = 50 + + input_ids = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda" + ) + seq_lengths = torch.tensor([seq_len], device="cuda") + + packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( + _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + ) + + # Verify single sequence packing + if packed_input_ids.shape != (1, seq_len): + return { + "success": False, + "error": f"Single sequence shape mismatch: expected (1, {seq_len}), got {packed_input_ids.shape}", + } + + expected_cu_seqlens = torch.tensor( + [0, seq_len], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens, expected_cu_seqlens): + return { + "success": False, + "error": f"Single sequence cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", + } + + # Test 2: Empty sequences (length 0) + batch_size = 3 + max_seq_len = 5 + input_ids = torch.randint( + 0, vocab_size, (batch_size, max_seq_len), device="cuda" + ) + seq_lengths = torch.tensor([3, 0, 2], device="cuda") + + packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( + _pack_sequences_for_megatron( + input_ids, seq_lengths, cp_rank=0, cp_size=1 + ) + ) + + # Should handle empty sequences gracefully + expected_total_tokens = 5 # 3 + 0 + 2 + if packed_input_ids.shape != (1, expected_total_tokens): + return { + "success": False, + "error": f"Empty sequence shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", + } + + expected_cu_seqlens = torch.tensor( + [0, 3, 3, 5], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens, expected_cu_seqlens): + return { + "success": False, + "error": f"Empty sequence cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", + } + + # Test 3: Large padding values + batch_size = 2 + seq_len = 4 + input_ids = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda" + ) + seq_lengths = torch.tensor([3, 2], device="cuda") + + packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( + _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=8, + cp_rank=0, + cp_size=1, + ) + ) + + # With padding to multiple of 8: [3, 2] -> [8, 8] = 16 tokens + expected_total_tokens = 16 + if packed_input_ids.shape != (1, expected_total_tokens): + return { + "success": False, + "error": f"Large padding shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", + } + + return {"success": True, "error": None} + + except Exception as e: + return {"success": False, "error": f"Edge cases test failed: {str(e)}"} + + def _test_context_parallel( + self, _pack_sequences_for_megatron, _get_tokens_on_this_cp_rank, rank + ): + """Test sequence packing with context parallelism.""" + # Test parameters + batch_size = 2 + seq_len = 16 # Ensure divisible by cp_size * 2 + vocab_size = 100 + + # Ensure sequence length is compatible with CP + if seq_len % (2 * self.cp_size) != 0: + seq_len = (seq_len // (2 * self.cp_size) + 1) * (2 * self.cp_size) + + # Create test data + torch.manual_seed(42) # For reproducibility + input_ids = torch.arange(seq_len * batch_size, device="cuda").reshape( + batch_size, seq_len + ) + seq_lengths = torch.tensor([seq_len, seq_len], device="cuda") + + # Test 1: CP packing with individual sequence padding + ( + packed_input_ids, + packed_input_ids_cp_sharded, + packed_seq_params, + cu_seqlens, + cu_seqlens_padded, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=self.cp_size * 2, + cp_rank=rank, + cp_size=self.cp_size, + ) + + # Verify the packed tensor shape + expected_tokens_per_rank = seq_len // self.cp_size + expected_total_tokens = batch_size * expected_tokens_per_rank + if packed_input_ids_cp_sharded.shape != (1, expected_total_tokens): + return { + "success": False, + "error": f"CP packing shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids_cp_sharded.shape}", + } + + # Verify cu_seqlens for original sequences + expected_cu_seqlens = torch.tensor( + [0, seq_len, seq_len * 2], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens, expected_cu_seqlens): + return { + "success": False, + "error": f"CP cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", + } + + # Verify PackedSeqParams + if packed_seq_params.qkv_format != "thd": + return { + "success": False, + "error": f"CP wrong qkv_format: expected 'thd', got {packed_seq_params.qkv_format}", + } + + # Test 2: CP packing with full sequence padding + pad_full_seq_to = (batch_size * seq_len) + 8 # Add some padding + ( + packed_input_ids_full, + packed_input_ids_cp_sharded, + packed_seq_params_full, + cu_seqlens_full, + cu_seqlens_padded_full, + ) = _pack_sequences_for_megatron( + input_ids, + seq_lengths, + pad_individual_seqs_to_multiple_of=self.cp_size * 2, + pad_packed_seq_to=pad_full_seq_to, + cp_rank=rank, + cp_size=self.cp_size, + ) + + # Verify the packed tensor shape with full padding + expected_tokens_per_rank_full = pad_full_seq_to // self.cp_size + if packed_input_ids_cp_sharded.shape != (1, expected_tokens_per_rank_full): + return { + "success": False, + "error": f"CP full padding shape mismatch: expected (1, {expected_tokens_per_rank_full}), got {packed_input_ids_cp_sharded.shape}", + } + + # Verify cu_seqlens_padded for full padding + expected_cu_seqlens_padded_full = torch.tensor( + [0, seq_len, pad_full_seq_to], device="cuda", dtype=torch.int32 + ) + if not torch.equal(cu_seqlens_padded_full, expected_cu_seqlens_padded_full): + return { + "success": False, + "error": f"CP full padding cu_seqlens_padded mismatch: expected {expected_cu_seqlens_padded_full}, got {cu_seqlens_padded_full}", + } + + correct_ids_0 = torch.tensor( + [0, 1, 2, 3, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 0, 0, 0, 0, 0, 0], + device="cuda", + ) + correct_ids_1 = torch.tensor( + [4, 5, 6, 7, 8, 9, 10, 11, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 0, 0], + device="cuda", + ) + + if ( + rank == 0 + and torch.sum(torch.abs(packed_input_ids_cp_sharded - correct_ids_0)).item() + != 0 + ): + return { + "success": False, + "error": f"CP full padding ids mismatch: expected {correct_ids_0}, got {packed_input_ids_cp_sharded[0, :20]}", + } + if ( + rank == 1 + and torch.sum(torch.abs(packed_input_ids_cp_sharded - correct_ids_1)).item() + != 0 + ): + return { + "success": False, + "error": f"CP full padding ids mismatch: expected {correct_ids_1}, got {packed_input_ids_cp_sharded[0, 20:]}", + } + + return {"success": True, "error": None} + + +@ray.remote(num_gpus=1) +class GetPackSequenceParametersTestActor: + def __init__(self): + pass + + def run_all_get_pack_sequence_parameters_for_megatron_tests(self): + """Test _get_pack_sequence_parameters_for_megatron function with various configurations.""" + from nemo_rl.models.megatron.data import ( + _get_pack_sequence_parameters_for_megatron, + ) + + # Test 1: Basic configuration - no parallelism, no FP8 + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + } + max_seq_len = 1023 + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 2: Context parallelism only + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 4, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 4 * 2 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=4*2, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 3: Tensor parallelism with sequence parallelism + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 2 # tp_size when SP is enabled + if pad_individual != 2 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=2, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 4: Tensor parallelism without sequence parallelism + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 5: Pipeline parallelism + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 4, + "context_parallel_size": 1, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to != max_seq_len: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 6: Combined CP and TP with SP + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 4, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size + if ( + pad_individual != expected_individual + or pad_packed != 1 + or pad_to is not None + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 7: FP8 enabled with default recipe + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "fp8_cfg": { + "enabled": True, + "fp8": "hybrid", + "fp8_recipe": "tensorwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 16 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=16, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 8: FP8 enabled with blockwise recipe + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 128 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=128, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 9: FP8 with CP and TP+SP + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 4, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size + expected_packed = 128 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size + if ( + pad_individual != expected_individual + or pad_packed != expected_packed + or pad_to is not None + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 10: All parallelism types with FP8 and PP + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 4, + "context_parallel_size": 2, + "fp8_cfg": { + "enabled": True, + "fp8": "hybrid", + "fp8_recipe": "tensorwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 2 * 2 * 2 # cp_size * 2 * tp_size + expected_packed = 16 * 2 * 2 * 2 # divisor * cp_size * 2 * tp_size + + def _round_up_to_multiple_of(x, y): + return (x + y - 1) // y * y + + if ( + pad_individual != expected_individual + or pad_packed != expected_packed + or pad_to != _round_up_to_multiple_of(max_seq_len, expected_packed) + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 11: FP8 disabled explicitly + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "fp8_cfg": { + "enabled": False, + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 12: Missing fp8_cfg (should default to disabled) + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + # No fp8_cfg key + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 13: Edge case - very large parallelism values + megatron_cfg = { + "tensor_model_parallel_size": 8, + "sequence_parallel": True, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 8, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 8 * 2 * 8 # cp_size * 2 * tp_size = 128 + expected_packed = 128 * 8 * 2 * 8 # divisor * cp_size * 2 * tp_size = 16384 + if ( + pad_individual != expected_individual + or pad_packed != expected_packed + or pad_to is not None + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 14: Edge case - different max_seq_len values with PP + for test_seq_len in [512, 2048, 4096]: + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 2, + "context_parallel_size": 1, + } + + pad_individual, pad_packed, pad_to = ( + _get_pack_sequence_parameters_for_megatron(megatron_cfg, test_seq_len) + ) + + if pad_individual != 1 or pad_packed != 1 or pad_to != test_seq_len: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=1, pad_to={test_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 15: FP8 with MXFP8 recipe + megatron_cfg = { + "tensor_model_parallel_size": 1, + "sequence_parallel": False, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "mxfp8", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + if pad_individual != 1 or pad_packed != 32 or pad_to is not None: + return { + "success": False, + "error": f"Expected pad_individual=1, pad_packed=32, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 16: FP8 with MXFP8 recipe, CP, and TP+SP + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 4, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "mxfp8", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size + expected_packed = 32 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size + + if ( + pad_individual != expected_individual + or pad_packed != expected_packed + or pad_to is not None + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + # Test 17: FP8 with MXFP8 recipe, CP, TP+SP, and PP + megatron_cfg = { + "tensor_model_parallel_size": 2, + "sequence_parallel": True, + "pipeline_model_parallel_size": 4, + "context_parallel_size": 4, + "fp8_cfg": { + "enabled": True, + "fp8": "e4m3", + "fp8_recipe": "mxfp8", + "fp8_param": False, + }, + } + + pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( + megatron_cfg, max_seq_len + ) + + expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size + expected_packed = 32 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size * pp_size + + if ( + pad_individual != expected_individual + or pad_packed != expected_packed + or pad_to != _round_up_to_multiple_of(max_seq_len, expected_packed) + ): + return { + "success": False, + "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", + } + + return {"success": True, "error": None} diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index 6e381d2933..3610b77d9f 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -23,7 +23,6 @@ - Sequence dimension validation """ -import os from unittest.mock import MagicMock, patch import pytest @@ -38,6 +37,10 @@ ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup +from tests.unit.models.megatron.megatron_data_actors import ( + GetPackSequenceParametersTestActor, + PackSequencesTestActor, +) @pytest.mark.mcore @@ -630,546 +633,6 @@ def test_make_processed_microbatch_iterator_with_packing(self, mock_process): assert call_kwargs["pad_full_seq_to"] == 1024 -@ray.remote(num_gpus=1) -class PackSequencesTestActor: - def __init__(self, cp_size): - self.cp_size = cp_size - self.env_vars = dict(os.environ) - - def run_all_pack_sequences_tests(self): - """Run all sequence packing tests in a single call to avoid expensive reinitializations.""" - from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank - from nemo_rl.models.megatron.data import _pack_sequences_for_megatron - - # Initialize process group if CP > 1 - if self.cp_size > 1: - torch.distributed.init_process_group(backend="nccl") - rank = int(os.environ["RANK"]) - else: - rank = 0 - - results = {} - - # Test 1: Basic packing functionality - results["basic"] = self._test_basic_packing(_pack_sequences_for_megatron) - if not results["basic"]["success"]: - return results["basic"] - - # Test 2: Variable sequence lengths - results["variable_lengths"] = self._test_variable_lengths( - _pack_sequences_for_megatron - ) - if not results["variable_lengths"]["success"]: - return results["variable_lengths"] - - # Test 3: Content preservation and consistency - results["consistency"] = self._test_consistency(_pack_sequences_for_megatron) - if not results["consistency"]["success"]: - return results["consistency"] - - # Test 4: Edge cases - results["edge_cases"] = self._test_edge_cases(_pack_sequences_for_megatron) - if not results["edge_cases"]["success"]: - return results["edge_cases"] - - # Test 5: Context parallelism (only if CP > 1) - if self.cp_size > 1: - results["context_parallel"] = self._test_context_parallel( - _pack_sequences_for_megatron, _get_tokens_on_this_cp_rank, rank - ) - if not results["context_parallel"]["success"]: - return results["context_parallel"] - else: - results["context_parallel"] = { - "success": True, - "error": None, - "skipped": "CP=1", - } - - return {"success": True, "error": None, "detailed_results": results} - - def _test_basic_packing(self, _pack_sequences_for_megatron): - """Test basic sequence packing without context parallelism.""" - try: - # Test parameters - batch_size = 3 - max_seq_len = 10 - vocab_size = 100 - - # Create test data with variable sequence lengths - input_ids = torch.randint( - 0, vocab_size, (batch_size, max_seq_len), device="cuda" - ) - seq_lengths = torch.tensor([8, 5, 7], device="cuda") - - # Test 1: Basic packing without CP - packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( - _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - ) - - # Verify shapes - expected_total_tokens = seq_lengths.sum().item() - if packed_input_ids.shape != (1, expected_total_tokens): - return { - "success": False, - "error": f"Basic packing shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", - } - - # Verify cu_seqlens - expected_cu_seqlens = torch.tensor( - [0, 8, 13, 20], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens, expected_cu_seqlens): - return { - "success": False, - "error": f"cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", - } - - # Verify PackedSeqParams - if packed_seq_params.qkv_format != "thd": - return { - "success": False, - "error": f"Wrong qkv_format: expected 'thd', got {packed_seq_params.qkv_format}", - } - - if packed_seq_params.max_seqlen_q != 8: - return { - "success": False, - "error": f"Wrong max_seqlen_q: expected 8, got {packed_seq_params.max_seqlen_q}", - } - - # Test 2: Packing with individual sequence padding - ( - packed_input_ids_pad, - _, - packed_seq_params_pad, - cu_seqlens_pad, - cu_seqlens_padded_pad, - ) = _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=4, - cp_rank=0, - cp_size=1, - ) - - # With padding to multiple of 4: [8, 5, 7] -> [8, 8, 8] = 24 tokens - expected_total_tokens_pad = 24 - if packed_input_ids_pad.shape != (1, expected_total_tokens_pad): - return { - "success": False, - "error": f"Padded packing shape mismatch: expected (1, {expected_total_tokens_pad}), got {packed_input_ids_pad.shape}", - } - - # Verify padded cu_seqlens - expected_cu_seqlens_padded = torch.tensor( - [0, 8, 16, 24], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens_padded_pad, expected_cu_seqlens_padded): - return { - "success": False, - "error": f"Padded cu_seqlens mismatch: expected {expected_cu_seqlens_padded}, got {cu_seqlens_padded_pad}", - } - - return {"success": True, "error": None} - - except Exception as e: - return {"success": False, "error": f"Basic packing test failed: {str(e)}"} - - def _test_variable_lengths(self, _pack_sequences_for_megatron): - """Test sequence packing with variable sequence lengths.""" - try: - # Test parameters - batch_size = 4 - max_seq_len = 12 - vocab_size = 50 - - # Create test data with highly variable sequence lengths - input_ids = torch.randint( - 0, vocab_size, (batch_size, max_seq_len), device="cuda" - ) - seq_lengths = torch.tensor([12, 3, 8, 1], device="cuda") - - # Test 1: Variable lengths without padding - packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( - _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - ) - - # Verify total tokens - expected_total_tokens = seq_lengths.sum().item() # 12 + 3 + 8 + 1 = 24 - if packed_input_ids.shape != (1, expected_total_tokens): - return { - "success": False, - "error": f"Variable lengths shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", - } - - # Verify cu_seqlens - expected_cu_seqlens = torch.tensor( - [0, 12, 15, 23, 24], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens, expected_cu_seqlens): - return { - "success": False, - "error": f"Variable lengths cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", - } - - # Test 2: Variable lengths with padding - ( - packed_input_ids_pad, - _, - packed_seq_params_pad, - cu_seqlens_pad, - cu_seqlens_padded_pad, - ) = _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=4, - cp_rank=0, - cp_size=1, - ) - - # With padding to multiple of 4: [12, 3, 8, 1] -> [12, 4, 8, 4] = 28 tokens - expected_total_tokens_pad = 28 - if packed_input_ids_pad.shape != (1, expected_total_tokens_pad): - return { - "success": False, - "error": f"Variable lengths padded shape mismatch: expected (1, {expected_total_tokens_pad}), got {packed_input_ids_pad.shape}", - } - - # Verify padded cu_seqlens - expected_cu_seqlens_padded = torch.tensor( - [0, 12, 16, 24, 28], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens_padded_pad, expected_cu_seqlens_padded): - return { - "success": False, - "error": f"Variable lengths padded cu_seqlens mismatch: expected {expected_cu_seqlens_padded}, got {cu_seqlens_padded_pad}", - } - - # Verify max_seqlen - if packed_seq_params.max_seqlen_q != 12: - return { - "success": False, - "error": f"Variable lengths wrong max_seqlen_q: expected 12, got {packed_seq_params.max_seqlen_q}", - } - - if packed_seq_params_pad.max_seqlen_q != 12: - return { - "success": False, - "error": f"Variable lengths padded wrong max_seqlen_q: expected 12, got {packed_seq_params_pad.max_seqlen_q}", - } - - return {"success": True, "error": None} - - except Exception as e: - return { - "success": False, - "error": f"Variable lengths test failed: {str(e)}", - } - - def _test_consistency(self, _pack_sequences_for_megatron): - """Test that packing produces consistent results and that content is preserved.""" - try: - # Test parameters - batch_size = 2 - seq_len = 8 - vocab_size = 20 - - # Create deterministic test data - torch.manual_seed(123) - input_ids = torch.randint( - 0, vocab_size, (batch_size, seq_len), device="cuda" - ) - seq_lengths = torch.tensor([6, 4], device="cuda") - - # Test consistency between multiple calls - ( - packed_input_ids_1, - _, - packed_seq_params_1, - cu_seqlens_1, - cu_seqlens_padded_1, - ) = _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - - ( - packed_input_ids_2, - _, - packed_seq_params_2, - cu_seqlens_2, - cu_seqlens_padded_2, - ) = _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - - # Verify consistency - if not torch.equal(packed_input_ids_1, packed_input_ids_2): - return { - "success": False, - "error": "Inconsistent packed_input_ids between calls", - } - - if not torch.equal(cu_seqlens_1, cu_seqlens_2): - return { - "success": False, - "error": "Inconsistent cu_seqlens between calls", - } - - # Verify content preservation - # Extract the first sequence (length 6) and compare with original - first_seq_packed = packed_input_ids_1[0, :6] - first_seq_original = input_ids[0, :6] - - if not torch.equal(first_seq_packed, first_seq_original): - return { - "success": False, - "error": "Content not preserved in first sequence", - } - - # Extract the second sequence (length 4) and compare with original - second_seq_packed = packed_input_ids_1[0, 6:10] - second_seq_original = input_ids[1, :4] - - if not torch.equal(second_seq_packed, second_seq_original): - return { - "success": False, - "error": "Content not preserved in second sequence", - } - - return {"success": True, "error": None} - - except Exception as e: - return {"success": False, "error": f"Consistency test failed: {str(e)}"} - - def _test_edge_cases(self, _pack_sequences_for_megatron): - """Test edge cases and error conditions.""" - try: - # Test 1: Single sequence - batch_size = 1 - seq_len = 10 - vocab_size = 50 - - input_ids = torch.randint( - 0, vocab_size, (batch_size, seq_len), device="cuda" - ) - seq_lengths = torch.tensor([seq_len], device="cuda") - - packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( - _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - ) - - # Verify single sequence packing - if packed_input_ids.shape != (1, seq_len): - return { - "success": False, - "error": f"Single sequence shape mismatch: expected (1, {seq_len}), got {packed_input_ids.shape}", - } - - expected_cu_seqlens = torch.tensor( - [0, seq_len], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens, expected_cu_seqlens): - return { - "success": False, - "error": f"Single sequence cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", - } - - # Test 2: Empty sequences (length 0) - batch_size = 3 - max_seq_len = 5 - input_ids = torch.randint( - 0, vocab_size, (batch_size, max_seq_len), device="cuda" - ) - seq_lengths = torch.tensor([3, 0, 2], device="cuda") - - packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( - _pack_sequences_for_megatron( - input_ids, seq_lengths, cp_rank=0, cp_size=1 - ) - ) - - # Should handle empty sequences gracefully - expected_total_tokens = 5 # 3 + 0 + 2 - if packed_input_ids.shape != (1, expected_total_tokens): - return { - "success": False, - "error": f"Empty sequence shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", - } - - expected_cu_seqlens = torch.tensor( - [0, 3, 3, 5], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens, expected_cu_seqlens): - return { - "success": False, - "error": f"Empty sequence cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", - } - - # Test 3: Large padding values - batch_size = 2 - seq_len = 4 - input_ids = torch.randint( - 0, vocab_size, (batch_size, seq_len), device="cuda" - ) - seq_lengths = torch.tensor([3, 2], device="cuda") - - packed_input_ids, _, packed_seq_params, cu_seqlens, cu_seqlens_padded = ( - _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=8, - cp_rank=0, - cp_size=1, - ) - ) - - # With padding to multiple of 8: [3, 2] -> [8, 8] = 16 tokens - expected_total_tokens = 16 - if packed_input_ids.shape != (1, expected_total_tokens): - return { - "success": False, - "error": f"Large padding shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids.shape}", - } - - return {"success": True, "error": None} - - except Exception as e: - return {"success": False, "error": f"Edge cases test failed: {str(e)}"} - - def _test_context_parallel( - self, _pack_sequences_for_megatron, _get_tokens_on_this_cp_rank, rank - ): - """Test sequence packing with context parallelism.""" - # Test parameters - batch_size = 2 - seq_len = 16 # Ensure divisible by cp_size * 2 - vocab_size = 100 - - # Ensure sequence length is compatible with CP - if seq_len % (2 * self.cp_size) != 0: - seq_len = (seq_len // (2 * self.cp_size) + 1) * (2 * self.cp_size) - - # Create test data - torch.manual_seed(42) # For reproducibility - input_ids = torch.arange(seq_len * batch_size, device="cuda").reshape( - batch_size, seq_len - ) - seq_lengths = torch.tensor([seq_len, seq_len], device="cuda") - - # Test 1: CP packing with individual sequence padding - ( - packed_input_ids, - packed_input_ids_cp_sharded, - packed_seq_params, - cu_seqlens, - cu_seqlens_padded, - ) = _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=self.cp_size * 2, - cp_rank=rank, - cp_size=self.cp_size, - ) - - # Verify the packed tensor shape - expected_tokens_per_rank = seq_len // self.cp_size - expected_total_tokens = batch_size * expected_tokens_per_rank - if packed_input_ids_cp_sharded.shape != (1, expected_total_tokens): - return { - "success": False, - "error": f"CP packing shape mismatch: expected (1, {expected_total_tokens}), got {packed_input_ids_cp_sharded.shape}", - } - - # Verify cu_seqlens for original sequences - expected_cu_seqlens = torch.tensor( - [0, seq_len, seq_len * 2], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens, expected_cu_seqlens): - return { - "success": False, - "error": f"CP cu_seqlens mismatch: expected {expected_cu_seqlens}, got {cu_seqlens}", - } - - # Verify PackedSeqParams - if packed_seq_params.qkv_format != "thd": - return { - "success": False, - "error": f"CP wrong qkv_format: expected 'thd', got {packed_seq_params.qkv_format}", - } - - # Test 2: CP packing with full sequence padding - pad_full_seq_to = (batch_size * seq_len) + 8 # Add some padding - ( - packed_input_ids_full, - packed_input_ids_cp_sharded, - packed_seq_params_full, - cu_seqlens_full, - cu_seqlens_padded_full, - ) = _pack_sequences_for_megatron( - input_ids, - seq_lengths, - pad_individual_seqs_to_multiple_of=self.cp_size * 2, - pad_packed_seq_to=pad_full_seq_to, - cp_rank=rank, - cp_size=self.cp_size, - ) - - # Verify the packed tensor shape with full padding - expected_tokens_per_rank_full = pad_full_seq_to // self.cp_size - if packed_input_ids_cp_sharded.shape != (1, expected_tokens_per_rank_full): - return { - "success": False, - "error": f"CP full padding shape mismatch: expected (1, {expected_tokens_per_rank_full}), got {packed_input_ids_cp_sharded.shape}", - } - - # Verify cu_seqlens_padded for full padding - expected_cu_seqlens_padded_full = torch.tensor( - [0, seq_len, pad_full_seq_to], device="cuda", dtype=torch.int32 - ) - if not torch.equal(cu_seqlens_padded_full, expected_cu_seqlens_padded_full): - return { - "success": False, - "error": f"CP full padding cu_seqlens_padded mismatch: expected {expected_cu_seqlens_padded_full}, got {cu_seqlens_padded_full}", - } - - correct_ids_0 = torch.tensor( - [0, 1, 2, 3, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 0, 0, 0, 0, 0, 0], - device="cuda", - ) - correct_ids_1 = torch.tensor( - [4, 5, 6, 7, 8, 9, 10, 11, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 0, 0], - device="cuda", - ) - - if ( - rank == 0 - and torch.sum(torch.abs(packed_input_ids_cp_sharded - correct_ids_0)).item() - != 0 - ): - return { - "success": False, - "error": f"CP full padding ids mismatch: expected {correct_ids_0}, got {packed_input_ids_cp_sharded[0, :20]}", - } - if ( - rank == 1 - and torch.sum(torch.abs(packed_input_ids_cp_sharded - correct_ids_1)).item() - != 0 - ): - return { - "success": False, - "error": f"CP full padding ids mismatch: expected {correct_ids_1}, got {packed_input_ids_cp_sharded[0, 20:]}", - } - - return {"success": True, "error": None} - - PACK_SEQUENCES_TEST_ACTOR_FQN = ( f"{PackSequencesTestActor.__module__}.PackSequencesTestActor" ) @@ -1311,425 +774,6 @@ def test_pack_sequences_with_context_parallel(pack_sequences_setup): print(f" Error: {test_result['error']}") -@ray.remote(num_gpus=1) -class GetPackSequenceParametersTestActor: - def __init__(self): - pass - - def run_all_get_pack_sequence_parameters_for_megatron_tests(self): - """Test _get_pack_sequence_parameters_for_megatron function with various configurations.""" - from nemo_rl.models.megatron.data import ( - _get_pack_sequence_parameters_for_megatron, - ) - - # Test 1: Basic configuration - no parallelism, no FP8 - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - } - max_seq_len = 1023 - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 2: Context parallelism only - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 4, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 4 * 2 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=4*2, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 3: Tensor parallelism with sequence parallelism - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 2 # tp_size when SP is enabled - if pad_individual != 2 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=2, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 4: Tensor parallelism without sequence parallelism - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 5: Pipeline parallelism - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 4, - "context_parallel_size": 1, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to != max_seq_len: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 6: Combined CP and TP with SP - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 4, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size - if ( - pad_individual != expected_individual - or pad_packed != 1 - or pad_to is not None - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 7: FP8 enabled with default recipe - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - "fp8_cfg": { - "enabled": True, - "fp8": "hybrid", - "fp8_recipe": "tensorwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 16 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=16, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 8: FP8 enabled with blockwise recipe - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "blockwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 128 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=128, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 9: FP8 with CP and TP+SP - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 4, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "blockwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size - expected_packed = 128 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size - if ( - pad_individual != expected_individual - or pad_packed != expected_packed - or pad_to is not None - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 10: All parallelism types with FP8 and PP - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 4, - "context_parallel_size": 2, - "fp8_cfg": { - "enabled": True, - "fp8": "hybrid", - "fp8_recipe": "tensorwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 2 * 2 * 2 # cp_size * 2 * tp_size - expected_packed = 16 * 2 * 2 * 2 # divisor * cp_size * 2 * tp_size - - def _round_up_to_multiple_of(x, y): - return (x + y - 1) // y * y - - if ( - pad_individual != expected_individual - or pad_packed != expected_packed - or pad_to != _round_up_to_multiple_of(max_seq_len, expected_packed) - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 11: FP8 disabled explicitly - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - "fp8_cfg": { - "enabled": False, - "fp8": "e4m3", - "fp8_recipe": "blockwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 12: Missing fp8_cfg (should default to disabled) - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - # No fp8_cfg key - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 13: Edge case - very large parallelism values - megatron_cfg = { - "tensor_model_parallel_size": 8, - "sequence_parallel": True, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 8, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "blockwise", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 8 * 2 * 8 # cp_size * 2 * tp_size = 128 - expected_packed = 128 * 8 * 2 * 8 # divisor * cp_size * 2 * tp_size = 16384 - if ( - pad_individual != expected_individual - or pad_packed != expected_packed - or pad_to is not None - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 14: Edge case - different max_seq_len values with PP - for test_seq_len in [512, 2048, 4096]: - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 2, - "context_parallel_size": 1, - } - - pad_individual, pad_packed, pad_to = ( - _get_pack_sequence_parameters_for_megatron(megatron_cfg, test_seq_len) - ) - - if pad_individual != 1 or pad_packed != 1 or pad_to != test_seq_len: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=1, pad_to={test_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 15: FP8 with MXFP8 recipe - megatron_cfg = { - "tensor_model_parallel_size": 1, - "sequence_parallel": False, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 1, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "mxfp8", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - if pad_individual != 1 or pad_packed != 32 or pad_to is not None: - return { - "success": False, - "error": f"Expected pad_individual=1, pad_packed=32, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 16: FP8 with MXFP8 recipe, CP, and TP+SP - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 1, - "context_parallel_size": 4, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "mxfp8", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size - expected_packed = 32 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size - - if ( - pad_individual != expected_individual - or pad_packed != expected_packed - or pad_to is not None - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to=None, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - # Test 17: FP8 with MXFP8 recipe, CP, TP+SP, and PP - megatron_cfg = { - "tensor_model_parallel_size": 2, - "sequence_parallel": True, - "pipeline_model_parallel_size": 4, - "context_parallel_size": 4, - "fp8_cfg": { - "enabled": True, - "fp8": "e4m3", - "fp8_recipe": "mxfp8", - "fp8_param": False, - }, - } - - pad_individual, pad_packed, pad_to = _get_pack_sequence_parameters_for_megatron( - megatron_cfg, max_seq_len - ) - - expected_individual = 4 * 2 * 2 # cp_size * 2 * tp_size - expected_packed = 32 * 4 * 2 * 2 # divisor * cp_size * 2 * tp_size * pp_size - - if ( - pad_individual != expected_individual - or pad_packed != expected_packed - or pad_to != _round_up_to_multiple_of(max_seq_len, expected_packed) - ): - return { - "success": False, - "error": f"Expected pad_individual={expected_individual}, pad_packed={expected_packed}, pad_to={max_seq_len}, got pad_individual={pad_individual}, pad_packed={pad_packed}, pad_to={pad_to}", - } - - return {"success": True, "error": None} - - GET_PACK_SEQUENCE_PARAMETERS_TEST_ACTOR_FQN = f"{GetPackSequenceParametersTestActor.__module__}.GetPackSequenceParametersTestActor" From 9527d565dbf72db82500c8b4bd73fbc2cfaf82ff Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 17:04:13 +0200 Subject: [PATCH 17/23] temporarily disabled test_megatron_checkpoint_save_kill_and_restore due to upstream bug Signed-off-by: Ahmad Kiswani --- tests/unit/models/policy/test_megatron_worker.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 7d329ab411..6386978754 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -1087,6 +1087,12 @@ def test_megatron_reference_policy_functionality(tiny_llama_model_path): @pytest.mark.timeout(400) @pytest.mark.hf_gated +@pytest.mark.xfail( + reason="MCore DistribOptimizer._set_main_param_and_optimizer_states passes non-tensor " + "optimizer state (e.g. bool) to TE FusedAdam.set_scaled_state which expects a tensor. " + "Needs fix in MCore distrib_optimizer.py.", + strict=False, +) @pytest.mark.parametrize( "num_gpus,tp,pp", [ From 8d0d33ecd1b7632a57984488a42dbbeccd269e6b Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Sun, 22 Feb 2026 17:29:42 +0200 Subject: [PATCH 18/23] temp: don't exit on first failure. Signed-off-by: Ahmad Kiswani --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 7b1ced085c..34006eb0ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -350,7 +350,7 @@ exclude = ''' ''' [tool.pytest.ini_options] -addopts = "--durations=15 -s -rA -x" +addopts = "--durations=15 -s -rA" testpaths = ["tests"] python_files = "test_*.py" markers = [ From 8a09bf800245671bc7b59c08c90fe1545d17ec38 Mon Sep 17 00:00:00 2001 From: Ahmad Kiswani Date: Mon, 23 Feb 2026 03:15:51 +0200 Subject: [PATCH 19/23] switched to a fork of mcore with fixed dist optim Signed-off-by: Ahmad Kiswani --- .gitmodules | 4 ++-- 3rdparty/Megatron-LM-workspace/Megatron-LM | 2 +- tests/unit/models/policy/test_megatron_worker.py | 6 ------ 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/.gitmodules b/.gitmodules index 8d7c7be7e5..86c2f47064 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM - url = https://github.com/NVIDIA/Megatron-LM.git - branch = main + url = https://github.com/ahmadki/Megatron-LM.git + branch = ahmadki/dist_optim_non_tensor_fix shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index a6d6dc6a85..e0dd64251f 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit a6d6dc6a853e94cf222881f9d67084383ddf5b65 +Subproject commit e0dd64251f2fde58606c5253280469b4dc81c75b diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 6386978754..7d329ab411 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -1087,12 +1087,6 @@ def test_megatron_reference_policy_functionality(tiny_llama_model_path): @pytest.mark.timeout(400) @pytest.mark.hf_gated -@pytest.mark.xfail( - reason="MCore DistribOptimizer._set_main_param_and_optimizer_states passes non-tensor " - "optimizer state (e.g. bool) to TE FusedAdam.set_scaled_state which expects a tensor. " - "Needs fix in MCore distrib_optimizer.py.", - strict=False, -) @pytest.mark.parametrize( "num_gpus,tp,pp", [ From 5e4a61f57a006cddbed168ed5db895ae92769f79 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Sun, 22 Feb 2026 21:20:07 -0800 Subject: [PATCH 20/23] Latest changes to remove inference_mode context manager --- examples/configs/grpo_math_1B_megatron.yaml | 1 + nemo_rl/models/megatron/setup.py | 4 + .../policy/workers/megatron_policy_worker.py | 156 +++++++----------- 3 files changed, 63 insertions(+), 98 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index f9eebea735..3a98abb283 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -105,6 +105,7 @@ policy: cuda_graph_scope: null use_te_rng_tracker: true inference_rng_tracker: true + batch_invariant_mode: false optimizer: optimizer: "adam" diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index e811f711a5..ce03878307 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -377,6 +377,10 @@ def _apply_cuda_graph_and_rng_tracker_config(model_cfg: Any, config: PolicyConfi model_cfg.cuda_graph_scope = config["megatron_cfg"]["cuda_graph_scope"] model_cfg.use_te_rng_tracker = config["megatron_cfg"]["use_te_rng_tracker"] model_cfg.inference_rng_tracker = config["megatron_cfg"]["inference_rng_tracker"] + model_cfg.batch_invariant_mode = config["megatron_cfg"]["batch_invariant_mode"] + if model_cfg.batch_invariant_mode: + from megatron.core.transformer.enums import AttnBackend + model_cfg.attention_backend = AttnBackend.flash def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c7d3e2feec..cc232aa9dd 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -876,92 +876,6 @@ async def _wake_engine(self): await self.dynamic_inference_engine.running.wait() self.dynamic_inference_engine.resume() - @contextmanager - def inference_mode(self, mcore_generation_config: dict): - """Context manager for inference mode, following Megatron RL's pattern. - - This mirrors megatron_rl_inference_mode from megatron/rl/rl_utils.py - - ENTER order: - 1. Put model in eval mode - 2. Clear rotary cache - 3. Toggle CUDA graphs ON - 4. Initialize inference engine (first time only) - 5. Resume engine (reallocates KV cache, recreates CUDA graphs as needed) - - EXIT order: - 1. Suspend engine (deallocates KV cache and GPU state) - 2. Toggle CUDA graphs OFF - 3. Clear rotary cache - 4. Put model back in train mode - - KV cache lifecycle is managed by the engine's suspend/resume mechanism - via KVCacheManagementMode in InferenceConfig. - - Yields: - The dynamic inference engine for use during inference. - """ - # Get the language module (unwrap from precision wrappers if needed) - lang_module = self._get_lang_module() - - # Get config settings - cuda_graph_impl = mcore_generation_config.get("cuda_graph_impl", "local") - - # Save training state - was_training = lang_module.training - - # === ENTER INFERENCE MODE === - - # 1. Put model in eval mode - lang_module.eval() - - # 2. Clear rotary position embedding caches (Megatron RL does this) - rotary_module = getattr(lang_module, "rotary_pos_emb", None) - has_lru_cache = rotary_module is not None and hasattr(rotary_module.forward, "cache_parameters") - if has_lru_cache: - rotary_module.forward.cache_clear() - - if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) - - # 4. Initialize inference engine if not already done - if not self._inference_engine_initialized: - self._initialize_inference_engine(mcore_generation_config) - # Start the coordinator and engine loop (first time only) - coordinator_port = self.cfg["generation"].get( - "inference_coordinator_port", 5995 - ) - self._run_async_coordinator_start(coordinator_port) - - if self._inference_engine_alseep: - self._wake() - - try: - # Yield the inference engine for use - yield self.dynamic_inference_engine - - finally: - - # 1. pause the inference engine - if self._inference_engine_initialized and not self._inference_engine_alseep: - self._sleep() - - # 2. Toggle CUDA graphs OFF - if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to="none") - - # 3. Clear rotary embedding cache again (Megatron RL does this on exit too) - if has_lru_cache: - rotary_module.forward.cache_clear() - - # 4. Restore training state - if was_training: - lang_module.train() - - # 6. Force garbage collection and CUDA memory cleanup - gc.collect() - torch.cuda.empty_cache() - @wrap_with_nvtx_name("megatron_policy_worker/generate") def generate( @@ -990,8 +904,23 @@ def generate( - logprobs: Log probabilities for each token - generation_lengths: Lengths of each response """ + from megatron.core.inference.sampling_params import SamplingParams + def _log_gpu_memory(tag: str): + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + allocated = torch.cuda.memory_allocated() / (1024 ** 3) + reserved = torch.cuda.memory_reserved() / (1024 ** 3) + free, total = torch.cuda.mem_get_info() + free_gb, total_gb = free / (1024 ** 3), total / (1024 ** 3) + print( + f"[GPU Rank {rank}] {tag} | " + f"Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB, " + f"Free: {free_gb:.2f} GB, Total: {total_gb:.2f} GB" + ) + + _log_gpu_memory("generate START") + self.model.config.flash_decode = False if self.should_disable_forward_pre_hook: self.model = self.move_model( @@ -1017,10 +946,32 @@ def generate( mcore_generation_config = self.cfg["generation"]["mcore_generation_config"] - # Use inference_mode context manager (mirrors megatron_rl_inference_mode from Megatron RL) - # This handles: eval mode, CUDA graph toggle, engine init/resume, and cleanup - with torch.no_grad(), self.inference_mode(mcore_generation_config) as inference_engine: - # Handle None values for top_k - convert to integer as required by Megatron + + lang_module = self._get_lang_module() + cuda_graph_impl = mcore_generation_config.get("cuda_graph_impl", "local") + was_training = lang_module.training + + lang_module.eval() + rotary_module = getattr(lang_module, "rotary_pos_emb", None) + has_lru_cache = rotary_module is not None and hasattr(rotary_module.forward, "cache_parameters") + if has_lru_cache: + rotary_module.forward.cache_clear() + + with torch.no_grad(): + + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) + + if not self._inference_engine_initialized: + self._initialize_inference_engine(mcore_generation_config) + coordinator_port = self.cfg["generation"].get( + "inference_coordinator_port", 5995 + ) + self._run_async_coordinator_start(coordinator_port) + + if self._inference_engine_alseep: + self._wake() + top_k_cfg = self.cfg["generation"]["top_k"] top_k_val = 1 if greedy else (int(top_k_cfg) if top_k_cfg is not None else 0) @@ -1040,7 +991,6 @@ def generate( termination_id=self.megatron_tokenizer.eod, ) - # Only rank 0 has actual data to submit if is_request_submitter: input_ids = data["input_ids"] print(f"[Rank {dist_rank}] input_ids: {input_ids.shape}") @@ -1051,24 +1001,33 @@ def generate( prompt_tokens_tensor = torch.empty(0, dtype=torch.long, device="cuda") prompt_lengths_tensor = torch.empty(0, dtype=torch.long, device="cuda") - # Run the coordinator-based generation using the persistent engine - # Rank 0 submits requests, other ranks participate in engine loop - # Results are broadcast to all ranks inside this method result = self._run_async_generation_with_persistent_engine( prompt_tokens_tensor, prompt_lengths_tensor, sampling_params, ) - self.model.config.flash_decode = False + if self._inference_engine_initialized and not self._inference_engine_alseep: + self._sleep() + + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to="none") + + if has_lru_cache: + rotary_module.forward.cache_clear() + + if was_training: + lang_module.train() - # Context manager has exited - CUDA graphs are now disabled, model is back in train mode + gc.collect() + torch.cuda.empty_cache() + + self.model.config.flash_decode = False # Only rank 0 needs to format and return results # Other ranks return None (their results are ignored due to output_is_replicated) if not is_request_submitter: - # Return empty result for non-submitter ranks - # Use BatchedDataDict directly instead of from_batches to avoid padding issues with empty tensors + _log_gpu_memory("generate END (non-submitter)") return BatchedDataDict({ "output_ids": torch.empty(0, 0, dtype=torch.long), "logprobs": torch.empty(0, 0, dtype=torch.float), @@ -1124,6 +1083,7 @@ def generate( "unpadded_sequence_lengths": unpadded_sequence_lengths, } + _log_gpu_memory("generate END") return BatchedDataDict.from_batches([out_dict]).to("cpu") def _start_inference_loop_thread(self): From 5258e35b967e7127ec523178ce1706b0b132a72f Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Tue, 24 Feb 2026 11:57:45 -0800 Subject: [PATCH 21/23] nccl timeout issue --- nccl_timeout.md | 156 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 nccl_timeout.md diff --git a/nccl_timeout.md b/nccl_timeout.md new file mode 100644 index 0000000000..052ea4e81a --- /dev/null +++ b/nccl_timeout.md @@ -0,0 +1,156 @@ +# NCCL Timeout During CUDA Graph Warmup in MoE RL Training + +## Symptom + +After several successful GRPO training steps (anywhere from step 5 to step 20+), the job crashes with NCCL collective operation timeouts during the generation phase. The errors look like: + +``` +[Rank 5] Watchdog caught collective operation timeout: + WorkNCCL(SeqNum=10707, OpType=COALESCED, ..., Timeout(ms)=600000) + ran for 600030 milliseconds before timing out. +``` + +Key signatures: +- Different ranks report **different NCCL operation types** (ALLTOALL_BASE, REDUCE_SCATTER_BASE, ALLREDUCE, COALESCED) -- a collective mismatch +- The crash always happens during `cuda graph warmup` at the start of generation +- A new NCCL communicator is lazily initialized at the failing step (`NCCL version 2.27.5+cuda12.9` printed mid-warmup) +- Steps 1 through N-1 complete normally; the crash is non-deterministic + +## Background: The Training-Inference Cycle + +In the RL training loop, each step does: + +``` +generate() { + _wake() // resume inference engine (realloc KV cache, rebuild CUDA graphs) + + _sleep() // suspend inference engine (dealloc KV cache, delete CUDA graphs) +} +``` + +With `static_kv_memory_pointers=false` and `kv_cache_management_mode=recompute`, every suspend/resume cycle **destroys and recreates CUDA graphs**. Graph warmup runs forward passes through the model, which for MoE models includes NCCL alltoall collectives across expert-parallel (EP) ranks. All EP/TP ranks must execute the same sequence of NCCL operations in lockstep during this warmup. + +## Architecture: Two Communication Systems on One Event Loop + +The `DynamicInferenceEngine` runs an async engine loop on a dedicated event loop thread. This single event loop handles two different communication systems: + +| System | Purpose | Mechanism | +|--------|---------|-----------| +| **EP consensus** (`_ep_group_has_work`) | Coordinate EP ranks on work availability | Async ZMQ all-reduce | +| **CUDA graph warmup** (inside `resume()`) | Capture model forward passes into graphs | Blocking NCCL collectives | + +Both run on the **same event loop thread**. This is the root of the problem. + +## Root Cause + +The engine loop has this structure (simplified from `run_engine_with_coordinator`): + +```python +while True: + self.schedule_requests() # read ZMQ messages + ep_group_has_work = await self._ep_group_has_work(...) # ZMQ all-reduce across EP ranks + if not ep_group_has_work: + if self.suspend_signal: + self.suspend() # no-op when already suspended + else: + self.resume() # CUDA graph warmup -- blocks with NCCL! + await asyncio.sleep(0.02) +``` + +When the coordinator sends `RESUME + UNPAUSE` to all engines, the signals arrive asynchronously. EP ranks process them at different times depending on ZMQ delivery and event loop scheduling. This leads to a **divergence**: + +``` +Rank A (received RESUME): suspend_signal=False --> calls resume() --> NCCL alltoall BLOCKS event loop +Rank B (not yet received): suspend_signal=True --> calls suspend() (no-op) --> sleeps 20ms +``` + +On the next iteration, Rank B calls `_ep_group_has_work()` which does an async ZMQ all-reduce. This requires Rank A to respond. But Rank A's event loop is **blocked inside NCCL** (graph warmup forward pass). Rank A can never respond to ZMQ while NCCL is blocking its event loop. + +**Deadlock: Rank A waits for Rank B in NCCL. Rank B waits for Rank A in ZMQ.** + +After 10 minutes, the NCCL watchdog times out and kills the process. + +### Why it's non-deterministic + +The deadlock only occurs when at least one EP rank enters `resume()` before all other EP ranks have received the `RESUME` signal. When all ranks happen to process the signals within the same ~20ms engine loop cycle, they all enter `resume()` together and the warmup succeeds. This timing depends on ZMQ delivery, event loop scheduling, and OS thread scheduling -- hence the non-determinism. + + +### Implementation 1 (This causes delay of 25%) + +```python +def _wake(self): + # Phase 1: Unpause the engine loop (async, event loop stays free for ZMQ) + asyncio.run_coroutine_threadsafe(self._unpause_engine(), self._inference_loop).result() + + # Phase 2: Synchronized resume on the main thread + self._synchronized_resume() + +async def _unpause_engine(self): + # Send only UNPAUSE (not RESUME) -- keeps suspend_signal=True so the engine + # loop never calls resume() on its own + if torch.distributed.get_rank() == 0: + self.inference_client.unpause_engines() + await self.dynamic_inference_engine.running.wait() + +def _synchronized_resume(self): + engine = self.dynamic_inference_engine + + # Guard: replace suspend() with a no-op while we resume + original_suspend = engine.suspend + engine.suspend = lambda: None + + try: + torch.distributed.barrier() # all ranks ready + engine.resume() # CUDA graph warmup (NCCL collectives) + engine.suspend_signal = False # let engine loop transition to normal mode + torch.distributed.barrier() # all ranks done + finally: + engine.suspend = original_suspend +``` + +### Why this works + +**No event-loop blocking.** The NCCL barriers and `resume()` run on the main thread. The event loop thread continues running the engine loop, freely handling ZMQ communication for EP consensus. No rank's ZMQ is ever starved. + +**No RESUME signal divergence.** We never send the `RESUME` header to the coordinator. Instead, we send only `UNPAUSE` (which restarts the engine loop) and keep `suspend_signal=True`. The engine loop sees `suspend_signal=True`, calls `suspend()` (no-op since already suspended), and idles. It never calls `resume()` on its own. We control exactly when `resume()` happens -- after the barrier on the main thread. + +**No thread-safety race.** When `resume()` runs on the main thread, it sets `is_suspended=False`. Without the guard, the engine loop's next `suspend()` call (on the event loop thread) would read `is_suspended=False`, enter the suspend body, and **deallocate buffers while the main thread is still creating CUDA graphs**. The `suspend()` guard (replacing it with `lambda: None`) prevents this. The guard is removed only after `suspend_signal=False` is set, so the engine loop transitions to calling `resume()` (which is a no-op since we already resumed) instead of `suspend()`. + +### Thread interaction timeline + +``` +Main Thread Event Loop Thread (engine loop) +----------- -------------------------------- +_unpause_engine() ----sends UNPAUSE----> + schedule_requests(): reads UNPAUSE + _ep_group_has_work(): ZMQ all-reduce + suspend_signal=True -> suspend() [no-op] + asyncio.sleep(0.02) + +barrier() ............all ranks sync.... (ZMQ continues running freely) + +engine.suspend = no-op suspend() -> no-op [guarded] +engine.resume() (ZMQ continues, no GPU conflict) + -> reinitialize buffers + -> create_cuda_graphs() [NCCL] +engine.suspend_signal = False (ZMQ continues) + +barrier() ............all ranks done.... + +engine.suspend = original suspend_signal=False -> resume() [no-op] + (engine is ready for requests) +``` + +## Affected Configuration + +This bug affects MoE models using the megatron generation backend with: +- `moe_token_dispatcher_type=alltoall` (NCCL alltoall inside CUDA graph warmup) +- `static_kv_memory_pointers=false` (CUDA graphs deleted/recreated each cycle) +- `kv_cache_management_mode=recompute` (full dealloc on suspend) +- `num_cuda_graphs > 0` +- Expert parallelism (EP) > 1 + +Dense models or configs with `static_kv_memory_pointers=true` are not affected because CUDA graphs are not recreated on resume. + +### Implementation 2 +In dynamic_engine.py you set asyncio.sleep(0) instead of 0.02. This works \ No newline at end of file From 3f935501f92e182a1a1f6d80d643f47f253e7f0e Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy Date: Wed, 25 Feb 2026 17:55:49 -0800 Subject: [PATCH 22/23] Using main branch --- .gitmodules | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitmodules b/.gitmodules index e6f53c03ad..832216f926 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/Megatron-LM"] path = 3rdparty/Megatron-LM-workspace/Megatron-LM url = https://github.com/shanmugamr1992/Megatron-LM.git - branch = fixes_latest + branch = main shallow = true [submodule "3rdparty/Megatron-Bridge"] path = 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge From a4f1dcda396ba7b9dae2fda4db084027243613ed Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Sat, 28 Feb 2026 07:45:08 -0800 Subject: [PATCH 23/23] Update nccl_timeout.md Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> --- nccl_timeout.md | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/nccl_timeout.md b/nccl_timeout.md index 052ea4e81a..86f5250a80 100644 --- a/nccl_timeout.md +++ b/nccl_timeout.md @@ -5,9 +5,14 @@ After several successful GRPO training steps (anywhere from step 5 to step 20+), the job crashes with NCCL collective operation timeouts during the generation phase. The errors look like: ``` -[Rank 5] Watchdog caught collective operation timeout: - WorkNCCL(SeqNum=10707, OpType=COALESCED, ..., Timeout(ms)=600000) - ran for 600030 milliseconds before timing out. +(MegatronPolicyWorker[rank=1] pid=151407) [rank1]:[E228 00:10:27.202494521 ProcessGroupNCCL.cpp:2057] [PG ID 12 PG GUID 99(EXPERT_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1145527, OpType=ALLTOALL_BASE, NumelIn=206438400, NumelOut=206438400, Timeout(ms)=600000) ran for 600016 milliseconds before timing out. +(MegatronPolicyWorker[rank=1] pid=151407) +(MegatronPolicyWorker[rank=1] pid=151407) [2026-02-28 00:10:27,258 E 151407 153013] logging.cc:118: Unhandled exception: N3c1016DistBackendErrorE. what(): [PG ID 12 PG GUID 99(EXPERT_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1145527, OpType=ALLTOALL_BASE, NumelIn=206438400, NumelOut=206438400, Timeout(ms)=600000) ran for 600016 milliseconds before timing out. +(MegatronPolicyWorker[rank=1] pid=151407) +(MegatronPolicyWorker[rank=1] pid=151407) +(MegatronPolicyWorker[rank=7] pid=151429) [rank7]:[E228 00:10:27.168785312 ProcessGroupNCCL.cpp:2057] [PG ID 5 PG GUID 36(TENSOR_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=2036074, OpType=_REDUCE_SCATTER_BASE, NumelIn=3225600, NumelOut=1612800, Timeout(ms)=600000) ran for 600000 milliseconds before timing out. +(MegatronPolicyWorker[rank=7] pid=151429) +(MegatronPolicyWorker[rank=7] pid=151429) [2026-02-28 00:10:27,224 E 151429 152962] logging.cc:118: Unhandled exception: N3c1016DistBackendErrorE. what(): [PG ID 5 PG GUID 36(TENSOR_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=2036074, OpType=_REDUCE_SCATTER_BASE, NumelIn=3225600, NumelOut=1612800, Timeout(ms)=600000) ran for 600000 milliseconds before timing out. ``` Key signatures: @@ -153,4 +158,4 @@ This bug affects MoE models using the megatron generation backend with: Dense models or configs with `static_kv_memory_pointers=true` are not affected because CUDA graphs are not recreated on resume. ### Implementation 2 -In dynamic_engine.py you set asyncio.sleep(0) instead of 0.02. This works \ No newline at end of file +In dynamic_engine.py you set asyncio.sleep(0) instead of 0.02. This works