diff --git a/experimental/train_inference_parity/README.md b/experimental/train_inference_parity/README.md new file mode 100644 index 000000000..5394357ad --- /dev/null +++ b/experimental/train_inference_parity/README.md @@ -0,0 +1,186 @@ +# 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.12, 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.12, Torch 2.13.0+cu130, 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` writes JSON by fsyncing a same-directory temporary file and +atomically replacing the destination. 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. 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..c750340b4 --- /dev/null +++ b/experimental/train_inference_parity/algorithm.py @@ -0,0 +1,192 @@ +"""Experimental GRPO with collective-safe exact replay/rollout parity.""" + +from __future__ import annotations + +from contextlib import ExitStack +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, + _grpo_clip_loss, + _resolve_clip_range_from_schedule, + 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: + super().prepare_segment(conditions=conditions, segment=segment) + 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") + 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 _build_policy_loss( + self, + *, + segment: "TextSegment", + new_logp: torch.Tensor, + advantages: torch.Tensor, + training_progress: float, + ) -> tuple[torch.Tensor, float, Dict[str, torch.Tensor]]: + old_logp = segment.rollout_log_probs.to(device=new_logp.device) + adv_per_token = self._expand_advantages_to_tokens( + advantages, + segment.lengths, + dtype=new_logp.dtype, + device=new_logp.device, + ) + clip_range = _resolve_clip_range_from_schedule(self.clip_range, self.clip_schedule, training_progress) + clip_high = ( + None + if self.clip_range_high is None + else _resolve_clip_range_from_schedule( + self.clip_range_high, + self.clip_schedule, + training_progress, + ) + ) + loss_per_elem, ratio_metrics = _grpo_clip_loss( + new_logp=new_logp, + old_logp=old_logp, + advantages=adv_per_token, + clip_range=clip_range, + clip_range_high=clip_high, + ) + mask: Optional[torch.Tensor] = None + if segment.loss_mask is not None: + mask = segment.loss_mask.to(dtype=loss_per_elem.dtype, device=loss_per_elem.device) + loss_per_elem = loss_per_elem * mask + + if self.loss_agg_mode in ("seq-mean-token-sum-norm", "seq-mean-token-mean"): + parts = torch.split(loss_per_elem, segment.lengths.tolist()) + if mask is None: + if self.loss_agg_mode == "seq-mean-token-sum-norm": + loss = torch.stack([part.sum() for part in parts]).mean() / float(self.horizon) + else: + loss = torch.stack([part.mean() if part.numel() else part.new_zeros(()) for part in parts]).mean() + else: + mask_parts = torch.split(mask, segment.lengths.tolist()) + valid_parts = [ + (part, float(mask_part.sum().item())) + for part, mask_part in zip(parts, mask_parts) + if bool(mask_part.any()) + ] + if self.loss_agg_mode == "seq-mean-token-sum-norm": + per_seq = [part.sum() / float(self.horizon) for part, _ in valid_parts] + else: + per_seq = [part.sum() / weight for part, weight in valid_parts] + loss = torch.stack(per_seq).mean() if per_seq else loss_per_elem.sum() * 0.0 + elif mask is None: + loss = loss_per_elem.mean() + else: + loss = loss_per_elem.sum() / mask.sum().clamp(min=1) + return loss, clip_range, ratio_metrics + + 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}" + # Replay and FSDP forwards have their own collectives. Every rank must + # agree to enter them, or every rank must fail before they start. + collective_error_barrier(setup_error, phase="pre-replay") + + # Keep exact mode active through backward. Non-reentrant activation + # checkpointing reruns forward operators from inside ``backward()``. + with ExitStack() as exact_stack: + exact_stack.enter_context(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 + + # No rank raises a local shape/dtype/finite/setup error before this + # fixed collective. Every rank either continues or fails together. + 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/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..dd7eca67f --- /dev/null +++ b/experimental/train_inference_parity/gate.py @@ -0,0 +1,260 @@ +"""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 a rank-local setup error before every rank enters the next collective. + + Replay and FSDP forwards perform their own collectives. A rank that raises + before those calls never reaches the later exact gate, so the remaining + ranks hang. This barrier is the pre-replay meeting point. + """ + 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..6c279fe0c --- /dev/null +++ b/experimental/train_inference_parity/models/qwen3_moe_30b_a3b/contract.py @@ -0,0 +1,113 @@ +"""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: + 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..9e1e21f82 --- /dev/null +++ b/experimental/train_inference_parity/runtime_contract.py @@ -0,0 +1,427 @@ +"""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 + 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("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("logging.report_to_wandb", False) + 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 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..d5d95c36f --- /dev/null +++ b/experimental/train_inference_parity/verification.py @@ -0,0 +1,740 @@ +"""Machine-readable evidence writer for the two-phase TP4 parity run.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import math +import os +import platform +import subprocess +import tempfile +from dataclasses import asdict, is_dataclass +from pathlib import Path +from typing import Any, Dict, Mapping, Optional + +from .runtime_contract import ( + EXPECTED_TP, + EXPECTED_VERSIONS, + PLUGIN_DISTRIBUTION, + 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", +) +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] + + +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: + """Write one durable JSON value via same-directory fsync + replace.""" + target = Path(path).expanduser().resolve() + target.parent.mkdir(parents=True, exist_ok=True) + tmp_path: Optional[Path] = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=target.parent, + prefix=f".{target.name}.", + suffix=".tmp", + delete=False, + ) as handle: + tmp_path = 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(tmp_path, target) + directory_fd = os.open(target.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if tmp_path is not None and tmp_path.exists(): + tmp_path.unlink() + + +def _run_git(*args: str) -> str: + result = subprocess.run( + ["git", "-C", str(_REPOSITORY_ROOT), *args], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=10, + ) + return result.stdout.strip() + + +_GIT_DIRTY_ALLOWLIST_PREFIXES = ( + "outputs/", + "experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm.egg-info/", +) +_GIT_DIRTY_ALLOWLIST_SUFFIXES = (".egg-info", ".egg-info/") + + +def _git_path_from_porcelain(line: str) -> str: + path = line[3:] if len(line) > 3 else line + if " -> " in path: + path = path.split(" -> ", 1)[1] + return path.strip() + + +def _git_path_is_allowlisted(path: str) -> bool: + normalized = path.replace("\\", "/").lstrip("./") + if any(normalized.startswith(prefix) for prefix in _GIT_DIRTY_ALLOWLIST_PREFIXES): + return True + return any( + normalized.endswith(suffix) or f"{normalized}/".endswith(suffix) for suffix in _GIT_DIRTY_ALLOWLIST_SUFFIXES + ) + + +def _git_state() -> Dict[str, Any]: + try: + commit = _run_git("rev-parse", "HEAD") + # Tracked edits always count. Untracked generated outputs are ignored; + # untracked Python/source next to the experiment still makes the tree dirty. + porcelain = _run_git("status", "--porcelain") + dirty_lines = [] + for line in porcelain.splitlines(): + if not line.strip(): + continue + path = _git_path_from_porcelain(line) + if line.startswith("??") and _git_path_is_allowlisted(path): + continue + if line.startswith("??") and not path.endswith((".py", ".yaml", ".yml", ".toml", ".json")): + continue + dirty_lines.append(line) + return { + "commit": commit, + "dirty": bool(dirty_lines), + "dirty_scope": "tracked_and_untracked_source", + "dirty_paths": dirty_lines, + } + except (FileNotFoundError, 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) -> Optional[str]: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def _software_state() -> Dict[str, Any]: + state: Dict[str, Any] = { + "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 + driver_query = getattr(torch._C, "_cuda_getDriverVersion", None) + if callable(driver_query): + try: + driver_version = driver_query() + except RuntimeError: + 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 _source_revision(path: Path) -> Optional[str]: + parts = path.resolve().parts + if "snapshots" in parts: + index = parts.index("snapshots") + if index + 1 < len(parts): + return parts[index + 1] + return None + + +def _model_fingerprint(raw_path: str) -> Dict[str, Any]: + """Fingerprint model identity without rereading every multi-GB weight byte.""" + 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()): + relative = item.relative_to(path).as_posix() + size = item.stat().st_size + entry: Dict[str, Any] = { + "path": relative, + "size_bytes": size, + "sha256": _sha256_file(item), + } + entries.append(entry) + digest.update(json.dumps(entry, sort_keys=True, separators=(",", ":")).encode("utf-8")) + digest.update(b"\n") + return { + "path": str(path), + "kind": "file_tree_bytes", + "algorithm": "sha256(relative_path,size,file_sha256)", + "sha256": digest.hexdigest(), + "revision": _source_revision(path), + "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, + } + + +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 () + if all(name in (getattr(micro, "metrics", {}) or {}) for name in CANONICAL_PARITY_METRICS) + ] + 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") + 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 = [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}" + 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/experimental/train_inference_parity/vllm_plugin/pyproject.toml b/experimental/train_inference_parity/vllm_plugin/pyproject.toml new file mode 100644 index 000000000..9678e4308 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/pyproject.toml @@ -0,0 +1,24 @@ +[build-system] +requires = ["setuptools>=61", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "unirl-train-inference-parity-vllm" +version = "0.1.0" +description = "Experimental UniRL train/inference parity plugin for vLLM" +requires-python = ">=3.12,<3.14" +license = { text = "Apache-2.0" } + +[project.optional-dependencies] +verified = [ + "vllm==0.27.0", + "torch==2.13.0", + "transformers==5.6.0", +] + +[project.entry-points."vllm.general_plugins"] +unirl_train_inference_parity = "unirl_train_inference_parity_vllm:register" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["unirl_train_inference_parity_vllm*"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/__init__.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/__init__.py new file mode 100644 index 000000000..62440935c --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/__init__.py @@ -0,0 +1,87 @@ +"""vLLM general-plugin entry point for UniRL train/inference parity.""" + +from __future__ import annotations + +import os + +_RUNTIME_MANIFEST: dict[str, object] | None = None + + +def _install_capability_manifest_bridge() -> None: + """Attach this worker's plugin manifest to UniRL's generic capability RPC.""" + from unirl.rollout.engine.vllm.worker_extension import UniRLWeightSyncExtension + + original = UniRLWeightSyncExtension.unirl_weight_sync_capabilities + if getattr(original, "_unirl_parity_manifest_bridge", False): + return + + def with_parity_manifest(self): + capabilities = dict(original(self)) + manifest = runtime_manifest() + if manifest is None: + raise RuntimeError("parity plugin manifest is unavailable after registration") + capabilities["train_inference_parity"] = manifest + return capabilities + + with_parity_manifest._unirl_parity_manifest_bridge = True + UniRLWeightSyncExtension.unirl_weight_sync_capabilities = with_parity_manifest + + +def register() -> tuple[str, ...]: + global _RUNTIME_MANIFEST + if os.environ.get("UNIRL_PARITY_ENABLE") != "1": + _RUNTIME_MANIFEST = None + return () + + import json + + from .compat import validate_runtime_versions + from .config import load_config + from .registry import Installer, install_selected + + config = load_config() + if config.model != "qwen3_moe_30b_a3b": + raise ValueError(f"unsupported parity model {config.model!r}") + versions = validate_runtime_versions() + from .common import install_common, preflight_common + from .models.qwen3_moe_30b_a3b import ( + install_qwen3_moe_patch, + preflight_qwen3_moe_patch, + ) + + installers = { + "common": Installer(preflight=preflight_common, install=install_common), + "qwen3_moe_30b_a3b": Installer( + preflight=preflight_qwen3_moe_patch, + install=install_qwen3_moe_patch, + ), + } + installed = install_selected( + config.patches, + strict=config.strict, + installers=installers, + ) + manifest = { + "entrypoint": "unirl_train_inference_parity", + "model": config.model, + "patches": [result.to_manifest() for result in installed], + "pid": os.getpid(), + "profile": config.profile, + "runtime": [result.to_manifest() for result in versions], + "source_file": __file__, + "strict": config.strict, + } + _RUNTIME_MANIFEST = manifest + _install_capability_manifest_bridge() + print( + "[unirl.parity.vllm] manifest=" + json.dumps(manifest, sort_keys=True, separators=(",", ":")), + flush=True, + ) + return tuple(result.name for result in installed) + + +def runtime_manifest() -> dict[str, object] | None: + return None if _RUNTIME_MANIFEST is None else dict(_RUNTIME_MANIFEST) + + +__all__ = ["register", "runtime_manifest"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/__init__.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/__init__.py new file mode 100644 index 000000000..a7c6f5a2d --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/__init__.py @@ -0,0 +1,36 @@ +"""Model-independent vLLM parity patches.""" + +from ..registry import PatchResult +from .attention import install_attention_contract, preflight_attention_contract +from .moe_combine import moe_combine +from .norm import install_norm_patch, preflight_norm_patch +from .precision import install_precision_contract, preflight_precision_contract +from .providers import preflight_providers +from .reductions import install_reduction_patch, preflight_reduction_patch + + +def preflight_common(strict: bool) -> None: + preflight_providers(strict=strict) + preflight_precision_contract(strict=strict) + preflight_reduction_patch(strict=strict) + preflight_norm_patch(strict=strict) + preflight_attention_contract(strict=strict) + + +def install_common(*, strict: bool) -> PatchResult: + if not strict: + raise ValueError("the public-reference common installer requires strict=True") + symbols = ( + *install_precision_contract(), + *install_reduction_patch(), + *install_norm_patch(), + *install_attention_contract(), + ) + return PatchResult(name="common", symbols=symbols) + + +__all__ = [ + "install_common", + "moe_combine", + "preflight_common", +] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/attention.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/attention.py new file mode 100644 index 000000000..4087ee646 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/attention.py @@ -0,0 +1,152 @@ +"""Activate only vLLM's attention batch-invariance gates.""" + +from __future__ import annotations + +import importlib +import types + +from ..compat import require_symbol, require_value +from ..registry import SymbolResult, symbol_result, value_result + +_INSTALLED = False +_ENV_MODULES = ( + "vllm.model_executor.layers.attention.attention", + "vllm.v1.attention.backends.flash_attn", + "vllm.v1.attention.backends.fa_utils", +) +_TRITON_MODULE = "vllm.v1.attention.ops.triton_unified_attention" +_FLASH_MODULE = "vllm.v1.attention.backends.flash_attn" + + +class _EnvProxy: + def __init__(self, original): + self._original = original + + def __getattr__(self, name): + if name == "VLLM_BATCH_INVARIANT": + return True + return getattr(self._original, name) + + +def _probe_fa3() -> tuple[object, type]: + backend = require_symbol( + _FLASH_MODULE, + "FlashAttentionBackend", + origin=f"{_FLASH_MODULE}.FlashAttentionBackend", + strict=True, + ) + get_version = require_symbol( + _FLASH_MODULE, + "get_flash_attn_version", + parameters=( + ("requires_alibi", "False"), + ("head_size", "None"), + ("head_size_v", "None"), + ("has_sinks", "False"), + ("requires_local_attention", "False"), + ), + origin="vllm.v1.attention.backends.fa_utils.get_flash_attn_version", + strict=True, + ) + if get_version() != 3: + raise RuntimeError(f"parity requires vLLM FlashAttention 3, got {get_version()!r}") + if backend.get_name() != "FLASH_ATTN" or backend.supports_batch_invariance() is not True: + raise RuntimeError("vLLM FlashAttention backend does not satisfy the batch-invariant contract") + implementation = backend.get_impl_cls() + if ( + getattr(implementation, "__module__", None) != _FLASH_MODULE + or getattr(implementation, "__name__", None) != "FlashAttentionImpl" + ): + raise RuntimeError(f"unexpected vLLM FlashAttention implementation {implementation!r}") + facade = importlib.import_module("vllm.vllm_flash_attn") + varlen = getattr(facade, "flash_attn_varlen_func", None) + if not callable(varlen) or getattr(varlen, "__module__", None) != ("vllm.vllm_flash_attn.flash_attn_interface"): + raise RuntimeError(f"unexpected vLLM FA3 varlen provider {varlen!r}") + return varlen, implementation + + +def preflight_attention_contract(*, strict: bool) -> None: + if not strict: + raise ValueError("the parity attention contract requires strict=True") + for module_name in _ENV_MODULES: + envs = require_value( + module_name, + "envs", + expected_type=types.ModuleType, + ) + if envs.__name__ != "vllm.envs": + raise RuntimeError( + f"conflicting attention env provider at {module_name}.envs: expected vllm.envs, got {envs.__name__}" + ) + require_value( + module_name, + "envs.VLLM_BATCH_INVARIANT", + expected_type=bool, + ) + require_value( + _TRITON_MODULE, + "is_batch_invariant", + expected_type=bool, + ) + _probe_fa3() + + +def install_attention_contract() -> tuple[SymbolResult, ...]: + global _INSTALLED + if _INSTALLED: + raise RuntimeError("attention parity contract installed twice") + fa3_varlen, fa3_implementation = _probe_fa3() + results = [] + for module_name in _ENV_MODULES: + module = importlib.import_module(module_name) + original = module.envs + proxy = _EnvProxy(original) + module.envs = proxy + results.append( + symbol_result( + f"{module_name}.envs", + proxy, + before=original, + actual=module.envs, + ) + ) + triton_module = importlib.import_module(_TRITON_MODULE) + original_batch_invariant = triton_module.is_batch_invariant + triton_module.is_batch_invariant = True + results.append( + value_result( + f"{_TRITON_MODULE}.is_batch_invariant", + "literal:true", + before=original_batch_invariant, + actual=triton_module.is_batch_invariant, + verified=triton_module.is_batch_invariant is True, + ) + ) + results.extend( + ( + symbol_result( + "vllm.vllm_flash_attn.flash_attn_varlen_func", + fa3_varlen, + before=fa3_varlen, + actual=fa3_varlen, + ), + symbol_result( + f"{_FLASH_MODULE}.FlashAttentionImpl", + fa3_implementation, + before=fa3_implementation, + actual=fa3_implementation, + ), + value_result( + f"{_FLASH_MODULE}.get_flash_attn_version()", + "literal:3", + before=3, + actual=3, + verified=True, + ), + ) + ) + _INSTALLED = True + return tuple(results) + + +__all__ = ["install_attention_contract", "preflight_attention_contract"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/moe_combine.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/moe_combine.py new file mode 100644 index 000000000..2a6bfee4c --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/moe_combine.py @@ -0,0 +1,31 @@ +"""Shared correctness-first MoE output combination.""" + +from __future__ import annotations + +import torch + + +def moe_combine( + contributions: torch.Tensor, + row_map: torch.Tensor, +) -> torch.Tensor: + """Combine top-k expert rows with the parity rounding contract.""" + if contributions.dim() != 2 or row_map.dim() != 2: + raise ValueError("moe_combine requires 2D contributions and row_map tensors") + tokens, topk = row_map.shape + hidden = int(contributions.shape[-1]) + total = torch.zeros( + (tokens, hidden), + dtype=torch.float32, + device=contributions.device, + ) + zero = torch.zeros_like(total) + for slot in range(topk): + rows = row_map[:, slot].long() + contribution = contributions.index_select(0, rows.clamp_min(0)).float() + total = total + torch.where((rows >= 0)[:, None], contribution, zero) + total = total.to(torch.bfloat16).float() + return total.to(torch.bfloat16) + + +__all__ = ["moe_combine"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/norm.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/norm.py new file mode 100644 index 000000000..17df9cd7f --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/norm.py @@ -0,0 +1,52 @@ +"""Route vLLM RMSNorm layers to the public BI primitive.""" + +from __future__ import annotations + +from ..compat import require_symbol +from ..registry import SymbolResult, symbol_result + +_INSTALLED = False + + +def preflight_norm_patch(*, strict: bool) -> None: + require_symbol( + "vllm.model_executor.layers.layernorm", + "RMSNorm.forward_cuda", + parameters=(("self", ""), ("x", ""), ("residual", "None")), + origin="vllm.model_executor.layers.layernorm.RMSNorm.forward_cuda", + strict=strict, + ) + + +def install_norm_patch() -> tuple[SymbolResult, ...]: + global _INSTALLED + if _INSTALLED: + raise RuntimeError("RMSNorm parity patch installed twice") + from vllm.model_executor.layers.layernorm import RMSNorm + + from .providers import rms_norm + + original = RMSNorm.forward_cuda + + def forward_cuda(self, x, residual=None): + if self.variance_size_override is not None: + raise RuntimeError("parity RMSNorm does not support variance_size_override") + if residual is not None: + added = x + residual + return rms_norm(added, self.weight.data, self.variance_epsilon), added + return rms_norm(x, self.weight.data, self.variance_epsilon) + + forward_cuda._unirl_parity_original = original + RMSNorm.forward_cuda = forward_cuda + _INSTALLED = True + return ( + symbol_result( + "vllm.model_executor.layers.layernorm.RMSNorm.forward_cuda", + forward_cuda, + before=original, + actual=RMSNorm.forward_cuda, + ), + ) + + +__all__ = ["install_norm_patch", "preflight_norm_patch"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/precision.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/precision.py new file mode 100644 index 000000000..cd732578d --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/precision.py @@ -0,0 +1,100 @@ +"""Process-global precision settings without enabling vLLM's global BI mode.""" + +from __future__ import annotations + +from ..compat import require_symbol, require_value +from ..registry import SymbolResult, value_result + +_INSTALLED = False + + +def preflight_precision_contract(*, strict: bool) -> None: + for symbol in ( + "matmul.fp32_precision", + "matmul.allow_bf16_reduced_precision_reduction", + "matmul.allow_fp16_reduced_precision_reduction", + ): + require_value( + "torch.backends.cuda", + symbol, + expected_type=str if symbol.endswith("fp32_precision") else bool, + ) + for symbol in ("conv.fp32_precision", "rnn.fp32_precision"): + require_value("torch.backends.cudnn", symbol, expected_type=str) + require_symbol( + "torch.backends.cuda", + "preferred_blas_library", + parameters=(("backend", "None"),), + origin="torch.backends.cuda.preferred_blas_library", + strict=strict, + ) + + +def install_precision_contract() -> tuple[SymbolResult, ...]: + global _INSTALLED + if _INSTALLED: + raise RuntimeError("precision parity contract installed twice") + import torch + + before = { + "matmul_fp32": torch.backends.cuda.matmul.fp32_precision, + "conv_fp32": torch.backends.cudnn.conv.fp32_precision, + "rnn_fp32": torch.backends.cudnn.rnn.fp32_precision, + "bf16_reduction": torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + "fp16_reduction": torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + "blas": torch.backends.cuda.preferred_blas_library(), + } + torch.backends.cuda.matmul.fp32_precision = "ieee" + torch.backends.cudnn.conv.fp32_precision = "ieee" + torch.backends.cudnn.rnn.fp32_precision = "ieee" + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + torch.backends.cuda.preferred_blas_library(backend="cublaslt") + _INSTALLED = True + return ( + value_result( + "torch.backends.cuda.matmul.fp32_precision", + "literal:'ieee'", + before=before["matmul_fp32"], + actual=torch.backends.cuda.matmul.fp32_precision, + verified=torch.backends.cuda.matmul.fp32_precision == "ieee", + ), + value_result( + "torch.backends.cudnn.conv.fp32_precision", + "literal:'ieee'", + before=before["conv_fp32"], + actual=torch.backends.cudnn.conv.fp32_precision, + verified=torch.backends.cudnn.conv.fp32_precision == "ieee", + ), + value_result( + "torch.backends.cudnn.rnn.fp32_precision", + "literal:'ieee'", + before=before["rnn_fp32"], + actual=torch.backends.cudnn.rnn.fp32_precision, + verified=torch.backends.cudnn.rnn.fp32_precision == "ieee", + ), + value_result( + "torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction", + "literal:false", + before=before["bf16_reduction"], + actual=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + verified=not torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + ), + value_result( + "torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction", + "literal:false", + before=before["fp16_reduction"], + actual=torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + verified=not torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + ), + value_result( + "torch.backends.cuda.preferred_blas_library", + "torch._C._BlasBackend.Cublaslt", + before=before["blas"], + actual=torch.backends.cuda.preferred_blas_library(), + verified=torch.backends.cuda.preferred_blas_library().name == "Cublaslt", + ), + ) + + +__all__ = ["install_precision_contract", "preflight_precision_contract"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/providers.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/providers.py new file mode 100644 index 000000000..aba737e06 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/providers.py @@ -0,0 +1,84 @@ +"""Public vLLM/Torch providers used by the reference profile.""" + +from __future__ import annotations + +import torch + +from ..compat import require_symbol + + +def preflight_providers(*, strict: bool) -> None: + module = "vllm.model_executor.layers.batch_invariant" + contracts = { + "linear_batch_invariant": ( + ("input", ""), + ("weight", ""), + ("bias", "None"), + ), + "rms_norm_batch_invariant": ( + ("input", ""), + ("weight", ""), + ("eps", "1e-06"), + ("residual", "None"), + ), + "log_softmax": (("input", ""), ("dim", "-1")), + "softmax_batch_invariant": ( + ("input", ""), + ("dim", ""), + ("dtype", "None"), + ), + "mean_batch_invariant": ( + ("input", ""), + ("dim", ""), + ("keepdim", "False"), + ("dtype", "None"), + ), + } + for name, parameters in contracts.items(): + require_symbol( + module, + name, + parameters=parameters, + origin=f"{module}.{name}", + strict=strict, + ) + + +def linear(input: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import linear_batch_invariant + + return linear_batch_invariant(input.contiguous(), weight.contiguous(), bias) + + +def rms_norm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import rms_norm_batch_invariant + + return rms_norm_batch_invariant(input.contiguous(), weight.contiguous(), eps) + + +def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import log_softmax as implementation + + return implementation(input, dim=dim) + + +def softmax(input: torch.Tensor, dim: int = -1, dtype=None) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import softmax_batch_invariant + + return softmax_batch_invariant(input, dim=dim, dtype=dtype) + + +def mean(input: torch.Tensor, dim, keepdim=False, dtype=None) -> torch.Tensor: + from vllm.model_executor.layers.batch_invariant import mean_batch_invariant + + return mean_batch_invariant(input, dim=dim, keepdim=keepdim, dtype=dtype) + + +__all__ = [ + "linear", + "log_softmax", + "mean", + "preflight_providers", + "rms_norm", + "softmax", +] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/reductions.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/reductions.py new file mode 100644 index 000000000..478f220ab --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/common/reductions.py @@ -0,0 +1,117 @@ +"""Install vLLM public batch-invariant reductions as CUDA ATen providers.""" + +from __future__ import annotations + +import inspect + +from ..compat import require_symbol +from ..registry import SymbolResult, symbol_result + +_LIBRARY = None +_NATIVE_CUDA_OPS = ( + "aten::_log_softmax", + "aten::softmax", + "aten::_softmax", + "aten::mean.dim", +) + + +def _cuda_registration(torch, operator: str) -> str | None: + return next( + (line for line in torch._C._dispatch_dump_table(operator).splitlines() if line.startswith("CUDA:")), + None, + ) + + +def preflight_reduction_patch(*, strict: bool) -> None: + import torch + + require_symbol( + "vllm.v1.worker.gpu.sample.logprob", + "compute_token_logprobs", + parameters=(("logits", ""), ("token_ids", "")), + origin=("vllm.v1.worker.gpu.sample.logprob.compute_token_logprobs"), + strict=strict, + ) + if strict: + conflicts = { + operator: registration + for operator in _NATIVE_CUDA_OPS + if (registration := _cuda_registration(torch, operator)) is not None and "/pytorch/" not in registration + } + if conflicts: + raise RuntimeError(f"unexpected pre-existing CUDA reduction providers: {conflicts}") + + +def install_reduction_patch() -> tuple[SymbolResult, ...]: + global _LIBRARY + if _LIBRARY is not None: + raise RuntimeError("reduction parity patch installed twice") + import torch + + from .providers import log_softmax, mean, softmax + + before_registrations = {operator: _cuda_registration(torch, operator) for operator in _NATIVE_CUDA_OPS} + import vllm.v1.worker.gpu.sample.logprob as sample_logprob + + original_compute_token_logprobs = sample_logprob.compute_token_logprobs + + def aten_log_softmax(input, dim, half_to_float): + return log_softmax(input.float() if half_to_float else input, dim=dim) + + library = torch.library.Library("aten", "IMPL") + library.impl("aten::_log_softmax", aten_log_softmax, "CUDA") + library.impl("aten::softmax", softmax, "CUDA") + + def aten_softmax(input, dim, half_to_float): + return softmax(input.float() if half_to_float else input, dim=dim) + + library.impl("aten::_softmax", aten_softmax, "CUDA") + library.impl("aten::mean.dim", mean, "CUDA") + + def compute_token_logprobs(logits, token_ids): + return log_softmax(logits.float(), dim=-1).gather( + -1, + token_ids.to(torch.int64), + ) + + sample_logprob.compute_token_logprobs = compute_token_logprobs + _LIBRARY = library + registered = ( + ("aten::_log_softmax[CUDA]", aten_log_softmax, "aten::_log_softmax"), + ("aten::softmax[CUDA]", softmax, "aten::softmax"), + ("aten::_softmax[CUDA]", aten_softmax, "aten::_softmax"), + ("aten::mean.dim[CUDA]", mean, "aten::mean.dim"), + ) + results = [] + for symbol, provider, operator in registered: + registration = _cuda_registration(torch, operator) + results.append( + SymbolResult( + symbol=symbol, + before_provider=before_registrations[operator] or "", + before_signature=None, + before_identity=None, + after_provider=registration or "", + after_signature=str(inspect.signature(provider)), + after_identity=id(provider), + verified=( + registration is not None + and "unirl_train_inference_parity_vllm/common/reductions.py" in registration + ), + ) + ) + if not all(result.verified for result in results): + raise RuntimeError(f"failed to verify CUDA reduction providers: {results}") + results.append( + symbol_result( + "vllm.v1.worker.gpu.sample.logprob.compute_token_logprobs", + compute_token_logprobs, + before=original_compute_token_logprobs, + actual=sample_logprob.compute_token_logprobs, + ) + ) + return tuple(results) + + +__all__ = ["install_reduction_patch", "preflight_reduction_patch"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/compat.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/compat.py new file mode 100644 index 000000000..1d1f20898 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/compat.py @@ -0,0 +1,131 @@ +"""Pinned runtime and private-symbol contracts for the verified stack.""" + +from __future__ import annotations + +import importlib +import inspect +from dataclasses import asdict, dataclass +from importlib.metadata import PackageNotFoundError, version + +_VERSION_ALLOWLIST = { + "vllm": ("0.27.0",), + "torch": ("2.13.0",), + "transformers": ("5.6.0",), +} +_REQUIRED = "" + + +@dataclass(frozen=True) +class RuntimeVersion: + package: str + installed: str + allowed: tuple[str, ...] + + def to_manifest(self) -> dict[str, object]: + return asdict(self) + + +def validate_runtime_versions() -> tuple[RuntimeVersion, ...]: + results = [] + failures = [] + for package, allowed in _VERSION_ALLOWLIST.items(): + try: + installed = version(package) + except PackageNotFoundError: + failures.append(f"{package}=; allowed={list(allowed)}") + continue + public_version = installed.split("+", 1)[0] + if public_version not in allowed: + failures.append(f"{package}={installed}; allowed={list(allowed)}") + results.append(RuntimeVersion(package=package, installed=installed, allowed=allowed)) + if failures: + raise RuntimeError("unsupported UniRL parity runtime: " + "; ".join(failures)) + return tuple(results) + + +def _signature_shape( + value: object, +) -> tuple[tuple[str, str, str], ...]: + parameters = inspect.signature(value).parameters.values() + return tuple( + ( + parameter.name, + parameter.kind.name, + _REQUIRED if parameter.default is inspect.Parameter.empty else repr(parameter.default), + ) + for parameter in parameters + ) + + +def require_symbol( + module_name: str, + symbol_path: str, + *, + parameters: tuple[tuple[str, str] | tuple[str, str, str], ...] | None = None, + origin: str | None = None, + strict: bool, +) -> object: + """Resolve a required symbol and optionally enforce its exact signature.""" + try: + value: object = importlib.import_module(module_name) + except ImportError as error: + raise RuntimeError(f"required parity module {module_name!r} could not be imported") from error + for component in symbol_path.split("."): + try: + value = getattr(value, component) + except AttributeError as error: + raise RuntimeError(f"required parity symbol {module_name}.{symbol_path} is missing") from error + + qualified_name = ".".join( + part + for part in ( + getattr(value, "__module__", None), + getattr(value, "__qualname__", None), + ) + if part + ) + if origin is not None and qualified_name != origin: + raise RuntimeError( + f"conflicting provider at {module_name}.{symbol_path}: " + f"expected {origin}, got {qualified_name or type(value).__name__}" + ) + + if parameters is not None: + try: + actual = _signature_shape(value) + except (TypeError, ValueError) as error: + raise RuntimeError(f"cannot inspect required parity symbol {module_name}.{symbol_path}") from error + expected = tuple( + ( + parameter[0], + (inspect.Parameter.POSITIONAL_OR_KEYWORD.name if len(parameter) == 2 else parameter[1]), + parameter[-1], + ) + for parameter in parameters + ) + if strict and actual != expected: + raise RuntimeError(f"signature drift at {module_name}.{symbol_path}: expected {expected}, got {actual}") + return value + + +def require_value( + module_name: str, + symbol_path: str, + *, + expected_type: type, +) -> object: + value = require_symbol(module_name, symbol_path, strict=True) + if not isinstance(value, expected_type): + raise RuntimeError( + f"invalid parity symbol {module_name}.{symbol_path}: " + f"expected {expected_type.__name__}, got {type(value).__name__}" + ) + return value + + +__all__ = [ + "RuntimeVersion", + "require_symbol", + "require_value", + "validate_runtime_versions", +] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/config.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/config.py new file mode 100644 index 000000000..eb20b4c4c --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/config.py @@ -0,0 +1,60 @@ +"""Environment-only configuration available in every spawned vLLM process.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +PUBLIC_REFERENCE = "public_reference" +QWEN3_MOE_MODEL = "qwen3_moe_30b_a3b" +REQUIRED_PATCHES = ("common", QWEN3_MOE_MODEL) + + +@dataclass(frozen=True) +class PluginConfig: + profile: str + model: str + patches: tuple[str, ...] + strict: bool + + +def _required_environment(name: str) -> str: + try: + value = os.environ[name].strip() + except KeyError as error: + raise RuntimeError(f"{name} is required when UNIRL_PARITY_ENABLE=1") from error + if not value: + raise RuntimeError(f"{name} must be non-empty when UNIRL_PARITY_ENABLE=1") + return value + + +def load_config() -> PluginConfig: + profile = _required_environment("UNIRL_PARITY_PROFILE").lower() + if profile != PUBLIC_REFERENCE: + raise ValueError(f"unknown UNIRL_PARITY_PROFILE={profile!r}") + model = _required_environment("UNIRL_PARITY_MODEL").lower() + if model != QWEN3_MOE_MODEL: + raise ValueError(f"unsupported UNIRL_PARITY_MODEL={model!r}") + raw_patches = _required_environment("UNIRL_PARITY_PATCHES") + patch_values = tuple(value.strip() for value in raw_patches.split(",")) + if any(not value for value in patch_values) or len(patch_values) != len(set(patch_values)): + raise ValueError("UNIRL_PARITY_PATCHES must contain unique, non-empty patch names") + if patch_values != REQUIRED_PATCHES: + raise ValueError( + f"UNIRL_PARITY_PATCHES must enable the complete ordered contract {REQUIRED_PATCHES!r}; got {patch_values!r}" + ) + patches = patch_values + raw_strict = _required_environment("UNIRL_PARITY_STRICT") + if raw_strict != "1": + raise ValueError("UNIRL_PARITY_STRICT must be exactly '1' for the public-reference contract") + strict = True + return PluginConfig(profile=profile, model=model, patches=patches, strict=strict) + + +__all__ = [ + "PUBLIC_REFERENCE", + "QWEN3_MOE_MODEL", + "REQUIRED_PATCHES", + "PluginConfig", + "load_config", +] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/__init__.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/__init__.py new file mode 100644 index 000000000..455439daf --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/__init__.py @@ -0,0 +1 @@ +"""Model-specific vLLM parity patches.""" diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/__init__.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/__init__.py new file mode 100644 index 000000000..8eeb19b76 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/__init__.py @@ -0,0 +1,5 @@ +"""Qwen3-30B-A3B vLLM model patch.""" + +from .patch import install_qwen3_moe_patch, preflight_qwen3_moe_patch + +__all__ = ["install_qwen3_moe_patch", "preflight_qwen3_moe_patch"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/dense.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/dense.py new file mode 100644 index 000000000..546dcc0f3 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/dense.py @@ -0,0 +1,125 @@ +"""Public-reference QKV, o_proj and LM-head providers.""" + +from __future__ import annotations + +import types + +import torch +import torch.nn as nn + +from ...common.providers import linear + +_PARITY_O_PROJECTION_CLASSES: dict[type, type] = {} + + +def _use_direct_layerwise_reload(layer) -> None: + """Keep the custom o-proj storage live and load it without meta staging. + + The parity projection replaces vLLM's row-parallel parameter with a + column-parallel one of equal element count but different shape. vLLM's + generic meta-staged reload cannot infer that custom layout and copied the + local q-projection shard into this storage. Skipping meta staging lets the + bound loader write the canonical full o-proj directly into its CuMem + weights-pool allocation. + """ + from vllm.model_executor.model_loader.reload.meta import SKIP_MODULES + + base_class = type(layer) + parity_class = _PARITY_O_PROJECTION_CLASSES.get(base_class) + if parity_class is None: + parity_class = type("UniRLParityOProjection", (base_class,), {}) + _PARITY_O_PROJECTION_CLASSES[base_class] = parity_class + layer.__class__ = parity_class + SKIP_MODULES.add(parity_class.__name__) + + +def _column_weight_loader(self, parameter, loaded_weight): + rank = int(self._unirl_parity_tp_rank) + world = int(self._unirl_parity_tp_world) + rows = int(loaded_weight.shape[0]) + shard = rows // world + loaded_shard = loaded_weight.narrow(0, rank * shard, shard).to( + device=parameter.device, + dtype=parameter.dtype, + ) + parameter.data.copy_(loaded_shard) + + +def _reshape_row_to_column(layer) -> None: + from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + rank = get_tensor_model_parallel_rank() + world = get_tensor_model_parallel_world_size() + _use_direct_layerwise_reload(layer) + old = layer.weight + replacement = nn.Parameter( + torch.empty( + layer.output_size // world, + layer.input_size, + device=old.device, + dtype=old.dtype, + ), + requires_grad=False, + ) + from vllm.model_executor.utils import set_weight_attrs + + layer._unirl_parity_tp_rank = int(rank) + layer._unirl_parity_tp_world = int(world) + loader = types.MethodType(_column_weight_loader, layer) + set_weight_attrs(replacement, {"weight_loader": loader}) + layer.weight = replacement + layer._unirl_parity_output_size = layer.output_size + + +def _qkv_forward(self, hidden): + return linear(hidden, self.weight, None), None + + +def _o_forward(self, hidden): + from vllm.distributed import get_tp_group + + group = get_tp_group() + full_input = group.all_gather(hidden.contiguous(), dim=-1) + local_output = linear(full_input, self.weight, None) + output = local_output if group.world_size == 1 else group.all_gather(local_output.contiguous(), dim=-1) + return output, None + + +def make_attention(base_class): + class ParityAttention(base_class): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.qkv_proj.forward = types.MethodType( + _qkv_forward, + self.qkv_proj, + ) + _reshape_row_to_column(self.o_proj) + self.o_proj.forward = types.MethodType( + _o_forward, + self.o_proj, + ) + + ParityAttention.__name__ = "ParityQwen3MoeAttention" + return ParityAttention + + +def make_logits_processor(base_class): + class ParityLogitsProcessor(base_class): + def _get_logits(self, hidden_states, lm_head, embedding_bias=None): + if embedding_bias is not None: + raise ValueError("Qwen3-MoE parity LM head does not support bias") + from vllm.distributed import get_tp_group + + group = get_tp_group() + local_logits = linear(hidden_states, lm_head.weight, None) + logits = local_logits if group.world_size == 1 else group.all_gather(local_logits.contiguous(), dim=-1) + return logits[..., : self.org_vocab_size] + + ParityLogitsProcessor.__name__ = "ParityQwen3MoeLogitsProcessor" + return ParityLogitsProcessor + + +__all__ = ["make_attention", "make_logits_processor"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/moe.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/moe.py new file mode 100644 index 000000000..74bd55f02 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/moe.py @@ -0,0 +1,252 @@ +"""Correctness-first Qwen3-MoE path using vLLM BI linear and Torch combine.""" + +from __future__ import annotations + +import types + +import torch +import torch.nn.functional as F + +from ...common.moe_combine import moe_combine +from ...common.providers import linear +from .router import install_gate, install_hf_router + + +def _routed_experts(experts): + return getattr(experts, "routed_experts", experts) + + +def _validate_tp_routing( + counts: torch.Tensor, + topk_ids: torch.Tensor, + *, + num_experts: int, + group, + world: int, +) -> None: + """Collectively fail before route-dependent expert collectives diverge.""" + shape = tuple(topk_ids.shape) + counts_valid = bool( + counts.dim() == 1 + and counts.numel() == num_experts + and torch.all(counts >= 0).item() + and int(counts.sum().item()) == int(topk_ids.numel()) + ) + ids_valid = bool( + topk_ids.dim() == 2 + and (topk_ids.numel() == 0 or (int(topk_ids.min().item()) >= 0 and int(topk_ids.max().item()) < num_experts)) + ) + if world == 1: + if not counts_valid or not ids_valid: + raise RuntimeError("Qwen3-MoE parity routing metadata is invalid") + return + metadata = torch.tensor( + ( + topk_ids.dim(), + topk_ids.numel(), + shape[0] if len(shape) > 0 else -1, + shape[1] if len(shape) > 1 else -1, + counts.numel(), + num_experts, + int(counts_valid), + int(ids_valid), + ), + dtype=torch.int64, + device=topk_ids.device, + ) + gathered_metadata = group.all_gather(metadata.contiguous(), dim=0).view( + world, + metadata.numel(), + ) + reference_metadata = gathered_metadata[0] + metadata_matches = bool(torch.all(gathered_metadata == reference_metadata).item()) + metadata_valid = bool( + torch.all(gathered_metadata[:, 0] == 2).item() + and torch.all(gathered_metadata[:, 1] == gathered_metadata[:, 2] * gathered_metadata[:, 3]).item() + and torch.all(gathered_metadata[:, 4] == num_experts).item() + and torch.all(gathered_metadata[:, 5] == num_experts).item() + and torch.all(gathered_metadata[:, 6] == 1).item() + and torch.all(gathered_metadata[:, 7] == 1).item() + ) + if not metadata_matches or not metadata_valid: + observed = gathered_metadata.detach().cpu().tolist() + raise RuntimeError(f"Qwen3-MoE parity TP routing metadata mismatch before expert collectives: {observed}") + + payload = torch.cat( + ( + counts.detach().to(device=topk_ids.device, dtype=torch.int64), + topk_ids.detach().reshape(-1).to(dtype=torch.int64), + ) + ).contiguous() + gathered_payload = group.all_gather(payload, dim=0).view(world, payload.numel()) + mismatch_ranks = ( + torch.any(gathered_payload != gathered_payload[0], dim=1).nonzero().flatten().detach().cpu().tolist() + ) + if mismatch_ranks: + raise RuntimeError( + "Qwen3-MoE parity TP routing mismatch before expert collectives; " + f"counts/topk_ids differ on ranks {mismatch_ranks}" + ) + + +def _get_w2_column(experts) -> torch.Tensor: + cached = getattr(experts, "_unirl_parity_w2_column", None) + if cached is not None: + return cached + from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, + ) + + world = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + stock = _routed_experts(experts).w2_weight + full = stock if world == 1 else get_tp_group().all_gather(stock.contiguous(), dim=-1) + hidden_local = int(full.shape[1]) // world + cached = full[:, rank * hidden_local : (rank + 1) * hidden_local].contiguous() + experts._unirl_parity_w2_column = cached + return cached + + +def _install_reload_invalidation(experts) -> None: + original = experts.load_weights + if getattr(original, "_unirl_parity_wrapper", False): + return + + def load_weights(self, weights): + self._unirl_parity_w2_column = None + return original(weights) + + load_weights._unirl_parity_wrapper = True + experts.load_weights = types.MethodType(load_weights, experts) + # Do not wrap w13/w2 Parameter loaders. vLLM's layerwise WTE reload + # introspects and replays those exact loaders; a variadic invalidation + # wrapper caused all 96 expert tensors to be reconstructed incorrectly. + # The module hook above and the worker pre-sleep hook cover cache lifetime. + + +def _expert_forward(block, hidden_states: torch.Tensor) -> torch.Tensor: + from vllm.distributed import ( + get_tensor_model_parallel_world_size, + get_tp_group, + ) + from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( + moe_permute, + ) + + input_was_1d = hidden_states.dim() == 1 + hidden = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + router_logits, _ = block.gate(hidden) + topk_weights, topk_ids = block.experts.router.select_experts( + hidden_states=hidden, + router_logits=router_logits, + ) + num_experts = int(block.n_routed_experts) + permuted, _scale, offsets, inverse, _indices = moe_permute( + hidden, + None, + topk_ids.to(torch.int32), + num_experts, + ) + counts_tensor = offsets[1:] - offsets[:-1] + world = get_tensor_model_parallel_world_size() + group = get_tp_group() + _validate_tp_routing( + counts_tensor, + topk_ids, + num_experts=num_experts, + group=group, + world=world, + ) + counts = counts_tensor.detach().cpu().tolist() + w2_column = _get_w2_column(block.experts) + routed_experts = _routed_experts(block.experts) + slots = int(topk_ids.numel()) + if slots == 0: + return torch.zeros_like(hidden) + max_count = max(int(value) for value in counts) + local_gate_width = int(routed_experts.w13_weight.shape[1]) + local_intermediate = local_gate_width // 2 + padded_gate_up = [] + offset = 0 + for expert, count_value in enumerate(counts): + count = int(count_value) + rows = permuted[offset : offset + count].contiguous() + offset += count + if count == 0: + local_gate_up = hidden.new_zeros((max_count, local_gate_width)) + else: + active = linear( + rows, + routed_experts.w13_weight[expert], + None, + ) + padding = hidden.new_zeros((max_count - count, local_gate_width)) + local_gate_up = torch.cat((active, padding), dim=0) + padded_gate_up.append(local_gate_up) + + local_gate_up = torch.stack(padded_gate_up, dim=0).contiguous() + gathered_gate_up = local_gate_up if world == 1 else group.all_gather(local_gate_up, dim=-1) + rank_packed = gathered_gate_up.view( + num_experts, + max_count, + world, + 2, + local_intermediate, + ) + gate = rank_packed[:, :, :, 0].reshape(num_experts, max_count, -1) + up = rank_packed[:, :, :, 1].reshape(num_experts, max_count, -1) + activation = F.silu(gate) * up + + local_hidden_width = int(w2_column.shape[1]) + padded_down = [] + for expert, count_value in enumerate(counts): + if int(count_value) == 0: + local_down = hidden.new_zeros((max_count, local_hidden_width)) + else: + local_down = linear( + activation[expert].contiguous(), + w2_column[expert], + None, + ) + padded_down.append(local_down) + local_down = torch.stack(padded_down, dim=0).contiguous() + gathered_down = local_down if world == 1 else group.all_gather(local_down, dim=-1) + expert_output = torch.cat( + [gathered_down[expert, : int(count)] for expert, count in enumerate(counts) if int(count) > 0], + dim=0, + ) + permuted_weights = torch.zeros( + slots, + dtype=topk_weights.dtype, + device=topk_weights.device, + ) + permuted_weights[inverse.long()] = topk_weights.reshape(-1) + contributions = (expert_output * permuted_weights[:, None]).to(torch.bfloat16) + result = moe_combine( + contributions, + inverse.view_as(topk_ids).to(torch.int32), + ) + return result.squeeze(0) if input_was_1d else result + + +def make_moe_block(base_class): + class ParityMoeBlock(base_class): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + install_gate(self.gate) + install_hf_router(self.experts) + _install_reload_invalidation(self.experts) + + def forward(self, hidden_states): + return _expert_forward(self, hidden_states) + + def _unirl_before_sleep(self): + self.experts._unirl_parity_w2_column = None + + ParityMoeBlock.__name__ = "ParityQwen3MoeSparseMoeBlock" + return ParityMoeBlock + + +__all__ = ["make_moe_block"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/patch.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/patch.py new file mode 100644 index 000000000..243919037 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/patch.py @@ -0,0 +1,304 @@ +"""Install the public-reference Qwen3-MoE vLLM patch.""" + +from __future__ import annotations + +import importlib + +import torch + +from ...common.providers import preflight_providers +from ...compat import require_symbol +from ...registry import PatchResult, SymbolResult, symbol_result, value_result +from .dense import make_attention, make_logits_processor +from .moe import make_moe_block + +_INSTALLED = False + + +def _install_rope() -> tuple[SymbolResult, ...]: + import vllm.model_executor.layers.rotary_embedding.base as rope + + original_inverse_frequency = rope.RotaryEmbeddingBase._compute_inv_freq + original_forward_cuda = rope.RotaryEmbedding.forward_cuda + + def inverse_frequency(self, base): + exponent = torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) + return 1.0 / (base ** (exponent / self.rotary_dim)) + + def apply(self, positions, query, key=None): + positions = positions.flatten() + cache = self._match_cos_sin_cache_dtype(query).index_select(0, positions) + cosine, sine = cache.chunk(2, dim=-1) + tokens = int(positions.shape[0]) + + def rotate(value): + original_shape = value.shape + value = value.view(tokens, -1, self.head_size) + first, second = torch.chunk(value[..., : self.rotary_dim], 2, dim=-1) + cos = cosine.unsqueeze(-2).to(value.dtype) + sin = sine.unsqueeze(-2).to(value.dtype) + rotated = torch.cat( + (first * cos - second * sin, second * cos + first * sin), + dim=-1, + ) + return torch.cat( + (rotated, value[..., self.rotary_dim :]), + dim=-1, + ).reshape(original_shape) + + return rotate(query), None if key is None else rotate(key) + + rope.RotaryEmbeddingBase._compute_inv_freq = inverse_frequency + rope.RotaryEmbedding.forward_cuda = apply + return ( + symbol_result( + ("vllm.model_executor.layers.rotary_embedding.base.RotaryEmbeddingBase._compute_inv_freq"), + inverse_frequency, + before=original_inverse_frequency, + actual=rope.RotaryEmbeddingBase._compute_inv_freq, + ), + symbol_result( + ("vllm.model_executor.layers.rotary_embedding.base.RotaryEmbedding.forward_cuda"), + apply, + before=original_forward_cuda, + actual=rope.RotaryEmbedding.forward_cuda, + ), + ) + + +def _install_sleep_cache_invalidation() -> SymbolResult: + """Drop parity-owned CUDA caches before vLLM releases their allocations.""" + import vllm.v1.worker.gpu_worker as gpu_worker + + original_sleep = gpu_worker.Worker.sleep + + def sleep(self, level: int = 1) -> None: + model = getattr(getattr(self, "model_runner", None), "model", None) + if model is not None: + for module in model.modules(): + before_sleep = getattr(module, "_unirl_before_sleep", None) + if callable(before_sleep): + before_sleep() + original_sleep(self, level=level) + + gpu_worker.Worker.sleep = sleep + return symbol_result( + "vllm.v1.worker.gpu_worker.Worker.sleep", + sleep, + before=original_sleep, + actual=gpu_worker.Worker.sleep, + ) + + +def preflight_qwen3_moe_patch(strict: bool) -> None: + preflight_providers(strict=strict) + contracts = ( + ( + "vllm.model_executor.layers.rotary_embedding.base", + "RotaryEmbeddingBase._compute_inv_freq", + (("self", ""), ("base", "")), + ("vllm.model_executor.layers.rotary_embedding.base.RotaryEmbeddingBase._compute_inv_freq"), + ), + ( + "vllm.model_executor.layers.rotary_embedding.base", + "RotaryEmbedding.forward_cuda", + ( + ("self", ""), + ("positions", ""), + ("query", ""), + ("key", "None"), + ), + ("vllm.model_executor.layers.rotary_embedding.base.RotaryEmbedding.forward_cuda"), + ), + ( + "vllm.model_executor.models.qwen3_moe", + "Qwen3MoeSparseMoeBlock", + (("vllm_config", ""), ("prefix", "''")), + "vllm.model_executor.models.qwen3_moe.Qwen3MoeSparseMoeBlock", + ), + ( + "vllm.model_executor.models.qwen3_moe", + "Qwen3MoeSparseMoeBlock.forward", + (("self", ""), ("hidden_states", "")), + ("vllm.model_executor.models.qwen3_moe.Qwen3MoeSparseMoeBlock.forward"), + ), + ( + "vllm.model_executor.models.qwen3_moe", + "Qwen3MoeAttention", + ( + ("hidden_size", ""), + ("num_heads", ""), + ("num_kv_heads", ""), + ("rope_parameters", ""), + ("max_position_embeddings", "8192"), + ("head_dim", "None"), + ("rms_norm_eps", "1e-06"), + ("qkv_bias", "False"), + ("cache_config", "None"), + ("quant_config", "None"), + ("prefix", "''"), + ("dual_chunk_attention_config", "None"), + ), + "vllm.model_executor.models.qwen3_moe.Qwen3MoeAttention", + ), + ( + "vllm.model_executor.models.qwen3_moe", + "Qwen3MoeAttention.forward", + ( + ("self", ""), + ("positions", ""), + ("hidden_states", ""), + ), + ("vllm.model_executor.models.qwen3_moe.Qwen3MoeAttention.forward"), + ), + ( + "vllm.model_executor.models.qwen3_moe", + "LogitsProcessor", + ( + ("vocab_size", ""), + ("org_vocab_size", "None"), + ("scale", "1.0"), + ("logits_as_input", "False"), + ("soft_cap", "None"), + ), + "vllm.model_executor.layers.logits_processor.LogitsProcessor", + ), + ( + "vllm.model_executor.models.qwen3_moe", + "LogitsProcessor._get_logits", + ( + ("self", ""), + ("hidden_states", ""), + ("lm_head", ""), + ("embedding_bias", ""), + ), + ("vllm.model_executor.layers.logits_processor.LogitsProcessor._get_logits"), + ), + ( + "vllm.model_executor.layers.fused_moe.moe_permute_unpermute", + "moe_permute", + ( + ("hidden_states", ""), + ("a1q_scale", ""), + ("topk_ids", ""), + ("n_expert", ""), + ("n_local_expert", "-1"), + ("expert_map", "None"), + ("permuted_hidden_states", "None"), + ("scratch", "None"), + ), + ("vllm.model_executor.layers.fused_moe.moe_permute_unpermute.moe_permute"), + ), + ( + "vllm.distributed.parallel_state", + "GroupCoordinator.all_gather", + ( + ("self", ""), + ("input_", ""), + ("dim", "-1"), + ), + "vllm.distributed.parallel_state.GroupCoordinator.all_gather", + ), + ( + "vllm.model_executor.layers.linear", + "ReplicatedLinear.forward", + (("self", ""), ("x", "")), + "vllm.model_executor.layers.linear.ReplicatedLinear.forward", + ), + ( + "vllm.model_executor.layers.fused_moe.router.base_router", + "BaseRouter.select_experts", + ( + ("self", ""), + ("hidden_states", ""), + ("router_logits", ""), + ("topk_indices_dtype", "None"), + ("input_ids", "KEYWORD_ONLY", "None"), + ), + "vllm.model_executor.layers.fused_moe.router.fused_moe_router.FusedMoERouter.select_experts", + ), + ( + "vllm.model_executor.layers.fused_moe.router.fused_topk_router", + "FusedTopKRouter._compute_routing", + ( + ("self", ""), + ("hidden_states", ""), + ("router_logits", ""), + ("indices_type", ""), + ("input_ids", "KEYWORD_ONLY", "None"), + ), + ("vllm.model_executor.layers.fused_moe.router.fused_topk_router.FusedTopKRouter._compute_routing"), + ), + ( + "vllm.v1.worker.gpu_worker", + "Worker.sleep", + (("self", ""), ("level", "1")), + "vllm.v1.worker.gpu_worker.Worker.sleep", + ), + ) + for module, symbol, parameters, origin in contracts: + require_symbol( + module, + symbol, + parameters=parameters, + origin=origin, + strict=strict, + ) + + +def install_qwen3_moe_patch(*, strict: bool) -> PatchResult: + if not strict: + raise ValueError("the Qwen3-MoE public-reference installer requires strict=True") + global _INSTALLED + if _INSTALLED: + raise RuntimeError("Qwen3-MoE parity patch installed twice") + module = importlib.import_module("vllm.model_executor.models.qwen3_moe") + original_moe_block = module.Qwen3MoeSparseMoeBlock + original_attention = module.Qwen3MoeAttention + original_logits_processor = module.LogitsProcessor + original_marker = getattr(module, "_unirl_train_inference_parity", "") + moe_block = make_moe_block(original_moe_block) + attention = make_attention(original_attention) + logits_processor = make_logits_processor(original_logits_processor) + rope_results = _install_rope() + sleep_result = _install_sleep_cache_invalidation() + module.Qwen3MoeSparseMoeBlock = moe_block + module.Qwen3MoeAttention = attention + module.LogitsProcessor = logits_processor + module._unirl_train_inference_parity = True + _INSTALLED = True + return PatchResult( + name="qwen3_moe_30b_a3b", + symbols=( + *rope_results, + sleep_result, + symbol_result( + ("vllm.model_executor.models.qwen3_moe.Qwen3MoeSparseMoeBlock"), + moe_block, + before=original_moe_block, + actual=module.Qwen3MoeSparseMoeBlock, + ), + symbol_result( + "vllm.model_executor.models.qwen3_moe.Qwen3MoeAttention", + attention, + before=original_attention, + actual=module.Qwen3MoeAttention, + ), + symbol_result( + "vllm.model_executor.models.qwen3_moe.LogitsProcessor", + logits_processor, + before=original_logits_processor, + actual=module.LogitsProcessor, + ), + value_result( + ("vllm.model_executor.models.qwen3_moe._unirl_train_inference_parity"), + "literal:true", + before=original_marker, + actual=module._unirl_train_inference_parity, + verified=module._unirl_train_inference_parity is True, + ), + ), + ) + + +__all__ = ["install_qwen3_moe_patch", "preflight_qwen3_moe_patch"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/router.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/router.py new file mode 100644 index 000000000..521ebb8dc --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/models/qwen3_moe_30b_a3b/router.py @@ -0,0 +1,45 @@ +"""HF-ordered router implemented with public vLLM BI providers.""" + +from __future__ import annotations + +import types + +import torch + +from ...common.providers import linear, softmax + + +def _gate_forward(self, hidden_states): + if self.bias is not None: + raise ValueError("Qwen3-MoE parity router gate requires bias=False") + hidden = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + return linear(hidden, self.weight, None), None + + +def install_gate(gate) -> None: + gate.forward = types.MethodType(_gate_forward, gate) + + +def _compute_routing( + self, + hidden_states, + router_logits, + indices_type, + **kwargs, +): + del hidden_states, kwargs + probabilities = softmax(router_logits.float().contiguous(), dim=-1) + weights, indices = torch.topk(probabilities, self.top_k, dim=-1) + if self.renormalize: + weights = weights / weights.sum(dim=-1, keepdim=True) + return weights.float(), indices.to(indices_type or torch.int32) + + +def install_hf_router(experts) -> None: + router = experts.router + if router.scoring_func != "softmax" or router.renormalize is not True: + raise ValueError("Qwen3-MoE parity requires softmax routing with renormalize=true") + router._compute_routing = types.MethodType(_compute_routing, router) + + +__all__ = ["install_gate", "install_hf_router"] diff --git a/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/registry.py b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/registry.py new file mode 100644 index 000000000..2af2aa257 --- /dev/null +++ b/experimental/train_inference_parity/vllm_plugin/src/unirl_train_inference_parity_vllm/registry.py @@ -0,0 +1,142 @@ +"""Fail-closed, two-phase registry for parity plugin installers.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from dataclasses import asdict, dataclass + + +@dataclass(frozen=True) +class SymbolResult: + """Before/after evidence for one concrete patched symbol.""" + + symbol: str + before_provider: str + before_signature: str | None + before_identity: int | None + after_provider: str + after_signature: str | None + after_identity: int | None + verified: bool + + +@dataclass(frozen=True) +class PatchResult: + """Structured result returned by every patch installer.""" + + name: str + symbols: tuple[SymbolResult, ...] + + def to_manifest(self) -> dict[str, object]: + return asdict(self) + + +@dataclass(frozen=True) +class Installer: + preflight: Callable[[bool], None] + install: Callable[[bool], PatchResult] + + +_INSTALLED: dict[str, PatchResult] = {} + + +def _provider_name(value: object) -> str: + module = getattr(value, "__module__", None) + qualname = getattr(value, "__qualname__", None) + if module and qualname: + return f"{module}.{qualname}" + if isinstance(value, str): + return value + return repr(value) + + +def _signature(value: object) -> str | None: + try: + return str(inspect.signature(value)) + except (TypeError, ValueError): + return None + + +def symbol_result( + symbol: str, + provider: object, + *, + before: object, + actual: object, +) -> SymbolResult: + """Describe and verify the provider currently bound to ``symbol``.""" + verified = actual is provider + if not verified: + raise RuntimeError( + f"parity provider verification failed for {symbol}: expected identity {provider!r}, got {actual!r}" + ) + return SymbolResult( + symbol=symbol, + before_provider=_provider_name(before), + before_signature=_signature(before), + before_identity=id(before), + after_provider=_provider_name(provider), + after_signature=_signature(provider), + after_identity=id(actual), + verified=True, + ) + + +def value_result( + symbol: str, + provider: str, + *, + before: object, + actual: object, + verified: bool, +) -> SymbolResult: + if not verified: + raise RuntimeError(f"parity value verification failed for {symbol}") + return SymbolResult( + symbol=symbol, + before_provider=_provider_name(before), + before_signature=None, + before_identity=id(before), + after_provider=provider, + after_signature=None, + after_identity=id(actual), + verified=True, + ) + + +def install_selected( + names: tuple[str, ...], + *, + strict: bool, + installers: Mapping[str, Installer], +) -> tuple[PatchResult, ...]: + unknown = sorted(set(names) - installers.keys()) + if unknown: + raise ValueError(f"unknown parity patches {unknown}; known={sorted(installers)}") + + pending = tuple(name for name in names if name not in _INSTALLED) + # Complete every compatibility check before mutating process-global state. + for name in pending: + installers[name].preflight(strict=strict) + for name in pending: + result = installers[name].install(strict=strict) + if result.name != name: + raise RuntimeError(f"parity installer {name!r} returned result for {result.name!r}") + _INSTALLED[name] = result + return tuple(_INSTALLED[name] for name in names) + + +def installed_manifest() -> tuple[PatchResult, ...]: + return tuple(_INSTALLED[name] for name in sorted(_INSTALLED)) + + +__all__ = [ + "Installer", + "PatchResult", + "SymbolResult", + "install_selected", + "installed_manifest", + "symbol_result", + "value_result", +] 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..583327e9a 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,45 @@ 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") + _EXACT_CONTEXT_FACTORY = context_factory + _EXACT_LOG_SOFTMAX = log_softmax + + def _replay_aware_forward( self: Any, *, @@ -80,8 +121,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 +133,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 +159,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 +287,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 +301,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 +401,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 +638,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..7e337a79e 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 @@ -74,9 +83,20 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": # 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/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(