fix(train): narrow the accelerate env guard, and fix a device-bound assert (#4347)

Two post-merge CI failures on main, both from #4010.

Benchmark Integration Tests (Libero) — `accelerate launch` exports whole groups
of variables unconditionally (the five ACCELERATE_DYNAMO_* it writes default the
backend to "no"), so matching on prefixes refused launches that configure
nothing, contradicting the documented flow where accelerate is supported as a
plain launcher. The guard now watches only the three switches that hand a
subsystem to the environment.

GPU Tests — `test_metrics_tracker_reduce_across_ranks_invokes_all_reduce`
compared the captured reduction buffer against a CPU tensor, so the assert
raised "Expected all tensors to be on the same device" wherever CUDA is
available. The expected tensor is built on the buffer's device instead.

src/lerobot/distributed/factory.py

@@ -33,11 +33,7 @@ if TYPE_CHECKING:
     from lerobot.policies.pretrained import PreTrainedPolicy
 
 # Env vars through which `accelerate launch --config_file` (or a stray shell) would configure
-# accelerate behind the config system's back. Plugin `__post_init__`s read these silently as
-# field fallbacks (ACCELERATE_DYNAMO_* enables torch.compile through the default
-# TorchDynamoPlugin; ACCELERATE_GRADIENT_ACCUMULATION_STEPS overrides the explicitly passed
-# value inside Accelerator.__init__), which would make train_config.json lie about what ran.
-_ACCELERATE_ENV_PREFIXES = ("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_DYNAMO_")
+# accelerate behind the config system's back, making train_config.json lie about what ran.
 _ACCELERATE_ENV_VARS = (
     "ACCELERATE_USE_FSDP",
     "ACCELERATE_USE_PARALLELISM_CONFIG",
@@ -59,11 +55,7 @@ def guard_against_env_interference() -> None:
     """
     if os.environ.get(_ENV_OVERRIDE):
         return
-    offending = sorted(
-        name
-        for name in os.environ
-        if name in _ACCELERATE_ENV_VARS or name.startswith(_ACCELERATE_ENV_PREFIXES)
-    )
+    offending = sorted(name for name in _ACCELERATE_ENV_VARS if name in os.environ)
     if offending:
         raise RuntimeError(
             f"Accelerate-configuring environment variables are set: {', '.join(offending)}. "

tests/distributed/test_parallel_dims_and_factory.py

@@ -79,13 +79,10 @@ class TestParallelDims:
 
 
 class TestEnvGuard:
-    # ACCELERATE_DYNAMO_*/ACCELERATE_GRADIENT_ACCUMULATION_STEPS are silent config overrides
-    # inside accelerate itself — the guard must catch them too.
+    # Silent config overrides inside accelerate itself — the guard must catch them.
     _POISON = (
         "ACCELERATE_USE_FSDP",
-        "FSDP_VERSION",
-        "PARALLELISM_CONFIG_DP_SHARD_SIZE",
-        "ACCELERATE_DYNAMO_BACKEND",
+        "ACCELERATE_USE_PARALLELISM_CONFIG",
         "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
     )
 
@@ -102,7 +99,7 @@ class TestEnvGuard:
             guard_against_env_interference()
 
     def test_override_acknowledges(self, monkeypatch):
-        monkeypatch.setenv("FSDP_VERSION", "2")
+        monkeypatch.setenv("ACCELERATE_USE_FSDP", "true")
         monkeypatch.setenv(_ENV_OVERRIDE, "1")
         guard_against_env_interference()
 

tests/utils/test_logging_utils.py

@@ -197,7 +197,7 @@ def test_metrics_tracker_reduce_across_ranks_invokes_all_reduce(monkeypatch):
     tracker.reduce_across_ranks()
 
     assert captured["op"] == logging_utils.dist.ReduceOp.MAX
-    assert torch.allclose(captured["values"], torch.tensor([0.4]))
+    assert torch.allclose(captured["values"], torch.tensor([0.4], device=captured["values"].device))
     assert tracker.update_s.avg == pytest.approx(0.9)
     # Metrics without a reduction stay untouched.
     assert tracker.loss.avg == 1.0