-
Notifications
You must be signed in to change notification settings - Fork 160
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
Add more algebra simplify rules #1291
Conversation
18b5eea
to
aac25a3
Compare
aac25a3
to
963b64e
Compare
edd31ee
to
a939e1c
Compare
} | ||
}; | ||
|
||
struct OptimizationBarrierSimplifierPattern |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
what's the purpose of this pattern rewriter?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We will have
B = All-Gather(A)
C = Reshape(B)
D = OptimizationBarrier(C)
E = Slice(D)
graph pattern lowered from TorchAcc
, to prevent standalone Reshape
Op, which will cause extra global memory read & write, we add this rule to change this graph into
B = All-Gather(A)
C = OptimizationBarrier(B)
D = Reshape(C)
E = Slice(D)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you for the clarification. After discussing with @eedalong in the office, it appears that the reshape->optimization_barrier->slice pattern is not a commonly simplified rule. In the XLA FSDP ( implementation, this pattern serves as a workaround for TPUs. However, it is not required for NVIDIA GPUs. Perhaps we could set _shard_size_multiple=1, which would eliminate the need for this pattern. Ref: :https://github.com/AlibabaPAI/xla/blob/67edb354372a8e5cbce41f20ecfca68b328635c6/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py#L1435 .
cc @anw90
a939e1c
to
73641ea
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
No description provided.