diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index fdcad6cc22..aa5b2e58c9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -747,7 +747,7 @@ def auto_setup_bench(self): @model_validator(mode="after") def validate_no_inference_broadcast(self): - """Static SFT uses the filesystem marker as its trainer handshake.""" + """NCCL broadcast requires a running inference deployment.""" if not self.needs_inference and self.weight_broadcast.type != "filesystem": raise ValueError( "dataset-backed SFT requires filesystem weight broadcast — there is no " diff --git a/src/prime_rl/orchestrator/static_sft.py b/src/prime_rl/orchestrator/static_sft.py index 1bcc4677bd..25cd91d61d 100644 --- a/src/prime_rl/orchestrator/static_sft.py +++ b/src/prime_rl/orchestrator/static_sft.py @@ -21,7 +21,7 @@ from prime_rl.utils.config import to_toml_dict from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import format_time, get_logger, setup_logger -from prime_rl.utils.pathing import get_broadcast_dir, get_step_path, wait_for_path +from prime_rl.utils.pathing import get_step_path, get_trainer_step_dir, wait_for_path from prime_rl.utils.utils import resolve_latest_ckpt_step if TYPE_CHECKING: @@ -260,7 +260,7 @@ async def start(self) -> None: samples, attempts = await source.build_batch(config.batch_size, config.token_batch_size) await self.sender.send(TrainingBatch(examples=samples, step=step)) - stable = get_step_path(get_broadcast_dir(config.output_dir), step) / "STABLE" + stable = get_step_path(get_trainer_step_dir(config.output_dir), step) / "STABLE" await wait_for_path(stable) num_tokens = sum(len(sample.token_ids) for sample in samples) diff --git a/src/prime_rl/trainer/rl/broadcast/filesystem.py b/src/prime_rl/trainer/rl/broadcast/filesystem.py index d0a453f0a1..ea6a4eebff 100644 --- a/src/prime_rl/trainer/rl/broadcast/filesystem.py +++ b/src/prime_rl/trainer/rl/broadcast/filesystem.py @@ -37,6 +37,9 @@ def __init__( def broadcast_weights(self, model: nn.Module, step: int) -> None: """Broadcast weights by saving a HF-compatible checkpoint to shared filesystem and notifies the orchestrator.""" + ready_idxs = self.multi_run_manager.ready_to_broadcast_idxs + if not ready_idxs: + return self.logger.debug("Starting broadcasting weights to inference engine via shared filesystem") start_time = time.perf_counter() adapter_only = self.lora_config is not None @@ -50,7 +53,7 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: state_dict = revert_weight_conversion(model, state_dict) - for idx in self.multi_run_manager.ready_to_update_idxs: + for idx in ready_idxs: self.logger.debug( f"Broadcasting weights for run {idx} (ready_to_update={self.multi_run_manager.ready_to_update[idx]})" ) @@ -99,8 +102,6 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: self.logger.warning(f"Run {idx} is deleted, skipping") except Exception as e: self.logger.error(f"Error broadcasting weights for run {idx}: {e}") - finally: - self.multi_run_manager.ready_to_update[idx] = False if self.world.is_master: self.logger.debug(f"Weights broadcasted in {time.perf_counter() - start_time:.2f}s") diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 75abf42195..2bba4a72fb 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -207,6 +207,8 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: # the orchestrator has paused inference, the collectives sit unmatched # until NCCL's watchdog kills the process after 10 min. notified_runs = self._compute_notified_runs() + if not notified_runs: + return if self.world.is_master: self._notify_orchestrator(notified_runs) self._wait_for_nccl_ready(notified_runs) @@ -223,9 +225,7 @@ def _compute_notified_runs(self) -> list[tuple[int, Path]]: non-master ranks agree on which NCCL_READY markers to wait for. """ notified_runs: list[tuple[int, Path]] = [] - for idx in self.multi_run_manager.used_idxs: - if not self.multi_run_manager.ready_to_update[idx]: - continue + for idx in self.multi_run_manager.ready_to_broadcast_idxs: try: # pack() already advanced progress to the next step, so the model we just # trained — policy v(step-1) — broadcasts to broadcasts/step_{step-1}. @@ -241,10 +241,10 @@ def _compute_notified_runs(self) -> list[tuple[int, Path]]: return notified_runs def _notify_orchestrator(self, notified_runs: list[tuple[int, Path]]) -> None: - """Create STABLE markers for each notified run and clear their ready flags. + """Create STABLE markers for each notified run. - Master-only side effects (filesystem writes + state mutation). Called - after `_compute_notified_runs`; non-master ranks skip this entirely. + Called after `_compute_notified_runs`; non-master ranks skip this + filesystem side effect entirely. """ for idx, save_dir in notified_runs: try: @@ -255,8 +255,6 @@ def _notify_orchestrator(self, notified_runs: list[tuple[int, Path]]) -> None: self.logger.warning(f"Run {idx} is deleted, skipping") except Exception as e: self.logger.error(f"Error broadcasting weights for run {idx}: {e}") - finally: - self.multi_run_manager.ready_to_update[idx] = False def _wait_for_nccl_ready(self, notified_runs: list[tuple[int, Path]]): """Wait for inference workers to signal they are ready to receive NCCL broadcast.""" diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 5bd015c0ca..a20ef36ee6 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -590,9 +590,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: weight_broadcast.maybe_clean(interval_to_keep) else: broadcast_weights_time = 0 - # Usually the broadcast will set this. If broadcast is skipped, we need to reset this here. - for idx in multi_run_manager.used_idxs: - multi_run_manager.ready_to_update[idx] = False + + multi_run_manager.mark_ready_steps_complete() # Checkpoint the step we just finished (model = policy v{progress.step}). if config.max_concurrent_runs > 1: diff --git a/src/prime_rl/trainer/runs.py b/src/prime_rl/trainer/runs.py index 5da62a39b5..49688f547e 100644 --- a/src/prime_rl/trainer/runs.py +++ b/src/prime_rl/trainer/runs.py @@ -12,7 +12,7 @@ from prime_rl.configs.trainer import LoRAConfig from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger -from prime_rl.utils.pathing import get_all_ckpt_steps, get_stable_ckpt_steps +from prime_rl.utils.pathing import get_all_ckpt_steps, get_stable_ckpt_steps, get_step_path, get_trainer_step_dir if TYPE_CHECKING: from prime_rl.configs.orchestrator import OrchestratorConfig @@ -486,6 +486,24 @@ def used_idxs(self): def ready_to_update_idxs(self): return [idx for idx, ready in enumerate(self.ready_to_update) if ready] + @property + def ready_to_broadcast_idxs(self) -> list[int]: + return [idx for idx in self.used_idxs if self.ready_to_update[idx] and self.config[idx].needs_inference] + + def mark_ready_steps_complete(self) -> None: + """Publish trainer-step completion and release every run updated this step.""" + completed_idxs = [idx for idx in self.used_idxs if self.ready_to_update[idx]] + if self.world.is_master: + for idx in completed_idxs: + run_dir = self.get_run_dir(idx) + if not run_dir.exists(): + continue + step_dir = get_step_path(get_trainer_step_dir(run_dir), self.progress[idx].step - 1) + step_dir.mkdir(parents=True, exist_ok=True) + (step_dir / "STABLE").touch() + for idx in completed_idxs: + self.ready_to_update[idx] = False + def run_dirs(self) -> list[Path]: return [self.output_dir / run_id for run_id in self.id_2_idx.keys()] diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index bad9fdcf05..cbc97f862d 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -88,6 +88,10 @@ def get_broadcast_dir(output_dir: Path) -> Path: return output_dir / "broadcasts" +def get_trainer_step_dir(output_dir: Path) -> Path: + return output_dir / "trainer_steps" + + def get_step_path(path: Path, step: int) -> Path: return path / f"step_{step}" @@ -157,7 +161,7 @@ def validate_output_dir(output_dir: Path, *, resuming: bool, clean: bool, ckpt_o def clean_future_steps(output_dir: Path, resume_step: int) -> None: - """Remove stale rollouts, broadcasts, and traces past ``resume_step``. + """Remove stale rollouts, broadcasts, trainer markers, and traces past ``resume_step``. Pass ``resume_step=-1`` to wipe every step directory (fresh runs). """ @@ -166,6 +170,7 @@ def clean_future_steps(output_dir: Path, resume_step: int) -> None: get_rollout_dir(output_dir), get_rollout_dir(run_default), get_broadcast_dir(run_default), + get_trainer_step_dir(run_default), ] for directory in dirs: diff --git a/tests/unit/train/test_runs.py b/tests/unit/train/test_runs.py index f650541be9..e921934f77 100644 --- a/tests/unit/train/test_runs.py +++ b/tests/unit/train/test_runs.py @@ -1,11 +1,12 @@ from pathlib import Path +from types import SimpleNamespace from typing import Generator import pytest import tomli_w import torch.distributed as dist -from prime_rl.trainer.runs import MultiRunManager +from prime_rl.trainer.runs import MultiRunManager, Progress @pytest.fixture(autouse=True, scope="module") @@ -64,6 +65,29 @@ def test_initial_state(tmp_path: Path) -> None: assert multi_run_manager.run_dirs() == [] +def test_step_completion_is_independent_from_weight_broadcast(tmp_path: Path) -> None: + multi_run_manager = MultiRunManager.__new__(MultiRunManager) + multi_run_manager.output_dir = tmp_path + multi_run_manager.world = SimpleNamespace(is_master=True) + multi_run_manager.idx_2_id = {0: "run_static", 1: "run_policy"} + multi_run_manager.config = { + 0: SimpleNamespace(needs_inference=False), + 1: SimpleNamespace(needs_inference=True), + } + multi_run_manager.progress = {0: Progress(step=4), 1: Progress(step=7)} + multi_run_manager.ready_to_update = [True, True] + for run_id in multi_run_manager.idx_2_id.values(): + (tmp_path / run_id).mkdir() + + assert multi_run_manager.ready_to_broadcast_idxs == [1] + + multi_run_manager.mark_ready_steps_complete() + + assert (tmp_path / "run_static" / "trainer_steps" / "step_3" / "STABLE").exists() + assert (tmp_path / "run_policy" / "trainer_steps" / "step_6" / "STABLE").exists() + assert multi_run_manager.ready_to_update == [False, False] + + def test_detect_new_runs(tmp_path: Path) -> None: """Test that new runs are detected correctly.""" multi_run_manager = MultiRunManager(output_dir=tmp_path, max_runs=5)