diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index ab27edd47..b47117bb8 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -206,6 +206,15 @@ class PI052Config(PI05Config): # at -4.5% step time on H100 (bench job 22161421); peak memory # unchanged. ``fused_linear_cross_entropy`` ships separately via # ``_shifted_lin_ce`` / ``_fast_lin_ce``. + use_flashrt_fp8_mlp: bool = False + """Opt-in: swap every Gemma GeGLU MLP (action expert + prefix LM) and the + SigLIP vision MLP to FlashRT fused FP8 kernels (Hugging Face Kernel Hub + ``flashrt/*``). The swap needs a one-time activation calibration on a real + observation, so it is applied explicitly via + ``PI052Policy.apply_flashrt_fp8_mlp(batch)`` after loading (not at build). + Degrades gracefully to BF16 if ``kernels`` / the FlashRT packages are + missing. Default off keeps behaviour identical to the BF16 path.""" + use_hf_kernels: bool = True """Deprecated. Liger HF kernels are patched unconditionally by ``_enable_hf_kernels`` — this field is retained as a no-op for diff --git a/src/lerobot/policies/pi052/flashrt_fp8.py b/src/lerobot/policies/pi052/flashrt_fp8.py new file mode 100644 index 000000000..e4496f9e0 --- /dev/null +++ b/src/lerobot/policies/pi052/flashrt_fp8.py @@ -0,0 +1,272 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Optional FP8 MLP swap for PI052 using the FlashRT Hugging Face Kernel Hub. + +Replaces every Gemma GeGLU MLP (action expert + prefix language model) with the +fused ``fp8_geglu_mlp_bf16`` kernel and the SigLIP vision-tower MLP with +``fp8_gelu_mlp_bf16``. Static activation scales are calibrated once on a real +observation; weights are quantized once. This is opt-in and degrades gracefully +to the BF16 path if ``kernels`` or the FlashRT packages are unavailable. + +Use: + policy = PI052Policy.from_pretrained(...) + batch = preprocessor(observation) # one representative observation + policy.apply_flashrt_fp8_mlp(batch) # calibrate + swap in place +""" + +from __future__ import annotations + +import logging + +import torch +import torch.nn as nn +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +_FP8_MAX = 448.0 + + +def _roundtrip_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Quantize->dequantize an activation through FP8 E4M3 at ``scale`` (f32).""" + q = torch.clamp(x.float() / scale.float(), -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + return q.float() * scale.float() +_SWIGLU_REPO = "flashrt/flashrt-fp8-swiglu-ffn" +_GELU_REPO = "flashrt/flashrt-fp8-ffn" +_GEMM_REPO = "flashrt/flashrt-gemm-epilogues" + + +def _get_kernel(repo: str): + """Load a FlashRT Hub package (cached). Returns None if unavailable.""" + from kernels import get_kernel + + return get_kernel(repo, version=1) + + +def _quantize_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = max(weight.detach().float().abs().max().item(), 1e-12) / _FP8_MAX + fp8 = torch.clamp(weight.float() / scale, -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + return fp8.contiguous(), torch.tensor([scale], dtype=torch.float32) + + +def _static_scale(amax: float, safety: float) -> torch.Tensor: + return torch.tensor([max(amax, 1e-12) / _FP8_MAX * safety], dtype=torch.float32) + + +class _FlashRTGeGLU(nn.Module): + """FP8 drop-in for a Gemma GeGLU MLP (gate/up/down, gelu_pytorch_tanh, no bias).""" + + def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety, fuse_weight=None): + super().__init__() + self.ffn_ops = ffn_ops + self.quant_ops = quant_ops + self.in_features = mlp.gate_proj.weight.shape[1] + device = mlp.gate_proj.weight.device + gate_up = torch.cat([mlp.gate_proj.weight, mlp.up_proj.weight], dim=0).float() + # Fold the preceding RMSNorm weight (1 + w) into the gate/up GEMM and feed + # the kernel channel_scale = 1/(1+w). This is exact (it just moves the + # per-channel (1+w) from the activation to the weight) and is what keeps + # FP8 accurate: the normed activation rms(x) is uniform, while + # rms(x)*(1+w) has per-channel outliers that per-tensor FP8 quantizes + # poorly. Mirrors the FlashRT runtime (norm runs with ones, weight folds + # 1+w). Only the fixed-weight (non-adaptive) norms fold; adaptive-RMSNorm + # layers pass fuse_weight=None (channel_scale = ones). + if fuse_weight is not None: + f = (1.0 + fuse_weight.detach().float()) + gate_up = gate_up * f[None, :] + channel_scale = (1.0 / f).to(torch.bfloat16) + else: + channel_scale = torch.ones(self.in_features, dtype=torch.bfloat16) + gate_up_fp8, gate_up_scale = _quantize_fp8(gate_up) + down_fp8, down_scale = _quantize_fp8(mlp.down_proj.weight) + self.register_buffer("gate_up_fp8", gate_up_fp8.to(device)) + self.register_buffer("down_fp8", down_fp8.to(device)) + self.register_buffer("gate_up_scale", gate_up_scale.to(device)) + self.register_buffer("down_scale", down_scale.to(device)) + self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device)) + self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device)) + self.register_buffer("channel_scale", channel_scale.to(device)) + self.safety = safety + self.calibrating = False + self._ia = 0.0 + self._ha = 0.0 + + def _calibrate_step(self, x): + # FP8-propagated calibration (FlashRT contract): measure input/hidden amax + # on the live (already-FP8-upstream) activation, running-max across steps. + flat = x.reshape(-1, self.in_features).to(torch.bfloat16) + xq = flat.float() * self.channel_scale.float() + self._ia = max(self._ia, xq.abs().max().item()) + self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device)) + xdq = _roundtrip_fp8(xq, self.input_scale) + wdq = self.gate_up_fp8.float() * self.gate_up_scale.float() + gate, up = (xdq @ wdq.t()).chunk(2, dim=-1) + hidden = F.gelu(gate, approximate="tanh") * up + self._ha = max(self._ha, hidden.abs().max().item()) + self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device)) + + def forward(self, x): + if self.calibrating: + self._calibrate_step(x) + shape = x.shape + flat = x.reshape(-1, self.in_features).to(torch.bfloat16) + x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale) + out = self.ffn_ops.fp8_geglu_mlp_bf16( + x_fp8, self.gate_up_fp8, self.down_fp8, + self.input_scale, self.gate_up_scale, self.hidden_scale, self.down_scale, + ) + return out.reshape(shape) + + +class _FlashRTGeluMLP(nn.Module): + """FP8 drop-in for a SigLIP MLP (fc1 -> gelu_tanh -> fc2, with bias).""" + + def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety): + super().__init__() + self.ffn_ops = ffn_ops + self.quant_ops = quant_ops + self.in_features = mlp.fc1.weight.shape[1] + self.out_features = mlp.fc2.weight.shape[0] + device = mlp.fc1.weight.device + up_fp8, up_scale = _quantize_fp8(mlp.fc1.weight) + down_fp8, down_scale = _quantize_fp8(mlp.fc2.weight) + self.register_buffer("up_fp8", up_fp8.to(device)) + self.register_buffer("down_fp8", down_fp8.to(device)) + self.register_buffer("up_scale", up_scale.to(device)) + self.register_buffer("down_scale", down_scale.to(device)) + self.register_buffer("up_bias", mlp.fc1.bias.detach().to(torch.bfloat16)) + self.register_buffer("down_bias", mlp.fc2.bias.detach().to(torch.bfloat16)) + self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device)) + self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device)) + self.register_buffer("channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16)) + self.safety = safety + self.calibrating = False + self._ia = 0.0 + self._ha = 0.0 + + def _calibrate_step(self, x): + flat = x.reshape(-1, self.in_features).to(torch.bfloat16) + self._ia = max(self._ia, flat.float().abs().max().item()) + self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device)) + xdq = _roundtrip_fp8(flat.float(), self.input_scale) + hid = (xdq @ (self.up_fp8.float() * self.up_scale.float()).t()) + self.up_bias.float() + hid = F.gelu(hid, approximate="tanh") + self._ha = max(self._ha, hid.abs().max().item()) + self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device)) + + def forward(self, x): + if self.calibrating: + self._calibrate_step(x) + shape = x.shape + dtype = x.dtype + flat = x.reshape(-1, self.in_features).to(torch.bfloat16) + x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale) + out = self.ffn_ops.fp8_gelu_mlp_bf16( + x_fp8, self.up_fp8, self.up_bias, self.down_fp8, self.down_bias, + self.input_scale, self.up_scale, self.hidden_scale, self.down_scale, + ) + return out.reshape(*shape[:-1], self.out_features).to(dtype) + + +def _siglip_mlps(model) -> list: + tower = model.paligemma_with_expert.paligemma.model.vision_tower + return [m for _, m in tower.named_modules() if type(m).__name__ == "SiglipMLP"] + + +def _run_forward(policy, batches) -> None: + """Run predict_action_chunk in eager mode (a compiled ``sample_actions`` + bypasses the Python module forwards, so drop it for the calibration pass).""" + model = policy.model + saved = {name: vars(model).pop(name) for name in ("sample_actions", "forward") if name in vars(model)} + with torch.inference_mode(): + for batch in batches: + policy.predict_action_chunk({k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()}) + torch.cuda.synchronize() + vars(model).update(saved) + + +def _fixed_norm_weight(norm): + """RMSNorm (1+w) fold weight if ``norm`` is fixed-weight; None if adaptive.""" + return norm.weight if getattr(norm, "dense", None) is None else None + + +def _fp8_supported(device) -> bool: + """FP8 E4M3 tensor cores require CUDA SM >= 8.9 (Ada / Hopper / Blackwell). + Older GPUs (e.g. A100 SM 8.0) and CPU have no FP8 path, so the kernels would + fail at runtime — gate here and keep BF16.""" + if device.type != "cuda" or not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability(device) + return (major, minor) >= (8, 9) + + +def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool: + """Swap every Gemma GeGLU MLP and SigLIP GELU MLP to the FlashRT FP8 kernels. + + Calibration is FP8-propagated (the FlashRT contract): the FP8 modules are + swapped in first, then a single forward on one representative frame measures + each GEMM's input/hidden amax on the *already-quantized* activations it sees + at runtime (not on a clean BF16 forward), running-max across denoise steps. + The preceding fixed RMSNorm weight (1+w) is folded into the GEMM so the + quantized activation is the uniform rms(x); adaptive-RMSNorm inputs (action + expert) do not fold. ``scale = amax/448 * safety``. A single frame is enough. + Returns True on success, False (no-op, BF16 kept) if the device lacks FP8 + support or the kernels are unavailable. + """ + device = next(policy.parameters()).device + if not _fp8_supported(device): + logger.warning( + "PI052: device %s has no FP8 (E4M3) support (needs CUDA SM>=8.9); keeping BF16.", + device, + ) + return False + batches = batch if isinstance(batch, (list, tuple)) else [batch] + try: + ffn_ops = _get_kernel(_SWIGLU_REPO) + gelu_ops = _get_kernel(_GELU_REPO) + quant_ops = _get_kernel(_GEMM_REPO) + except Exception as exc: # noqa: BLE001 + logger.warning("PI052: FlashRT FP8 kernels unavailable (%s); keeping BF16.", exc) + return False + + model = policy.model + calibrating = [] + + gemma_layers = ( + list(model.paligemma_with_expert.gemma_expert.model.layers) + + list(model.paligemma_with_expert.paligemma.model.language_model.layers) + ) + for layer in gemma_layers: + fw = _fixed_norm_weight(layer.post_attention_layernorm) + layer.mlp = _FlashRTGeGLU(layer.mlp, 1.0, 1.0, ffn_ops, quant_ops, safety, fuse_weight=fw).to(device) + calibrating.append(layer.mlp) + + siglip = _siglip_mlps(model) + for mlp_parent in model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers: + mlp_parent.mlp = _FlashRTGeluMLP(mlp_parent.mlp, 1.0, 1.0, gelu_ops, quant_ops, safety).to(device) + calibrating.append(mlp_parent.mlp) + + # FP8-propagated calibration: one forward with every module in calibrate mode. + for m in calibrating: + m.calibrating = True + _run_forward(policy, batches) + for m in calibrating: + m.calibrating = False + + logger.info( + "PI052: FlashRT FP8 enabled (%d Gemma + %d SigLIP MLPs).", + len(gemma_layers), len(siglip), + ) + return True diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index 38358de30..0161e588b 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -38,7 +38,6 @@ for the LM head. from __future__ import annotations import builtins -import copy import logging import math import types @@ -416,8 +415,13 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` embs = torch.cat(embs, dim=1) pad_masks = torch.cat(pad_masks, dim=1) - att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device) - att_masks = att_masks[None, :].expand(bsize, len(att_masks)) + # The suffix mask is the constant [1, 0, ..., 0]; build it on-device + # rather than via torch.tensor(python_list, device=cuda), which is a + # host->device sync on every denoise step. + n = len(att_masks) + att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device) + att_masks[0] = 1 + att_masks = att_masks[None, :].expand(bsize, n) return embs, pad_masks, att_masks, adarms_cond @@ -517,10 +521,17 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` dt = -1.0 / num_steps + # Precompute the whole timestep schedule on-device once, instead of + # rebuilding a tensor from a Python float every step + # (``torch.tensor(time, device=cuda)`` is a host->device sync ×num_steps). + times = torch.tensor( + [1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device + ) + x_t = noise for step in range(num_steps): - time = 1.0 + step * dt - time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize) + time = 1.0 + step * dt # Python float kept for the RTC branch below + time_tensor = times[step].expand(bsize) def denoise_step_partial_call(input_x_t, current_timestep=time_tensor): return self.denoise_step( @@ -579,7 +590,12 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` ) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 - past_key_values = copy.deepcopy(past_key_values) + # The expert forward appends the suffix K/V to the prefix cache in-place + # (GemmaAttention.update runs even with use_cache=False), so each step + # must start from a prefix-only cache. Instead of deep-copying the whole + # cache every step, let it append and crop back to the prefix length + # afterwards (the prefix K/V are read-only, so this is exact and keeps + # the loop a single graph). outputs_embeds, _ = self.paligemma_with_expert.forward( attention_mask=full_att_2d_masks_4d, position_ids=position_ids, @@ -588,6 +604,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` use_cache=False, adarms_cond=[None, adarms_cond], ) + past_key_values.crop(prefix_len) suffix_out = outputs_embeds[1] suffix_out = suffix_out[:, -self.config.chunk_size :] @@ -1063,6 +1080,18 @@ class PI052Policy(PreTrainedPolicy): # subtask can be held across several chunks (see subtask_replan_steps). self._subtask_chunk_counter = 0 + def apply_flashrt_fp8_mlp(self, batch: dict[str, Tensor], *, safety: float = 1.05) -> bool: + """Opt-in: swap every Gemma + SigLIP MLP to FlashRT fused FP8 kernels. + + Calibrates static activation scales once on ``batch`` (one representative + observation, already through the preprocessor) and swaps the MLP modules + in place. Returns False (no-op, BF16 kept) if the kernels are missing. + Gated by ``config.use_flashrt_fp8_mlp`` — see flashrt_fp8.py. + """ + from .flashrt_fp8 import apply_fp8_mlp # noqa: PLC0415 + + return apply_fp8_mlp(self, batch, safety=safety) + # ------------------------------------------------------------------ # Head unfreeze helper # ------------------------------------------------------------------ @@ -2360,6 +2389,13 @@ class PI052Policy(PreTrainedPolicy): """Predict a chunk of actions given environment observations.""" self.eval() + # Opt-in FlashRT FP8: calibrate static scales on the first real observation + # and swap the MLPs in place. Guard set before the call so the calibration + # forward (which re-enters predict_action_chunk) does not recurse. + if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False): + self._fp8_applied = True + self.apply_flashrt_fp8_mlp(batch) + # Prepare inputs images, img_masks = self._preprocess_images(batch) tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]