[rollout, vllm, sglang, trainer] feat: quantized-layer audit

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:

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

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