diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index 34be2ec600..d8fca8b160 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -148,12 +148,11 @@ get_backends() { } Group register_group(std::shared_ptr group, std::string bk) { - auto& backends = get_backends(); if (group == nullptr) { - group = std::make_shared(); - } else { - backends.insert({"any", group}); + return Group(std::make_shared()); } + auto& backends = get_backends(); + backends.insert({"any", group}); backends.insert({std::move(bk), group}); return Group(group); } diff --git a/python/tests/ring_test_distributed.py b/python/tests/ring_test_distributed.py index 31db2ed5f7..8b4de9ccae 100644 --- a/python/tests/ring_test_distributed.py +++ b/python/tests/ring_test_distributed.py @@ -23,6 +23,16 @@ def test_groups(self): with self.assertRaises(RuntimeError): sub = world.split(world.rank() % 2) + def test_strict_init_after_fallback(self): + with mlx_tests.scoped_env( + JACCL_IBV_DEVICES=None, + MLX_IBV_DEVICES=None, + ): + self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) + + with self.assertRaises(RuntimeError): + mx.distributed.init(backend="jaccl", strict=True) + def test_all_reduce_extra(self): world = mx.distributed.init() dtypes = [