diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 25e26533fe..f6a3a19b09 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2928,6 +2928,27 @@ def miles_validate_args(args): _maybe_apply_dumper_overrides(args) +def validate_async_off_policy_correction(args) -> None: + """Require an explicit behavior-policy choice for async PPO training. + + In the async train loop the next rollout is generated before the current + weight update is published, so samples can come from a stale policy. With + the default flags the PPO ratio denominator (``log_probs``) is recomputed + by the *current* actor, silently anchoring clipping (and KL-shaped + advantages) to a policy that never generated the trajectory; the recorded + ``weight_versions`` are a metric, not an enforcement mechanism. + """ + if not args.use_critic: + return + assert args.use_rollout_logprobs or args.use_tis or args.keep_old_actor, ( + "Async PPO training requires an explicit behavior-policy correction, because rollouts are " + "generated before the current weight update while log probs are recomputed by the current " + "actor by default. Pass one of: --use-rollout-logprobs (use the rollout engine's log probs " + "as the ratio denominator), --use-tis (truncated importance sampling correction), or " + "--keep-old-actor (recompute the denominator with the weights the rollout engines used)." + ) + + def _maybe_apply_dumper_overrides(args) -> None: if not args.dumper_enable: return diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 4f84cb5dd4..09f38999d5 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -13,6 +13,7 @@ _resolve_ft_components, get_miles_extra_args_provider, miles_validate_args, + validate_async_off_policy_correction, ) from miles.utils.misc import function_registry @@ -356,3 +357,27 @@ def test_sglang_parallel_size_aliases_keep_last_value(): args = parser.parse_args(["--sglang-data-parallel-size", "2", "--sglang-dp-size", "3"]) assert args.sglang_dp_size == 3 + + +def _make_async_ppo_args(**overrides) -> SimpleNamespace: + defaults = dict( + use_critic=True, + use_rollout_logprobs=False, + use_tis=False, + keep_old_actor=False, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +class TestValidateAsyncOffPolicyCorrection: + def test_ppo_without_correction_is_rejected(self): + with pytest.raises(AssertionError, match="behavior-policy correction"): + validate_async_off_policy_correction(_make_async_ppo_args()) + + @pytest.mark.parametrize("flag", ["use_rollout_logprobs", "use_tis", "keep_old_actor"]) + def test_ppo_with_any_correction_passes(self, flag): + validate_async_off_policy_correction(_make_async_ppo_args(**{flag: True})) + + def test_non_ppo_estimators_are_unaffected(self): + validate_async_off_policy_correction(_make_async_ppo_args(use_critic=False)) diff --git a/train_async.py b/train_async.py index e9c23293d0..16d9a8d561 100644 --- a/train_async.py +++ b/train_async.py @@ -4,7 +4,7 @@ from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from miles.utils import object_store -from miles.utils.arguments import parse_args +from miles.utils.arguments import parse_args, validate_async_off_policy_correction from miles.utils.async_utils import eager_create_task from miles.utils.audit_utils.process_identity import MainProcessIdentity from miles.utils.data import remove_rollout_data_refs @@ -21,6 +21,7 @@ # The framework supports other asynchronous approaches such as fully async (which is shown in examples/full_async). async def train(args): assert not args.colocate, "Colocation is not supported for async training." + validate_async_off_policy_correction(args) configure_logger(args, source=MainProcessIdentity()) maybe_start_periodic_pyspy_dump() # allocate the GPUs