Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions tests/test_npu_910b4_prep.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
7 changes: 7 additions & 0 deletions yeto/rl/adapters/verl/island_entry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
Loading