diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index bd517d598..d3061d3b4 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -1150,9 +1150,11 @@ class PI052Policy(PreTrainedPolicy): ): return self._pi05_flow_forward(batch, reduction=reduction) - run_flow = self.config.flow_loss_weight > 0 and ( - predict_actions_t is None or bool(predict_actions_t.any().item()) - ) + # Whether any sample in the batch wants actions predicted. This is a data-dependent branch, so + # it needs a host-side bool (one CUDA sync); compute it once and reuse for both flow and FAST + # instead of syncing twice. + predict_any = predict_actions_t is None or bool(predict_actions_t.any().item()) + run_flow = self.config.flow_loss_weight > 0 and predict_any run_text = self.config.text_loss_weight > 0 and text_labels is not None loss_dict: dict[str, Any] = {} @@ -1162,7 +1164,7 @@ class PI052Policy(PreTrainedPolicy): run_fast = ( getattr(self.config, "enable_fast_action_loss", False) and self.config.fast_action_loss_weight > 0 - and (predict_actions_t is None or bool(predict_actions_t.any().item())) + and predict_any ) action_tokens = action_mask = action_code_mask = None if run_fast: @@ -1200,13 +1202,13 @@ class PI052Policy(PreTrainedPolicy): action_code_mask=action_code_mask if run_fast else None, predict_actions_t=predict_actions_t, ) - loss_dict["flow_loss"] = float(flow_loss.detach().item()) + loss_dict["flow_loss"] = flow_loss.detach() total = self.config.flow_loss_weight * flow_loss if text_loss is not None: - loss_dict["text_loss"] = float(text_loss.detach().item()) + loss_dict["text_loss"] = text_loss.detach() total = total + self.config.text_loss_weight * text_loss if fast_loss is not None: - loss_dict["fast_action_loss"] = float(fast_loss.detach().item()) + loss_dict["fast_action_loss"] = fast_loss.detach() total = total + self.config.fast_action_loss_weight * fast_loss elif run_text or run_fast: text_loss, fast_loss = self._compute_text_and_fast_loss( @@ -1218,11 +1220,11 @@ class PI052Policy(PreTrainedPolicy): predict_actions_t=predict_actions_t, ) if text_loss is not None: - loss_dict["text_loss"] = float(text_loss.detach().item()) + loss_dict["text_loss"] = text_loss.detach() weighted = self.config.text_loss_weight * text_loss total = weighted if total is None else total + weighted if fast_loss is not None: - loss_dict["fast_action_loss"] = float(fast_loss.detach().item()) + loss_dict["fast_action_loss"] = fast_loss.detach() weighted = self.config.fast_action_loss_weight * fast_loss total = weighted if total is None else total + weighted @@ -1235,7 +1237,9 @@ class PI052Policy(PreTrainedPolicy): "nothing to train." ) - loss_dict["loss"] = float(total.detach().item()) if total.dim() == 0 else float("nan") + # Keep loss components as detached tensors (no CUDA sync here); the training loop converts + # them to python floats only on logging steps (see update_policy's log_metrics gate). + loss_dict["loss"] = total.detach() if total.dim() == 0 else float("nan") if reduction == "none": return total.expand(batch[OBS_LANGUAGE_TOKENS].shape[0]), loss_dict return total, loss_dict diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index f28e8404e..e51654792 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -82,6 +82,7 @@ def update_policy( lr_scheduler=None, lock=None, sample_weighter=None, + log_metrics: bool = True, ) -> tuple[MetricsTracker, dict | None]: """ Performs a single training step to update the policy's weights. @@ -99,6 +100,9 @@ def update_policy( lr_scheduler: An optional learning rate scheduler. lock: An optional lock for thread-safe optimizer updates. sample_weighter: Optional SampleWeighter instance for per-sample loss weighting. + log_metrics: When True, read loss/grad_norm/update_s off the GPU via `.item()` (a CUDA sync). + On non-logging steps set False so the step stays fully async — letting the next batch's + dataloading and enqueue overlap GPU compute instead of stalling on a per-step sync. Returns: A tuple containing: @@ -166,12 +170,24 @@ def update_policy( if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"): accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update() - train_metrics.loss = loss.item() - train_metrics.grad_norm = grad_norm.item() train_metrics.lr = optimizer.param_groups[0]["lr"] - train_metrics.update_s = time.perf_counter() - start_time if torch.cuda.is_available(): train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3) + # `loss.item()` / `grad_norm.item()` each block on a CUDA sync. Only pay that on logging steps; + # on the other steps the readouts are never consumed, so skipping them keeps the step async and + # lets CPU-side dataloading/enqueue overlap GPU compute. update_s is only accurate under that + # sync, so it too is recorded on logging steps only. + if log_metrics: + train_metrics.loss = loss.item() + train_metrics.grad_norm = grad_norm.item() + train_metrics.update_s = time.perf_counter() - start_time + # Policies may hand back loss components as detached tensors to keep the forward async on + # non-logging steps (e.g. pi052). Materialize them to python floats here, on logging steps only, + # so the per-step CUDA sync is paid 1-in-log_freq rather than every step. + if output_dict: + output_dict = { + k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items() + } return train_metrics, output_dict @@ -642,6 +658,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time + # This update produces step number `step + 1`; only sync metrics off the GPU when that step + # will be logged (mirrors the is_log_step gate below). Everything else stays async. + log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0 + train_tracker, output_dict = update_policy( train_tracker, policy, @@ -651,6 +671,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): accelerator=accelerator, lr_scheduler=lr_scheduler, sample_weighter=sample_weighter, + log_metrics=log_metrics, ) # EMA update: pull one step of the live weights into the shadow.