Skip to content

Fix sharded-to-all bias and integer split positions on quantized layers - #4715

Open
ronaldmannak wants to merge 3 commits into
ml-explore:mainfrom
PicoMLX:bias-quantized-integer
Open

ronaldmannak wants to merge 3 commits into
ml-explore:mainfrom
PicoMLX:bias-quantized-integer

Conversation

@ronaldmannak

@ronaldmannak ronaldmannak commented Oct 17, 2026 •

Copy link
Copy Markdown
Contributor

This PR has three fixes:

  • ShardedToAllLinear ranks start with a different random bias, and the bias is added after the sum, resulting in different values being returned. Bias now starts at zeros.
  • Integer split positions on QuantizedShardedToAllLinear can split in wrong places. Fix converts integer positions to exact fractions of the input size, raises ValueError when segment can't be split correctly
  • Same fix for shard_inplace for quantized layers

The first two fixes are already fixed in MLX-Swift. I'll open a separate PR for the third issue on MLX-Swift.

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: Issues found and fixed by Opus 5.5 while porting distributed to MLX-Swift. Code reviewed by Codex

@ronaldmannak ronaldmannak changed the title Initialize ShardedToAllLinear bias to zeros and integer split position in QuantizedShardedToAllLinear Fix sharded-to-all bias and integer split positions on quantized layers Oct 17, 2026
@zcbenz zcbenz added low priority await verification This pull request is non-trivial and requires a human expert to verify its correctness. labels Oct 19, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness. low priority

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants