feat(trainer): add multi-objective policy distillation by yulangz · Pull Request #1592 · areal-project/AReaL

yulangz · 2026-08-11T04:18:47Z

Description

Add multi-objective policy distillation (MOPD) to the v1 Megatron RL training stack. The implementation scores routed rollout subsets with persistent teacher companions, materializes actor-owned distillation targets, and composes the MOPD objective with the existing policy loss.

Key changes:

Add structured MOPD teacher, route, and loss configuration with topology and model compatibility checks.
Add weighted reverse-KL target construction, behavioral-mask integration, and configurable importance-ratio truncation.
Keep disk-backed teacher workers alive across phases while offloading their Megatron flat buffers, and strictly drain remote tensors before release.
Release SGLang KV cache, weights, and CUDA graphs at the actor/rollout boundary; preserve physical GPU and transfer-rank identity for multi-node AWEX synchronization.
Propagate role-specific environments and complete port groups to forked local, Ray, and Slurm workers.
Add reusable data hooks and an instruction-following gap reward for experience-distillation workflows.
Add English and Chinese algorithm documentation and one local GSM8K Qwen3-14B-to-0.6B example.

The change set intentionally excludes alternate checkpoint staging, generic scheduler transaction hardening, optional finite-weight diagnostics, duplicate AWEX compatibility code now provided by AWEX 0.8.1, v2/FSDP/Archon MOPD support, and redundant examples/tests.

Related Issue

N/A

Type of Change

Checklist

I have read the Contributing Guide
Pre-commit hooks pass (pre-commit run --all-files)
Relevant tests pass; new tests added for new functionality
Documentation updated
Branch is up to date with main
Self-reviewed via /review-pr command
This PR is a breaking change

Validation

Full pre-commit stack passed after merging the latest main.
Focused Python 3.12 Slurm suite: 312 passed, 2 skipped; the remaining mock assertion was corrected and its targeted rerun passed.
MOPD loss/config CPU suite: 48 passed.
Single-node local GSM8K long run completed all 290 steps without CUDA OOM.
Four-node Qwen3-Coder experience-OPD smoke reached 32 actor workers and eight TP=4 inference endpoints, accepted 64 samples, completed the teacher forward/offload phase, and entered actor PPO.

Full non-slow SGLang and vLLM Slurm suites are running separately; their results will be added when complete.

Breaking Change Details

None.

Enable policy distillation from multiple teacher objectives while sharing colocated training and rollout resources. Key changes: - Add MOPD objectives, routing, losses, and teacher lifecycle management - Support forked colocated workers across local, Ray, and Slurm schedulers - Extend AWEX weight synchronization and memory residency handling - Add bilingual documentation, a local GSM8K example, and focused tests

Copilot

Pull request overview

This PR adds Multi-Teacher On-Policy Distillation (MOPD) to AReaL’s RL training stack, including trainer-side distillation targets/losses and the supporting infrastructure for colocated teacher/actor/rollout lifecycles (notably AWEX colocation and forked worker management).

Changes:

Implement MOPD core logic (compatibility checks, target aggregation, phase machine, loss composition) plus docs and a local GSM8K example.
Extend AWEX/v2 weight-update gateway + adapters for colocated initialization, preflight validation, better error propagation, and teardown semantics.
Add robustness/observability improvements (finite-weight validation, safer distributed stat reductions, process cleanup race handling, data hooks) with broad unit/torchrun test coverage.

Reviewed changes

Copilot reviewed 96 out of 96 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/v2/weight_update/test_wu_controller.py	Adds regression test ensuring connect failures don’t persist pair state.
tests/v2/weight_update/test_pair_registry.py	Adds reservation/commit/release tests for atomic pair-name ownership.
tests/v2/weight_update/test_megatron_adapter.py	Adds unit tests for colocated AWEX preflight/init behavior and backend errors.
tests/v2/weight_update/test_gateway_colocate.py	Adds gateway tests for colocated connect validation, rollback, and duplicate connect handling.
tests/v2/inference_service/test_guard.py	Tightens forked-child cleanup expectations with kill-side effects.
tests/torchrun/test_mopd_fork_lifecycle.py	Adds a hardware-gated fork lifecycle soak test scaffold.
tests/torchrun/test_megatron_mopd_cp2.py	Adds a 2-GPU torchrun regression wrapper for CP scalar reassembly.
tests/torchrun/run_mopd_teacher_residency.py	Adds CUDA regression script for teacher residency/offload behavior.
tests/torchrun/run_mopd_loss_logits.py	Adds CUDA regression script for MOPD loss edge cases and gradients.
tests/torchrun/run_megatron_mopd_cp2.py	Adds torchrun script for CP packed logprob split/reassembly validation.
tests/test_weight_finite.py	Adds unit tests for finite-weight scanning diagnostics and distributed failure propagation.
tests/test_train_controller.py	Adds tests for worker-role rollback semantics and stricter destroy error handling.
tests/test_stats_tracker.py	Adds unit test ensuring NCCL reductions don’t receive CPU tensors.
tests/test_sglang_pp_unit.py	Adds test asserting pause request ordering (abort then retract).
tests/test_saver.py	Adds tests for base-model asset copy inference/skip behavior in periodic saves.
tests/test_rollout_controller.py	Adds tests for colocation port selection/GRES behavior and stricter destroy failures.
tests/test_ray_scheduler.py	Adds retry semantics test for Ray launcher stop failures.
tests/test_proc_utils.py	Adds tests for pidfd EINVAL race handling and zombie/liveness edge cases in process cleanup.
tests/test_mopd_rtensor.py	Adds tests for MOPD target aggregation and strict RTensor drain/clear behavior.
tests/test_mopd_routing.py	Adds tests for route extraction/propagation and eval behavior under MOPD routing.
tests/test_mopd_compatibility.py	Adds tests for teacher-architecture and tokenizer-ID compatibility validation.
tests/test_local_scheduler.py	Updates scheduler cleanup tests to reflect retryable cleanup failures and SIGKILL escalation behavior.
tests/test_inference_engines.py	Adds tests for server-wait failures and SGLang child env sanitization under AWEX.
tests/test_eval_dispatch.py	Adds tests for eval padding alignment and RTensor dummy materialization semantics.
tests/test_data_hooks.py	Adds tests for DataHookManager ordering, rollback, and actor role constraints.
tests/infra/data_service/test_guard.py	Updates guard tests to reflect kill-side effects during cleanup.
tests/experimental/openai/test_tool_call_parser.py	Adds tests for disabling reasoning parsing and for better config-key error messages.
tests/experimental/openai/test_streaming_chat_completions.py	Adds tests ensuring response `model` mirrors request model in streaming/non-streaming paths.
tests/experimental/openai/test_proxy_rollout_server.py	Adds tests ensuring proxy `/health` reports worker identity and SLURM fallback logic.
examples/mopd/README.md	Adds quickstart instructions for a local GSM8K MOPD example.
examples/mopd/gsm8k_qwen3_14b_to_0_6b.py	Adds runnable local MOPD GSM8K distillation entrypoint and dry-run validator.
examples/mopd/gsm8k_qwen3_14b_to_0_6b_local.yaml	Adds a single-node 8-GPU local config for Qwen3 teacher→actor MOPD.
docs/zh/algorithms/mopd.md	Adds Chinese documentation describing MOPD lifecycle/config/operational notes.
docs/en/algorithms/mopd.md	Adds English documentation describing MOPD lifecycle/config/operational notes.
areal/v2/weight_update/training_adapter.py	Extends training adapter contract with colocated preflight method.
areal/v2/weight_update/gateway/pair_registry.py	Adds atomic reservation/commit API to prevent overlapping connect races.
areal/v2/weight_update/gateway/app.py	Adds remote worker error type, better error propagation, colocated preflight, rollback/teardown, and reservation-based connect.
areal/v2/weight_update/controller/controller.py	Improves connect error detail preservation and avoids mutating state on failed connect.
areal/v2/weight_update/awex/megatron_adapter.py	Adds colocated preflight and switches weight offload/reload to MCore flat-buffer semantics.
areal/v2/weight_update/awex/fsdp_adapter.py	Adds explicit failure for colocated preflight on unsupported FSDP backend.
areal/v2/training_service/worker/awex.py	Adds colocated preflight endpoint wiring.
areal/v2/training_service/controller/controller.py	Threads worker `role` / `data_hook_role` through forked worker initialization.
areal/v2/inference_service/sglang/scheduler.py	Adds teardown RPC for weight-update group cleanup.
areal/v2/inference_service/sglang/awex.py	Adds inference-side `/awex/teardown` endpoint for rollback support.
areal/utils/testing_utils.py	Makes model-path resolution lazy to avoid requiring offline model resolution at test collection time.
areal/utils/stats_tracker.py	Adds backend-aware `_all_reduce` helper to avoid invalid device/backend combinations.
areal/utils/saver.py	Infers base HF asset source path from engine config and avoids self-copy on resumed checkpoints.
areal/utils/functional/functional.py	Adds `train_loss_mask` to actor loss stats outputs for downstream consumers.
areal/utils/data.py	Extends `make_dummy_eval_item` to preserve group dims and support active-attention dummies + RTensor templates.
areal/utils/data_hook.py	Adds a general worker-local data hook manager for pre/post lifecycle transformations with rollback/cleanup behavior.
areal/trainer/mopd/targets.py	Adds teacher contribution aggregation into per-trajectory tensors and removes routing metadata.
areal/trainer/mopd/phase.py	Adds a strict phase machine for rollout→teacher→drain→train transitions.
areal/trainer/mopd/loss.py	Adds MOPD reverse-KL surrogate with capped importance sampling and composition with RL loss.
areal/trainer/mopd/compatibility.py	Adds checkpoint/tokenizer fingerprinting and compatibility validation for persistent multi-teacher control.
areal/trainer/mopd/init.py	Exposes MOPD public API surface.
areal/reward/if_gap.py	Adds instruction-following “gap reward” with optional dependency loading.
areal/reward/init.py	Exports `if_gap_reward_fn` via reward registry.
areal/infra/workflow_executor.py	Adds rollout-side DataHookManager integration and ensures hooks are closed on destroy.
areal/infra/utils/slurm.py	Extends cancel signal options and makes cancellation failures propagate.
areal/infra/utils/proc.py	Adds pidfd EINVAL race fallback and zombie-aware liveness checks during process tree termination.
areal/infra/scheduler/fork_utils.py	Adds reusable fork-role lifecycle/ownership helpers (reservation, release, reconcile).
areal/infra/scheduler/exceptions.py	Adds `WorkerCleanupError` to represent partial cleanup that must be retried.
areal/infra/rpc/rtensor.py	Strengthens shard delete error handling, returns cleared counts, and adds fetch-buffer matching stats helper.
areal/infra/remote_inf_engine.py	Documents two-stage pause semantics for SGLang in colocated workflows.
areal/infra/controller/rollout_controller.py	Adds MOPD route extraction/propagation, colocation GPU/port handling, and stricter destroy rollback behavior.
areal/experimental/openai/tool_call_parser.py	Adds sglang-missing fallback behavior, reasoning splitting helper, and more actionable config-key errors.
areal/experimental/openai/proxy/proxy_rollout_server.py	Adds worker identity reporting, safer worker index resolution, and preserves request model forwarding.
areal/experimental/openai/client.py	Propagates request `model` into completion/stream/response objects (fixes downstream adapters).
areal/experimental/engine/archon_engine.py	Adds data hook setup/cleanup integration and exposes `aggregate_mopd_targets`.
areal/engine/weight_finite.py	Adds environment-gated finite-weight scanning with chunking and distributed “fail all ranks” semantics.
areal/engine/sglang_remote.py	Switches readiness check away from `/health`, adds env sanitization, and supports AWEX meta addr via env var.
areal/engine/fsdp_engine.py	Adds data hook setup/cleanup integration and exposes `aggregate_mopd_targets`.
areal/engine/awex/colocate_writer.py	Adds finite-weight checks around conversion + improves inference completion error propagation.
areal/api/scheduler_api.py	Tightens delete_workers contract and adds helpers for resolving/deleting colocated role chains safely.

Suppressed comments (1)

areal/v2/weight_update/awex/megatron_adapter.py:583

The fallback reload path depends on buf.param_data.cpu_data, which has the same attribute-assignment risk as the offload path. After switching to a buffer-owned CPU mirror (e.g. buf.cpu_param_data), reload should read from that buffer attribute instead.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+                        if buf.param_data.storage().size() == 0:
+                            continue
+                        if not hasattr(buf.param_data, "cpu_data"):
+                            buf.param_data.cpu_data = torch.empty_like(
+                                buf.param_data, device="cpu", pin_memory=True
+                            )
+                        buf.param_data.cpu_data.copy_(buf.param_data)
+                        buf.param_data_size = buf.param_data.storage().size()
+                        buf.param_data.storage().resize_(0)


yulangz · 2026-08-24T09:44:38Z

@@ -265,8 +265,13 @@ async def delete(self, node_addr: str, shard_ids: list[str]) -> None:
            async with session.delete(


Fixed on the current branch. HttpRTensorBackend.delete is annotated as returning int, matching the protocol and its cleared-count result.

Forked colocated roles inherited only the owner environment, so rollout allocator settings could not differ from actor and teacher settings. Pass per-task overrides through every scheduler fork path and expose the active allocator configuration in residency logs.

Free CUDA graph allocations alongside weights and KV cache while colocated training owns the GPU. Enable SGLang's graph memory saver before capture and restore graphs after AWEX weight publication. Key changes: - preserve the preload hook required for pauseable CUDA graphs - pass the cuda_graph tag through rollout offload and restore - cover child environment and restore ordering with unit tests

sitabulaixizawaluduo · 2026-08-24T06:13:55Z

 logger = getLogger("AwexColocateReader")


+def _patch_awex_qwen3_attention_names() -> None:


The awex release version should already support the co-card of Qwen3, so this converter is no longer needed.

Fixed in e035e94. The Qwen3 tied-lm-head compatibility converter was removed and the released AWEX implementation is used directly. AWEX 0.8.1 API tests passed on Slurm.

sitabulaixizawaluduo · 2026-08-24T06:44:47Z

+_patch_awex_qwen3_attention_names()
+
+
+def _add_tied_lm_head_meta_alias(


The same as before, remove it

Fixed in e035e94. The tied lm_head metadata alias and its call site were removed.

sitabulaixizawaluduo · 2026-08-24T07:04:04Z

-                        )
-                        param.data = torch.empty(0, device="cpu")
-                        count += 1
+                raise RuntimeError(


If the current AWEX colocate mode doesn't support wrap_with_ddp=False, the check should be performed in MegatronEngineConfig.post_init().

Fixed in e035e94. Validation is performed in TrainEngineConfig.post_init, which owns both weight_update_mode and the nested Megatron configuration. AWEX now rejects megatron.wrap_with_ddp=false before engine startup.

sitabulaixizawaluduo · 2026-08-24T07:05:28Z

        return default


+def _bool_env(name: str, default: bool = False) -> bool:


The get_bool_env_var function is already in areal/utils/environ.py, you can call it directly.

Fixed in e035e94. The plugin now uses the shared boolean environment helper from areal.utils.environ; the helper also recognizes the common true and false spellings consistently.

sitabulaixizawaluduo · 2026-08-24T07:06:57Z

+    return value.strip().lower() in {"1", "true", "yes", "y", "on"}
+
+
+def _int_env(name: str, default: int) -> int:


It can be written in areal/utils/environ.py to avoid similar repetitive constructions in the future.

Fixed in e035e94. Shared integer and float environment parsers were added to areal.utils.environ and the AWEX code now uses them.

sitabulaixizawaluduo · 2026-08-24T07:09:42Z

+        return default
+
+
+def _scheduler_int_attr(scheduler: Any, name: str, default: int) -> int:


Write the current and several following methods into the AwexSchedulerPlugin, except there should be no other places to call these functions here.

Fixed in e035e94. Scheduler-specific environment, callable, instance, and device mapping helpers are now encapsulated in AwexSchedulerPlugin.

sitabulaixizawaluduo · 2026-08-24T07:22:13Z


        self._eager_publish_awex_train_info(meta_server_addr)

+    def init_weight_residency_adapter(self) -> None:


It's a bit strange to use AwexMegatronAdapter to manage the onload/offload of flat buffers, as it has no relation with weight synchronization. A more elegant approach would be to create a new MegatronWeightResidencyAdapter to manage these, and let AwexMegatronAdapter manage the memory release and loading under colocate mode through it. Here, it is sufficient to call MegatronWeightResidencyAdapter, but this does not block the current PR, and it can be merged in and addressed later.

Agreed. This coupling predates MOPD, but MOPD introduces a residency-only consumer and therefore makes the boundary part of this change: init_weight_residency_adapter() names a separate concern yet still constructs AwexMegatronAdapter, and _awex_adapter is not None changes save, train_batch, and offload/onload behavior. Please split a generic MegatronWeightResidency from AwexWeightPublisher, with the publisher composing the residency service. That lets the teacher depend only on residency without acquiring AWEX publication state.

Acknowledged. Per the current review scope, this non-blocking residency-adapter separation is intentionally deferred to a follow-up and is not mixed into this PR update.

Implemented in 34a4535. MegatronEngine now owns a single MegatronWeightResidency. AwexWeightPublisher composes that same instance and contains only AWEX publication, MetaServer, converter, and IPC state. init_weight_residency_adapter creates no AWEX state, so the MOPD teacher depends only on residency. The AwexMegatronAdapter compatibility alias is kept for the existing v2 caller. Native and fallback flat-buffer round trips plus publication ordering are covered by tests.

sitabulaixizawaluduo · 2026-08-24T07:29:02Z

Actually returns int, inconsistent with the protocol

Fixed on the current branch. HttpRTensorBackend.delete now returns and is annotated as int, consistent with RTensorBackend.

sitabulaixizawaluduo · 2026-08-24T07:54:24Z

+                # rollout owns its worker and can use its second port. Always
+                # override the global server argument because replicas on the
+                # same node cannot safely share one explicit NCCL port.
+                port_index = 1 if self.config.scheduling_strategy.fork else 2


Here, the fork decision determines whether to use worker_ports[1] or [2], which is reasonable under the current port layout: non-fork mode must avoid the actor's TCPStore, while fork mode can use the second port of an independent worker. However, this relies on an implicit convention that the indices in worker_ports correspond to port usage, and fork worker ports are allocated by different parent guards. Currently, allocated_ports are a process-level state, and find_free_ports() checks without persistently occupying sockets. Therefore, when multiple guards on the same node allocate concurrently, can it guarantee that these nccl_ports are globally unique?
It is recommended that the scheduler explicitly allocates/retains named NCCL ports, rather than hardcoding indices in the RolloutController. Additionally, non-fork colocation requires at least three ports, and it's best to fail fast during configuration validation, rather than reporting an error in initialize() after creating workers.

Fixed in e035e94. Fork ports now have explicit owner records, node-wide reservation locks, idempotent release, and cleanup on spawn/configuration/kill failures. Local, Slurm, Ray, and v2 service paths allocate the target-specific port count and release the same reservation lifecycle.

sitabulaixizawaluduo · 2026-08-24T08:00:14Z

-        """Get SGLang health check request."""
-        return HttpRequest(endpoint="/health", payload={}, method="GET")
+        """Get SGLang readiness check request."""
+        # SGLang's /health is a functional 1-token generation probe in recent


The modification will affect all SGLang paths, not just AWEX. The normal separation mode will no longer perform functional generation checks, so some issues that are only revealed during actual reasoning will be delayed until the first request.
Suggestion: AWEX startup: use /model_info, retain /health for normal SGLang

Fixed in e035e94. Normal SGLang startup keeps the /health generation check; only the AWEX launch instance uses /model_info.

fishcrap · 2026-08-24T08:11:34Z

+            loss,
+            config=mopd_config,
+            logprobs=logprobs,
+            old_logprobs=old_logp,


Please keep the rollout behavior log-probability separate here. With the currently valid recompute_logprob=true and use_decoupled_loss=false configuration, the advantage-preparation path replaces data["logprobs"] with recomputed proximal log-probabilities, so old_logp here is no longer $\pi_{behavior}$. The MOPD importance ratio therefore becomes $\pi_\theta / \pi_{prox}$ (often close to 1) instead of $\pi_\theta / \pi_{behavior}$, silently disabling the intended off-policy correction. Preserve an immutable behavior_logp for MOPD, or reject this configuration combination.

Fixed in e035e94. Rollout behavior log-probabilities are captured in a dedicated immutable tensor before proximal recomputation and are used for the MOPD importance ratio.

fishcrap · 2026-08-24T08:11:34Z

+            # Drain the old IDs now: the caller replaces ``rollout_batch`` with
+            # ``aggregated`` and its step-end cleanup can no longer discover
+            # them. Teachers also localized the same rollout shards while
+            # scoring, and a configured critic localized them before this phase,


This drain fan-out also needs to include the configured reference engine. self.ref.compute_logp(rollout_batch) localizes these original shard IDs earlier in the step, but this method returns fresh aggregated shards and the caller replaces rollout_batch, so the step-end ref cleanup can no longer identify the old shards. With KL/ref enabled, this leaks ref fetch-buffer entries on every MOPD step. Drain every consumer, including ref, before release; preferably centralize consumer ownership instead of hard-coding role-specific drains.

Fixed in e035e94. The reference engine is included in both normal and emergency strict drains before original shard release, with release gated on complete consumer receipts.

fishcrap · 2026-08-24T08:11:34Z

+            nnodes=nnodes,
+            instance_world_size=_scheduler_instance_world_size(self._scheduler),
+        )
+        physical_gpu_id = _resolve_physical_gpu_id(


Please do not overwrite the CVD-derived physical GPU identity from transfer_rank. On Ray multi-node runs, SLURM_NODEID/SLURM_NNODES fall back to 0/1; with non-contiguous CUDA_VISIBLE_DEVICES, this modulo mapping also produces logical IDs rather than the physical IDs used by the writer. The reader and writer can therefore publish/look up different MetaServer keys, causing collisions or a permanent wait. Pass an explicit scheduler-derived topology descriptor (node_index, node_count, physical_gpu_id, transfer_rank), or reject unsupported scheduler/topology combinations.

Fixed in e035e94. Transfer rank remains a dense logical identity, while the physical GPU identity remains derived from CUDA_VISIBLE_DEVICES. Unsupported UUID and non-numeric mappings are rejected instead of silently remapped.

fishcrap · 2026-08-24T08:11:34Z

+        trajectory["mopd_teacher_logp_sum"] = logp_sum
+        trajectory["mopd_teacher_weight_sum"] = weight_sum
+        trajectory["mopd_rl_coefficient"] = rl_coefficient
+        trajectory["mopd_distillation_coefficient"] = distillation_coefficient


Please avoid serializing static loss configuration into every trajectory. These string-key scalars cross the generic batching path, where non-tensor values are taken from the first trajectory without consistency validation, and the actor then reconstructs an incomplete MOPDLossConfig; this is how importance_ratio_cap was silently dropped. Bind the complete MOPDLossConfig once to the actor/loss, and keep only per-sample teacher target tensors in the batch.

Fixed in e035e94. Static MOPD loss configuration is bound once to the actor, while trajectory batches now contain only per-sample teacher target tensors.

sitabulaixizawaluduo · 2026-08-24T08:12:18Z

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+


Is there any correlation between this feature and the features of mopd? If so, please reply to me. If not, please propose the relevant changes as a separate feature.

There is no required MOPD dependency on DataHook. The generic DataHook implementation, workflow integration, and tests have therefore been removed from this PR in e035e94.

Make routing an explicit property of each configured dataset source so workflows and persisted samples no longer need routing fields. This preserves arbitrary teacher fan-out while keeping rollout inputs clean. Key changes: - Support local and remote routed dataset mixtures - Carry routes as controller task metadata into trajectories - Apply remaining MOPD review fixes and AWEX 0.8.1 cleanup - Update examples, bilingual docs, and regression coverage

sitabulaixizawaluduo · 2026-08-24T08:27:28Z

scheduler requested the full port_cnt from the parent guard's /alloc_ports, but these ports were only added to GuardState.allocated_ports without establishing ownership with (role, worker_index). The /kill_forked_worker only terminates child processes and does not remove the corresponding ports from allocated_ports. Therefore, during the same lifecycle of a parent guard, logical ports may continue to leak regardless of normal destruction, failed startup rollback, or repeated forks, potentially exhausting available ports.
It is recommended to save the port list in the fork child record and atomically release it during kill/rollback; or to add an explicit /release_ports interface and ensure all cleanup paths in Local/Ray/Slurm call it. LocalScheduler also has its own _allocated_ports, so it's best to unify port ownership to avoid scheduler and guard using two separate, unaware allocators.

fishcrap

This PR now combines a high-risk algorithm/lifecycle change with several reusable or unrelated infrastructure changes (64 files and roughly 8.5k added lines). I recommend splitting it into five reviewable changesets:

Ancillary independent changes: extract the OpenAI parser/client/proxy fixes, IF-gap reward, and DataHook framework from the MOPD stack. None of these should be coupled to MOPD review or rollout.
Generic routed dataset mixture: introduce multiple dataset sources, source/route provenance, a typed task/result route, and an explicit sampling policy. This layer should not contain MOPD-specific dataset or controller APIs.
Generic runtime foundations: add explicit RuntimeTopology, transactional/idempotent fork handles, port ownership, and a typed RTensor batch-lease/drain contract.
Megatron residency refactor: separate MegatronWeightResidency from AwexWeightPublisher, and keep rollout/actor ownership transitions at the controller boundary.
MOPD core + v1 orchestration: keep the algorithm and its first supported vertical slice together—typed/static loss configuration, immutable behavior log-probabilities, target aggregation, execution planning, teacher compatibility/manager, a narrow MOPDTeacherPhase, trainer integration, docs, examples, and tests.

For the fifth changeset, I would deliberately keep the first support matrix narrow: Megatron v1 + SGLang + AWEX, process-isolated actor/teacher/rollout workers, and the disk-backed teacher provider. Local-memory staging, same-process rollout support, and additional backends can follow once their generic ownership contracts exist.

Keeping MOPD core and orchestration together makes sense because they form one testable end-to-end feature. The reusable prerequisites should land separately so generic controllers, schedulers, and engines do not acquire MOPD-specific APIs or hidden data contracts.

fishcrap · 2026-08-24T09:12:18Z

+            raise
+        return controller
+
+    def _run_mopd_teacher_phase(


Please extract this transaction from PPOTrainer into a narrow MOPDTeacherPhase instead of making the trainer own teacher selection, loading, scoring, target materialization, RTensor drain, release, and rollback. A small materialize(batch) / close() component with injected routes, TeacherManager, actor target aggregation, and BatchDrainers would reduce the current lifecycle cross-product. AWEX rollout/actor ownership transitions should remain in the outer trainer; passing the entire trainer into a broad MOPDCoordinator would only create a second god object.

Implemented in 34a4535. The transaction is now an MOPDTeacherPhase with materialize and close methods and injected config, manager, actor, and optional drainers. PPOTrainer retains only AWEX ownership transitions and the factory. The phase tracks actor, critic, reference, and unique-teacher receipts, releases original shards only after all required ACKs, and drains missing consumers before closing on failure.

fishcrap · 2026-08-24T09:12:18Z

                self._worker_role,
            )
+
+    def strict_clear_batches(self, *targets: Any) -> dict[str, int | bool]:


Please make this lifecycle boundary typed and generic. Returning an anonymous dict and reconstructing DrainReceipt(**result) in the MOPD trainer is a fragile cross-layer contract: a field rename can fail after destructive drain side effects have already completed, and actor-specific field names leak into generic infrastructure. Define a frozen infra-owned RTensorDrainReceipt and pair it with an expected-consumer/lease contract, so completeness proves that every registered consumer (actor, critic, ref, and teachers) was drained rather than relying on a hard-coded role list.

Implemented the typed boundary in 34a4535. RTensorDrainReceipt is a frozen infra-level receipt carrying consumer_role, shard_count, source_node_count, and consumer_dp_head_count. strict_clear returns only after DELETE acknowledgements and all DP heads report clean, and the teacher phase keeps an explicit fixed consumer ledger. I did not introduce a full system-wide lease/consumer registry in this PR: current consumers are statically enumerable, while a dynamic lease system would require coordinated changes to RTensor creation, localization, and every consumer. That broader protocol should be designed separately if dynamic consumers are introduced.

fishcrap · 2026-08-24T09:12:18Z

+        if (
+            teacher_schedule.type != SchedulingStrategyType.colocation.value
+            or teacher_schedule.target != "actor"
+            or not teacher_schedule.fork


Please frame target='actor' plus fork=true as a supported v1 colocated-runtime topology, not as an inherent MOPD requirement. MOPD can mathematically use separately scheduled or remote teachers; this restriction exists because the current implementation reuses actor GPUs and needs a separate process to isolate CUDA, Megatron MPU/process-group state, and other process-global state. Prefer validating a named runtime capability/topology here and document the v1 limitation explicitly, leaving room for dedicated teacher resources later.

Accepted the capability framing in 34a4535: validation, documentation, and errors now describe the supported MOPD v1 colocated-runtime topology. I did not add a RuntimeTopology class. SchedulingStrategy, ModelAllocation and ParallelStrategy, and live Megatron process-group assertions already own the declarative, scheduled, and runtime topology respectively; another wrapper for the single supported shape would duplicate state and could become stale. It is worth revisiting when a second topology or capability negotiation is added.

fishcrap · 2026-08-24T09:12:18Z

+MOPD_ROUTE_METADATA_KEY = "__areal_mopd_route"
+
+
+class MOPDDataset:


Source-based routing is the right direction, but this should be a generic routed dataset/mixture abstraction rather than an MOPD-owned dataset. _DatasetConfig.sources now looks generic, while get_custom_dataset rejects it and only MOPDDataset, MOPD_ROUTE_METADATA_KEY, and MOPD-aware controller code can consume it. Please extract a generic RoutedDataset/DatasetMixture with typed out-of-band source/route provenance and an explicit sampling policy; MOPD should only interpret the route. The current concatenation is implicitly length-proportional, so larger sources also dominate training without a configurable mixture policy.

Implemented in 34a4535. RoutedDataset and get_routed_dataset are generic dataset APIs, DatasetRoute is immutable provenance stripped at the rollout boundary, every DatasetSourceConfig requires a route, and only the MOPD layer interprets route values. Mixture behavior is now explicit: proportional preserves source cardinalities, while uniform deterministically cycles shorter sources before the distributed sampler shuffle. get_custom_dataset accepts sources generically. I did not add arbitrary weighted sampling in this PR; the two explicit policies remove the implicit behavior without introducing a new stateful sampler framework.

fishcrap · 2026-08-24T09:12:18Z

+
+
+@dataclass
+class MOPDTeacherEngineConfig(PPOActorConfig):


Please avoid modeling a forward-only teacher as an empty subclass of the full PPOActorConfig. This advertises optimizer, PPO objective, advantage/reward normalization, and other training-only settings that the teacher does not support, forcing validators to reject inherited capabilities later. Prefer a purpose-built ScoringEngineConfig, composed from the model/runtime/parallel/residency settings the teacher actually needs.

Implemented in 34a4535. MOPDTeacherEngineConfig now derives from TrainEngineConfig and no longer exposes PPO objective or advantage/reward-normalization fields; its optimizer is disabled and dropout defaults are scoring-safe. MegatronScoringEngine provides only eval and forward log-probability computation, while MOPDTeacherController preserves pipeline padding and active dummy batches. The teacher factory no longer constructs MegatronPPOActor.

fishcrap · 2026-08-24T09:12:18Z

+                logger.info("[AWEX] colocate: offload cuda_graph...")
+                self.rollout.offload(tags=["cuda_graph"])
+                try:
+                    if self.config.mopd is not None:


Please derive an explicit MOPD execution plan from the configured coefficients before entering the infrastructure phases. With distillation_coefficient=0, this still loads, scores, drains, and releases teachers; with rl_coefficient=0, the later path still computes advantages and the full RL objective before discarding it. A disabled objective should not require its infrastructure to succeed. Bind a static MOPDExecutionPlan (for example requires_teacher_scoring, requires_rl, and requires_ref) and short-circuit the corresponding phases.

Implemented in 34a4535 with a frozen MOPDExecutionPlan derived before infrastructure setup. Pure RL skips compatibility checks, teacher manager/fork/scoring/drain, and teacher tensors. Pure distillation skips critic/reference setup, reward/GAE computation, and the PPO surrogate, preparing only aligned MOPD tensors. Joint training retains the existing behavior; proximal log-probabilities are requested only when RL, M2 weighting, or rejection filtering needs them.

fishcrap · 2026-08-24T09:18:10Z

+            or rollout_schedule.target != "actor"
+        ):
+            raise ValueError("mopd rollout must use colocation target='actor'")
+        if not rollout_schedule.fork:


Please do not accept a supported topology while warning that its metrics may be incorrect. With fork=false, rollout and actor reuse the same worker process and therefore the same module-global stats_tracker registry; actor-side export_stats(reset=True) can export and clear rollout scopes before the rollout controller reads them, making telemetry order-dependent or empty. For the initial MOPD support matrix, require fork=true. If same-process rollout is needed later, first introduce owner-scoped tracker registries (or explicit export/reset ownership) so actor and rollout metrics have independent lifecycles.

Implemented in 34a4535. The current MOPD v1 rollout configuration now rejects fork=false and requires the colocated actor target to fork. Minimum actor port validation follows the forked path as well. This removes the unsupported shared-process stats_tracker ownership case instead of allowing it with a warning.

Preserve rollout behavior probabilities and make teacher data draining explicit so MOPD loss inputs remain correct across phase transitions. Key changes: - validate AWEX topology and normalize environment parsing - make fork port ownership and cleanup explicit across schedulers - remove the unused generic DataHook layer and add regression tests - restore vLLM parser fallback when SGLang is unavailable Refs: #1592

yulangz · 2026-08-24T09:45:00Z

Addressed the fork-port lifecycle concern in e035e94. Reservations are now owner-bound, guarded across parent Guard processes on the same node, released atomically on child termination and startup rollback, and cleaned through explicit idempotent release paths. Scheduler and v2 service call sites now use target-specific port counts and preserve ownership through fork/configuration cleanup. Regression coverage is included in the guard and scheduler tests.

Make teacher scoring, RTensor draining, and Megatron residency explicit so disabled objectives do not initialize unused infrastructure and AWEX publication no longer owns generic residency state. Key changes: - Extract teacher phase, execution plan, and scoring-only engine - Add typed drain receipts and explicit routed mixture policies - Separate Megatron residency from AWEX publication

Resolve overlapping MOPD and mainline lifecycle changes while preserving the split Megatron residency ownership and main's DTE runtime support. Key changes: - Compose MOPD execution planning with DTE initialization - Keep generic environment parsers and typed numeric helpers - Regenerate CLI docs from the merged config schema

Match inference server reservations with Guard ownership so strict fork validation succeeds, and roll back partial concurrent startup.

Copilot AI lite review requested due to automatic review settings August 11, 2026 04:18

yulangz requested review from CormickKneey, HwVanICI, PrometheusComing, TaoZex, fishcrap, garrett4wade, geshi001, guozhihao-224, nuzant, rchardx and sitabulaixizawaluduo as code owners August 11, 2026 04:18

Copilot started reviewing on behalf of yulangz August 11, 2026 04:19 View session

Copilot AI reviewed Aug 11, 2026

View reviewed changes

yulangz added the safe-to-test Ready to run unit-tests in a PR. label Aug 11, 2026

yulangz had a problem deploying to AReaL-unittests August 11, 2026 05:24 — with GitHub Actions Failure

yulangz added 2 commits August 11, 2026 18:37

yulangz added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 11, 2026

Merge branch 'main' into feature/mopd

a0be494

yulangz added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 11, 2026

yulangz had a problem deploying to AReaL-unittests August 11, 2026 12:59 — with GitHub Actions Failure

yulangz removed the safe-to-test Ready to run unit-tests in a PR. label Aug 12, 2026

fix: restore MOPD staging and Qwen3 weight sync

6dc8fe4

sitabulaixizawaluduo reviewed Aug 24, 2026

View reviewed changes

fishcrap reviewed Aug 24, 2026

View reviewed changes

sitabulaixizawaluduo reviewed Aug 24, 2026

View reviewed changes

fishcrap reviewed Aug 24, 2026

View reviewed changes

yulangz added 2 commits August 24, 2026 20:34

yulangz added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 25, 2026

yulangz had a problem deploying to AReaL-unittests August 25, 2026 04:04 — with GitHub Actions Failure

fix(infra): bind inference forks to port owners

5624252

Match inference server reservations with Guard ownership so strict fork validation succeeds, and roll back partial concurrent startup.

yulangz added safe-to-test Ready to run unit-tests in a PR. and removed safe-to-test Ready to run unit-tests in a PR. labels Aug 25, 2026

yulangz temporarily deployed to AReaL-unittests August 25, 2026 09:51 — with GitHub Actions Inactive

		@@ -265,8 +265,13 @@ async def delete(self, node_addr: str, shard_ids: list[str]) -> None:
		async with session.delete(

		logger = getLogger("AwexColocateReader")


		def _patch_awex_qwen3_attention_names() -> None:

		_patch_awex_qwen3_attention_names()


		def _add_tied_lm_head_meta_alias(

		return default


		def _bool_env(name: str, default: bool = False) -> bool:

		return value.strip().lower() in {"1", "true", "yes", "y", "on"}


		def _int_env(name: str, default: int) -> int:

		return default


		def _scheduler_int_attr(scheduler: Any, name: str, default: int) -> int:


		self._eager_publish_awex_train_info(meta_server_addr)

		def init_weight_residency_adapter(self) -> None:

		MOPD_ROUTE_METADATA_KEY = "__areal_mopd_route"


		class MOPDDataset:

Uh oh!

Conversation

yulangz commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

Checklist

Validation

Breaking Change Details

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yulangz commented Aug 11, 2026 •

edited

Loading