Any plans to support Apple Silicon GPUs?

[thomasmarsh]

I know that Torch added support for Apple silicon (device = 'mps'). I tried naively using this in a few places, but got some type mismatches between native Tensor types and MPSTensor. I thought I would check here if there were any plans to support this before I spent too much time playing with it.

[puyuan1996]

Hello,

Sorry for the late reply. We regret to inform you that we do not currently have plans to support the integration of LightZero with Metal Performance Shaders (MPS). We will keep a close eye on the development and adoption of MPS and consider providing support once its maturity has increased. If you are willing to attempt adapting LightZero to MPS and submit a Pull Request, we would be immensely appreciative.

If you encounter data type mismatch issues when attempting to use the 'mps' device for Tensor operations, here are several potential solutions (partially provided by GPT4):

• Version Compatibility: Please ensure that you are using the latest version of PyTorch that supports Apple Silicon. Over time, the support for MPS may improve.
• Data Type Conversion: Before transferring a Tensor to an MPS device, make sure its data type is compatible. In some cases, manual data type conversion may be necessary.
• Official Documentation: Refer to the latest official documentation of PyTorch for its support and recommendations for Apple Silicon. Official documentation usually contains the most updated information and solutions.
• Community Resources: Explore the PyTorch community forums or GitHub repositories to find if other users have encountered similar problems and their solutions.
• Code Updates: If your code is based on an older version of PyTorch, it may need to be updated to meet the requirements of the MPS backend.
• Patience: If the current support for MPS is not yet perfect, you may choose to wait for future versions of PyTorch which might offer better integration and support.

Should you have any questions or require further discussion, feel free to continue the conversation here. Best wishes.

[Eternity402]

I've tested mps with muzero, and found it's working.

if you set like below in lzero/entry/train_muzero.py, it works with mps

cfg.policy.device = 'mps'
policy._eval_model.to(torch.device('mps'))
policy._device='mps'

[thomasmarsh]

Thanks. That seems to have gotten it farther than my attempts. I also had to make the following change:

--- a/lzero/policy/utils.py
+++ b/lzero/policy/utils.py
@@ -614,7 +614,8 @@ def to_torch_float_tensor(data_list: Union[np.ndarray, List[np.ndarray]], device
     elif isinstance(data_list, list) and all(isinstance(data, np.ndarray) for data in data_list):
         output_data_list = []
         for data in data_list:
-            output_data_list.append(torch.from_numpy(data).to(device).float())
+            output_data_list.append(torch.from_numpy(data.astype(np.float32)).to(device))
         return output_data_list
     else:
         raise TypeError("The type of input must be np.ndarray or List[np.ndarray]")

It's probable that there's a cleaner way to do that. I'll keep playing with this. Thanks for the help.