Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion src/lightning/fabric/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
_CLICK_AVAILABLE = RequirementCache("click")
_LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")

_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu")
_SUPPORTED_ACCELERATORS = ("cpu", "gpu", "cuda", "mps", "tpu", "auto")


def _get_supported_strategies() -> list[str]:
Expand Down Expand Up @@ -187,6 +187,19 @@ def _set_env_variables(args: Namespace) -> None:

def _get_num_processes(accelerator: str, devices: str) -> int:
"""Parse the `devices` argument to determine how many processes need to be launched on the current machine."""
if accelerator == "auto" or accelerator is None:
if torch.cuda.is_available() and torch.cuda.device_count() > 0:
accelerator = "cuda"
elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
accelerator = "mps"
else:
accelerator = "cpu"

if devices == "auto":
if accelerator == "cuda" and torch.cuda.device_count() > 0 or accelerator == "mps" or accelerator == "cpu":
devices = "1"
else:
raise ValueError(f"Cannot default to '1' device for accelerator='{accelerator}'")
if accelerator == "gpu":
parsed_devices = _parse_gpu_ids(devices, include_cuda=True, include_mps=True)
elif accelerator == "cuda":
Expand Down
2 changes: 1 addition & 1 deletion tests/tests_fabric/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def test_run_env_vars_defaults(monkeypatch, fake_script):
assert "LT_PRECISION" not in os.environ


@pytest.mark.parametrize("accelerator", ["cpu", "gpu", "cuda", pytest.param("mps", marks=RunIf(mps=True))])
@pytest.mark.parametrize("accelerator", ["cpu", "gpu", "cuda", "auto", pytest.param("mps", marks=RunIf(mps=True))])
@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
@mock.patch("lightning.fabric.accelerators.cuda.num_cuda_devices", return_value=2)
def test_run_env_vars_accelerator(_, accelerator, monkeypatch, fake_script):
Expand Down
Loading