From 99443a936d73ad3c7cb62b126ef36e9b0045b154 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 31 Jul 2026 16:46:54 +0200 Subject: [PATCH] fix(train): honor policy.use_amp when the policy has no dtype field (#4274) * 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 --- src/lerobot/scripts/lerobot_train.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 178a68cfa..8a3222df1 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -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,