From b93896c3350397632ed9742228b2f337233874f7 Mon Sep 17 00:00:00 2001 From: bruceszchen Date: Tue, 22 Sep 2026 21:49:02 +0800 Subject: [PATCH] feat(parity): add Qwen3-MoE training and verification experiment --- experimental/train_inference_parity/README.md | 206 ++++++++ .../train_inference_parity/__init__.py | 5 + .../train_inference_parity/algorithm.py | 122 +++++ .../train_inference_parity/artifacts.py | 236 +++++++++ .../examples/qwen3_moe_30b_a3b_fsdp_tp4.yaml | 173 +++++++ experimental/train_inference_parity/gate.py | 255 ++++++++++ .../train_inference_parity/lifecycle.py | 68 +++ .../train_inference_parity/models/__init__.py | 1 + .../models/qwen3_moe_30b_a3b/__init__.py | 6 + .../models/qwen3_moe_30b_a3b/actor.py | 20 + .../models/qwen3_moe_30b_a3b/aten.py | 50 ++ .../models/qwen3_moe_30b_a3b/attention.py | 215 ++++++++ .../models/qwen3_moe_30b_a3b/context.py | 49 ++ .../models/qwen3_moe_30b_a3b/contract.py | 115 +++++ .../models/qwen3_moe_30b_a3b/fsdp.py | 5 + .../models/qwen3_moe_30b_a3b/install.py | 216 ++++++++ .../models/qwen3_moe_30b_a3b/linear_ops.py | 90 ++++ .../models/qwen3_moe_30b_a3b/moe.py | 164 ++++++ .../models/qwen3_moe_30b_a3b/norm.py | 62 +++ .../models/qwen3_moe_30b_a3b/rope.py | 34 ++ .../models/qwen3_moe_30b_a3b/router.py | 27 + experimental/train_inference_parity/run.py | 77 +++ .../runtime_contract.py | 433 ++++++++++++++++ experimental/train_inference_parity/sync.py | 196 ++++++++ .../train_inference_parity/verification.py | 475 ++++++++++++++++++ unirl/algorithms/grpo.py | 57 ++- unirl/distributed/weight_sync/full/ipc.py | 10 + unirl/models/qwen3/ar.py | 72 ++- unirl/models/qwen3/bundle.py | 28 +- unirl/models/qwen3/config.py | 5 + unirl/models/qwen3/pipeline.py | 12 +- unirl/models/types/meta_init.py | 5 +- unirl/rollout/engine/vllm/engine.py | 13 +- unirl/trainer/ar.py | 10 +- 34 files changed, 3472 insertions(+), 40 deletions(-) create mode 100644 experimental/train_inference_parity/README.md create mode 100644 experimental/train_inference_parity/__init__.py create mode 100644 experimental/train_inference_parity/algorithm.py create mode 100644 experimental/train_inference_parity/artifacts.py create mode 100644 experimental/train_inference_parity/examples/qwen3_moe_30b_a3b_fsdp_tp4.yaml create mode 100644 experimental/train_inference_parity/gate.py create mode 100644 experimental/train_inference_parity/lifecycle.py create mode 100644 experimental/train_inference_parity/models/__init__.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/__init__.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/actor.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/aten.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/attention.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/context.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/contract.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/fsdp.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/install.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/linear_ops.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/moe.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/norm.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/rope.py create mode 100644 experimental/train_inference_parity/models/qwen3_moe_30b_a3b/router.py create mode 100644 experimental/train_inference_parity/run.py create mode 100644 experimental/train_inference_parity/runtime_contract.py create mode 100644 experimental/train_inference_parity/sync.py create mode 100644 experimental/train_inference_parity/verification.py diff --git a/experimental/train_inference_parity/README.md b/experimental/train_inference_parity/README.md new file mode 100644 index 000000000..028b20eeb --- /dev/null +++ b/experimental/train_inference_parity/README.md @@ -0,0 +1,206 @@ +# experimental/train_inference_parity + +Owner: UniRL Qwen3-MoE parity maintainers + +This package incubates an exact rollout-versus-gradient-replay selected-token +log-probability contract without shipping experimental kernels in the UniRL +wheel. A normal baseline does not load this package or its nested vLLM plugin. + +## Current status: UNVERIFIED, not PASS + +There is no checked-in verification artifact produced by the current recipe, +so this repository makes no PASS claim. The recipe uses the dedicated +`experimental.train_inference_parity.sync.ParityIPCWeightSync` evidence adapter +over UniRL's production vLLM native IPC sync and will fail rather than falling +back to a different transport. + +A future result may be called PASS only when the generated JSON artifact has +`status: "PASS"` and contains both `initial_checkpoint` and +`post_update_reload`. Console output, W&B history, and a first-rollout-only run +are not substitutes. + +## Frozen target contract + +- Model: reviewed local mirror of `Qwen/Qwen3-30B-A3B` revision + `ad44e777bcd18fa416d9da3bd8f70d33ebb85d39` +- Topology: one FSDP actor world of 4 and one direct vLLM TP4 engine on exactly + 4 visible GPUs +- Software: Python 3.13, vLLM 0.27.0, Torch 2.13.0, Transformers 5.6.0 +- Actor/model compute: BF16; selected-token comparison and stored log-probs: + FP32 +- Sampling: temperature 1, top-p 1, top-k 0, processed log-probs +- Lengths: 256 prompt tokens plus 1024 response tokens; actor and rollout must + both use suffix-preserving (left) truncation +- Data: complete GSM8K JSONL, shuffled with OS entropy (`seed: null`); each + rollout consumes the next non-repeating batch until the epoch rolls over +- Updates: one optimizer update per rollout; two rollouts total + +Bitwise means the selected-token FP32 rollout and differentiable replay +log-probs on this frozen matrix. It does not claim bitwise gradients, optimizer +state, or portability to other GPUs, drivers, kernels, or package versions. + +`public_reference` is the only profile constant. There is intentionally no +provider registry or provider environment-variable map for one implementation. + +Numerical gotchas: + +- The FSDP actor is not actually tensor-parallel. Its exact MoE path reproduces + vLLM TP4's column/row slicing and concatenation order. `train_tp` has no code + default, the runtime contract checks divisibility, and actor/vLLM call the + same nested-plugin MoE combine precision function. +- The exact context remains active through loss and backward so non-reentrant + activation-checkpoint recomputation selects the same forward providers. + RMSNorm, MoE, attention, and log-softmax still use documented Torch surrogate + backward formulas; gradients must be finite but are not bitwise claims. +- Actor monkey patches require `UNIRL_PARITY_ENABLE=1` and a dedicated process. + Recoverable Python symbols retain their originals. PyTorch does not support + unregistering the CUDA ATen overrides, so process exit is the unload boundary. + +### Verification matrix + +| Recipe | Hardware target | Driver/runtime | Python packages | Git/artifact | Status | +|---|---|---|---|---|---| +| `qwen3_moe_30b_a3b_fsdp_tp4.yaml` | 4 × NVIDIA H20 96 GB, one node | driver 535.161.08 + CUDA-13 forward compatibility, CUDA 13.0 runtime; NCCL recorded at run time | Python 3.13, Torch 2.13.0, Transformers 5.6.0, vLLM 0.27.0 | must be the clean HEAD recorded in the generated JSON | **UNVERIFIED** | + +This row remains UNVERIFIED until a repository-generated artifact is checked +in or linked. It is a target matrix, not a historical PASS claim. + +## Install and preflight + +From a clean repository checkout on the target single-node four-GPU host: + +```bash +uv sync --extra vllm --extra train +uv pip install -e "experimental/train_inference_parity/vllm_plugin[verified]" + +export QWEN3_MOE_PATH=/dev/shm/Qwen3-30B-A3B +export QWEN3_MOE_REVISION=ad44e777bcd18fa416d9da3bd8f70d33ebb85d39 +export DATA_PATH=/path/to/gsm8k/train.jsonl +export EVAL_DATA_PATH=/path/to/gsm8k/test.jsonl +CUDA_VISIBLE_DEVICES=0,1,2,3 uv run ray start --head --num-gpus=4 +``` + +Driver 535 requires the host's CUDA-13 forward-compatibility library for the +CUDA-13.0 wheel stack. Configure that library before launch; the artifact +records the driver, runtime, and NCCL versions actually observed. + +`run.py` constructs one `ParityRuntimeContract`, installs its complete explicit +environment, then aggregates topology, precision, sampling, length, version, +GPU-count, trust, data/model-path, and Ray-preinitialization errors before +starting a model or connecting to Ray. In particular, Ray must not already be +initialized in the driver process. + +The lifecycle sets `UNIRL_PARITY_ENABLE=1`, the vLLM plugin allowlist, +deterministic CUDA/NCCL flags, the model/profile/patch manifest, and propagates +the same values into every Ray role. `trust_remote_code=true` is explicit on +both actor and rollout; use only a reviewed local checkpoint. +The plugin is a complete no-op when the enable flag is absent. When enabled it +requires strict mode, the complete ordered patch set, the three pinned package +versions, and the verified private-symbol signatures; it does not silently +skip a missing provider. + +## Quick smoke (no parity claim) + +These checks validate imports and the plugin's disabled no-op behavior. They do +not replace the TP4 run: + +```bash +uv run python -m compileall -q experimental/train_inference_parity \ + unirl/rollout/engine/vllm unirl/distributed/weight_sync +UNIRL_PARITY_ENABLE=0 uv run python -c \ + 'import unirl_train_inference_parity_vllm as p; assert p.register() == ()' +``` + +## Full two-phase run + +On the frozen target host: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 uv run python -m experimental.train_inference_parity.run \ + --config-name=qwen3_moe_30b_a3b_fsdp_tp4 +``` + +Defaults: + +- W&B disabled; opt in with `logging.report_to_wandb=true` +- vLLM native packed-IPC sync with a 640 MiB bucket, 2 GiB GPU safety margin, + and a Gloo control timeout (2400s) that outlives the 1800s weight-update RPC +- atomic JSON artifact: + `outputs/train_inference_parity/qwen3-moe-tp4.json` (override + `UNIRL_PARITY_ARTIFACT_PATH`) + +`ignore_eos=true` is a verification choice that forces the 1024-token boundary +and avoids EOS-dependent cases. It is not a recommended production RL default. + +## Verification artifact + +`verification.py` validates phase metrics and publication evidence; +`artifacts.py` owns provenance and atomic JSON persistence. The schema records: + +- git commit and tracked-file dirty state +- software, CUDA/driver/NCCL, and per-GPU hardware +- full model-tree and data-file SHA-256 fingerprints +- all numerical flags from the runtime contract +- the vLLM worker's strict before/after symbol-provider manifest +- canonical parity metrics, replay token count, grad norm, optimizer updates, + and parameter-change evidence for each phase +- the native IPC publication manifest, per-TP-worker receipts, and sampled + Actor parameter-change fingerprint + +The canonical metric names are `token_count`, `torch_equal_fp32`, +`mismatch_count`, `max_absdiff_fp32`, `k3_mean`, and `k3_max`. + +## Integration requirements + +1. Keep the recipe on `ParityIPCWeightSync`, which subclasses the production + `IPCWeightSync` data plane and adds evidence only. +2. Preserve `record_initial_actor_fingerprint()`, `mark_actor_updated()`, and + `verification_receipts()` on the parity sync adapter. Actor-change evidence + samples large structural shards (embed, Q/O, MoE, LM head), not the smallest + RMSNorms. Each publication must include the native structural manifest, one + receipt for every TP rank, matching publication/model identities, unique CUDA + UUIDs, `parameter_changed: true`, and `prefix_cache_reset: true`. +3. `TrainInferenceParityGRPO` must reach a pre-replay error barrier before any + rank enters FSDP/vLLM collectives, then the exact log-prob gate after replay. +4. Preserve the plugin's reload lifecycle: clear the derived MoE column cache + before vLLM sleep, leave the stock expert Parameter loaders unobstructed, and + load the custom column-parallel `o_proj` directly into its CuMem allocation + instead of passing it through vLLM's generic meta-staged reload. +5. Keep actor and direct-vLLM prompt handling suffix-preserving at the same + configured `max_prompt_length`; do not re-enable tokenizer-default + chat-template truncation before the shared engine limit. +6. Each phase must observe `batch_size × samples_per_prompt × max_new_tokens` + replay tokens (`4 × 8 × 1024 = 32768` on the frozen recipe). +7. Run from a clean source worktree on the frozen matrix and publish the + resulting artifact before changing this status to PASS. Generated `outputs/` + and plugin egg-info do not count as dirty source. + +## Package boundaries + +- `unirl/` never imports this package. +- Recipes under this package may target + `experimental.train_inference_parity.*`. +- The nested plugin imports only UniRL's stable vLLM worker-extension seam to + attach its runtime manifest; core UniRL never imports the plugin. +- Common code needed by a second experiment graduates into core rather than + being imported sideways. + +## Gotchas + +- The training-side PR requires the separately installed vLLM plugin for its + shared numerical providers. Its core UniRL changes remain inert on baseline + recipes, but the experimental recipe needs both PRs. +- GRPO loss construction is shared with the baseline implementation. The + experiment accepts only rollout-sourced anchors and keeps exact mode active + through backward, including activation-checkpoint recomputation. +- A fixed pre-replay barrier synchronizes local setup failures before FSDP + forwards. It cannot recover a worker that dies inside a model collective. +- Every micro-batch must report all canonical metrics. Missing evidence is an + error, not a micro-batch that can be omitted from aggregation. +- Source provenance includes individual untracked Python/config files inside + new directories. Untracked notes are excluded; tracked changes always count. +- The frozen 256/1024 token lengths cannot be shortened while retaining this + recipe's verification claim. W&B logging may be enabled independently. +- The Qwen3 meta-init path resolves weight shards with `model_revision`, just + like config and model loading. Local checkpoints still need external review + against that revision; setting the revision does not authenticate local bytes. diff --git a/experimental/train_inference_parity/__init__.py b/experimental/train_inference_parity/__init__.py new file mode 100644 index 000000000..f65e0e57a --- /dev/null +++ b/experimental/train_inference_parity/__init__.py @@ -0,0 +1,5 @@ +"""Incubating train/inference parity contracts and Qwen3-MoE integration.""" + +from .runtime_contract import PUBLIC_REFERENCE + +__all__ = ["PUBLIC_REFERENCE"] diff --git a/experimental/train_inference_parity/algorithm.py b/experimental/train_inference_parity/algorithm.py new file mode 100644 index 000000000..7a97738a4 --- /dev/null +++ b/experimental/train_inference_parity/algorithm.py @@ -0,0 +1,122 @@ +"""Experimental GRPO with collective-safe exact replay/rollout parity.""" + +from __future__ import annotations + +from typing import Any, Dict, Mapping, Optional + +import torch + +from experimental.train_inference_parity.gate import collective_error_barrier, exact_parity_gate +from unirl.algorithms.base import ( + AlgorithmStepResult, + typed_conditions, +) +from unirl.algorithms.grpo import GRPO +from unirl.types.conditions import Condition +from unirl.types.segments.text import TextSegment + + +class TrainInferenceParityGRPO(GRPO): + """GRPO whose optimization anchor is the exact rollout log-probability.""" + + supports_multi_update = False + + def prepare_segment( + self, + *, + conditions: Mapping[str, Condition], + segment: "TextSegment", + ) -> None: + if self.old_logp_source != "rollout": + raise ValueError("exact parity requires old_logp_source='rollout'") + if not isinstance(segment.log_probs, torch.Tensor): + raise TypeError("TrainInferenceParityGRPO requires rollout log_probs") + segment.rollout_log_probs = segment.log_probs.detach().cpu().clone() + + def _prepare_exact_replay( + self, + *, + conditions: Mapping[str, Condition], + segment: "TextSegment", + ) -> Any: + if segment.tokens is None or segment.lengths is None: + raise ValueError("TrainInferenceParityGRPO requires packed tokens and lengths on every rank") + if not isinstance(segment.rollout_log_probs, torch.Tensor): + raise TypeError("TrainInferenceParityGRPO requires segment.rollout_log_probs") + segment.log_probs = segment.rollout_log_probs + typed_conds = typed_conditions(conditions, self.conditions_cls) + context_factory = getattr(self.stage, "exact_context", None) + if not callable(context_factory): + raise RuntimeError("TrainInferenceParityGRPO requires stage.exact_context()") + return typed_conds, context_factory + + def compute_loss_and_backward( + self, + *, + conditions: Mapping[str, Condition], + segment: "TextSegment", + advantages: torch.Tensor, + training_progress: float, + loss_scale: float, + ) -> AlgorithmStepResult: + setup_error: Optional[str] = None + typed_conds = None + context_factory = None + try: + typed_conds, context_factory = self._prepare_exact_replay( + conditions=conditions, + segment=segment, + ) + except Exception as error: + setup_error = f"{type(error).__name__}: {error}" + collective_error_barrier(setup_error, phase="pre-replay") + + with context_factory(): + replay_error: Optional[str] = None + new_logp: Optional[torch.Tensor] = None + rollout_anchor: Optional[torch.Tensor] = None + try: + new_logp = self.stage.replay(typed_conds, segment=segment, temperature=self.sampling_temperature) + rollout_anchor = segment.rollout_log_probs + except Exception as error: + replay_error = f"{type(error).__name__}: {error}" + new_logp = None + rollout_anchor = None + + parity_metrics = exact_parity_gate( + new_logp if isinstance(new_logp, torch.Tensor) else None, + rollout_anchor, + local_error=replay_error, + ) + if new_logp is None or new_logp.numel() == 0: + return AlgorithmStepResult( + loss=0.0, + metrics=parity_metrics, + num_steps_or_tokens=0, + has_backward=False, + ) + loss, clip_range, ratio_metrics = self._build_policy_loss( + segment=segment, + new_logp=new_logp, + advantages=advantages, + training_progress=training_progress, + ) + (loss * loss_scale).backward() + + metrics: Dict[str, Any] = dict(parity_metrics) + metrics.update( + { + "policy_loss": float(loss.detach().item()), + "clip_range": float(clip_range), + **{key: float(value.item()) for key, value in ratio_metrics.items()}, + } + ) + return AlgorithmStepResult( + loss=float(loss.detach().item()), + metrics=metrics, + num_steps_or_tokens=int(new_logp.shape[0]), + has_backward=True, + ) + + +__all__ = ["TrainInferenceParityGRPO"] diff --git a/experimental/train_inference_parity/artifacts.py b/experimental/train_inference_parity/artifacts.py new file mode 100644 index 000000000..197474401 --- /dev/null +++ b/experimental/train_inference_parity/artifacts.py @@ -0,0 +1,236 @@ +"""Atomic JSON output and source/runtime provenance for parity verification.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +import platform +import subprocess +import tempfile +from dataclasses import asdict, is_dataclass +from pathlib import Path +from typing import Any, Mapping + +from .runtime_contract import EXPECTED_VERSIONS, PLUGIN_DISTRIBUTION + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +_SOURCE_SUFFIXES = {".py", ".yaml", ".yml", ".toml", ".json"} + + +def jsonable(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if is_dataclass(value): + return jsonable(asdict(value)) + if isinstance(value, Mapping): + return {str(key): jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [jsonable(item) for item in value] + item = getattr(value, "item", None) + if callable(item): + try: + return jsonable(item()) + except (TypeError, ValueError): + pass + return repr(value) + + +def atomic_write_json(path: str | os.PathLike[str], payload: Mapping[str, Any]) -> None: + """Persist strict JSON with same-directory fsync and atomic replacement.""" + target = Path(path).expanduser().resolve() + target.parent.mkdir(parents=True, exist_ok=True) + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=target.parent, prefix=f".{target.name}.", suffix=".tmp", delete=False + ) as handle: + temporary = Path(handle.name) + json.dump(jsonable(payload), handle, indent=2, sort_keys=True, allow_nan=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, target) + directory_fd = os.open(target.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def _run_git(*args: str) -> str: + return subprocess.run( + ["git", "-C", str(_REPOSITORY_ROOT), *args], + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout + + +def git_state() -> dict[str, Any]: + """Include tracked edits and individual untracked source files, including new directories.""" + try: + commit = _run_git("rev-parse", "HEAD").strip() + records = iter(_run_git("status", "--porcelain=v1", "-z", "--untracked-files=all").split("\0")) + dirty = [] + for record in records: + if not record: + continue + status, path = record[:2], record[3:] + if "R" in status or "C" in status: + next(records) # Porcelain -z puts the original rename/copy path second. + if status == "??": + parts = Path(path).parts + if parts[0] == "outputs" or any(part.endswith(".egg-info") for part in parts): + continue + if Path(path).suffix not in _SOURCE_SUFFIXES: + continue + dirty.append(record) + return { + "commit": commit, + "dirty": bool(dirty), + "dirty_scope": "tracked_and_untracked_source", + "dirty_paths": dirty, + } + except (OSError, subprocess.SubprocessError) as error: + return { + "commit": None, + "dirty": None, + "dirty_scope": "tracked_and_untracked_source", + "error": f"{type(error).__name__}: {error}", + } + + +def _package_version(distribution: str) -> str | None: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def software_state() -> dict[str, Any]: + state = { + "python": platform.python_version(), + "platform": platform.platform(), + "packages": { + name: _package_version(name) for name in ("torch", "transformers", "vllm", "ray", PLUGIN_DISTRIBUTION) + }, + "expected_version_lines": dict(EXPECTED_VERSIONS), + } + try: + import torch + except ImportError: + return state + nccl_version = None + try: + nccl_version = torch.cuda.nccl.version() + except (AttributeError, RuntimeError): + pass + driver_version = None + try: + driver_version = subprocess.run( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout.splitlines()[0] + except (OSError, subprocess.SubprocessError, IndexError): + pass + state.update( + cuda_runtime=torch.version.cuda, + cuda_driver=driver_version, + cudnn=torch.backends.cudnn.version(), + nccl=nccl_version, + ) + return state + + +def hardware_state() -> dict[str, Any]: + try: + import torch + except ImportError: + return {"cuda_available": False, "cuda_device_count": 0, "gpus": []} + count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + gpus = [] + for index in range(count): + properties = torch.cuda.get_device_properties(index) + gpus.append( + { + "index": index, + "name": properties.name, + "total_memory_bytes": int(properties.total_memory), + "compute_capability": [int(properties.major), int(properties.minor)], + "multi_processor_count": int(properties.multi_processor_count), + } + ) + return { + "hostname": platform.node(), + "machine": platform.machine(), + "cuda_available": bool(torch.cuda.is_available()), + "cuda_device_count": count, + "gpus": gpus, + } + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def model_fingerprint(raw_path: str) -> dict[str, Any]: + """Hash every checkpoint file byte, including model weights.""" + path = Path(raw_path).expanduser().resolve() + if path.is_file(): + return { + "path": str(path), + "kind": "file_bytes", + "sha256": _sha256_file(path), + "size_bytes": path.stat().st_size, + } + if not path.is_dir(): + return {"path": str(path), "kind": "missing", "sha256": None} + entries = [] + digest = hashlib.sha256() + for item in sorted(candidate for candidate in path.rglob("*") if candidate.is_file()): + entry = { + "path": item.relative_to(path).as_posix(), + "size_bytes": item.stat().st_size, + "sha256": _sha256_file(item), + } + entries.append(entry) + digest.update(json.dumps(entry, sort_keys=True, separators=(",", ":")).encode("utf-8")) + digest.update(b"\n") + parts = path.parts + revision_index = parts.index("snapshots") + 1 if "snapshots" in parts else len(parts) + return { + "path": str(path), + "kind": "file_tree_bytes", + "algorithm": "sha256(relative_path,size,file_sha256)", + "sha256": digest.hexdigest(), + "revision": parts[revision_index] if revision_index < len(parts) else None, + "files": entries, + } + + +def data_fingerprint(raw_path: str) -> dict[str, Any]: + path = Path(raw_path).expanduser().resolve() + if not path.is_file(): + return {"path": str(path), "kind": "missing", "sha256": None} + with path.open("rb") as handle: + rows = sum(1 for line in handle if line.strip()) + return { + "path": str(path), + "kind": "file_bytes", + "sha256": _sha256_file(path), + "size_bytes": path.stat().st_size, + "nonempty_rows": rows, + } diff --git a/experimental/train_inference_parity/examples/qwen3_moe_30b_a3b_fsdp_tp4.yaml b/experimental/train_inference_parity/examples/qwen3_moe_30b_a3b_fsdp_tp4.yaml new file mode 100644 index 000000000..65bd0d443 --- /dev/null +++ b/experimental/train_inference_parity/examples/qwen3_moe_30b_a3b_fsdp_tp4.yaml @@ -0,0 +1,173 @@ +# @package _global_ + +parity: + profile: public_reference + train_tp: 4 + max_prompt_length: 256 + max_new_tokens: 1024 + artifact_path: ${oc.env:UNIRL_PARITY_ARTIFACT_PATH,outputs/train_inference_parity/qwen3-moe-tp4.json} + +num_devices: 4 +devices_per_node: 4 +batch_size: 4 +num_rollouts: 2 +weight_sync_interval: 1 +eval_interval: 0 +enable_fsdp_offload: true +adv_normalization_scope: group +normalize_adv_by_std: true + +logging: + report_to_wandb: false + project_name: unirl-qwen3-moe-parity + run_name: qwen3-30b-a3b-tp4-bitwise-n8-r1024 + entity: null + tags: [experimental, qwen3-moe, bitwise, tp4, gsm8k, n8, response-1024] + +bundle: + _target_: unirl.models.qwen3.bundle.Qwen3Bundle.from_config + config: + _target_: unirl.models.qwen3.config.Qwen3PipelineConfig + pretrained_model_ckpt_path: ${oc.env:QWEN3_MOE_PATH,/dev/shm/Qwen3-30B-A3B} + model_revision: ${oc.env:QWEN3_MOE_REVISION,ad44e777bcd18fa416d9da3bd8f70d33ebb85d39} + trust_remote_code: true + meta_init_transformer: true + model_precision: bf16 + attn_implementation: flash_attention_3 + autocast_precision: bf16 + logprob_precision: fp32 + exact_actor_logprobs: true + external_libs: + - experimental.train_inference_parity.models.qwen3_moe_30b_a3b.fsdp + use_gradient_checkpointing: true + max_prompt_length: 256 + +pipeline: + _target_: experimental.train_inference_parity.models.qwen3_moe_30b_a3b.actor.ParityQwen3Pipeline.from_bundle + enable_thinking: true + max_prompt_length: 256 + autocast_precision: bf16 + logprob_precision: fp32 + exact_actor_logprobs: true + +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["Qwen3MoeDecoderLayer"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + forward_prefetch: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 1.0e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.01 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 200 + +rollout: + _target_: unirl.rollout.engine.vllm.engine.VLLMRolloutEngine + config: + _target_: unirl.rollout.engine.vllm.config.VLLMEngineConfig + pretrained_model_ckpt_path: ${oc.env:QWEN3_MOE_PATH,/dev/shm/Qwen3-30B-A3B} + model_revision: ${oc.env:QWEN3_MOE_REVISION,ad44e777bcd18fa416d9da3bd8f70d33ebb85d39} + tp_size: 4 + trust_remote_code: true + dtype: bfloat16 + moe_backend: triton + quantization: null + enforce_eager: true + enable_prefix_caching: false + enable_chunked_prefill: false + logprobs_mode: processed_logprobs + max_prompt_length: ${parity.max_prompt_length} + request_timeout_s: 7200 + startup_timeout_s: 300 + generate_timeout_s: 7200 + weight_update_timeout_s: 1800 + health_timeout_s: 60 + sleep_timeout_s: 180 + wake_up_timeout_s: 180 + shutdown_timeout_s: 120 + max_new_tokens: 1024 + temperature: 1.0 + top_p: 1.0 + top_k: 0 + # Verification-only: fixed-length responses avoid the EOS boundary. This is + # not a recommended general-purpose RL training default. + ignore_eos: true + chat_template_kwargs: + enable_thinking: true + engine_kwargs: + max_model_len: 1280 + max_num_batched_tokens: 1280 + max_num_seqs: 16 + gpu_memory_utilization: 0.6 + disable_log_stats: true + +reward: + _target_: unirl.reward.service.RewardService + truncated_reward: keep + backend: + _target_: unirl.reward.local.mathverify.MathVerifyRewardScorer + base_device: cpu + config: + _target_: unirl.reward.local.mathverify.MathVerifySpec + +algorithm: + _target_: experimental.train_inference_parity.algorithm.TrainInferenceParityGRPO + stage_attr: ar + clip_range: 0.2 + clip_range_high: 0.28 + clip_schedule: constant + loss_agg_mode: seq-mean-token-sum-norm + horizon: 1024 + sampling_temperature: 1.0 + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.qwen3.conditions.Qwen3ARConditions + +sync: + _target_: experimental.train_inference_parity.sync.ParityIPCWeightSync + lora_merged: false + bucket_size_mb: 640 + flush_cache: true + gpu_memory_headroom_mb: 2048 + control_timeout_s: 2400 + +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 1 + max_grad_norm: 1.0 + num_updates_per_batch: 1 + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH} + eval_data_path: ${oc.env:EVAL_DATA_PATH} + seed: null + shuffle: true + algorithm: + prompts_per_rollout: 4 + +sampling: + _target_: unirl.types.sampling.ARSamplingParams + samples_per_prompt: 8 + temperature: 1.0 + top_p: 1.0 + top_k: 0 + max_new_tokens: 1024 diff --git a/experimental/train_inference_parity/gate.py b/experimental/train_inference_parity/gate.py new file mode 100644 index 000000000..a1bfc597b --- /dev/null +++ b/experimental/train_inference_parity/gate.py @@ -0,0 +1,255 @@ +"""Collective-safe exact log-probability parity gate.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist + +_STATUS_SIZE = 8 + + +@dataclass(frozen=True) +class _LocalStatus: + shape_ok: bool + dtype_ok: bool + finite_ok: bool + mismatch_count: int + max_absdiff_fp32: float + token_count: int + k3_sum: float + k3_max: float + diagnostic: Dict[str, Any] + + def values(self) -> list[float]: + max_absdiff = self.max_absdiff_fp32 + if not math.isfinite(max_absdiff): + max_absdiff = math.inf + k3_sum = self.k3_sum + if not math.isfinite(k3_sum): + k3_sum = math.inf + k3_max = self.k3_max + if not math.isfinite(k3_max): + k3_max = math.inf + return [ + float(self.shape_ok), + float(self.dtype_ok), + float(self.finite_ok), + float(self.mismatch_count), + max_absdiff, + float(self.token_count), + k3_sum, + k3_max, + ] + + +def local_gate_status( + new_logp: Optional[torch.Tensor], + rollout_logp: Optional[torch.Tensor], + *, + local_error: Optional[str] = None, +) -> _LocalStatus: + """Build local status without raising before every rank reaches the collective.""" + diagnostic: Dict[str, Any] = {} + token_count = 0 + try: + new_shape = tuple(new_logp.shape) if isinstance(new_logp, torch.Tensor) else None + rollout_shape = tuple(rollout_logp.shape) if isinstance(rollout_logp, torch.Tensor) else None + token_count = int(new_logp.numel()) if isinstance(new_logp, torch.Tensor) else 0 + diagnostic.update( + { + "new_shape": new_shape, + "rollout_shape": rollout_shape, + "new_dtype": str(new_logp.dtype) if isinstance(new_logp, torch.Tensor) else None, + "rollout_dtype": str(rollout_logp.dtype) if isinstance(rollout_logp, torch.Tensor) else None, + "token_count": token_count, + } + ) + if local_error is not None: + diagnostic["error"] = str(local_error) + return _LocalStatus(False, False, False, 0, math.inf, token_count, math.inf, math.inf, diagnostic) + if not isinstance(new_logp, torch.Tensor) or not isinstance(rollout_logp, torch.Tensor): + diagnostic["error"] = "new_logp and rollout_logp must both be tensors" + return _LocalStatus(False, False, False, 0, math.inf, token_count, math.inf, math.inf, diagnostic) + + dtype_ok = new_logp.dtype == torch.float32 and rollout_logp.dtype == torch.float32 + diagnostic["dtype_ok"] = dtype_ok + if not dtype_ok: + diagnostic["error"] = "exact parity requires replay and the unquantized rollout anchor to both be FP32" + return _LocalStatus(False, False, False, 0, math.inf, token_count, math.inf, math.inf, diagnostic) + + new_fp32 = new_logp.detach() + rollout_fp32 = rollout_logp.detach().to(device=new_fp32.device) + new_finite = bool(torch.isfinite(new_fp32).all().item()) + rollout_finite = bool(torch.isfinite(rollout_fp32).all().item()) + shape_ok = new_fp32.shape == rollout_fp32.shape + diagnostic.update( + { + "new_finite": new_finite, + "rollout_finite": rollout_finite, + "shape_ok": shape_ok, + } + ) + if not shape_ok: + return _LocalStatus( + False, + dtype_ok, + new_finite and rollout_finite, + 0, + math.inf, + token_count, + math.inf, + math.inf, + diagnostic, + ) + + mismatch = new_fp32 != rollout_fp32 + mismatch_count = int(mismatch.sum().item()) + mismatch_indices = torch.nonzero(mismatch.reshape(-1), as_tuple=False).reshape(-1) + first_mismatch = int(mismatch_indices[0].item()) if mismatch_indices.numel() else None + if token_count: + absdiff = (new_fp32 - rollout_fp32).abs() + max_absdiff = float(absdiff.max().item()) + log_r = (new_fp32 - rollout_fp32).clamp(min=-20.0, max=20.0) + k3 = torch.expm1(log_r) - log_r + k3_sum = float(k3.sum().item()) + k3_max = float(k3.max().item()) + else: + max_absdiff = 0.0 + k3_sum = 0.0 + k3_max = 0.0 + diagnostic.update( + { + "mismatch_count": mismatch_count, + "first_mismatch": first_mismatch, + "max_absdiff_fp32": max_absdiff, + "first_new_fp32": ( + float(new_fp32.reshape(-1)[first_mismatch].item()) if first_mismatch is not None else None + ), + "first_rollout_fp32": ( + float(rollout_fp32.reshape(-1)[first_mismatch].item()) if first_mismatch is not None else None + ), + } + ) + return _LocalStatus( + True, + dtype_ok, + new_finite and rollout_finite, + mismatch_count, + max_absdiff, + token_count, + k3_sum, + k3_max, + diagnostic, + ) + except Exception as error: + diagnostic["error"] = f"{type(error).__name__}: {error}" + return _LocalStatus(False, False, False, 0, math.inf, token_count, math.inf, math.inf, diagnostic) + + +def _collective_device(new_logp: Optional[torch.Tensor]) -> torch.device: + if isinstance(new_logp, torch.Tensor) and new_logp.is_cuda: + return new_logp.device + # A multi-backend default group may report itself as ``undefined`` even + # though its CUDA backend is NCCL. Query that device backend directly: + # NCCL rejects CPU status rows, while a Gloo-only test should stay on CPU. + if torch.cuda.is_available() and dist.is_available() and dist.is_initialized(): + try: + group = dist.distributed_c10d._get_default_group() + cuda_backend = group._get_backend(torch.device("cuda")) + if "nccl" in type(cuda_backend).__name__.lower(): + return torch.device("cuda", torch.cuda.current_device()) + except RuntimeError: + pass + return torch.device("cpu") + + +def collective_fail_closed( + local: _LocalStatus, + *, + collective_tensor: Optional[torch.Tensor] = None, +) -> Dict[str, Any]: + """Collect one fixed status row per rank and make every rank fail together.""" + rows = [local.values()] + distributed = dist.is_available() and dist.is_initialized() + if distributed: + device = _collective_device(collective_tensor) + status = torch.tensor(local.values(), dtype=torch.float64, device=device) + gathered = [torch.empty(_STATUS_SIZE, dtype=status.dtype, device=device) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, status) + rows = torch.stack(gathered).cpu().tolist() + + shape_ok = all(bool(row[0]) for row in rows) + dtype_ok = all(bool(row[1]) for row in rows) + finite_ok = all(bool(row[2]) for row in rows) + mismatch_count = int(sum(row[3] for row in rows)) + max_absdiff = max(row[4] for row in rows) + token_count = int(sum(row[5] for row in rows)) + k3_sum = sum(row[6] for row in rows) + k3_max = max(row[7] for row in rows) + local_token_counts = [int(row[5]) for row in rows] + mixed_empty = any(count == 0 for count in local_token_counts) and any(count > 0 for count in local_token_counts) + failed = not shape_ok or not dtype_ok or not finite_ok or mismatch_count != 0 or mixed_empty + + if failed: + local_diagnostic = dict(local.diagnostic) + local_diagnostic["rank"] = dist.get_rank() if distributed else 0 + diagnostics = [local_diagnostic] + if distributed: + diagnostics = [None] * dist.get_world_size() + dist.all_gather_object(diagnostics, local_diagnostic) + raise RuntimeError( + "train/inference exact parity gate failed: " + f"shape_ok={shape_ok} dtype_ok={dtype_ok} finite_ok={finite_ok} " + f"mismatch_count={mismatch_count} max_absdiff_fp32={max_absdiff!r} " + f"token_count={token_count} local_token_counts={local_token_counts} " + f"diagnostics={diagnostics!r}" + ) + + return { + "token_count": token_count, + "per_rank_token_count": local_token_counts, + "torch_equal_fp32": True, + "mismatch_count": mismatch_count, + "max_absdiff_fp32": max_absdiff, + "k3_mean": k3_sum / token_count if token_count else 0.0, + "k3_max": k3_max, + } + + +def exact_parity_gate( + new_logp: Optional[torch.Tensor], + rollout_logp: Optional[torch.Tensor], + *, + local_error: Optional[str] = None, +) -> Dict[str, Any]: + """Build local status, synchronize it once, and fail closed before backward.""" + local = local_gate_status(new_logp, rollout_logp, local_error=local_error) + return collective_fail_closed(local, collective_tensor=new_logp) + + +def collective_error_barrier(local_error: Optional[str], *, phase: str) -> None: + """Synchronize setup errors before entering replay collectives; see README.md.""" + local = _LocalStatus( + shape_ok=local_error is None, + dtype_ok=local_error is None, + finite_ok=local_error is None, + mismatch_count=0 if local_error is None else 1, + max_absdiff_fp32=0.0 if local_error is None else math.inf, + token_count=0, + k3_sum=0.0 if local_error is None else math.inf, + k3_max=0.0 if local_error is None else math.inf, + diagnostic={"phase": str(phase), "error": None if local_error is None else str(local_error)}, + ) + collective_fail_closed(local) + + +__all__ = [ + "collective_error_barrier", + "collective_fail_closed", + "exact_parity_gate", + "local_gate_status", +] diff --git a/experimental/train_inference_parity/lifecycle.py b/experimental/train_inference_parity/lifecycle.py new file mode 100644 index 000000000..552ff3533 --- /dev/null +++ b/experimental/train_inference_parity/lifecycle.py @@ -0,0 +1,68 @@ +"""Explicit environment propagation and plugin discovery checks.""" + +from __future__ import annotations + +import importlib.metadata +import os +from typing import Mapping + +from .runtime_contract import ( + PLUGIN_DISTRIBUTION, + PLUGIN_ENTRYPOINT, + required_environment, +) + + +def apply_parity_environment(environment: Mapping[str, str]) -> None: + """Install the complete, explicit environment before validating or starting Ray.""" + expected_names = set(required_environment()) + missing = sorted(expected_names - set(environment)) + if missing: + raise ValueError(f"incomplete parity environment; missing {missing}") + for name, value in environment.items(): + os.environ[str(name)] = str(value) + + +def require_vllm_plugin_installed() -> None: + entries = importlib.metadata.entry_points().select(group="vllm.general_plugins") + names = {entry.name for entry in entries} + if PLUGIN_ENTRYPOINT not in names: + raise RuntimeError( + f"missing vLLM plugin entry point {PLUGIN_ENTRYPOINT!r}; install it with " + "`pip install -e experimental/train_inference_parity/vllm_plugin`" + ) + + +def connect_ray_with_parity_environment() -> None: + """Connect the driver and propagate parity variables to every Ray role.""" + import ray + + if ray.is_initialized(): + raise RuntimeError("Ray was initialized before the parity runtime_env could be installed") + explicit = set(required_environment()) + optional_passthrough = { + "CUDA_VISIBLE_DEVICES", + "DATA_PATH", + "EVAL_DATA_PATH", + "LD_LIBRARY_PATH", + "LD_PRELOAD", + "PYTHONPATH", + "PYTORCH_CUDA_ALLOC_CONF", + "QWEN3_MOE_PATH", + "QWEN3_MOE_REVISION", + "UNIRL_PARITY_ARTIFACT_PATH", + } + env_vars = {name: os.environ[name] for name in sorted(explicit | optional_passthrough) if name in os.environ} + ray.init( + address=os.environ.get("RAY_ADDRESS", "auto"), + runtime_env={"env_vars": env_vars}, + ) + + +__all__ = [ + "PLUGIN_DISTRIBUTION", + "PLUGIN_ENTRYPOINT", + "apply_parity_environment", + "connect_ray_with_parity_environment", + "require_vllm_plugin_installed", +] diff --git a/experimental/train_inference_parity/models/__init__.py b/experimental/train_inference_parity/models/__init__.py new file mode 100644 index 000000000..13ec053ae --- /dev/null +++ b/experimental/train_inference_parity/models/__init__.py @@ -0,0 +1 @@ +"""Model-specific parity experiment packages.""" diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/__init__.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/__init__.py new file mode 100644 index 000000000..c9905dc89 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/__init__.py @@ -0,0 +1,6 @@ +"""Qwen3-30B-A3B train/inference parity profile.""" + +from .actor import ParityQwen3Pipeline +from .contract import QWEN3_MOE_30B_A3B_CONTRACT + +__all__ = ["ParityQwen3Pipeline", "QWEN3_MOE_30B_A3B_CONTRACT"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/actor.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/actor.py new file mode 100644 index 000000000..fa215edc9 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/actor.py @@ -0,0 +1,20 @@ +"""Qwen3 pipeline wrapper that validates the experiment's frozen model contract.""" + +from __future__ import annotations + +from unirl.models.qwen3.pipeline import Qwen3Pipeline + +from .contract import validate_runtime_contract + + +class ParityQwen3Pipeline(Qwen3Pipeline): + @classmethod + def from_bundle(cls, bundle, **kwargs): + validate_runtime_contract( + bundle.transformer.config, + dtype=bundle.dtype, + ) + return super().from_bundle(bundle, **kwargs) + + +__all__ = ["ParityQwen3Pipeline"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/aten.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/aten.py new file mode 100644 index 000000000..cc1f138cc --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/aten.py @@ -0,0 +1,50 @@ +"""Irreversible ATen registration for a dedicated parity-training process.""" + +from __future__ import annotations + +import torch + +from .context import require_parity_opt_in +from .linear_ops import linear_backward_cuda, linear_cuda + +_LIBRARY = None + + +def _cuda_registration(operator: str) -> str | None: + return next( + (line for line in torch._C._dispatch_dump_table(operator).splitlines() if line.startswith("CUDA:")), + None, + ) + + +def preflight_aten() -> None: + """Reject another extension's process-global CUDA overrides.""" + conflicts = {} + for operator in ("aten::linear", "aten::linear_backward"): + registration = _cuda_registration(operator) + if registration is not None and not any( + marker in registration for marker in ("/pytorch/", "site-packages/torch/") + ): + conflicts[operator] = registration + if conflicts: + raise RuntimeError(f"unexpected pre-existing CUDA linear providers: {conflicts}") + + +def register_aten() -> None: + """Register guarded linear kernels for the lifetime of this process.""" + global _LIBRARY + require_parity_opt_in() + if _LIBRARY is not None: + return + preflight_aten() + library = torch.library.Library("aten", "IMPL") + library.impl("aten::linear", linear_cuda, "CUDA") + library.impl("aten::linear_backward", linear_backward_cuda, "CUDA") + _LIBRARY = library + + +def aten_registered() -> bool: + return _LIBRARY is not None + + +__all__ = ["aten_registered", "preflight_aten", "register_aten"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/attention.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/attention.py new file mode 100644 index 000000000..70ff10df0 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/attention.py @@ -0,0 +1,215 @@ +"""vLLM FA3 facade with deterministic forward and finite surrogate backward.""" + +from __future__ import annotations + +import functools +import importlib +import importlib.machinery +import math +import types + +import torch +import torch.nn.functional as F + +from .context import exact_enabled + + +@functools.lru_cache(maxsize=1) +def probe_fa3_backend(): + """Import and validate the real bundled backend before advertising FA3.""" + try: + backend = importlib.import_module("vllm.vllm_flash_attn") + except Exception as error: + raise RuntimeError("Qwen3 parity FA3 requires importable vllm.vllm_flash_attn") from error + if not callable(getattr(backend, "flash_attn_varlen_func", None)): + raise RuntimeError("vllm.vllm_flash_attn does not expose flash_attn_varlen_func") + return backend + + +def fa3_available() -> bool: + try: + probe_fa3_backend() + except RuntimeError: + return False + return torch.cuda.is_available() + + +def _require_bfloat16(*tensors: torch.Tensor) -> None: + invalid = [str(tensor.dtype) for tensor in tensors if tensor.dtype != torch.bfloat16] + if invalid: + raise ValueError(f"Qwen3 parity FA3 requires BF16 tensors, got {invalid}") + + +def _math_attention_backward(q, k, v, grad_output, *, scale, causal): + with torch.enable_grad(): + qr = q.detach().requires_grad_(q.requires_grad) + kr = k.detach().requires_grad_(k.requires_grad) + vr = v.detach().requires_grad_(v.requires_grad) + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH): + output = F.scaled_dot_product_attention( + qr.transpose(1, 2), + kr.transpose(1, 2), + vr.transpose(1, 2), + is_causal=causal, + scale=scale, + enable_gqa=qr.shape[-2] != kr.shape[-2], + ).transpose(1, 2) + gradients = torch.autograd.grad( + output, + (qr, kr, vr), + grad_output, + allow_unused=True, + ) + return gradients + + +class _Fa3Fixed(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, scale, causal): + _require_bfloat16(q, k, v) + flash_attn_varlen_func = probe_fa3_backend().flash_attn_varlen_func + batch, sequence = q.shape[:2] + cu = torch.arange( + 0, + (batch + 1) * sequence, + sequence, + dtype=torch.int32, + device=q.device, + ) + output = flash_attn_varlen_func( + q=q.reshape(-1, *q.shape[2:]), + k=k.reshape(-1, *k.shape[2:]), + v=v.reshape(-1, *v.shape[2:]), + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=sequence, + max_seqlen_k=sequence, + softmax_scale=scale, + causal=causal, + deterministic=True, + num_splits=1, + fa_version=3, + ).view_as(q) + ctx.save_for_backward(q, k, v) + ctx.scale, ctx.causal = scale, causal + return output + + @staticmethod + def backward(ctx, grad_output): + return ( + *_math_attention_backward( + *ctx.saved_tensors, + grad_output, + scale=ctx.scale, + causal=ctx.causal, + ), + None, + None, + ) + + +class _Fa3Varlen(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, cu_q, cu_k, max_q, max_k, scale, causal): + _require_bfloat16(q, k, v) + flash_attn_varlen_func = probe_fa3_backend().flash_attn_varlen_func + output = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + softmax_scale=scale, + causal=causal, + deterministic=True, + num_splits=1, + fa_version=3, + ) + ctx.save_for_backward(q, k, v, cu_q, cu_k) + ctx.scale, ctx.causal = scale, causal + return output + + @staticmethod + def backward(ctx, grad_output): + q, k, v, cu_q, cu_k = ctx.saved_tensors + q_parts, k_parts, v_parts = [], [], [] + for index in range(int(cu_q.numel()) - 1): + qb, qe = int(cu_q[index]), int(cu_q[index + 1]) + kb, ke = int(cu_k[index]), int(cu_k[index + 1]) + gradients = _math_attention_backward( + q[qb:qe].unsqueeze(0), + k[kb:ke].unsqueeze(0), + v[kb:ke].unsqueeze(0), + grad_output[qb:qe].unsqueeze(0), + scale=ctx.scale, + causal=ctx.causal, + ) + q_parts.append(gradients[0].squeeze(0)) + k_parts.append(gradients[1].squeeze(0)) + v_parts.append(gradients[2].squeeze(0)) + return ( + torch.cat(q_parts), + torch.cat(k_parts), + torch.cat(v_parts), + None, + None, + None, + None, + None, + None, + ) + + +def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False, **_kwargs): + if not exact_enabled(): + raise RuntimeError("Qwen3 parity FA3 facade may only run inside stage.exact_context()") + if dropout_p: + raise ValueError("parity FA3 requires dropout=0") + scale = float(softmax_scale or (1.0 / math.sqrt(q.shape[-1]))) + return _Fa3Fixed.apply(q, k, v, scale, bool(causal)) + + +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + softmax_scale=None, + causal=False, + **kwargs, +): + if not exact_enabled(): + raise RuntimeError("Qwen3 parity FA3 facade may only run inside stage.exact_context()") + if float(kwargs.get("dropout_p", 0.0) or 0.0): + raise ValueError("parity FA3 requires dropout=0") + scale = float(softmax_scale or (1.0 / math.sqrt(q.shape[-1]))) + return _Fa3Varlen.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + int(max_seqlen_q), + int(max_seqlen_k), + scale, + bool(causal), + ) + + +def make_flash_attn_interface() -> types.ModuleType: + """Build the Transformers-facing facade after the backend probe passed.""" + probe_fa3_backend() + module = types.ModuleType("flash_attn_interface") + module.__spec__ = importlib.machinery.ModuleSpec("flash_attn_interface", loader=None) + module.flash_attn_func = flash_attn_func + module.flash_attn_varlen_func = flash_attn_varlen_func + module.flash_attn_with_kvcache = None + return module + + +__all__ = ["fa3_available", "make_flash_attn_interface", "probe_fa3_backend"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/context.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/context.py new file mode 100644 index 000000000..7f46548ea --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/context.py @@ -0,0 +1,49 @@ +"""Process-local exact-mode state for the Qwen3 parity actor.""" + +from __future__ import annotations + +import contextlib +import contextvars +import os +from collections.abc import Iterator + +_EXACT = contextvars.ContextVar("unirl_parity_exact", default=False) +_PROCESS_EXACT_DEPTH = 0 + + +def require_parity_opt_in() -> None: + """Fail closed unless this process was dedicated to parity overrides.""" + if os.environ.get("UNIRL_PARITY_ENABLE") != "1": + raise RuntimeError( + "Qwen3 parity overrides require UNIRL_PARITY_ENABLE=1. Set it only for a dedicated parity-training process." + ) + + +def exact_enabled() -> bool: + # Non-reentrant HF/FSDP checkpointing may replay under a copied Context + # whose ContextVar value predates the outer algorithm scope. The process is + # parity-dedicated, so the outer scope also holds a nested process counter. + return bool(_EXACT.get()) or _PROCESS_EXACT_DEPTH > 0 + + +@contextlib.contextmanager +def exact_mode(enabled: bool = True) -> Iterator[None]: + global _PROCESS_EXACT_DEPTH + token = _EXACT.set(bool(enabled)) + increment = int(bool(enabled)) + _PROCESS_EXACT_DEPTH += increment + try: + yield + finally: + _PROCESS_EXACT_DEPTH -= increment + if _PROCESS_EXACT_DEPTH < 0: + raise AssertionError("Qwen3 parity exact-mode process depth became negative") + _EXACT.reset(token) + + +def exact_context(): + """Enter exact mode; installation remains an explicit, separate action.""" + return exact_mode(True) + + +__all__ = ["exact_context", "exact_enabled", "exact_mode", "require_parity_opt_in"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/contract.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/contract.py new file mode 100644 index 000000000..f2a1248ab --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/contract.py @@ -0,0 +1,115 @@ +"""Frozen structural contract for the first parity model profile.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class Qwen3MoeContract: + model_name: str + hidden_size: int + attention_heads: int + key_value_heads: int + head_dim: int + experts: int + topk: int + intermediate_size: int + vocab_size: int + dtype: str + + +QWEN3_MOE_30B_A3B_CONTRACT = Qwen3MoeContract( + model_name="Qwen3-30B-A3B", + hidden_size=2048, + attention_heads=32, + key_value_heads=4, + head_dim=128, + experts=128, + topk=8, + intermediate_size=768, + vocab_size=151936, + dtype="bfloat16", +) + + +def _is_bfloat16(dtype: object) -> bool: + if dtype is torch.bfloat16: + return True + return str(dtype).lower().removeprefix("torch.") in {"bf16", "bfloat16"} + + +def require_train_tp() -> int: + """Read the explicit training TP contract; there is intentionally no default.""" + raw = os.environ.get("UNIRL_PARITY_TRAIN_TP") + if raw is None or not raw.strip(): + raise RuntimeError("Qwen3 parity requires UNIRL_PARITY_TRAIN_TP to be set explicitly") + try: + train_tp = int(raw) + except ValueError as error: + raise ValueError(f"UNIRL_PARITY_TRAIN_TP must be a positive integer, got {raw!r}") from error + if train_tp <= 0: + raise ValueError(f"UNIRL_PARITY_TRAIN_TP must be a positive integer, got {raw!r}") + return train_tp + + +def validate_parallel_dimensions(hidden_size: int, intermediate_size: int, train_tp: int) -> None: + if train_tp <= 0: + raise ValueError(f"train_tp must be positive, got {train_tp}") + mismatches = {} + if hidden_size % train_tp: + mismatches["hidden_size"] = (hidden_size, train_tp) + if intermediate_size % train_tp: + mismatches["intermediate_size"] = (intermediate_size, train_tp) + if mismatches: + raise ValueError(f"Qwen3 parity dimensions must be divisible by train_tp: {mismatches}") + + +def validate_model_config(config, *, dtype: object | None = None) -> None: + contract = QWEN3_MOE_30B_A3B_CONTRACT + checks = { + "hidden_size": contract.hidden_size, + "num_attention_heads": contract.attention_heads, + "num_key_value_heads": contract.key_value_heads, + "head_dim": contract.head_dim, + "num_experts": contract.experts, + "num_experts_per_tok": contract.topk, + "moe_intermediate_size": contract.intermediate_size, + "vocab_size": contract.vocab_size, + } + mismatches = { + name: (getattr(config, name, None), expected) + for name, expected in checks.items() + if getattr(config, name, None) != expected + } + if mismatches: + raise ValueError(f"Qwen3-30B-A3B parity contract mismatch: {mismatches}") + resolved_dtype = dtype + if resolved_dtype is None: + resolved_dtype = getattr(config, "dtype", getattr(config, "torch_dtype", None)) + if not _is_bfloat16(resolved_dtype): + raise ValueError(f"Qwen3-30B-A3B parity requires BF16 model dtype, got {resolved_dtype!r}") + + +def validate_runtime_contract(config, *, dtype: object) -> int: + validate_model_config(config, dtype=dtype) + train_tp = require_train_tp() + validate_parallel_dimensions( + int(getattr(config, "hidden_size")), + int(getattr(config, "moe_intermediate_size")), + train_tp, + ) + return train_tp + + +__all__ = [ + "QWEN3_MOE_30B_A3B_CONTRACT", + "Qwen3MoeContract", + "require_train_tp", + "validate_model_config", + "validate_parallel_dimensions", + "validate_runtime_contract", +] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/fsdp.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/fsdp.py new file mode 100644 index 000000000..21d9e663e --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/fsdp.py @@ -0,0 +1,5 @@ +"""Thin compatibility facade for the split Qwen3 parity FSDP adapter.""" + +from .install import install + +__all__ = ["install"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/install.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/install.py new file mode 100644 index 000000000..a2aea31a5 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/install.py @@ -0,0 +1,216 @@ +"""Explicit installer for process-wide Qwen3 parity overrides.""" + +from __future__ import annotations + +import inspect +import sys +from collections.abc import MutableMapping +from dataclasses import dataclass +from typing import Any + +from unirl.models.qwen3.ar import register_exact_actor_provider + +from .aten import preflight_aten, register_aten +from .attention import fa3_available, make_flash_attn_interface, probe_fa3_backend +from .context import exact_context, require_parity_opt_in +from .linear_ops import exact_log_softmax +from .moe import sparse_moe_forward +from .norm import rmsnorm_forward +from .rope import rotary_forward +from .router import router_forward + +_INSTALLED = False + + +@dataclass +class _AttributePatch: + target: Any + name: str + existed_directly: bool + original: Any + + +@dataclass +class _MappingPatch: + target: MutableMapping + key: Any + existed: bool + original: Any + + +_PYTHON_PATCHES: list[_AttributePatch | _MappingPatch] = [] + + +def _patch_attribute(target: Any, name: str, replacement: Any) -> None: + _PYTHON_PATCHES.append( + _AttributePatch( + target=target, + name=name, + existed_directly=name in vars(target), + original=getattr(target, name, None), + ) + ) + setattr(target, name, replacement) + + +def _patch_mapping(target: MutableMapping, key: Any, replacement: Any) -> None: + _PYTHON_PATCHES.append( + _MappingPatch( + target=target, + key=key, + existed=key in target, + original=target.get(key), + ) + ) + target[key] = replacement + + +def _clear_cache(value: Any) -> None: + clear = getattr(value, "cache_clear", None) + if callable(clear): + clear() + + +def _install_fa3_python_patches() -> None: + probe_fa3_backend() + _patch_mapping(sys.modules, "flash_attn_interface", make_flash_attn_interface()) + + import transformers.modeling_flash_attention_utils as flash_utils + import transformers.modeling_utils as modeling_utils + import transformers.utils as transformers_utils + import transformers.utils.import_utils as import_utils + + previous = getattr(import_utils, "is_flash_attn_3_available", None) + _clear_cache(previous) + + def available() -> bool: + return fa3_available() + + available.cache_clear = lambda: None + _patch_attribute(import_utils, "is_flash_attn_3_available", available) + _patch_attribute(transformers_utils, "is_flash_attn_3_available", available) + _patch_attribute(flash_utils, "is_flash_attn_3_available", available) + + matrix = modeling_utils.FLASH_ATTENTION_COMPATIBILITY_MATRIX[3] + _patch_mapping(matrix, "general_availability_check", available) + _patch_mapping(matrix, "pkg_availability_check", lambda *args, **kwargs: fa3_available()) + + +def _install_qwen_python_patches(modeling) -> None: + original_attention = modeling.Qwen3MoeAttention.forward + + def attention_forward( + self, + hidden_states, + position_embeddings, + attention_mask, + past_key_values=None, + **kwargs, + ): + # Transformers 5.6 calls ``s_aux.to(...)`` unconditionally in its FA + # integration. Qwen3-MoE has no attention sinks; the parity facade + # ignores this empty placeholder. + kwargs.setdefault("s_aux", hidden_states.new_empty(0)) + return original_attention( + self, + hidden_states, + position_embeddings, + attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + + setattr(router_forward, "_unirl_parity_original", modeling.Qwen3MoeTopKRouter.forward) + setattr(rmsnorm_forward, "_unirl_parity_original", modeling.Qwen3MoeRMSNorm.forward) + setattr(sparse_moe_forward, "_unirl_parity_original", modeling.Qwen3MoeSparseMoeBlock.forward) + setattr(rotary_forward, "_unirl_parity_original", modeling.Qwen3MoeRotaryEmbedding.forward) + setattr(attention_forward, "_unirl_parity_original", original_attention) + _patch_attribute(modeling.Qwen3MoeAttention, "forward", attention_forward) + _patch_attribute(modeling.Qwen3MoeTopKRouter, "forward", router_forward) + _patch_attribute(modeling.Qwen3MoeRMSNorm, "forward", rmsnorm_forward) + _patch_attribute(modeling.Qwen3MoeSparseMoeBlock, "forward", sparse_moe_forward) + _patch_attribute(modeling.Qwen3MoeRotaryEmbedding, "forward", rotary_forward) + + +def restore_python_patches() -> None: + """Restore Python objects while leaving irreversible ATen registrations active.""" + global _INSTALLED + while _PYTHON_PATCHES: + patch = _PYTHON_PATCHES.pop() + if isinstance(patch, _MappingPatch): + if patch.existed: + patch.target[patch.key] = patch.original + else: + patch.target.pop(patch.key, None) + continue + if patch.existed_directly: + setattr(patch.target, patch.name, patch.original) + else: + delattr(patch.target, patch.name) + _clear_cache(patch.original) + _INSTALLED = False + + +def install() -> None: + """Install the adapter after an explicit dedicated-process opt-in.""" + global _INSTALLED + require_parity_opt_in() + if _INSTALLED: + return + + # Probe all external Python surfaces before the irreversible ATen step. + from unirl_train_inference_parity_vllm.common.providers import preflight_providers + from unirl_train_inference_parity_vllm.compat import validate_runtime_versions + + validate_runtime_versions() + preflight_providers(strict=True) + probe_fa3_backend() + from transformers.models.qwen3_moe import modeling_qwen3_moe as modeling + + method_contracts = { + "Qwen3MoeAttention": ( + "self", + "hidden_states", + "position_embeddings", + "attention_mask", + "past_key_values", + "kwargs", + ), + "Qwen3MoeTopKRouter": ("self", "hidden_states"), + "Qwen3MoeRMSNorm": ("self", "hidden_states"), + "Qwen3MoeSparseMoeBlock": ("self", "hidden_states"), + "Qwen3MoeRotaryEmbedding": ("self", "x", "position_ids"), + } + for class_name, expected_parameters in method_contracts.items(): + if not hasattr(modeling, class_name): + raise RuntimeError(f"transformers Qwen3 MoE class is missing: {class_name}") + method = getattr(getattr(modeling, class_name), "forward", None) + if not callable(method): + raise RuntimeError(f"transformers Qwen3 MoE class has no callable forward: {class_name}") + actual_parameters = tuple(inspect.signature(method).parameters) + if actual_parameters != expected_parameters: + raise RuntimeError( + f"transformers Qwen3 MoE signature drift for {class_name}.forward: " + f"expected {expected_parameters}, got {actual_parameters}" + ) + preflight_aten() + + try: + _install_fa3_python_patches() + _install_qwen_python_patches(modeling) + register_aten() + register_exact_actor_provider( + context_factory=exact_context, + log_softmax=exact_log_softmax, + ) + except Exception: + restore_python_patches() + raise + _INSTALLED = True + + +def installed() -> bool: + return _INSTALLED + + +__all__ = ["install", "installed", "restore_python_patches"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/linear_ops.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/linear_ops.py new file mode 100644 index 000000000..a7b4b7b43 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/linear_ops.py @@ -0,0 +1,90 @@ +"""Guarded linear and direct reduction providers for exact parity.""" + +from __future__ import annotations + +import torch + +from .context import exact_enabled + + +def providers(): + from unirl_train_inference_parity_vllm.common import providers as implementations + + return implementations + + +def native_linear(input, weight, bias=None): + output = torch.matmul(input, weight.t()) + return output if bias is None else output + bias + + +def linear_cuda(input, weight, bias=None): + """Use the BI provider only inside an exact BF16 CUDA region.""" + if not ( + exact_enabled() + and input.is_cuda + and weight.is_cuda + and input.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and input.ndim in (2, 3) + and weight.ndim == 2 + ): + return native_linear(input, weight, bias) + original_shape = input.shape + output = providers().linear( + input.reshape(-1, original_shape[-1]), + weight, + bias, + ) + return output.reshape(*original_shape[:-1], output.shape[-1]) + + +def linear_backward_cuda(input, grad_output, weight, output_mask): + input_2d = input.reshape(-1, input.shape[-1]) + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + grad_input = ( + torch.matmul(grad_output, weight) if output_mask[0] else torch.empty(0, device=input.device, dtype=input.dtype) + ) + grad_weight = ( + torch.matmul(grad_2d.t(), input_2d) + if output_mask[1] + else torch.empty(0, device=weight.device, dtype=weight.dtype) + ) + grad_bias = ( + grad_2d.sum(dim=0) if output_mask[2] else torch.empty(0, device=grad_output.device, dtype=grad_output.dtype) + ) + return grad_input, grad_weight, grad_bias + + +class _LogSoftmax(torch.autograd.Function): + @staticmethod + def forward(ctx, input, dim): + output = providers().log_softmax(input, dim=int(dim)) + ctx.save_for_backward(output) + ctx.dim = int(dim) + return output + + @staticmethod + def backward(ctx, grad_output): + (output,) = ctx.saved_tensors + grad_input = grad_output - output.exp() * grad_output.sum( + dim=ctx.dim, + keepdim=True, + ) + return grad_input, None + + +def exact_log_softmax(input, dim=-1): + """Bitwise selected-token forward provider with an analytic gradient.""" + if not exact_enabled(): + raise RuntimeError("Qwen3 parity log-softmax may only run inside stage.exact_context()") + return _LogSoftmax.apply(input, dim) + + +__all__ = [ + "exact_log_softmax", + "linear_backward_cuda", + "linear_cuda", + "native_linear", + "providers", +] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/moe.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/moe.py new file mode 100644 index 000000000..28d5b859d --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/moe.py @@ -0,0 +1,164 @@ +"""Exact selected-token MoE forward with a Torch surrogate backward.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from .context import exact_enabled +from .contract import require_train_tp, validate_parallel_dimensions +from .linear_ops import providers + + +def _dispatch(hidden, topk_ids, experts): + tokens, topk = topk_ids.shape + flat_ids = topk_ids.reshape(-1).long() + order = torch.argsort(flat_ids, stable=True) + rows = ( + hidden[:, None, :] + .expand(tokens, topk, hidden.shape[-1]) + .reshape(-1, hidden.shape[-1]) + .index_select(0, order) + .contiguous() + ) + counts = torch.bincount(flat_ids, minlength=experts) + return rows, counts, order + + +def _grouped_linear(inputs, weights, counts, *, public): + outputs = [] + offset = 0 + for expert in range(int(weights.shape[0])): + count = int(counts[expert].item()) + if count: + rows = inputs[offset : offset + count].contiguous() + outputs.append( + providers().linear(rows, weights[expert], None) if public else F.linear(rows, weights[expert]) + ) + offset += count + return torch.cat(outputs, dim=0) + + +def _combine(contributions, order, tokens, topk): + from unirl_train_inference_parity_vllm.common.moe_combine import moe_combine + + row_map = torch.empty(tokens * topk, dtype=torch.long, device=order.device) + row_map[order] = torch.arange(tokens * topk, device=order.device) + return moe_combine(contributions, row_map.view(tokens, topk)) + + +def _validate_bfloat16(*tensors: torch.Tensor) -> None: + invalid = [str(tensor.dtype) for tensor in tensors if tensor.dtype != torch.bfloat16] + if invalid: + raise ValueError(f"Qwen3 parity MoE requires BF16 tensors, got {invalid}") + + +def _moe_forward(hidden, topk_ids, topk_weights, w13, w2, *, public): + _validate_bfloat16(hidden, w13, w2) + tokens, hidden_size = hidden.shape + topk = int(topk_ids.shape[1]) + experts = int(w13.shape[0]) + if int(w13.shape[1]) % 2: + raise ValueError(f"Qwen3 parity gate/up width must be even, got {int(w13.shape[1])}") + intermediate = int(w13.shape[1]) // 2 + if tuple(w2.shape) != (experts, hidden_size, intermediate): + raise ValueError( + "Qwen3 parity down-projection shape mismatch: " + f"got {tuple(w2.shape)}, expected {(experts, hidden_size, intermediate)}" + ) + train_tp = require_train_tp() + validate_parallel_dimensions(hidden_size, intermediate, train_tp) + dispatched, counts, order = _dispatch(hidden, topk_ids, experts) + local_intermediate = intermediate // train_tp + fc1_parts = [] + for rank in range(train_tp): + begin, end = rank * local_intermediate, (rank + 1) * local_intermediate + rank_weight = torch.cat( + (w13[:, begin:end], w13[:, intermediate + begin : intermediate + end]), + dim=1, + ).contiguous() + fc1_parts.append(_grouped_linear(dispatched, rank_weight, counts, public=public)) + packed = torch.cat(fc1_parts, dim=-1).view(-1, train_tp, 2, local_intermediate) + activation = (F.silu(packed[:, :, 0]) * packed[:, :, 1]).reshape(-1, intermediate) + local_hidden = hidden_size // train_tp + fc2_parts = [] + for rank in range(train_tp): + rank_weight = w2[ + :, + rank * local_hidden : (rank + 1) * local_hidden, + :, + ].contiguous() + fc2_parts.append(_grouped_linear(activation, rank_weight, counts, public=public)) + fc2 = torch.cat(fc2_parts, dim=-1) + sorted_weights = topk_weights.reshape(-1).index_select(0, order) + contributions = (fc2 * sorted_weights[:, None]).to(torch.bfloat16) + return _combine(contributions, order, tokens, topk) + + +class _Moe(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, topk_ids, topk_weights, w13, w2): + ctx.save_for_backward(hidden, topk_ids, topk_weights, w13, w2) + return _moe_forward( + hidden, + topk_ids, + topk_weights, + w13, + w2, + public=True, + ) + + @staticmethod + def backward(ctx, grad_output): + hidden, topk_ids, topk_weights, w13, w2 = ctx.saved_tensors + needs = ctx.needs_input_grad + sources = (hidden, topk_weights, w13, w2) + rebuilt = [ + tensor.detach().requires_grad_(needed) + for tensor, needed in zip( + sources, + (needs[0], needs[2], needs[3], needs[4]), + strict=True, + ) + ] + with torch.enable_grad(): + output = _moe_forward( + rebuilt[0], + topk_ids, + rebuilt[1], + rebuilt[2], + rebuilt[3], + public=False, + ) + requested = [tensor for tensor in rebuilt if tensor.requires_grad] + computed = torch.autograd.grad( + output, + requested, + grad_output, + allow_unused=True, + ) + iterator = iter(computed) + gradients = [next(iterator) if tensor.requires_grad else None for tensor in rebuilt] + return gradients[0], None, gradients[1], gradients[2], gradients[3] + + +def sparse_moe_forward(self, hidden_states): + if not exact_enabled(): + original = getattr(sparse_moe_forward, "_unirl_parity_original", None) + if not callable(original): + raise RuntimeError("Qwen3 parity MoE original is unavailable") + return original(self, hidden_states) + shape = hidden_states.shape + hidden = hidden_states.reshape(-1, shape[-1]) + _logits, weights, indices = self.gate(hidden) + output = _Moe.apply( + hidden, + indices, + weights, + self.experts.gate_up_proj, + self.experts.down_proj, + ) + return output.view(shape) + + +__all__ = ["sparse_moe_forward"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/norm.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/norm.py new file mode 100644 index 000000000..499f19070 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/norm.py @@ -0,0 +1,62 @@ +"""Exact RMSNorm forward with a finite Torch surrogate gradient.""" + +from __future__ import annotations + +import torch + +from .context import exact_enabled +from .linear_ops import providers + + +class _RmsNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, weight, eps): + ctx.save_for_backward(hidden, weight) + ctx.eps = float(eps) + return providers().rms_norm(hidden, weight.to(hidden.dtype), float(eps)) + + @staticmethod + def backward(ctx, grad_output): + hidden, weight = ctx.saved_tensors + needs_hidden, needs_weight = ctx.needs_input_grad[:2] + with torch.enable_grad(): + hidden_ref = hidden.detach().requires_grad_(needs_hidden) + weight_ref = weight.detach().requires_grad_(needs_weight) + variance = hidden_ref.float().square().mean(dim=-1, keepdim=True) + output = (hidden_ref.float() * torch.rsqrt(variance + ctx.eps) * weight_ref.float()).to(hidden.dtype) + requested = [ + tensor + for tensor, needed in ( + (hidden_ref, needs_hidden), + (weight_ref, needs_weight), + ) + if needed + ] + computed = torch.autograd.grad( + output, + requested, + grad_output, + allow_unused=True, + ) + iterator = iter(computed) + return ( + next(iterator) if needs_hidden else None, + next(iterator) if needs_weight else None, + None, + ) + + +def rmsnorm_forward(self, hidden_states): + if not exact_enabled(): + original = getattr(rmsnorm_forward, "_unirl_parity_original", None) + if not callable(original): + raise RuntimeError("Qwen3 parity RMSNorm original is unavailable") + return original(self, hidden_states) + return _RmsNorm.apply( + hidden_states, + self.weight, + float(self.variance_epsilon), + ) + + +__all__ = ["rmsnorm_forward"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/rope.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/rope.py new file mode 100644 index 000000000..80af8a580 --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/rope.py @@ -0,0 +1,34 @@ +"""Deterministic Qwen3 rotary embedding implementation.""" + +from __future__ import annotations + +import torch + +from .context import exact_enabled + + +def rotary_forward(self, x, position_ids): + if not exact_enabled(): + original = getattr(rotary_forward, "_unirl_parity_original", None) + if not callable(original): + raise RuntimeError("Qwen3 parity RoPE original is unavailable") + return original(self, x, position_ids) + base = self.config.rope_parameters["rope_theta"] + dimension = getattr(self.config, "head_dim", None) or (self.config.hidden_size // self.config.num_attention_heads) + inverse = 1.0 / (base ** (torch.arange(0, dimension, 2, dtype=torch.float32, device=x.device) / dimension)) + expanded_frequency = inverse[None, :, None].expand( + position_ids.shape[0], + -1, + 1, + ) + expanded_positions = position_ids[:, None, :].float() + with torch.autocast(device_type=x.device.type, enabled=False): + frequencies = (expanded_frequency.float() @ expanded_positions.float()).transpose(1, 2) + embedding = torch.cat((frequencies, frequencies), dim=-1) + return ( + (embedding.cos() * self.attention_scaling).to(x.dtype), + (embedding.sin() * self.attention_scaling).to(x.dtype), + ) + + +__all__ = ["rotary_forward"] diff --git a/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/router.py b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/router.py new file mode 100644 index 000000000..a9cc8172c --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/router.py @@ -0,0 +1,27 @@ +"""Qwen3 MoE router using an explicit parity softmax provider.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from .context import exact_enabled +from .linear_ops import providers + + +def router_forward(self, hidden_states): + if not exact_enabled(): + original = getattr(router_forward, "_unirl_parity_original", None) + if not callable(original): + raise RuntimeError("Qwen3 parity router original is unavailable") + return original(self, hidden_states) + hidden = hidden_states.reshape(-1, self.hidden_dim) + logits = F.linear(hidden, self.weight) + probabilities = providers().softmax(logits.float(), dim=-1) + weights, indices = torch.topk(probabilities, self.top_k, dim=-1) + if self.norm_topk_prob: + weights = weights / weights.sum(dim=-1, keepdim=True) + return logits, weights.float(), indices + + +__all__ = ["router_forward"] diff --git a/experimental/train_inference_parity/run.py b/experimental/train_inference_parity/run.py new file mode 100644 index 000000000..f20d33706 --- /dev/null +++ b/experimental/train_inference_parity/run.py @@ -0,0 +1,77 @@ +"""Hydra entry point for the train/inference parity experiment.""" + +from __future__ import annotations + +import hydra +from omegaconf import DictConfig + +from experimental.train_inference_parity.lifecycle import ( + apply_parity_environment, + connect_ray_with_parity_environment, + require_vllm_plugin_installed, +) +from experimental.train_inference_parity.runtime_contract import ( + ParityRuntimeContract, +) +from experimental.train_inference_parity.verification import ParityVerification +from unirl.trainer.ar import ARTrainer +from unirl.utils.graceful_shutdown import GracefulShutdown + + +@hydra.main( + version_base=None, + config_path="examples", + config_name="qwen3_moe_30b_a3b_fsdp_tp4", +) +def main(cfg: DictConfig) -> None: + contract = ParityRuntimeContract.from_config(cfg) + apply_parity_environment(contract.environment) + contract.validate() + verification = ParityVerification(contract) + verification.initialize() + trainer = None + + def teardown() -> None: + if trainer is not None: + trainer.shutdown() + + try: + require_vllm_plugin_installed() + connect_ray_with_parity_environment() + with GracefulShutdown(teardown, name="train-inference-parity") as guard: + trainer = ARTrainer( + cfg=cfg, + batch_size=cfg.batch_size, + bundle_cfg=cfg.bundle, + pipeline_cfg=cfg.pipeline, + backend_cfg=cfg.backend, + rollout_cfg=cfg.rollout, + reward_cfg=cfg.reward, + algorithm_cfg=cfg.algorithm, + stack_cfg=cfg.stack, + data_source_cfg=cfg.data_source, + sampling_cfg=cfg.sampling, + sync_cfg=cfg.get("sync"), + logging_cfg=cfg.get("logging"), + adv_normalization_scope=cfg.get("adv_normalization_scope", "group"), + normalize_adv_by_std=cfg.get("normalize_adv_by_std", True), + enable_fsdp_offload=cfg.get("enable_fsdp_offload", True), + ) + verification.bind(trainer) + guard.claim_signals() + trainer.train( + num_rollouts=cfg.get("num_rollouts", 2), + weight_sync_interval=cfg.get("weight_sync_interval", 1), + on_rollout_complete=verification.on_rollout_complete, + ) + verification.finalize() + except BaseException as error: + try: + verification.record_failure("run", error) + except BaseException as artifact_error: + error.add_note(f"could not persist parity failure artifact: {artifact_error!r}") + raise + + +if __name__ == "__main__": + main() diff --git a/experimental/train_inference_parity/runtime_contract.py b/experimental/train_inference_parity/runtime_contract.py new file mode 100644 index 000000000..d33915693 --- /dev/null +++ b/experimental/train_inference_parity/runtime_contract.py @@ -0,0 +1,433 @@ +"""Fail-closed runtime contract for the public TP4 parity recipe.""" + +from __future__ import annotations + +import importlib.metadata +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Mapping + +from omegaconf import DictConfig, OmegaConf + +PUBLIC_REFERENCE = "public_reference" +EXPECTED_TP = 4 +EXPECTED_VERSIONS = { + "torch": "2.13.0", + "transformers": "5.6.0", + "vllm": "0.27.0", +} +EXPECTED_MODEL_REVISION = "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" +EXPECTED_SYNC_TARGET = "experimental.train_inference_parity.sync.ParityIPCWeightSync" +PLUGIN_DISTRIBUTION = "unirl-train-inference-parity-vllm" +PLUGIN_ENTRYPOINT = "unirl_train_inference_parity" + +_REQUIRED_ENVIRONMENT = { + "CUBLAS_WORKSPACE_CONFIG": ":16:8", + "FLASH_ATTENTION_DETERMINISTIC": "1", + "HF_HUB_OFFLINE": "1", + "NCCL_ALGO": "Ring", + "NCCL_PROTO": "Simple", + "PYTHONHASHSEED": "0", + "PYTHONNOUSERSITE": "1", + "TOKENIZERS_PARALLELISM": "false", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + "TRANSFORMERS_OFFLINE": "1", + "UNIRL_PARITY_ENABLE": "1", + "UNIRL_PARITY_MODEL": "qwen3_moe_30b_a3b", + "UNIRL_PARITY_PATCHES": "common,qwen3_moe_30b_a3b", + "UNIRL_PARITY_PROFILE": PUBLIC_REFERENCE, + "UNIRL_PARITY_STRICT": "1", + "VLLM_BATCH_INVARIANT": "0", + "VLLM_PLUGINS": PLUGIN_ENTRYPOINT, + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", +} + + +def _select(cfg: DictConfig, path: str, default: Any = None) -> Any: + return OmegaConf.select(cfg, path, default=default) + + +def _int(value: Any) -> int: + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _same(actual: Any, expected: Any) -> bool: + if isinstance(expected, float): + try: + return float(actual) == expected + except (TypeError, ValueError): + return False + return actual == expected + + +def _validate_gsm8k_jsonl(path: Path, *, min_rows: int) -> list[str]: + """Return GSM8K schema/uniqueness errors for one JSONL file.""" + errors: list[str] = [] + prompt_ids: set[str] = set() + rows = 0 + try: + with path.open("r", encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + continue + row = json.loads(line) + rows += 1 + if not isinstance(row, dict): + errors.append(f"{path}:{line_number} must contain a JSON object") + break + prompt_id = row.get("prompt_id") + prompt = row.get("prompt") + metadata = row.get("metadata") + if prompt_id is None: + prompt_id = f"{path}:{line_number}" + elif not isinstance(prompt_id, str) or not prompt_id.strip(): + errors.append(f"{path}:{line_number} has an empty prompt_id") + break + if prompt_id in prompt_ids: + errors.append(f"{path}:{line_number} duplicates prompt_id={prompt_id!r}") + break + prompt_ids.add(prompt_id) + if not isinstance(prompt, str) or not prompt.strip(): + errors.append(f"{path}:{line_number} has an empty prompt") + break + if not isinstance(metadata, dict) or metadata.get("answer") is None: + errors.append(f"{path}:{line_number} has no metadata.answer") + break + except (OSError, json.JSONDecodeError, TypeError) as error: + errors.append(f"could not validate GSM8K JSONL {path}: {type(error).__name__}: {error}") + return errors + if rows < min_rows: + errors.append(f"GSM8K JSONL {path} has {rows} rows; requires at least {min_rows}") + return errors + + +@dataclass(frozen=True) +class ParityRuntimeContract: + """The one pre-launch source of truth for this experiment.""" + + cfg: DictConfig + train_tp: int + num_devices: int + rollout_tp: int + max_prompt_length: int + max_new_tokens: int + model_path: str + model_revision: str + data_path: str + artifact_path: str + + @classmethod + def from_config(cls, cfg: DictConfig) -> "ParityRuntimeContract": + return cls( + cfg=cfg, + train_tp=_int(_select(cfg, "parity.train_tp", 0)), + num_devices=_int(_select(cfg, "num_devices", 0)), + rollout_tp=_int(_select(cfg, "rollout.config.tp_size", 0)), + max_prompt_length=_int(_select(cfg, "parity.max_prompt_length", 0)), + max_new_tokens=_int(_select(cfg, "parity.max_new_tokens", 0)), + model_path=str(_select(cfg, "bundle.config.pretrained_model_ckpt_path", "")), + model_revision=str(_select(cfg, "bundle.config.model_revision", "")), + data_path=str(_select(cfg, "data_source.args.run.data_path", "")), + artifact_path=str(_select(cfg, "parity.artifact_path", "")), + ) + + @property + def environment(self) -> Dict[str, str]: + values = dict(_REQUIRED_ENVIRONMENT) + values["UNIRL_PARITY_TRAIN_TP"] = str(self.train_tp) + return values + + def flags(self) -> Dict[str, Any]: + """Return the resolved numerical switches written into verification JSON.""" + return { + "profile": _select(self.cfg, "parity.profile"), + "train_tp": self.train_tp, + "rollout_tp": self.rollout_tp, + "num_devices": self.num_devices, + "model_precision": _select(self.cfg, "bundle.config.model_precision"), + "model_revision": self.model_revision, + "autocast_precision": _select(self.cfg, "pipeline.autocast_precision"), + "logprob_precision": _select(self.cfg, "pipeline.logprob_precision"), + "exact_actor_logprobs": _select(self.cfg, "pipeline.exact_actor_logprobs"), + "temperature": _select(self.cfg, "sampling.temperature"), + "top_p": _select(self.cfg, "sampling.top_p"), + "top_k": _select(self.cfg, "sampling.top_k"), + "logprobs_mode": _select(self.cfg, "rollout.config.logprobs_mode"), + "max_prompt_length": self.max_prompt_length, + "max_new_tokens": self.max_new_tokens, + "max_model_len": _select(self.cfg, "rollout.config.engine_kwargs.max_model_len"), + "num_updates_per_batch": _select(self.cfg, "stack.num_updates_per_batch"), + "data_shuffle": _select(self.cfg, "data_source.args.run.shuffle"), + "data_seed": _select(self.cfg, "data_source.args.run.seed"), + "batch_size": _select(self.cfg, "batch_size"), + "samples_per_prompt": _select(self.cfg, "sampling.samples_per_prompt"), + "expected_token_count": ( + _int(_select(self.cfg, "batch_size", 0)) + * _int(_select(self.cfg, "sampling.samples_per_prompt", 0)) + * self.max_new_tokens + ), + "ignore_eos": _select(self.cfg, "rollout.config.ignore_eos"), + "enable_prefix_caching": _select(self.cfg, "rollout.config.enable_prefix_caching"), + "enable_chunked_prefill": _select(self.cfg, "rollout.config.enable_chunked_prefill"), + "trust_remote_code": { + "actor": _select(self.cfg, "bundle.config.trust_remote_code"), + "rollout": _select(self.cfg, "rollout.config.trust_remote_code"), + }, + } + + def validate(self) -> None: + """Aggregate every pre-launch mismatch and raise once.""" + errors: list[str] = [] + + def expect(path: str, expected: Any) -> None: + actual = _select(self.cfg, path) + if not _same(actual, expected): + errors.append(f"{path}={actual!r}; expected {expected!r}") + + expect("parity.profile", PUBLIC_REFERENCE) + expect("parity.train_tp", EXPECTED_TP) + expect("parity.max_prompt_length", 256) + expect("parity.max_new_tokens", 1024) + expect("num_devices", EXPECTED_TP) + expect("devices_per_node", EXPECTED_TP) + expect("rollout.config.tp_size", EXPECTED_TP) + expect("bundle.config.model_revision", EXPECTED_MODEL_REVISION) + expect("rollout.config.model_revision", EXPECTED_MODEL_REVISION) + expect( + "rollout._target_", + "unirl.rollout.engine.vllm.engine.VLLMRolloutEngine", + ) + + expect("bundle.config.model_precision", "bf16") + expect("bundle.config.autocast_precision", "bf16") + expect("pipeline.autocast_precision", "bf16") + expect("bundle.config.logprob_precision", "fp32") + expect("pipeline.logprob_precision", "fp32") + expect("bundle.config.exact_actor_logprobs", True) + expect("pipeline.exact_actor_logprobs", True) + expect( + "bundle.config.external_libs", + ["experimental.train_inference_parity.models.qwen3_moe_30b_a3b.fsdp"], + ) + expect( + "algorithm._target_", + "experimental.train_inference_parity.algorithm.TrainInferenceParityGRPO", + ) + + for path in ( + "rollout.config.temperature", + "sampling.temperature", + "algorithm.sampling_temperature", + ): + expect(path, 1.0) + for path in ("rollout.config.top_p", "sampling.top_p"): + expect(path, 1.0) + for path in ("rollout.config.top_k", "sampling.top_k"): + expect(path, 0) + expect( + "rollout.config.logprobs_mode", + "processed_logprobs", + ) + expect("rollout.config.dtype", "bfloat16") + expect("rollout.config.enforce_eager", True) + expect("rollout.config.moe_backend", "triton") + expect("rollout.config.enable_prefix_caching", False) + expect("rollout.config.enable_chunked_prefill", False) + expect("rollout.config.ignore_eos", True) + expect("rollout.config.startup_timeout_s", 300) + expect("rollout.config.generate_timeout_s", 7200) + expect("rollout.config.weight_update_timeout_s", 1800) + expect("rollout.config.health_timeout_s", 60) + expect("rollout.config.sleep_timeout_s", 180) + expect("rollout.config.wake_up_timeout_s", 180) + expect("rollout.config.shutdown_timeout_s", 120) + + for path in ( + "bundle.config.max_prompt_length", + "pipeline.max_prompt_length", + "rollout.config.max_prompt_length", + ): + expect(path, self.max_prompt_length) + chat_truncation = _select(self.cfg, "rollout.config.chat_template_kwargs.truncation") + if chat_truncation not in (None, False): + errors.append( + "rollout chat-template truncation must be unset; " + "VLLMRolloutEngine applies the shared suffix limit after tokenization" + ) + for path in ( + "rollout.config.max_new_tokens", + "sampling.max_new_tokens", + "algorithm.horizon", + ): + expect(path, self.max_new_tokens) + expect( + "rollout.config.engine_kwargs.max_model_len", + self.max_prompt_length + self.max_new_tokens, + ) + expect( + "rollout.config.engine_kwargs.max_num_batched_tokens", + self.max_prompt_length + self.max_new_tokens, + ) + + expect("stack.num_updates_per_batch", 1) + expect("stack.micro_batch_size", 1) + expect("num_rollouts", 2) + expect("bundle.config.trust_remote_code", True) + expect("rollout.config.trust_remote_code", True) + expect("sync._target_", EXPECTED_SYNC_TARGET) + expect("sync.flush_cache", True) + expect("sync.bucket_size_mb", 640) + expect("sync.gpu_memory_headroom_mb", 2048) + expect("sync.control_timeout_s", 2400) + weight_timeout = float(_select(self.cfg, "rollout.config.weight_update_timeout_s") or 0) + control_timeout = float(_select(self.cfg, "sync.control_timeout_s") or 0) + if control_timeout < weight_timeout + 300.0: + errors.append( + "sync.control_timeout_s must exceed rollout.config.weight_update_timeout_s " + f"by at least 300s; got control={control_timeout}, weight_update={weight_timeout}" + ) + expect("backend.fsdp_cfg.activation_checkpointing", True) + expect("bundle.config.use_gradient_checkpointing", True) + expect("bundle.config.attn_implementation", "flash_attention_3") + expect("sampling.samples_per_prompt", 8) + expect("batch_size", 4) + expect("data_source.args.algorithm.prompts_per_rollout", 4) + expect("data_source.args.run.shuffle", True) + expect("data_source.args.run.seed", None) + + actor_path = str(_select(self.cfg, "bundle.config.pretrained_model_ckpt_path", "")) + rollout_path = str(_select(self.cfg, "rollout.config.pretrained_model_ckpt_path", "")) + if not actor_path or actor_path != rollout_path: + errors.append( + "actor and rollout pretrained_model_ckpt_path must resolve to the same non-empty value; " + f"got actor={actor_path!r}, rollout={rollout_path!r}" + ) + actor_revision = str(_select(self.cfg, "bundle.config.model_revision", "")) + rollout_revision = str(_select(self.cfg, "rollout.config.model_revision", "")) + if actor_revision != rollout_revision: + errors.append(f"actor and rollout model revisions differ: {actor_revision!r} vs {rollout_revision!r}") + + batch_size = _int(_select(self.cfg, "batch_size", 0)) + samples_per_prompt = _int(_select(self.cfg, "sampling.samples_per_prompt", 0)) + micro_batch_size = _int(_select(self.cfg, "stack.micro_batch_size", 0)) + total_samples = batch_size * samples_per_prompt + if batch_size <= 0 or samples_per_prompt <= 0: + errors.append("batch_size and sampling.samples_per_prompt must both be positive") + elif self.num_devices <= 0: + errors.append("num_devices must be positive") + elif total_samples % self.num_devices: + errors.append( + f"batch_size*samples_per_prompt={total_samples} is not divisible by num_devices={self.num_devices}" + ) + elif micro_batch_size <= 0 or (total_samples // self.num_devices) % micro_batch_size: + errors.append( + "per-rank generated samples must be divisible by stack.micro_batch_size; " + f"got total={total_samples}, devices={self.num_devices}, micro={micro_batch_size}" + ) + + eval_data_path = str(_select(self.cfg, "data_source.args.run.eval_data_path", "")) + for label, raw_path in ( + ("model checkpoint", self.model_path), + ("GSM8K training data", self.data_path), + ("GSM8K evaluation data", eval_data_path), + ): + if not raw_path or not Path(raw_path).expanduser().exists(): + errors.append(f"{label} path does not exist: {raw_path!r}") + required_unique_prompts = batch_size * _int(_select(self.cfg, "num_rollouts", 0)) + for raw_path in (self.data_path, eval_data_path): + path = Path(raw_path).expanduser() + if raw_path and path.is_file(): + errors.extend(_validate_gsm8k_jsonl(path, min_rows=required_unique_prompts)) + if not self.artifact_path: + errors.append("parity.artifact_path must be explicit") + + if self.model_path and Path(self.model_path).expanduser().exists(): + try: + from transformers import AutoConfig + + from experimental.train_inference_parity.models.qwen3_moe_30b_a3b.contract import ( + validate_model_config, + ) + + model_config = AutoConfig.from_pretrained( + self.model_path, + revision=self.model_revision, + local_files_only=True, + trust_remote_code=True, + ) + validate_model_config(model_config, dtype="bfloat16") + if getattr(model_config, "model_type", None) != "qwen3_moe": + errors.append( + f"model config type={getattr(model_config, 'model_type', None)!r}; expected 'qwen3_moe'" + ) + architectures = tuple(getattr(model_config, "architectures", ()) or ()) + if "Qwen3MoeForCausalLM" not in architectures: + errors.append(f"model config architectures={architectures!r}; expected Qwen3MoeForCausalLM") + for field in ("attention_dropout", "hidden_dropout", "hidden_dropout_prob"): + value = getattr(model_config, field, None) + if value not in (None, 0, 0.0): + errors.append(f"model config {field}={value!r}; exact parity requires dropout=0") + except Exception as error: + errors.append(f"could not validate frozen model structure: {type(error).__name__}: {error}") + + for name, expected in self.environment.items(): + actual = os.environ.get(name) + if actual != expected: + errors.append(f"environment {name}={actual!r}; expected {expected!r}") + + for distribution, expected in EXPECTED_VERSIONS.items(): + try: + actual = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + errors.append(f"{distribution} is not installed; expected {expected}") + continue + if actual.split("+", 1)[0] != expected: + errors.append(f"{distribution}=={actual}; expected public version {expected}") + try: + importlib.metadata.version(PLUGIN_DISTRIBUTION) + except importlib.metadata.PackageNotFoundError: + errors.append(f"{PLUGIN_DISTRIBUTION} is not installed from vllm_plugin/") + + try: + import ray + except ImportError: + errors.append("ray is not installed") + else: + if ray.is_initialized(): + errors.append("Ray is already initialized; parity runtime_env propagation cannot be proven") + + try: + import torch + except ImportError: + errors.append("torch is not importable") + else: + visible_devices = torch.cuda.device_count() if torch.cuda.is_available() else 0 + if visible_devices != EXPECTED_TP: + errors.append(f"visible CUDA devices={visible_devices}; expected exactly {EXPECTED_TP}") + + if errors: + detail = "\n".join(f" - {error}" for error in errors) + raise RuntimeError(f"Parity runtime contract failed:\n{detail}") + + +def required_environment(train_tp: int = EXPECTED_TP) -> Mapping[str, str]: + """Expose the explicit environment for launchers and documentation tooling.""" + return { + **_REQUIRED_ENVIRONMENT, + "UNIRL_PARITY_TRAIN_TP": str(int(train_tp)), + } + + +__all__ = [ + "EXPECTED_SYNC_TARGET", + "EXPECTED_TP", + "EXPECTED_VERSIONS", + "PUBLIC_REFERENCE", + "ParityRuntimeContract", + "required_environment", +] diff --git a/experimental/train_inference_parity/sync.py b/experimental/train_inference_parity/sync.py new file mode 100644 index 000000000..754017b3f --- /dev/null +++ b/experimental/train_inference_parity/sync.py @@ -0,0 +1,196 @@ +"""Parity evidence adapter over UniRL's production vLLM native IPC sync.""" + +from __future__ import annotations + +import hashlib +import json +from typing import Any, Optional + +from unirl.distributed.group.dispatch import Dispatch, distributed +from unirl.distributed.weight_sync.full.ipc import IPCWeightSync + +# Prefer shards that a 1e-6 BF16 AdamW step is likely to change. Sampling the +# smallest RMSNorms first produced false "no actor shard changed" failures. +_FINGERPRINT_NAME_SUFFIXES = ( + "embed_tokens.weight", + "lm_head.weight", + "self_attn.q_proj.weight", + "self_attn.o_proj.weight", + "mlp.gate.weight", + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", +) + + +class ParityIPCWeightSync(IPCWeightSync): + """Record actor-change and native publication evidence without owning transport.""" + + def __init__( + self, + *args, + fingerprint_param_samples: int = 8, + fingerprint_element_samples: int = 16, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self._fingerprint_param_samples = int(fingerprint_param_samples) + self._fingerprint_element_samples = int(fingerprint_element_samples) + if self._fingerprint_param_samples <= 0 or self._fingerprint_element_samples <= 0: + raise ValueError("parity fingerprint sample counts must be positive") + self._actor_baseline: Optional[dict[str, str]] = None + self._actor_baseline_source: Optional[dict[str, Any]] = None + self._actor_candidate: Optional[dict[str, str]] = None + self._actor_candidate_source: Optional[dict[str, Any]] = None + self._actor_update_pending = False + self._publication_history: list[dict[str, Any]] = [] + + def _select_fingerprint_parameters(self, candidates: list[tuple[str, Any]]) -> list[tuple[str, Any]]: + """Pick large structural shards instead of the smallest RMSNorms.""" + selected: list[tuple[str, Any]] = [] + seen: set[str] = set() + for suffix in _FINGERPRINT_NAME_SUFFIXES: + for name, parameter in candidates: + if name.endswith(suffix) and name not in seen: + selected.append((name, parameter)) + seen.add(name) + break + remaining = sorted( + ((name, parameter) for name, parameter in candidates if name not in seen), + key=lambda item: (-int(item[1].numel()), item[0]), + ) + for name, parameter in remaining: + if len(selected) >= self._fingerprint_param_samples: + break + selected.append((name, parameter)) + seen.add(name) + selected.sort(key=lambda item: item[0]) + return selected[: self._fingerprint_param_samples] + + def _capture_local_shard_fingerprints(self, phase: str) -> dict[str, str]: + fingerprints: dict[str, str] = {} + error: Optional[BaseException] = None + try: + import torch + from torch.distributed.tensor import DTensor + + candidates = [ + (name, parameter) + for name, parameter in self._backend.model.named_parameters() + if parameter.requires_grad and not parameter.is_meta + ] + for name, parameter in self._select_fingerprint_parameters(candidates): + local = parameter.detach() + if isinstance(local, DTensor): + local = local.to_local() + flat = local.reshape(-1) + if flat.numel() == 0: + continue + count = min(int(flat.numel()), self._fingerprint_element_samples) + indices = ( + [0] if count == 1 else [index * (int(flat.numel()) - 1) // (count - 1) for index in range(count)] + ) + sample = flat.index_select( + 0, + torch.tensor(indices, dtype=torch.long, device=flat.device), + ).contiguous() + digest = hashlib.sha256() + digest.update(name.encode()) + digest.update(str(tuple(parameter.shape)).encode()) + digest.update(str(parameter.dtype).encode()) + digest.update(sample.cpu().view(torch.uint8).numpy().tobytes()) + fingerprints[name] = digest.hexdigest() + if not fingerprints: + raise RuntimeError("no trainable local actor shards were available for fingerprinting") + except BaseException as exc: + error = exc + self._consensus(error, phase) + return fingerprints + + def _aggregate_source_fingerprint(self, fingerprints: dict[str, str]) -> dict[str, Any]: + reports = self._all_gather({"rank": self._global_rank, "parameters": dict(fingerprints)}) + reports.sort(key=lambda report: int(report["rank"])) + canonical = json.dumps(reports, sort_keys=True, separators=(",", ":")).encode() + return { + "algorithm": "sha256(sampled-local-fsdp-shards-v1)", + "sha256": hashlib.sha256(canonical).hexdigest(), + "world_size": len(reports), + "parameter_samples_per_rank": self._fingerprint_param_samples, + "element_samples_per_parameter": self._fingerprint_element_samples, + "rank_reports": reports, + } + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def record_initial_actor_fingerprint(self) -> dict[str, Any]: + fingerprints = self._capture_local_shard_fingerprints("initial-actor-fingerprint") + self._actor_baseline = fingerprints + self._actor_baseline_source = self._aggregate_source_fingerprint(fingerprints) + self._actor_candidate = None + self._actor_candidate_source = None + self._actor_update_pending = False + return { + "baseline_recorded": True, + "rank": self._global_rank, + "sampled_parameters": sorted(fingerprints), + "source_model_fingerprint": dict(self._actor_baseline_source), + } + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def mark_actor_updated(self) -> None: + if self._actor_baseline is None: + raise RuntimeError("actor fingerprint baseline was not recorded before training") + self._actor_update_pending = True + + def _capture_pending_actor_update(self) -> None: + if not self._actor_update_pending: + return + if self._actor_baseline is None: + raise RuntimeError("actor fingerprint baseline is missing") + candidate = self._capture_local_shard_fingerprints("post-update-actor-fingerprint") + local_changed = sorted(name for name, digest in candidate.items() if self._actor_baseline.get(name) != digest) + changed_reports = self._all_gather({"rank": self._global_rank, "changed_parameters": local_changed}) + if not any(report["changed_parameters"] for report in changed_reports): + raise RuntimeError("no sampled actor local shard changed after the optimizer update") + self._actor_candidate = candidate + self._actor_candidate_source = self._aggregate_source_fingerprint(candidate) + + def _sync_with_vllm_native_engine(self) -> None: + self._capture_pending_actor_update() + super()._sync_with_vllm_native_engine() + + def _on_native_publication_committed( + self, + *, + header: dict[str, Any], + result: dict[str, Any], + ) -> None: + source = self._actor_candidate_source or self._actor_baseline_source + if source is None: + raise RuntimeError("parity publication has no actor source fingerprint") + changed = bool( + self._actor_baseline_source and source.get("sha256") != self._actor_baseline_source.get("sha256") + ) + publication = { + "publication_id": str(header["publication_id"]), + "model_version": int(header["model_version"]), + "tp_world_size": int(header["tp_world_size"]), + "payload": dict(header["expected_manifest"]), + "workers": [dict(item) for item in result["worker_receipts"]], + "prefix_cache_reset": bool(result["prefix_cache_reset"]), + "source_model_fingerprint": dict(source), + "parameter_changed": changed, + "reload_applied": True, + } + self._publication_history.append(publication) + if self._actor_update_pending: + self._actor_baseline = self._actor_candidate + self._actor_baseline_source = self._actor_candidate_source + self._actor_candidate = None + self._actor_candidate_source = None + self._actor_update_pending = False + + @distributed(dispatch_mode=Dispatch.BROADCAST) + def verification_receipts(self) -> list[dict[str, Any]]: + return [dict(publication) for publication in self._publication_history] + + +__all__ = ["ParityIPCWeightSync"] diff --git a/experimental/train_inference_parity/verification.py b/experimental/train_inference_parity/verification.py new file mode 100644 index 000000000..36df55df7 --- /dev/null +++ b/experimental/train_inference_parity/verification.py @@ -0,0 +1,475 @@ +"""Machine-readable evidence writer for the two-phase TP4 parity run.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Any, Dict, Mapping, Optional + +from .artifacts import ( + atomic_write_json, + data_fingerprint, + git_state, + hardware_state, + jsonable, + model_fingerprint, + software_state, +) +from .runtime_contract import ( + EXPECTED_TP, + EXPECTED_VERSIONS, + ParityRuntimeContract, +) + +SCHEMA_VERSION = "unirl.train_inference_parity.verification.v1" +PHASES = ("initial_checkpoint", "post_update_reload") +CANONICAL_PARITY_METRICS = ( + "token_count", + "torch_equal_fp32", + "mismatch_count", + "max_absdiff_fp32", + "k3_mean", + "k3_max", +) + + +def _fingerprint_digest(value: Mapping[str, Any]) -> Optional[str]: + raw = value.get("sha256", value.get("digest")) + return str(raw) if raw else None + + +def _normalize_source_fingerprint(value: Any) -> Dict[str, Any]: + value = jsonable(value) + if isinstance(value, Mapping): + normalized = dict(value) + if not _fingerprint_digest(normalized): + raise ValueError("source fingerprint mapping must contain sha256 or digest") + return normalized + if not isinstance(value, list) or not value: + raise TypeError("source_model_fingerprint() must return a mapping or rank-report list") + if not all(isinstance(item, Mapping) for item in value): + raise TypeError("source fingerprint rank reports must all be mappings") + reports = [dict(item) for item in value] + digests = {_fingerprint_digest(report) for report in reports} + if None in digests or len(digests) != 1: + raise ValueError(f"source fingerprint ranks disagree: {digests!r}") + normalized = {key: item for key, item in reports[0].items() if key != "rank"} + normalized["rank_reports"] = reports + return normalized + + +def _baseline_from_initial_state(value: Any) -> Dict[str, Any]: + """Extract one consensus source fingerprint from the initial sync-state RPC.""" + value = jsonable(value) + states = value if isinstance(value, list) else [value] + if not states or not all(isinstance(state, Mapping) for state in states): + raise TypeError("record_initial_actor_fingerprint() must return mapping states") + fingerprints = [state.get("source_model_fingerprint") for state in states] + normalized = [_normalize_source_fingerprint(fingerprint) for fingerprint in fingerprints] + digests = {_fingerprint_digest(fingerprint) for fingerprint in normalized} + if None in digests or len(digests) != 1: + raise ValueError(f"initial actor fingerprint states disagree: {digests!r}") + return normalized[0] + + +def _call_required(target: Any, method_name: str) -> Any: + method = getattr(target, method_name, None) + if not callable(method): + raise RuntimeError(f"parity integration is incomplete: missing callable {method_name}()") + return method() + + +def _normalize_receipts(value: Any) -> list[Dict[str, Any]]: + value = jsonable(value) + if isinstance(value, list) and value and all(isinstance(item, list) for item in value): + canonical = {json.dumps(item, sort_keys=True, separators=(",", ":")) for item in value} + if len(canonical) != 1: + raise ValueError("weight-sync ranks disagree on verification receipts") + value = value[0] + if isinstance(value, Mapping) and "receipts" in value: + value = value["receipts"] + if isinstance(value, Mapping): + value = [value] + if not isinstance(value, list): + raise TypeError(f"verification_receipts() returned {type(value).__name__}, expected list") + receipts = [] + for item in value: + if not isinstance(item, Mapping): + raise TypeError("every weight receipt must be a mapping") + receipts.append(dict(item)) + return receipts + + +def _validate_plugin_manifest(value: Any) -> Dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"vLLM plugin manifest must be a mapping, got {type(value).__name__}") + manifest = dict(value) + if ( + manifest.get("profile") != "public_reference" + or manifest.get("model") != "qwen3_moe_30b_a3b" + or manifest.get("strict") is not True + ): + raise ValueError(f"vLLM plugin manifest contract mismatch: {manifest!r}") + patches = manifest.get("patches") + if not isinstance(patches, list) or [patch.get("name") for patch in patches] != [ + "common", + "qwen3_moe_30b_a3b", + ]: + raise ValueError(f"vLLM plugin manifest has the wrong patch set: {patches!r}") + for patch in patches: + symbols = patch.get("symbols") + if not isinstance(symbols, list) or not symbols: + raise ValueError(f"vLLM plugin patch has no symbol evidence: {patch!r}") + if not all(isinstance(symbol, Mapping) and symbol.get("verified") is True for symbol in symbols): + raise ValueError(f"vLLM plugin patch contains unverified symbols: {patch!r}") + versions = manifest.get("runtime") + if not isinstance(versions, list): + raise ValueError("vLLM plugin manifest omitted runtime versions") + installed_versions = { + str(item.get("package")): str(item.get("installed")).split("+", 1)[0] + for item in versions + if isinstance(item, Mapping) + } + if installed_versions != EXPECTED_VERSIONS: + raise ValueError(f"vLLM worker versions={installed_versions!r}; expected {EXPECTED_VERSIONS!r}") + return manifest + + +def _normalize_plugin_manifests(value: Any) -> Dict[str, Any]: + value = jsonable(value) + if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list): + value = value[0] + if not isinstance(value, list): + raise TypeError(f"vLLM runtime capabilities must be a list, got {type(value).__name__}") + manifests = [] + for tp_rank, capability in enumerate(value): + if not isinstance(capability, Mapping): + raise TypeError(f"vLLM TP capability {tp_rank} must be a mapping") + manifest = capability.get("train_inference_parity") + if manifest is None: + raise ValueError(f"vLLM TP capability {tp_rank} omitted the parity plugin manifest") + manifests.append( + { + "tp_rank": int(capability.get("tp_rank", -1)), + "manifest": _validate_plugin_manifest(manifest), + } + ) + ranks = sorted(item["tp_rank"] for item in manifests) + if ranks != list(range(EXPECTED_TP)): + raise ValueError(f"vLLM plugin manifest TP ranks={ranks}; expected {list(range(EXPECTED_TP))}") + return {"tp_world_size": len(manifests), "workers": manifests} + + +def _per_rank_token_counts(result: Any) -> list[int]: + totals: list[int] = [] + for micro in getattr(result, "micros", ()) or (): + counts = (getattr(micro, "metrics", {}) or {}).get("per_rank_token_count") + if not isinstance(counts, (list, tuple)): + continue + if not totals: + totals = [0] * len(counts) + if len(counts) != len(totals): + raise ValueError("parity micro-batches reported inconsistent rank counts") + totals = [total + int(count) for total, count in zip(totals, counts, strict=True)] + return totals + + +def _aggregate_parity_metrics(result: Any) -> tuple[Dict[str, Any], list[int]]: + micros = [dict(getattr(micro, "metrics", {}) or {}) for micro in getattr(result, "micros", ()) or ()] + for index, metrics in enumerate(micros): + missing = [name for name in CANONICAL_PARITY_METRICS if name not in metrics] + if missing: + raise RuntimeError(f"micro-batch {index} omitted canonical parity metrics: {missing}") + if not micros: + metrics = dict(getattr(result, "metrics", {}) or {}) + missing = [name for name in CANONICAL_PARITY_METRICS if name not in metrics] + if missing: + raise RuntimeError(f"missing canonical parity metrics: {missing}") + return {name: jsonable(metrics[name]) for name in CANONICAL_PARITY_METRICS}, [] + + token_count = sum(int(metrics["token_count"]) for metrics in micros) + k3_sum = sum(float(metrics["k3_mean"]) * int(metrics["token_count"]) for metrics in micros) + canonical = { + "token_count": token_count, + "torch_equal_fp32": all(bool(metrics["torch_equal_fp32"]) for metrics in micros), + "mismatch_count": sum(int(metrics["mismatch_count"]) for metrics in micros), + "max_absdiff_fp32": max(float(metrics["max_absdiff_fp32"]) for metrics in micros), + "k3_mean": k3_sum / token_count if token_count else 0.0, + "k3_max": max(float(metrics["k3_max"]) for metrics in micros), + } + return canonical, _per_rank_token_counts(result) + + +class ParityVerification: + """Collect evidence and emit PASS only after both required phases.""" + + def __init__(self, contract: ParityRuntimeContract) -> None: + self.contract = contract + self.artifact_path = Path(contract.artifact_path).expanduser().resolve() + self._weight_sync: Any = None + self._baseline_source_fingerprint: Optional[Dict[str, Any]] = None + self._document: Dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "status": "INCOMPLETE", + "git": git_state(), + "software": software_state(), + "hardware": hardware_state(), + "model_fingerprint": model_fingerprint(contract.model_path), + "data_fingerprint": data_fingerprint(contract.data_path), + "flags": contract.flags(), + "environment": dict(contract.environment), + "phases": [], + "failure_reasons": ["both required phases have not completed"], + } + + def initialize(self) -> None: + """Persist launch metadata before Ray or model construction begins.""" + self._write() + + def bind(self, trainer: Any) -> None: + """Capture the initial actor fingerprint before the first optimizer update.""" + self._weight_sync = getattr(trainer, "weight_sync", None) + if self._weight_sync is None: + self._fail("trainer has no dedicated weight-sync object") + try: + rollout = getattr(trainer, "rollout", None) + capability_method = getattr(rollout, "runtime_capabilities", None) + if not callable(capability_method): + raise RuntimeError("direct vLLM rollout does not expose runtime_capabilities()") + self._document["vllm_plugin_manifest"] = _normalize_plugin_manifests(capability_method()) + initial_state = _call_required(self._weight_sync, "record_initial_actor_fingerprint") + self._document["initial_actor_fingerprint_state"] = jsonable(initial_state) + baseline = _baseline_from_initial_state(initial_state) + self._baseline_source_fingerprint = baseline + self._document["initial_actor_fingerprint"] = baseline + self._write() + except Exception as error: + self._fail(f"could not capture initial actor fingerprint: {error}", cause=error) + + def on_rollout_complete( + self, + *, + rollout_id: int, + result: Any, + mean_reward: float, + sync_weights: bool, + ) -> None: + """ARTrainer callback: record one canonical parity phase.""" + try: + if rollout_id < 0 or rollout_id >= len(PHASES): + raise ValueError(f"unexpected rollout_id={rollout_id}; the contract requires exactly two") + if rollout_id != len(self._document["phases"]): + raise ValueError(f"out-of-order parity phase: rollout_id={rollout_id}") + phase_name = PHASES[rollout_id] + expected_reload = phase_name == "post_update_reload" + if bool(sync_weights) != expected_reload: + raise RuntimeError(f"{phase_name} sync_weights={sync_weights}; expected {expected_reload}") + + canonical, per_rank_token_count = _aggregate_parity_metrics(result) + reported_token_count = int(canonical["token_count"]) + token_count = sum(per_rank_token_count) if per_rank_token_count else reported_token_count + if per_rank_token_count and reported_token_count != token_count: + raise RuntimeError( + f"aggregated token_count={reported_token_count} does not match per-rank total={token_count}" + ) + canonical["token_count"] = token_count + expected_tokens = int(self.contract.flags().get("expected_token_count") or 0) + if token_count <= 0: + raise RuntimeError(f"{phase_name} recorded no replay tokens") + if expected_tokens > 0 and token_count != expected_tokens: + raise RuntimeError( + f"{phase_name} token_count={token_count}; expected {expected_tokens} " + "(batch_size × samples_per_prompt × max_new_tokens with ignore_eos)" + ) + + receipts: list[Dict[str, Any]] = [] + parameter_changed: Optional[bool] = None + receipt_errors: list[str] = [] + if expected_reload: + receipts = _normalize_receipts(_call_required(self._weight_sync, "verification_receipts")) + parameter_changed, receipt_errors = self._validate_receipts(receipts) + + exact_metrics = ( + canonical["torch_equal_fp32"] in (True, 1, 1.0) + and canonical["mismatch_count"] == 0 + and canonical["max_absdiff_fp32"] == 0.0 + and canonical["k3_mean"] == 0.0 + and canonical["k3_max"] == 0.0 + ) + grad_norm = float(getattr(result, "grad_norm")) + optimizer_updates = int(getattr(result, "optimizer_updates", 0)) + if phase_name == "initial_checkpoint" and optimizer_updates == 1: + _call_required(self._weight_sync, "mark_actor_updated") + phase_errors = list(receipt_errors) + if not exact_metrics: + phase_errors.append("canonical parity metrics are not exact") + if not math.isfinite(grad_norm): + phase_errors.append(f"grad_norm is not finite: {grad_norm!r}") + if optimizer_updates != 1: + phase_errors.append(f"optimizer_updates={optimizer_updates}; expected exactly 1") + if expected_reload and parameter_changed is not True: + phase_errors.append("source parameters did not change from the initial actor") + + phase = { + "phase": phase_name, + "rollout_id": int(rollout_id), + "status": "PASS" if not phase_errors else "FAIL", + "token_count": token_count, + "per_rank_token_count": per_rank_token_count, + "canonical_parity_metrics": canonical, + "grad_norm": grad_norm, + "optimizer_updates": optimizer_updates, + "mean_reward": float(mean_reward), + "parameter_changed": parameter_changed, + "weight_receipts": receipts, + "failure_reasons": phase_errors, + } + self._document["phases"].append(phase) + self._refresh_status() + self._write() + if phase_errors: + raise RuntimeError(f"{phase_name} verification failed: {'; '.join(phase_errors)}") + except Exception as error: + if not self._document["phases"] or self._document["phases"][-1].get("rollout_id") != rollout_id: + self._document["status"] = "FAIL" + self._document["failure_reasons"] = [ + f"{PHASES[rollout_id] if 0 <= rollout_id < len(PHASES) else rollout_id}: {error}" + ] + self._write() + raise + + def _validate_receipts(self, receipts: list[Dict[str, Any]]) -> tuple[bool, list[str]]: + errors: list[str] = [] + if not receipts: + return False, ["post-update reload produced no weight receipt"] + baseline = _fingerprint_digest(self._baseline_source_fingerprint or {}) + source_digests = set() + for receipt_index, receipt in enumerate(receipts): + payload = receipt.get("payload") + workers = receipt.get("workers") + source = receipt.get("source_model_fingerprint") + if not isinstance(payload, Mapping): + errors.append(f"receipt[{receipt_index}].payload is missing") + continue + if int(payload.get("tensor_count", 0)) <= 0: + errors.append(f"receipt[{receipt_index}] has an empty tensor payload") + if int(payload.get("byte_count", 0)) <= 0: + errors.append(f"receipt[{receipt_index}] has zero payload bytes") + expected_fingerprint = payload.get("fingerprint") + if not isinstance(expected_fingerprint, str) or len(expected_fingerprint) != 64: + errors.append(f"receipt[{receipt_index}] has an invalid payload fingerprint") + if int(receipt.get("tp_world_size", 0)) != EXPECTED_TP: + errors.append( + f"receipt[{receipt_index}] tp_world_size={receipt.get('tp_world_size')!r}; expected {EXPECTED_TP}" + ) + if not isinstance(source, Mapping) or not _fingerprint_digest(source): + errors.append(f"receipt[{receipt_index}].source_model_fingerprint is missing") + else: + source_digests.add(_fingerprint_digest(source)) + if not isinstance(workers, list): + errors.append(f"receipt[{receipt_index}].workers is missing") + continue + ranks = sorted(int(worker.get("tp_rank", -1)) for worker in workers if isinstance(worker, Mapping)) + if ranks != list(range(EXPECTED_TP)): + errors.append(f"receipt[{receipt_index}] TP ranks={ranks}; expected {list(range(EXPECTED_TP))}") + versions = set() + device_uuids = set() + for worker in workers: + if not isinstance(worker, Mapping): + errors.append(f"receipt[{receipt_index}] has a non-mapping worker") + continue + if worker.get("publication_id") != receipt.get("publication_id"): + errors.append(f"receipt[{receipt_index}] tp_rank={worker.get('tp_rank')} publication_id mismatch") + if worker.get("expected_manifest_fingerprint") != expected_fingerprint: + errors.append( + f"receipt[{receipt_index}] tp_rank={worker.get('tp_rank')} manifest fingerprint mismatch" + ) + if worker.get("ready") is not True or int(worker.get("resident_parameter_count", 0)) <= 0: + errors.append(f"receipt[{receipt_index}] tp_rank={worker.get('tp_rank')} is not ready") + versions.add(worker.get("model_version")) + device_uuid = worker.get("device_uuid") + if not isinstance(device_uuid, str) or not device_uuid: + errors.append(f"receipt[{receipt_index}] tp_rank={worker.get('tp_rank')} has no CUDA UUID") + else: + device_uuids.add(device_uuid) + expected_model_version = receipt.get("model_version") + if versions != {expected_model_version}: + errors.append( + f"receipt[{receipt_index}] worker model versions={versions}; " + f"expected publication model_version={expected_model_version!r}" + ) + if len(device_uuids) != EXPECTED_TP: + errors.append( + f"receipt[{receipt_index}] unique worker CUDA UUIDs={len(device_uuids)}; expected {EXPECTED_TP}" + ) + if receipt.get("reload_applied") is not True: + errors.append(f"receipt[{receipt_index}] reload_applied is not true") + if receipt.get("prefix_cache_reset") is not True: + errors.append(f"receipt[{receipt_index}] prefix_cache_reset is not true") + if receipt.get("parameter_changed") is not True: + errors.append(f"receipt[{receipt_index}] parameter_changed is not true") + changed = bool( + baseline + and len(source_digests) == 1 + and next(iter(source_digests), None) != baseline + and all(receipt.get("parameter_changed") is True for receipt in receipts) + ) + if len(source_digests) > 1: + errors.append(f"receipts disagree on source fingerprint: {source_digests}") + return changed, errors + + def _refresh_status(self) -> None: + phases = self._document["phases"] + complete = len(phases) == len(PHASES) and [phase.get("phase") for phase in phases] == list(PHASES) + failures = list(self._document.get("runtime_failures", ())) + failures.extend(reason for phase in phases for reason in phase.get("failure_reasons", [])) + git_state = self._document["git"] + if git_state.get("dirty") is not False: + failures.append("git source tree is dirty or unknown") + if complete and not failures and all(phase.get("status") == "PASS" for phase in phases): + self._document["status"] = "PASS" + self._document["failure_reasons"] = [] + else: + self._document["status"] = "FAIL" if failures else "INCOMPLETE" + self._document["failure_reasons"] = failures or ["both required phases have not completed"] + + def record_failure(self, stage: str, error: BaseException) -> None: + """Persist an external launch/runtime failure without masking it.""" + reason = f"{stage}: {type(error).__name__}: {error}" + self._document.setdefault("runtime_failures", []).append(reason) + existing = list(self._document.get("failure_reasons", ())) + if reason not in existing: + existing.append(reason) + self._document["status"] = "FAIL" + self._document["failure_reasons"] = existing + self._write() + + def finalize(self) -> None: + """Require a complete, clean, two-phase PASS artifact.""" + self._refresh_status() + self._write() + if self._document["status"] != "PASS": + raise RuntimeError( + f"parity verification did not produce PASS: {self._document.get('failure_reasons', [])!r}" + ) + + def _fail(self, message: str, *, cause: Optional[Exception] = None) -> None: + self._document["status"] = "FAIL" + self._document["failure_reasons"] = [message] + self._write() + if cause is not None: + raise RuntimeError(message) from cause + raise RuntimeError(message) + + def _write(self) -> None: + atomic_write_json(self.artifact_path, self._document) + + +__all__ = [ + "CANONICAL_PARITY_METRICS", + "PHASES", + "SCHEMA_VERSION", + "ParityVerification", + "atomic_write_json", +] diff --git a/unirl/algorithms/grpo.py b/unirl/algorithms/grpo.py index 15d940d30..590e71c51 100644 --- a/unirl/algorithms/grpo.py +++ b/unirl/algorithms/grpo.py @@ -117,10 +117,40 @@ def compute_loss_and_backward( new_logp = self.stage.replay( typed_conds, segment=segment, temperature=self.sampling_temperature ) # [total_tokens] - # old_logp = the frozen π_old anchor established by prepare_segment: - # the rollout log-prob by default, or a train-side replay under - # old_logp_source='replay'. Either way it stays frozen across all - # num_updates_per_batch steps (see the supports_multi_update comment). + loss, clip_range, ratio_metrics = self._build_policy_loss( + segment=segment, + new_logp=new_logp, + advantages=advantages, + training_progress=training_progress, + ) + (loss * loss_scale).backward() + + rollout_logp = (segment.rollout_log_probs if segment.rollout_log_probs is not None else segment.log_probs).to( + dtype=new_logp.dtype, device=new_logp.device + ) + metrics: Dict[str, Any] = { + "policy_loss": float(loss.detach().item()), + "clip_range": float(clip_range), + **rollout_replay_logp_absdiff(new_logp, rollout_logp), + **rollout_replay_k3(new_logp, rollout_logp), + **{k: float(v.item()) for k, v in ratio_metrics.items()}, + } + return AlgorithmStepResult( + loss=float(loss.detach().item()), + metrics=metrics, + num_steps_or_tokens=int(new_logp.shape[0]), + has_backward=True, + ) + + def _build_policy_loss( + self, + *, + segment: "TextSegment", + new_logp: torch.Tensor, + advantages: torch.Tensor, + training_progress: float, + ) -> tuple[torch.Tensor, float, Dict[str, torch.Tensor]]: + """Build the clipped loss against the frozen anchor without running replay or backward.""" old_logp = segment.log_probs.to(dtype=new_logp.dtype, device=new_logp.device) adv_per_token = self._expand_advantages_to_tokens( advantages, segment.lengths, dtype=new_logp.dtype, device=new_logp.device @@ -164,24 +194,7 @@ def compute_loss_and_backward( loss = loss_per_elem.mean() else: loss = loss_per_elem.sum() / mask.sum().clamp(min=1) - (loss * loss_scale).backward() - - rollout_logp = (segment.rollout_log_probs if segment.rollout_log_probs is not None else segment.log_probs).to( - dtype=new_logp.dtype, device=new_logp.device - ) - metrics: Dict[str, Any] = { - "policy_loss": float(loss.detach().item()), - "clip_range": float(clip_range), - **rollout_replay_logp_absdiff(new_logp, rollout_logp), - **rollout_replay_k3(new_logp, rollout_logp), - **{k: float(v.item()) for k, v in ratio_metrics.items()}, - } - return AlgorithmStepResult( - loss=float(loss.detach().item()), - metrics=metrics, - num_steps_or_tokens=int(new_logp.shape[0]), - has_backward=True, - ) + return loss, clip_range, ratio_metrics @staticmethod def _expand_advantages_to_tokens( diff --git a/unirl/distributed/weight_sync/full/ipc.py b/unirl/distributed/weight_sync/full/ipc.py index cacd98ba3..8fd7d0d1d 100644 --- a/unirl/distributed/weight_sync/full/ipc.py +++ b/unirl/distributed/weight_sync/full/ipc.py @@ -354,6 +354,15 @@ def _poison_rollout(self, reason: str) -> None: except BaseException: logger.exception("Failed to poison vLLM native IPC runtime") + def _on_native_publication_committed( + self, + *, + header: dict[str, Any], + result: dict[str, Any], + ) -> None: + """Optional subclass hook after native publication validation.""" + del header, result + @distributed(dispatch_mode=Dispatch.BROADCAST) def sync(self) -> None: """Select the IPC protocol from the concrete rollout engine.""" @@ -457,6 +466,7 @@ def _verify_actor_source() -> None: fanout=tp_world_size, flush_cache=self._flush_cache, ) + self._on_native_publication_committed(header=header, result=result) self._next_model_version += 1 success = True logger.info( diff --git a/unirl/models/qwen3/ar.py b/unirl/models/qwen3/ar.py index be2c329a0..ab17f3ce7 100644 --- a/unirl/models/qwen3/ar.py +++ b/unirl/models/qwen3/ar.py @@ -24,6 +24,8 @@ from .conditions import Qwen3ARConditions logger = logging.getLogger(__name__) +_EXACT_CONTEXT_FACTORY = None +_EXACT_LOG_SOFTMAX = None _SPARSE_PACKED_ATTN = ("flex_attention", "flash_attention_2", "flash_attention_3", "flash_attention_4") @@ -52,6 +54,49 @@ def _packed_replay_supported(attn_impl: Optional[str]) -> bool: return True +def _exact_actor_enabled(model: Any) -> bool: + return bool(getattr(model, "_unirl_exact_actor_logprobs", False)) + + +def _exact_actor_context(model: Any): + if not _exact_actor_enabled(model): + return nullcontext() + if _EXACT_CONTEXT_FACTORY is None: + raise RuntimeError( + "exact_actor_logprobs requires an exact actor provider; load the train-inference parity FSDP adaptor first" + ) + return _EXACT_CONTEXT_FACTORY() + + +def _selected_token_log_probs( + model: Any, + hidden: torch.Tensor, + tokens: torch.Tensor, + temperature: float, +) -> torch.Tensor: + with _exact_actor_context(model): + logits = model.lm_head(hidden).float() / temperature + if _exact_actor_enabled(model): + if _EXACT_LOG_SOFTMAX is None: + raise RuntimeError("exact_actor_logprobs requires an exact log-softmax provider") + return _EXACT_LOG_SOFTMAX(logits, dim=-1).gather(-1, tokens.unsqueeze(-1)).squeeze(-1) + chosen = logits.gather(-1, tokens.unsqueeze(-1)).squeeze(-1) + return chosen - torch.logsumexp(logits, dim=-1) + + +def register_exact_actor_provider(*, context_factory, log_softmax) -> None: + """Register a process-local exact provider without importing experimental code.""" + global _EXACT_CONTEXT_FACTORY, _EXACT_LOG_SOFTMAX + if not callable(context_factory) or not callable(log_softmax): + raise TypeError("exact actor providers must be callable") + if _EXACT_CONTEXT_FACTORY is not None and ( + _EXACT_CONTEXT_FACTORY is not context_factory or _EXACT_LOG_SOFTMAX is not log_softmax + ): + raise RuntimeError("a different exact actor provider is already registered in this process") + _EXACT_CONTEXT_FACTORY = context_factory + _EXACT_LOG_SOFTMAX = log_softmax + + def _replay_aware_forward( self: Any, *, @@ -80,8 +125,9 @@ def _replay_aware_forward( autocast_ctx = ( torch.autocast("cuda", autocast_dtype) if autocast_dtype in (torch.float16, torch.bfloat16) else nullcontext() ) - with autocast_ctx: - hidden = self.model(**kw, use_cache=False, return_dict=True).last_hidden_state + with _exact_actor_context(self): + with autocast_ctx: + hidden = self.model(**kw, use_cache=False, return_dict=True).last_hidden_state T = float(temperature) if float(temperature) > 0.0 else 1.0 value_head = getattr(self, "value_head", None) if return_values else None @@ -91,8 +137,7 @@ def _replay_aware_forward( targets = response_tokens def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: - lf = self.lm_head(h).float() / T - return lf.gather(-1, tok.unsqueeze(-1)).squeeze(-1) - torch.logsumexp(lf, dim=-1) + return _selected_token_log_probs(self, h, tok, T) flat_parts: List[torch.Tensor] = [] flat_chunk = 2048 @@ -118,9 +163,7 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: resp_hidden = hidden[:, prompt_len - 1 : prompt_len - 1 + T_max, :] def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: - lf = self.lm_head(h).float() / T - chosen = lf.gather(-1, tok.unsqueeze(-1)).squeeze(-1) - return chosen - torch.logsumexp(lf, dim=-1) + return _selected_token_log_probs(self, h, tok, T) bsz = resp_hidden.size(0) chunk = max(64, 2048 // max(1, bsz)) @@ -248,11 +291,13 @@ def __init__( model: Qwen3Bundle, autocast_precision: str = "bf16", logprob_precision: str = "fp32", + exact_actor_logprobs: bool = False, ) -> None: self.model = model self.autocast_dtype = parse_torch_dtype(autocast_precision, field_name="Qwen3ARStage.autocast_precision") self.logprob_dtype = parse_torch_dtype(logprob_precision, field_name="Qwen3ARStage.logprob_precision") transformer = model.transformer + transformer._unirl_exact_actor_logprobs = bool(exact_actor_logprobs) if getattr(transformer.forward, "__func__", None) is not _replay_aware_forward: transformer.forward = MethodType(_replay_aware_forward, transformer) @@ -260,6 +305,10 @@ def trainable_module(self) -> "torch.nn.Module": """Return the HF causal LM module — the FSDP/LoRA wrap target.""" return self.model.transformer + def exact_context(self): + """Scope replay, parity gate, loss, backward, and checkpoint recomputation.""" + return _exact_actor_context(self.model.transformer) + def autoregress( self, conditions: Qwen3ARConditions, @@ -356,7 +405,7 @@ def replay( temperature: float = 1.0, return_values: bool = False, ) -> Union[torch.Tensor, ReplayResult]: - """Per-token log-prob replay; falls back to the dense ``[B, P_max + T_max]`` :meth:`padding_replay`.""" + """Score fixed response tokens with a full-sequence teacher-forcing forward.""" _require_value_head_for_replay(self.model.transformer, return_values) attn_impl = getattr(getattr(self.model.transformer, "config", None), "_attn_implementation", None) if _packed_replay_supported(attn_impl): @@ -593,4 +642,9 @@ def _pack_text_segment( ) -__all__ = ["Qwen3ARParams", "Qwen3ARStage", "Qwen3ARStep"] +__all__ = [ + "Qwen3ARParams", + "Qwen3ARStage", + "Qwen3ARStep", + "register_exact_actor_provider", +] diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py index 38a47c0b9..d8c3298bf 100644 --- a/unirl/models/qwen3/bundle.py +++ b/unirl/models/qwen3/bundle.py @@ -2,6 +2,7 @@ from __future__ import annotations +import importlib import logging from typing import Any @@ -58,6 +59,14 @@ def __init__( @classmethod def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": """Load the Qwen3 transformer + tokenizer from a HuggingFace-layout checkpoint.""" + for module_name in config.external_libs or (): + module = importlib.import_module(str(module_name)) + installer = getattr(module, "install", None) + if installer is not None: + if not callable(installer): + raise TypeError(f"Qwen3 external library {module_name!r} exposes a non-callable install") + installer() + from transformers import AutoModelForCausalLM, AutoTokenizer path = config.pretrained_model_ckpt_path @@ -70,13 +79,24 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": dtype = parse_torch_dtype(config.model_precision, field_name="model_precision") if config.meta_init_transformer: - transformer_weights_path = resolve_meta_init_weights(path) + transformer_weights_path = resolve_meta_init_weights(path, revision=config.model_revision) # Restore non-persistent RoPE buffers after meta initialization. from transformers import AutoConfig - hf_config = AutoConfig.from_pretrained(path, trust_remote_code=bool(config.trust_remote_code)) + hf_config = AutoConfig.from_pretrained( + path, + revision=config.model_revision, + trust_remote_code=bool(config.trust_remote_code), + ) + model_kwargs = {} + if getattr(config, "attn_implementation", None): + model_kwargs["attn_implementation"] = str(config.attn_implementation) transformer, meta_init_state = build_meta_init_transformer( - lambda: AutoModelForCausalLM.from_config(hf_config, trust_remote_code=bool(config.trust_remote_code)), + lambda: AutoModelForCausalLM.from_config( + hf_config, + trust_remote_code=bool(config.trust_remote_code), + **model_kwargs, + ), dtype=dtype, ) else: @@ -85,6 +105,7 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": load_kwargs["attn_implementation"] = str(config.attn_implementation) transformer = AutoModelForCausalLM.from_pretrained( path, + revision=config.model_revision, torch_dtype=dtype, trust_remote_code=bool(config.trust_remote_code), **load_kwargs, @@ -102,6 +123,7 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": tokenizer = AutoTokenizer.from_pretrained( tokenizer_path, + revision=config.model_revision, trust_remote_code=bool(config.trust_remote_code), ) if tokenizer.pad_token is None and tokenizer.eos_token is not None: diff --git a/unirl/models/qwen3/config.py b/unirl/models/qwen3/config.py index 43996b26d..0b90305b0 100644 --- a/unirl/models/qwen3/config.py +++ b/unirl/models/qwen3/config.py @@ -15,6 +15,7 @@ class Qwen3PipelineConfig: pretrained_model_ckpt_path: str tokenizer_ckpt_path: Optional[str] = None trust_remote_code: bool = True + model_revision: Optional[str] = None model_precision: Any = "bf16" attn_implementation: Optional[str] = None @@ -22,6 +23,7 @@ class Qwen3PipelineConfig: autocast_precision: str = "bf16" logprob_precision: str = "fp32" + exact_actor_logprobs: bool = False use_gradient_checkpointing: bool = False @@ -33,6 +35,7 @@ class Qwen3PipelineConfig: lora_target_modules: Optional[List[str]] = None use_value_head: bool = False + external_libs: Optional[List[str]] = None system_instruction: Optional[str] = None enable_thinking: bool = False @@ -40,6 +43,8 @@ class Qwen3PipelineConfig: def __post_init__(self) -> None: validate_precision_type(self.model_precision, field="Qwen3PipelineConfig.model_precision") + if self.model_revision is not None and not str(self.model_revision).strip(): + raise ValueError("Qwen3PipelineConfig.model_revision must be non-empty when set") __all__ = ["Qwen3PipelineConfig"] diff --git a/unirl/models/qwen3/pipeline.py b/unirl/models/qwen3/pipeline.py index ebfbdbec7..53d192c1d 100644 --- a/unirl/models/qwen3/pipeline.py +++ b/unirl/models/qwen3/pipeline.py @@ -27,6 +27,7 @@ def __init__( ar: Optional[Qwen3ARStage] = None, autocast_precision: str = "bf16", logprob_precision: str = "fp32", + exact_actor_logprobs: bool = False, ) -> None: super().__init__() self.bundle = bundle @@ -34,7 +35,12 @@ def __init__( self.ar = ( ar if ar is not None - else Qwen3ARStage(model=bundle, autocast_precision=autocast_precision, logprob_precision=logprob_precision) + else Qwen3ARStage( + model=bundle, + autocast_precision=autocast_precision, + logprob_precision=logprob_precision, + exact_actor_logprobs=exact_actor_logprobs, + ) ) @classmethod @@ -47,6 +53,7 @@ def from_bundle( logprob_precision: str = "fp32", enable_thinking: bool = False, max_prompt_length: int = 4096, + exact_actor_logprobs: bool = False, ) -> "Qwen3Pipeline": """Wire chat-template + AR stages around an already-loaded bundle.""" chat_template = Qwen3ChatTemplateStage( @@ -59,6 +66,7 @@ def from_bundle( model=bundle, autocast_precision=autocast_precision, logprob_precision=logprob_precision, + exact_actor_logprobs=exact_actor_logprobs, ) return cls( bundle=bundle, @@ -66,6 +74,7 @@ def from_bundle( ar=ar, autocast_precision=autocast_precision, logprob_precision=logprob_precision, + exact_actor_logprobs=exact_actor_logprobs, ) @classmethod @@ -82,6 +91,7 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Pipeline": model=bundle, autocast_precision=config.autocast_precision, logprob_precision=config.logprob_precision, + exact_actor_logprobs=config.exact_actor_logprobs, ) return cls(bundle=bundle, chat_template=chat_template, ar=ar) diff --git a/unirl/models/types/meta_init.py b/unirl/models/types/meta_init.py index d84f2cd62..3a08d692d 100644 --- a/unirl/models/types/meta_init.py +++ b/unirl/models/types/meta_init.py @@ -60,7 +60,9 @@ def _validate_safetensors(weights_path: str, checkpoint_path: str, expected: str ) -def resolve_meta_init_weights(checkpoint_path: str, *, component: Optional[str] = None) -> str: +def resolve_meta_init_weights( + checkpoint_path: str, *, component: Optional[str] = None, revision: Optional[str] = None +) -> str: """Resolve and validate the local safetensors directory for a meta-init bundle.""" snapshot_path = checkpoint_path expected = os.path.join(component, "*.safetensors") if component else "*.safetensors" @@ -70,6 +72,7 @@ def resolve_meta_init_weights(checkpoint_path: str, *, component: Optional[str] try: snapshot_path = snapshot_download( repo_id=checkpoint_path, + revision=revision, allow_patterns=[expected, f"{expected}.index.json"], ) except Exception as exc: diff --git a/unirl/rollout/engine/vllm/engine.py b/unirl/rollout/engine/vllm/engine.py index 8d68fe054..a1418acd1 100644 --- a/unirl/rollout/engine/vllm/engine.py +++ b/unirl/rollout/engine/vllm/engine.py @@ -15,7 +15,7 @@ import torch from unirl.config.require import require -from unirl.distributed.group.dispatch import Dispatch, distributed +from unirl.distributed.group.dispatch import Dispatch, Execute, distributed from unirl.rollout.engine.base import BaseRolloutEngine from unirl.rollout.engine.sglang.adapters.text import TextLMAdapter from unirl.rollout.engine.vllm.config import VLLMEngineConfig @@ -105,6 +105,7 @@ def __init__( self._next_request_id = 1 self._pending_native_publication = None self._ipc_worker_device_uuids: List[str] = [] + self._worker_capabilities: List[Dict[str, Any]] = [] self.adapter = None if not self._is_tp_zero: @@ -162,13 +163,14 @@ def __init__( if not isinstance(ready, dict) or ready.get("event") != "ready": raise _ProtocolError(f"vLLM returned invalid startup result: {ready!r}") self._validate_runtime_manifest(ready) - self._ipc_worker_device_uuids = [ - str(item["cuda_device_uuid"]) + self._worker_capabilities = [ + dict(item) for item in sorted( ready["worker_capabilities"], key=lambda capability: int(capability["tp_rank"]), ) ] + self._ipc_worker_device_uuids = [str(item["cuda_device_uuid"]) for item in self._worker_capabilities] self._process_group_id = int(ready["process_group_id"]) except BaseException: child.close() @@ -190,6 +192,11 @@ def weight_payload_fanout(self) -> int: def ipc_worker_device_uuids(self) -> List[str]: return list(self._ipc_worker_device_uuids) + @distributed(dispatch_mode=Dispatch.BROADCAST, execute_mode=Execute.RANK_ZERO) + def runtime_capabilities(self) -> List[Dict[str, Any]]: + """Return startup capabilities reported by every native vLLM worker.""" + return [dict(item) for item in self._worker_capabilities] + @property def connection_state(self) -> VLLMConnectionState: with self._lock: diff --git a/unirl/trainer/ar.py b/unirl/trainer/ar.py index 18623478b..90fc8ac5b 100644 --- a/unirl/trainer/ar.py +++ b/unirl/trainer/ar.py @@ -5,7 +5,7 @@ import math import time from contextlib import contextmanager, nullcontext -from typing import Dict, Iterator, Optional, Set, Tuple +from typing import Callable, Dict, Iterator, Optional, Set, Tuple import torch from hydra.utils import get_class, get_object, instantiate @@ -677,6 +677,7 @@ def train( save_dir: Optional[str] = None, load_dir: Optional[str] = None, save_mode: str = "auto", + on_rollout_complete: Optional[Callable[..., None]] = None, ) -> None: """Minimal training loop: ``num_rollouts`` iterations of ``train_step``.""" interval = max(1, weight_sync_interval) @@ -720,6 +721,13 @@ def train( preserve_rollout_weights=preserve_rollout_weights, ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) + if on_rollout_complete is not None: + on_rollout_complete( + rollout_id=rollout_id, + result=result, + mean_reward=mean_reward, + sync_weights=sync_weights, + ) if self.eval_interval > 0 and (rollout_id + 1) % self.eval_interval == 0: next_rollout_syncs = next_rollout_id < num_rollouts and next_rollout_id % interval == 0 self.evaluate(