Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[XLA:SPMD] Remove LookaheadUserSharding in sharding propagation. #17267

Merged
merged 1 commit into from
Sep 18, 2024
Merged

[XLA:SPMD] Remove LookaheadUserSharding in sharding propagation. #17267

merged 1 commit into from
Sep 18, 2024

Commits on Sep 18, 2024

  1. [XLA:SPMD] Remove LookaheadUserSharding in sharding propagation. 

    When we infer the dot sharding from its operands, it is possible that both operands can improve the dot sharding. LookaheadUserSharding iterates the dot users and decides which dot operand sharding is preferred. This cl removes it for two reasons.

1. It is unnecessary. If we can predict the sharding from dot users, we can wait the sharding to be propagated from users. The propagted sharding from users can still help us make choice between dot operands.
2. The lookhead sharding may be wrong. LookaheadUserSharding is a heuristics. We cannot guarantee that the predicted sharding will hold in the dot users.

Reverts b4ea979

PiperOrigin-RevId: 676140549
    @ZixuanJiang @Google-ML-Automation
    ZixuanJiang authored and Google-ML-Automation committed Sep 18, 2024
    Configuration menu
    Copy the full SHA
    7b61179 View commit details
    Browse the repository at this point in the history