FP8 RL Training and Rollout Weight Synchronization

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -413,12 +413,11 @@ def preprocess_packed_seqs(
       per row. This is the historical SkyRL behavior used by the RL path
       and the existing SFT path without mini-batch packing.
     - ``sub_seq_lengths is not None``: each row may contain multiple
-      sub-sequences concatenated end-to-end. ``sub_seq_lengths[r]`` lists
-      the per-sub-sequence valid token counts for row ``r``. Tokens
-      ``input_ids[r, :sum(sub_seq_lengths[r])]`` are assumed to be the
-      concatenated sub-sequences in order; any trailing tokens in the row
-      are pad. ``cu_seqlens`` enumerates every sub-sequence across every
-      row.
+      sub-sequences. ``sub_seq_lengths[r]`` lists their valid token counts.
+      Each sub-sequence begins at the next ``align_size`` boundary, so internal
+      alignment padding may separate adjacent sub-sequences; any remaining
+      trailing tokens are pad. ``cu_seqlens`` enumerates every sub-sequence
+      across every row.
 
     CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0
     gets first and last chunks, GPU1 gets second and second last chunks,

skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py

@@ -5,23 +5,29 @@
 def is_fp8_enabled(fp8: Any) -> bool:
     """Return whether a Megatron/TE fp8 config value enables FP8 execution."""
     if isinstance(fp8, str):
-        return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no"}
+        return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no", "off"}
     return bool(fp8)
 
 
 def get_packed_seq_align_size(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int:
-    """Return global per-subsequence padding needed for TP/CP layout."""
+    """Return the global alignment unit for packed TP/CP/FP8 sequences."""
+    if tp_size < 1 or cp_size < 1:
+        raise ValueError(f"tp_size and cp_size must be positive, got tp_size={tp_size}, cp_size={cp_size}")
     if cp_size > 1:
         layout_align = tp_size * cp_size * 2
     else:
         layout_align = tp_size
     if not fp8_enabled:
         return layout_align
-    return math.lcm(layout_align, 16 * cp_size)
+    fp8_token_align = 128 * tp_size * cp_size if tp_size > 1 else 16 * cp_size
+    return math.lcm(layout_align, fp8_token_align)
 
 
 def get_unpacked_seq_align_size(tp_size: int, fp8_enabled: bool = False) -> int:
-    """Return sequence padding needed when removing microbatch padding without CP."""
+    """Return the alignment unit for unpacked TP/FP8 sequences without CP."""
+    if tp_size < 1:
+        raise ValueError(f"tp_size must be positive, got {tp_size}")
     if not fp8_enabled:
         return tp_size
-    return math.lcm(tp_size, 16)
+    fp8_token_align = 128 * tp_size if tp_size > 1 else 16
+    return math.lcm(tp_size, fp8_token_align)

skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py

@@ -24,15 +24,136 @@
         skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap.NewInferenceWorkerWrap
 """
 
+from typing import Any
+
 import torch
 
 from skyrl.backends.skyrl_train.inference_servers.layerwise_reload import (
     LayerwiseReloadWorkerMixin,
 )
+from skyrl.backends.skyrl_train.weight_sync.base import cuda_uuid_to_str
+from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import (
+    SKYRL_BATCHED_MOE_FP8_PREFIX,
+)
 
 VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS = f"{__name__}.NewInferenceWorkerWrap"
 
 
+_BATCHED_MOE_TARGETS = {
+    ".experts.gate_proj.weight": (".experts.w13_weight", "w1"),
+    ".experts.up_proj.weight": (".experts.w13_weight", "w3"),
+    ".experts.down_proj.weight": (".experts.w2_weight", "w2"),
+    ".experts.gate_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w1"),
+    ".experts.up_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w3"),
+    ".experts.down_proj.weight_scale_inv": (".experts.w2_weight_scale_inv", "w2"),
+}
+
+
+def _map_hf_weight_name(model: torch.nn.Module, name: str) -> str:
+    """Apply a top-level vLLM model's HF-to-runtime prefix mapping."""
+    mapper = getattr(model, "hf_to_vllm_mapper", None)
+    if mapper is None:
+        return name
+    mapped = mapper.apply_list([name])
+    if len(mapped) != 1:
+        raise ValueError(f"Unable to map batched MoE checkpoint name {name!r}")
+    return mapped[0]
+
+
+def _load_batched_moe_fp8_tensor(
+    model: torch.nn.Module,
+    params_dict: dict[str, torch.nn.Parameter],
+    wire_name: str,
+    loaded_weight: torch.Tensor,
+) -> bool:
+    """Load one expert-batched FP8 weight/scale through FusedMoE's loader.
+
+    Returns ``False`` for ordinary checkpoint tensors. Marked tensors are
+    required to resolve successfully so a protocol mismatch cannot silently
+    leave stale rollout weights behind.
+    """
+    if not wire_name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX):
+        return False
+    if loaded_weight.ndim != 3:
+        raise ValueError(
+            f"Batched MoE wire tensor must be 3D, got name={wire_name!r}, shape={tuple(loaded_weight.shape)}"
+        )
+
+    checkpoint_name = wire_name.removeprefix(SKYRL_BATCHED_MOE_FP8_PREFIX)
+    mapped_name = _map_hf_weight_name(model, checkpoint_name)
+    target_name = None
+    shard_id = None
+    for checkpoint_suffix, (
+        target_suffix,
+        candidate_shard_id,
+    ) in _BATCHED_MOE_TARGETS.items():
+        if mapped_name.endswith(checkpoint_suffix):
+            target_name = mapped_name[: -len(checkpoint_suffix)] + target_suffix
+            shard_id = candidate_shard_id
+            break
+    if target_name is None or shard_id is None:
+        raise ValueError(f"Unsupported batched MoE wire tensor name {wire_name!r}")
+    if target_name not in params_dict:
+        raise ValueError(f"Batched MoE target parameter {target_name!r} was not found for wire tensor {wire_name!r}")
+
+    param = params_dict[target_name]
+    weight_loader = getattr(param, "weight_loader", None)
+    if weight_loader is None or not getattr(weight_loader, "supports_moe_loading", False):
+        # Layerwise reload wraps the loader with functools.wraps, which copies
+        # this marker from FusedMoE.weight_loader onto the wrapper.
+        raise ValueError(f"Parameter {target_name!r} does not expose a FusedMoE weight loader")
+
+    if param.shape[0] == loaded_weight.shape[0]:
+        success = weight_loader(
+            param,
+            loaded_weight,
+            target_name,
+            shard_id=shard_id,
+            expert_id=0,
+            return_success=True,
+        )
+        if not success:
+            raise ValueError(f"Fused loading failed for batched MoE tensor {wire_name!r}")
+        return True
+
+    # Expert-parallel vLLM keeps only a subset locally. Retain the compact wire
+    # format, but let the loader map global expert IDs one view at a time.
+    loaded_any = False
+    for expert_id, expert_weight in enumerate(loaded_weight.unbind(0)):
+        loaded_any = (
+            bool(
+                weight_loader(
+                    param,
+                    expert_weight,
+                    target_name,
+                    shard_id=shard_id,
+                    expert_id=expert_id,
+                    return_success=True,
+                )
+            )
+            or loaded_any
+        )
+    if not loaded_any:
+        raise ValueError(f"No local expert accepted batched MoE tensor {wire_name!r}")
+    return True
+
+
+def _load_checkpoint_weights(model: torch.nn.Module, weights: list[tuple[str, torch.Tensor]]) -> Any:
+    """Load ordinary HF tensors plus SkyRL's compact batched-MoE tensors."""
+    params_dict: dict[str, torch.nn.Parameter] | None = None
+    ordinary_weights: list[tuple[str, torch.Tensor]] = []
+    for name, weight in weights:
+        if name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX):
+            if params_dict is None:
+                params_dict = dict(model.named_parameters())
+            _load_batched_moe_fp8_tensor(model, params_dict, name, weight)
+        else:
+            ordinary_weights.append((name, weight))
+    if ordinary_weights:
+        return model.load_weights(weights=ordinary_weights)
+    return set()
+
+
 class NewInferenceWorkerWrap(LayerwiseReloadWorkerMixin):
     """
     vLLM worker extension for chunked weight sync (new inference path).
@@ -86,7 +207,7 @@ def update_weights_ipc(self, update_info: dict) -> None:
         handles = pickle.loads(base64.b64decode(pickled))
 
         device_index = torch.cuda.current_device()
-        physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid)
+        physical_gpu_id = cuda_uuid_to_str(torch.cuda.get_device_properties(device_index).uuid)
         if physical_gpu_id not in handles:
             raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}")
         func, args = handles[physical_gpu_id]
@@ -109,7 +230,7 @@ def update_weights_ipc(self, update_info: dict) -> None:
         model = self.model_runner.model
         with set_current_vllm_config(self.vllm_config), torch.device(self.device):
             if self._skyrl_is_checkpoint_format:
-                model.load_weights(weights=weights)
+                _load_checkpoint_weights(model, weights)
                 # vLLM's load only updates the main model; the spec-decode (MTP/Eagle)
                 # drafter is a separate module and must be reloaded from the same
                 # checkpoint-format weights (see spec_decode_utils).
@@ -165,7 +286,7 @@ def update_weights_nccl(self, update_info: dict) -> None:
 
         def _load_weights(weights):
             weights = list(weights)
-            loaded = model.load_weights(weights=weights)
+            loaded = _load_checkpoint_weights(model, weights)
             _reload_spec_decode_drafter(self.model_runner, weights)
             return loaded
 

skyrl/backends/skyrl_train/inference_servers/utils.py

@@ -11,6 +11,13 @@
     SKYRL_LORA_ADAPTER_NAME,
 )
 from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy
+from skyrl.backends.skyrl_train.weight_sync.serialized_fp8 import (
+    SERIALIZED_BLOCKWISE_FP8,
+    get_qwen35_fp8_ignored_layers,
+    get_serialized_fp8_quantization_config,
+    is_qwen35_config,
+    should_use_serialized_fp8,
+)
 from skyrl.train.config import (
     InferenceEngineConfig,
     SkyRLTrainConfig,
@@ -20,6 +27,88 @@
 logger = logging.getLogger(__name__)
 
 
+def _serialized_fp8_ignored_layers(model_path: Optional[str]) -> list[str]:
+    if not model_path:
+        raise ValueError("A model path is required when serialized FP8 weight sync is enabled")
+    try:
+        from transformers import AutoConfig
+
+        hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    except Exception as exc:
+        raise RuntimeError(
+            "Could not inspect the model config required to derive serialized FP8 ignored layers: "
+            f"model_path={model_path!r}"
+        ) from exc
+    if not is_qwen35_config(hf_config):
+        raise ValueError(
+            "Serialized FP8 weight sync currently supports only Qwen3.5 checkpoint layouts; "
+            f"model_path={model_path!r}"
+        )
+    return get_qwen35_fp8_ignored_layers(hf_config)
+
+
+def _set_or_validate(mapping: Dict[str, Any], key: str, expected: Any, *, context: str) -> None:
+    if key in mapping and mapping[key] != expected:
+        raise ValueError(
+            f"{context}.{key} must be {expected!r} when serialized FP8 weight sync is enabled, " f"got {mapping[key]!r}"
+        )
+    mapping[key] = copy.deepcopy(expected)
+
+
+def _apply_serialized_fp8_weight_sync_defaults(
+    ie_cfg: InferenceEngineConfig,
+    engine_kwargs: Dict[str, Any],
+    model_path: Optional[str] = None,
+) -> None:
+    """Configure vLLM for checkpoint-format blockwise FP8 weight reloads."""
+
+    mode = ie_cfg.fp8_weight_sync_mode
+    if mode is None:
+        return
+    if not should_use_serialized_fp8(mode):
+        raise ValueError(
+            f"Unsupported fp8_weight_sync_mode={mode!r}. " f"Supported value: {SERIALIZED_BLOCKWISE_FP8!r}."
+        )
+
+    _set_or_validate(engine_kwargs, "quantization", "fp8", context="engine_init_kwargs")
+    # Build FP8 modules without a bootstrap checkpoint; the first full-weight
+    # sync replaces the dummy values.
+    _set_or_validate(engine_kwargs, "load_format", "dummy", context="engine_init_kwargs")
+
+    hf_overrides_value = engine_kwargs.get("hf_overrides")
+    hf_overrides = {} if hf_overrides_value is None else copy.deepcopy(hf_overrides_value)
+    if not isinstance(hf_overrides, dict):
+        raise ValueError("engine_init_kwargs.hf_overrides must be a dict when serialized FP8 weight sync is enabled")
+
+    qcfg_value = hf_overrides.get("quantization_config")
+    qcfg = {} if qcfg_value is None else copy.deepcopy(qcfg_value)
+    if not isinstance(qcfg, dict):
+        raise ValueError(
+            "engine_init_kwargs.hf_overrides.quantization_config must be a dict "
+            "when serialized FP8 weight sync is enabled"
+        )
+
+    ignored_layers = _serialized_fp8_ignored_layers(model_path)
+    if ignored_layers:
+        logger.info(
+            "Serialized FP8 weight sync will leave %d vLLM modules unquantized "
+            "to match the Qwen3.5 serialized FP8 policy.",
+            len(ignored_layers),
+        )
+
+    for key, value in get_serialized_fp8_quantization_config(
+        ignored_layers=ignored_layers,
+    ).items():
+        _set_or_validate(
+            qcfg,
+            key,
+            value,
+            context="engine_init_kwargs.hf_overrides.quantization_config",
+        )
+    hf_overrides["quantization_config"] = qcfg
+    engine_kwargs["hf_overrides"] = hf_overrides
+
+
 def _uses_lora_weight_sync(cfg: SkyRLTrainConfig) -> bool:
     """Return True when the trainer syncs LoRA adapters (not merged weights).
 
@@ -153,6 +242,11 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
         logger.info(f"vLLM speculative decoding enabled: speculative_config={spec_cfg}")
 
     engine_kwargs = get_config_as_dict(ie_cfg.engine_init_kwargs)
+    _apply_serialized_fp8_weight_sync_defaults(
+        ie_cfg,
+        engine_kwargs,
+        cfg.trainer.policy.model.path,
+    )
     for key, value in engine_kwargs.items():
         setattr(args, key, value)
 

skyrl/backends/skyrl_train/weight_sync/base.py

@@ -2,7 +2,7 @@
 
 from dataclasses import asdict, dataclass, field
 from functools import cached_property
-from typing import Any, Dict, List
+from typing import Any, Dict, Iterator, List
 
 import torch
 
@@ -104,3 +104,53 @@ def total_numel(self) -> int:
     def total_size_bytes(self) -> int:
         """Calculate total memory footprint in bytes."""
         return sum(t.numel() * t.element_size() for t in self.tensors)
+
+
+def torch_dtype_name(dtype: torch.dtype) -> str:
+    """Return the dtype spelling expected by vLLM weight-transfer metadata."""
+    return str(dtype).split(".")[-1]
+
+
+def cuda_uuid_to_str(uuid: str | bytes) -> str:
+    """Normalize CUDA UUIDs identically on both sides of an IPC transfer."""
+    return uuid.decode("ascii") if isinstance(uuid, bytes) else str(uuid)
+
+
+def iter_single_dtype_chunks(chunk: WeightChunk) -> Iterator[WeightChunk]:
+    """Yield dtype-homogeneous subchunks in first-seen dtype order.
+
+    CUDA IPC packs tensors into a typed buffer, while serialized FP8 chunks mix
+    FP8 weights, FP32 scales, and unquantized BF16 tensors. vLLM's NCCL path
+    byte-packs mixed dtypes and does not require this split.
+    """
+    by_dtype: Dict[torch.dtype, Dict[str, list]] = {}
+    dtype_order: List[torch.dtype] = []
+
+    for name, tensor in zip(chunk.names, chunk.tensors):
+        dtype = tensor.dtype
+        if dtype not in by_dtype:
+            dtype_order.append(dtype)
+            by_dtype[dtype] = {"names": [], "dtypes": [], "shapes": [], "tensors": []}
+        group = by_dtype[dtype]
+        group["names"].append(name)
+        group["dtypes"].append(str(dtype))
+        group["shapes"].append(list(tensor.shape))
+        group["tensors"].append(tensor)
+
+    for dtype in dtype_order:
+        group = by_dtype[dtype]
+        yield WeightChunk(
+            names=group["names"],
+            dtypes=group["dtypes"],
+            shapes=group["shapes"],
+            tensors=group["tensors"],
+        )
+
+
+def get_weight_chunk_metadata(chunk: WeightChunk) -> Dict[str, List]:
+    """Return vLLM metadata for the tensors in a transfer chunk."""
+    return {
+        "names": list(chunk.names),
+        "dtype_names": [torch_dtype_name(tensor.dtype) for tensor in chunk.tensors],
+        "shapes": [list(tensor.shape) for tensor in chunk.tensors],
+    }