From bc01191661e7841f80539c715c20acb6c9115d70 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Mon, 9 Nov 2026 15:52:30 +0800 Subject: [PATCH] Support multi-node multi-Draft training Share node-local Target features across independent cross-node Draft groups, with per-Draft placement, checkpoint recovery, preflight and reader fencing. Add the multinode-shared recipe, usage documentation and CPU regression coverage. --- README.md | 4 +- README_zh-CN.md | 2 +- ROADMAP.md | 3 +- aspec/commands/recipes.py | 4 +- aspec/controller/inference_manager.py | 50 ++++- aspec/controller/shared_source.py | 175 ++++++++++++++-- aspec/multi_draft.py | 162 +++++++++----- aspec/train_entry.py | 11 +- configs/recipes/README.md | 2 + ...n3_8b_4k_multinode_multi_draft_nvlink.yaml | 155 ++++++++++++++ docs/aspec.zh-CN.md | 2 +- docs/contributing/validation.md | 13 +- docs/framework/ARCHITECTURE.md | 8 +- docs/framework/DECISIONS.md | 12 +- docs/framework/ROADMAP.md | 2 +- docs/support.md | 39 ++-- docs/user-guide/multinode.md | 3 +- docs/user-guide/recipes.md | 5 + docs/user-guide/sharing.md | 158 ++++++++++++-- .../test_inference_manager_lifecycle.py | 32 ++- tests/{ => cpu}/test_multi_draft_plan.py | 71 ++++++- tests/{ => cpu}/test_multi_draft_runtime.py | 112 ++++++++-- tests/cpu/test_release_recipes.py | 8 +- tests/{ => cpu}/test_shared_source.py | 197 ++++++++++++++++++ tests/test_startup_cleanup.py | 33 ++- tools/recipe_catalog.json | 15 +- tools/recipes.py | 5 +- 27 files changed, 1097 insertions(+), 186 deletions(-) create mode 100644 configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml rename tests/{ => cpu}/test_inference_manager_lifecycle.py (94%) rename tests/{ => cpu}/test_multi_draft_plan.py (79%) rename tests/{ => cpu}/test_multi_draft_runtime.py (78%) rename tests/{ => cpu}/test_shared_source.py (75%) diff --git a/README.md b/README.md index a58f221..2a90796 100644 --- a/README.md +++ b/README.md @@ -25,8 +25,8 @@ independent experiments, and take a checkpoint through evaluation and export. - **One Target, multiple concurrent Draft experiments.** Compare algorithms, structures or hyperparameters while generating compatible Target features once. Each Draft keeps its own weights, optimizer, progress and checkpoints, without a shared optimizer-step barrier. - See [shared Target training](docs/user-guide/sharing.md) and the - [comparison and recovery walkthrough](examples/qwen3-8b-shared-target/README.md). + See [shared Target training](docs/user-guide/sharing.md) for single-node and static multi-node workflows, + and the [comparison and recovery walkthrough](examples/qwen3-8b-shared-target/README.md). - **Make long-context training practical.** Combine sequence packing, context parallelism and optimized attention execution to reduce padding waste and distribute long sequences across GPUs. - **Explore without rebuilding the workflow.** Switch between algorithm and model recipes with diff --git a/README_zh-CN.md b/README_zh-CN.md index a46b68a..5ffd9c2 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -22,7 +22,7 @@ ASpec 基于 [TorchSpec](https://github.com/lightseekorg/TorchSpec) 源码开发 按有效工作量组织训练，减少两端等待。 - **一个 Target，多个 Draft 并行实验。** 比较算法、结构或超参数时，兼容的 Target 特征只生成 一次。各 Draft 独立维护权重、优化器、进度和 checkpoint，无需每次更新都相互等待。 - 当前单机用法见[共享训练指南](docs/user-guide/sharing.md)。 + 单机及静态多机的用法见[共享训练指南](docs/user-guide/sharing.md)。 - **高效处理长上下文。** 结合序列打包、上下文并行与优化的注意力执行，减少 padding 浪费， 将长序列训练分布到多张 GPU。 - **换算法，不重搭工作流。** 通过 recipe 选择算法与结构，沿用一致的配置、报告、恢复、评估和 diff --git a/ROADMAP.md b/ROADMAP.md index 56c15b9..3c512cd 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -16,7 +16,8 @@ The authoritative details and limitations are in [supported paths](docs/support. - [x] Sequence packing, context parallelism and independent Drafts sharing compatible Target data. - [x] Primary SGLang + Qwen3/Ling path on single-node multi-GPU systems. - [x] Bounded static multi-node, single-Draft execution. -- [ ] Multi-node multi-Draft training: share compatible Target features across independent Draft +- [x] Static multi-node multi-Draft training with independent packing: + share compatible Target features across independent Draft experiments, preserving per-Draft progress and checkpoint recovery without requiring a shared filesystem. - [ ] Elastic membership and broader failure recovery. diff --git a/aspec/commands/recipes.py b/aspec/commands/recipes.py index d1a6c23..459f917 100644 --- a/aspec/commands/recipes.py +++ b/aspec/commands/recipes.py @@ -391,7 +391,7 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None run = stage / label overrides = [*base, f"training.num_train_steps={limit}"] if label == "resumed": - if recipe == "shared-target": + if entry.get("draft_experiments", 1) > 1: from omegaconf import OmegaConf experiments = OmegaConf.to_container( @@ -425,7 +425,7 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None limit, label == "resumed", require_eval=recipe != "quickstart", - expected_ranks=entry["draft_gpus"] // (2 if recipe == "shared-target" else 1), + expected_ranks=entry["draft_gpus"] // entry.get("draft_experiments", 1), ) if recipe == "qwen27b-long": from aspec.commands._recipe_capacity import audit_long_run diff --git a/aspec/controller/inference_manager.py b/aspec/controller/inference_manager.py index 9018e41..3ed6be0 100644 --- a/aspec/controller/inference_manager.py +++ b/aspec/controller/inference_manager.py @@ -56,7 +56,7 @@ ) from aspec.utils.logging import logger from aspec.utils.types import InferenceInput, InferenceOutput -from aspec.work_lifecycle import WORK_IDENTITY_KEYS, require_work_identity +from aspec.work_lifecycle import EXECUTION_ID_KEY, WORK_IDENTITY_KEYS, require_work_identity MOONCAKE_BACKPRESSURE_POLL_INTERVAL = 0.5 # seconds MOONCAKE_BACKPRESSURE_LOG_INTERVAL = 5.0 # seconds @@ -774,6 +774,11 @@ async def get_status(self) -> dict: if self._node_ready_watermarks: status.update( { + "ready_credit_scope": ( + "shared_admission_audience" + if getattr(getattr(self, "args", None), "aspec_shared_source", False) + else "controller" + ), "ready_watermark_by_node": dict(self._node_ready_watermarks), "sample_budget_by_node": dict(self._sample_budget_by_node), "scheduled_samples_by_node": dict( @@ -868,6 +873,15 @@ async def _start_dispatch(self, dispatch_size: int) -> bool: if inflight + dispatch_size > self._batch_size * self._engines.capacity_by_node[node] } if self._node_ready_watermarks else set() + pinned_source = bool(getattr(getattr(self, "args", None), "aspec_shared_source", False)) and bool(self._node_ready_watermarks) + if pinned_source: + for node in self._node_ready_watermarks: + available = sum( + entry.metadata.get("shared_source_node") == node + for entry in self._prompt_buffer + ) + if available < dispatch_size: + excluded_nodes.add(node) while True: lease = self._engines.acquire_nowait(excluded_nodes=excluded_nodes) if lease is None: @@ -883,9 +897,13 @@ async def _start_dispatch(self, dispatch_size: int) -> bool: excluded_nodes.add(node) continue try: - admitted = await self._reserve_pool_capacity( - dispatch_size, node=node - ) + source_execution_ids = [ + entry.metadata[EXECUTION_ID_KEY] + for entry in self._prompt_buffer + if entry.metadata.get("shared_source_node") == node + ][:dispatch_size] if pinned_source else None + admission_options = {"source_execution_ids": source_execution_ids} if pinned_source else {} + admitted = await self._reserve_pool_capacity(dispatch_size, node=node, **admission_options) except BaseException: lease.release() raise @@ -894,7 +912,17 @@ async def _start_dispatch(self, dispatch_size: int) -> bool: lease.release() excluded_nodes.add(lease.node) - entries = self._take_batch(dispatch_size) + if pinned_source: + entries = [] + remaining = deque() + for entry in self._prompt_buffer: + if len(entries) < dispatch_size and entry.metadata.get("shared_source_node") == node: + entries.append(entry) + else: + remaining.append(entry) + self._prompt_buffer = remaining + else: + entries = self._take_batch(dispatch_size) self._scheduled_samples += len(entries) if node is not None and self._sample_budget_by_node: self._scheduled_samples_by_node[node] += len(entries) @@ -972,7 +1000,11 @@ async def _refill_buffer(self) -> None: if len(self._prompt_buffer) >= self._buffer_low_watermark: return - if self._max_pool_size > 0: + # Shared readers can have opposite node pressures. The global maximum + # can be full everywhere while a particular admission still has credit. + # Fetch bounded prompts here and gate their actual audiences at dispatch. + pinned_source = bool(getattr(getattr(self, "args", None), "aspec_shared_source", False)) and bool(self._node_ready_watermarks) + if self._max_pool_size > 0 and not pinned_source: if self._node_ready_watermarks: pool_sizes = await self.controller.get_ready_pool_size_by_node.remote() pool_full = all( @@ -1018,6 +1050,7 @@ async def _reserve_pool_capacity( incoming_samples: int, *, node: str | None = None, + source_execution_ids: list[str] | None = None, ) -> bool: """Gate on ready outputs while accounting for independently bounded requests. @@ -1041,7 +1074,10 @@ async def _reserve_pool_capacity( raise RuntimeError( f"engine lease has no planned locality credit: node={node!r}" ) - pool_sizes = await self.controller.get_ready_pool_size_by_node.remote() + if source_execution_ids is None: + pool_sizes = await self.controller.get_ready_pool_size_by_node.remote() + else: + pool_sizes = await self.controller.get_ready_pool_size_by_node.remote(source_execution_ids) unknown_nodes = set(pool_sizes) - set(self._node_ready_watermarks) if unknown_nodes: raise RuntimeError( diff --git a/aspec/controller/shared_source.py b/aspec/controller/shared_source.py index 6ea7bc2..7b194c2 100644 --- a/aspec/controller/shared_source.py +++ b/aspec/controller/shared_source.py @@ -26,6 +26,7 @@ from aspec.config.draft_execution import is_direct_cp_backend from aspec.controller.training_controller import TrainingControllerState +from aspec.locality import allocate_node_ready_watermarks from aspec.transfer.mooncake.p2p import ( MooncakeP2PDescriptor, MooncakeP2PTransfer, @@ -87,9 +88,26 @@ def __getattr__(self, method: str): class SharedSourceState: def __init__(self, source_args, experiment_args: dict): + scheduler_args = copy.copy(source_args) + # The source selects documents, while each child owns its DP queues. + scheduler_args.mooncake_feature_affinity = "any" self.source = TrainingControllerState( - source_args, 1, with_train_queues=False, with_eval_queues=False + scheduler_args, 1, with_train_queues=False, with_eval_queues=False ) + self._source_nodes = ( + list(source_args.training_node_ips) + if getattr(source_args, "mooncake_feature_affinity", "any") == "node_local_required" + else [] + ) + self._node_cursor = 0 + self._node_watermarks = ( + allocate_node_ready_watermarks( + int(getattr(source_args, "max_sample_pool_size", 0) or 0), + dict.fromkeys(self._source_nodes, 1), + ) if self._source_nodes else {} + ) + self._pending_entries = deque() + self._remaining_by_node = {} self.experiments = { name: TrainingControllerState(args, args.dp_size, with_eval_queues=bool(getattr(args, "eval_data_path", None))) for name, args in experiment_args.items() @@ -110,6 +128,7 @@ def __init__(self, source_args, experiment_args: dict): self._lock = threading.RLock() self._dispatch_locks = {name: threading.Lock() for name in self.experiments} self._audiences: dict[str, tuple[str, ...]] = {} + self._admitted_nodes: dict[str, str] = {} self._order: deque[str] = deque() self._completed = {} self._descriptors = {name: {} for name in self.experiments} @@ -225,6 +244,13 @@ def invoke(self, name, method, args, kwargs): # rather than cumulative sample count, define this attempt. steps = child.args.num_train_steps - completed_steps self.remaining[name] = None if getattr(child.args, "dflash_packing", False) and steps else steps * child.args.global_batch_size + if self._source_nodes: + budget = None + if self.remaining[name] is not None: + budget, remainder = divmod(self.remaining[name], len(self._source_nodes)) + if remainder: + raise ValueError("shared source demand must divide into complete node-local rounds") + self._remaining_by_node[name] = dict.fromkeys(self._source_nodes, budget) self.ready.add(name) self._start_epochs[name] = child._dataset_epoch self._committed_epoch[name] = child._dataset_epoch @@ -257,7 +283,12 @@ def invoke(self, name, method, args, kwargs): return result if method == "get_checkpoint_work_state": state = child.get_checkpoint_work_state() - state["dataset_epoch"] = self._committed_epoch[name] + # Packing may commit later documents before an older epoch's + # tail. Restart at the oldest outstanding epoch and skip the + # cumulative committed work set, rather than losing that tail. + state["dataset_epoch"] = min( + [self._committed_epoch[name], *self._epochs[name].values()] + ) return state return getattr(child, method)(*args, **kwargs) @@ -272,6 +303,9 @@ def get_prompts(self, count): for index, entry in enumerate(entries, self._eval_cursor): identity = entry.metadata[EXECUTION_ID_KEY] audience = tuple(n for n in self._eval_names if index < len(self.experiments[n]._stored_eval_dataset)) + if self._source_nodes: + entry.metadata["shared_source_node"] = self._source_nodes[index % len(self._source_nodes)] + self._admitted_nodes[identity] = entry.metadata["shared_source_node"] self._audiences[identity] = audience self._order.append(identity) for name in audience: @@ -281,33 +315,88 @@ def get_prompts(self, count): if not self.active.issubset(self.ready): return [] entries = [] + # Include requests not yet published: fetching ahead must not send + # an entire prompt window to the node that currently looks empty. + pressure = {} + if self._source_nodes: + for name in self.active: + ready = self.experiments[name].get_ready_pool_size_by_node() + pressure[name] = { + node: ready.get(node, 0) + sum( + name in readers and self._admitted_nodes[key] == node + for key, readers in self._audiences.items() + ) for node in self._source_nodes + } + while self._pending_entries and len(entries) < count: + entries.append(self._pending_entries.popleft()) while len(entries) < count and any(self.remaining[name] is None or self.remaining[name] > 0 for name in self.active): + # Reserve ready credit before fetching a logical document. All + # not-yet-published requests count here, so a prefetched shared + # audience can never block behind an already full reader. + if self._source_nodes and any( + (self.remaining[name] is None or self.remaining[name] > 0) + and not any( + (self._remaining_by_node[name][node] is None or self._remaining_by_node[name][node] > 0) + and (not self._node_watermarks[node] or pressure[name][node] < self._node_watermarks[node]) + for node in self._source_nodes + ) for name in self.active + ): + break candidates = self.source.get_prompts(1) if not candidates: self.source.reload_dataset() continue entry = candidates[0] work_id = entry.metadata[WORK_ID_KEY] - audience = tuple( - sorted( - name - for name in self.active - if (self.remaining[name] is None or self.remaining[name] > 0) - and self.source._dataset_epoch >= self._start_epochs[name] - and work_id not in self.experiments[name]._restored_committed_work_ids - ) - ) - if not audience: - continue - entry.metadata["shared_source_epoch"] = self.source._dataset_epoch - execution_id = entry.metadata[EXECUTION_ID_KEY] - self._audiences[execution_id] = audience - self._order.append(execution_id) - for name in audience: - if self.remaining[name] is not None: - self.remaining[name] -= 1 - self._epochs[name][execution_id] = self.source._dataset_epoch - entries.append(entry) + audience = { + name for name in self.active + if (self.remaining[name] is None or self.remaining[name] > 0) + and self.source._dataset_epoch >= self._start_epochs[name] + and work_id not in self.experiments[name]._restored_committed_work_ids + } + while audience: + node = None + selected = audience + if self._source_nodes: + # Readers can consume different numbers of documents per + # packed row. Route each reader to its least loaded node; + # share the capture when their placements agree. + order = self._source_nodes[self._node_cursor:] + self._source_nodes[:self._node_cursor] + preferred = {} + for name in sorted(audience): + eligible = [candidate for candidate in order + if (self._remaining_by_node[name][candidate] is None + or self._remaining_by_node[name][candidate] > 0) + and (not self._node_watermarks[candidate] + or pressure[name][candidate] < self._node_watermarks[candidate])] + if not eligible: + raise RuntimeError("shared node-local demand disagrees with remaining work") + preferred[name] = min(eligible, key=pressure[name].__getitem__) + node = max(order, key=lambda candidate: sum(value == candidate for value in preferred.values())) + self._node_cursor = (self._source_nodes.index(node) + 1) % len(self._source_nodes) + selected = {name for name in audience if preferred[name] == node} + admitted = copy.copy(entry) + admitted.metadata = dict(entry.metadata, shared_source_epoch=self.source._dataset_epoch) + execution_id = entry.metadata[EXECUTION_ID_KEY] + if node is not None: + admitted.metadata["shared_source_node"] = node + execution_id = f"{execution_id}/node/{self._source_nodes.index(node)}" + admitted.metadata[EXECUTION_ID_KEY] = execution_id + self._admitted_nodes[execution_id] = node + self._audiences[execution_id] = tuple(sorted(selected)) + self._order.append(execution_id) + for name in selected: + if node is not None: + pressure[name][node] += 1 + if self.remaining[name] is not None: + self.remaining[name] -= 1 + if node is not None: + self._remaining_by_node[name][node] -= 1 + self._epochs[name][execution_id] = self.source._dataset_epoch + self._pending_entries.append(admitted) + audience = audience - selected + while self._pending_entries and len(entries) < count: + entries.append(self._pending_entries.popleft()) return entries def push_inference_results(self, results): @@ -324,12 +413,19 @@ def push_inference_results(self, results): execution_id = self._order[0] output = self._completed[execution_id] descriptor = MooncakeP2PDescriptor.from_key(output.mooncake_key) + if self._source_nodes and ( + output.metadata.get("feature_affinity") != "node_local_required" + or output.metadata.get("source_node_ip") != self._admitted_nodes[execution_id] + or output.metadata.get("shared_source_node") != self._admitted_nodes[execution_id] + ): + raise RuntimeError("shared output violates its admitted source-node placement") audience = [n for n in self._audiences[execution_id] if n in self.active] if not audience: source_control(descriptor) self._order.popleft() del self._completed[execution_id] del self._audiences[execution_id] + self._admitted_nodes.pop(execution_id, None) continue views, readers, payloads = [], [], {} for name in audience: @@ -396,6 +492,7 @@ def push_inference_results(self, results): self._order.popleft() del self._completed[execution_id] del self._audiences[execution_id] + self._admitted_nodes.pop(execution_id, None) # Every registered right is tracked before publishing to any # controller, so a partial fanout failure cannot orphan readers. for name, sample in views: @@ -434,6 +531,32 @@ def get_ready_pool_size(self): default=0, ) + def get_ready_pool_size_by_node(self, execution_ids=None): + """Count the slowest reader, optionally restricted to an admission's audience.""" + with self._lock: + counts = dict.fromkeys(self._source_nodes, 0) + readers = self.active + if execution_ids is not None: + readers = set() + for identity in execution_ids: + if identity not in self._audiences: + raise ValueError("unknown shared source admission") + readers.update(self._audiences[identity]) + readers &= self.active + child_counts = { + name: self.experiments[name].get_ready_pool_size_by_node() + for name in readers + } + for name, ready in child_counts.items(): + for node in counts: + unpublished = sum( + name in self._audiences[key] + and output.metadata.get("source_node_ip") == node + for key, output in self._completed.items() + ) + counts[node] = max(counts[node], ready.get(node, 0) + unpublished) + return counts + def get_pool_size(self): with self._lock: return len( @@ -470,6 +593,8 @@ def request_cancel(self, name, reason="requested"): if name in self.active: self._cancelled[name] = str(reason) self.remaining[name] = 0 + if name in self._remaining_by_node: + self._remaining_by_node[name] = dict.fromkeys(self._source_nodes, 0) return {"experiment": name, "requested": name in self._cancelled} def cancellation_reason(self, name): @@ -486,6 +611,8 @@ def retire(self, name, *, owners_fenced: bool): return self.active.discard(name) self.remaining[name] = 0 + if name in self._remaining_by_node: + self._remaining_by_node[name] = dict.fromkeys(self._source_nodes, 0) for execution_id, descriptor in self._descriptors[name].items(): # Replaying the frozen subscription is idempotent whether the # first RPC was rejected or only its ACK was lost. Do not @@ -541,11 +668,14 @@ def finish_production(self, *, inflight_samples): del self._completed[execution_id] self._order.remove(execution_id) del self._audiences[execution_id] + self._admitted_nodes.pop(execution_id, None) self._discarded_unpublished += 1 self._abandoned_admissions += len(self._order) self._order.clear() + self._pending_entries.clear() self._completed.clear() self._audiences.clear() + self._admitted_nodes.clear() return self.status() def status(self): @@ -563,6 +693,7 @@ def status(self): "pending_source_subscriptions": len(self._pending_subscriptions), "peak_completed_unpublished_observations": self._peak_unpublished, "remaining_source_demand": dict(self.remaining), + "remaining_source_demand_by_node": copy.deepcopy(self._remaining_by_node), "error": self._error, "cancelled_experiments": dict(self._cancelled), "abandoned_admissions_without_output": self._abandoned_admissions, diff --git a/aspec/multi_draft.py b/aspec/multi_draft.py index 6ff79ca..8cf2f48 100644 --- a/aspec/multi_draft.py +++ b/aspec/multi_draft.py @@ -23,6 +23,7 @@ import time import uuid from concurrent.futures import ThreadPoolExecutor, as_completed +from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path @@ -38,13 +39,14 @@ from aspec.evidence import RunEvidenceRecorder from aspec.inference.factory import prepare_inference_engines from aspec.planning import ExecutionPlan, compile_execution_plan, save_execution_plan -from aspec.ray.owner_fence import _OwnerFenceProbe +from aspec.ray.cluster_preflight import run_cluster_preflight +from aspec.ray.owner_fence import SourceTeardownBlockedError, wait_for_owner_exit from aspec.ray.placement_group import allocate_train_group, create_placement_groups from aspec.reporting import _is_safe_experiment_name from aspec.training.trainer_actor import TrainerActor from aspec.transfer.mooncake.utils import launch_mooncake_master from aspec.utils.env import get_aspec_env_vars -from aspec.utils.logging import ExperimentTracking, logger +from aspec.utils.logging import ExperimentTracking @dataclass @@ -141,16 +143,21 @@ def compile_multi_draft_plan(config) -> MultiDraftPlan: if args.report_to not in {"none", "tensorboard", "wandb"}: raise ValueError("logging.report_to must be none, tensorboard or wandb") unsupported = { - "multi-node": args.cluster_mode != "single_node" or args.training_num_nodes != 1, "colocated GPU roles": args.colocate, "legacy sequence parallel": args.attention_backend == "usp", "decode": args.train_with_decode, "service profiling": bool(getattr(args, "topology_producer_profile_samples", 0)), + "finite source_work_budget": args.source_work_budget is not None, "Store data plane": args.mooncake_data_plane != "transfer_engine", } rejected = [key for key, value in unsupported.items() if value] if rejected: raise NotImplementedError(f"shared Target training does not support: {rejected}") + if args.cluster_mode == "static_multi_node": + if args.mooncake_feature_affinity != "node_local_required": + raise ValueError("shared multi-node training requires node_local_required feature affinity") + if args.inference_batch_size != 1: + raise NotImplementedError("shared multi-node training requires inference_batch_size=1") if child.training.max_seq_length != config.training.max_seq_length: raise ValueError("shared experiments require compatible source sequence lengths") for key in ("placement_strategy", "training_node_ips", "training_node_selectors"): @@ -174,13 +181,20 @@ def compile_multi_draft_plan(config) -> MultiDraftPlan: features = shared_feature_union(plan.workload.features for plan in plans.values()) heterogeneous = len({p.workload.features["contract_id"] for p in plans.values()}) > 1 source_plan = None - if heterogeneous or any(plan.to_runtime_args().draft_context_parallel_size > 1 for plan in plans.values()): + multi_node = config.cluster.mode == "static_multi_node" + if multi_node or heterogeneous or any(plan.to_runtime_args().draft_context_parallel_size > 1 for plan in plans.values()): source_config = load_config(base_config=OmegaConf.merge( next(iter(configs.values())), - {"training": {"draft_context_parallel_size": 1, "training_num_gpus_per_node": 1}}, + {"training": { + "draft_context_parallel_size": 1, + "training_num_gpus_per_node": ( + sum(child.training.training_num_gpus_per_node for child in configs.values()) + if multi_node else 1 + ), + "load_path": None, + }}, ), save_snapshot=False) - # This plan describes producer staging, not a phantom CP1 trainer - # with the first experiment\'s GPU count and DP dispatch demand. + # Size producer staging for the union of consumer layouts and GPU allocations. source_plan = compile_execution_plan(source_config, source_feature_contract=features) source_context = (source_plan or next(iter(plans.values()))).to_runtime_args().aspec_context_contract for child in plans.values(): @@ -190,6 +204,25 @@ def compile_multi_draft_plan(config) -> MultiDraftPlan: return MultiDraftPlan(configs, plans, int(source_seed), source_plan) +def experiment_gpu_allocation(training_pg, experiment_args, name): + """Select a child's node-major ranks from the shared node-major reservation.""" + pg, bundles, gpu_ids = training_pg + children = list(experiment_args) + nodes = int(getattr(experiment_args[name], "training_num_nodes", 1)) + widths = { + key: int(getattr(value, "training_num_gpus_per_node", value.world_size)) + for key, value in experiment_args.items() + } + stride = sum(widths.values()) + offset = sum(widths[key] for key in children[:children.index(name)]) + indexes = [ + node * stride + offset + rank + for node in range(nodes) + for rank in range(widths[name]) + ] + return pg, [bundles[i] for i in indexes], [gpu_ids[i] for i in indexes] + + class _ConsumerObservations: def __init__(self, statuses): self.statuses = statuses @@ -321,36 +354,11 @@ def experiment_contract(contract, name): def _fence_group(group, statuses, timeout=30): - # Single-node scope: the probe is pinned to the trainer's actual node, not - # the driver's placement or a hostname copied from a configuration. - from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy - - node_ip = statuses[0]["node_ip"] - node_id = next( - node["NodeID"] - for node in ray.nodes() - if node["Alive"] and node["NodeManagerAddress"] == node_ip - ) - probe = _OwnerFenceProbe.options( - scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) - ).remote() - deadline = time.monotonic() + timeout - try: - for actor in group._actor_handlers: - ray.kill(actor, no_restart=True) - while (remaining := deadline - time.monotonic()) > 0: - if ray.get( - probe.exited.remote([item["pid"] for item in statuses]), - timeout=min(5, remaining), - ): - return - time.sleep(0.05) - raise TimeoutError("experiment read owners did not exit; source withdrawal is unsafe") - finally: - try: - ray.kill(probe, no_restart=True) - except Exception: - logger.exception("Failed to close experiment owner probe") + for actor in group._actor_handlers: + ray.kill(actor, no_restart=True) + # Every node is probed using the initialized actor's immutable node identity. + # With shared readers, an unobserved remote exit keeps source ownership live. + return wait_for_owner_exit(statuses, timeout=timeout) def _finish_experiment(group, statuses, hub, name, *, graceful): @@ -366,6 +374,35 @@ def _finish_experiment(group, statuses, hub, name, *, graceful): return close_error +def preflight_multi_draft(args, plan): + """Validate every child and the total reservation before starting models.""" + from aspec.ray.node_local_files import close_stores + from aspec.train_entry import _prepare_run_directory + + if args.cluster_mode != "static_multi_node": + raise ValueError("shared --preflight-only requires static_multi_node") + run_id = f"uuid:{uuid.uuid4()}" + attempt_id = f"uuid:{uuid.uuid4()}" + experiments = {name: child.to_runtime_args() for name, child in plan.plans.items()} + source = copy.copy(args) + report = {"schema_version": "aspec.multi-draft-preflight/v1", "plan_id": plan.plan_id, + "state": "validated", "experiments": {}} + with ExitStack() as cleanup: + for child in [*experiments.values(), source]: + child.aspec_run_id, child.aspec_attempt_id = run_id, attempt_id + cleanup.callback(close_stores, child) + for name, child in experiments.items(): + _prepare_run_directory(child.output_dir) + report["experiments"][name] = run_cluster_preflight(child) + report["shared_source"] = run_cluster_preflight(source) + destination = Path(args.output_dir) / "cluster-preflight.json" + temporary = destination.with_suffix(".tmp") + temporary.write_text(json.dumps(report, indent=2) + "\n") + os.replace(temporary, destination) + print(json.dumps(report, indent=2)) + return report + + def run_multi_draft(args, plan: MultiDraftPlan): from aspec.train_entry import ( _get_draft_model_config, @@ -430,26 +467,32 @@ def persist_session(): _prepare_run_directory(child.output_dir) _save_config_snapshot(plan.configs[name]) save_execution_plan(child_plan, Path(child.output_dir) / "plan.json") - child.draft_model_config_obj = _get_draft_model_config(child) - _validate_and_configure_dflash(child, child.draft_model_config_obj) - experiment_args[name] = child recorders[name] = RunEvidenceRecorder( child_plan, child.output_dir, run_id=run_id, attempt_id=attempt_id ) + experiment_args[name] = child + preflight = run_cluster_preflight(child) + if preflight is not None: + recorders[name].record_cluster_preflight(preflight) + session.setdefault("cluster_preflight", {})[name] = preflight + child.draft_model_config_obj = _get_draft_model_config(child) + _validate_and_configure_dflash(child, child.draft_model_config_obj) - # One immutable full-layout source; each experiment chooses CP views. - # No duplicate Target forward or full-feature destination copy. + # Capture full-layout features; each experiment reads its own CP view. if any(getattr(child, "draft_context_parallel_size", 1) > 1 for child in experiment_args.values()): source_args.draft_context_parallel_size = 1 source_args.sp_size = 1 source_args.training_num_gpus_per_node = sum( - child.world_size for child in experiment_args.values() + int(getattr(child, "training_num_gpus_per_node", child.world_size)) + for child in experiment_args.values() ) - source_args.world_size = source_args.training_num_gpus_per_node + source_args.world_size = sum(child.world_size for child in experiment_args.values()) + if getattr(source_args, "cluster_mode", "single_node") == "static_multi_node": + session["source_cluster_preflight"] = run_cluster_preflight(source_args) source_args.num_train_steps = None # demand comes from the per-experiment budgets source_args.aspec_shared_source = True hub = SharedSourceController.options( - runtime_env={"env_vars": get_aspec_env_vars()}, + runtime_env={"env_vars": get_aspec_env_vars(source_args)}, name=f"shared-source-{run_id.removeprefix('uuid:')}", namespace="aspec", ).remote(source_args, experiment_args) @@ -464,30 +507,29 @@ def persist_session(): for child in experiment_args.values(): auto_calculate_training_steps(child, dataset_size) pgs = create_placement_groups(source_args) - session["gpu_node_id"] = require_single_gpu_node(pgs) + if getattr(source_args, "cluster_mode", "single_node") == "single_node": + session["gpu_node_id"] = require_single_gpu_node(pgs) launch_mooncake_master(source_args) mooncake = build_mooncake_config(source_args) source_args.draft_model_config_obj = next( iter(experiment_args.values()) ).draft_model_config_obj _validate_and_configure_dflash(source_args, source_args.draft_model_config_obj) - training_pg, bundles, gpu_ids = pgs["training"] - offset = 0 + init_refs = {} for name, child in experiment_args.items(): width = child.world_size allocation_in_progress = True group = allocate_train_group( child, - 1, - width, - (training_pg, bundles[offset : offset + width], gpu_ids[offset : offset + width]), + int(getattr(child, "training_num_nodes", 1)), + int(getattr(child, "training_num_gpus_per_node", width)), + experiment_gpu_allocation(pgs["training"], experiment_args, name), TrainerActor, ) groups[name] = group allocation_in_progress = False actor_statuses[name] = group.get_status() - offset += width destination = copy.deepcopy(mooncake) destination.gpu_buffer_size = build_mooncake_config(child).gpu_buffer_size destination_configs[name] = destination @@ -557,12 +599,17 @@ def persist_session(): if not running_names: raise RuntimeError("all Draft experiments failed to initialize") manager = AsyncInferenceManager.options( - runtime_env={"env_vars": get_aspec_env_vars()} + runtime_env={"env_vars": get_aspec_env_vars(source_args)} ).remote( source_args, hub, inference_engines=engines, max_concurrent_batches=source_args.max_concurrent_batches, + inference_engine_nodes=( + [status["node_ip"] for status in producer_placements] + if getattr(source_args, "cluster_mode", "single_node") == "static_multi_node" + else None + ), ) manager_future = manager.run.remote() training_started = True # Eval capture can already publish registered source readers. @@ -806,6 +853,10 @@ def fence(name=name, group=group): if hub is not None: cleanup(lambda: ray.get(hub.close_queues.remote(), timeout=10)) cleanup(lambda: ray.kill(hub, no_restart=True)) + from aspec.ray.node_local_files import close_stores + + for child in [source_args, *experiment_args.values()]: + cleanup(lambda child=child: close_stores(child)) unique_pgs = {item[0].id: item[0] for item in pgs.values()} for pg in unique_pgs.values(): cleanup(lambda pg=pg: ray.util.remove_placement_group(pg)) @@ -822,6 +873,11 @@ def fence(name=name, group=group): session["budget_gpu_hours"] = allocated_gpus * session["wall_seconds"] / 3600 try: persist_session() + if training_started and not owners_fenced: + raise SourceTeardownBlockedError([ + status for name, statuses in actor_statuses.items() + if name not in fenced_groups for status in statuses + ]) finally: for key, value in previous_environment.items(): if value is None: diff --git a/aspec/train_entry.py b/aspec/train_entry.py index cf0f852..c5bba3f 100644 --- a/aspec/train_entry.py +++ b/aspec/train_entry.py @@ -223,8 +223,6 @@ def _prepare_config() -> tuple[argparse.Namespace, ExecutionPlan]: if config.experiments: from aspec.multi_draft import compile_multi_draft_plan - if args.preflight_only: - parser.error("shared Target training does not support --preflight-only") plan = compile_multi_draft_plan(config) else: plan = compile_execution_plan( @@ -528,16 +526,19 @@ def main() -> None: def _main() -> None: args, plan = _prepare_config() - from aspec.multi_draft import MultiDraftPlan, run_multi_draft + from aspec.multi_draft import MultiDraftPlan, preflight_multi_draft, run_multi_draft if isinstance(plan, MultiDraftPlan): try: with termination_signal_handlers(): ensure_ray_initialized() install_termination_signal_handler() - run_multi_draft(args, plan) + if getattr(args, "preflight_only", False) is True: + preflight_multi_draft(args, plan) + else: + run_multi_draft(args, plan) finally: - if ray.is_initialized(): + if not isinstance(sys.exception(), SourceTeardownBlockedError) and ray.is_initialized(): ray.shutdown() return # ``preflight_only`` is an argparse boolean. Requiring the literal value diff --git a/configs/recipes/README.md b/configs/recipes/README.md index 56e6628..fcd1755 100644 --- a/configs/recipes/README.md +++ b/configs/recipes/README.md @@ -11,6 +11,8 @@ whose supported export boundary is native-only; `false` does not disable the nat Use each recipe's documented inputs and check capacity on your own hardware. The [shared Target walkthrough](../../examples/qwen3-8b-shared-target/README.md) uses the recipe for experiment comparison, separate curves and recovery from different checkpoints. +The `multinode-shared` workflow has a complete +[startup and independent recovery example](../../docs/user-guide/sharing.md#start-two-drafts-and-resume-from-different-checkpoints). Supplemental variants belong under `configs/examples/`; benchmark comparisons remain under `configs/benchmarks/`. diff --git a/configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml b/configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml new file mode 100644 index 0000000..2beaf39 --- /dev/null +++ b/configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml @@ -0,0 +1,155 @@ +# MIT License +# +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# Adapted from the TorchSpec configuration examples at 6c042a87140a84d13839e341ece2c5c3ada918bc. +# ASpec updates schema bindings, portable inputs and execution settings. + +# Two-node Qwen3-8B / DSpark 4K: each node uses 4 Target + 2 Draft A + 2 Draft B GPUs. +# Set node IPs and model/data paths as described in docs/user-guide/multinode.md. +# Target features are required to stay on their producer node and use NVLink; +# Draft gradients still synchronize across the independent cross-node DP groups. + +model: + target_model_path: ${oc.env:ASPEC_TARGET_MODEL} + native_context_length: 32768 + trust_remote_code: false + draft_model_config: ${oc.env:ASPEC_DRAFT_CONFIG} + +dataset: + eval_interval: 100 + eval_micro_batch_size: 1 + train_data_path: ${oc.env:ASPEC_TRAIN_DATA} + eval_data_path: ${oc.env:ASPEC_EVAL_DATA} + chat_template: qwen + shuffle_dataset: true + min_loss_tokens: 32 + +training: + attention_backend: cudnn_flex_attention + draft_context_parallel_size: 1 + micro_batch_size: 1 + draft_accumulation_steps: 64 + learning_rate: 6e-4 + min_lr: 6e-5 + weight_decay: 0.0 + max_concurrent_batches: 1 + streaming_dispatch_with_training: true + pipeline_dispatch_with_training: true + max_grad_norm: 1.0 + max_seq_length: 4096 + num_epochs: 3 + seed: 42 + training_num_gpus_per_node: 2 + training_num_nodes: 2 + placement_strategy: custom + training_node_ips: + - ${oc.env:ASPEC_NODE_0_IP} + - ${oc.env:ASPEC_NODE_1_IP} + # Optional per-run network overrides: training.train_env_vars JSON. + # With no override NCCL/Gloo select the route from the live environment. + ttt_length: 7 + fsdp_strategy: REPLICATE + fsdp_reduce_dtype: bfloat16 + prefetch_depth: 1 + save_interval: 100 + save_per_epoch: false + save_final_checkpoint: true + warmup_ratio: 0.04 + dflash_block_size: 7 + dflash_packing: true + dspark_num_anchors: 512 + dspark_num_target_layers: 5 + dspark_loss_decay_gamma: 4.0 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + +inference: + inference_engine_type: sgl + store_last_hidden_states: true + target_prefill_memory_optimized: true + inference_num_gpus: 8 + inference_num_gpus_per_engine: 1 + inference_num_gpus_per_node: 4 + inference_node_ips: + - ${oc.env:ASPEC_NODE_0_IP} + - ${oc.env:ASPEC_NODE_1_IP} + inference_fetch_batch: 16 + inference_buffer_threshold: 16 + max_sample_pool_size: 16 + inference_batch_size: 1 + sglang: + tp_size: 1 + mem_fraction_static: 0.5 + extra_args: + max_running_requests: 2 + max_total_tokens: 8704 + max_prefill_tokens: 8704 + +mooncake: + master_server_address: null + metadata_server: null + data_plane: transfer_engine + protocol: nvlink_intra + feature_affinity: node_local_required + require_nvlink: true + enable_gpu_direct: true + allow_host_fallback: false + gpu_ring_slots: 2 + # Retain a complete ready window until the Draft consumers finish their reads. + source_gpu_ring_slots: 16 + destination_gpu_ring_slots: 1 + global_segment_size: 4GB + local_buffer_size: 512MB + max_seq_len: 4096 + hidden_dim: 4096 + kv_lease_ttl_s: 60 + get_retry_max_wait_seconds: 120 + +cluster: + mode: static_multi_node + storage_mode: node_local + preflight_timeout_seconds: 900 + file_transfer_timeout_seconds: 600 + require_homogeneous_gpus: true + +logging: + report_to: none + +debug: + enable_perf_metrics: true + performance_measurement_required: false + performance_warmup_steps: 2 + performance_drain_steps: 1 + +experiments: + - name: draft_a + training: + seed: 42 + - name: draft_b + training: + seed: 43 + learning_rate: 3e-4 + min_lr: 3e-5 + +output_dir: ${oc.env:ASPEC_OUTPUT_DIR} +cache_dir: ${oc.env:ASPEC_CACHE_DIR} diff --git a/docs/aspec.zh-CN.md b/docs/aspec.zh-CN.md index 9315c3b..d04468a 100644 --- a/docs/aspec.zh-CN.md +++ b/docs/aspec.zh-CN.md @@ -122,7 +122,7 @@ ASpec 允许多个独立 Draft 共享兼容的 Target 特征。不同实验需 ![ASpec：异步生成与传输 Target 特征，供多个独立 Draft 实验使用；各实验保留自己的训练状态。](assets/architecture.svg) -系统会限制待消费的数据量，避免较慢的实验让特征无限积压。一个实验结束或退出后，其他实验可以继续运行；恢复时，每个实验从各自已提交的训练进度继续。当前多 Draft 共享模式支持单机，配置与恢复方式见[使用指南](user-guide/sharing.md)。 +系统会限制待消费的数据量，避免较慢的实验让特征无限积压。一个实验结束或退出后，其他实验可以继续运行；恢复时，每个实验从各自已提交的训练进度继续。多 Draft 共享模式支持单机和静态多机训练，各 Draft 可独立启用 packing。配置与恢复方式见[使用指南](user-guide/sharing.md)。 ## 接下来 diff --git a/docs/contributing/validation.md b/docs/contributing/validation.md index f8041ed..df419b0 100644 --- a/docs/contributing/validation.md +++ b/docs/contributing/validation.md @@ -15,7 +15,7 @@ python -m tools.check_docs python -m ruff check aspec tests tools ``` -`contracts` checks all 14 configs/13 workflows, resource budgets and training defaults, negative +`contracts` checks all 15 configs/14 workflows, resource budgets and training defaults, negative acceptance cases, and existing algorithm/packing/CP/sharing/multi-node contracts. Model names, data paths and node IPs in CPU tests are fixtures, never real GPU inputs. No models are downloaded by this command. CPU contracts do not establish GPU kernel parity or available capacity. @@ -152,6 +152,17 @@ exit; it does not stop the cluster. Input replicas are hashed, checkpoint shards collected, and resume state is explicitly distributed. This test needs **two compatible nodes, eight GPUs each**; it cannot be replaced by single-node file-exchange tests. +For two independent Drafts sharing the Target, use the same cluster and replicated inputs: + +```bash +bash tools/run_in_env.sh python -m tools.recipes accept multinode-shared \ + --steps 64 --output /data/aspec/acceptance/multinode-shared-001 +``` + +The launcher checks each Draft's training, checkpoint recovery and native export separately. +The [shared training guide](../user-guide/sharing.md#start-two-drafts-and-resume-from-different-checkpoints) +also shows how to resume experiments from different optimizer steps. + ## Whole worker-node loss and recovery Use a dedicated cluster and an external infrastructure controller. Keep the coordinator and its diff --git a/docs/framework/ARCHITECTURE.md b/docs/framework/ARCHITECTURE.md index 97f57e7..3775b5c 100644 --- a/docs/framework/ARCHITECTURE.md +++ b/docs/framework/ARCHITECTURE.md @@ -75,7 +75,13 @@ These identities are not interchangeable. [Supported paths](../support.md) descr Each shared Draft owns weights, optimizer, schedule, RNG, committed work, evaluation and checkpoints. There is no shared optimizer-step barrier. Slow readers apply bounded backpressure; safely fenced -failed readers need not stop healthy siblings. Shared multi-node training is outside the current path. +failed readers need not stop healthy siblings. Static multi-node sharing keeps features node-local +and exchanges checkpoints per Draft. Each Draft packs independently within each node; Target requests +capture one original document at a time. Pending captures reserve each reader’s node-local ready credits. + +The shared source tracks each Draft's remaining demand per node. Compatible readers on the same node +share a capture. Different restored frontiers can require a document to be captured on separate nodes +for disjoint readers; physical capture identity remains separate from each Draft's committed work. Exact resume restores all required training state and identities. Weights-only initialization starts a new phase. Export preserves actual algorithm/structure/stage, vocabulary mapping, frozen Target diff --git a/docs/framework/DECISIONS.md b/docs/framework/DECISIONS.md index f41f89e..2e7ed90 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -32,8 +32,9 @@ These decisions describe the maintained implementation and its public behavior. ASpec-managed artifacts. Exact resume uses a new run directory and a complete checkpoint. Unrelated files do not force rejection; implicit overwrite is unsupported. - **Sharing and clusters.** Compatible Target capture can serve independently progressing Drafts. - Shared mode is single-node. Multi-node single-Draft mode uses static homogeneous membership, - node-local model parallelism and fail-stop attempt recovery; node-local feature affinity is a + Shared mode supports single-node and static multi-node execution with independent packing. + Multi-node Target requests capture one original document at a time. Multi-node execution uses static + homogeneous membership, node-local model parallelism and checkpoint-based restart; node-local feature affinity is a required capability, with no remote feature substitution. - **Algorithms and assets.** EAGLE3/4, DFlash2 and native MTP retain their own objectives and state identities. CSA/QSA/DSA stages use explicit weights-only transitions. DSA uses per-query raw-token @@ -52,9 +53,10 @@ These decisions describe the maintained implementation and its public behavior. Public behavior changes require corresponding guide/config updates and executable verification. - **Training recipes.** Maintain the [recipe catalog](../user-guide/recipes.md), including - Ling tiny/flash 3.0/flash 3.1, Native MTP, EAGLE3/4, one Qwen3.8-27B long-context CP2 lane and one - two-node node-local lane. Separate full training defaults, bounded validation overrides and - supplementary examples. Validate training, recovery and asset behavior when changing a recipe. + Ling tiny/flash 3.0/flash 3.1, Native MTP, EAGLE3/4, one Qwen3.8-27B long-context CP2 lane and + two-node node-local lanes for single and multiple Drafts. Separate full training defaults, + bounded validation overrides and supplementary examples. Validate training, recovery and asset + behavior when changing a recipe. - **Multi-node portability.** The release recipe requires no shared filesystem. `cluster.storage_mode=node_local` checks input replicas and exchanges checkpoint shards explicitly; `shared` retains the older optional diff --git a/docs/framework/ROADMAP.md b/docs/framework/ROADMAP.md index f973e4d..5dbbbba 100644 --- a/docs/framework/ROADMAP.md +++ b/docs/framework/ROADMAP.md @@ -7,7 +7,7 @@ current combinations and limits. - Preserve training objectives, document boundaries, gradients and frozen Target bindings across optimized backends. Use numerical comparisons before interpreting performance. - Extend reliability beyond whole worker-process loss to additional concrete failure mechanisms; - keep shared multi-node execution and elastic membership explicit roadmap items. + extend shared multi-node workload coverage and keep elastic membership an explicit roadmap item. - Keep checkpoint recovery, native assets and downstream formats faithful as algorithms evolve. - Compare normal complete workflows using [benchmark methodology](../benchmarks/methodology.md). Maintain reusable [validation tools](../contributing/validation.md) without making historical diff --git a/docs/support.md b/docs/support.md index 58bd775..10be4b7 100644 --- a/docs/support.md +++ b/docs/support.md @@ -1,9 +1,9 @@ # Supported paths and limitations ASpec is a source release for trusted Linux GPU training environments. The primary performance -path is SGLang + DSpark/DFlash + Qwen3 on multi-GPU systems. The -[installation guide](user-guide/installation.md) defines the tested dependency lane. CPU planning -and tests do not imply CPU training support. +path is SGLang + DSpark/DFlash + Qwen3 on multi-GPU systems. Training requires CUDA GPUs; +configuration planning can run on a CPU. See the [installation guide](user-guide/installation.md) +for the supported dependencies. ## Algorithms and structures @@ -29,12 +29,13 @@ passing Plan validation establishes legality, not performance, convergence or do With packing enabled, DSpark/DFlash full and sliding attention preserve document semantics with FA4 or native cuDNN flattened-token execution at CP1 and above. Selected-entry CSA/QSA/DSA packing/CP remains FA4-only. -- Shared Target capture supports independent single-node Drafts, progress, packing/CP views, - evaluation, recovery and tracking. Shared multi-node execution and online weight synchronization - are outside scope. +- Shared Target capture supports independent Drafts, progress, evaluation, recovery and tracking. + Both single-node and static homogeneous multi-node sharing support independent packing/CP views. + Multi-node sharing requires one original document per Target request and node-local feature affinity. + Draft weights and optimizer state remain independent; online weight synchronization is unsupported. - Static homogeneous multi-node single-Draft training supports node-local model parallelism, explicit transports and required node-local feature affinity. Elastic membership and cross-node - Target TP are outside scope. Preflight and throughput checks do not certify all failure modes. + Target TP are unsupported. - Offline feature replay is single-node, single-Draft, CP1 with CPU Store/TCP transport. - SGLang export is accepted only when the actual structure can be faithfully represented. CSA/QSA/DSA, materialized MLA and EAGLE4 use native assets where required; strict loading/proposal components @@ -44,21 +45,17 @@ passing Plan validation establishes legality, not performance, convergence or do Measured training-loss and verifier-acceptance results are available in the [Qwen3-8B DSpark comparison](benchmarks/dspark-alignment.md) and -[CP comparison](benchmarks/dspark-cp-alignment.md). Their configurations and -validation scope are recorded in those reports. +[CP comparison](benchmarks/dspark-cp-alignment.md). Each report lists its hardware, +runtime settings and workloads. -- Short functional training and exact-recovery tests do not certify long-run model quality, every - hardware/backend combination, or CUDA driver/DMA/RDMA hardware-hang recovery. -- Static node-local multi-node acceptance covers whole worker Ray/GPU process loss, fencing and - exact restart with a new node/attempt identity. Host power loss, arbitrary partitions and - unreturning native/driver operations remain unverified; see [fault validation](contributing/validation.md#whole-worker-node-loss-and-recovery). -- Some short runs do not observe a requested pipeline event and correctly report - `incomplete_evidence`; absence must not be relabeled successful activation. -- Repeated SGLang engine startup/shutdown has produced upstream `resource_tracker` warnings. - No GPU-process residue in a bounded check does not certify every OS shared-memory resource. -- Original benchmark environments remain attached to their results. They do not automatically - certify the current dependency lane. [Benchmark guidance](user-guide/benchmarks.md) distinguishes - input capacity, sampled supervision and quality. +- Automatic recovery from CUDA driver, DMA or RDMA hangs is unsupported. +- In single-Draft static multi-node training, worker loss stops the current attempt. Fence the old + processes, restore the worker and resume a new attempt from a complete checkpoint; see + [worker recovery](contributing/validation.md#whole-worker-node-loss-and-recovery). +- Reports use `incomplete_evidence` when a run ends before required pipeline activity is observed. +- SGLang may emit `resource_tracker` warnings during repeated engine startup and shutdown. +- Benchmark results apply to the software versions and settings listed in each report. See + [benchmark guidance](user-guide/benchmarks.md) for throughput and quality measurements. ### Decode reproducibility diff --git a/docs/user-guide/multinode.md b/docs/user-guide/multinode.md index 9cf79d5..687916b 100644 --- a/docs/user-guide/multinode.md +++ b/docs/user-guide/multinode.md @@ -3,7 +3,8 @@ The node-local recipe uses **independent local storage**, not a cross-node shared filesystem. It places 4 Target + 4 Draft GPUs on each of two compatible eight-GPU nodes. Features remain on their production node over NVLink; Draft gradients synchronize across nodes. Use a dedicated Ray cluster on a trusted -private network. Shared multi-Draft execution remains single-node. +private network. For independent Draft experiments sharing Target features, use the +[multi-Draft recipe](sharing.md#static-multi-node-sharing). Unused-output disposal follows the same locality rule: the runtime selects a consumer on the producer's node and rejects missing or unmatched placement metadata. diff --git a/docs/user-guide/recipes.md b/docs/user-guide/recipes.md index d16df31..693eb5b 100644 --- a/docs/user-guide/recipes.md +++ b/docs/user-guide/recipes.md @@ -27,6 +27,7 @@ connections on your hardware before training. | `shared-target` | Qwen3-8B / two independent DSpark Drafts | 1 node, 8 GPUs: 4T + 2D × 2; 4K | Compare independent Drafts sharing Target features | | `qwen27b-long` | Qwen3.8-27B / DSpark | 1 node, 8 GPUs: 3 × TP2 Target + 2D (CP2); 256K-class | Distribute long-context training with CP | | `multinode` | Qwen3-8B / DSpark | 2 nodes × 8 GPUs; each 4T + 4D; 4K | Keep features node-local and synchronize gradients across nodes | +| `multinode-shared` | Qwen3-8B / two independent DSpark Drafts | 2 nodes × 8 GPUs; each 4T + 2D × 2; 4K | Share node-local Target features across independently progressing cross-node Drafts | The [27B recipe](../../configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml) uses a 261,888-token cap, native cuDNN THD attention and CP2. Each TP2 Target replica admits one request, @@ -84,6 +85,10 @@ output using `--set training.load_path=/absolute/path/to/checkpoints`; see [reco The optional [runtime validation command](../contributing/validation.md) runs a short training, recovery and export check without changing normal training defaults. +For `multinode-shared`, follow the complete +[two-Draft start and independent resume example](sharing.md#start-two-drafts-and-resume-from-different-checkpoints). +It preserves the full experiment list and each Draft's independent checkpoint and LR schedule. + ## DSpark performance settings The `dspark` recipe is a general FA4/4T+4D training starting point. For online DSpark full attention, diff --git a/docs/user-guide/sharing.md b/docs/user-guide/sharing.md index d0158b0..94f32d8 100644 --- a/docs/user-guide/sharing.md +++ b/docs/user-guide/sharing.md @@ -1,8 +1,8 @@ # Share a Target across Drafts Shared mode captures compatible Target features once and trains independent Draft experiments. -Each experiment owns its weights, optimizer, LR, RNG, evaluation, work frontier, checkpoint and -tracking output. Shared execution is currently single-node. +Each experiment owns its weights, optimizer, LR, RNG, evaluation, training progress, checkpoint and +tracking output. Shared execution supports one node or a static homogeneous multi-node cluster. The [comparison and independent recovery walkthrough](../../examples/qwen3-8b-shared-target/README.md) derives two learning-rate experiments from the shared-target recipe, records separate @@ -16,14 +16,16 @@ training state and progress rather than participating in one Draft's data-parall | Example | Purpose | Starting GPU budget | | --- | --- | --- | | `configs/recipes/qwen3_8b_4k_multi_draft_nvlink.yaml` | Two independent DSpark experiments | 4 Target + 4 Draft | +| `configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml` | Two independent cross-node DSpark experiments | 2 nodes, each 4 Target + 2 GPUs per Draft | | `configs/examples/qwen3_8b_multi_algorithm_feature_views.yaml` | Independent DSpark/EAGLE3/DFlash feature views and packing/CP | 1 Target + 7 Draft | | `configs/examples/ling3_flash_multi_structure_dspark.yaml` | Dense, Qwen MoE and Ling MoE Drafts | TP2 Target + 3 Draft | | `configs/examples/ling3_flash_swa_compare_dspark.yaml` | Full context versus SWA | TP2 Target + 2 Draft | -Supply training and held-out evaluation data and inspect all child Plans before starting. The Qwen 4K -shared template uses `Qwen/Qwen3-8B` by default; set `ASPEC_TARGET_MODEL` to use a local copy, -and provide `ASPEC_TRAIN_DATA` and `ASPEC_EVAL_DATA`. The Ling functional example below uses -bundled smoke data: +Supply training and held-out evaluation data and inspect each Draft's Plan before starting. +The single-node Qwen 4K recipe uses `Qwen/Qwen3-8B` by default; set `ASPEC_TARGET_MODEL` +to use a local copy, and provide `ASPEC_TRAIN_DATA` and `ASPEC_EVAL_DATA`. +The multi-node recipe requires local model and data replicas, as described below. +The Ling example uses bundled sample conversations: ```bash aspec --config configs/examples/ling3_flash_multi_structure_dspark.yaml --dry-run \ @@ -47,15 +49,23 @@ each child Plan still validates its own execution requirements. ## Progress, resources and failure -Experiments progress independently without a step barrier. Slow consumers apply bounded -backpressure; completed, cancelled or failed readers have their rights fenced before reclamation. -Safe healthy siblings can continue, but a failed/cancelled experiment makes the session exit nonzero. +Experiments progress independently without a step barrier. Slow experiments limit feature production +to keep queues bounded. When an experiment finishes, is cancelled or fails, its trainer processes +stop reading before shared features are released. Other experiments can continue, but a failed or +cancelled experiment makes the session exit nonzero. Reports distinguish shared source counts and each experiment's committed work. Do not add child consumption counts and label the sum unique Target production. -Check each child's required capabilities, fallback, lifecycle and evaluation. Per-child throughput -is coupled to shared supply and backpressure, not isolated compute capacity. Root wall time and GPU -budget describe the whole session; different objectives/auxiliary losses do not form a quality ranking. +Feature and transport status is finalized when the shared session ends. Check the session and +experiment reports after all experiments finish. + +Use each Draft's report to inspect training progress, evaluation and resource use. Its throughput +depends on shared Target supply and queue pressure. Root wall time and GPU budget describe the +whole session. Compare experiments using the same held-out evaluation metric rather than their +algorithm-specific training losses. + +Set `training.num_train_steps` for each Draft's update limit. Shared mode supplies documents until +those limits are reached. `training.source_work_budget` is available only for single-Draft runs. ## Resume and tracking @@ -69,6 +79,126 @@ for every child, keep the full `lr_total_steps`, choose a fresh root output directory, and set `num_train_steps` to the intended total frontier. Children may resume from different committed frontiers. One child's progress never licenses skipping another child's work. +## Static multi-node sharing + +Prepare each node and start a task-owned cluster using [multi-node setup](multinode.md). +The multi-Draft recipe uses the same portable model/data/output environment variables and +independent local filesystems. Run it through the normal configuration entry point: + +```bash +aspec --config configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml --dry-run +aspec --config configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml +``` + +All experiments use the same ordered training nodes and Target allocation. Each experiment's +`training_num_gpus_per_node` reserves its own GPUs on every node. GPU budgets add across experiments; +the Target allocation is shared. Each Draft synchronizes gradients within its own cross-node group. +Model-parallel groups and required NVLink feature reads stay node-local. + +Each Draft can enable `training.dflash_packing` independently. Packed rows contain complete documents +from the same producer node, with document boundaries and local positions preserved. The Target still +captures one original document per request (`inference_batch_size=1`); packing happens in each Draft’s +controller. Supply follows each Draft’s remaining optimizer steps and node-local queue space. + +Multi-node sharing requires TransferEngine and `feature_affinity=node_local_required`. +The normal launch preflights each child and the aggregate allocation. Use `--preflight-only` +with a fresh output directory for a separate check. +The `multinode-shared` recipe launcher also provides a complete training, recovery and export workflow. + +### Start two Drafts and resume from different checkpoints + +Run the following from the coordinator's installed checkout after completing +[node preparation and cluster startup](multinode.md#prepare-each-node-independently). +Keep the owned cluster running across both segments. Set the input environment from +[multi-node configuration](multinode.md#configure-preflight-and-train), including +`ASPEC_NODE_0_IP`, `ASPEC_NODE_1_IP`, `ASPEC_TARGET_MODEL`, `ASPEC_DRAFT_CONFIG`, +`ASPEC_TRAIN_DATA`, `ASPEC_EVAL_DATA`, `ASPEC_CACHE_DIR` and an explicit `RAY_ADDRESS`. +Model, data and Draft JSON replicas use the same absolute logical paths on both nodes. +Reserve eight GPUs per node: four Target GPUs, two for Draft A and two for Draft B. + +This example gives A an initial limit of 32 optimizer steps and B a limit of 33, with a full +LR horizon of 64 from the beginning. The recipe's accumulation of 64 remains in effect: +with DP4 and micro-batch 1, each Draft's global batch is 256 packed rows. The number of original +documents per row varies with sequence length. Supply enough usable supervised documents for +complete optimizer batches. The source repeats the dataset as needed to finish the configured +optimizer steps. Without an explicit step limit, the horizon is derived from `num_epochs` and +the dataset size divided by the global batch; with packing, this is not a fixed number of +original-document passes. + +Create a complete first-segment configuration without editing the recipe: + +```bash +export ASPEC_OUTPUT_DIR=/data/aspec/runs/compare-first +python - <<'PY' +from pathlib import Path +from omegaconf import OmegaConf + +config = OmegaConf.load("configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml") +config.training.num_train_steps = 33 +config.training.lr_total_steps = 64 +config.training.continual_training = False +config.training.continue_source_cursor = False +config.training.save_interval = 16 +config.dataset.eval_interval = 16 +config.logging.report_to = "tensorboard" +limits = {"draft_a": 32, "draft_b": 33} +for experiment in config.experiments: + experiment.training.num_train_steps = limits[experiment.name] +path = Path("run-configs/compare-first.yaml") +path.parent.mkdir(parents=True, exist_ok=True) +with path.open("x") as stream: + OmegaConf.save(config, stream, resolve=True) +PY +bash tools/run_in_env.sh aspec --config run-configs/compare-first.yaml --dry-run +bash tools/run_in_env.sh aspec --config run-configs/compare-first.yaml +aspec --report /data/aspec/runs/compare-first +``` + +Wait for both experiments and the shared session to finish successfully. On the coordinator, +`compare-first/draft_a/checkpoints` and `compare-first/draft_b/checkpoints` then hold the independently +published complete checkpoints. Derive the resumed configuration from the session's saved +`config.yaml`, preserving every experiment entry: + +```bash +python - <<'PY' +from pathlib import Path +from omegaconf import OmegaConf + +first = Path("/data/aspec/runs/compare-first") +resumed = Path("/data/aspec/runs/compare-resumed") +config = OmegaConf.load(first / "config.yaml") +assert config.training.lr_total_steps == 64 +config.output_dir = str(resumed) +config.training.num_train_steps = 64 +config.training.continual_training = False +config.training.continue_source_cursor = False +for experiment in config.experiments: + experiment.training.load_path = str(first / experiment.name / "checkpoints") + experiment.training.num_train_steps = 64 + experiment.training.continual_training = False + experiment.training.continue_source_cursor = False +with Path("run-configs/compare-resumed.yaml").open("x") as stream: + OmegaConf.save(config, stream, resolve=True) +PY +bash tools/run_in_env.sh aspec --config run-configs/compare-resumed.yaml --dry-run +bash tools/run_in_env.sh aspec --config run-configs/compare-resumed.yaml +aspec --report /data/aspec/runs/compare-resumed +aspec --report /data/aspec/runs/compare-resumed/draft_a +aspec --report /data/aspec/runs/compare-resumed/draft_b +``` + +The new YAML contains its own absolute output directory. Both run directories and generated YAML +filenames must be fresh. A resumes at 32 and performs 32 more updates; B resumes at 33 and performs +31 more. Their restored LR schedules retain the original horizon of 64. The source can apply +backpressure while supplying work needed by the less advanced reader; each reader preserves its +own committed progress. Use the root report for session status and each Draft's report for its +training and evaluation results. Each Draft has its own checkpoint to +[export and evaluate](evaluation.md#export); the session root is not a Draft checkpoint. + +Before releasing shared features, the runtime verifies that the stopped Draft's trainer processes +have exited on every node. If a node becomes unreachable, keep the driver and source running until +those processes have been stopped. + TensorBoard and W&B use separate runs/writers per experiment. See -[training and tracking](training.md#tensorboard-and-wb). Online weight synchronization, multi-node -shared execution and arbitrary multi-tenant isolation are outside the current contract. +[training and tracking](training.md#tensorboard-and-wb). Online weight synchronization, elastic membership +and isolation between untrusted tenants are unsupported. diff --git a/tests/test_inference_manager_lifecycle.py b/tests/cpu/test_inference_manager_lifecycle.py similarity index 94% rename from tests/test_inference_manager_lifecycle.py rename to tests/cpu/test_inference_manager_lifecycle.py index 3303115..b346781 100644 --- a/tests/test_inference_manager_lifecycle.py +++ b/tests/cpu/test_inference_manager_lifecycle.py @@ -59,6 +59,24 @@ def _manager(): return manager +def test_shared_multinode_refill_uses_admission_audience_instead_of_global_maximum(): + async def scenario(): + manager = _manager() + manager.args = SimpleNamespace(aspec_shared_source=True) + manager._sample_budget, manager._scheduled_samples = None, 0 + manager._prompt_buffer = deque() + manager._buffer_low_watermark = manager._refill_size = 4 + manager._max_pool_size = 4 + manager._node_ready_watermarks = {'a': 2, 'b': 2} + manager.controller.get_ready_pool_size_by_node = _AsyncRemote({'a': 2, 'b': 2}) + manager.controller.get_prompts = _AsyncRemote([InferenceInput('row')]) + await manager._refill_buffer() + assert list(manager._prompt_buffer)[0].data_id == 'row' + assert manager.controller.get_prompts.calls == [((4,), {})] + assert manager.controller.get_ready_pool_size_by_node.calls == [] + asyncio.run(scenario()) + + def test_forward_results_preserves_success_prefix_then_fails_closed(): async def scenario(): manager = _manager() @@ -336,7 +354,8 @@ async def scenario(): asyncio.run(scenario()) -def test_node_local_admission_skips_saturated_node_and_preserves_healthy_progress(): +@pytest.mark.parametrize("shared", [False, True]) +def test_node_local_admission_skips_saturated_node_and_preserves_healthy_progress(shared): async def scenario(): manager = _manager() manager._max_pool_size = 8 @@ -377,7 +396,11 @@ async def scenario(): ] ) - async def ready_by_node(): + async def ready_by_node(execution_ids=None): + if shared: + assert execution_ids == [manager._prompt_buffer[0].metadata["aspec_execution_id"]] + else: + assert execution_ids is None return {"node-a": 4, "node-b": 0} manager.controller.get_ready_pool_size_by_node = SimpleNamespace( @@ -390,12 +413,15 @@ async def dispatch(entries, *, engine_lease=None): observed_nodes.append(lease.node) return [(entries[0], ValueError("injected producer failure"))] + manager.args = SimpleNamespace(aspec_shared_source=shared) + if shared: + manager._prompt_buffer[0].metadata["shared_source_node"] = "node-b" manager._dispatch_batch = dispatch await manager._dispatch_ready_batches() await asyncio.sleep(0) assert observed_nodes == ["node-b"] - assert manager._capacity_rejections_by_node == {"node-a": 1, "node-b": 0} + assert manager._capacity_rejections_by_node == {"node-a": 0 if shared else 1, "node-b": 0} assert manager._inflight_samples_by_node == {"node-a": 0, "node-b": 1} assert manager._scheduled_samples_by_node == {"node-a": 0, "node-b": 1} with pytest.raises(RuntimeError, match="injected producer failure"): diff --git a/tests/test_multi_draft_plan.py b/tests/cpu/test_multi_draft_plan.py similarity index 79% rename from tests/test_multi_draft_plan.py rename to tests/cpu/test_multi_draft_plan.py index 268a2c9..1eb1ed4 100644 --- a/tests/test_multi_draft_plan.py +++ b/tests/cpu/test_multi_draft_plan.py @@ -26,7 +26,7 @@ def _config(): - root = Path(__file__).parents[1] + root = Path(__file__).parents[2] base = OmegaConf.load(root / "configs/recipes/qwen3_8b_dspark.yaml") base.dataset.train_data_path = str(root / "examples/data/sample_conversations.jsonl") base.model.target_model_path = "Qwen/Qwen3-8B" @@ -228,7 +228,7 @@ def test_shared_context_semantics_are_validated_before_gpu_allocation(): def test_shared_source_preserves_first_draft_attention_backend(tmp_path, attention_backend): import json - root = Path(__file__).parents[1] + root = Path(__file__).parents[2] raw = json.loads((root / "aspec/config/dflash2_draft_config.json").read_text()) raw["is_causal"] = True draft = tmp_path / "causal.json" @@ -288,3 +288,70 @@ def test_common_cp_config_validates_the_single_physical_source_before_compilatio loaded.experiments = [] with pytest.raises(ValueError, match="source.*ring"): load_config(base_config=loaded, save_snapshot=False) + + +def _multinode_config(): + root = Path(__file__).parents[2] + base = OmegaConf.load(root / 'configs/recipes/qwen3_8b_4k_multinode_node_local_nvlink.yaml') + base.training.training_node_ips = ['10.0.0.1', '10.0.0.2'] + base.inference.inference_node_ips = ['10.0.0.1', '10.0.0.2'] + base.inference.inference_batch_size = 1 + base.model.target_model_path = '/tmp/aspec-test/Qwen3-8B' + base.experiments = [ + {'name': 'a', 'training': {'training_num_gpus_per_node': 1}}, + {'name': 'b', 'training': {'training_num_gpus_per_node': 2}}, + ] + return load_config(base_config=base, save_snapshot=False) + + +def test_shared_multinode_plan_reserves_all_drafts_on_each_node(): + plan = compile_multi_draft_plan(_multinode_config()) + args = plan.to_runtime_args() + assert args.training_num_nodes == 2 + assert args.training_num_gpus_per_node == 3 + assert args.world_size == 6 + assert plan.to_dict()['resources']['draft_gpus'] == {'a': 2, 'b': 4} + domains = plan.source_plan.policy.resources['locality_domains'] + assert [d['training_world_ranks'] for d in domains] == [[0, 1, 2], [3, 4, 5]] + assert args.load_path is None + + +def test_child_allocation_preserves_node_major_rank_order(): + from types import SimpleNamespace + + from aspec.multi_draft import experiment_gpu_allocation + + children = { + 'a': SimpleNamespace(training_num_nodes=2, training_num_gpus_per_node=1, world_size=2), + 'b': SimpleNamespace(training_num_nodes=2, training_num_gpus_per_node=2, world_size=4), + } + reservation = ('pg', [10, 11, 12, 20, 21, 22], [0, 1, 2, 0, 1, 2]) + assert experiment_gpu_allocation(reservation, children, 'a') == ('pg', [10, 20], [0, 0]) + assert experiment_gpu_allocation(reservation, children, 'b') == ('pg', [11, 12, 21, 22], [1, 2, 1, 2]) + + +@pytest.mark.parametrize('packing', [True, False]) +def test_shared_multinode_requires_single_sample_target_requests(packing): + config = _multinode_config() + config.training.dflash_packing = packing + config.inference.inference_batch_size = 2 + with pytest.raises(NotImplementedError, match='inference_batch_size=1'): + compile_multi_draft_plan(config) + + +def test_shared_multinode_supports_independent_packing(): + config = _multinode_config() + config.experiments[1].training.dflash_packing = True + plan = compile_multi_draft_plan(config) + assert plan.plans['b'].to_runtime_args().dflash_packing is True + assert plan.plans['a'].to_runtime_args().dflash_packing is False + + +@pytest.mark.parametrize('multinode', [False, True]) +@pytest.mark.parametrize('packing', [False, True]) +def test_shared_source_rejects_finite_document_budget(multinode, packing): + config = _multinode_config() if multinode else _config() + config.training.dflash_packing = packing + config.training.source_work_budget = 8192 + with pytest.raises(NotImplementedError, match='source_work_budget'): + compile_multi_draft_plan(config) diff --git a/tests/test_multi_draft_runtime.py b/tests/cpu/test_multi_draft_runtime.py similarity index 78% rename from tests/test_multi_draft_runtime.py rename to tests/cpu/test_multi_draft_runtime.py index 9e8ea4c..dcff686 100644 --- a/tests/test_multi_draft_runtime.py +++ b/tests/cpu/test_multi_draft_runtime.py @@ -43,23 +43,19 @@ def test_failed_fence_does_not_withdraw_reader(monkeypatch): hub.retire.remote.assert_not_called() -def test_fence_probe_uses_live_node_not_stale_record_for_same_ip(monkeypatch): - group, probe_actor = Mock(), Mock() - group._actor_handlers = [Mock()] - monkeypatch.setattr(multi_draft, "_OwnerFenceProbe", probe_actor) - monkeypatch.setattr( - multi_draft.ray, - "nodes", - lambda: [ - {"NodeID": "0" * 56, "NodeManagerAddress": "host", "Alive": False}, - {"NodeID": "1" * 56, "NodeManagerAddress": "host", "Alive": True}, - ], - ) - monkeypatch.setattr(multi_draft.ray, "get", Mock(return_value=True)) +def test_fence_waits_for_every_initialized_node_identity(monkeypatch): + group = Mock() + group._actor_handlers = [Mock(), Mock()] + statuses = [ + {"node_id": "a", "node_ip": "host-a", "pid": 1}, + {"node_id": "b", "node_ip": "host-b", "pid": 2}, + ] monkeypatch.setattr(multi_draft.ray, "kill", Mock()) - multi_draft._fence_group(group, [{"node_ip": "host", "pid": 1}]) - placement = probe_actor.options.call_args.kwargs["scheduling_strategy"] - assert placement.node_id == "1" * 56 + wait = Mock(return_value={"state": "verified"}) + monkeypatch.setattr(multi_draft, "wait_for_owner_exit", wait) + multi_draft._fence_group(group, statuses) + wait.assert_called_once_with(statuses, timeout=30) + assert multi_draft.ray.kill.call_count == 2 def _session( @@ -280,7 +276,7 @@ def fail_close(**kwargs): def test_unsafe_fence_withholds_source_teardown_and_placement_release(monkeypatch, tmp_path): args, plan, recorders, engine = _session(monkeypatch, tmp_path, fence_error=True) - with pytest.raises(TimeoutError, match="owner alive"): + with pytest.raises(multi_draft.SourceTeardownBlockedError): multi_draft.run_multi_draft(args, plan) engine.shutdown.remote.assert_not_called() multi_draft.ray.util.remove_placement_group.assert_not_called() @@ -346,7 +342,7 @@ class Captured(Exception): pass parent = recipe_gpu_smoke.AuditedTrainerActor - config = Path(__file__).parents[1] / "configs/examples/qwen3_8b_multi_algorithm_feature_views.yaml" + config = Path(__file__).parents[2] / "configs/examples/qwen3_8b_multi_algorithm_feature_views.yaml" monkeypatch.setattr(sys, "argv", [ "lifecycle", "--config", str(config), "--output", str(tmp_path / "new"), "--resume-from", str(tmp_path / "old"), "--gpus", "8", @@ -391,3 +387,83 @@ def observe(observed_args, _group, _engines): assert evidence["state"] == "completed_with_failed_experiments" assert evidence["cleanup_errors"] == [] recorders["a"].fail.assert_not_called() + + +def test_shared_cli_retains_ray_driver_when_remote_owner_exit_is_unverified(monkeypatch): + from contextlib import nullcontext + + args = SimpleNamespace() + plan = multi_draft.MultiDraftPlan({}, {}, 42) + blocked = multi_draft.SourceTeardownBlockedError([ + {'node_id': 'lost', 'node_ip': '10.0.0.2', 'pid': 123}, + ]) + monkeypatch.setattr(train_entry, '_prepare_config', lambda: (args, plan)) + monkeypatch.setattr(train_entry, 'termination_signal_handlers', nullcontext) + monkeypatch.setattr(train_entry, 'ensure_ray_initialized', lambda: None) + monkeypatch.setattr(train_entry, 'install_termination_signal_handler', lambda: None) + monkeypatch.setattr(multi_draft, 'run_multi_draft', Mock(side_effect=blocked)) + shutdown = Mock() + monkeypatch.setattr(train_entry.ray, 'shutdown', shutdown) + monkeypatch.setattr(train_entry.ray, 'is_initialized', lambda: True) + with pytest.raises(multi_draft.SourceTeardownBlockedError): + train_entry._main() + shutdown.assert_not_called() + + +@pytest.mark.parametrize("first_name", ["a", "shared_source"]) +def test_shared_preflight_writes_child_and_aggregate_results(monkeypatch, tmp_path, first_name): + args = SimpleNamespace(cluster_mode='static_multi_node', output_dir=str(tmp_path), width=3) + children = { + name: SimpleNamespace(output_dir=str(tmp_path / name), width=width) + for name, width in ((first_name, 1), ('b', 2)) + } + plan = SimpleNamespace(plan_id='plan', plans={ + name: SimpleNamespace(to_runtime_args=lambda child=child: child) + for name, child in children.items() + }) + from aspec.ray import node_local_files + + close = Mock() + monkeypatch.setattr(node_local_files, 'close_stores', close) + monkeypatch.setattr(train_entry, '_prepare_run_directory', lambda path: None) + probe = Mock(side_effect=lambda child: {'state': 'validated', 'width': child.width}) + monkeypatch.setattr(multi_draft, 'run_cluster_preflight', probe) + report = multi_draft.preflight_multi_draft(args, plan) + assert probe.call_count == 3 + assert report['experiments'] == { + first_name: {'state': 'validated', 'width': 1}, + 'b': {'state': 'validated', 'width': 2}, + } + assert report['shared_source'] == {'state': 'validated', 'width': 3} + assert len({call.args[0].aspec_run_id for call in probe.call_args_list}) == 1 + assert close.call_count == 3 + assert json.loads((tmp_path / 'cluster-preflight.json').read_text()) == report + + +@pytest.mark.parametrize("failure", ["preflight", "cleanup"]) +def test_shared_preflight_attempts_every_store_cleanup(monkeypatch, tmp_path, failure): + from aspec.ray import node_local_files + + args = SimpleNamespace(cluster_mode="static_multi_node", output_dir=str(tmp_path)) + children = { + name: SimpleNamespace(output_dir=str(tmp_path / name)) for name in ("a", "b") + } + plan = SimpleNamespace(plan_id="plan", plans={ + name: SimpleNamespace(to_runtime_args=lambda child=child: child) + for name, child in children.items() + }) + monkeypatch.setattr(train_entry, "_prepare_run_directory", lambda path: None) + probe = Mock(return_value={"state": "validated"}) + close = Mock() + if failure == "preflight": + probe.side_effect = ValueError("preflight failed") + else: + close.side_effect = [RuntimeError("cleanup failed"), None, None] + monkeypatch.setattr(multi_draft, "run_cluster_preflight", probe) + monkeypatch.setattr(node_local_files, "close_stores", close) + with pytest.raises((ValueError, RuntimeError), match=f"{failure} failed"): + multi_draft.preflight_multi_draft(args, plan) + assert {call.args[0].output_dir for call in close.call_args_list} == { + str(tmp_path), str(tmp_path / "a"), str(tmp_path / "b"), + } + assert close.call_count == 3 diff --git a/tests/cpu/test_release_recipes.py b/tests/cpu/test_release_recipes.py index 98a1a4c..9f8baa2 100644 --- a/tests/cpu/test_release_recipes.py +++ b/tests/cpu/test_release_recipes.py @@ -37,14 +37,14 @@ def test_catalog_plans(name, monkeypatch): monkeypatch.setenv("ASPEC_NODE_0_IP", "192.0.2.1") monkeypatch.setenv("ASPEC_NODE_1_IP", "192.0.2.2") monkeypatch.setenv("ASPEC_SHARED_ROOT", "/tmp/aspec-test") - if name == "multinode": + if name.startswith("multinode"): monkeypatch.setenv("ASPEC_SHARED_ROOT", str(recipes.ROOT.parent)) monkeypatch.setenv("ASPEC_TARGET_MODEL", str(recipes.ROOT.parent / "models/Qwen3-8B")) monkeypatch.setenv("ASPEC_TRAIN_DATA", str(recipes.ROOT.parent / "train.parquet")) monkeypatch.setenv("ASPEC_EVAL_DATA", str(recipes.ROOT.parent / "eval.parquet")) for path in recipes.catalog()[name]["configs"]: config = load_config(str(recipes.ROOT / path)) - if name == "shared-target": + if config.experiments: from aspec.multi_draft import compile_multi_draft_plan compile_multi_draft_plan(config) @@ -73,7 +73,7 @@ def test_qwen8b_recipes_plan_without_target_override(name, monkeypatch): for path in recipes.catalog()[name]["configs"]: config = load_config(str(recipes.ROOT / path)) assert config.model.target_model_path == "Qwen/Qwen3-8B" - if name == "shared-target": + if config.experiments: from aspec.multi_draft import compile_multi_draft_plan compile_multi_draft_plan(config) @@ -83,7 +83,7 @@ def test_qwen8b_recipes_plan_without_target_override(name, monkeypatch): def test_catalog_exactly_covers_formal_files(): entries = recipes.catalog() - assert len(entries) == 13 + assert len(entries) == 14 declared = {p for entry in entries.values() for p in entry["configs"]} assert declared == { str(p.relative_to(recipes.ROOT)) for p in (recipes.ROOT / "configs/recipes").glob("*.yaml") diff --git a/tests/test_shared_source.py b/tests/cpu/test_shared_source.py similarity index 75% rename from tests/test_shared_source.py rename to tests/cpu/test_shared_source.py index 4846b1f..edf9472 100644 --- a/tests/test_shared_source.py +++ b/tests/cpu/test_shared_source.py @@ -15,6 +15,7 @@ from collections import deque from concurrent.futures import ThreadPoolExecutor from copy import copy +from queue import Queue from threading import Event from types import SimpleNamespace from unittest.mock import Mock @@ -136,6 +137,10 @@ def set_inference_error(self, message): def get_ready_pool_size(self): return len(self.outputs) + def get_ready_pool_size_by_node(self): + return {node: sum(output.metadata.get("source_node_ip") == node for output in self.sample_pool) + for node in getattr(self.args, "training_node_ips", [])} + def test_shared_eval_fanout_does_not_spend_training_demand(monkeypatch): monkeypatch.setattr(shared_source, "TrainingControllerState", _Controller) @@ -665,3 +670,195 @@ def test_packed_horizon_is_rows_not_a_fixed_sample_budget(monkeypatch): hub.retire("packed", owners_fenced=True) assert hub.remaining["packed"] == 0 assert hub.get_prompts(1) == [] + + +def test_shared_multinode_admission_reserves_each_child_tail_on_every_node(monkeypatch): + monkeypatch.setattr(shared_source, 'TrainingControllerState', _Controller) + args = SimpleNamespace(dp_size=2, num_train_steps=2, global_batch_size=2, + training_node_ips=['a', 'b'], mooncake_feature_affinity='node_local_required') + shorter = copy(args) + shorter.num_train_steps = 1 + hub = shared_source.SharedSourceState(args, {'long': args, 'short': shorter}) + hub.load_dataset(args) + for name in ('long', 'short'): + hub.invoke(name, 'submit_training_dataset', (), {}) + entries = hub.get_prompts(10) + assert len(entries) == 4 + assert [e.metadata['shared_source_node'] for e in entries] == ['a', 'b', 'a', 'b'] + assert [hub._audiences[e.metadata[EXECUTION_ID_KEY]] for e in entries] == [ + ('long', 'short'), ('long', 'short'), ('long',), ('long',), + ] + assert hub.remaining == {'long': 0, 'short': 0} + assert hub._remaining_by_node == {'long': {'a': 0, 'b': 0}, 'short': {'a': 0, 'b': 0}} + assert hub.source.args.mooncake_feature_affinity == 'any' + + +def test_shared_node_credits_include_unpublished_outputs_per_slowest_reader(monkeypatch): + monkeypatch.setattr(shared_source, 'TrainingControllerState', _Controller) + args = SimpleNamespace(dp_size=2, num_train_steps=2, global_batch_size=2, + training_node_ips=['a', 'b'], mooncake_feature_affinity='node_local_required') + hub = shared_source.SharedSourceState(args, {'fast': args, 'slow': args}) + hub.experiments['fast'].get_ready_pool_size_by_node = lambda: {'a': 0, 'b': 2} + hub.experiments['slow'].get_ready_pool_size_by_node = lambda: {'a': 3, 'b': 0} + hub._completed['pending'] = SimpleNamespace(metadata={'source_node_ip': 'b'}) + hub._audiences['pending'] = ('fast', 'slow') + assert hub.get_ready_pool_size_by_node() == {'a': 3, 'b': 3} + + +def test_disjoint_node_tail_demands_preserve_each_readers_next_work(monkeypatch): + monkeypatch.setattr(shared_source, 'TrainingControllerState', _Controller) + args = SimpleNamespace(dp_size=2, num_train_steps=1, global_batch_size=2, + training_node_ips=['a', 'b'], mooncake_feature_affinity='node_local_required') + hub = shared_source.SharedSourceState(args, {'left': args, 'right': args}) + hub.load_dataset(args) + for name in hub.experiments: + hub.invoke(name, 'submit_training_dataset', (), {}) + # The two independent restored readers need the same next document on + # different nodes. Neither may silently skip it to fill its local tail. + hub.remaining = {'left': 1, 'right': 1} + hub._remaining_by_node = {'left': {'a': 1, 'b': 0}, 'right': {'a': 0, 'b': 1}} + first, = hub.get_prompts(1) + second, = hub.get_prompts(1) + assert first.metadata['aspec_work_id'] == second.metadata['aspec_work_id'] + assert first.metadata[EXECUTION_ID_KEY] != second.metadata[EXECUTION_ID_KEY] + assert hub._audiences[first.metadata[EXECUTION_ID_KEY]] == ('left',) + assert hub._audiences[second.metadata[EXECUTION_ID_KEY]] == ('right',) + assert hub.remaining == {'left': 0, 'right': 0} + + +def _packed_multinode_hub(monkeypatch, *, plain=False): + monkeypatch.setattr(shared_source, 'TrainingControllerState', _Controller) + args = SimpleNamespace(dp_size=2, num_train_steps=2, global_batch_size=2, + dflash_packing=True, max_sample_pool_size=4, + training_node_ips=['a', 'b'], mooncake_feature_affinity='node_local_required') + other = copy(args) + other.dflash_packing = not plain + hub = shared_source.SharedSourceState(args, {'left': args, 'right': other}) + hub.load_dataset(args) + for name in hub.experiments: + hub.invoke(name, 'submit_training_dataset', (), {}) + return hub + + +@pytest.mark.parametrize('plain', [False, True]) +def test_packed_multinode_reserves_prefetch_credit_and_resumes_supply(monkeypatch, plain): + hub = _packed_multinode_hub(monkeypatch, plain=plain) + assert hub._remaining_by_node['left'] == {'a': None, 'b': None} + entries = hub.get_prompts(100) + assert len(entries) == 4 + assert hub.get_prompts(1) == [] + assert len(hub.source.prompt_buffer) == 0 + assert hub.source._dataset_epoch == 0 # No logical work lost while blocked. + controls = Mock() + monkeypatch.setattr(shared_source, 'source_control', controls) + outputs = [] + for index, entry in enumerate(entries): + output = _output(entry, index) + output.metadata.update(feature_affinity='node_local_required', + source_node_ip=entry.metadata['shared_source_node']) + outputs.append(output) + hub.push_inference_results(outputs[::-1]) + assert hub.get_prompts(1) == [] + for child in hub.experiments.values(): + child.sample_pool.clear() + new = hub.get_prompts(100) + assert len(new) == 4 + assert hub.remaining['left'] is None + if plain: + assert all(hub._audiences[entry.metadata[EXECUTION_ID_KEY]] == ('left',) for entry in new) + + +def test_packed_multinode_splits_crossed_reader_pressure_with_scoped_credit(monkeypatch): + hub = _packed_multinode_hub(monkeypatch) + hub.experiments['left'].get_ready_pool_size_by_node = lambda: {'a': 2, 'b': 0} + hub.experiments['right'].get_ready_pool_size_by_node = lambda: {'a': 0, 'b': 2} + first, second = hub.get_prompts(2) + assert first.metadata['aspec_work_id'] == second.metadata['aspec_work_id'] + assert first.metadata[EXECUTION_ID_KEY] != second.metadata[EXECUTION_ID_KEY] + assert hub.get_ready_pool_size_by_node() == {'a': 2, 'b': 2} + for entry in (first, second): + identity, node = entry.metadata[EXECUTION_ID_KEY], entry.metadata['shared_source_node'] + assert hub.get_ready_pool_size_by_node([identity])[node] == 0 + assert len(hub.get_prompts(100)) == 2 # Includes the already reserved credits. + assert hub.get_prompts(1) == [] + with pytest.raises(ValueError, match='unknown shared source admission'): + hub.get_ready_pool_size_by_node(['unknown']) + + +def test_shared_multinode_output_cannot_rewrite_its_admitted_node(monkeypatch): + hub = _packed_multinode_hub(monkeypatch) + entry, = hub.get_prompts(1) + output = _output(entry, 0) + other_node = 'b' if entry.metadata['shared_source_node'] == 'a' else 'a' + output.metadata.update(feature_affinity='node_local_required', + source_node_ip=other_node, shared_source_node=other_node) + with pytest.raises(RuntimeError, match='admitted source-node placement'): + hub.push_inference_results([output]) + + +@pytest.mark.parametrize('cancel', [True, False]) +def test_packed_multinode_stops_node_demand_on_cancel_or_retirement(monkeypatch, cancel): + hub = _packed_multinode_hub(monkeypatch) + hub.get_prompts(2) + if cancel: + hub.request_cancel('left') + else: + hub.retire('left', owners_fenced=True) + assert hub.remaining['left'] == 0 + assert hub.status()['remaining_source_demand_by_node']['left'] == {'a': 0, 'b': 0} + assert hub.remaining['right'] is None + + +def test_shared_packing_checkpoint_replays_an_uncommitted_earlier_epoch(monkeypatch): + hub = _packed_multinode_hub(monkeypatch) + hub._committed_epoch['left'] = 3 + hub._epochs['left'] = {'old-tail': 2, 'later': 3} + assert hub.invoke('left', 'get_checkpoint_work_state', (), {})['dataset_epoch'] == 2 + del hub._epochs['left']['old-tail'] + assert hub.invoke('left', 'get_checkpoint_work_state', (), {})['dataset_epoch'] == 3 + + +@pytest.mark.parametrize('plain', [False, True]) +def test_shared_outputs_use_independent_real_node_local_packing_controllers(monkeypatch, plain): + from aspec.controller import training_controller + + real_controller = shared_source.TrainingControllerState + hub = _packed_multinode_hub(monkeypatch, plain=plain) + monkeypatch.setattr(training_controller, 'Queue', Queue) + monkeypatch.setattr(shared_source, 'source_control', Mock()) + hub._node_watermarks = {'a': 4, 'b': 4} + for name, old in list(hub.experiments.items()): + args = copy(old.args) + args.training_num_gpus_per_node = 1 + args.max_seq_length, args.dflash_block_size = 32, 2 + args.packing_min_samples_per_dispatch = 4 + hub.experiments[name] = real_controller(args, 2, with_eval_queues=False) + entries = hub.get_prompts(8) + outputs = [] + for index, entry in enumerate(entries): + output = _output(entry, index) + output.tensor_shapes = {'input_ids': (1, 7), 'hidden_states': (1, 7, 3)} + output.packed_loss_mask = '0,7' + output.metadata.update(feature_affinity='node_local_required', + source_node_ip=entry.metadata['shared_source_node']) + outputs.append(output) + hub.push_inference_results(outputs[::-1]) + selected = {} + for name, child in hub.experiments.items(): + assert hub.invoke(name, 'try_dispatch_batch', (), {'optimizer_step': 0, 'microstep': 0}) + selected[name] = [] + for node, queue in zip(('a', 'b'), child.train_queues, strict=True): + row = [] + while not queue.empty(): + sample = queue.get_nowait() + assert sample.metadata['source_node_ip'] == sample.metadata['assigned_draft_node_ip'] == node + row.append(sample) + assert len(row) == (1 if name == 'right' and plain else 4) + if child.packing_enabled: + assert [sample.last_in_row for sample in row] == [False, False, False, True] + assert len({sample.row_id for sample in row}) == 1 + selected[name].extend(sample.metadata[EXECUTION_ID_KEY] for sample in row) + assert len(selected[name]) == len(set(selected[name])) + assert child.get_scheduler_evidence()['cross_node_dispatch_samples'] == 0 + if not plain: + assert set(selected['left']) == set(selected['right']) diff --git a/tests/test_startup_cleanup.py b/tests/test_startup_cleanup.py index 0111fb7..c480680 100644 --- a/tests/test_startup_cleanup.py +++ b/tests/test_startup_cleanup.py @@ -18,29 +18,24 @@ import pytest -def test_multi_draft_probe_is_closed_if_trainer_kill_fails(monkeypatch): +def test_multi_draft_kill_failure_keeps_reader_registered(monkeypatch): from aspec import multi_draft as module - actor, probe = object(), object() - creator = MagicMock() - creator.options.return_value.remote.return_value = probe - monkeypatch.setattr(module, "_OwnerFenceProbe", creator) - monkeypatch.setattr( - module.ray, - "nodes", - lambda: [{"Alive": True, "NodeID": "1" * 56, "NodeManagerAddress": "ip"}], - ) - killed = [] - - def kill(target, **kw): - killed.append(target) - if target is actor: - raise RuntimeError("kill dispatch failed") - + actor = object() + kill = MagicMock(side_effect=RuntimeError("kill dispatch failed")) + wait = MagicMock() + hub = SimpleNamespace(retire=SimpleNamespace(remote=MagicMock())) monkeypatch.setattr(module.ray, "kill", kill) + monkeypatch.setattr(module, "wait_for_owner_exit", wait) with pytest.raises(RuntimeError, match="kill dispatch failed"): - module._fence_group(SimpleNamespace(_actor_handlers=[actor]), [{"node_ip": "ip", "pid": 1}]) - assert probe in killed + module._finish_experiment( + SimpleNamespace(_actor_handlers=[actor]), + [{"node_id": "1" * 56, "node_ip": "ip", "pid": 1}], + hub, "draft", graceful=False, + ) + kill.assert_called_once_with(actor, no_restart=True) + wait.assert_not_called() + hub.retire.remote.assert_not_called() def test_materializer_startup_failure_enters_cleanup(monkeypatch, tmp_path): diff --git a/tools/recipe_catalog.json b/tools/recipe_catalog.json index 6c0873f..3e3b82e 100644 --- a/tools/recipe_catalog.json +++ b/tools/recipe_catalog.json @@ -111,7 +111,8 @@ "gpus": 8, "target_gpus": 4, "draft_gpus": 4, - "native_only_export": false + "native_only_export": false, + "draft_experiments": 2 }, "qwen27b-long": { "configs": [ @@ -146,5 +147,17 @@ "target_gpus": 4, "draft_gpus": 4, "native_only_export": true + }, + "multinode-shared": { + "configs": [ + "configs/recipes/qwen3_8b_4k_multinode_multi_draft_nvlink.yaml" + ], + "model": "Qwen3-8B", + "nodes": 2, + "gpus": 16, + "target_gpus": 8, + "draft_gpus": 8, + "native_only_export": false, + "draft_experiments": 2 } } diff --git a/tools/recipes.py b/tools/recipes.py index f34fbc2..70e943e 100644 --- a/tools/recipes.py +++ b/tools/recipes.py @@ -208,7 +208,10 @@ def main(): "tests/cpu/test_training_resume_rng.py", "tests/cpu/test_evaluation_contract.py", "tests/test_ling_moe.py", - "tests/test_multi_draft_plan.py", + "tests/cpu/test_multi_draft_plan.py", + "tests/cpu/test_multi_draft_runtime.py", + "tests/cpu/test_shared_source.py", + "tests/cpu/test_inference_manager_lifecycle.py", "tests/test_multinode_tcp_reference.py", ] )