feat(flux): core runtime + Megatron adapter scaffolding (#807)

Base of an 18-PR series splitting the Flux diffusion-training feature
(training Flux, a DiT text-to-image diffusion model, on Primus/Megatron)
out of one large branch for reviewability. Targets `feat/flux/ci-env`
and auto-retargets to `main` once that merges. Every other content PR
stacks on this one.

## What this changes
The shared runtime scaffolding the rest of the feature builds on: core
runtime state + train-runtime wiring, the Megatron adapter and
base/pretrain trainers, and a patch auto-loader (`patches/__init__.py`)
that imports every `*_patches.py` in the package on import, so each
later layer only drops in its own patch file with no registry edit. Also
carries the shared test root (`tests/conftest.py`, `tests/utils.py`) and
a one-line `.gitignore` change (the bare `data` ignore → root-anchored
`/data/*`, and nothing else) that stops Git from ignoring the in-repo
`data/` source and config directories the later layers add. Kept
deliberately small so it can land first.

## Dependencies
Sequenced after the CI-pins PR (`feat/flux/ci-env`); no functional
dependency on the turbo bump.

## Test plan
`pytest tests/unit_tests/core tests/unit_tests/backends/megatron`;
lint/pre-commit clean. Validated locally on an AMD GPU container: 58
passed.

## Files
39 (core runtime, Megatron adapter/trainers, base patch loader, shared
test root, `.gitignore`).

---------

Co-authored-by: Flux Split <flux-split@local>
Co-authored-by: WangLingxun <linxwang@amd.com>
Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com>
Co-authored-by: luiza-amd <Luiza.Sayfullina@amd.com>
Co-authored-by: Flux Split Trial <flux-split-trial@local>

Author: 6 people

.gitignore

@@ -11,5 +11,5 @@ local/
 .gitmodules
 output
 experiment
-data
+/data/*
 pp_simulation_result

primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py

@@ -11,8 +11,9 @@
 
 
 class HummingbirdXTPosttrainTrainer(BaseTrainer):
-    def __init__(self, backend_args: Any):
-        super().__init__(backend_args=backend_args)
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # Accept and forward runtime context kwargs so BaseTrainer can filter them.
+        super().__init__(backend_args=backend_args, **kwargs)
 
     def setup(self):
         pass

primus/backends/maxtext/maxtext_pretrain_trainer.py

@@ -41,8 +41,11 @@ class MaxTextPretrainTrainer(BaseTrainer):
     Trainer class for MaxText pre-training.
     """
 
-    def __init__(self, backend_args: Any):
-        super().__init__(backend_args=backend_args)
+    def __init__(self, backend_args: Any = None, **kwargs):
+        # The core runtime instantiates every trainer with BaseModule-style
+        # context kwargs (module_name, primus_config, module_rank, ...). MaxText
+        # does not need them; accept and forward so BaseTrainer can filter them.
+        super().__init__(backend_args=backend_args, **kwargs)
 
         # Training state (populated in init())
         self.train_config: Optional[Any] = None

primus/backends/megatron/__init__.py

@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -12,3 +12,22 @@
 BackendRegistry.register_adapter("megatron", MegatronAdapter)
 BackendRegistry.register_trainer_class(MegatronPretrainTrainer, "megatron")
 BackendRegistry.register_trainer_class(MegatronSFTTrainer, "megatron", "sft")
+
+# Export trainers for convenience
+# Use lazy import for FluxPretrainTrainer to avoid Megatron dependency
+# when importing data pipeline components
+__all__ = [
+    "MegatronAdapter",
+    "FluxPretrainTrainer",
+    "MegatronPretrainTrainer",
+    "MegatronSFTTrainer",
+]
+
+
+def __getattr__(name):
+    """Lazy import for trainer classes to avoid Megatron dependency on module import."""
+    if name == "FluxPretrainTrainer":
+        from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer
+
+        return FluxPretrainTrainer
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

primus/backends/megatron/megatron_adapter.py

@@ -1,9 +1,12 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
 
+import importlib
+from typing import Optional
+
 import primus.backends.megatron.patches  # noqa: F401  # Register patches
 from primus.backends.megatron.argument_builder import MegatronArgBuilder
 from primus.core.backend.backend_adapter import BackendAdapter
@@ -18,8 +21,27 @@ def __init__(self, framework="megatron"):
         super().__init__(framework)
         self.third_party_dir_name = "Megatron-LM"
 
-    def load_trainer_class(self, stage: str = "pretrain"):
-        """Return the trainer class for the specified training stage."""
+    def load_trainer_class(self, stage: str = "pretrain", trainer_class: Optional[str] = None):
+        """
+        Return the Trainer class for the specified training stage or trainer class name.
+
+        Args:
+            stage: Training stage (e.g., "pretrain", "sft"). Defaults to "pretrain".
+            trainer_class: Optional specific trainer class name for dynamic loading.
+                          If provided, this takes precedence over stage-based selection.
+
+        Returns:
+            Trainer class
+
+        Raises:
+            ImportError: If trainer class cannot be imported
+            ValueError: If stage is invalid or trainer class is not found
+        """
+        # If trainer_class is specified, attempt dynamic loading
+        if trainer_class:
+            return self._load_trainer_class_by_name(trainer_class)
+
+        # Fallback to stage-based selection via the backend registry
         try:
             trainer_cls = BackendRegistry.get_trainer_class(self.framework, stage=stage)
         except (ValueError, AssertionError) as exc:
@@ -31,6 +53,76 @@ def load_trainer_class(self, stage: str = "pretrain"):
         log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer class: {trainer_cls.__name__}")
         return trainer_cls
 
+    def _load_trainer_class_by_name(self, trainer_class: str):
+        """
+        Dynamically load trainer class by name.
+
+        Args:
+            trainer_class: Trainer class name (e.g., "FluxPretrainTrainer", "MegatronPretrainTrainer")
+
+        Returns:
+            Trainer class
+
+        Raises:
+            ImportError: If trainer class cannot be imported
+            ValueError: If trainer class is not found
+        """
+        # Define trainer registry for Megatron backend
+        # This maps trainer class names to their module paths
+        MEGATRON_TRAINERS = {
+            "MegatronPretrainTrainer": "primus.backends.megatron.megatron_pretrain_trainer.MegatronPretrainTrainer",
+            "FluxPretrainTrainer": "primus.backends.megatron.flux_pretrain_trainer.FluxPretrainTrainer",
+        }
+
+        if trainer_class not in MEGATRON_TRAINERS:
+            # Try to load from common locations as fallback
+            log_rank_0(
+                f"[Primus:MegatronAdapter] Trainer '{trainer_class}' not in registry, attempting dynamic import..."
+            )
+
+            possible_paths = [
+                f"primus.backends.megatron.{trainer_class.lower()}.{trainer_class}",
+                f"primus.modules.trainer.megatron.{trainer_class.lower()}.{trainer_class}",
+                f"primus.backends.megatron.{trainer_class}",
+            ]
+
+            for module_path in possible_paths:
+                try:
+                    module_name, class_name = module_path.rsplit(".", 1)
+                    module = importlib.import_module(module_name)
+                    trainer_cls = getattr(module, class_name)
+                    log_rank_0(
+                        f"[Primus:MegatronAdapter] Successfully loaded trainer: {trainer_class} from {module_name}"
+                    )
+                    return trainer_cls
+                except (ImportError, AttributeError):
+                    continue
+
+            # If all attempts failed, provide helpful error message
+            available = list(MEGATRON_TRAINERS.keys())
+            raise ValueError(
+                f"Trainer class '{trainer_class}' not found.\n"
+                f"Available trainers: {', '.join(available) if available else 'none'}\n"
+                f"Hint: Set 'trainer_class' in your experiment YAML to a registered trainer "
+                f"(e.g. trainer_class: FluxPretrainTrainer), register it in MEGATRON_TRAINERS, "
+                f"or ensure it is importable from standard locations."
+            )
+
+        # Load from registry
+        trainer_path = MEGATRON_TRAINERS[trainer_class]
+        module_path, class_name = trainer_path.rsplit(".", 1)
+
+        try:
+            module = importlib.import_module(module_path)
+            trainer_cls = getattr(module, class_name)
+            log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer: {trainer_class} from {module_path}")
+            return trainer_cls
+        except (ImportError, AttributeError) as e:
+            raise ImportError(
+                f"Failed to load trainer class '{trainer_class}' from '{module_path}': {e}\n"
+                f"Hint: Check that the module exists and the class name is correct."
+            ) from e
+
     def detect_backend_version(self) -> str:
         """Detect Megatron-LM version via AST parsing (avoids __init__.py execution)."""
         import ast

primus/backends/megatron/megatron_base_trainer.py

@@ -1,5 +1,5 @@
 ###############################################################################
-# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 ###############################################################################
@@ -14,7 +14,6 @@
 )
 from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts
 from primus.core.trainer.base_trainer import BaseTrainer
-from primus.core.utils.env import flush_before_hard_exit
 from primus.modules.module_utils import log_rank_0, warning_rank_0
 
 
@@ -23,9 +22,84 @@ class MegatronBaseTrainer(BaseTrainer):
 
     def setup(self):
         """Setup Megatron runtime: set global vars and patch parse_args."""
+        self._ensure_megatron_path()
         set_primus_global_variables(self.backend_args)
         self._patch_parse_args()
 
+    def _ensure_megatron_path(self):
+        """Ensure Megatron-LM path is in sys.path before any megatron imports."""
+        import os
+        import sys
+        from pathlib import Path
+
+        # First, check if megatron is already importable
+        try:
+            import megatron  # type: ignore
+
+            return  # Path already set correctly
+        except ImportError:
+            pass
+
+        # Try multiple methods to find the Megatron-LM path
+        megatron_paths = []
+
+        # Method 1: From PRIMUS_PATH environment variable (most reliable)
+        primus_path = os.getenv("PRIMUS_PATH")
+        if primus_path:
+            megatron_path = Path(primus_path) / "third_party" / "Megatron-LM"
+            if megatron_path.exists():
+                megatron_paths.append(str(megatron_path))
+
+        # Method 2: From current working directory (works in container)
+        try:
+            cwd = Path.cwd()
+            # Check if we're in /workspace/Primus or a subdirectory
+            if "Primus" in str(cwd):
+                # Find the Primus root
+                primus_root = cwd
+                while primus_root.name != "Primus" and primus_root != primus_root.parent:
+                    primus_root = primus_root.parent
+                if primus_root.name == "Primus":
+                    megatron_path = primus_root / "third_party" / "Megatron-LM"
+                    if megatron_path.exists():
+                        megatron_paths.append(str(megatron_path))
+        except Exception:
+            pass
+
+        # Method 3: From current file location
+        try:
+            repo_root = Path(__file__).resolve().parents[3]
+            megatron_path = repo_root / "third_party" / "Megatron-LM"
+            if megatron_path.exists():
+                megatron_paths.append(str(megatron_path))
+        except Exception:
+            pass
+
+        # Method 4: Check if already in sys.path
+        for path in sys.path:
+            path_obj = Path(path)
+            if path_obj.exists():
+                megatron_pkg = path_obj / "megatron"
+                if megatron_pkg.exists() and megatron_pkg.is_dir():
+                    return  # Path already set correctly
+
+        # Add paths to sys.path if not already present
+        for path in megatron_paths:
+            if path not in sys.path:
+                sys.path.insert(0, path)
+                log_rank_0(f"[Primus:MegatronBaseTrainer] Added Megatron-LM to sys.path: {path}")
+
+        # Verify the path was set correctly by trying to import
+        try:
+            import megatron  # type: ignore
+
+            log_rank_0("[Primus:MegatronBaseTrainer] Successfully verified megatron import")
+        except ImportError as e:
+            log_rank_0(
+                f"[Primus:MegatronBaseTrainer] WARNING: Failed to import megatron after path setup: {e}"
+            )
+            log_rank_0(f"[Primus:MegatronBaseTrainer] sys.path: {sys.path[:5]}")
+
     def init(self):
         """Initialize Megatron training components."""
         log_rank_0("Initializing Megatron training...")
@@ -53,7 +127,14 @@ def cleanup(self, on_error: bool = False):
 
         if exit_fast and not on_error:
             log_rank_0("[MegatronBaseTrainer] PRIMUS_EXIT_FAST=1 -> os._exit(0)")
-            flush_before_hard_exit()
+            # Flush stdout/stderr so the final log lines are not lost.
+            try:
+                import sys
+
+                sys.stdout.flush()
+                sys.stderr.flush()
+            except Exception:  # pragma: no cover
+                pass
             os._exit(0)
 
     def _finalize_mlflow_artifacts(self):