mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +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."
|
"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(
|
p.add_argument(
|
||||||
"--subtask_chunks_per_gen",
|
"--subtask_chunks_per_gen",
|
||||||
type=int,
|
type=int,
|
||||||
@@ -461,6 +472,7 @@ def _load_policy_and_preprocessor(
|
|||||||
dataset_repo_id: str | None,
|
dataset_repo_id: str | None,
|
||||||
*,
|
*,
|
||||||
load_processors_from_checkpoint: bool = False,
|
load_processors_from_checkpoint: bool = False,
|
||||||
|
fp8: bool = False,
|
||||||
) -> tuple[Any, Any, Any, Any]:
|
) -> tuple[Any, Any, Any, Any]:
|
||||||
"""Load a policy checkpoint (local path or Hub repo id).
|
"""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):
|
if getattr(cfg, "gradient_checkpointing", False):
|
||||||
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
|
ds_meta = None
|
||||||
preprocessor = None
|
preprocessor = None
|
||||||
postprocessor = None
|
postprocessor = None
|
||||||
@@ -1529,6 +1553,7 @@ def run(
|
|||||||
args.policy_path,
|
args.policy_path,
|
||||||
args.dataset_repo_id,
|
args.dataset_repo_id,
|
||||||
load_processors_from_checkpoint=sim_mode,
|
load_processors_from_checkpoint=sim_mode,
|
||||||
|
fp8=args.fp8,
|
||||||
)
|
)
|
||||||
|
|
||||||
policy_type = getattr(policy.config, "type", None)
|
policy_type = getattr(policy.config, "type", None)
|
||||||
|
|||||||
Reference in New Issue
Block a user