Buckets:
| """Persistent per-GPU rollout workers using one current-policy shared queue.""" | |
| from __future__ import annotations | |
| import copy | |
| import multiprocessing as mp | |
| from pathlib import Path | |
| import time | |
| import traceback | |
| import uuid | |
| import torch | |
| import torch.distributed as dist | |
| from shared.artifacts import rng_state, restore_rng_state, write_json | |
| from shared.data import PackedTrajectoryCache, answer_reference | |
| from shared.modeling import load_model, load_processor | |
| from .contracts import VerifierInput | |
| from .rollout_artifacts import save_rollout | |
| from .verifier import build_verifier, evaluate_verifier | |
| from .worker_control import WorkerControl, WorkerStageError, initialize_worker_progress | |
| def claim_indices(store, indices): | |
| while True: | |
| position = store.add("next", 1) - 1 | |
| if position >= len(indices): | |
| return | |
| yield indices[position] | |
| def drain(model, processor, cache, verifier, config, task, worker): | |
| from .trainer import _build_rollout_runtime, _rollout_sampling_seed | |
| controller = _build_rollout_runtime(model, processor, config) | |
| controller.backend.cache_autocast = task.get("cache_autocast", config.rollout.cache_autocast) | |
| store = dist.FileStore(task["queue_path"], -1) | |
| control = WorkerControl(task, worker=worker, store=store) | |
| controller.backend.progress_callback = control.progress | |
| directory = Path(task["directory"]) | |
| model.eval() | |
| started = time.monotonic() | |
| tokens = count = 0 | |
| torch.cuda.reset_peak_memory_stats() | |
| for index in claim_indices(store, task["indices"]): | |
| control.progress(phase="load_sample", sample_index=index, generated_tokens=0, sequence_tokens=0, force=True) | |
| item = cache.load(index) | |
| rollout = controller.collect(index, item, sampling_seed=_rollout_sampling_seed( | |
| config, phase=task["phase"], update=task["update"], sample_id=cache.sample_id(index), | |
| )) | |
| bundle = directory / "samples" / f"{index:08d}" | |
| control.progress(phase="save_rollout", force=True) | |
| save_rollout(bundle, rollout, processor.tokenizer) | |
| row = { | |
| "index": index, "sample_id": rollout.sample_id, | |
| "protocol_valid": rollout.protocol_valid, "stop_reason": rollout.stop_reason, | |
| "task": str(item.get("config") or "unknown"), "worker": worker, | |
| "trajectory": str(bundle.relative_to(directory) / "trajectory.pt"), | |
| } | |
| if task["phase"] == "validation": | |
| row["correct"] = bool(rollout.protocol_valid and rollout.has_scored_window and evaluate_verifier( | |
| verifier, VerifierInput.from_rollout(rollout, reference=answer_reference(item)), | |
| timeout_seconds=config.verifier.timeout_seconds, | |
| )) | |
| write_json(directory / "samples" / f"{index:08d}.json", row) | |
| tokens += sum(owner == "student" for owner in rollout.token_owners) | |
| count += 1 | |
| print(f"OPD_{task['phase'].upper()}_PROGRESS worker={worker} completed={count} tokens={tokens} seconds={time.monotonic()-started:.1f}", flush=True) | |
| torch.cuda.synchronize() | |
| control.progress(phase="done", force=True) | |
| return {"samples": count, "tokens": tokens, "seconds": time.monotonic()-started, | |
| "peak_gpu_memory_gib": torch.cuda.max_memory_allocated()/1024**3} | |
| def _child(config, device, shared, connection): | |
| task = None | |
| try: | |
| torch.set_num_threads(2) | |
| torch.cuda.set_device(device) | |
| torch.backends.cuda.matmul.allow_tf32 = config.trainer.tf32 | |
| child_config = copy.deepcopy(config) | |
| child_config.checkpoint.resume_from = None | |
| child_config.model.gradient_checkpointing = False | |
| processor = load_processor(config.model.path) | |
| # Match parent PEFT dtypes and requires_grad flags, so autocast can cache | |
| # FP32 master-weight casts even though the rollout itself is no-grad. | |
| model = load_model(child_config, torch.device("cuda", device), for_training=True) | |
| parameters = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} | |
| if set(parameters) != set(shared): | |
| raise ValueError("rollout worker trainable parameter schema mismatch") | |
| cache = PackedTrajectoryCache(config.data.cache_root) | |
| verifier = build_verifier(config.verifier, processor.tokenizer) | |
| connection.send({"ready": True}) | |
| while True: | |
| task = connection.recv() | |
| if task is None: | |
| break | |
| control = WorkerControl(task, worker=task["worker"]) | |
| control.progress(phase="load_weights", force=True) | |
| with torch.no_grad(): | |
| for name, source in shared.items(): | |
| parameters[name].copy_(source) | |
| result = drain(model, processor, cache, verifier, config, task, task["worker"]) | |
| # Other CUDA processes cannot reuse this allocator's idle blocks. | |
| # Release rollout KV/workspace reservations before parent backward. | |
| control.progress(phase="release_workspace", force=True) | |
| torch.cuda.empty_cache() | |
| control.progress(phase="done", force=True) | |
| connection.send({"result": result}) | |
| except Exception: | |
| error = traceback.format_exc() | |
| if task is not None: | |
| WorkerControl(task, worker=task["worker"]).fail(error) | |
| connection.send({"error": error}) | |
| finally: | |
| connection.close() | |
| class RolloutWorkers: | |
| """Reuse models across updates; refresh weights before each frozen batch.""" | |
| def __init__(self, model, processor, cache, verifier, config, parallel): | |
| self.model, self.processor, self.cache = model, processor, cache | |
| self.verifier, self.config, self.parallel = verifier, config, parallel | |
| self.parameters = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} | |
| self.shared, self.children = {}, [] | |
| count = max(config.rollout.workers_per_gpu, config.evaluation.workers_per_gpu) | |
| if count > 1: | |
| self.shared = {name: parameter.detach().cpu().clone().share_memory_() for name, parameter in self.parameters.items()} | |
| for _ in range(count - 1): | |
| parent, child = mp.get_context("spawn").Pipe() | |
| process = mp.get_context("spawn").Process(target=_child, args=(config, parallel.local_rank, self.shared, child), daemon=True) | |
| process.start() | |
| child.close() | |
| self.children.append((process, parent)) | |
| for process, connection in self.children: | |
| self._receive(process, connection) | |
| def _receive(self, process, connection, *, control=None, worker=None): | |
| deadline = time.monotonic() + self.config.trainer.process_group_timeout_minutes * 60 | |
| while not connection.poll(1): | |
| if control is not None: | |
| control.check_health(waiting_worker=worker) | |
| if not process.is_alive() or (control is None and time.monotonic() >= deadline): | |
| raise RuntimeError(f"rollout worker {process.pid} exited or timed out") | |
| result = connection.recv() | |
| if "error" in result: | |
| raise RuntimeError(result["error"]) | |
| return result | |
| def run(self, indices, directory, *, phase, update, workers_per_gpu): | |
| parallel = self.parallel | |
| started = time.monotonic() | |
| saved_rng = rng_state() | |
| modes = [(module, module.training) for module in self.model.modules()] | |
| setup, store = None, None | |
| if parallel.is_main: | |
| directory = Path(directory) | |
| directory.mkdir(parents=True, exist_ok=True) | |
| (directory / "samples").mkdir(exist_ok=True) | |
| queue_path = str(directory / f".queue-{uuid.uuid4().hex}") | |
| store = dist.FileStore(queue_path, -1) | |
| store.set("next", "0") | |
| worker_ids = [f"{rank}:{slot}" for rank in range(parallel.world_size) for slot in range(workers_per_gpu)] | |
| initialize_worker_progress(store, directory, worker_ids) | |
| setup = {"directory": str(directory), "queue_path": queue_path, "indices": tuple(indices), "phase": phase, "update": update, "cache_autocast": self.config.rollout.cache_autocast, "worker_ids": worker_ids, "idle_timeout_seconds": self.config.trainer.process_group_timeout_minutes * 60} | |
| task = parallel.broadcast_object(setup) | |
| control = WorkerControl(task, worker=f"{parallel.rank}:0", store=store) | |
| active = self.children[:workers_per_gpu - 1] | |
| results, error = [], None | |
| try: | |
| if indices: | |
| with torch.no_grad(): | |
| for name, buffer in self.shared.items(): | |
| buffer.copy_(self.parameters[name].detach().cpu()) | |
| torch.cuda.empty_cache() | |
| for slot, (_process, connection) in enumerate(active, 1): | |
| connection.send(task | {"worker": f"{parallel.rank}:{slot}"}) | |
| results.append(drain(self.model, self.processor, self.cache, self.verifier, self.config, task, f"{parallel.rank}:0")) | |
| for slot, (process, connection) in enumerate(active, 1): | |
| results.append(self._receive(process, connection, control=control, worker=f"{parallel.rank}:{slot}")["result"]) | |
| except Exception: | |
| error = traceback.format_exc() | |
| control.fail(error) | |
| raise WorkerStageError(error) | |
| finally: | |
| restore_rng_state(saved_rng) | |
| for module, training in modes: | |
| module.training = training | |
| # Long but healthy generation does not spend the Gloo collective timeout. | |
| # A failure leaves its original traceback in the shared progress files. | |
| control.wait_for_ranks(parallel.rank, parallel.world_size) | |
| gathered = parallel.gather_stage(results, None, stage=f"{phase} worker queue") | |
| rows = [row for group in gathered for row in group] | |
| seconds = max(parallel.all_gather_object(time.monotonic() - started)) | |
| del store | |
| return { | |
| "seconds": seconds, | |
| "tokens": sum(row["tokens"] for row in rows), | |
| "tokens_per_second": sum(row["tokens"] for row in rows)/max(seconds, 1e-9), | |
| "samples_per_second": len(indices)/max(seconds, 1e-9), | |
| "worker_peak_memory_gib_sum": sum(row["peak_gpu_memory_gib"] for row in rows), | |
| "peak_worker_memory_gib_per_gpu": max((sum(row["peak_gpu_memory_gib"] for row in group) for group in gathered), default=0.), | |
| "workers": workers_per_gpu * parallel.world_size, | |
| } | |
| def close(self): | |
| for process, connection in self.children: | |
| if process.is_alive(): | |
| try: | |
| connection.send(None) | |
| except (BrokenPipeError, OSError): | |
| pass | |
| process.join(5) | |
| if process.is_alive(): | |
| process.terminate() | |
| process.join(5) | |
| connection.close() | |
Xet Storage Details
- Size:
- 11.1 kB
- Xet hash:
- ad1148d36577f7f95c19f77dfccf56d4dc50e618ce4ea9d6a41653449664ef88
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.