diff --git a/docs/low_precision/mxfp8.md b/docs/low_precision/mxfp8.md index d074003e484..860def5bbc7 100644 --- a/docs/low_precision/mxfp8.md +++ b/docs/low_precision/mxfp8.md @@ -52,13 +52,17 @@ verl defaults the MXFP8 dense and MoE kernels of vLLM and SGLang to DeepGEMM; fo ## Checks +![The quantized-rollout sentinel, the post-sync self-check and the quantized-layer audit](https://github.com/user-attachments/assets/031f882a-9379-4ef7-bbce-3a746959d317) + | Check | Runs | Reports | Switch | |---|---|---|---| | Quantized-rollout sentinel | every training step | NaN log-probs, KL above 1.0, every response at the length cap | `VERL_QUANT_SENTINEL` | | Post-sync self-check | after every weight sync | kernel scale layouts that did not follow the sync | `VERL_MXFP8_REFIT_CHECK` | +| Quantized-layer audit | training vs export: first weight sync after a training step; export vs engine: first weight sync, before each bucket is written | layers that training, the exported weights and the engine quantize differently | `VERL_QUANT_LAYER_AUDIT` | Each switch takes `warn` (the default: log the finding and keep going), `raise` (stop the run or the sync) -or `0` (off). The sentinel applies to every quantized rollout, FP8 and FP4 included. +or `0` (off). The sentinel applies to every quantized rollout, FP8 and FP4 included, and so does the audit's +export-vs-engine comparison (on SGLang, to the formats of its refit loader: MXFP8, MXFP4, NVFP4). ### Quantized-rollout sentinel @@ -99,6 +103,31 @@ What it does not show: updated rests on indirect evidence (the re-derived layouts are copied back in place, and the KL stays in its healthy range), not on a graph replay. +### Quantized-layer audit + +Three parties decide which layers are quantized, each on its own: training implicitly (TE linear modules +inside `fp8_autocast`), and by name the trainer-side exporter and the engine, both from the same quantization +config (`rollout.quantization`, or the checkpoint's own `quantization_config`). The audit compares them in two +steps, each once: + +- **Training vs export**, when training runs a TE FP8 recipe, at the first weight sync after a training step + (the sync before it has no FP8 trace yet): each decoder layer's FP8 trace (TE's FP8 weight workspaces) + against the format the sync sends for that layer's GEMM projections (attention, MLP, experts, + linear-attention and SSM in / out projections). Weights the exporter keeps in BF16 by design are skipped. + Other decoder-layer weights, such as a convolution, are not judged, and the audit reports how many it left + out. +- **Export vs engine**, whenever the rollout is quantized, during the first weight sync: inside the engine's + weight loader, before each bucket is written, every incoming weight's format (quantized or high + precision) against the engine parameter that receives it. vLLM checks every quantized format; SGLang + checks the formats of its refit loader (MXFP8, MXFP4, NVFP4). Later syncs send the same formats and are + not checked again. + +`VERL_QUANT_LAYER_AUDIT` switches both. `VERL_QUANT_LAYER_AUDIT_TRACE=1` logs each step of the training-side +comparison. The training-side trace is per layer, +so a TE module kept in BF16 inside an FP8 layer would go unnoticed (Megatron has no such per-module option +today). With `disable_parameter_transpose_cache=True` TE keeps no FP8 workspace, and the audit warns once +that it cannot run. + ### Configurations SGLang refuses Two configurations are refused before a weight sync writes anything: diff --git a/docs/low_precision/qat.md b/docs/low_precision/qat.md index b7ca5fdcbd9..b3ac86b1d0e 100644 --- a/docs/low_precision/qat.md +++ b/docs/low_precision/qat.md @@ -247,6 +247,17 @@ SGLang uses the default CUDA Graph configuration): | Blackwell | SGLang | ❌ | NVFP4 MoE | FlashInfer CUTLASS (MoE; BF16 Linear) | 0.00946 | | Blackwell | SGLang | ❌ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS (MoE; BF16 Linear) | 0.00653 | +## Format check at the first weight sync + +When the rollout is quantized (QAT included), the engine compares, during the first weight sync and before each +bucket is written, every synced weight's format (quantized or high precision) with the engine parameter that +receives it. The trainer-side exporter and the engine read the same quantization config, so a disagreement means +the two matched a module differently, for example an `ignore` entry one side reads as a prefix and the other as an +exact name; loading would cast the weight into a parameter of the other format. vLLM checks every quantized +format; SGLang checks the formats of its refit loader (MXFP8, MXFP4, NVFP4), not `FP8_BLOCK`. +`VERL_QUANT_LAYER_AUDIT` switches it: `warn` (the default) logs the disagreement, `raise` stops the sync, `0` +turns it off. See also the [checks of the MXFP8 recipe](mxfp8.md#checks). + ## Limitations - SGLang MXFP8, MXFP4 and NVFP4 custom refit requires the default bucketed diff --git a/verl/utils/megatron_utils.py b/verl/utils/megatron_utils.py index e47ec921b86..ad8a762ed57 100644 --- a/verl/utils/megatron_utils.py +++ b/verl/utils/megatron_utils.py @@ -445,12 +445,20 @@ def _clear_te_fp8_weight_workspaces(model_chunk): rebuilds the workspace on the next forward, so dropping it here is safe. This is a no-op for bf16/fp16 models (the attribute is absent or empty). + The workspace is also the quantized-layer audit's evidence that a module ran + FP8 GEMMs (``verl.utils.quant_layer_audit``); each module is marked as having + held one before the cache is dropped, so the audit still works with param + offload on. + Returns the number of cached workspace entries cleared. """ + from verl.utils.quant_layer_audit import remember_fp8_workspaces + cleared = 0 for submodule in model_chunk.modules(): workspaces = getattr(submodule, "_fp8_workspaces", None) if isinstance(workspaces, dict) and workspaces: + remember_fp8_workspaces(submodule) cleared += len(workspaces) workspaces.clear() return cleared diff --git a/verl/utils/quant_layer_audit.py b/verl/utils/quant_layer_audit.py new file mode 100644 index 00000000000..888ad5a25cb --- /dev/null +++ b/verl/utils/quant_layer_audit.py @@ -0,0 +1,468 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Training vs export: did training run each decoder layer in the format the weight sync sends it in? + +This is the training half of the quantized-layer audit. The rollout half, export vs engine, runs inside each +rollout engine's weight loader (``verl.utils.quant_sync_format_check``); ``VERL_QUANT_LAYER_AUDIT`` switches both. + +Why. Matched train/rollout quantization presumes both sides quantize the same +layers, but each side decides on its own: training implicitly (a layer runs FP8 +GEMMs iff its TE linear modules run inside ``fp8_autocast``: the decoder layers, +minus ``first_last_layers_bf16``; never the embedding, the mcore output layer or +the torch-implemented MoE router), rollout by name (the trainer-side exporter +from ``rollout.quantization``, the engine from the configuration it was started +with). On B200, vLLM 0.24's ModelOpt MXFP8 config matched the router exclusion +``mlp.gate`` as a prefix of the dense ``mlp.gate_up_proj``: 56 of 311 parameters +of Qwen3-1.7B were served in BF16 while training ran them in FP8, and nothing +reported it. There training and the export agreed and the engine did not, which +is the rollout half's to catch; this module covers the training side. + +What. Two sources: + +- training: which decoder layers hold a TE FP8 weight workspace + (``module._fp8_workspaces``), per mcore ``TransformerLayer``; +- export: what the sync actually sends, read from the trainer's weight stream + after the rollout weight quantizer (``dispatch_rollout_weight_quantizer``): a + quantized weight arrives in FP8 (or packed FP4, as ``weight_packed``), every + other weight in high precision. + +The comparison skips the modules the exporter keeps in BF16 by design (its +``DEFAULT_IGNORE``: ``lm_head``, embeddings, vision towers, MTP layers, routers, +the shared-expert gate) and judges the GEMM projections of each decoder layer +(attention, MLP, experts, linear-attention / SSM in and out projections) against +that layer's FP8 trace. A decoder-layer weight of any other kind, such as a +convolution, is not judged, since the layer-level trace says nothing about it; +the audit reports how many it left out. + +It runs once, when training runs a TE FP8 recipe and the rollout is quantized, at +the first weight sync that happens after an FP8 forward (the first sync, before any +step, has no workspaces yet). ``VERL_QUANT_LAYER_AUDIT``: ``warn`` (default) logs the +disagreements, ``raise`` raises ``RuntimeError``, ``0`` disables. + +Two things can remove the training-side signal, and the audit says so instead of +staying silent. verl drops the workspaces when it offloads the training model +(``offload_megatron_model_to_cpu``, at the end of every ``train_mode()`` when +``param_offload`` is on), so the clearing helper first marks every module that +held one (``remember_fp8_workspaces``) and the audit accepts the marker. And TE +only caches the FP8 weight when mcore passes ``is_first_microbatch``, which +``disable_parameter_transpose_cache=True`` turns off; with that flag there is +nothing to read, and the audit warns once that it cannot run. +""" + +from __future__ import annotations + +import logging +import os +import re +import sys +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass, field +from typing import Optional + +import torch + + +def _configure_audit_logger() -> logging.Logger: + """Route the audit's own output to stderr, past whatever the root logger was set to. + + The audit's evidence lines are INFO/WARNING. Inside the trainer worker process a vLLM transitive + dependency (``model_hosting_container_standards.logging_config.configure_root_logger``, pulled in + for SageMaker) raises the root logger and every root handler to ERROR at import time; anything that + reaches this logger and then propagates to the root handlers is silently dropped - which is why the + audit was invisible on the first B200 runs even though it was executing. A dedicated stderr handler + on this logger with ``propagate=False`` carries the audit's INFO/WARNING lines out on their own, + independent of the root configuration. + """ + log = logging.getLogger(__name__) + log.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO").upper()) + log.propagate = False + handler_name = "verl.quant_audit.stderr" + if not any(getattr(h, "name", None) == handler_name for h in log.handlers): + handler = logging.StreamHandler(sys.stderr) + handler.set_name(handler_name) + handler.setLevel(logging.NOTSET) # let the logger's own level decide + handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s [quant-audit pid=%(process)d] %(message)s")) + log.addHandler(handler) + return log + + +logger = _configure_audit_logger() + +_LAYER_RE = re.compile(r"(?:^|\.)layers\.(\d+)\.") +# Routers (and mcore's shared-expert gate) are plain torch linears on the training side, never TE, +# regardless of how the HF checkpoint names them. +_ROUTER_RE = re.compile(r"(?:^|\.)(gate|router|shared_expert_gate)\.weight$") +# Non-linear leaves inside a decoder layer that never see a TE GEMM. +_NON_LINEAR_RE = re.compile(r"(norm|embed|rotary|bias$)", re.IGNORECASE) +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) +# compressed-tensors' name for a packed FP4 weight; the module's ``weight`` in every other format. +_PACKED_SUFFIX = ".weight_packed" +_MAX_LISTED = 40 + + +def audit_mode() -> str: + return os.environ.get("VERL_QUANT_LAYER_AUDIT", "warn").strip().lower() + + +def trace(where: str, **state) -> None: + """Log one checkpoint of the audit's own execution, at WARNING, when tracing is on. + + A guard that produces nothing is impossible to diagnose from the outside: telling "never + installed" from "installed but never reached" from "reached but found nothing" otherwise needs + instrumented runs. ``VERL_QUANT_LAYER_AUDIT_TRACE=1`` makes each of those states visible in one + run, at a level no logging configuration drops. + """ + if os.environ.get("VERL_QUANT_LAYER_AUDIT_TRACE", "0") != "1": + return + detail = " ".join(f"{k}={v}" for k, v in state.items()) + logger.warning("quantized-layer audit trace: %s%s", where, f" | {detail}" if detail else "") + + +# --------------------------------------------------------------------------- +# Training side: which decoder layers ran FP8 GEMMs +# --------------------------------------------------------------------------- + + +@dataclass +class TrainFp8Report: + """Which decoder layers ran FP8 GEMMs on this rank, read from TE's workspaces.""" + + seen_layers: set[int] = field(default_factory=set) + quantized_layers: set[int] = field(default_factory=set) + outside_layers_quantized: list[str] = field(default_factory=list) + # mcore's TE wrappers carry config.disable_parameter_transpose_cache; True means TE never caches + # the FP8 weight (is_first_microbatch=None), so no workspace can ever appear. + transpose_cache_disabled: bool = False + + @property + def active(self) -> bool: + return bool(self.quantized_layers) or bool(self.outside_layers_quantized) + + +def fp8_recipe_of(engine_config) -> Optional[str]: + """The Transformer Engine FP8 recipe a Megatron engine trains with, or None for BF16 GEMMs. + + Read from ``override_transformer_config``, the keys the Megatron provider receives: ``fp8`` + turns TE's FP8 GEMMs on and ``fp8_recipe`` picks the scaling recipe (Megatron's default is + ``delayed``). Engines without that config return None. + """ + overrides = getattr(engine_config, "override_transformer_config", None) + if not overrides or not overrides.get("fp8"): + return None + return str(overrides.get("fp8_recipe") or "delayed").lower() + + +# Set on a TE module whose FP8 weight workspace verl has dropped (param offload); survives the offload. +FP8_WORKSPACE_SEEN_ATTR = "_verl_fp8_workspace_seen" + + +def remember_fp8_workspaces(module: torch.nn.Module) -> bool: + """Mark ``module`` as having held an FP8 weight workspace. Call before clearing ``_fp8_workspaces``.""" + ws = getattr(module, "_fp8_workspaces", None) + if isinstance(ws, dict) and len(ws) > 0: + setattr(module, FP8_WORKSPACE_SEEN_ATTR, True) + return True + return False + + +def _has_fp8_workspace(module: torch.nn.Module) -> bool: + ws = getattr(module, "_fp8_workspaces", None) + return (isinstance(ws, dict) and len(ws) > 0) or bool(getattr(module, FP8_WORKSPACE_SEEN_ATTR, False)) + + +def _is_transformer_layer(module: torch.nn.Module) -> bool: + return isinstance(getattr(module, "layer_number", None), int) and "TransformerLayer" in type(module).__name__ + + +def collect_train_fp8_report(modules: Iterable[torch.nn.Module] | torch.nn.Module) -> TrainFp8Report: + """Walk Megatron model chunks and record, per global decoder layer, whether any TE module quantized.""" + if isinstance(modules, torch.nn.Module): + modules = [modules] + report = TrainFp8Report() + for chunk in modules: + in_layer: set[int] = set() + for _, layer in chunk.named_modules(): + if not _is_transformer_layer(layer): + continue + idx = layer.layer_number - 1 # mcore layer_number is 1-based and global across PP + report.seen_layers.add(idx) + for sub in layer.modules(): + in_layer.add(id(sub)) + if getattr(sub, "disable_parameter_transpose_cache", False): + report.transpose_cache_disabled = True + if _has_fp8_workspace(sub): + report.quantized_layers.add(idx) + for name, sub in chunk.named_modules(): + if id(sub) not in in_layer and _has_fp8_workspace(sub): + report.outside_layers_quantized.append(name) + return report + + +# transformers >= 5 saves MoE experts fused and without a ``.weight`` suffix +# (``mlp.experts.gate_up_proj`` [E, 2I, H], ``mlp.experts.down_proj`` [E, H, I]). +_FUSED_EXPERT_RE = re.compile(r"\.experts\.(gate_up_proj|down_proj)$") + + +def is_linear_weight_name(name: str) -> bool: + return name.endswith(".weight") or bool(_FUSED_EXPERT_RE.search(name)) + + +# HF leaf modules of the projections a decoder layer runs as TE GEMMs: attention (including MLA), MLP, +# experts (Mixtral's w1/w2/w3 too) and the in / out projections of linear-attention and SSM mixers. +_GEMM_WEIGHT_RE = re.compile( + r"(?:^|\.)(?:q_proj|k_proj|v_proj|o_proj|qkv_proj|query_key_value|q_a_proj|q_b_proj|kv_a_proj_with_mqa|" + r"kv_b_proj|gate_proj|up_proj|down_proj|gate_up_proj|in_proj\w*|out_proj|fc1|fc2|w1|w2|w3)\.weight$" +) + + +def is_gemm_weight_name(name: str) -> bool: + return bool(_GEMM_WEIGHT_RE.search(name) or _FUSED_EXPERT_RE.search(name)) + + +def _of_unjudged_kind(name: str) -> bool: + """A decoder-layer weight that is neither a known GEMM projection nor a known non-GEMM leaf.""" + return ( + is_linear_weight_name(name) + and _LAYER_RE.search(name) is not None + and not (_ROUTER_RE.search(name) or _NON_LINEAR_RE.search(name)) + and not is_gemm_weight_name(name) + ) + + +def expected_train_quantized(name: str, report: TrainFp8Report) -> Optional[bool]: + """Training-side expectation for one HF weight name; ``None`` when this rank cannot judge.""" + if not is_linear_weight_name(name): + return False + m = _LAYER_RE.search(name) + if m is None: + # embeddings, lm_head, vision towers, ...: mcore quantizes nothing outside the decoder stack + return None if report.outside_layers_quantized else False + idx = int(m.group(1)) + if idx not in report.seen_layers: + return None # another pipeline stage owns this layer + if _ROUTER_RE.search(name) or _NON_LINEAR_RE.search(name): + return False + if not is_gemm_weight_name(name): + return None # e.g. a convolution: the layer's FP8 trace does not say how training ran it + return idx in report.quantized_layers + + +# --------------------------------------------------------------------------- +# Export side: what the weight sync sends +# --------------------------------------------------------------------------- + + +def exported_weight(name: str, tensor: torch.Tensor) -> Optional[tuple[str, bool, str]]: + """``(weight name, quantized, format)`` for one entry of the exported stream; ``None`` for scales and the rest.""" + if name.endswith(_PACKED_SUFFIX): + return name[: -len("_packed")], True, "packed FP4" + if not is_linear_weight_name(name): + return None + dtype = getattr(tensor, "dtype", None) + return name, dtype in _FP8_DTYPES, str(dtype).removeprefix("torch.") + + +def kept_high_precision_by_design(name: str) -> bool: + """Whether the exporter keeps this weight's module in BF16 whatever the preset (its ``DEFAULT_IGNORE``).""" + from verl.utils.qat.compressed_tensors import DEFAULT_IGNORE, _matches + + module = name[: -len(".weight")] if name.endswith(".weight") else name + return _matches(module, DEFAULT_IGNORE) + + +def unjudged_weights(exported: Mapping[str, tuple[bool, str]]) -> list[str]: + """Exported decoder-layer weights of a kind the training-vs-export comparison does not judge.""" + return [name for name in exported if _of_unjudged_kind(name) and not kept_high_precision_by_design(name)] + + +def _unjudged_note(unjudged: list[str]) -> str: + if not unjudged: + return "" + shown = ", ".join(unjudged[:3]) + (", ..." if len(unjudged) > 3 else "") + return f"; {len(unjudged)} decoder-layer weights of a kind the audit does not judge, e.g. {shown}" + + +def audit_train_vs_export(report: TrainFp8Report, exported: Mapping[str, tuple[bool, str]]) -> tuple[list[str], int]: + """One line per exported weight on which training and the export disagree, and the by-design skip count.""" + problems = [] + by_design = 0 + for name, (quantized, fmt) in exported.items(): + if kept_high_precision_by_design(name): + by_design += 1 + continue + expected = expected_train_quantized(name, report) + if expected is None: + continue + if quantized and not expected: + problems.append(f"{name}: the export sends it in {fmt}, training ran it in high precision") + elif expected and not quantized: + problems.append(f"{name}: training ran it in FP8, the export sends it in {fmt}") + return problems, by_design + + +def _listing(problems: list[str]) -> str: + return ( + "\n - " + + "\n - ".join(problems[:_MAX_LISTED]) + + (f"\n - ... and {len(problems) - _MAX_LISTED} more" if len(problems) > _MAX_LISTED else "") + ) + + +# --------------------------------------------------------------------------- +# Training vs export, once per run +# --------------------------------------------------------------------------- + + +class TrainExportAudit: + """Training vs export: did training run each decoder layer in the format the weight sync sends it in? + + ``wrap_export`` wraps the training engine's weight export after the rollout weight quantizer, the one + point every sync route passes through (the colocated worker, the checkpoint engine, server replicas): + it records the format of every exported weight and, when the stream ends, calls + ``compare_with_training``. The comparison needs TE's FP8 weight workspaces, which exist only after a + training step, so the sync before any step waits and the first sync after one decides; from then on + the export passes through untouched. + """ + + def __init__(self, mode: str, engine: str = "", enabled: bool = True): + self.mode = mode + self.engine = engine + self.enabled = enabled and mode not in ("0", "off", "false") + self.finished = False # compared, or gave up for lack of a training-side FP8 trace + # HF weight name -> (sent quantized, format), from the last recorded sync + self.exported: dict[str, tuple[bool, str]] = {} + self.syncs_seen = 0 # distinct weight syncs wrapped; the give-up threshold counts these + self._explained = False # one-shot "why this guard produced nothing" warning + + @classmethod + def create(cls, *, rollout_quantized: bool, trainer_fp8_recipe: Optional[str], rollout_name: str): + """Arm the comparison for a TE FP8 trainer whose rollout is quantized; otherwise return a disabled audit.""" + mode = audit_mode() + if not rollout_quantized or trainer_fp8_recipe is None: + logger.info( + "quantized-layer audit, training vs export: not armed (rollout quantized=%s, training FP8 recipe=%s): " + "no FP8 training to compare", + rollout_quantized, + trainer_fp8_recipe, + ) + return cls(mode, rollout_name, enabled=False) + logger.info( + "quantized-layer audit, training vs export: armed (engine=%s, training FP8 recipe=%s, mode=%s); runs at " + "the first weight sync after a training step", + rollout_name, + trainer_fp8_recipe, + mode, + ) + return cls(mode, rollout_name) + + def wrap_export( + self, weights: Iterable[tuple[str, torch.Tensor]], training_modules=None + ) -> Iterator[tuple[str, torch.Tensor]]: + """Pass the exported stream through, recording each weight's format; compare with training when it ends. + + ``training_modules`` (the Megatron model chunks) makes the comparison run as soon as the stream is + exhausted, wherever the stream is consumed. + """ + trace("wrap_export", enabled=self.enabled, finished=self.finished, with_modules=training_modules is not None) + if not self.enabled or self.finished: + yield from weights + return + self.exported = {} + self.syncs_seen += 1 # one wrapped weight stream == one sync + for name, tensor in weights: + entry = exported_weight(name, tensor) + if entry is not None: + self.exported[entry[0]] = entry[1:] + yield name, tensor + trace("wrap_export.stream_end", weights=len(self.exported), will_compare=training_modules is not None) + if training_modules is not None: + self.compare_with_training(training_modules) + + def _explain_silence(self, reason: str) -> None: + """Warn once that the comparison produced no verdict: a guard that silently does nothing is the failure + this module exists to prevent, so the reason is reported at WARNING.""" + if self._explained: + return + self._explained = True + logger.warning( + "quantized-layer audit, training vs export, produced no verdict (%s). engine=%s weights=%d. Training " + "vs export is unverified for this run; VERL_QUANT_LAYER_AUDIT=0 silences the audit entirely.", + reason, + self.engine or "?", + len(self.exported), + ) + + def _report(self, msg: str) -> None: + if self.mode == "raise": + raise RuntimeError(msg) + logger.warning(msg) + + def compare_with_training(self, training_modules) -> Optional[list[str]]: + """Compare the training-side FP8 trace with the recorded export. ``None`` if it did not run.""" + trace("compare_with_training", enabled=self.enabled, finished=self.finished, weights=len(self.exported)) + if self.finished or not self.enabled: + return None + if not self.exported: + self._explain_silence("the weight sync streamed no weights through the audit") + return None + report = collect_train_fp8_report(training_modules) + if not report.active: + # Give up only after two distinct syncs left no FP8 trace, counted by wrap_export(), so a second + # comparison on the same sync cannot trip this before the first training step. + if report.transpose_cache_disabled or self.syncs_seen >= 2: + self.finished = True + cause = ( + "disable_parameter_transpose_cache=True: TE caches the FP8 weight only when mcore passes " + "is_first_microbatch, which that flag turns off" + if report.transpose_cache_disabled + else "no TE module holds an FP8 weight workspace after a training step; check that " + "override_transformer_config.fp8 is set on the model that ran, and that nothing other " + "than offload_megatron_model_to_cpu clears module._fp8_workspaces" + ) + self._explain_silence( + f"no training-side FP8 trace on this rank: 0 of {len(report.seen_layers)} decoder layers " + f"left an FP8 weight workspace. Cause: {cause}" + ) + return None + logger.info( + "quantized-layer audit, training vs export: no FP8 trace yet on this rank (%d decoder layers seen, " + "%d weights recorded); expected before the first training step, will retry at the next sync", + len(report.seen_layers), + len(self.exported), + ) + return None + problems, by_design = audit_train_vs_export(report, self.exported) + unjudged = unjudged_weights(self.exported) + self.finished = True + n_q, n_seen = len(report.quantized_layers), len(report.seen_layers) + if not problems: + logger.info( + "quantized-layer audit: training and the exported weights agree (%d of %d decoder layers on this " + "rank ran FP8 GEMMs; %d weights checked, %d kept in BF16 on the rollout by design%s)", + n_q, + n_seen, + len(self.exported) - by_design - len(unjudged), + by_design, + _unjudged_note(unjudged), + ) + return [] + self._report( + f"quantized-layer audit: training and the exported weights disagree on {len(problems)} weight(s) " + f"({n_q} of {n_seen} decoder layers on this rank ran FP8 GEMMs{_unjudged_note(unjudged)}):" + + _listing(problems) + + "\nThe trainer ran these layers in one precision and the rollout is served them in another. Fix with " + "rollout.quantization.ignore (rollout) or first_last_layers_bf16 / model wiring (training)." + + " VERL_QUANT_LAYER_AUDIT=raise turns this into an error, =0 silences it." + ) + return problems diff --git a/verl/utils/quant_sync_format_check.py b/verl/utils/quant_sync_format_check.py new file mode 100644 index 00000000000..ed98e772357 --- /dev/null +++ b/verl/utils/quant_sync_format_check.py @@ -0,0 +1,149 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Export vs engine: is each synced weight sent in the format of the engine parameter that receives it? + +This is the rollout half of the quantized-layer audit; ``verl.utils.quant_layer_audit`` holds the training +half (training vs export). Each rollout engine runs it in its own processes, inside the loader that writes a +weight sync, so it reads the very parameters the sync is about to write: + +- vLLM: the colocated worker extension (``verl.workers.rollout.vllm_rollout.utils``), with the name + resolution of ``verl.utils.vllm.vllm_sync_format_check``; +- SGLang: the MXFP8/MXFP4/NVFP4 refit loader (``verl.utils.sglang.sglang_quant_utils.load_and_reprocess``), + with the name resolution of ``verl.utils.sglang.sglang_sync_format_check``. + +It compares quantized against high precision, not exact dtypes. A quantized weight (FP8, packed FP4) sent to a +BF16 parameter, or a BF16 weight sent to a quantized one, would be cast into a buffer of the other format, and a +quantized layer would never receive its scale. The trainer-side exporter and the engine read the same +quantization config (``rollout.quantization``, or the checkpoint's own ``quantization_config``), so a +disagreement means the two matched a module differently, for example an ``ignore`` entry one side reads as a +prefix and the other as an exact name. + +What a sync sends follows that config, and the engine keeps the formats it was built with, so a later sync cannot +disagree where the first one agreed. The check therefore runs on every bucket of the first weight sync, before +the bucket is written, and reports its totals once when that sync finishes; a sync that does not finish is +checked again in full. + +``VERL_QUANT_LAYER_AUDIT``, shared with the training half: ``warn`` (default) logs the disagreements and lets the +sync go on, ``raise`` stops the sync before the bucket is written, ``0`` disables. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from typing import Optional + +import torch + +from verl.utils.quant_layer_audit import audit_mode, logger + +# HF leaf names that can be linear weights: ``.weight``, compressed-tensors' packed FP4 ``.weight_packed``, and +# the fused expert tensors transformers >= 5 writes without a suffix (``mlp.experts.gate_up_proj``). +_CHECKED_LEAVES = ("weight", "weight_packed", "gate_up_proj", "down_proj") +_HIGH_PRECISION = (torch.bfloat16, torch.float16, torch.float32) +_OFF = ("0", "off", "false") +_MAX_LISTED = 40 + +# (engine model, HF weight name) -> the engine parameter that receives it, or None when the name does not resolve +ResolveParam = Callable[[torch.nn.Module, str], Optional[torch.Tensor]] + + +def _format(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +class SyncFormatCheck: + """Export vs engine for one engine model. It lives on the model, so each engine process keeps its own.""" + + _ATTR = "_verl_sync_format_check" + + def __init__(self, engine: str, resolve_param: ResolveParam): + self.engine = engine + self.resolve_param = resolve_param + self.finished = False # the first weight sync has been checked end to end + # totals of the sync being checked + self.checked = 0 # weights compared with the parameter that receives them + self.disagreeing = 0 # of those, sent in another format than the parameter's + self.unmapped = 0 # names that resolve onto no engine parameter: not judged + self._buckets = 0 + + @classmethod + def of(cls, model: torch.nn.Module, engine: str, resolve_param: ResolveParam) -> SyncFormatCheck: + """The check attached to ``model``, created on first use.""" + check = model.__dict__.get(cls._ATTR) + if check is None: + check = cls(engine, resolve_param) + model.__dict__[cls._ATTR] = check + return check + + @property + def pending(self) -> bool: + return not self.finished and audit_mode() not in _OFF + + def begin_sync(self) -> None: + """Start a weight sync. Totals left by a sync that did not finish are dropped: it is checked again in full.""" + if self.finished: + return + self.checked = self.disagreeing = self.unmapped = self._buckets = 0 + + def check_bucket(self, model: torch.nn.Module, named_tensors: Iterable[tuple[str, torch.Tensor]]) -> None: + """Report the bucket's weights sent in another format than the receiving parameter's. Call before writing.""" + if not self.pending: + return + problems = [] + for name, tensor in named_tensors: + if name.rsplit(".", 1)[-1] not in _CHECKED_LEAVES: + continue + target = self.resolve_param(model, name) + if target is None: + self.unmapped += 1 + continue + self.checked += 1 + if (tensor.dtype in _HIGH_PRECISION) != (target.dtype in _HIGH_PRECISION): + problems.append( + f"{name}: the sync sends {_format(tensor.dtype)}, the engine parameter is {_format(target.dtype)}" + ) + self._buckets += 1 + self.disagreeing += len(problems) + if problems: + self._report(problems) + + def end_sync(self) -> None: + """After the first finished sync: report the totals once and stop checking.""" + if not self.pending or self._buckets == 0: + return + self.finished = True + logger.info( + "quantized-layer audit, export vs %s engine, first weight sync: %d weights checked, %d in a format other " + "than the receiving parameter's, %d names not mapped onto an engine parameter (not judged); later syncs " + "send the same formats and are not checked again", + self.engine, + self.checked, + self.disagreeing, + self.unmapped, + ) + + def _report(self, problems: list[str]) -> None: + msg = ( + f"quantized-layer audit: the weight sync and the {self.engine} engine disagree on the format of " + f"{len(problems)} weight(s):\n - " + + "\n - ".join(problems[:_MAX_LISTED]) + + (f"\n - ... and {len(problems) - _MAX_LISTED} more" if len(problems) > _MAX_LISTED else "") + + "\nThe trainer-side exporter and the engine read the same quantization config (rollout.quantization, " + "or the checkpoint's own quantization_config) but matched these modules differently. Loading would cast " + "each weight into a parameter of the other format, and a quantized layer would never receive its scale. " + "Check rollout.quantization.ignore against docs/low_precision/mxfp8.md (Checks)." + ) + if audit_mode() == "raise": + raise RuntimeError(msg + " The sync is stopped before this bucket is written.") + logger.warning("%s The sync goes on; VERL_QUANT_LAYER_AUDIT=raise stops it here, =0 disables the check.", msg) diff --git a/verl/utils/sglang/sglang_quant_utils.py b/verl/utils/sglang/sglang_quant_utils.py index 0928d68b26b..9d2262c9230 100644 --- a/verl/utils/sglang/sglang_quant_utils.py +++ b/verl/utils/sglang/sglang_quant_utils.py @@ -13,6 +13,7 @@ # limitations under the License. """SGLang MX low-precision config translation, weight export, and refit registration.""" +import logging import os from collections.abc import Iterable, Iterator, Mapping from copy import deepcopy @@ -34,6 +35,10 @@ _SYNC_NVFP4_ATTR = "_verl_sglang_nvfp4_refit_staged" _SGLangRefitKind = Literal["mxfp8", "mxfp4", "nvfp4"] +logger = logging.getLogger(__name__) +# SGLang's scheduler processes fall back to the root WARNING level and would drop the INFO evidence lines. +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + _REGEX_IGNORE_TO_SGLANG = { r"re:.*lm_head$": "lm_head", r"re:.*embed_tokens$": "embed_tokens", @@ -232,16 +237,26 @@ def finalize_tensors() -> list[tuple[str, torch.Tensor]]: def load_and_reprocess(model: torch.nn.Module, named_tensors: Iterable[tuple[str, torch.Tensor]]) -> None: - """Load one SGLang quantized-refit bucket and finalize kernel layouts once per sync.""" + """Load one SGLang quantized-refit bucket and finalize kernel layouts once per sync. + + Before a bucket is written, the export-vs-engine format check of the first weight sync reads it + (``verl.utils.sglang.sglang_sync_format_check``). + """ from verl.utils.sglang.sglang_fp4_utils import stage_mxfp4_moe_params, stage_nvfp4_moe_params from verl.utils.sglang.sglang_fp8_utils import refuse_relaid_bf16_moe_layers, stage_mxfp8_moe_scales + from verl.utils.sglang.sglang_sync_format_check import sglang_sync_format_check named_tensors = list(named_tensors) finalize = any(name == FINALIZE_MARKER for name, _ in named_tensors) named_tensors = [(name, tensor) for name, tensor in named_tensors if name != FINALIZE_MARKER] + format_check = sglang_sync_format_check(model) try: if named_tensors: - if _SYNC_FP8_ATTR not in model.__dict__: + first_request = _SYNC_FP8_ATTR not in model.__dict__ # first request of this sync + if first_request: + format_check.begin_sync() + format_check.check_bucket(model, named_tensors) # before anything is written into the engine + if first_request: refuse_relaid_bf16_moe_layers(model) model.__dict__[_SYNC_FP8_ATTR] = stage_mxfp8_moe_scales(model) model.__dict__[_SYNC_FP4_ATTR] = stage_mxfp4_moe_params(model) @@ -249,6 +264,7 @@ def load_and_reprocess(model: torch.nn.Module, named_tensors: Iterable[tuple[str model.load_weights(named_tensors) if finalize: _finish_sync(model) + format_check.end_sync() except BaseException: _abandon_sync(model) raise diff --git a/verl/utils/sglang/sglang_sync_format_check.py b/verl/utils/sglang/sglang_sync_format_check.py new file mode 100644 index 00000000000..cea34236dad --- /dev/null +++ b/verl/utils/sglang/sglang_sync_format_check.py @@ -0,0 +1,77 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Export vs engine on SGLang: the parameter that receives each synced HF weight. + +The check itself is ``verl.utils.quant_sync_format_check.SyncFormatCheck``. The MXFP8/MXFP4/NVFP4 refit loader +(``verl.utils.sglang.sglang_quant_utils.load_and_reprocess``) calls it: ``begin_sync`` on the first request of a +sync, ``check_bucket`` before each bucket is loaded and ``end_sync`` after the kernel layouts are re-derived. +Formats SGLang loads without that loader (FP8 blockwise) are not checked. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from verl.utils.quant_sync_format_check import SyncFormatCheck + +# HF projections SGLang fuses into one module (the model's stacked_params_mapping, restated). +_STACKED = { + "q_proj": "qkv_proj", + "k_proj": "qkv_proj", + "v_proj": "qkv_proj", + "gate_proj": "gate_up_proj", + "up_proj": "gate_up_proj", +} +_W2_LEAVES = ("down_proj", "w2") +_PARAM_LEAVES = ("weight", "weight_packed") + + +def resolve_sglang_param(model: torch.nn.Module, hf_name: str) -> Optional[torch.Tensor]: + """The SGLang parameter that receives HF weight ``hf_name``, or ``None`` when the name does not resolve. + + Walks the module tree by name, trying the fused module name for projections SGLang stacks + (``q_proj`` -> ``qkv_proj``, ``gate_proj`` -> ``gate_up_proj``) and stopping at a fused-MoE block + (a module holding ``w13_weight``), where the remaining per-expert parts select ``w13`` or ``w2``. + A packed FP4 ``.weight_packed`` lands in the module's ``weight_packed``, or in ``weight`` on a layer + the engine keeps in BF16. + """ + parts = hf_name.split(".") + leaf = parts[-1] if parts[-1] in _PARAM_LEAVES else None + module_path = parts[:-1] if leaf else parts + cur: object = model + for i, part in enumerate(module_path): + if hasattr(cur, "w13_weight"): + return cur.w2_weight if any(r in _W2_LEAVES for r in module_path[i:]) else cur.w13_weight + nxt = getattr(cur, part, None) + if nxt is None and part in _STACKED: + nxt = getattr(cur, _STACKED[part], None) + if nxt is None: + return None + cur = nxt + if hasattr(cur, "w13_weight"): + return cur.w2_weight if parts[-1] in _W2_LEAVES else cur.w13_weight + if isinstance(cur, torch.nn.Module): + for attr in ("weight_packed", "weight") if leaf == "weight_packed" else ("weight",): + param = getattr(cur, attr, None) + if isinstance(param, torch.Tensor): + return param + return None + return cur if isinstance(cur, torch.Tensor) else None + + +def sglang_sync_format_check(model: torch.nn.Module) -> SyncFormatCheck: + """The export-vs-engine check of an SGLang model.""" + return SyncFormatCheck.of(model, "sglang", resolve_sglang_param) diff --git a/verl/utils/vllm/vllm_sync_format_check.py b/verl/utils/vllm/vllm_sync_format_check.py new file mode 100644 index 00000000000..6c0d85a0d4b --- /dev/null +++ b/verl/utils/vllm/vllm_sync_format_check.py @@ -0,0 +1,97 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Export vs engine on vLLM: the parameter that receives each synced HF weight. + +The check itself is ``verl.utils.quant_sync_format_check.SyncFormatCheck``. The colocated worker extension +(``verl.workers.rollout.vllm_rollout.utils``) calls it around its quantized weight loading: ``begin_sync`` +before the first bucket, ``check_bucket`` before each bucket is loaded and ``end_sync`` after the kernel +layouts are re-derived. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from verl.utils.quant_sync_format_check import SyncFormatCheck + + +def _vllm_receiving_module(model: torch.nn.Module, name: str) -> Optional[torch.nn.Module]: + """The vLLM module that receives HF weight ``name``, or ``None`` when the name does not resolve. + + Applies the model's ``hf_to_vllm_mapper``, then its ``packed_modules_mapping`` for the last + module name (``q_proj`` -> ``qkv_proj``), and stops at a fused-MoE block, whose remaining + per-expert parts (``experts.3.gate_proj``) its own weight loader resolves. + """ + mapper = getattr(model, "hf_to_vllm_mapper", None) + if mapper is not None: + mapped = mapper.apply_list([name]) + if not mapped: + return None + name = mapped[0] + parts = name.split(".") + path = parts[:-1] if parts[-1] in ("weight", "weight_packed") else parts + packed = getattr(model, "packed_modules_mapping", None) or {} + owner = {sub: fused for fused, subs in packed.items() for sub in subs} + module = model + for i, part in enumerate(path): + if hasattr(module, "routed_experts") or hasattr(module, "w13_weight"): + break + child = getattr(module, part, None) + if child is None and i == len(path) - 1 and part in owner: + child = getattr(module, owner[part], None) + if not isinstance(child, torch.nn.Module): + return None + module = child + for holder in ("routed_experts", "base_layer"): # the MoE weight holder; the layer under a LoRA wrapper + inner = getattr(module, holder, None) + if isinstance(inner, torch.nn.Module): + module = inner + return module + + +def resolve_vllm_param(model: torch.nn.Module, hf_name: str) -> Optional[torch.Tensor]: + """The vLLM parameter that receives HF weight ``hf_name``, or ``None`` when the name does not resolve. + + A linear layer receives ``.weight`` in its ``weight`` and compressed-tensors' packed FP4 ``.weight_packed`` + in its ``weight_packed``; a layer the engine keeps in BF16 has only ``weight``, which then receives the + packed weight. A fused-MoE block receives every expert's gate and up projections in ``w13_weight`` and + its down projection in ``w2_weight``. Layers another pipeline stage owns are unresolved on this rank. + """ + module = _vllm_receiving_module(model, hf_name) + if module is None or type(module).__name__ == "PPMissingLayer": + return None + if hasattr(module, "w13_weight"): + parts = hf_name.split(".") + expert_parts = parts[parts.index("experts") + 1 :] if "experts" in parts else parts + return module.w2_weight if any(p in ("down_proj", "w2") for p in expert_parts) else module.w13_weight + candidates = ("weight_packed", "weight") if hf_name.endswith(".weight_packed") else ("weight",) + for attr in candidates: + param = getattr(module, attr, None) + if isinstance(param, torch.Tensor): + return param + return None + + +def vllm_sync_format_check(model: torch.nn.Module, quant_config: Any = None) -> Optional[SyncFormatCheck]: + """The export-vs-engine check of a vLLM model, or ``None`` when the engine quantizes what it receives. + + verl syncs weights the trainer-side exporter has already quantized into the checkpoint format of + ``rollout.quantization`` or of the model's own ``quantization_config``. An engine started with online + quantization (an unquantized checkpoint, ``is_checkpoint_fp8_serialized=False``) receives BF16 by design. + """ + if getattr(quant_config, "is_checkpoint_fp8_serialized", True) is False: + return None + return SyncFormatCheck.of(model, "vllm", resolve_vllm_param) diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index d473ce66c4e..0202313e1c7 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -212,6 +212,7 @@ def __init__( ) logger.info(f"QAT enabled in MegatronEngine: mode={self._qat_config.mode}") self._rollout_weight_quantizer = None + self._train_export_audit = None # Router replay configuration for MoE models self.enable_routing_replay = self.engine_config.router_replay.mode != "disabled" @@ -1015,12 +1016,20 @@ def get_per_tensor_param(self, base_sync_done=False, **kwargs): if refresh_scales is not None: refresh_scales(self.module) per_tensor_param = self._rollout_weight_quantizer(per_tensor_param) + if self._train_export_audit is not None: + # Every sync route passes through this export, so the audit reads what the rollout is sent + # here; it compares with the training side when the stream ends. + per_tensor_param = self._train_export_audit.wrap_export(per_tensor_param, training_modules=self.module) return per_tensor_param, peft_config def set_rollout_weight_quantizer(self, quantizer) -> None: self._rollout_weight_quantizer = quantizer + def set_train_export_audit(self, audit) -> None: + """Compare the quantized export with the layers that ran FP8 in training (``verl.utils.quant_layer_audit``).""" + self._train_export_audit = audit + def _mcore_export_index(self): """Build (once) the per-parameter delta export index: geometry specs and form-B probes derived from the bridge's conversion tasks. See diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index b724eae7c49..5246a9c85ad 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -737,7 +737,17 @@ def init_model(self): qat_config=qat, ) if weight_quantizer is not None: + from verl.utils.quant_layer_audit import TrainExportAudit, fp8_recipe_of + self.actor.engine.set_rollout_weight_quantizer(weight_quantizer) + # Training vs export; export vs engine runs in the rollout engine's weight loader. + self.actor.engine.set_train_export_audit( + TrainExportAudit.create( + rollout_quantized=True, + trainer_fp8_recipe=fp8_recipe_of(getattr(self.actor.engine, "engine_config", None)), + rollout_name=self.config.rollout.name, + ) + ) # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo aggressive_empty_cache(force_sync=True) diff --git a/verl/workers/rollout/vllm_rollout/utils.py b/verl/workers/rollout/vllm_rollout/utils.py index 241eac8fbcd..73c6ce48f28 100644 --- a/verl/workers/rollout/vllm_rollout/utils.py +++ b/verl/workers/rollout/vllm_rollout/utils.py @@ -245,6 +245,12 @@ def monkey_patch_model(self, vocab_size: int, banned_token_ids: Optional[list[in # patch weight loader to support MoE model patch_vllm_moe_model_weight_loader(model) + def _sync_format_check(self): + """Export vs engine for the main model, or ``None`` (``verl.utils.vllm.vllm_sync_format_check``).""" + from verl.utils.vllm.vllm_sync_format_check import vllm_sync_format_check + + return vllm_sync_format_check(self.model_runner.model, self.model_runner.vllm_config.quant_config) + def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False, use_shm: bool = False): """Update the weights of the rollout model.""" if self.device is None: @@ -271,6 +277,9 @@ def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False quant_reload_states = [ (model, prepare_quanted_weights_for_loading(model)) for model in self._iter_all_models() ] + format_check = self._sync_format_check() + if format_check is not None: + format_check.begin_sync() else: # TODO(wuxibin): not need anymore for newer vllm version. for model in self._iter_all_models(): @@ -329,7 +338,24 @@ def on_bucket_received(weights: list[tuple[str, torch.Tensor]], is_last: bool) - base_sync_done=base_sync_done, ) - receiver.receive_weights(on_bucket_received=on_bucket_received) + # Finish the receiver handshake before reporting a bucket error, as the delta path does: the sender waits + # for every bucket's ack, so an error raised inside the receiver would leave it, and the trainer behind it, + # blocked for good. Buckets after the failed one are acked unread; the error is raised once the sender is done. + # Interrupts are not caught, so they still stop the worker at once. + bucket_error: Exception | None = None + + def load_or_skip(weights: list[tuple[str, torch.Tensor]], is_last: bool) -> None: + nonlocal bucket_error + if bucket_error is not None: + return + try: + on_bucket_received(weights, is_last) + except Exception as exc: + bucket_error = exc + + receiver.receive_weights(on_bucket_received=load_or_skip) + if bucket_error is not None: + raise bucket_error # =========================== step 3: process weights after loading =========================== if peft_config and base_sync_done: @@ -337,6 +363,9 @@ def on_bucket_received(weights: list[tuple[str, torch.Tensor]], is_last: bool) - elif is_quantized_model(self.model_runner.vllm_config): for model, reload_state in quant_reload_states: process_quanted_weights_after_loading(model, reload_state) + format_check = self._sync_format_check() + if format_check is not None: + format_check.end_sync() else: # Some post-load transforms are non-idempotent; run once after all buckets. from vllm.model_executor.model_loader.utils import process_weights_after_loading @@ -381,8 +410,11 @@ def _update_weights( # Check if FP8 quantization is enabled and apply appropriate weight loading if is_quantized_model(self.model_runner.vllm_config): logger.info(f"FP8 model detected (async): {self.model_runner.vllm_config.quant_config}") - # Convert bf16 weights to fp8 format before loading model_updates = drop_tied_alias_updates(self.model_runner.model, param_updates) + format_check = self._sync_format_check() + if format_check is not None and model_updates: + # export vs engine, first sync only: before the bucket is written + format_check.check_bucket(self.model_runner.model, model_updates) loaded_params = load_quanted_weights(model_updates, self.model_runner) if model_updates else [] # Keep the draft model in sync when present. if self._use_mtp_drafter_weight_sync() and model_updates: