Add versioned trainer graphs and logical-rank APIs

.github/workflows/prek.yml

@@ -213,7 +213,7 @@ jobs:
       - name: Run Megatron lightweight tests
         run: |
           megatron_runtime/.venv/bin/python -c "import megatron.core.packed_seq_params"
-          megatron_runtime/.venv/bin/python -m pytest --nbval --current-env --tb=short \
+          megatron_runtime/.venv/bin/python -m pytest -v --nbval --current-env --tb=short \
             tests/unit/test_megatron_reference_logprobs.py \
             tests/unit/test_preprocessing_tokenize.py::test_gemma4_normalizes_json_tool_arguments_for_mapping_template \
             tests/unit/test_moe_routing_replay.py \
@@ -224,10 +224,14 @@ jobs:
             tests/unit/test_prefix_tree_grad_parity.py \
             tests/unit/test_prefix_tree_packing.py \
             tests/unit/test_trainer_rank_handoff_budget.py \
+            tests/unit/test_trainer_rank_graph_backward_work.py \
             tests/unit/test_qwen35_adapter_config.py \
             tests/unit/test_trainer_rank_physical_reserve.py \
             tests/unit/test_trainer_rank_validation.py \
             tests/unit/test_trainer_rank_weird_shapes.py \
+            tests/unit/test_trainer_rank_slot_graph_lifetime.py \
+            tests/unit/test_trainer_rank_forward_handoff.py \
+            tests/unit/test_trainer_rank_commands.py::test_gloo_dp2_tp2_participation_and_gradients \
             tests/unit/test_trainer_rank_admission_inputs.py \
             tests/unit/test_trainer_rank_checkpoint_memory.py \
             tests/unit/test_trainer_rank_profile_warm.py \
@@ -249,6 +253,7 @@ jobs:
             tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \
             tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_refuses_after_recovery \
             tests/acceptance/trainer_rank_planner \
+            tests/integration/megatron/lora/test_lora_disk_codecs.py::test_export_preparation_failure_is_collective_and_retryable \
             tests/integration/megatron/test_sft_packing.py::test_sft_packing_preserves_training_targets \
             tests/integration/megatron/model_support/test_dispatcher_graph_retention.py \
             tests/integration/megatron/gdn_shared_prefix/test_gdn_planner_runtime_model.py \
@@ -264,6 +269,7 @@ jobs:
           uv run --no-sync pytest --nbval --current-env --tb=short tests/unit \
             --deselect=tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \
             --deselect=tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_refuses_after_recovery \
+            --deselect=tests/unit/test_trainer_rank_commands.py::test_gloo_dp2_tp2_participation_and_gradients \
             --ignore=tests/unit/test_megatron_reference_logprobs.py \
             --ignore=tests/unit/test_moe_routing_replay.py \
             --ignore=tests/unit/test_moe_routing_real_path.py \
@@ -273,10 +279,13 @@ jobs:
             --ignore=tests/unit/test_prefix_tree_grad_parity.py \
             --ignore=tests/unit/test_prefix_tree_packing.py \
             --ignore=tests/unit/test_trainer_rank_handoff_budget.py \
+            --ignore=tests/unit/test_trainer_rank_graph_backward_work.py \
             --ignore=tests/unit/test_qwen35_adapter_config.py \
             --ignore=tests/unit/test_trainer_rank_physical_reserve.py \
             --ignore=tests/unit/test_trainer_rank_validation.py \
             --ignore=tests/unit/test_trainer_rank_weird_shapes.py \
+            --ignore=tests/unit/test_trainer_rank_slot_graph_lifetime.py \
+            --ignore=tests/unit/test_trainer_rank_forward_handoff.py \
             --ignore=tests/unit/test_trainer_rank_admission_inputs.py \
             --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \
             --ignore=tests/unit/test_trainer_rank_profile_warm.py \

dev/trainer_rank.py

@@ -70,17 +70,17 @@ def main(
         for step in range(steps):
             loss_sum = torch.tensor(0.0, device=rank.device)
             token_count = torch.tensor(0.0, device=rank.device)
-            for micro in rank.forward_micro_batches(inputs, checkpoint=slot):
+            for micro in rank.forward_batches(inputs, checkpoint=slot):
                 loss = torch.tensor(0.0, device=rank.device)
                 for output in micro.outputs:
                     assert output.target_logprobs is not None
                     loss = loss - output.target_logprobs.sum()
                     token_count += output.target_logprobs.numel()
-                loss.backward()
+                rank.backward(loss)
                 loss_sum += loss.detach()
 
-            rank.dp_reduce(loss_sum)
-            rank.dp_reduce(token_count)
+            rank.reduce(loss_sum)
+            rank.reduce(token_count)
             scale = 1.0 / max(float(token_count.item()), 1.0)
             metrics = rank.optim_step(
                 params=AdamParams(learning_rate=lr),

dev/trainer_rank_check.py

@@ -226,7 +226,7 @@ def _local_outputs(
     rank: TrainerRank,
     indexed_requests: Sequence[tuple[int, ForwardInput]],
 ) -> list[dict[str, object]]:
-    outputs = rank.dp_rank_forward([request for _, request in indexed_requests])
+    outputs = rank.forward([request for _, request in indexed_requests])
     return [
         _output_record(index, torch.arange(request.input_tokens.numel()), output)
         for (index, request), output in zip(indexed_requests, outputs, strict=True)
@@ -502,7 +502,7 @@ def _head_backward_chunk_parity(
             outputs = rank._project_head(items, prepared, candidate)
             loss = _output_loss(outputs)
             logprob_sums.append(loss.detach())
-            loss.backward()
+            rank.backward(loss)
             assert candidate.grad is not None
             gradients.append(candidate.grad)
         finally:
@@ -584,7 +584,7 @@ def _slot_gradients(
     slots: Sequence[str],
 ) -> dict[str, list[torch.Tensor]]:
     rank.zero_grad()
-    _output_loss(rank.dp_rank_forward(requests)).backward()
+    rank.backward(_output_loss(rank.forward(requests)))
     return {
         slot: [
             torch.zeros_like(parameter, dtype=torch.float32, device="cpu")
@@ -665,14 +665,16 @@ def step() -> list[MicroBatchStats]:
         rank.zero_grad()
         stats: list[MicroBatchStats] = []
         if adaptive:
-            for micro in rank.forward_micro_batches(requests):
-                _output_loss(cast(Sequence[ForwardOutput], micro.outputs)).backward()
+            for micro in rank.forward_batches(requests):
+                rank.backward(
+                    _output_loss(cast(Sequence[ForwardOutput], micro.outputs))
+                )
                 stats.append(micro.stats)
         else:
-            outputs = rank.dp_rank_forward(requests[dp_rank::dp_size])
+            outputs = rank.forward(requests[dp_rank::dp_size])
             if workload == "unequal_slots":
                 _trace_unequal_slots("forward_ready")
-            _output_loss(outputs).backward()
+            rank.backward(_output_loss(outputs))
             if workload == "unequal_slots":
                 _trace_unequal_slots("backward_ready")
         if optimizer_step:

dev/trainer_rank_collective_trace.py

@@ -141,7 +141,7 @@ def logged(*args, **kwargs):
     # gdn.layout imported the function by name; rebind it to the logged one.
     gdn_layout.all_to_all_single = dist.all_to_all_single
 
-    original_forward = TrainerRank.dp_rank_forward
+    original_forward = TrainerRank.forward
 
     @functools.wraps(original_forward)
     def logged_forward(self, *args, **kwargs):
@@ -166,7 +166,7 @@ def logged_forward(self, *args, **kwargs):
         )
         return outputs
 
-    TrainerRank.dp_rank_forward = logged_forward  # type: ignore[method-assign]
+    TrainerRank.forward = logged_forward  # type: ignore[method-assign]
 
 
 def _start_watchdog(logger: _Logger, log_dir: Path, stall_seconds: float) -> None:

dev/trainer_rank_landing_acceptance.py

@@ -137,17 +137,35 @@ def phase_contract() -> None:
 
     problems: list[str] = []
     constructor = _public_parameters(trainer_rank.TrainerRank.__init__)
-    if list(constructor) != ["runtime"]:
+    if (
+        list(constructor) != ["runtime", "options"]
+        or constructor["options"].kind is not inspect.Parameter.KEYWORD_ONLY
+        or constructor["options"].default is not None
+    ):
         problems.append(
-            f"TrainerRank must accept exactly (runtime); found {sorted(constructor)}"
+            f"TrainerRank must accept (runtime, *, options=None); found {constructor}"
         )
-    for method_name in ("forward_micro_batches", "dp_rank_forward"):
+    for method_name in ("forward_batches", "forward"):
         parameters = _public_parameters(getattr(trainer_rank.TrainerRank, method_name))
         # ``yield_empty`` (PR #864) is a keyword-only flag defaulting to False:
         # off, the contract the acceptance suite pins is unchanged.
-        extra = set(parameters) - {"inputs", "checkpoint", "no_grad", "yield_empty"}
+        extra = set(parameters) - {
+            "inputs",
+            "checkpoint",
+            "no_grad",
+            "yield_empty",
+            "options",
+        }
         if extra:
             problems.append(f"{method_name} has extra parameters {sorted(extra)}")
+        options = parameters.get("options")
+        if options is None or (
+            options.kind is not inspect.Parameter.KEYWORD_ONLY
+            or options.default is not None
+        ):
+            problems.append(
+                f"{method_name}: options must be keyword-only and default to None"
+            )
         flag = parameters.get("yield_empty")
         if flag is not None and (
             flag.kind is not inspect.Parameter.KEYWORD_ONLY or flag.default is not False
@@ -374,8 +392,8 @@ def phase_measure(cell: str, arm: str, output_jsonl: str, repeat: int) -> None:
 
         # Behavior smoke from the contract: empty inputs are valid zero-work
         # calls that must not disturb subsequent planning.
-        empty = rank.dp_rank_forward([])
-        assert len(list(empty)) == 0, "dp_rank_forward([]) must return no outputs"
+        empty = rank.forward([])
+        assert len(list(empty)) == 0, "forward([]) must return no outputs"
 
         rows: list[dict[str, object]] = []
         for sample in range(repeat + 4):  # 1 cold + 3 warmup + repeat measured
@@ -387,8 +405,8 @@ def phase_measure(cell: str, arm: str, output_jsonl: str, repeat: int) -> None:
             start.record()
             admission_failed = False
             try:
-                outputs = rank.dp_rank_forward(requests)
-                _output_loss(outputs).backward()
+                outputs = rank.forward(requests)
+                rank.backward(_output_loss(outputs))
             except TrainerRankMemoryError as error:
                 admission_failed = True
                 rows.append(
@@ -629,7 +647,7 @@ def forward(
         ) -> tuple[list[torch.Tensor], dict[str, object]]:
             torch.cuda.synchronize()
             torch.cuda.reset_peak_memory_stats()
-            outputs = rank.dp_rank_forward(requests(), no_grad=no_grad)
+            outputs = rank.forward(requests(), no_grad=no_grad)
             telemetry = rank.last_forward_telemetry()
             logprobs = [
                 output.target_logprobs.detach().float().clone() for output in outputs
@@ -638,9 +656,11 @@ def forward(
 
         def combined_backward(info: dict[str, object]) -> None:
             outputs = cast(list, info["outputs"])
-            torch.stack(
-                [output.target_logprobs.float().sum() for output in outputs]
-            ).sum().backward()
+            rank.backward(
+                torch.stack(
+                    [output.target_logprobs.float().sum() for output in outputs]
+                ).sum()
+            )
             torch.cuda.synchronize()
 
         def unsplit_requirement() -> int:
@@ -684,7 +704,7 @@ def expect_decline(arm: str) -> None:
             torch.cuda.synchronize()
             before = int(torch.cuda.memory_allocated())
             try:
-                rank.dp_rank_forward(requests())
+                rank.forward(requests())
             except TrainerRankMemoryError as error:
                 message = str(error).lower()
                 if "unable to find a feasible split" not in message:
@@ -773,9 +793,14 @@ def pressured_cap() -> int:
             if len(partition) < 2:
                 _fail("reverse-order arm expected a split plan")
             for indices in reversed(partition):
-                torch.stack(
-                    [outputs[index].target_logprobs.float().sum() for index in indices]
-                ).sum().backward()
+                rank.backward(
+                    torch.stack(
+                        [
+                            outputs[index].target_logprobs.float().sum()
+                            for index in indices
+                        ]
+                    ).sum()
+                )
             rank.zero_grad()
             del info, outputs
 
@@ -972,7 +997,7 @@ def conversion(
 # backward through an active LoRA slot, compared against the depth-one arm.
 # ``dp2-tp2-waves`` (4 ranks, DP2 x TP2) exercises the global wave planner's
 # collectives (world scope) composed with TP execution collectives (pair scope)
-# through public ``forward_micro_batches``, including an empty DP slot.
+# through public ``forward_batches``, including an empty DP slot.
 
 TP_GATES: dict[str, float] = {
     # bf16 kernels reorder reductions across packings (same metric/tolerance as
@@ -1337,7 +1362,7 @@ def _write_rows(evidence: str | None, rows: list[dict[str, object]], name: str)
 def phase_tp2_public(
     evidence: str | None, repeat: int, *, tp: int, dump_dir: str | None
 ) -> None:
-    """DP1 x TP{tp} x CP1 public ``dp_rank_forward`` cell (Qwen3.5-4B full model).
+    """DP1 x TP{tp} x CP1 public ``forward`` cell (Qwen3.5-4B full model).
 
     Run at ``--tp 2`` (the gate) and at ``--tp 1`` (the control: the identical
     cell on one GPU). Structural gates run here; the numerics gates compare
@@ -1407,9 +1432,9 @@ def run(arm: str) -> dict[str, Any]:
             start = torch.cuda.Event(enable_timing=True)
             end = torch.cuda.Event(enable_timing=True)
             start.record()
-            outputs = rank.dp_rank_forward(requests, checkpoint=slot)
+            outputs = rank.forward(requests, checkpoint=slot)
             loss = _output_loss(outputs)
-            loss.backward()
+            rank.backward(loss)
             end.record()
             torch.cuda.synchronize()
             telemetry = rank.last_forward_telemetry()
@@ -1759,7 +1784,7 @@ def structured(label: str, profile: dict[str, float]) -> None:
 
 
 def phase_dp2_tp2_waves(evidence: str | None) -> None:
-    """DP2 x TP2 public ``forward_micro_batches`` gate (4 ranks, Qwen3.5-4B).
+    """DP2 x TP2 public ``forward_batches`` gate (4 ranks, Qwen3.5-4B).
 
     Arm A (branchy): six hierarchical GRPO groups as top-level items, the
     test-only memory cap sized so the stream needs at least two waves, forward
@@ -1816,7 +1841,7 @@ def stream(arm: str, cap_bytes: int | None) -> dict[str, Any]:
             logprobs: dict[int, list[torch.Tensor]] = {}
             loss_total = 0.0
             depths: list[int] = []
-            for batch in rank.forward_micro_batches(items, checkpoint=slot):
+            for batch in rank.forward_batches(items, checkpoint=slot):
                 seen.extend(int(index) for index in batch.indices)
                 telemetry = rank.last_forward_telemetry()
                 depths.append(int(telemetry["selected_max_depth"]))
@@ -1840,7 +1865,7 @@ def stream(arm: str, cap_bytes: int | None) -> dict[str, Any]:
                 if flat:
                     loss = _output_loss(flat)
                     loss_total += float(loss.detach().float().item())
-                    loss.backward()
+                    rank.backward(loss)
             torch.cuda.synchronize()
             result = {
                 "arm": arm,
@@ -1955,12 +1980,12 @@ def stream(arm: str, cap_bytes: int | None) -> dict[str, Any]:
         rank.zero_grad()
         waves = 0
         local_outputs = 0
-        for batch in rank.forward_micro_batches([single], checkpoint=slot):
+        for batch in rank.forward_batches([single], checkpoint=slot):
             waves += 1
             flat = [output for group in batch.outputs for output in group]
             local_outputs += len(flat)
             if flat:
-                _output_loss(flat).backward()
+                rank.backward(_output_loss(flat))
         torch.cuda.synchronize()
         counts = _gather_objects((dp_rank, waves, local_outputs))
         rows.append({"arm": "dp2-tp2-empty-slot", "per_rank": counts})
@@ -2625,9 +2650,9 @@ def run(label: str) -> dict[str, object]:
             start.record()
             failed = 0
             try:
-                outputs = rank.dp_rank_forward(requests, checkpoint=slot)
+                outputs = rank.forward(requests, checkpoint=slot)
                 loss = _output_loss(outputs)
-                loss.backward()
+                rank.backward(loss)
             except TrainerRankMemoryError as error:
                 failed = 1
                 message = str(error)