Movatterモバイル変換


[0]ホーム

URL:


Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Introduces a new utility function,MatchPermutedSliceAndPartitionOffset, to detect a pattern where aDynamicSlice consumes the output of anAllGather with a permuted set of offsets. This pattern is equivalent to aCollectivePermute and can be optimized accordingly.#97189

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 ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Draft
copybara-service wants to merge1 commit intomaster
base:master
Choose a base branch
Loading
fromexported_pr_783030292

Conversation

copybara-service[bot]
Copy link

Introduces a new utility function,MatchPermutedSliceAndPartitionOffset, to detect a pattern where aDynamicSlice consumes the output of anAllGather with a permuted set of offsets. This pattern is equivalent to aCollectivePermute and can be optimized accordingly.

The logic is divided into four main sections:

  1. Initial Checks: Verifies that theAllGather is a suitable candidate for this optimization. It ensures the operation is performed across multiple partitions with a single replica and uses flattened IDs (i.e.,use_global_device_ids is enabled).
  2. Locate DynamicSlice: Finds theDynamicSlice user of theAllGather, correctly traversing through any interveningReshape orBitcast operations that do not alter the data.
  3. Extract Offset Specifications: Derives the memory offset-to-partition ID mappings from both theAllGather (the source layout) and theDynamicSlice (the destination/permuted layout).
  4. Match Specifications: Compares the source and destination offset maps to derive the permutation. For each memory offset, it pairs the source partition ID from theAllGather with the destination partition ID from theDynamicSlice.

This CL also introduces a few key data structures to support this logic:

  • PartitionOffsetSpec: Represents the mapping from a memory offset to a partition ID for each replica group. This is used to model the data layout produced by theAllGather and the permuted access pattern of theDynamicSlice.
  • PermutationPairs: A type alias for a list of(source_id, destination_id) pairs, representing a permute operation for CP.

…set`, to detect a pattern where a `DynamicSlice` consumes the output of an `AllGather` with a permuted set of offsets. This pattern is equivalent to a `CollectivePermute` and can be optimized accordingly.The logic is divided into four main sections:1.  **Initial Checks:** Verifies that the `AllGather` is a suitable candidate for this optimization. It ensures the operation is performed across multiple partitions with a single replica and uses flattened IDs (i.e., `use_global_device_ids` is enabled).2.  **Locate DynamicSlice:** Finds the `DynamicSlice` user of the `AllGather`, correctly traversing through any intervening `Reshape` or `Bitcast` operations that do not alter the data.3.  **Extract Offset Specifications:** Derives the memory offset-to-partition ID mappings from both the `AllGather` (the source layout) and the `DynamicSlice` (the destination/permuted layout).4.  **Match Specifications:** Compares the source and destination offset maps to derive the permutation. For each memory offset, it pairs the source partition ID from the `AllGather` with the destination partition ID from the `DynamicSlice`.This CL also introduces a few key data structures to support this logic:*   **`PartitionOffsetSpec`**: Represents the mapping from a memory offset to a partition ID for each replica group. This is used to model the data layout produced by the `AllGather` and the permuted access pattern of the `DynamicSlice`.*   **`PermutationPairs`**: A type alias for a list of `(source_id, destination_id)` pairs, representing a permute operation for CP.PiperOrigin-RevId: 783030292
Sign up for freeto join this conversation on GitHub. Already have an account?Sign in to comment
Reviewers
No reviews
Assignees

@felixwqpfelixwqp

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

1 participant
@felixwqp

[8]ページ先頭

©2009-2025 Movatter.jp