feat(config): add multiprocessing option to DataLoader context and sets spawn as default (#4139)

* Add dataloader_multiprocessing_context, default to spawn

Make the DataLoader multiprocessing start method configurable on
TrainPipelineConfig and default it to 'spawn'.

The previous default (fork on Linux) is unsafe with libraries that hold
non-fork-safe state in the parent process — common ones in this codebase
are PyAV, torchcodec, and the ffmpeg shared libs they wrap. Symptoms
reported in #2488, #2209, and observed locally include:

- multiprocessing.context.AuthenticationError: digest received was wrong
- RuntimeError: Pin memory thread exited unexpectedly
- RuntimeError: DataLoader worker exited unexpectedly
- Random SIGSEGV inside worker processes during video decode

Switching to spawn re-imports modules cleanly in each worker and
eliminates these failure modes. Added the setting as a config field
rather than hard-coding so users on platforms where fork is preferred
can opt back in via --dataloader-multiprocessing-context=fork.

* Address review: shorten config comment, note spawn startup tradeoff

Per @jashshah999, mention that spawn workers re-import modules and so
add some startup time vs fork. Also trim the failure-mode dump from
the inline comment — the linked issue covers the symptoms in detail.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore(scripts): add multiprocessing_context safeguards

* chore(config): add libs note

---------

Co-authored-by: 0o8o0-blip <0o8o0-blip@users.noreply.github.com>
This commit is contained in:
Steven Palma
2026-07-28 00:42:55 +02:00
committed by GitHub
parent 95256d766d
commit 95211b98f1
2 changed files with 30 additions and 4 deletions
+18
View File
@@ -14,6 +14,7 @@
import builtins import builtins
import datetime as dt import datetime as dt
import json import json
import multiprocessing
import os import os
import tempfile import tempfile
from dataclasses import dataclass, field from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
batch_size: int = 8 batch_size: int = 8
prefetch_factor: int = 4 prefetch_factor: int = 4
persistent_workers: bool = True persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000 steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled). # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000 env_eval_freq: int = 20_000
@@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir) self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None: def validate(self) -> None:
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
and self.dataloader_multiprocessing_context not in available_contexts
):
raise ValueError(
"`dataloader_multiprocessing_context` must be None or one of "
f"{available_contexts} on this platform, got "
f"{self.dataloader_multiprocessing_context!r}."
)
self._resolve_pretrained_from_cli() self._resolve_pretrained_from_cli()
if self.policy is None and self.reward_model is None: if self.policy is None and self.reward_model is None:
+12 -4
View File
@@ -71,6 +71,16 @@ from lerobot.utils.utils import (
from .lerobot_eval import eval_policy_all from .lerobot_eval import eval_policy_all
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
}
def update_policy( def update_policy(
train_metrics: MetricsTracker, train_metrics: MetricsTracker,
policy: PreTrainedPolicy, policy: PreTrainedPolicy,
@@ -473,8 +483,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda", pin_memory=device.type == "cuda",
drop_last=False, drop_last=False,
collate_fn=collate_fn, collate_fn=collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, **_dataloader_worker_kwargs(cfg),
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
) )
# Build eval dataloader if a held-out split exists # Build eval dataloader if a held-out split exists
@@ -500,8 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda", pin_memory=device.type == "cuda",
drop_last=False, drop_last=False,
collate_fn=eval_collate_fn, collate_fn=eval_collate_fn,
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, **_dataloader_worker_kwargs(cfg),
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
) )
# Prepare everything with accelerator # Prepare everything with accelerator