diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 95110681d..f8b6878be 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -42,7 +42,8 @@ env: # Includes fused expert wgrad output overwrite support (#534), MegaMoE-owned FlyDSL # kernels (#525), the repaired DeepEP intranode cheap fence, on by default (#539), # and MegaMoE SwiGLU-OAI activations (#538). - PRIMUS_TURBO_COMMIT: 8f19706aa1c656b11a146eb00db32a6c2b46822d + # Turbo main includes merged MoE gather fusion (#546) and fused cross entropy (#550). + PRIMUS_TURBO_COMMIT: 4e6bd90b4fd9aa40c7945674d83975ee150a7b68 PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py #ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) diff --git a/docs/README.md b/docs/README.md index 62b5a03e2..02a053ac0 100644 --- a/docs/README.md +++ b/docs/README.md @@ -68,6 +68,7 @@ Deep technical topics for advanced users. - [Parallelism configuration](./04-technical-guides/parallelism-configuration.md): per-backend parallelism setup and batch size relationships - [Collective operations](./04-technical-guides/collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy - [Performance tuning](./04-technical-guides/performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization +- [Turbo cross entropy](./features/turbo_cross_entropy.md): TP1 fused loss and configuration - [MoE training deep-dive](./04-technical-guides/moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models - [MegaMoE fused MoE layer](./04-technical-guides/mega-moe.md): FlyDSL-based fused MoE layer for EP-only bf16 training, setup and reproduction - [Data preparation](./04-technical-guides/data-preparation.md): tokenization, data formats, mock data diff --git a/docs/features/turbo_cross_entropy.md b/docs/features/turbo_cross_entropy.md new file mode 100644 index 000000000..64eef1031 --- /dev/null +++ b/docs/features/turbo_cross_entropy.md @@ -0,0 +1,42 @@ +# Opt-in Turbo cross entropy for Megatron training + +Install a Primus-Turbo version with `primus_turbo.pytorch.ops.cross_entropy`, then +set the following Megatron training arguments in the experiment YAML: + +```yaml +use_turbo_cross_entropy: true +turbo_ce_overwrite_input: false +cross_entropy_loss_fusion: true +``` + +These options are declared in `primus/configs/modules/megatron/primus_turbo.yaml` +and can also be supplied as Primus overrides, for example +`--use_turbo_cross_entropy true --turbo_ce_overwrite_input false`. + +The patch replaces +Megatron's language-model loss for TP=1, BF16/FP32 logits when +`cross_entropy_loss_fusion` is enabled. TP>1, other dtypes, and an unfused loss +configuration use the existing configured implementation. + +The patch preserves Core's layout and reduction contract: sequence-first logits, +batch-first labels and output per-token loss. Training still applies its mask, +normalization and gradient scaling. Do not set `cross_entropy_fusion_impl=turbo`; +the opt-in patch provides the integration while retaining the original Core +configuration for fallback. + +`turbo_ce_overwrite_input: true` additionally permits backward to reuse +contiguous logits storage. Enable it only when the language-model logits have +no other consumer that needs their original values. Noncontiguous logits use +the preserving path. Both arguments default to false. Set +`use_turbo_cross_entropy: false` to restore the configured loss path without +changing packages. + +This integration changes no installed Transformer Engine files. It requires no +TE native rebuild; Triton JIT compilation occurs on first use. The operator +supports one backward per forward and no higher-order gradients. Language-model recipes need +numerical and convergence validation because its FP32 gradient arithmetic +differs from older TE's intermediate BF16 rounding. + +The former `PRIMUS_TURBO_CROSS_ENTROPY` and `PRIMUS_TURBO_CE_OVERWRITE_INPUT` +environment variables are no longer read by the patch. Launchers that retain +environment overrides must resolve them into these YAML arguments. diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index 24ebc2e94..92335eebb 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -90,6 +90,8 @@ subtrees: title: Data preparation - file: 04-technical-guides/checkpoint-management.md title: Checkpoint management + - file: features/turbo_cross_entropy.md + title: Turbo cross entropy - file: 04-technical-guides/moe-training.md title: MoE training - file: 04-technical-guides/mega-moe.md diff --git a/primus/backends/megatron/patches/turbo/cross_entropy_patches.py b/primus/backends/megatron/patches/turbo/cross_entropy_patches.py new file mode 100644 index 000000000..bcc7c6a69 --- /dev/null +++ b/primus/backends/megatron/patches/turbo/cross_entropy_patches.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. + +"""Opt-in Turbo TP=1 vocabulary CE without replacing installed TE files.""" + +from functools import wraps + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _enabled(ctx: PatchContext) -> bool: + return bool(getattr(get_args(ctx), "use_turbo_cross_entropy", False)) + + +@register_patch( + "megatron.turbo.cross_entropy", + backend="megatron", + phase="before_train", + description="Use Turbo fused CE with deferred backward for TP=1 BF16/FP32 logits", + condition=_enabled, +) +def patch_cross_entropy(ctx: PatchContext) -> None: + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + from primus_turbo.pytorch.ops.cross_entropy import cross_entropy + + original = LanguageModule.compute_language_model_loss + if getattr(original, "_primus_turbo_ce", False): + return + overwrite_input = bool(getattr(get_args(ctx), "turbo_ce_overwrite_input", False)) + + @wraps(original) + def compute_language_model_loss(self, labels, logits): + if ( + self.config.tensor_model_parallel_size != 1 + or logits.dtype not in (torch.bfloat16, torch.float32) + or not self.config.cross_entropy_loss_fusion + ): + return original(self, labels, logits) + # Core supplies sequence-first logits and batch-first targets. Keep + # the original per-token loss contract; training owns masks/scaling. + target = labels.transpose(0, 1).contiguous() + loss = cross_entropy( + logits, + target, + overwrite_input=overwrite_input and logits.is_contiguous(), + ) + return loss.transpose(0, 1).contiguous() + + compute_language_model_loss._primus_turbo_ce = True + LanguageModule.compute_language_model_loss = compute_language_model_loss + log_rank_0( + "[Patch:megatron.turbo.cross_entropy] Enabled Turbo CE for TP=1 BF16/FP32; " + f"overwrite_input={overwrite_input}; other cases retain the configured loss" + ) diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index 437685e0e..672b6d10c 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -103,3 +103,11 @@ moe_permute_quant_fusion: false moe_backward_permute_quant_fusion: false # Preserve activation storage across identity expert sorting. moe_skip_identity_sort: true + +# ===== Cross entropy ===== +# Use Turbo fused cross entropy for TP=1 BF16/FP32 logits when +# cross_entropy_loss_fusion is enabled. Other cases keep the configured loss. +use_turbo_cross_entropy: false +# Allow backward to overwrite contiguous logits with their gradient. Only enable +# when no other consumer needs the original logits. +turbo_ce_overwrite_input: false diff --git a/tests/trainer/test_megatron_trainer.py b/tests/trainer/test_megatron_trainer.py index 20b46cabf..0448e2d8d 100644 --- a/tests/trainer/test_megatron_trainer.py +++ b/tests/trainer/test_megatron_trainer.py @@ -7,23 +7,13 @@ import os import re -import socket import subprocess import sys import unittest import pytest -from tests.utils import PrimusUT, run_training_script - - -def _find_free_port() -> int: - """Ask the kernel for a free TCP port (bind to 0); avoids EADDRINUSE from - guessing a port that overlaps the ephemeral range or a not-yet-released one.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - +from tests.utils import PrimusUT, isolated_torchrun_env, run_training_script _GFX_TO_PLATFORM = { "gfx942": "MI300X", @@ -1031,11 +1021,7 @@ def _run_benchmark(self, suite: str, mode_flag: str) -> str: "--target-nodes", "2", ] - env = os.environ.copy() - # Kernel-assigned free port so back-to-back GPU benches don't hit - # EADDRINUSE (a fixed/random port can clash with a not-yet-released - # previous run or the ephemeral range). - env["MASTER_PORT"] = str(_find_free_port()) + env = isolated_torchrun_env() return self._check(subprocess.run(cmd, capture_output=True, text=True, env=env)) def test_performance_simulate_dense(self): diff --git a/tests/unit_tests/backends/megatron/test_pp_warmup_patches.py b/tests/unit_tests/backends/megatron/test_pp_warmup_patches.py index 9d9fd67f1..de7e2bb8a 100644 --- a/tests/unit_tests/backends/megatron/test_pp_warmup_patches.py +++ b/tests/unit_tests/backends/megatron/test_pp_warmup_patches.py @@ -56,6 +56,8 @@ from pathlib import Path from typing import List, Optional, Tuple +from tests.utils import isolated_torchrun_env + _MBS = 1 _GBS = 16 _PRIMUS_TOTAL_LAYERS = 4 @@ -210,7 +212,7 @@ def _build_training_cmd( "False", ] - env = os.environ.copy() + env = isolated_torchrun_env() env.update( { "MBS": str(_MBS), diff --git a/tests/unit_tests/test_torchrun_rendezvous.py b/tests/unit_tests/test_torchrun_rendezvous.py new file mode 100644 index 000000000..5df2fd6df --- /dev/null +++ b/tests/unit_tests/test_torchrun_rendezvous.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. + +"""Training test launches must not reuse a parent process's rendezvous port.""" + +import os +import re +import socket +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor + +from tests.utils import isolated_torchrun_env + + +def test_concurrent_launches_with_occupied_inherited_port(tmp_path): + worker = tmp_path / "worker.py" + worker.write_text( + "import os, torch\n" + "import torch.distributed as dist\n" + "dist.init_process_group('gloo')\n" + "value = torch.tensor([dist.get_rank() + int(os.environ['TEST_JOB']) * 10])\n" + "dist.all_reduce(value)\n" + "assert value.item() == 1 + int(os.environ['TEST_JOB']) * 20\n" + "print('DIST_PORT=' + os.environ['MASTER_PORT'], flush=True)\n" + "dist.destroy_process_group()\n" + ) + with socket.socket() as occupied: + occupied.bind(("127.0.0.1", 0)) + occupied.listen() + port = occupied.getsockname()[1] + parent_env = dict(os.environ, NNODES="1", MASTER_ADDR="127.0.0.1", MASTER_PORT=str(port)) + original_env = parent_env.copy() + + def launch(job): + env = isolated_torchrun_env(parent_env) + env["TEST_JOB"] = str(job) + return subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--nnodes=1", + "--nproc_per_node=2", + "--master_addr=127.0.0.1", + f"--master_port={port}", + str(worker), + ], + env=env, + capture_output=True, + text=True, + timeout=90, + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(launch, (0, 1))) + + for result in results: + assert result.returncode == 0, result.stdout + result.stderr + ports = re.findall(r"DIST_PORT=(\d+)", result.stdout) + assert len(ports) == 2, result.stdout + assert len(set(ports)) == 1 + assert ports[0] != str(port) + assert parent_env == original_env + + +def test_multinode_rendezvous_is_preserved(): + env = {"NNODES": "2", "MASTER_ADDR": "node0", "MASTER_PORT": "29500"} + assert isolated_torchrun_env(env) == env + + +def test_slurm_multinode_rendezvous_is_preserved(): + env = {"SLURM_NNODES": "2", "MASTER_ADDR": "node0", "MASTER_PORT": "29500"} + assert isolated_torchrun_env(env) == env diff --git a/tests/unit_tests/test_turbo_cross_entropy_config.py b/tests/unit_tests/test_turbo_cross_entropy_config.py new file mode 100644 index 000000000..8e8181de3 --- /dev/null +++ b/tests/unit_tests/test_turbo_cross_entropy_config.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. + +"""CPU dispatch tests; numerical GPU coverage lives in the integration suite.""" + +import importlib.util +import sys +from collections import defaultdict +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + + +def context(**kwargs): + return PatchContext( + backend="megatron", + phase="before_train", + extra={"module_config": SimpleNamespace(params=SimpleNamespace(**kwargs))}, + ) + + +@pytest.fixture +def patched(monkeypatch): + monkeypatch.setattr(PatchRegistry, "_all_patches", []) + monkeypatch.setattr(PatchRegistry, "_patches_by_backend_phase", defaultdict(lambda: defaultdict(list))) + path = Path(__file__).parents[2] / "primus/backends/megatron/patches/turbo/cross_entropy_patches.py" + spec = importlib.util.spec_from_file_location("ce_config_under_test", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + monkeypatch.setattr(mod, "log_rank_0", lambda *args: None) + original = Mock(return_value="fallback") + original._primus_turbo_ce = False + cls = type("LanguageModule", (), {"compute_language_model_loss": original}) + lm = ModuleType("megatron.core.models.common.language_module.language_module") + lm.LanguageModule = cls + monkeypatch.setitem(sys.modules, lm.__name__, lm) + ce = ModuleType("primus_turbo.pytorch.ops.cross_entropy") + ce.cross_entropy = Mock(side_effect=lambda logits, target, **kw: torch.zeros_like(target)) + monkeypatch.setitem(sys.modules, ce.__name__, ce) + return mod, cls, original, ce.cross_entropy + + +def test_explicit_arguments_override_legacy_environment(patched, monkeypatch): + mod, *_ = patched + monkeypatch.setenv("PRIMUS_TURBO_CROSS_ENTROPY", "1") + assert not mod._enabled(context()) + assert not mod._enabled(context(use_turbo_cross_entropy=False)) + monkeypatch.setenv("PRIMUS_TURBO_CROSS_ENTROPY", "0") + assert mod._enabled(context(use_turbo_cross_entropy=True)) + + +@pytest.mark.parametrize("overwrite", [False, True]) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_overwrite_argument_and_layout(patched, monkeypatch, overwrite, contiguous): + mod, cls, original, ce = patched + monkeypatch.setenv("PRIMUS_TURBO_CE_OVERWRITE_INPUT", str(int(not overwrite))) + mod.patch_cross_entropy(context(turbo_ce_overwrite_input=overwrite)) + installed = cls.compute_language_model_loss + mod.patch_cross_entropy(context(turbo_ce_overwrite_input=overwrite)) + assert cls.compute_language_model_loss is installed + logits = torch.empty(7, 3, 11) if contiguous else torch.empty(3, 7, 11).transpose(0, 1) + labels = torch.zeros(3, 7, dtype=torch.long) + model = SimpleNamespace( + config=SimpleNamespace(tensor_model_parallel_size=1, cross_entropy_loss_fusion=True) + ) + assert cls.compute_language_model_loss(model, labels, logits).shape == labels.shape + assert ce.call_args.kwargs["overwrite_input"] == (overwrite and contiguous) + torch.testing.assert_close(ce.call_args.args[1], labels.T.contiguous()) + original.assert_not_called() + + +@pytest.mark.parametrize( + "tp,fused,dtype", [(2, True, torch.float32), (1, False, torch.float32), (1, True, torch.float16)] +) +def test_fallback(patched, tp, fused, dtype): + mod, cls, original, ce = patched + mod.patch_cross_entropy(context()) + model = SimpleNamespace( + config=SimpleNamespace(tensor_model_parallel_size=tp, cross_entropy_loss_fusion=fused) + ) + assert ( + cls.compute_language_model_loss(model, torch.zeros(3, 7), torch.empty(7, 3, 11, dtype=dtype)) + == "fallback" + ) + original.assert_called_once() + ce.assert_not_called() diff --git a/tests/unit_tests/test_turbo_cross_entropy_patch.py b/tests/unit_tests/test_turbo_cross_entropy_patch.py new file mode 100644 index 000000000..15029860a --- /dev/null +++ b/tests/unit_tests/test_turbo_cross_entropy_patch.py @@ -0,0 +1,105 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. + +"""GPU integration contract for Core's batch/sequence layout and Turbo CE.""" + +import importlib.util +from collections import defaultdict +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from primus.core.patches import PatchContext + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") + + +def _context(**kwargs): + return PatchContext( + backend="megatron", + phase="before_train", + extra={"module_config": SimpleNamespace(params=SimpleNamespace(**kwargs))}, + ) + + +@pytest.fixture +def installed_patch(monkeypatch): + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + + from primus.core.patches.patch_registry import PatchRegistry + + # Loading the module registers its patch; keep the suite's global registry intact. + monkeypatch.setattr(PatchRegistry, "_all_patches", []) + monkeypatch.setattr( + PatchRegistry, + "_patches_by_backend_phase", + defaultdict(lambda: defaultdict(list)), + ) + calls = [] + + def original(self, labels, logits): + calls.append((labels, logits)) + return "fallback" + + monkeypatch.setattr(LanguageModule, "compute_language_model_loss", original) + path = Path(__file__).parents[2] / "primus/backends/megatron/patches/turbo/cross_entropy_patches.py" + spec = importlib.util.spec_from_file_location("ce_patch_under_test", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod, LanguageModule, calls + + +@pytest.mark.parametrize("overwrite", [False, True]) +def test_lm_head_backward_and_layout(installed_patch, monkeypatch, overwrite): + mod, cls, calls = installed_patch + mod.patch_cross_entropy(_context(turbo_ce_overwrite_input=overwrite)) + # Exercise the real linear->loss->linear backward ownership contract. + torch.manual_seed(312) + h = torch.randn(7, 3, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True) + w = torch.randn(1009, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True) + logits = F.linear(h, w) + ref_logits = logits.detach().float().requires_grad_() + labels = torch.randint(1009, (3, 7), device="cuda") + labels[0, 0] = -100 + weights = torch.rand(3, 7, device="cuda") + fake = SimpleNamespace( + config=SimpleNamespace(tensor_model_parallel_size=1, cross_entropy_loss_fusion=True) + ) + actual = cls.compute_language_model_loss(fake, labels, logits) + ref = ( + F.cross_entropy(ref_logits.flatten(0, 1), labels.T.contiguous().flatten(), reduction="none") + .reshape(7, 3) + .T.contiguous() + ) + torch.testing.assert_close(actual, ref, rtol=2e-5, atol=2e-5) + (actual * weights).sum().backward() + (dlogits,) = torch.autograd.grad(ref, ref_logits, weights) + expected_h = dlogits.to(h.dtype) @ w.detach() + expected_w = dlogits.to(h.dtype).flatten(0, 1).T @ h.detach().flatten(0, 1) + torch.testing.assert_close(h.grad, expected_h, rtol=0.008, atol=0.02) + torch.testing.assert_close(w.grad, expected_w, rtol=0.008, atol=0.02) + assert not calls + + +@pytest.mark.parametrize( + "tp,fused,dtype", + [(2, True, torch.bfloat16), (1, False, torch.bfloat16), (1, True, torch.float16)], +) +def test_fallback(installed_patch, tp, fused, dtype): + mod, cls, calls = installed_patch + mod.patch_cross_entropy(_context()) + first = cls.compute_language_model_loss + mod.patch_cross_entropy(_context()) + assert cls.compute_language_model_loss is first + fake = SimpleNamespace( + config=SimpleNamespace(tensor_model_parallel_size=tp, cross_entropy_loss_fusion=fused) + ) + logits = torch.empty(7, 3, 1009, device="cuda", dtype=dtype) + labels = torch.zeros(3, 7, device="cuda", dtype=torch.long) + assert cls.compute_language_model_loss(fake, labels, logits) == "fallback" + assert len(calls) == 1 diff --git a/tests/utils.py b/tests/utils.py index c285de0ed..2379401f5 100755 --- a/tests/utils.py +++ b/tests/utils.py @@ -12,12 +12,36 @@ import sys import time import unittest +import uuid from typing import Optional from primus.core.utils import logger TRAINING_COMPLETED_MARKER = "Training completed." + +def isolated_torchrun_env(env: Optional[dict] = None) -> dict: + """Give each single-node test launch its own kernel-assigned rendezvous port. + + Distributed fixtures can leave MASTER_PORT in the parent environment. A + bind/close probe also releases its chosen port before torchrun can use it. + Torchrun's c10d endpoint accepts port 0 and holds the socket it allocates, + avoiding both stale inherited ports and the probe-to-launch race. + """ + env = dict(os.environ if env is None else env) + if int(env.get("NNODES", env.get("SLURM_NNODES", "1"))) != 1: + return env + # PET_* are torchrun's native argument defaults. The production launcher + # and multi-node rendezvous settings do not need to change for these tests. + env.update( + PET_RDZV_BACKEND="c10d", + PET_RDZV_ENDPOINT="127.0.0.1:0", + PET_RDZV_ID=uuid.uuid4().hex, + PET_LOCAL_ADDR="127.0.0.1", + ) + return env + + _ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m") # Per-step metric lines, one pattern per backend log format: @@ -178,8 +202,8 @@ def run_training_script( # ~20s per end-to-end test. The marker we rely on ("Training completed.") # is emitted well before cleanup(), so this does not affect assertions. # Developers can opt out locally via PRIMUS_EXIT_FAST=0. - if env is not None: - env.setdefault("PRIMUS_EXIT_FAST", "1") + env = isolated_torchrun_env(env) + env.setdefault("PRIMUS_EXIT_FAST", "1") try: logger.info(f"[{tag}] Begin run...")