Repository navigation
fix(train): pick the local training device instead of assuming cuda - #130
Open
nicolas-rabault wants to merge 1 commit into
Open
nicolas-rabault wants to merge 1 commit into
nicolas-rabault wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Copilot review overview
馃煛 Changes recommended
HF Job resumes can retain a checkpoint鈥檚 CPU device and consequently ignore an available cloud GPU.
Review effort: Balanced
Findings: 1
Open (1)
What changed in this PR
Updates training and inference to use LeRobot鈥檚 automatic device selection across CUDA, MPS, XPU, and CPU.
Changes:
- Auto-selects devices for local training and inference.
- Limits local device choices to Auto/CPU and hides them for HF Jobs.
- Updates tests and the committed frontend bundle.
| File | Description |
|---|---|
lelab/鈥媡rain.py |
Selects local devices and omits cloud device flags. |
lelab/鈥媟ollout.py |
Uses LeRobot鈥檚 shared device selection. |
frontend/鈥媠rc/鈥媝ages/鈥婽raining.tsx |
Removes the CUDA default. |
frontend/鈥媠rc/鈥媍omponents/鈥媡raining/鈥媍onfig/鈥婣dvancedCard.tsx |
Adds Auto/CPU selection for local runs. |
tests/鈥媡est_train.py |
Tests local and cloud device arguments. |
tests/鈥媡est_rollout.py |
Removes obsolete detector tests. |
frontend/鈥媎ist/鈥媔ndex.html |
References the rebuilt frontend bundle. |
馃挕 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
156
to
+158
| is_cloud = job_target is not None and job_target.runner == "hf_cloud" | ||
| if not is_cloud: | ||
| cmd.extend(["--policy.device", request.policy_device or auto_select_torch_device().type]) |
TrainingRequest and the training form defaulted policy.device to "cuda", so every local run on a Mac logged LeRobot's "Device 'cuda' is not available. Switching to 'mps'." warning. Local runs now pass the device LeRobot's auto_select_torch_device picks (cuda, mps, xpu or cpu), which inference uses too in place of its own cuda/mps/cpu probe. Accelerate ignores policy.device except to force the CPU, so the form offers "Auto" and "CPU" only. HF Jobs runs no longer send a device: submit_to_hf resets policy.device so the pod picks its own, and the form hides the field for that target.
nicolas-rabault
force-pushed
the
fix/training-device
branch
from
October 9, 2026 08:15
c103437 to
b3b6f9d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Training defaulted
policy.devicetocuda, so every local run on a Mac loggedDevice 'cuda' is not available. Switching to 'mps'.Fix
auto_select_torch_devicepicks (cuda, mps, xpu or cpu). Inference uses the same function in place of its own cuda/mps/cpu probe, so XPU hosts get XPU there too.policy.deviceexcept to force the CPU (lerobot_train.py), so the Device select offers Auto and CPU only.submit_to_hfresetspolicy.deviceso the pod picks its own, and the select is hidden for that target.Not fixed here
On HF Jobs the pod still logs
Device 'None' is not available. Switching to 'cuda'.LeRobot stagespolicy.device = Nonefor the pod andPreTrainedConfig.__post_init__warns onNone. That fix belongs in LeRobot: only warn when a requested device is unavailable.Testing
tscshows only the 5 errors already on main, eslint clean on the changed files.--policy.device mpsand no warning after.Related: #30 and #32 (the opposite case: an NVIDIA GPU with a CPU-only torch).