diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 405d619c8..dd36f4a39 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -13,6 +13,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import gc import logging import time from contextlib import nullcontext @@ -304,29 +305,27 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): ) # Compute per-timestep normalizer for relative actions - # Only main process computes, then broadcasts to avoid video decoder issues + # Compute BEFORE accelerator.prepare() to avoid NCCL timeouts during long computation relative_normalizer = None if cfg.use_relative_actions: mode = "actions + state" if cfg.use_relative_state else "actions only" if is_main_process: logging.info(colored(f"Relative mode: {mode}", "cyan", attrs=["bold"])) logging.info("Computing per-timestep stats from dataset (first 1000 batches)...") - temp_loader = torch.utils.data.DataLoader( - dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0 - ) - mean, std = compute_relative_action_stats(temp_loader, num_batches=1000) - del temp_loader + + # All ranks compute independently to avoid NCCL timeout (computation takes ~10min) + temp_loader = torch.utils.data.DataLoader( + dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=0 + ) + mean, std = compute_relative_action_stats(temp_loader, num_batches=1000) + del temp_loader + gc.collect() # Clean up any video decoder resources + relative_normalizer = PerTimestepNormalizer(mean, std) + + if is_main_process: cfg.output_dir.mkdir(parents=True, exist_ok=True) - stats_path = cfg.output_dir / "relative_stats.pt" - torch.save({"mean": mean, "std": std}, stats_path) - logging.info(f"Saved stats to: {stats_path}") - - accelerator.wait_for_everyone() - - # All ranks load from saved file - stats_path = cfg.output_dir / "relative_stats.pt" - data = torch.load(stats_path, weights_only=True, map_location="cpu") - relative_normalizer = PerTimestepNormalizer(data["mean"], data["std"]) + relative_normalizer.save(cfg.output_dir / "relative_stats.pt") + logging.info(f"Saved stats to: {cfg.output_dir / 'relative_stats.pt'}") step = 0 # number of policy updates (forward + backward + optim)