From d2bef3c6bbb7acd8b311aeab5deeefc4dff3c246 Mon Sep 17 00:00:00 2001 From: gufengc Date: Sat, 17 Jan 2026 22:51:42 +0800 Subject: [PATCH] fix distributed all_to_sharded bias shard axis from -2 to -1 (#2987) --- python/mlx/nn/layers/distributed.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index 92acde8f..f33e3cbc 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -86,6 +86,8 @@ def _all_to_sharded(segments): representation becomes a sharded representation.""" def _shard_fn(path, weight): + if path.endswith("bias"): + return -1, segments return max(weight.ndim - 2, 0), segments return _shard_fn