extend mergekit to make it work on xpu #580
Conversation
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
|
All contributors have signed the CLA ✍️ ✅ |
|
I have read the CLA Document and I hereby sign the CLA |
|
recheck |
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
yao-matrix
changed the title
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
|
@cg123 , could you pls help review? Thx very much |
| @@ -529,7 +529,7 @@ def _move_tensors( | |||
| self, value: Any, device: torch.device, non_blocking: Optional[bool] = None | |||
| ) -> Any: | |||
| if non_blocking is None: | |||
| non_blocking = device.type == "cuda" | |||
| non_blocking = device.type in ["cuda", "xpu"] | |||
| stream = torch.cuda.Stream(device=device) | ||
| with torch.cuda.stream(stream): | ||
| stream = ( | ||
| torch.Stream(device=device) |
There was a problem hiding this comment.
Maybe we should just use torch.Stream in all cases?
There was a problem hiding this comment.
we can. The only concern is torch.Stream is only available from 2.5, but mergekit's torch dependency is >= 2.0 now. So if users installs torch 2.4, it will crash. Pls let me know your insights.
There was a problem hiding this comment.
Okay, fair! This is good for now then. I'll look at it if I bump the minimum torch version in the future.
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
|
@cg123 , thx very much for you comments. I've updated per your comments, pls help review again, thx. |
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
|
@cg123 , thx very much for you comments. I've updated per your comments, pls help review again, thx. |
|
@cg123 , could you pls help review again, thx. |
|
Think this is good to go now. Thanks for your patience and the pr! |
since PyTorch 2.5, xpu has been a built-in device of PyTorch.
In this PR, I extend the accelerator device of
mergekitfrom CUDA only to Intel XPU(which is the name of Intel's GPU) and potentially to more devices, by adding a newdevicefield inMergeOption.I've passed the UT cases with
@cg123 , pls help review and let me know what I need do next step, thx very much.