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.

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

Change Manual shard_type to Identity and remove unnecessary mesh_shar… #2373

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

wooseokTT
Copy link
Contributor

@wooseokTT wooseokTT commented Mar 5, 2025

…d op

Ticket

#2337
Part of #2367

Problem description

  1. Current description of shard_type is mainly from stablehlo notations. In particular, "Manual" mode in shard_type of mesh_shard op provides some confusion in interpreting the corresponding sharded tensor. It basically indicates that the input and output shards are identical such that we don't need any further sharding operation. To clarify the term, we need to update the "Manual" shard type.
  2. Current mesh_shard ops are costly in that the operations should happen on system memory, so we would like to minimize the usage.

What's changed

  1. Replace "Manual" with "Identity" as it better describes the sharded tensor state.
  2. If frontend provides pre-sharded tensor and expects sharded return, (1) check the sharding attribute of input and return and (2) create mesh_shard op only if it is necessary to maintain shape matching purpose.

Checklist

  • New/Existing tests provide coverage for changes

@wooseokTT wooseokTT added this to the [Multi Device 1] milestone Mar 5, 2025
@wooseokTT wooseokTT changed the title Change Manual shard_type to Identical and remove replicated mesh_shar… [WIP] Change Manual shard_type to Identical and remove replicated mesh_shar… Mar 5, 2025
@wooseokTT wooseokTT force-pushed the wooseok/rename_manual_mesh_sharding branch from dc582fb to c6c985c Compare March 6, 2025 00:36
@wooseokTT wooseokTT force-pushed the wooseok/rename_manual_mesh_sharding branch from c6c985c to 17be714 Compare March 6, 2025 19:11
@wooseokTT wooseokTT changed the title [WIP] Change Manual shard_type to Identical and remove replicated mesh_shar… Change Manual shard_type to Identical and remove replicated mesh_shar… Mar 6, 2025
@wooseokTT wooseokTT changed the title Change Manual shard_type to Identical and remove replicated mesh_shar… Change Manual shard_type to Identity and remove unnecessary mesh_shar… Mar 6, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[Multi-device] Update mesh_shard shardy_type enum to avoid confusion
2 participants