Fix fixed-batch elastic training with a fixed Core layout

engine/heterogeneous_engine.py

@@ -327,7 +327,7 @@ def wait_for_ready(self, timeout: float = 300.0):
 
     def wait_until_idle(self, timeout: float = 3600.0, poll_interval: float = 0.5):
         """Block until no rollout requests are active before reconfiguration."""
-        deadline = time.time() + max(0.0, timeout)
+        deadline = time.time() + timeout if timeout > 0 else float("inf")
         while True:
             with self._lock:
                 handles = list(getattr(self.scheduler, "_instances", []))

engine/megatron_core_checkpointing.py

@@ -168,7 +168,9 @@ def load(
         )
         self._load_model_state(self._unwrap_model(model), loaded)
         if optimizer is not None and "optimizer" in loaded:
-            optimizer.load_state_dict(loaded["optimizer"])
+            from RL_Framework.engine.optimizer_state import optimizer_state_load_context
+            with optimizer_state_load_context(optimizer):
+                optimizer.load_state_dict(loaded["optimizer"])
         self._barrier()
 
     def _build_sharded_state_dict(

engine/optimizer_memory.py

@@ -0,0 +1,42 @@
+"""Bound temporary storage in the native PyTorch Adam fallback."""
+
+import torch
+
+
+def bound_default_torch_adam_memory(optimizer) -> int:
+    """Disable automatic foreach selection in native Adam/AdamW leaves.
+
+    CUDA's automatic foreach Adam path creates a full parameter-sized list of
+    temporaries (for example, sqrt of every second moment at once). MCore's
+    non-fused fallback can therefore exceed memory even when model, master
+    parameters, gradients and Adam moments fit. The per-tensor implementation
+    retains the same precision and bounds those temporaries to one parameter.
+
+    Explicit foreach choices and fused implementations are preserved. This
+    walks the MCore chained/single optimizer wrappers without accessing a
+    ChainedOptimizer's single-child-only ``optimizer`` property. Return the
+    number of parameter groups changed; repeated calls are idempotent.
+    """
+    pending = [optimizer]
+    visited = set()
+    changed = 0
+    while pending:
+        current = pending.pop()
+        if current is None or id(current) in visited:
+            continue
+        visited.add(id(current))
+        if type(current) in (torch.optim.Adam, torch.optim.AdamW):
+            for group in current.param_groups:
+                if group.get("foreach") is None and group.get("fused") is not True:
+                    group["foreach"] = False
+                    changed += 1
+            if (current.defaults.get("foreach") is None
+                    and current.defaults.get("fused") is not True):
+                current.defaults["foreach"] = False
+            continue
+        children = getattr(current, "chained_optimizers", None)
+        if children is not None:
+            pending.extend(children)
+        else:
+            pending.append(getattr(current, "optimizer", None))
+    return changed

engine/optimizer_state.py

@@ -0,0 +1,58 @@
+"""Strict native MCore optimizer metadata restoration for fixed-topology EHP."""
+
+from contextlib import contextmanager
+
+import torch
+
+
+@contextmanager
+def optimizer_state_load_context(optimizer):
+    """Preserve native metadata's Torch Adam step during MCore tensor loads.
+
+    MCore 0.14's model-space loader deliberately excludes ``step`` from tensor
+    state, but its low-level setter iterates every Torch Adam state key. Supply
+    the step already loaded from common metadata to that setter only. Other
+    missing master/moment tensors still fail through the native strict copy.
+    This temporary instance adapter never edits installed MCore dependencies.
+    """
+    parts = getattr(optimizer, "chained_optimizers", (optimizer,))
+    restores = []
+    absent = object()
+    completed = False
+    try:
+        for part in parts:
+            inner = getattr(part, "optimizer", None)
+            native = getattr(part, "_set_main_param_and_optimizer_states", None)
+            if not isinstance(inner, (torch.optim.Adam, torch.optim.AdamW)) or not callable(native):
+                continue
+            name = "_set_main_param_and_optimizer_states"
+            previous = vars(part).get(name, absent)
+
+            def copy_with_metadata_step(model_param, tensors, *, part=part, inner=inner, native=native):
+                group, order = part.model_param_group_index_map[model_param]
+                main_param = inner.param_groups[group]["params"][order]
+                step = inner.state[main_param].get("step")
+                if "step" not in tensors and step is not None:
+                    tensors = dict(tensors, step=step)
+                return native(model_param, tensors)
+
+            setattr(part, name, copy_with_metadata_step)
+            restores.append((part, name, previous))
+        yield
+        completed = True
+    finally:
+        for part, name, previous in reversed(restores):
+            if previous is absent:
+                delattr(part, name)
+            else:
+                setattr(part, name, previous)
+            # Native MCore's Torch fallback assigns one shared scalar tensor
+            # to every Adam parameter during metadata loading. Torch Adam
+            # increments each parameter's step in-place; shared references
+            # would advance it once per parameter and change the next update.
+            if completed:
+                for state in part.optimizer.state.values():
+                    step = state.get("step")
+                    if isinstance(step, torch.Tensor):
+                        state["step"] = step.detach().clone()
+

infra/cost_model/global_resource_planner.py

@@ -15,6 +15,7 @@
     RequestInfo,
     RolloutClusterConfig,
     TrainParallelConfig,
+    representative_global_batch,
 )
 from RL_Framework.infra.cost_model.optimizer import (
     OptimizationResult,
@@ -25,6 +26,51 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class CoreRuntimeCapabilities:
+    """Fixed Core layout on the already initialized physical ranks.
+
+    Hybrid replicas use a separate process group and do not enlarge this rank
+    set. No capability is supplied to startup GRP, which may choose its size.
+    """
+    world_size: int
+    tp: int
+    pp: int
+    cp: int
+    ep: int
+    dp: int
+    rebuild_tp_dp: bool = field(default=False, init=False)
+
+    @classmethod
+    def from_engine(cls, engine):
+        world = getattr(engine, "world_size", None)
+        if not isinstance(world, int) or world < 1:
+            return None
+        getter = getattr(engine, "get_parallel_state", None)
+        state = getter() if callable(getter) else {}
+        dimensions = {axis: int(state.get(f"train_{axis}",
+            getattr(engine, f"train_{axis}_size", 1))) for axis in ("tp", "pp", "cp", "ep", "dp")}
+        return cls(world_size=world, **dimensions)
+
+    def incompatibility(self, plan):
+        topology = plan.train_config
+        if int(plan.train_gpus) != self.world_size:
+            return "physical_core_rank_set_changed"
+        for axis in ("pp", "cp", "ep"):
+            if int(getattr(topology, axis, 1)) != getattr(self, axis):
+                return f"core_{axis}_rebuild_unsupported"
+        if not self.rebuild_tp_dp and (int(topology.tp) != self.tp
+                or int(topology.dp) * int(getattr(topology, "ep", 1)) != self.dp):
+            return "core_tp_dp_rebuild_unsupported"
+        return ""
+
+    def to_dict(self):
+        return {"physical_rank_count": self.world_size, "tp": self.tp, "pp": self.pp,
+            "cp": self.cp, "ep": self.ep, "dp": self.dp,
+            "rebuild_tp_dp": self.rebuild_tp_dp,
+            "physical_rank_resize": False, "source": "initialized_training_engine"}
+
+
 @dataclass
 class GlobalResourcePlan:
     """Concrete cross-stage plan selected by the planner."""
@@ -188,11 +234,19 @@ class PlannerDecision:
     trigger: str = ""
     runtime_metrics: RuntimeResourceMetrics | None = None
     elastic_hybrid_signal: ElasticHybridSignal | None = None
+    # None preserves legacy hand-constructed decisions. Planner-produced
+    # decisions explicitly distinguish full repartition from EHP adjustment.
+    resource_should_reconfigure: bool | None = None
 
     def to_dict(self) -> dict[str, Any]:
         return {
             "step": self.step,
             "should_reconfigure": self.should_reconfigure,
+            "resource_should_reconfigure": self.resource_should_reconfigure,
+            "elastic_should_adjust": bool(
+                self.elastic_hybrid_signal is not None
+                and self.elastic_hybrid_signal.should_adjust
+            ),
             "reason": self.reason,
             "trigger": self.trigger,
             "num_requests": self.num_requests,
@@ -316,6 +370,7 @@ def __init__(
             1, int(elastic_hybrid_signal_ttl_steps)
         )
         self.verbose = verbose
+        self.runtime_capabilities: CoreRuntimeCapabilities | None = None
 
         self.optimizer = TwoLevelNestedOptimizer(
             cost_model=evaluator,
@@ -523,6 +578,7 @@ def plan_if_needed(
                     num_requests=len(self._history),
                     trigger="warmup",
                     runtime_metrics=runtime_metrics,
+                    resource_should_reconfigure=False,
                 )
             )
 
@@ -537,6 +593,7 @@ def plan_if_needed(
                     num_requests=len(self._history),
                     trigger=trigger,
                     runtime_metrics=runtime_metrics,
+                    resource_should_reconfigure=False,
                 )
             )
         if trigger == "interval_skip":
@@ -549,6 +606,7 @@ def plan_if_needed(
                     num_requests=len(self._history),
                     trigger=trigger,
                     runtime_metrics=runtime_metrics,
+                    resource_should_reconfigure=False,
                 )
             )
         if len(self._history) < self.min_history_size:
@@ -561,10 +619,12 @@ def plan_if_needed(
                     num_requests=len(self._history),
                     trigger=trigger,
                     runtime_metrics=runtime_metrics,
+                    resource_should_reconfigure=False,
                 )
             )
 
-        current = self._build_current_plan(config, self._history)
+        workload = representative_global_batch(self._history, int(config.batch_size))
+        current = self._build_current_plan(config, workload)
         forced_rollout_tp = self._forced_rollout_tp_list(config)
         forced_train = self._forced_train_config(
             config,
@@ -578,7 +638,7 @@ def plan_if_needed(
         forced_candidate = forced_rollout_tp is not None or forced_train is not None
         result = self.optimizer.optimize(
             n_total_gpus=self.n_total_gpus,
-            requests=list(self._history),
+            requests=workload,
             B_global=config.batch_size,
             allowed_train_tp=self.allowed_train_tp,
             allowed_train_ep=self.allowed_train_ep,
@@ -612,19 +672,20 @@ def plan_if_needed(
                     num_requests=len(self._history),
                     trigger=trigger,
                     runtime_metrics=runtime_metrics,
+                    resource_should_reconfigure=False,
                 )
             )
 
         if forced_rollout_tp is not None:
             forced_rollout = RolloutClusterConfig(tp_list=forced_rollout_tp)
             forced_rollout_time, forced_rollout_details = self.evaluator.evaluate_rollout(
                 forced_rollout,
-                list(self._history),
+                workload,
             )
             forced_train_time = current.t_train
             forced_train_details = current.metadata.get("training", {})
             if forced_train is not None:
-                sequence_lengths = [r.total_length for r in self._history] or [1024]
+                sequence_lengths = [r.total_length for r in workload] or [1024]
                 forced_train_time, forced_train_details = self.evaluator.evaluate_training(
                     forced_train,
                     config.batch_size,
@@ -649,7 +710,7 @@ def plan_if_needed(
                 },
             )
         elif forced_train is not None:
-            sequence_lengths = [r.total_length for r in self._history] or [1024]
+            sequence_lengths = [r.total_length for r in workload] or [1024]
             forced_train_time, forced_train_details = self.evaluator.evaluate_training(
                 forced_train,
                 config.batch_size,
@@ -669,6 +730,42 @@ def plan_if_needed(
                 },
             )
 
+        # Keep the unconstrained capacity recommendation as EHP demand, while
+        # selecting Core/rollout layouts that the live backend can execute.
+        # Startup has no live capabilities and retains the full 4/8/... search.
+        elastic_candidate = candidate
+        capabilities = self.runtime_capabilities
+        rejected = capabilities.incompatibility(candidate) if capabilities is not None else ""
+        if rejected:
+            unconstrained = candidate.to_dict()
+            if not forced_candidate:
+                feasible = self.optimizer.optimize(
+                    n_total_gpus=self.n_total_gpus, requests=workload,
+                    B_global=config.batch_size,
+                    allowed_train_tp=(self.allowed_train_tp if capabilities.rebuild_tp_dp else [capabilities.tp]),
+                    allowed_train_ep=[capabilities.ep], allowed_train_pp=[capabilities.pp],
+                    micro_batch_sizes=self.micro_batch_sizes,
+                    fixed_train_gpus=capabilities.world_size,
+                    allocation_granularity_gpus=self.allocation_granularity_gpus,
+                    min_train_gpus=self.min_train_gpus, min_rollout_gpus=self.min_rollout_gpus,
+                    rollout_node_tp_pattern=list(getattr(config.global_resource_planner,
+                        "rollout_node_tp_pattern", []) or []),
+                    train_cp_size=capabilities.cp,
+                    dp_batch_group_size=max(1, int(getattr(config, "n_samples", 1) or 1)))
+                candidate = self._plan_from_result(feasible, config) or current
+                result = feasible
+            else:
+                # Forced test commands also cannot authorize inventing ranks.
+                candidate = current
+            candidate.metadata["runtime_candidate_filter"] = {
+                "reason": rejected, "capabilities": capabilities.to_dict(),
+                "unconstrained_candidate": unconstrained,
+                "elastic_demand_preserved": True}
+        if capabilities is not None:
+            candidate.metadata["runtime_core_capabilities"] = capabilities.to_dict()
+        candidate.metadata["workload"] = {
+            "global_batch_size": int(config.batch_size), "history_size": len(self._history),
+            "representative_requests": len(workload), "selection": "deterministic_length_quantiles"}
         raw_gain = current.t_global - candidate.t_global
         net_gain = raw_gain - self.reconfiguration_cost_s
         gain_ratio = raw_gain / max(current.t_global, 1e-6)
@@ -681,7 +778,7 @@ def plan_if_needed(
             step=step,
             config=config,
             current=current,
-            candidate=candidate,
+            candidate=elastic_candidate,
             runtime_metrics=runtime_metrics,
         )
         elastic_adjust = bool(
@@ -692,6 +789,7 @@ def plan_if_needed(
         )
         resource_should = (
             (not same_plan or force_runtime_reconfigure)
+            and not (rejected and forced_candidate)
             and (
                 forced_candidate
                 or force_runtime_reconfigure
@@ -725,12 +823,16 @@ def plan_if_needed(
             trigger=trigger,
             runtime_metrics=runtime_metrics,
             elastic_hybrid_signal=elastic_signal,
+            resource_should_reconfigure=resource_should,
         )
-        if should:
+        if resource_should:
             self._active_plan = candidate
-            self._last_replan_step = int(step)
-        elif self._active_plan is None:
+        else:
+            # An EHP join/release does not apply the rejected candidate's
+            # training topology or rollout TP layout.
             self._active_plan = current
+        if should:
+            self._last_replan_step = int(step)
         self._mark_rejected_trigger_observed(trigger, runtime_metrics)
         return self._remember(decision)
 
@@ -1142,6 +1244,7 @@ def _build_current_plan(
         config: AsyncRLConfig,
         requests: list[RequestInfo],
     ) -> GlobalResourcePlan:
+        requests = representative_global_batch(requests, int(config.batch_size))
         train = TrainParallelConfig(
             tp=config.train_tp_size,
             ep=max(1, int(getattr(config, "train_ep_size", 1) or 1)),