[guihong-nv]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Gives vocab-parallel cross entropy and normalization a home in megatron/core/ops, ops/loss and ops/norm, following the layout and rules of #7380 and of the SSM and attention moves in #7453 and #8348. Kernel integrations such as #5438 (Liger-Kernel RMSNorm and vocab-parallel cross entropy) then add an adapter module and a selector branch in these families instead of new files in fusions/ and extensions/. Every old module path remains as a deprecated forwarder.

Based on main, ahead of the SSM move (#7453). It has six commits and a merge of main:

1. ops: write patches through deprecated forwarders and forward re-exported names: Move SSM helpers, Mamba2, GDN/GDN2 and GDP into megatron/core/ops with deprecated aliases (2/4) #7453's first commit, unchanged. This move needs it too: without it, a patch written through an old path such as megatron.core.fusions.fused_layer_norm misses the implementation, and a --spec or checkpoint spec recorded under an old norm path stops loading. Move SSM helpers, Mamba2, GDN/GDN2 and GDP into megatron/core/ops with deprecated aliases (2/4) #7453 drops the commit once this PR lands.
2. ops/loss, ops/norm: move vocab-parallel cross entropy and the norm modules into megatron/core/ops: the move.
3. Review follow-ups, kept as separate commits on the reviewed head:
   • ops: keep a patched name on its forwarder so nested patches undo by assignment. A nested patch that set the outer patch's original back for a while lost the outer patch's record, and its exit deleted the implementation's attribute. A forwarder now holds a written name until the write ends, as a patched module does, so mock.patch and monkeypatch undo inner patches by assignment.
   • training: keep private spec names in the inference-counterpart lookup. The forwarder fallback skipped every underscore name, so a private --spec such as _foo_stack_spec no longer found _foo_inference_stack_spec. It now skips only dunder names.
   • CODEOWNERS: give @NVIDIA/transformer all of megatron/core/ops/norm, rather than torch_norm.py alone.
   • ops: let a forwarder's first same-value write start a write. Patching a name through an old path with the value it already held, as in mock.patch.object(old_mod, "NAME", old_mod.NAME), raised KeyError. A write now ends only when one is in progress and the value is the one it replaced.
4. Merge main into ops-move/loss-norm: keeps both imports where Reduce DSv4 output RoPE activation memory #8399 added one next to a rewritten norm import in deepseek_v4_hybrid_attention.py. Since review started, main is merged in rather than regenerated, so the reviewed commits stay as they are.

It is independent of #8348. The moved code is unchanged apart from imports, so open fixes to the old files, such as #7466's label-smoothing fix in tensor_parallel/cross_entropy.py, apply to the new paths once this lands.

Like the other moves, the branch is generated by a script on main plus the helper commit rather than rebased, and is regenerated whenever main changes the moved files. Re-running the script reproduces the move commit apart from the manual steps listed below.

Layout

Old (megatron/core/)	New (megatron/core/)
tensor_parallel/cross_entropy.py	ops/loss/cross_entropy.py
fusions/fused_cross_entropy.py	ops/loss/fused_cross_entropy.py
models/backends.py: CrossEntropyTarget, select_cross_entropy	ops/loss/backends.py
fusions/fused_layer_norm.py	ops/norm/fused_layer_norm.py
transformer/torch_norm.py	ops/norm/torch_norm.py

ops/loss/__init__.py and ops/norm/__init__.py hold each family's contract and re-export nothing. Transformer Engine's cross entropy and norms stay in extensions/transformer_engine.py, and fused_gated_norm.py (GDN's output norm) stays in fusions/.

Where #5438's pieces go

#5438 adds	With this PR
fusions/liger_cross_entropy.py	ops/loss/liger_cross_entropy.py: an adapter with the CrossEntropyTarget signature target(logits, labels, tp_group)
a 'liger' arm in LanguageModule.compute_language_model_loss	a 'liger' branch in select_cross_entropy (ops/loss/backends.py), the selector every provider asks. Since #6766, LanguageModule binds whatever that selector returns, so the model needs no change
'liger' in cross_entropy_fusion_impl	unchanged; select_cross_entropy rejects values it does not know, so the value is added there too
the LigerMegatronRMSNorm import in extensions/liger_kernel_spec_provider.py	ops/norm/liger_norm.py, which imports liger_kernel at module top; LigerSpecProvider.layer_norm returns it, and the provider and the use_liger spec plumbing stay in extensions/ and models/gpt
HAVE_LIGER flags and try/except ImportError	require("liger_kernel.megatron", ..., needed_by=...) in the selector branch, so the other values never import Liger and a missing package fails at construction

The new README section "Adding a cross-entropy or normalization backend" describes this without naming Liger.

Compatibility

• The four old module paths warn once and forward lazily. Patches written through them reach the implementation, and pickles and (module, name) specs recorded under them load the canonical objects (covered by the identity and pickle tests).
• from megatron.core.tensor_parallel import vocab_parallel_cross_entropy and from megatron.core.models.backends import CrossEntropyTarget, select_cross_entropy keep working and return the same objects.
• Importing megatron.core.tensor_parallel no longer binds the cross_entropy submodule as a package attribute. An explicit import megatron.core.tensor_parallel.cross_entropy still works, through the forwarder. No in-tree code relied on the attribute.
• select_cross_entropy now reads is_te_min_version from ops/loss/backends.py, so a test that patched megatron.core.models.backends.is_te_min_version to steer it patches the new module instead. The two in-tree tests that did are updated.

How it was generated

git switch -c ops-move/loss-norm upstream/main
git cherry-pick c80af85c5        # #7453's forwarder helper commit
python ops_move_loss_norm.py     # script below; template was the #8348 script

Manual steps after the script:

• moved CrossEntropyTarget and select_cross_entropy verbatim from models/backends.py to ops/loss/backends.py; models/backends.py, the TE spec provider and LanguageModule import them from there;
• ops/norm/torch_norm.py imports TransformerConfig for type checking only (it is used in annotations only). Outside the transformer package, the runtime import ran the package __init__, which reached transformer_layer, which imports torch_norm before it had finished loading, and import megatron.core failed. Inside the package, the package always initialized first;
• the moved fused_layer_norm.py gets its four pylint fixes (an unused importlib import, a long comment, two method docstrings), because the linting job checks it as a new file;
• tools/check_kernel_determinism_coverage.py does not count a deprecated forwarder left in a kernel directory as kernel-bearing. fusions/ is one, and the forwarders there hold no kernel; the kernels are registered at their new paths. The alternative, manifest entries naming the old paths, would contradict the canonical-path test. A module with anything besides the forwarding call keeps the directory rule (tested). The manifest's KERNEL_DIRECTORIES comment also says why megatron/core/ops is not a kernel directory;
• the README layout and the backend section, a CODEOWNERS line keeping @NVIDIA/transformer on ops/norm/torch_norm.py, and the re-sorted allowlist key in tools/process_group_usage_allowlist.json;
• tests: identity and pickle cases for VocabParallelCrossEntropy, FusedLayerNorm, WrappedTorchNorm and L2Norm; the package export and the selector's old import path; the forwarder rule of the coverage tool; the torch_norm.py path in the test_fused_activations.py docstring;
• isort on the changed Python files and black on those under megatron/core and tests.

Verify the move: git diff -M -C --find-copies-harder 577586474~1 577586474 shows the moved modules as C092–C099. Two have one import line changed; fused_layer_norm.py and torch_norm.py carry the edits above.

ops_move_loss_norm.py

#!/usr/bin/env python3
"""Move vocab-parallel cross entropy and the normalization modules into ``megatron/core/ops``.

Run from the repository root on a clean checkout of ``main`` with the forwarder helper commit
(``ops: write patches through deprecated forwarders ...``) applied::

    python ops_move_loss_norm.py

The script only moves files (``git mv``), rewrites import paths and path strings, and writes
the forwarder modules, the new package docstrings and the ``FORWARDED`` entries. Formatting
(isort, black) runs afterwards.
"""

from __future__ import annotations

import ast
import importlib.util
import re
import subprocess
from pathlib import Path

ROOT = Path.cwd()
CORE = "megatron/core"

# Old module file -> new module file, relative to megatron/core. Only whole modules move; no
# old package moves, so no old package needs a forwarder of its own.
FILE_MOVES = {
    "tensor_parallel/cross_entropy.py": "ops/loss/cross_entropy.py",
    "fusions/fused_cross_entropy.py": "ops/loss/fused_cross_entropy.py",
    "fusions/fused_layer_norm.py": "ops/norm/fused_layer_norm.py",
    "transformer/torch_norm.py": "ops/norm/torch_norm.py",
}

# Docstrings of the new canonical packages; each holds its contract and re-exports nothing.
NEW_PACKAGES = {
    "ops/loss": (
        "Vocab-parallel cross entropy.\n\n"
        "A target is called as ``loss = target(logits, labels, tp_group)`` with ``logits``\n"
        "``[s, b, vocab / tp]`` and ``labels`` ``[s, b]``, and returns the per-token loss\n"
        "``[s, b]``. The target owns every reduction over ``tp_group`` and consumes ``logits``\n"
        "in place; ``backends.CrossEntropyTarget`` states the contract. Targets hold no\n"
        "parameters or state.\n\n"
        "* ``cross_entropy`` -- the eager custom-autograd implementation, with label smoothing,\n"
        "  and the ``VocabParallelCrossEntropy`` steps the fused path shares.\n"
        "* ``fused_cross_entropy`` -- the same computation compiled with ``jit_fuser``.\n"
        "* ``backends`` -- ``select_cross_entropy``, which turns ``cross_entropy_loss_fusion``\n"
        "  and ``cross_entropy_fusion_impl`` into the target a model binds once.\n\n"
        "Transformer Engine's cross entropy stays in ``megatron.core.extensions.transformer_engine``.\n"
        "Packages here do not re-export; import each name from the module that defines it.\n"
    ),
    "ops/norm": (
        "Normalization modules.\n\n"
        "A norm is built through a ``torch_norm.LayerNormBuilder``,\n"
        "``target(config=..., hidden_size=..., eps=...)``, and satisfies\n"
        "``torch_norm.LayerNormInterface``: ``forward(x)`` normalizes over the last dimension and\n"
        "``returns_residual`` says whether it also returns the residual it fused. A model's\n"
        "``BackendSpecProvider.layer_norm`` chooses the builder. A norm owns its parameters\n"
        "(``weight``, plus ``bias`` for LayerNorm; ``L2Norm`` has none) and no process groups.\n\n"
        "* ``torch_norm`` -- the builder protocols, ``WrappedTorchNorm`` (PyTorch LayerNorm and\n"
        "  RMSNorm) and ``L2Norm``.\n"
        "* ``fused_layer_norm`` -- Apex's fused LayerNorm, ``FusedLayerNorm``.\n\n"
        "Transformer Engine's norms stay in ``megatron.core.extensions.transformer_engine``.\n"
        "Packages here do not re-export; import each name from the module that defines it.\n"
    ),
}

# Text files whose path strings and dotted module names are rewritten.
REWRITE_SUFFIXES = {".py", ".md", ".sh", ".json", ".yaml", ".yml", ".rst", ".txt"}
# The deprecated-path regression suite names the old paths on purpose; it is edited by hand.
SKIP_PREFIXES = ("tests/unit_tests/ops/", "CHANGELOG.md")

COPYRIGHT = "# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.\n"


def dotted(path: str) -> str:
    """``megatron/core/a/b/__init__.py`` -> ``megatron.core.a.b``."""
    path = path[: -len(".py")] if path.endswith(".py") else path
    path = path[: -len("/__init__")] if path.endswith("/__init__") else path
    return path.replace("/", ".")


MOVES = {f"{CORE}/{old}": f"{CORE}/{new}" for old, new in FILE_MOVES.items()}
RENAMES = {dotted(old): dotted(new) for old, new in MOVES.items()}
# Packages that lose a module; their files may import it relatively, without naming megatron.
OLD_PACKAGE_DIRS = tuple(sorted({str(Path(old).parent) + "/" for old in MOVES}))


def map_module(name: str) -> str:
    """Map a dotted name through the longest moved module prefix."""
    parts = name.split(".")
    for i in range(len(parts), 0, -1):
        prefix = ".".join(parts[:i])
        if prefix in RENAMES:
            return ".".join([RENAMES[prefix], *parts[i:]])
    return name


_OLD = "|".join(re.escape(old) for old in sorted(RENAMES, key=len, reverse=True))
_OLD_PATHS = "|".join(re.escape(old.replace(".", "/")) for old in sorted(RENAMES, key=len, reverse=True))
DOTTED_RE = re.compile(rf"(?<![\w.])(?:{_OLD})(?:\.[A-Za-z_]\w*)*(?!\w)")
PATH_RE = re.compile(rf"(?<![\w/])(?:{_OLD_PATHS})(?:\.py)?(?![\w/])")


def map_path(match: re.Match) -> str:
    text = match.group(0)
    is_file = text.endswith(".py")
    stem = text[: -len(".py")] if is_file else text
    new = RENAMES[stem.replace("/", ".")].replace(".", "/")
    return f"{new}.py" if is_file else new


def package_of(path: str) -> str:
    return dotted(str(Path(path).parent))


def resolve(node: ast.ImportFrom, package: str) -> str:
    if not node.level:
        return node.module or ""
    return importlib.util.resolve_name("." * node.level + (node.module or ""), package)


def relative_to(source: str, package: str) -> str:
    """Spell ``source`` relative to ``package`` when it lives in that package."""
    if source == package:
        return "."
    if source.startswith(package + "."):
        return "." + source[len(package) + 1 :]
    return source


def rewrite_import_statements(path: str, text: str, new_path: str | None) -> str:
    """Rewrite import statements that a textual prefix substitution cannot express.

    * ``from <old package> import moved_module`` binds the moved module under its old name.
    * Relative imports keep their form when they still resolve after the move; otherwise
      they become relative to the file's new package, or absolute when the target is outside.
    """
    try:
        tree = ast.parse(text)
    except SyntaxError:
        return text
    lines = text.splitlines(keepends=True)
    old_package = package_of(path)
    new_package = package_of(new_path) if new_path else old_package
    edits = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.ImportFrom):
            continue
        module = resolve(node, old_package)
        indent = re.match(r"\s*", lines[node.lineno - 1]).group(0)
        groups: dict[str, list[str]] = {}
        if node.level and map_module(module) == module and all(
            f"{resolve(node, new_package)}.{alias.name}" == map_module(f"{module}.{alias.name}")
            for alias in node.names
        ):
            # The relative import still names the same modules and objects after the move.
            continue
        if any(f"{module}.{alias.name}" in RENAMES for alias in node.names):
            for alias in node.names:
                full = f"{module}.{alias.name}"
                if full in RENAMES:
                    parent, _, base = RENAMES[full].rpartition(".")
                    local = alias.asname or alias.name
                    groups.setdefault(parent, []).append(
                        base if base == local else f"{base} as {local}"
                    )
                else:
                    groups.setdefault(map_module(module), []).append(
                        alias.name + (f" as {alias.asname}" if alias.asname else "")
                    )
        elif node.level:
            names = [a.name + (f" as {a.asname}" if a.asname else "") for a in node.names]
            groups[map_module(module)] = names
        else:
            continue
        if node.level:
            groups = {relative_to(source, new_package): names for source, names in groups.items()}
        statement = "".join(
            f"{indent}from {source} import {', '.join(names)}\n" for source, names in groups.items()
        )
        edits.append((node.lineno, node.end_lineno, statement))
    for start, end, statement in sorted(edits, reverse=True):
        lines[start - 1 : end] = [statement]
    return "".join(lines)


# ``from <old package> import <moved module>`` written out in text, such as a code string run
# in a subprocess; real import statements were already rewritten from the syntax tree.
_OLD_PARENTS = "|".join(re.escape(p) for p in sorted({o.rpartition(".")[0] for o in RENAMES}))
FROM_IMPORT_RE = re.compile(
    rf"\bfrom ({_OLD_PARENTS}) import (\w+)(?: as (\w+))?(?=\s*$)", re.MULTILINE
)


def map_from_import(match: re.Match) -> str:
    module, name, alias = match.groups()
    if f"{module}.{name}" not in RENAMES:
        return match.group(0)
    parent, _, base = RENAMES[f"{module}.{name}"].rpartition(".")
    local = alias or name
    return f"from {parent} import {base}" + ("" if base == local else f" as {local}")


def rewrite_text(text: str) -> str:
    text = FROM_IMPORT_RE.sub(map_from_import, text)
    text = DOTTED_RE.sub(lambda m: map_module(m.group(0)), text)
    return PATH_RE.sub(map_path, text)


def forwarder(old: str, new: str) -> str:
    doc = f'"""Deprecated import path; use ``{new}``."""'
    if len(doc) > 100:
        doc = f'"""Deprecated import path.\n\nUse ``{new}``.\n"""'
    return (
        f"{COPYRIGHT}\n{doc}\n\n"
        "from megatron.core.ops._compat import deprecated_module\n\n"
        f'__getattr__, __dir__ = deprecated_module(__name__, "{new}")\n'
    )


def package_init(doc: str) -> str:
    # A one-line docstring closes on its own line; a longer one closes on the line after.
    doc = doc.rstrip("\n") if "\n" not in doc.rstrip("\n") else doc
    return f'{COPYRIGHT}\n"""{doc}"""\n'


def forwarded_entries() -> str:
    return "".join(f'    "{old}": ("{new}",),\n' for old, new in RENAMES.items())


def git(*args: str) -> str:
    return subprocess.run(["git", *args], check=True, capture_output=True, text=True).stdout


def main() -> None:
    assert not git("status", "--porcelain"), "run on a clean checkout"
    tracked = git("ls-files").splitlines()

    # 1. Rewrite every tracked text file in place (moved files keep their old location for now).
    for path in tracked:
        if path.startswith(SKIP_PREFIXES) or Path(path).suffix not in REWRITE_SUFFIXES:
            continue
        file = ROOT / path
        original = file.read_text()
        if "megatron" not in original and not path.startswith(OLD_PACKAGE_DIRS):
            continue
        text = original
        if path.endswith(".py"):
            text = rewrite_import_statements(path, text, MOVES.get(path))
        text = rewrite_text(text)
        if text != original:
            file.write_text(text)

    # 2. Create the new packages, then move.
    for package, doc in NEW_PACKAGES.items():
        init = ROOT / CORE / package / "__init__.py"
        init.parent.mkdir(parents=True, exist_ok=True)
        init.write_text(package_init(doc))
        git("add", str(init.relative_to(ROOT)))
    for old, new in MOVES.items():
        git("mv", old, new)

    # 3. Leave a forwarder at every old path.
    for old, new in MOVES.items():
        (ROOT / old).write_text(forwarder(dotted(old), dotted(new)))
        git("add", old)

    # 4. Record the forwarders in the regression table.
    table = ROOT / "tests/unit_tests/ops/deprecated_paths.py"
    # The first family to move finds the table empty, written on one line.
    text = table.read_text().replace("FORWARDED = {}\n", "FORWARDED = {\n}\n", 1)
    head, sep, tail = text.partition("\n}\n")
    assert sep and head.lstrip().startswith("#") and "FORWARDED = {" in head, "unexpected table"
    table.write_text(head + "\n" + forwarded_entries().rstrip("\n") + sep + tail)


if __name__ == "__main__":
    main()

Testing

Run locally on CPU with a Triton stub:

• tests/unit_tests/ops and tests/unit_tests/determinism/kernels/test_manifest.py: 385 passed, 7 skipped. Five skipped tests cover the SSM move and run once Move SSM helpers, Mamba2, GDN/GDN2 and GDP into megatron/core/ops with deprecated aliases (2/4) #7453 lands; the other two are the same-value write tests for a forwarder with several targets, which refuses writes.
• The forwarder write-through and spec-loading tests pass under pytest 9.1.1, the pinned version, and 9.2.0, whose monkeypatch chooses how to undo from the target's __dict__. With the follow-up fixes reverted, the extended tests fail on all three reported cases: a nested patch that sets the original back, a private spec name, and a first write of the value a name already holds.
• The suites that exercise the moved code (test_backend_provider_contracts.py, tensor_parallel/test_cross_entropy.py, fusions/test_fused_cross_entropy.py, test_kernel_warmup.py, test_mamba_layer.py, and the determinism kernel tests for cross entropy and norm) give identical results on this branch and on main, plus one new test of the coverage tool's forwarder rule, which passes. What does not pass needs NCCL or a GPU.
• A cold import of every new and old module path, and of megatron.core.
• black, isort, pylint and ruff on the changed files, as the linting job runs them.
• tools/check_kernel_determinism_coverage.py: no unregistered kernel files. It asks for the determinism-exempt label because import lines change in registered files (extensions/transformer_engine.py, transformer/multi_latent_attention.py).
• tools/check_process_group_usage.py fails on main itself. On this branch its report is identical.

The GPU suites run in CI.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: Related to #5438

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[guihong-nv]

/ok to test 82590ed

[svcnemo-autobot]

Something went wrong on our side and the automatic retries didn't recover — please try again. (ref: github:8d2d29f0-d314-11f1-9016-5a17d671f910)

[guihong-nv]

/ok to test 324e72f

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[guihong-nv]

/ok to test 5775864

[balasaajay]

/review

[svcnvidia-nemo-ci]

Found two localized compatibility regressions: nested patch restoration can delete a canonical attribute, and private HybridStack spec names no longer resolve their inference counterparts. Details and acceptance checks are inline.

Completed the light review across all changed files and directly affected paths. This was static inspection; no tests were executed.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[guihong-nv]

/ok to test 13c6295

[guihong-nv]

/ok to test 6ad0da6

[kevalmorabia97]

The post_training path changes are import-only: ModelOpt's gpu_megatron export/distill/speculative (158 passed) and Megatron-Bridge examples (106 passed) give identical results before and after on nemo:26.08.