diff --git a/tests/test_npu_910b4_prep.py b/tests/test_npu_910b4_prep.py index 622f578c..866e83d2 100644 --- a/tests/test_npu_910b4_prep.py +++ b/tests/test_npu_910b4_prep.py @@ -289,3 +289,22 @@ def test_verl_main_keeps_npus_visible_only_on_ascend(monkeypatch): monkeypatch.setattr(cat, "device_family", lambda *_a: "ascend") verl_main._keep_npu_visible_in_cpu_actors() assert os.environ[verl_main.NPU_NOSET_ENV] == "1" + + +def test_fully_async_plan_sets_trainer_device_npu_on_ascend(monkeypatch): + from types import SimpleNamespace as NS + + from yeto.hw import catalog as cat + from yeto.rl.adapters.verl import island_entry + + from yeto.rl.engine.algorithm import AlgorithmSpec, CorrectionSpec, ExecutionSpec + + spec = AlgorithmSpec(execution=ExecutionSpec(max_policy_staleness=1), + correction=CorrectionSpec(method="tis", tis_clip=2.0, tis_clip_low=0.0)).to_dict() + args = NS(groups_per_round=8, samples_per_group=4, global_rounds=3) + monkeypatch.setattr(cat, "device_family", lambda *_a: "ascend") + out, _ = island_entry.fully_async_plan(1, [], spec, args, "none") + assert "trainer.device=npu" in out + monkeypatch.setattr(cat, "device_family", lambda *_a: "nvidia") + out, _ = island_entry.fully_async_plan(1, [], spec, args, "none") + assert not any(o.startswith("trainer.device=") for o in out) diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index 5ac2f8ba..3fc953ed 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -383,6 +383,13 @@ def fully_async_plan(limit: int, overrides: list[str], spec_dict: dict | None, a check_algorithm_contract(profile, spec) out = fully_async_run_overrides(overrides, limit, groups_per_round=args.groups_per_round, rounds=args.global_rounds) + from yeto.hw.catalog import device_family + + if device_family() == "ascend": + # fully_async_ppo_trainer.yaml defaults trainer.device=cuda; on an NPU + # node its resource pools then ask Ray for "GPU" and never schedule + # (S20 ModelArts 2x910B4, runBA_try2_gpures). + out = [o for o in out if not o.startswith("trainer.device=")] + ["trainer.device=npu"] return out, tuple(vconf.ASSERTED_KEYS) + FULLY_ASYNC_ASSERTED_KEYS