fix distributed all_to_sharded bias shard axis from -2 to -1 (#2987)

2026-01-17 22:51:42 +08:00
parent 3fe7794f22
commit d2bef3c6bb
1 changed files with 2 additions and 0 deletions
@@ -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