From 7ce9f33b5a1d03cf8ae1a9b41f7fe97b61e2f0b7 Mon Sep 17 00:00:00 2001 From: Haizhong Date: Tue, 14 Jul 2026 11:07:08 -0400 Subject: [PATCH 1/3] feat(dataflow): earliest-deadline-first rollout buffer consumption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FIFO (completion-order) consumption has a difficulty bias: long generations span several weight versions, so their min_version is old on arrival, and queueing behind fresher short samples pushes them past max_staleness — the buffer preferentially expires exactly the long/hard prompts. Measured under rollout over-provisioning (6 inference GPUs vs 2 trainer GPUs, Qwen3-4B math RL): dropped samples are 1.3-1.7x longer than consumed ones and 17% of production expires. Replace the fresh-queue deque with a priority heap ordered by dataflow.buffer.queue_order: - "edf" (new default): ascending min_version — consume the sample closest to its staleness deadline first. The per-token invariant is unchanged (current_version - min_version <= max_staleness for everything trained on); ties break by arrival, so rollout groups stay contiguous. - "fifo": arrival order, bit-for-bit the historical behavior, kept for comparison runs. Plumb the option through AgentConfig -> AstraFlow -> data serving -> RolloutBuffer (per-model override supported), and surface the bias in metrics: buffer/consumed_len_avg, buffer/skipped_stale_len_avg, plus the previously computed-but-unreported eviction count. state_dict stays backward compatible with deque-era checkpoints (priorities are recomputed from metadata on load, so checkpoints also survive switching queue_order). A/B on 8xH100 (only queue_order differs): edf halves the drop rate, drops length-fairly (1.1x vs 1.7x), trains on ~25% longer samples, and reaches 57.6 overall avg@k at step 100 vs fifo's 52.5 (+5.1, reproduced across two independent edf runs; largest gain on the hardest set, AIME25 +9.4). Both arms converge by step 350 — edf is a sample-efficiency win, reaching fifo's plateau ~150 steps earlier. Covered by 14 new unit tests including a deterministic mixed-length load simulation asserting the bias exists under fifo and disappears under edf with no staleness violations in either mode. --- astraflow/core/config/loader.py | 1 + astraflow/dataflow/astraflow.py | 2 + astraflow/dataflow/data_serving.py | 6 +- astraflow/dataflow/rollout_buffer.py | 250 +++++++++++++----- astraflow/dataflow/service.py | 14 + astraflow/dataflow/service_config.py | 7 + .../tests/test_rollout_buffer_queue_order.py | 249 +++++++++++++++++ 7 files changed, 467 insertions(+), 62 deletions(-) create mode 100644 astraflow/dataflow/tests/test_rollout_buffer_queue_order.py diff --git a/astraflow/core/config/loader.py b/astraflow/core/config/loader.py index 6dd4c9d..7d62a57 100644 --- a/astraflow/core/config/loader.py +++ b/astraflow/core/config/loader.py @@ -294,6 +294,7 @@ def load_dataflow_config(raw: dict) -> dict: _set_if_missing(agent_fields, "replay_size", buffer.get("replay_size")) _set_if_missing(agent_fields, "replay_ratio", buffer.get("replay_ratio")) _set_if_missing(agent_fields, "max_staleness", buffer.get("max_staleness")) + _set_if_missing(agent_fields, "queue_order", buffer.get("queue_order")) _set_if_missing(agent_fields, "filter_function", buffer.get("filter_function")) # Extract service-level config from agent_fields (they're in the diff --git a/astraflow/dataflow/astraflow.py b/astraflow/dataflow/astraflow.py index b554bf3..d6f5aba 100644 --- a/astraflow/dataflow/astraflow.py +++ b/astraflow/dataflow/astraflow.py @@ -48,6 +48,7 @@ def __init__( max_staleness: int | None = None, replay_max_size: int | None = None, replay_selection_fn: ReplaySelectionFn | str | None = None, + queue_order: str = "edf", producer_error_backoff: float = 0.5, data_serving: DataServing | None = None, data_acquisition: DataAcquisition | None = None, @@ -71,6 +72,7 @@ def __init__( max_staleness=max_staleness, replay_max_size=replay_max_size, replay_selection_fn=replay_selection_fn, + queue_order=queue_order, per_model_config=per_model_buffer_config, ) self.data_serving = data_serving diff --git a/astraflow/dataflow/data_serving.py b/astraflow/dataflow/data_serving.py index f368897..2f6e9e0 100644 --- a/astraflow/dataflow/data_serving.py +++ b/astraflow/dataflow/data_serving.py @@ -128,6 +128,7 @@ def __init__( max_staleness: int | None = None, replay_max_size: int | None = None, replay_selection_fn: ReplaySelectionFn | str | None = None, + queue_order: str = "edf", ): resolved_replay_selection_fn = self._resolve_replay_selection_fn( replay_selection_fn @@ -139,6 +140,7 @@ def __init__( max_staleness=max_staleness, replay_max_size=replay_max_size, replay_selection_fn=resolved_replay_selection_fn, + queue_order=queue_order, ) self._acq_stats_lock = threading.Lock() self._acq_stats: dict[str, float | int] = _empty_acquisition_stats() @@ -392,7 +394,7 @@ class MultiModelDataServing: per_model_config : dict[str, dict] | None Optional per-model overrides. Keys are model_ids, values are dicts with optional keys ``buffer_size``, ``max_staleness``, - ``replay_max_size``. + ``replay_max_size``, ``queue_order``. """ def __init__( @@ -405,6 +407,7 @@ def __init__( max_staleness: int | None = None, replay_max_size: int | None = None, replay_selection_fn: ReplaySelectionFn | str | None = None, + queue_order: str = "edf", per_model_config: dict[str, dict] | None = None, ): if model_ids is not None and len(model_ids) > 0: @@ -424,6 +427,7 @@ def __init__( max_staleness=cfg.get("max_staleness", max_staleness), replay_max_size=cfg.get("replay_max_size", replay_max_size), replay_selection_fn=replay_selection_fn, + queue_order=cfg.get("queue_order", queue_order), ) buf.buffer.label = f"[{mid}] " if self._is_multi_model else "" self.buffers[mid] = buf diff --git a/astraflow/dataflow/rollout_buffer.py b/astraflow/dataflow/rollout_buffer.py index 271b56b..0f3dd57 100644 --- a/astraflow/dataflow/rollout_buffer.py +++ b/astraflow/dataflow/rollout_buffer.py @@ -7,6 +7,7 @@ from __future__ import annotations import copy +import heapq import logging import threading from collections import deque @@ -44,8 +45,34 @@ def _slice_tensor_dict(data: dict[str, Any], start: int, end: int) -> dict[str, return sliced_data +#: Valid fresh-queue consumption orders. +QUEUE_ORDERS = ("fifo", "edf") + + class RolloutBuffer: - """Thread-safe storage for fresh and replay rollout examples.""" + """Thread-safe storage for fresh and replay rollout examples. + + The fresh queue is a priority heap whose consumption order is set by + ``queue_order``: + + - ``"edf"`` (default, earliest deadline first): ascending + ``min_version``, i.e. + the sample closest to its staleness deadline + (``min_version + max_staleness``) is consumed first. + - ``"fifo"``: arrival order — the historical deque behavior; set + explicitly for comparison runs. + + Why edf is the default: long generations + enter the buffer with the least remaining staleness budget; FIFO makes + them wait behind fresher short samples and expire, biasing training + data toward short/easy prompts. EDF consumes them first instead. The + per-token staleness invariant is unchanged: any consumed sample still + satisfies ``current_version - min_version <= max_staleness``. + + Ties (same priority) are broken by arrival order, so samples of one + rollout group — ingested contiguously with the same ``min_version`` — + stay contiguous under both orders. + """ def __init__( self, @@ -56,6 +83,7 @@ def __init__( max_staleness: int | None = None, replay_max_size: int | None = None, replay_selection_fn: ReplaySelectionFn | None = None, + queue_order: str = "edf", ): # Filtering has been moved to data-acquisition layer. Keep this argument # for backward compatibility with older call sites. @@ -67,8 +95,20 @@ def __init__( self.debug = debug self.label = "" - self._buffer: deque[dict[str, Any]] = deque(maxlen=max_size) - self._metadata: deque[dict[str, Any]] = deque(maxlen=max_size) + if queue_order not in QUEUE_ORDERS: + logger.warning( + "Unknown queue_order %r; falling back to 'edf'. Valid: %s", + queue_order, + QUEUE_ORDERS, + ) + queue_order = "edf" + self.queue_order = queue_order + + # Fresh queue: heap of (priority, seq_no, example, metadata). + # priority is 0.0 for fifo (order degenerates to arrival seq_no) and + # min_version for edf. seq_no is a monotonic arrival counter. + self._heap: list[tuple[float, int, dict[str, Any], dict[str, Any]]] = [] + self._seq_counter = 0 self._replay_buffer: deque[dict[str, Any]] = deque(maxlen=replay_max_size) self._replay_metadata: deque[dict[str, Any]] = deque(maxlen=replay_max_size) @@ -82,8 +122,36 @@ def __init__( self.replay_selection_fn = ( replay_selection_fn if replay_selection_fn is not None else select_latest ) - self._put_stats = {"accepted": 0, "filtered": 0, "total": 0} - self._consume_stats = {"consumed": 0, "skipped_stale": 0} + self._put_stats = {"accepted": 0, "filtered": 0, "total": 0, "evicted": 0} + self._consume_stats = { + "consumed": 0, + "skipped_stale": 0, + "consumed_len_sum": 0, + "skipped_stale_len_sum": 0, + } + + def _priority(self, metadata: dict[str, Any]) -> float: + """Heap priority for a fresh example (lower = consumed earlier).""" + if self.queue_order == "edf": + min_v = metadata.get("min_version") + if isinstance(min_v, (int, float)): + return float(min_v) + # No version info: the sample can never be staleness-dropped, so + # consume it eagerly rather than risk starving it behind an + # endless stream of versioned samples. + return float("-inf") + return 0.0 + + @staticmethod + def _seq_len(example: dict[str, Any]) -> int: + """Token length of a single-example dict (0 if unknown).""" + mask = example.get("attention_mask") + try: + if torch.is_tensor(mask): + return int(mask.sum().item()) + except Exception: + pass + return 0 def _normalize_rewards(self, rewards: torch.Tensor) -> torch.Tensor: reward_score = rewards @@ -123,18 +191,25 @@ def put( inserted_count = 0 evicted_count = 0 + priority = self._priority(example_metadata) for i in range(batch_size): example = _slice_tensor_dict(batch, i, i + 1) - if len(self._buffer) >= self.max_size: - self._buffer.popleft() - self._metadata.popleft() + if len(self._heap) >= self.max_size: + # Evict the head: under fifo the oldest arrival (historic + # behavior), under edf the most staleness-critical sample + # (the one that would expire soonest anyway). + heapq.heappop(self._heap) evicted_count += 1 - self._buffer.append(example) - self._metadata.append(example_metadata) + heapq.heappush( + self._heap, + (priority, self._seq_counter, example, example_metadata), + ) + self._seq_counter += 1 inserted_count += 1 self._put_stats["accepted"] += inserted_count self._put_stats["total"] += batch_size + self._put_stats["evicted"] += evicted_count if inserted_count > 0: self._not_empty.notify() @@ -143,13 +218,18 @@ def put( def get_and_reset_put_stats(self) -> dict[str, int]: with self._lock: stats = dict(self._put_stats) - self._put_stats = {"accepted": 0, "filtered": 0, "total": 0} + self._put_stats = {"accepted": 0, "filtered": 0, "total": 0, "evicted": 0} return stats def get_and_reset_consume_stats(self) -> dict[str, int]: with self._lock: stats = dict(self._consume_stats) - self._consume_stats = {"consumed": 0, "skipped_stale": 0} + self._consume_stats = { + "consumed": 0, + "skipped_stale": 0, + "consumed_len_sum": 0, + "skipped_stale_len_sum": 0, + } return stats def get( @@ -165,11 +245,17 @@ def get( batch, _ = result return batch - def get_with_metadata( + def _pop_one( self, timeout: float | None = None, current_version: int | None = None, - ) -> tuple[dict[str, Any], dict[str, Any]] | None: + ) -> tuple[float, int, dict[str, Any], dict[str, Any]] | None: + """Pop the highest-priority fresh entry, dropping expired ones. + + Returns the full heap entry ``(priority, seq_no, example, metadata)`` + so callers can push it back verbatim. Blocks while the queue is + empty; returns None only when the buffer is closed. + """ del timeout import time as time_module @@ -180,22 +266,24 @@ def get_with_metadata( print("RolloutBuffer.get: Buffer is closed, returning None") return None - if len(self._buffer) == 0: + if len(self._heap) == 0: if self.debug: print("RolloutBuffer.get: Buffer is empty, waiting...") self._not_empty.wait() time_module.sleep(0.1) continue - metadata = self._metadata[0] + _, _, example, metadata = self._heap[0] if current_version is not None and self.max_staleness is not None: min_v = metadata.get("min_version") if min_v is not None and isinstance(min_v, (int, float)): version_gap = current_version - int(min_v) if version_gap > self.max_staleness: - self._buffer.popleft() - self._metadata.popleft() + heapq.heappop(self._heap) self._consume_stats["skipped_stale"] += 1 + self._consume_stats["skipped_stale_len_sum"] += ( + self._seq_len(example) + ) if self.debug: print( f"{self.label}RolloutBuffer: skipped stale example " @@ -205,11 +293,22 @@ def get_with_metadata( ) continue - example = self._buffer.popleft() - metadata = self._metadata.popleft() + entry = heapq.heappop(self._heap) self._consume_stats["consumed"] += 1 + self._consume_stats["consumed_len_sum"] += self._seq_len(entry[2]) self._not_full.notify() - return example, metadata + return entry + + def get_with_metadata( + self, + timeout: float | None = None, + current_version: int | None = None, + ) -> tuple[dict[str, Any], dict[str, Any]] | None: + entry = self._pop_one(timeout=timeout, current_version=current_version) + if entry is None: + return None + _, _, example, metadata = entry + return example, metadata def get_batch( self, @@ -225,17 +324,17 @@ def get_batch( start_time = time_module.time() if timeout is not None else None if timeout is not None: - first_result = self.get_with_metadata( + first_entry = self._pop_one( timeout=min(timeout, 5.0), current_version=current_version, ) else: - first_result = self.get_with_metadata( + first_entry = self._pop_one( timeout=None, current_version=current_version, ) - if first_result is None: + if first_entry is None: if self.debug: print( "RolloutBuffer.get_batch: No examples available, returning None", @@ -243,12 +342,10 @@ def get_batch( ) return None - first_example, first_metadata = first_result - examples = [first_example] - metadatas = [first_metadata] + entries = [first_entry] if self.debug: print( - f"{self.label}RolloutBuffer.get_batch: Collected {len(examples)}/{batch_size} fresh examples", + f"{self.label}RolloutBuffer.get_batch: Collected {len(entries)}/{batch_size} fresh examples", flush=True, ) @@ -263,43 +360,45 @@ def get_batch( if remaining_timeout <= 0: break - result = self.get_with_metadata( + entry = self._pop_one( timeout=remaining_timeout, current_version=current_version, ) - if result is None: + if entry is None: if self.debug: print( - "RolloutBuffer.get_batch: get_with_metadata returned None while collecting", + "RolloutBuffer.get_batch: _pop_one returned None while collecting", flush=True, ) break - example, metadata = result - examples.append(example) - metadatas.append(metadata) - if self.debug and (len(examples) == batch_size or len(examples) % 10 == 0): + entries.append(entry) + if self.debug and (len(entries) == batch_size or len(entries) % 10 == 0): print( - f"{self.label}RolloutBuffer.get_batch: Collected {len(examples)}/{batch_size} fresh examples", + f"{self.label}RolloutBuffer.get_batch: Collected {len(entries)}/{batch_size} fresh examples", flush=True, ) - if len(examples) < batch_size and len(examples) % 8 == 0: + if len(entries) < batch_size and len(entries) % 8 == 0: time_module.sleep(0.05) - if len(examples) < batch_size: + if len(entries) < batch_size: if self.debug: print( - f"RolloutBuffer.get_batch: Only collected {len(examples)}/{batch_size}, putting back", + f"RolloutBuffer.get_batch: Only collected {len(entries)}/{batch_size}, putting back", flush=True, ) with self._lock: - for example, metadata in zip(reversed(examples), reversed(metadatas)): - self._buffer.appendleft(example) - self._metadata.appendleft(metadata) + # Push entries back verbatim: original (priority, seq_no) keys + # restore the exact consumption order under both queue modes. + for entry in entries: + heapq.heappush(self._heap, entry) self._not_empty.notify_all() return None + examples = [entry[2] for entry in entries] + metadatas = [entry[3] for entry in entries] + with self._lock: for example, metadata in zip(examples, metadatas): self._replay_buffer.append(example) @@ -481,15 +580,23 @@ def _aggregate_metadata(self, metadatas: list[dict[str, Any]]) -> dict[str, Any] def size(self) -> int: with self._lock: - return len(self._buffer) + return len(self._heap) def state_dict(self) -> dict[str, Any]: with self._lock: + # Persist fresh entries in consumption order. Priorities are NOT + # persisted — they are recomputed from metadata on load, so a + # checkpoint saved under one queue_order loads correctly under + # another. + ordered = sorted(self._heap, key=lambda e: (e[0], e[1])) return { "max_size": self.max_size, "replay_max_size": self.replay_max_size, - "buffer": [self._clone_for_state(x) for x in self._buffer], - "metadata": [self._clone_for_state(x) for x in self._metadata], + "queue_order": self.queue_order, + "buffer": [self._clone_for_state(e[2]) for e in ordered], + "metadata": [self._clone_for_state(e[3]) for e in ordered], + "seq_nos": [e[1] for e in ordered], + "seq_counter": self._seq_counter, "replay_buffer": [ self._clone_for_state(x) for x in self._replay_buffer ], @@ -507,14 +614,27 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None: metadata_items = state_dict.get("metadata", []) replay_items = state_dict.get("replay_buffer", []) replay_metadata_items = state_dict.get("replay_metadata", []) - - self._buffer = deque( - [self._clone_for_state(x) for x in buffer_items], - maxlen=self.max_size, - ) - self._metadata = deque( - [self._clone_for_state(x) for x in metadata_items], - maxlen=self.max_size, + # Old checkpoints (deque era) have no seq_nos: assign arrival + # counters in list order, which preserves their FIFO order. + seq_nos = state_dict.get("seq_nos") + if not seq_nos or len(seq_nos) != len(buffer_items): + seq_nos = list(range(len(buffer_items))) + + self._heap = [] + for seq_no, example, metadata in zip( + seq_nos, buffer_items, metadata_items + ): + example = self._clone_for_state(example) + metadata = self._clone_for_state(metadata) + heapq.heappush( + self._heap, + (self._priority(metadata), int(seq_no), example, metadata), + ) + self._seq_counter = int( + state_dict.get( + "seq_counter", + (max(seq_nos) + 1) if seq_nos else 0, + ) ) self._replay_buffer = deque( [self._clone_for_state(x) for x in replay_items], @@ -525,16 +645,24 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None: maxlen=self.replay_max_size, ) self._closed = bool(state_dict.get("closed", False)) - self._put_stats = dict( - state_dict.get("put_stats", {"accepted": 0, "filtered": 0, "total": 0}) - ) - self._consume_stats = dict( - state_dict.get("consume_stats", {"consumed": 0, "skipped_stale": 0}) - ) + self._put_stats = { + "accepted": 0, + "filtered": 0, + "total": 0, + "evicted": 0, + **dict(state_dict.get("put_stats", {})), + } + self._consume_stats = { + "consumed": 0, + "skipped_stale": 0, + "consumed_len_sum": 0, + "skipped_stale_len_sum": 0, + **dict(state_dict.get("consume_stats", {})), + } - if len(self._buffer) > 0: + if len(self._heap) > 0: self._not_empty.notify_all() - if len(self._buffer) < self.max_size: + if len(self._heap) < self.max_size: self._not_full.notify_all() def close(self) -> None: diff --git a/astraflow/dataflow/service.py b/astraflow/dataflow/service.py index e2bf184..ec93c4d 100644 --- a/astraflow/dataflow/service.py +++ b/astraflow/dataflow/service.py @@ -176,6 +176,7 @@ def register_agent( buffer_debug=True, filter_fn=agent_config.filter_function, max_staleness=agent_config.max_staleness, + queue_order=agent_config.queue_order, replay_max_size=agent_config.replay_size, replay_selection_fn=agent_config.selection_fn, reward_norm=reward_norm, @@ -663,6 +664,19 @@ def _pull_batch( "buffer/consumed": float(consume.get("consumed", 0)), "buffer/skipped_stale": float(consume.get("skipped_stale", 0)), } + # Length breakdown of consumed vs staleness-dropped samples — the + # direct signal for the length/difficulty bias that queue_order=edf + # addresses (long generations expiring more often under fifo). + consumed_n = int(consume.get("consumed", 0)) + if consumed_n > 0: + buffer_stats["buffer/consumed_len_avg"] = ( + float(consume.get("consumed_len_sum", 0)) / consumed_n + ) + skipped_n = int(consume.get("skipped_stale", 0)) + if skipped_n > 0: + buffer_stats["buffer/skipped_stale_len_avg"] = ( + float(consume.get("skipped_stale_len_sum", 0)) / skipped_n + ) # Only include reward means when we have data — avoids logging # misleading 0.0 on the first step after recovery (buffer serves # pre-loaded data before new acquisition stats accumulate). diff --git a/astraflow/dataflow/service_config.py b/astraflow/dataflow/service_config.py index d44969b..4869f9e 100644 --- a/astraflow/dataflow/service_config.py +++ b/astraflow/dataflow/service_config.py @@ -22,6 +22,13 @@ class AgentConfig: max_staleness: int | None = None """Maximum version staleness for samples; older samples are discarded.""" + queue_order: str = "edf" + """Fresh-queue consumption order: "edf" (default; earliest staleness + deadline first, i.e. ascending min_version) or "fifo" (arrival order, + the historical behavior — set explicitly for comparison runs). EDF + removes the difficulty bias where long generations expire in the queue + more often than short ones. See RolloutBuffer.""" + replay_size: int | None = None """Maximum number of samples in the replay buffer.""" diff --git a/astraflow/dataflow/tests/test_rollout_buffer_queue_order.py b/astraflow/dataflow/tests/test_rollout_buffer_queue_order.py new file mode 100644 index 0000000..243009d --- /dev/null +++ b/astraflow/dataflow/tests/test_rollout_buffer_queue_order.py @@ -0,0 +1,249 @@ +"""Tests for RolloutBuffer queue ordering (fifo vs edf) and staleness bias. + +The edf order exists to remove a difficulty bias: long generations span more +weight versions, so their ``min_version`` is older on arrival and FIFO +(completion-order) consumption lets them expire in the queue far more often +than short generations. EDF consumes the most staleness-critical samples +first while preserving the exact ``max_staleness`` invariant. +""" + +from __future__ import annotations + +import torch + +from astraflow.dataflow.rollout_buffer import RolloutBuffer + + +def _example(length: int, tag: int = 0) -> dict: + return { + "attention_mask": torch.ones(1, length, dtype=torch.long), + "input_ids": torch.full((1, length), tag, dtype=torch.long), + "tag": torch.tensor([tag], dtype=torch.long), + } + + +def _put(buf: RolloutBuffer, length: int, min_version: int, tag: int = 0) -> None: + buf.put(_example(length, tag=tag), metadata={"min_version": min_version}) + + +class TestFifoParity: + """queue_order='fifo' must reproduce the historical deque behavior.""" + + def test_arrival_order(self): + buf = RolloutBuffer(max_size=16, queue_order="fifo") + for tag in range(5): + _put(buf, length=4, min_version=tag, tag=tag) + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(5) + ] + assert tags == [0, 1, 2, 3, 4] + + def test_stale_head_dropped(self): + buf = RolloutBuffer(max_size=16, max_staleness=2, queue_order="fifo") + _put(buf, length=4, min_version=0, tag=0) # stale at version 5 + _put(buf, length=4, min_version=5, tag=1) + example, _ = buf.get_with_metadata(current_version=5) + assert int(example["tag"].item()) == 1 + stats = buf.get_and_reset_consume_stats() + assert stats["skipped_stale"] == 1 + assert stats["consumed"] == 1 + + def test_partial_batch_putback_preserves_order(self): + # timeout=0 makes get_batch collect exactly one entry, hit the + # deadline, and put it back (the blocking-on-empty quirk of + # _pop_one predates this change, so a bigger shortfall would hang). + buf = RolloutBuffer(max_size=16, queue_order="fifo") + for tag in range(3): + _put(buf, length=4, min_version=0, tag=tag) + assert buf.get_batch(batch_size=5, timeout=0.0) is None + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(3) + ] + assert tags == [0, 1, 2] + + def test_multi_entry_putback_restores_order(self): + """White-box: pushing popped entries back restores exact order.""" + import heapq as _hq + + buf = RolloutBuffer(max_size=16, queue_order="fifo") + for tag in range(4): + _put(buf, length=4, min_version=0, tag=tag) + entries = [buf._pop_one() for _ in range(3)] + with buf._lock: + for entry in entries: + _hq.heappush(buf._heap, entry) + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(4) + ] + assert tags == [0, 1, 2, 3] + + def test_eviction_drops_oldest_and_is_counted(self): + buf = RolloutBuffer(max_size=2, queue_order="fifo") + for tag in range(3): + _put(buf, length=4, min_version=tag, tag=tag) + stats = buf.get_and_reset_put_stats() + assert stats["evicted"] == 1 + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(2) + ] + assert tags == [1, 2] + + +class TestEdfOrder: + def test_consumes_oldest_min_version_first(self): + buf = RolloutBuffer(max_size=16, queue_order="edf") + _put(buf, length=4, min_version=7, tag=0) + _put(buf, length=4, min_version=3, tag=1) + _put(buf, length=4, min_version=5, tag=2) + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(3) + ] + assert tags == [1, 2, 0] + + def test_ties_keep_arrival_order(self): + """Samples of one group (same min_version) stay contiguous.""" + buf = RolloutBuffer(max_size=16, queue_order="edf") + for tag in range(4): + _put(buf, length=4, min_version=2, tag=tag) + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(4) + ] + assert tags == [0, 1, 2, 3] + + def test_partial_batch_putback_preserves_priority(self): + buf = RolloutBuffer(max_size=16, queue_order="edf") + _put(buf, length=4, min_version=9, tag=0) + _put(buf, length=4, min_version=1, tag=1) + # timeout=0: collects the min_version=1 entry, then puts it back. + assert buf.get_batch(batch_size=5, timeout=0.0) is None + example, _ = buf.get_with_metadata() + assert int(example["tag"].item()) == 1 + + def test_unversioned_samples_consumed_eagerly(self): + buf = RolloutBuffer(max_size=16, queue_order="edf") + _put(buf, length=4, min_version=0, tag=0) + buf.put(_example(4, tag=1), metadata={}) # no min_version + example, _ = buf.get_with_metadata() + assert int(example["tag"].item()) == 1 + + def test_unknown_order_falls_back_to_edf(self): + buf = RolloutBuffer(max_size=16, queue_order="lifo") + assert buf.queue_order == "edf" + + def test_default_is_edf(self): + assert RolloutBuffer(max_size=16).queue_order == "edf" + + +def _run_mixed_load(queue_order: str) -> dict[str, int]: + """Deterministic mixed-length load with advancing weight versions. + + Per version tick v: 4 short samples arrive (generated within v, + min_version=v) and 4 long samples arrive that STARTED 3 versions ago + (min_version=v-3) — modelling generations that span several weight + updates. The consumer drains 6 samples per tick, so a backlog builds and + samples queue for multiple ticks. max_staleness=4 means shorts tolerate + ~4 ticks of queueing while longs tolerate only ~1. + """ + gen_span = 3 + max_staleness = 4 + buf = RolloutBuffer( + max_size=4096, max_staleness=max_staleness, queue_order=queue_order + ) + counts = { + "consumed_long": 0, + "consumed_short": 0, + "dropped_long": 0, + "dropped_short": 0, + "violations": 0, + } + short_len, long_len = 8, 64 + + for v in range(30): + # Fresh shorts arrive first each tick, so a consuming pop always + # finds a valid entry and never blocks on an all-stale queue. + for _ in range(4): + _put(buf, length=short_len, min_version=v) + if v >= gen_span: + for _ in range(4): + _put(buf, length=long_len, min_version=v - gen_span) + + for _ in range(6): + if buf.size() == 0: + break + result = buf.get_with_metadata(current_version=v) + if result is None: + break + example, metadata = result + length = int(example["attention_mask"].sum().item()) + is_long = length == long_len + counts["consumed_long" if is_long else "consumed_short"] += 1 + if v - int(metadata["min_version"]) > max_staleness: + counts["violations"] += 1 + stats = buf.get_and_reset_consume_stats() + dropped = int(stats["skipped_stale"]) + if dropped: + # Exact length attribution: with n dropped totalling len_sum + # tokens, n = a + b and len_sum = long*a + short*b, so + # a = (len_sum - short*n) / (long - short). + dropped_len = int(stats["skipped_stale_len_sum"]) + n_long = (dropped_len - short_len * dropped) // ( + long_len - short_len + ) + counts["dropped_long"] += n_long + counts["dropped_short"] += dropped - n_long + return counts + + +class TestDifficultyBias: + def test_fifo_expires_long_generations_more(self): + fifo = _run_mixed_load("fifo") + edf = _run_mixed_load("edf") + + # Invariant holds in both modes: nothing stale is ever consumed. + assert fifo["violations"] == 0 + assert edf["violations"] == 0 + + # The bias: FIFO drops long generations, EDF rescues them. + assert fifo["dropped_long"] > edf["dropped_long"] + assert edf["consumed_long"] > fifo["consumed_long"] + + # And EDF's rescue does not come from dropping shorts instead: + # shorts have budget to wait, so total drops go down. + total_fifo = fifo["dropped_long"] + fifo["dropped_short"] + total_edf = edf["dropped_long"] + edf["dropped_short"] + assert total_edf <= total_fifo + + +class TestStateDict: + def test_roundtrip_preserves_order(self): + for order in ("fifo", "edf"): + buf = RolloutBuffer(max_size=16, queue_order=order) + _put(buf, length=4, min_version=9, tag=0) + _put(buf, length=4, min_version=1, tag=1) + _put(buf, length=4, min_version=5, tag=2) + state = buf.state_dict() + + restored = RolloutBuffer(max_size=16, queue_order=order) + restored.load_state_dict(state) + tags = [ + int(restored.get_with_metadata()[0]["tag"].item()) + for _ in range(3) + ] + expected = [0, 1, 2] if order == "fifo" else [1, 2, 0] + assert tags == expected, (order, tags) + + def test_legacy_state_without_seq_nos(self): + """Old deque-era checkpoints load in their stored (FIFO) order.""" + buf = RolloutBuffer(max_size=16, queue_order="fifo") + legacy = { + "buffer": [_example(4, tag=0), _example(4, tag=1)], + "metadata": [{"min_version": 3}, {"min_version": 1}], + "replay_buffer": [], + "replay_metadata": [], + "closed": False, + } + buf.load_state_dict(legacy) + tags = [ + int(buf.get_with_metadata()[0]["tag"].item()) for _ in range(2) + ] + assert tags == [0, 1] From a151832d4f618c9cff396b43a6eb47294401bc38 Mon Sep 17 00:00:00 2001 From: Haizhong Date: Tue, 14 Jul 2026 11:07:27 -0400 Subject: [PATCH 2/3] recipes: add qwen3-4b staleness A/B pair (edf vs fifo queue order) Two Qwen3-4B math-RL recipes identical except dataflow.buffer.queue_order (and trial_name), for reproducing the rollout-buffer consumption-order comparison. The 6-RaaS/2-trainer GPU split deliberately over-provisions rollout so the fresh buffer builds real staleness pressure; ctx 16k / max_new_tokens 14000 makes generations span multiple weight versions. READMEs document the setup and the measured results (edf: +5.1 overall avg@k at step 100, +9.4 on AIME25, ~150 steps faster to the eval plateau; converged by step 350). --- examples/math/qwen3-4b-m2po-edf/README.md | 56 ++++++ .../qwen3-4b-m2po-edf/scripts/1_astraflow.sh | 36 ++++ .../math/qwen3-4b-m2po-edf/scripts/2_raas.sh | 44 +++++ .../scripts/3_trainer_model0.sh | 47 +++++ .../scripts/run_qwen3-4b-m2po-edf.sh | 104 +++++++++++ .../qwen3-4b-m2po-edf/yaml/experiment.yaml | 165 ++++++++++++++++++ .../math/qwen3-4b-m2po-edf/yaml/raas.yaml | 37 ++++ examples/math/qwen3-4b-m2po-fifo/README.md | 22 +++ .../qwen3-4b-m2po-fifo/scripts/1_astraflow.sh | 36 ++++ .../math/qwen3-4b-m2po-fifo/scripts/2_raas.sh | 44 +++++ .../scripts/3_trainer_model0.sh | 47 +++++ .../scripts/run_qwen3-4b-m2po-fifo.sh | 104 +++++++++++ .../qwen3-4b-m2po-fifo/yaml/experiment.yaml | 165 ++++++++++++++++++ .../math/qwen3-4b-m2po-fifo/yaml/raas.yaml | 37 ++++ 14 files changed, 944 insertions(+) create mode 100644 examples/math/qwen3-4b-m2po-edf/README.md create mode 100755 examples/math/qwen3-4b-m2po-edf/scripts/1_astraflow.sh create mode 100755 examples/math/qwen3-4b-m2po-edf/scripts/2_raas.sh create mode 100755 examples/math/qwen3-4b-m2po-edf/scripts/3_trainer_model0.sh create mode 100755 examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh create mode 100644 examples/math/qwen3-4b-m2po-edf/yaml/experiment.yaml create mode 100644 examples/math/qwen3-4b-m2po-edf/yaml/raas.yaml create mode 100644 examples/math/qwen3-4b-m2po-fifo/README.md create mode 100755 examples/math/qwen3-4b-m2po-fifo/scripts/1_astraflow.sh create mode 100755 examples/math/qwen3-4b-m2po-fifo/scripts/2_raas.sh create mode 100755 examples/math/qwen3-4b-m2po-fifo/scripts/3_trainer_model0.sh create mode 100755 examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh create mode 100644 examples/math/qwen3-4b-m2po-fifo/yaml/experiment.yaml create mode 100644 examples/math/qwen3-4b-m2po-fifo/yaml/raas.yaml diff --git a/examples/math/qwen3-4b-m2po-edf/README.md b/examples/math/qwen3-4b-m2po-edf/README.md new file mode 100644 index 0000000..13bc067 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/README.md @@ -0,0 +1,56 @@ +# Qwen3-4B Math RL — staleness A/B (queue_order=edf) + +One arm of a controlled A/B on the rollout buffer's **consumption order**. The +sibling recipe [`qwen3-4b-m2po-fifo`](../qwen3-4b-m2po-fifo/) is identical +except `dataflow.buffer.queue_order` (and `trial_name`). + +## Why this experiment exists + +Long generations span several weight versions, so their `min_version` is old +the moment they finish. Consuming in completion order (FIFO — the historical +behavior) parks them behind fresher short samples until they exceed +`max_staleness` and are dropped — a systematic **difficulty bias**: the +hard/long prompts a model most needs to learn are the ones discarded. + +`queue_order: edf` (earliest deadline first, now the default) consumes by +ascending `min_version` instead, training on the most staleness-critical +samples before they expire. The per-token staleness bound is unchanged. + +## Setup + +The GPU split deliberately over-provisions rollout so the buffer builds real +staleness pressure: + +```text +SERVICE_CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 -> RaaS (SGLang, DP=6) +TRAINER_MODEL0_GPUS=6,7 -> Trainer (FSDP, DP=2) +``` + +Model Qwen/Qwen3-4B, M2PO, ctx 16k / max_new_tokens 14000, batch 256, +`max_staleness: 8` — everything except `queue_order` matches the sibling. + +## Run + +```bash +bash examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh +# sibling arm: +bash examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh +``` + +## Results (8×H100, overall avg@k across AIME24/25, AMC, Minerva, MATH500) + +| step | edf | fifo | Δ | +|---|---|---|---| +| 0 | 42.2 | 43.1 | — | +| 100 | **57.6** | 52.5 | **+5.1** | +| 200 | **59.2** | 56.9 | +2.3 | +| 250 | **59.5** | 58.1 | +1.4 | +| 350 | 58.6 | 59.6 | −1.0 (converged) | + +- The @100 gap reproduced across two independent edf runs (57.8 and 57.6). +- Gains are largest on the hardest sets (@100: AIME25 +9.4, AIME24 +5.4). +- Mechanism metrics: fifo's staleness drops are length-biased (dropped + samples 1.3–1.7× longer than consumed); edf's are fair (~1.1×) with a + lower overall drop rate. +- edf is a **sample-efficiency** win (reaches fifo's plateau ~150 steps + earlier); with enough steps both arms converge. diff --git a/examples/math/qwen3-4b-m2po-edf/scripts/1_astraflow.sh b/examples/math/qwen3-4b-m2po-edf/scripts/1_astraflow.sh new file mode 100755 index 0000000..f888c19 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/scripts/1_astraflow.sh @@ -0,0 +1,36 @@ +#!/bin/bash +set -euo pipefail +# [1/3] Launch AstraFlow HTTP service +# +# Usage (terminal 1): +# bash examples/math/qwen3-4b-m2po-edf/scripts/1_astraflow.sh + +export CUDA_VISIBLE_DEVICES="" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== AstraFlow HTTP Service ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "Port : ${ASTRAFLOW_PORT}" +echo "===============================" + +python3 -u -m astraflow \ + --config "${EXPERIMENT_CONFIG}" \ + --port "${ASTRAFLOW_PORT}" \ + --host "${ASTRAFLOW_HOST}" \ + 2>&1 | tee "${LOG_DIR}/astraflow.log" diff --git a/examples/math/qwen3-4b-m2po-edf/scripts/2_raas.sh b/examples/math/qwen3-4b-m2po-edf/scripts/2_raas.sh new file mode 100755 index 0000000..3deb07c --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/scripts/2_raas.sh @@ -0,0 +1,44 @@ +#!/bin/bash +set -euo pipefail +# [2/3] Launch RaaS inference server (SGLang + TCP receiver) +# +# Usage (terminal 2, after AstraFlow is ready): +# bash examples/math/qwen3-4b-m2po-edf/scripts/2_raas.sh + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5}" +export RAAS_HOST="${RAAS_HOST:-0.0.0.0}" +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="${ASTRAFLOW_URL:-http://127.0.0.1:${ASTRAFLOW_PORT}}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== RaaS Inference Server (SGLang + TCP receiver) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "RaaS config : ${RAAS_CONFIG}" +echo "GPUs : ${CUDA_VISIBLE_DEVICES}" +echo "Port : ${RAAS_PORT}" +echo "AstraFlow URL : ${ASTRAFLOW_URL}" +echo "=======================================================" + +python3 -u -m astraflow.raas.server \ + --host "${RAAS_HOST}" \ + --port "${RAAS_PORT}" \ + --config "${EXPERIMENT_CONFIG}" \ + --config "${RAAS_CONFIG}" \ + --engine-id "${ENGINE_ID:-default}" \ + --astraflow-url "${ASTRAFLOW_URL}" \ + 2>&1 | tee "${LOG_DIR}/raas.log" diff --git a/examples/math/qwen3-4b-m2po-edf/scripts/3_trainer_model0.sh b/examples/math/qwen3-4b-m2po-edf/scripts/3_trainer_model0.sh new file mode 100755 index 0000000..17d7b3b --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/scripts/3_trainer_model0.sh @@ -0,0 +1,47 @@ +#!/bin/bash +set -euo pipefail +# [3/3] Launch Trainer for model0 (TCP, sender_agent on local_rank 0) +# +# Usage (terminal 3, after AstraFlow and RaaS are ready): +# bash examples/math/qwen3-4b-m2po-edf/scripts/3_trainer_model0.sh + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS:-6,7}" +TRAINER0_NPROC="$(echo "${CUDA_VISIBLE_DEVICES}" | awk -F',' '{print NF}')" + +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}" +export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}" + +# sender_agent (in trainer) listens on this HTTP port +export WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== Trainer model0 (TCP) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "GPUs : ${CUDA_VISIBLE_DEVICES} (FSDP dp${TRAINER0_NPROC})" +echo "AstraFlow : ${ASTRAFLOW_URL}" +echo "RaaS : ${ASTRAFLOW_RAAS_URL}" +echo "Sender HTTP : ${WEIGHT_TRANSFER_HTTP_PORT}" +echo "WANDB mode : ${WANDB_MODE:-online}" +echo "==========================================" + +torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \ + --master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \ + examples/launch_trainer.py \ + --config "${EXPERIMENT_CONFIG}" \ + --trainer trainer_model0 \ + "$@" 2>&1 | tee "${LOG_DIR}/trainer_model0.log" diff --git a/examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh b/examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh new file mode 100755 index 0000000..7476b3c --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh @@ -0,0 +1,104 @@ +#!/bin/bash +set -euo pipefail +# All-in-one launcher for AstraFlow v2 math training (Qwen3-4B, M2PO, TCP). +# +# Launches 3 processes: +# 1. AstraFlow HTTP service (CPU-only) +# 2. RaaS inference server (SGLang, SERVICE_CUDA_VISIBLE_DEVICES) +# 3. Trainer model0 (math, TRAINER_MODEL0_GPUS) +# +# Usage: +# bash examples/math/qwen3-4b-m2po-edf/scripts/run_qwen3-4b-m2po-edf.sh + +# ============================================================================= +# Part 1: Load env and settings +# ============================================================================= +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +# Defined in examples/_common/utils.sh. +astraflow_load_experiment_env + +# ============================================================================= +# Part 2: Set up env +# ============================================================================= +# GPU assignments (default: 4 GPUs for inference, 4 for training) +export SERVICE_CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5}" +export TRAINER_MODEL0_GPUS="${TRAINER_MODEL0_GPUS:-6,7}" +# Ports / URLs (each component gets its own port) +export RAAS_HOST="${RAAS_HOST:-0.0.0.0}" +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}" +export WEIGHT_TRANSFER_HTTP_PORT_MODEL0="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}" + +TRAINER0_NPROC="$(echo "${TRAINER_MODEL0_GPUS}" | awk -F',' '{print NF}')" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. +# Defined in examples/_common/utils.sh. +astraflow_setup_env + +# ============================================================================= +# Part 3: Print info and clean up +# ============================================================================= +echo "=== AstraFlow v2 (Qwen3-4B, math, M2PO, ctx16k, TCP) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "RaaS config : ${RAAS_CONFIG}" +echo "RaaS GPUs : ${SERVICE_CUDA_VISIBLE_DEVICES}" +echo "Trainer model0 GPUs : ${TRAINER_MODEL0_GPUS} (FSDP dp${TRAINER0_NPROC})" +echo "RaaS port : ${RAAS_PORT}" +echo "AstraFlow port : ${ASTRAFLOW_PORT}" +echo "Sender HTTP model0 : ${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" +echo "WANDB mode : ${WANDB_MODE:-online}" +echo "==========================================================" + +trap astraflow_cleanup_trap EXIT INT TERM + +# Kill leftover processes and shared memory from prior runs. +# Defined in examples/_common/utils.sh. +astraflow_kill_stale + +# ============================================================================= +# Part 4: Launch training +# ============================================================================= +echo "[1/3] Starting AstraFlow HTTP service..." +CUDA_VISIBLE_DEVICES="" \ + python3 -u -m astraflow \ + --config "${EXPERIMENT_CONFIG}" \ + --port "${ASTRAFLOW_PORT}" \ + --host "${ASTRAFLOW_HOST}" \ + 2>&1 | tee "${LOG_DIR}/astraflow.log" & +sleep 5 + +echo "[2/3] Starting RaaS inference server (SGLang + TCP receiver)..." +CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES}" \ + python3 -u -m astraflow.raas.server \ + --host "${RAAS_HOST}" \ + --port "${RAAS_PORT}" \ + --config "${EXPERIMENT_CONFIG}" \ + --config "${RAAS_CONFIG}" \ + --engine-id "${ENGINE_ID:-default}" \ + --astraflow-url "${ASTRAFLOW_URL}" \ + 2>&1 | tee "${LOG_DIR}/raas.log" & +sleep 15 + +export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}" + +echo "[3/3] Starting trainer model0..." +CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS}" \ +WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" \ + torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \ + --master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \ + examples/launch_trainer.py \ + --config "${EXPERIMENT_CONFIG}" \ + --trainer trainer_model0 \ + "$@" \ + 2>&1 | tee "${LOG_DIR}/trainer_model0.log" diff --git a/examples/math/qwen3-4b-m2po-edf/yaml/experiment.yaml b/examples/math/qwen3-4b-m2po-edf/yaml/experiment.yaml new file mode 100644 index 0000000..65177ec --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/yaml/experiment.yaml @@ -0,0 +1,165 @@ +# ============================================================================ +# Experiment config — AstraFlow service + Trainer +# Experiment: math / qwen3-4b-m2po-edf +# +# Staleness A/B recipe: 6 RaaS GPUs vs 2 trainer GPUs deliberately +# over-provisions rollout so the fresh buffer builds real staleness +# pressure. queue_order=edf — compare against the sibling recipe. +# +# Qwen3-4B math RL with M2PO, ctx 16k, lr 5e-6, full TCP weight transfer. +# GPU layout (default, 8 GPUs): +# SERVICE_CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 -> RaaS (model0 dp=6) +# TRAINER_MODEL0_GPUS=6,7 -> Trainer model0 (FSDP, 2 GPUs) +# ============================================================================ + +# ── Experiment: identity, model, shared settings ── +experiment: + experiment_name: astraflow-math + trial_name: qwen3-4b-m2po-edf + fileroot: ./data-experiments/${experiment.experiment_name}/${experiment.trial_name} + + model_path: "Qwen/Qwen3-4B" + tokenizer_path: "Qwen/Qwen3-4B" + seed: 1 + dtype: bfloat16 + weight_transfer_mode: tcp + weight_transfer_strategies: full + +# ── RaaS: what to generate (inference-level config) ── +# model keys here also determine expected_model_ids for AstraFlow service +raas: + models: + model0: + backend: sglang + gconfig: + n_samples: 8 + temperature: 1.0 + max_new_tokens: 14000 + min_new_tokens: 0 + +# ── AstraFlow: data pipeline ── +# auto-derives: expected_model_ids from raas.models keys +# auto-derives: dump_dir from experiment.fileroot +dataflow: + host: "0.0.0.0" + port: 8000 + + buffer: + queue_order: edf + size: 10000 + replay_size: 10000 + replay_ratio: 0 + max_staleness: 8 + filter_function: filter_zero_adv + + rollout_dataset: + dataset_fn: "astraflow.dataflow.dataset.deepscaler:get_deepscaler_rl_dataset" + max_length: 2000 + + workflow_spec: + workflow_cls: "rlvr" + reward_fn: "math_verify" + enable_thinking: false + + eval_workflows: + math_eval: + workflow_cls: "rlvr" + reward_fn: "math_verify" + enable_thinking: false + gconfig_overrides: + temperature: 0.6 + n_samples: 1 + + eval_datasets: + aime24: + dataset_fn: "astraflow.dataflow.dataset.aime24x4:get_aime_2024x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + aime25: + dataset_fn: "astraflow.dataflow.dataset.aime25x4:get_aime_2025x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + amc: + dataset_fn: "astraflow.dataflow.dataset.amc24:get_amc_2024x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + minerva: + dataset_fn: "astraflow.dataflow.dataset.minervamath:get_minerva_math_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + math500: + dataset_fn: "astraflow.dataflow.dataset.math500:get_math500_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + +# ── Trainer base: shared config ── +# auto-derives from experiment: experiment_name, trial_name, fileroot, +# tokenizer_path, seed, dtype, weight_transfer_mode +# auto-derives from raas.models.: actor.path, actor.max_new_tokens, +# ref.path +# auto-derives: saver, recover, stats_logger fields from experiment section +# auto-derives: cluster.name_resolve from experiment.fileroot +# auto-derives: trial_name suffix from model_id (e.g. trial_name-model0) +trainer_base: + total_train_steps: 800 + train_batch_size: 256 + n_samples: 8 + engine: + backend: fsdp + data_parallel_size: 2 + + actor: + attn_impl: kernels-community/flash-attn2 + gradient_checkpointing: true + mb_spec: + max_tokens_per_mb: 17408 + optimizer: + type: adam + lr: 5e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + # PPO / M2PO algorithm + m2_threshold: 0.01 + eps_clip: 100.0 + eps_clip_higher: 100.0 + reward_scaling: 1 + reward_bias: 0 + kl_ctl: 0.00 + kl_penalty_coef: 0.001 + ppo_n_minibatches: 4 + reward_norm: { mean_level: group, std_level: group } + adv_norm: { mean_level: batch, std_level: batch } + + ref: + attn_impl: kernels-community/flash-attn2 + mb_spec: + max_tokens_per_mb: 17408 + + recover: + mode: auto + freq_steps: 25 + + evaluator: + eval_at_start: false + freq_steps: 25 + + stats_logger: + wandb: + mode: online + id_suffix: "uid" + +# ── Trainer for model0 — only overrides ── +trainer_model0: + model_id: model0 + stats_logger: + wandb: + tags: ["m2po", "math", "astraflow-v2", "qwen3-4b", "tcp", "ctx16k"] diff --git a/examples/math/qwen3-4b-m2po-edf/yaml/raas.yaml b/examples/math/qwen3-4b-m2po-edf/yaml/raas.yaml new file mode 100644 index 0000000..16ec814 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-edf/yaml/raas.yaml @@ -0,0 +1,37 @@ +# ============================================================================ +# RaaS config — Inference serving instance (hardware/resources) +# Experiment: math / qwen3-4b-m2po-edf +# +# Hardware: 6x GPU, TP=1 +# model0: DP=6, TP=1 +# +# Merged with experiment.yaml at launch (--config experiment.yaml --config raas.yaml) +# experiment.yaml provides: model_path, tokenizer_path, seed, dtype, models/gconfig +# ============================================================================ + +rollout: + max_concurrent_rollouts: 512 + # Cap concurrent eval prefills to bound peak KV pressure during the + # ~3.5k-item eval burst (5 datasets x repeat=4) — default 128 OOMs sglang. + max_concurrent_evals: 64 + pause_grace_period: 3 + # Adaptive availability — drive /availability off sglang /get_load. + enable_adaptive_availability: true + target_waiting_queue_per_dp: 4 + adaptive_step_size: 4 + load_cache_ttl_ms: 100 + +engine: + model0: + backend: sglang + data_parallel_size: 6 + +sglang: + context_length: 16384 + mem_fraction_static: 0.7 + # Chunk prefill so the LM-head fp32 logits upcast can't spike to ~16 GiB + # and OOM a worker (root cause of the step-101 hang). Bounds per-forward + # tokens to 8192. + chunked_prefill_size: 8192 + max_running_requests: null + skip_tokenizer_init: true diff --git a/examples/math/qwen3-4b-m2po-fifo/README.md b/examples/math/qwen3-4b-m2po-fifo/README.md new file mode 100644 index 0000000..36272c9 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/README.md @@ -0,0 +1,22 @@ +# Qwen3-4B Math RL — staleness A/B (queue_order=fifo) + +The FIFO arm of the rollout-buffer consumption-order A/B. Identical to +[`qwen3-4b-m2po-edf`](../qwen3-4b-m2po-edf/) except +`dataflow.buffer.queue_order: fifo` (the historical completion-order +behavior) and `trial_name`. + +See the sibling recipe's README for the motivation, setup (6 RaaS GPUs vs +2 FSDP trainer GPUs to build staleness pressure), run commands, and the +A/B results. Summary: FIFO's staleness drops are length-biased (it +preferentially expires long/hard generations — dropped samples ~1.3–1.7× +longer than consumed ones), which costs ~5 overall avg@k points at step 100 +versus edf and delays reaching the eval plateau by ~150 steps. + +`queue_order: edf` is the default; this recipe exists to reproduce the +comparison. + +## Run + +```bash +bash examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh +``` diff --git a/examples/math/qwen3-4b-m2po-fifo/scripts/1_astraflow.sh b/examples/math/qwen3-4b-m2po-fifo/scripts/1_astraflow.sh new file mode 100755 index 0000000..0bc61d2 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/scripts/1_astraflow.sh @@ -0,0 +1,36 @@ +#!/bin/bash +set -euo pipefail +# [1/3] Launch AstraFlow HTTP service +# +# Usage (terminal 1): +# bash examples/math/qwen3-4b-m2po-fifo/scripts/1_astraflow.sh + +export CUDA_VISIBLE_DEVICES="" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== AstraFlow HTTP Service ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "Port : ${ASTRAFLOW_PORT}" +echo "===============================" + +python3 -u -m astraflow \ + --config "${EXPERIMENT_CONFIG}" \ + --port "${ASTRAFLOW_PORT}" \ + --host "${ASTRAFLOW_HOST}" \ + 2>&1 | tee "${LOG_DIR}/astraflow.log" diff --git a/examples/math/qwen3-4b-m2po-fifo/scripts/2_raas.sh b/examples/math/qwen3-4b-m2po-fifo/scripts/2_raas.sh new file mode 100755 index 0000000..e4c8411 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/scripts/2_raas.sh @@ -0,0 +1,44 @@ +#!/bin/bash +set -euo pipefail +# [2/3] Launch RaaS inference server (SGLang + TCP receiver) +# +# Usage (terminal 2, after AstraFlow is ready): +# bash examples/math/qwen3-4b-m2po-fifo/scripts/2_raas.sh + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5}" +export RAAS_HOST="${RAAS_HOST:-0.0.0.0}" +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="${ASTRAFLOW_URL:-http://127.0.0.1:${ASTRAFLOW_PORT}}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== RaaS Inference Server (SGLang + TCP receiver) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "RaaS config : ${RAAS_CONFIG}" +echo "GPUs : ${CUDA_VISIBLE_DEVICES}" +echo "Port : ${RAAS_PORT}" +echo "AstraFlow URL : ${ASTRAFLOW_URL}" +echo "=======================================================" + +python3 -u -m astraflow.raas.server \ + --host "${RAAS_HOST}" \ + --port "${RAAS_PORT}" \ + --config "${EXPERIMENT_CONFIG}" \ + --config "${RAAS_CONFIG}" \ + --engine-id "${ENGINE_ID:-default}" \ + --astraflow-url "${ASTRAFLOW_URL}" \ + 2>&1 | tee "${LOG_DIR}/raas.log" diff --git a/examples/math/qwen3-4b-m2po-fifo/scripts/3_trainer_model0.sh b/examples/math/qwen3-4b-m2po-fifo/scripts/3_trainer_model0.sh new file mode 100755 index 0000000..6468d94 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/scripts/3_trainer_model0.sh @@ -0,0 +1,47 @@ +#!/bin/bash +set -euo pipefail +# [3/3] Launch Trainer for model0 (TCP, sender_agent on local_rank 0) +# +# Usage (terminal 3, after AstraFlow and RaaS are ready): +# bash examples/math/qwen3-4b-m2po-fifo/scripts/3_trainer_model0.sh + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +astraflow_load_experiment_env + +export CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS:-6,7}" +TRAINER0_NPROC="$(echo "${CUDA_VISIBLE_DEVICES}" | awk -F',' '{print NF}')" + +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}" +export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}" + +# sender_agent (in trainer) listens on this HTTP port +export WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh. +astraflow_setup_env + +echo "=== Trainer model0 (TCP) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "GPUs : ${CUDA_VISIBLE_DEVICES} (FSDP dp${TRAINER0_NPROC})" +echo "AstraFlow : ${ASTRAFLOW_URL}" +echo "RaaS : ${ASTRAFLOW_RAAS_URL}" +echo "Sender HTTP : ${WEIGHT_TRANSFER_HTTP_PORT}" +echo "WANDB mode : ${WANDB_MODE:-online}" +echo "==========================================" + +torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \ + --master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \ + examples/launch_trainer.py \ + --config "${EXPERIMENT_CONFIG}" \ + --trainer trainer_model0 \ + "$@" 2>&1 | tee "${LOG_DIR}/trainer_model0.log" diff --git a/examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh b/examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh new file mode 100755 index 0000000..bbe6db2 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh @@ -0,0 +1,104 @@ +#!/bin/bash +set -euo pipefail +# All-in-one launcher for AstraFlow v2 math training (Qwen3-4B, M2PO, TCP). +# +# Launches 3 processes: +# 1. AstraFlow HTTP service (CPU-only) +# 2. RaaS inference server (SGLang, SERVICE_CUDA_VISIBLE_DEVICES) +# 3. Trainer model0 (math, TRAINER_MODEL0_GPUS) +# +# Usage: +# bash examples/math/qwen3-4b-m2po-fifo/scripts/run_qwen3-4b-m2po-fifo.sh + +# ============================================================================= +# Part 1: Load env and settings +# ============================================================================= +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${REPO_ROOT}" +export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +YAML_DIR="${SCRIPT_DIR}/yaml" +export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}" +export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}" +source "${REPO_ROOT}/examples/_common/utils.sh" +# Export EXP_NAME and TRIAL_NAME from the experiment YAML. +# Defined in examples/_common/utils.sh. +astraflow_load_experiment_env + +# ============================================================================= +# Part 2: Set up env +# ============================================================================= +# GPU assignments (default: 4 GPUs for inference, 4 for training) +export SERVICE_CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5}" +export TRAINER_MODEL0_GPUS="${TRAINER_MODEL0_GPUS:-6,7}" +# Ports / URLs (each component gets its own port) +export RAAS_HOST="${RAAS_HOST:-0.0.0.0}" +export RAAS_PORT="${RAAS_PORT:-19190}" +export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}" +export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}" +export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}" +export WEIGHT_TRANSFER_HTTP_PORT_MODEL0="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}" + +TRAINER0_NPROC="$(echo "${TRAINER_MODEL0_GPUS}" | awk -F',' '{print NF}')" + +# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. +# Defined in examples/_common/utils.sh. +astraflow_setup_env + +# ============================================================================= +# Part 3: Print info and clean up +# ============================================================================= +echo "=== AstraFlow v2 (Qwen3-4B, math, M2PO, ctx16k, TCP) ===" +echo "Experiment config : ${EXPERIMENT_CONFIG}" +echo "RaaS config : ${RAAS_CONFIG}" +echo "RaaS GPUs : ${SERVICE_CUDA_VISIBLE_DEVICES}" +echo "Trainer model0 GPUs : ${TRAINER_MODEL0_GPUS} (FSDP dp${TRAINER0_NPROC})" +echo "RaaS port : ${RAAS_PORT}" +echo "AstraFlow port : ${ASTRAFLOW_PORT}" +echo "Sender HTTP model0 : ${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" +echo "WANDB mode : ${WANDB_MODE:-online}" +echo "==========================================================" + +trap astraflow_cleanup_trap EXIT INT TERM + +# Kill leftover processes and shared memory from prior runs. +# Defined in examples/_common/utils.sh. +astraflow_kill_stale + +# ============================================================================= +# Part 4: Launch training +# ============================================================================= +echo "[1/3] Starting AstraFlow HTTP service..." +CUDA_VISIBLE_DEVICES="" \ + python3 -u -m astraflow \ + --config "${EXPERIMENT_CONFIG}" \ + --port "${ASTRAFLOW_PORT}" \ + --host "${ASTRAFLOW_HOST}" \ + 2>&1 | tee "${LOG_DIR}/astraflow.log" & +sleep 5 + +echo "[2/3] Starting RaaS inference server (SGLang + TCP receiver)..." +CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES}" \ + python3 -u -m astraflow.raas.server \ + --host "${RAAS_HOST}" \ + --port "${RAAS_PORT}" \ + --config "${EXPERIMENT_CONFIG}" \ + --config "${RAAS_CONFIG}" \ + --engine-id "${ENGINE_ID:-default}" \ + --astraflow-url "${ASTRAFLOW_URL}" \ + 2>&1 | tee "${LOG_DIR}/raas.log" & +sleep 15 + +export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}" + +echo "[3/3] Starting trainer model0..." +CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS}" \ +WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" \ + torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \ + --master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \ + examples/launch_trainer.py \ + --config "${EXPERIMENT_CONFIG}" \ + --trainer trainer_model0 \ + "$@" \ + 2>&1 | tee "${LOG_DIR}/trainer_model0.log" diff --git a/examples/math/qwen3-4b-m2po-fifo/yaml/experiment.yaml b/examples/math/qwen3-4b-m2po-fifo/yaml/experiment.yaml new file mode 100644 index 0000000..71dcb1a --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/yaml/experiment.yaml @@ -0,0 +1,165 @@ +# ============================================================================ +# Experiment config — AstraFlow service + Trainer +# Experiment: math / qwen3-4b-m2po-fifo +# +# Staleness A/B recipe: 6 RaaS GPUs vs 2 trainer GPUs deliberately +# over-provisions rollout so the fresh buffer builds real staleness +# pressure. queue_order=fifo — compare against the sibling recipe. +# +# Qwen3-4B math RL with M2PO, ctx 16k, lr 5e-6, full TCP weight transfer. +# GPU layout (default, 8 GPUs): +# SERVICE_CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 -> RaaS (model0 dp=6) +# TRAINER_MODEL0_GPUS=6,7 -> Trainer model0 (FSDP, 2 GPUs) +# ============================================================================ + +# ── Experiment: identity, model, shared settings ── +experiment: + experiment_name: astraflow-math + trial_name: qwen3-4b-m2po-fifo + fileroot: ./data-experiments/${experiment.experiment_name}/${experiment.trial_name} + + model_path: "Qwen/Qwen3-4B" + tokenizer_path: "Qwen/Qwen3-4B" + seed: 1 + dtype: bfloat16 + weight_transfer_mode: tcp + weight_transfer_strategies: full + +# ── RaaS: what to generate (inference-level config) ── +# model keys here also determine expected_model_ids for AstraFlow service +raas: + models: + model0: + backend: sglang + gconfig: + n_samples: 8 + temperature: 1.0 + max_new_tokens: 14000 + min_new_tokens: 0 + +# ── AstraFlow: data pipeline ── +# auto-derives: expected_model_ids from raas.models keys +# auto-derives: dump_dir from experiment.fileroot +dataflow: + host: "0.0.0.0" + port: 8000 + + buffer: + queue_order: fifo + size: 10000 + replay_size: 10000 + replay_ratio: 0 + max_staleness: 8 + filter_function: filter_zero_adv + + rollout_dataset: + dataset_fn: "astraflow.dataflow.dataset.deepscaler:get_deepscaler_rl_dataset" + max_length: 2000 + + workflow_spec: + workflow_cls: "rlvr" + reward_fn: "math_verify" + enable_thinking: false + + eval_workflows: + math_eval: + workflow_cls: "rlvr" + reward_fn: "math_verify" + enable_thinking: false + gconfig_overrides: + temperature: 0.6 + n_samples: 1 + + eval_datasets: + aime24: + dataset_fn: "astraflow.dataflow.dataset.aime24x4:get_aime_2024x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + aime25: + dataset_fn: "astraflow.dataflow.dataset.aime25x4:get_aime_2025x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + amc: + dataset_fn: "astraflow.dataflow.dataset.amc24:get_amc_2024x4_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + minerva: + dataset_fn: "astraflow.dataflow.dataset.minervamath:get_minerva_math_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + math500: + dataset_fn: "astraflow.dataflow.dataset.math500:get_math500_test_dataset" + max_length: 2000 + repeat: 4 + eval_workflow: math_eval + +# ── Trainer base: shared config ── +# auto-derives from experiment: experiment_name, trial_name, fileroot, +# tokenizer_path, seed, dtype, weight_transfer_mode +# auto-derives from raas.models.: actor.path, actor.max_new_tokens, +# ref.path +# auto-derives: saver, recover, stats_logger fields from experiment section +# auto-derives: cluster.name_resolve from experiment.fileroot +# auto-derives: trial_name suffix from model_id (e.g. trial_name-model0) +trainer_base: + total_train_steps: 800 + train_batch_size: 256 + n_samples: 8 + engine: + backend: fsdp + data_parallel_size: 2 + + actor: + attn_impl: kernels-community/flash-attn2 + gradient_checkpointing: true + mb_spec: + max_tokens_per_mb: 17408 + optimizer: + type: adam + lr: 5e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + # PPO / M2PO algorithm + m2_threshold: 0.01 + eps_clip: 100.0 + eps_clip_higher: 100.0 + reward_scaling: 1 + reward_bias: 0 + kl_ctl: 0.00 + kl_penalty_coef: 0.001 + ppo_n_minibatches: 4 + reward_norm: { mean_level: group, std_level: group } + adv_norm: { mean_level: batch, std_level: batch } + + ref: + attn_impl: kernels-community/flash-attn2 + mb_spec: + max_tokens_per_mb: 17408 + + recover: + mode: auto + freq_steps: 25 + + evaluator: + eval_at_start: false + freq_steps: 25 + + stats_logger: + wandb: + mode: online + id_suffix: "uid" + +# ── Trainer for model0 — only overrides ── +trainer_model0: + model_id: model0 + stats_logger: + wandb: + tags: ["m2po", "math", "astraflow-v2", "qwen3-4b", "tcp", "ctx16k"] diff --git a/examples/math/qwen3-4b-m2po-fifo/yaml/raas.yaml b/examples/math/qwen3-4b-m2po-fifo/yaml/raas.yaml new file mode 100644 index 0000000..9b2f5e7 --- /dev/null +++ b/examples/math/qwen3-4b-m2po-fifo/yaml/raas.yaml @@ -0,0 +1,37 @@ +# ============================================================================ +# RaaS config — Inference serving instance (hardware/resources) +# Experiment: math / qwen3-4b-m2po-fifo +# +# Hardware: 6x GPU, TP=1 +# model0: DP=6, TP=1 +# +# Merged with experiment.yaml at launch (--config experiment.yaml --config raas.yaml) +# experiment.yaml provides: model_path, tokenizer_path, seed, dtype, models/gconfig +# ============================================================================ + +rollout: + max_concurrent_rollouts: 512 + # Cap concurrent eval prefills to bound peak KV pressure during the + # ~3.5k-item eval burst (5 datasets x repeat=4) — default 128 OOMs sglang. + max_concurrent_evals: 64 + pause_grace_period: 3 + # Adaptive availability — drive /availability off sglang /get_load. + enable_adaptive_availability: true + target_waiting_queue_per_dp: 4 + adaptive_step_size: 4 + load_cache_ttl_ms: 100 + +engine: + model0: + backend: sglang + data_parallel_size: 6 + +sglang: + context_length: 16384 + mem_fraction_static: 0.7 + # Chunk prefill so the LM-head fp32 logits upcast can't spike to ~16 GiB + # and OOM a worker (root cause of the step-101 hang). Bounds per-forward + # tokens to 8192. + chunked_prefill_size: 8192 + max_running_requests: null + skip_tokenizer_init: true From d38680af267d29d9b0ef2f9db16532607a338914 Mon Sep 17 00:00:00 2001 From: Haizhong Date: Tue, 14 Jul 2026 11:07:27 -0400 Subject: [PATCH 3/3] docs: warn that flash-attn must be installed from outside the repo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Following the conda install steps verbatim leaves flash-attn 2.8.3 silently uninstalled: pyproject's [tool.uv] exclude-dependencies=["flash-attn"] (needed so project installs don't try to resolve/build it) also applies to the explicit `uv pip install flash-attn` in Step 4 when run from the repo directory — the command exits 0 without installing anything. The failure is then masked because flash-attn-4 (sglang's attention backend) leaves a namespace stub at site-packages/flash_attn/, so `import flash_attn` appears to work and Step 6's `flash_attn.__version__` check fails with a confusing AttributeError instead of a clear signal. Step 4 now cds to /tmp before installing (mirroring the Dockerfile fix from d192bbe) and explains why; Step 6 verifies by importing flash_attn_func and reading the wheel metadata version, with the failure signature and remedy spelled out. Both envs built from these docs on this cluster had hit the trap unnoticed. --- docs/en/get-started/installation.md | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/docs/en/get-started/installation.md b/docs/en/get-started/installation.md index 6d2a7aa..1d7030e 100644 --- a/docs/en/get-started/installation.md +++ b/docs/en/get-started/installation.md @@ -51,6 +51,15 @@ This is FlashAttention-**2** (`import flash_attn`), used by the FSDP trainer. It is excluded from uv resolution (see `pyproject.toml` `[tool.uv]`) and built from source, so it needs the CUDA 13 toolchain and a roomy build-temp directory: +> **Run this install from OUTSIDE the repo directory.** Inside the repo, uv +> picks up `pyproject.toml`'s `[tool.uv] exclude-dependencies = ["flash-attn"]` +> (which exists so project installs don't try to resolve/build flash-attn) and +> applies it to this explicit install too — the command **exits 0 without +> installing anything**. Worse, `import flash_attn` still appears to succeed +> afterwards, because the `flash-attn-4` package (pulled in by sglang) leaves a +> namespace stub at `site-packages/flash_attn/` — see the verification in +> Step 6 that actually catches this. + ```bash # nvcc must be on PATH and match torch's CUDA (13.0 for torch 2.11+cu130) export CUDA_HOME=/usr/local/cuda-13.0 @@ -61,7 +70,11 @@ export PATH="$CUDA_HOME/bin:$PATH" # "nvFatbin error: empty input" or "Disk quota exceeded" from truncated temps. export TMPDIR=/tmp/fa-build && mkdir -p "$TMPDIR" +# IMPORTANT: leave the repo so uv ignores the project's [tool.uv] table +# (same reason the Dockerfile installs flash-attn from /tmp). +cd /tmp uv pip install "flash-attn==2.8.3" --no-build-isolation +cd - ``` > On a single-GPU-arch box you can speed up the build and shrink its footprint @@ -155,12 +168,22 @@ Verify Flash Attention and SGLang: ```bash python -c " -import flash_attn, sglang -print(f'flash-attn: {flash_attn.__version__}') +# NOTE: a bare 'import flash_attn' succeeding is NOT proof FA2 is installed — +# flash-attn-4 (sglang's attention backend) leaves a namespace stub at +# site-packages/flash_attn/ that imports fine but is empty. Import a real +# symbol and read the wheel metadata instead. +from importlib.metadata import version +from flash_attn import flash_attn_func # fails if FA2 was silently skipped +import sglang +print(f'flash-attn: {version(\"flash-attn\")}') print(f'sglang: {sglang.__version__}') " ``` +If the `flash_attn_func` import fails with `cannot import name ... (unknown +location)`, flash-attn was silently skipped — re-run the Step 4 install from a +directory **outside** the repo. + ## Option B: Docker Pre-built images are published on Docker Hub — they skip the from-source steps above