admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-01 05:59:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(train): honor policy.use_amp when the policy has no dtype field (#4274)

Browse Source 
* fix(train): honor policy.use_amp when the policy has no dtype field

Since the Accelerate migration, accelerator.autocast() has been a no-op for
policies that do not expose a string dtype field (act, diffusion,
multi_task_dit, ...): PR #3912 wires mixed_precision from policy.dtype, so
these policies resolve to None and always train in full fp32 regardless of
--policy.use_amp=true. Meanwhile lerobot-eval does honor use_amp, and
PreTrainedConfig documents the flag as applying to training and evaluation.

Fall back to use_amp when no dtype field is present: bf16 where supported,
fp16 otherwise (matching torch.autocast's cuda default used by eval).

Measured on multi_task_dit / pusht, batch 256, H100: peak allocated drops
from ~75 GiB (fp32) to ~46.5 GiB under bf16 autocast.

* fix(train): conservative check

---------

Co-authored-by: Reece O'Mahoney <reece.omahoney3@gmail.com>
This commit is contained in:
Steven Palma
2026-07-31 16:46:54 +02:00
committed by GitHub
parent 7a3298ea26
commit 99443a936d
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+9 -1
View File
@@ -243,9 +243,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
# Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
force_cpu = cfg.trainable_config.device == "cpu"
# Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32/absent -> launcher default).
# Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32 -> full precision).
has_policy_dtype = hasattr(cfg.trainable_config, "dtype")
policy_dtype = getattr(cfg.trainable_config, "dtype", None)
mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype)
# Policies without a `dtype` field fall back to `use_amp`, which would otherwise be
# silently ignored here while lerobot-eval honors it. Follow torch.autocast's default
# for the configured device so training and evaluation use the same precision.
if not has_policy_dtype and getattr(cfg.trainable_config, "use_amp", False):
device_type = torch.device(cfg.trainable_config.device).type
autocast_dtype = torch.get_autocast_dtype(device_type)
mixed_precision = {torch.bfloat16: "bf16", torch.float16: "fp16"}[autocast_dtype]
accelerator = Accelerator(
step_scheduler_with_optimizer=False,
mixed_precision=mixed_precision,