mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-13 04:51:59 +00:00
feat(runtime): add --fp8 flag to enable PI052 FlashRT FP8 MLPs
Wire the existing (but previously unreachable from the runtime) PI052 FlashRT FP8 MLP swap into the language runtime. --fp8 sets config.use_flashrt_fp8_mlp before load; the policy calibrates and swaps every Gemma + SigLIP MLP to fused FP8 on its first predict_action_chunk. Ignored with a warning for policies without the flag (PI052 only). Measured ~1.12x faster action-chunk inference (124 -> 111 ms) on an RTX 5090; needs the `kernels` package (pin <0.13 for transformers) and CUDA SM>=8.9, else it degrades to BF16. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -383,6 +383,17 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar
|
||||
"decode and a meaningless memory line."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--fp8",
|
||||
action="store_true",
|
||||
help=(
|
||||
"PI052 only: enable FlashRT FP8 MLP kernels. Calibrates on the "
|
||||
"first inference call and swaps every Gemma + SigLIP MLP to fused "
|
||||
"FP8 (needs the `kernels` package and CUDA SM>=8.9; degrades to "
|
||||
"BF16 otherwise). Speeds up the forward pass at a small accuracy "
|
||||
"cost — see policies/pi052/flashrt_fp8.py."
|
||||
),
|
||||
)
|
||||
p.add_argument(
|
||||
"--subtask_chunks_per_gen",
|
||||
type=int,
|
||||
@@ -461,6 +472,7 @@ def _load_policy_and_preprocessor(
|
||||
dataset_repo_id: str | None,
|
||||
*,
|
||||
load_processors_from_checkpoint: bool = False,
|
||||
fp8: bool = False,
|
||||
) -> tuple[Any, Any, Any, Any]:
|
||||
"""Load a policy checkpoint (local path or Hub repo id).
|
||||
|
||||
@@ -490,6 +502,18 @@ def _load_policy_and_preprocessor(
|
||||
if getattr(cfg, "gradient_checkpointing", False):
|
||||
cfg.gradient_checkpointing = False
|
||||
|
||||
# Opt-in FP8: only PI052 has the FlashRT MLP swap. The policy calibrates and
|
||||
# swaps on its first predict_action_chunk when this flag is set.
|
||||
if fp8:
|
||||
if hasattr(cfg, "use_flashrt_fp8_mlp"):
|
||||
cfg.use_flashrt_fp8_mlp = True
|
||||
logging.info("[runtime] FP8 MLP kernels enabled (FlashRT) for %s", cfg.type)
|
||||
else:
|
||||
logging.warning(
|
||||
"[runtime] --fp8 ignored: %s has no use_flashrt_fp8_mlp (PI052 only).",
|
||||
cfg.type,
|
||||
)
|
||||
|
||||
ds_meta = None
|
||||
preprocessor = None
|
||||
postprocessor = None
|
||||
@@ -1529,6 +1553,7 @@ def run(
|
||||
args.policy_path,
|
||||
args.dataset_repo_id,
|
||||
load_processors_from_checkpoint=sim_mode,
|
||||
fp8=args.fp8,
|
||||
)
|
||||
|
||||
policy_type = getattr(policy.config, "type", None)
|
||||
|
||||
Reference in New Issue
Block a user