Open ryanrudes opened 2 years ago
Miffyli
commented
2 years ago Ideally yes, SB3 should support that device too (not a big change), but seems like it would, at the moment, require some operation-call changes to fully support. Those need to be addressed first (or wait till torch has equal functions for all platforms), but the changes should not interfere with existing code at all; this could spur up lots of hidden changes otherwise.
araffin
commented
2 years ago @ryanrudes could you test again? (there was a PyTorch release recently)
And maybe test with PyTorch nightly build, it apparently works: https://github.com/DLR-RM/rl-baselines3-zoo/issues/267
I will update the "auto" device behavior in case it does ;)
qgallouedec
commented
2 years ago from stable_baselines3 import PPO
import gym
env = gym.make("Pendulum-v1")
ppo = PPO("MlpPolicy", env, device="mps")
ppo.learn(total_timesteps=1000)With current pip version of PyTorch (1.12), it raises the following exception:
Traceback (most recent call last):
File "/Users/quentingallouedec/stable-baselines3/try_mps.py", line 7, in <module>
ppo.learn(total_timesteps=1000)
File "/Users/quentingallouedec/stable-baselines3/stable_baselines3/ppo/ppo.py", line 310, in learn
return super().learn(
File "/Users/quentingallouedec/stable-baselines3/stable_baselines3/common/on_policy_algorithm.py", line 267, in learn
self.train()
File "/Users/quentingallouedec/stable-baselines3/stable_baselines3/ppo/ppo.py", line 270, in train
loss.backward()
File "/Users/quentingallouedec/stable-baselines3/env/lib/python3.9/site-packages/torch/_tensor.py", line 396, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/Users/quentingallouedec/stable-baselines3/env/lib/python3.9/site-packages/torch/autograd/__init__.py", line 173, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
NotImplementedError: The operator 'aten::logical_and.out' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.Following the suggestion of the traceback (export PYTORCH_ENABLE_MPS_FALLBACK=1) works for me.
So it is not completely stable for the moment. We may have to wait until the next release...
araffin
commented
2 years ago thanks @qgallouedec for the feedback =)
We do need to wait for more coverage yes, issue is here: https://github.com/pytorch/pytorch/issues/77764
qgallouedec
commented
2 years ago So it is not completely stable for the moment. We may have to wait until the next release...
The bug no longer occurs with the new version of PyTorch (1.12.1)
traderpedroso
commented
2 years ago TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
MPS framework only support float32 add dtype=th.float32 on utils.py
def obs_as_tensor(
obs: Union[np.ndarray, Dict[Union[str, int], np.ndarray]], device: th.device
) -> Union[th.Tensor, TensorDict]:
"""
Moves the observation to the given device.
:param obs:
:param device: PyTorch device
:return: PyTorch tensor of the observation on a desired device.
"""
if isinstance(obs, np.ndarray):
return th.as_tensor(obs, dtype=th.float32).to(device)
elif isinstance(obs, dict):
return {key: th.as_tensor(_obs).to(device) for (key, _obs) in obs.items()}
else:
raise Exception(f"Unrecognized type of observation {type(obs)}")its fixed the problem converting float64 to float32! but after on traing model
NotImplementedError: The operator 'aten::multinomial' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variablePYTORCH_ENABLE_MPS_FALLBACK=1to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
so than add export PYTORCH_ENABLE_MPS_FALLBACK=1 and same error continues
qgallouedec
commented
2 years ago For the moment, consider that SB3 is not compatible with MPS. But we are working on it: #951
Have you seen something in the documentation about float64 and MPS? (You can answer in the PR conversation.)
traderpedroso
commented
2 years ago For the moment, consider that SB3 is not compatible with MPS. But we are working on it: #951
Have you seen something in the documentation about float64 and MPS? (You can answer in the PR conversation.)
yes MPS framework
For a more extensive list of which data types do and don’t run:
Avoid Float64 on all Apple devices. Even if the hardware supports Double physically (AMD or Intel), the Metal API doesn’t let you access it. Avoid BFloat16. That is natively supported by the latest Nvidia GPUs, but not supported in Metal. Also don’t try to use TF18/TF32 or Int4. All standard integer types (UInt8, UInt16, UInt32, UInt64) and their signed counterparts work natively on Apple devices. Not exactly 8-bit integers, which are cast to 16-bit integers before being stored into registers, but those aren’t going to harm performance. Yes, they run 64-bit integers on the Apple GPU and not 64-bit floats. Metal allows you to use 64-bit integers in shaders on AMD and Intel, but the arithmetic there might just happen through emulation (slow). I think that’s where I experienced the crash in MPSGraph previously - trying to run an operation on UInt64 on my Intel Mac mini.
https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
traderpedroso
commented
1 year ago This modified version of obs_as_tensor should work. Make these changes in the stable_baselines3/common/utils.py The modified obs_as_tensor function should now automatically convert the observation to float32 if the device is an MPS device.
Original
def obs_as_tensor(
obs: Union[np.ndarray, Dict[Union[str, int], np.ndarray]], device: th.device
) -> Union[th.Tensor, TensorDict]:
"""
Moves the observation to the given device.
:param obs:
:param device: PyTorch device
:return: PyTorch tensor of the observation on a desired device.
"""
if isinstance(obs, np.ndarray):
return th.as_tensor(obs, device=device)
elif isinstance(obs, dict):
return {key: th.as_tensor(_obs, device=device) for (key, _obs) in obs.items()}
else:
raise Exception(f"Unrecognized type of observation {type(obs)}")a workaround
def obs_as_tensor(obs: Union[np.ndarray, Dict[Union[str, int], np.ndarray]], device: th.device) -> Union[th.Tensor, TensorDict]:
"""
Moves the observation to the given device.
:param obs:
:param device: PyTorch device
:return: PyTorch tensor of the observation on a desired device.
"""
dtype = th.float32 if device.type == "mps" else None
if isinstance(obs, np.ndarray):
return th.as_tensor(obs, device=device, dtype=dtype)
elif isinstance(obs, dict):
return {key: th.as_tensor(_obs, device=device, dtype=dtype) for (key, _obs) in obs.items()}
else:
raise ValueError(f"Unsupported observation format: {obs}")Although it works normally, the CPU continues to be more performant than the MPS. Honestly, I don't know if this workaround is worth it, but it worked nonetheless.
ScharanCysne
commented
1 year ago Sorry @traderpedroso, but I don't see the difference between the workaround and the original code.
traderpedroso
commented
1 year ago Sorry @traderpedroso, but I don't see the difference between the workaround and the original code.
Apologies, I hadn't noticed that I duplicated the functions. I have now updated the code. Thank you for pointing it out.
hom-bahrani
commented
1 year ago @traderpedroso thanks for creating this issue. The M1 pro already comes with quite a lot of CPU's which distribute training nicely. I was wondering if you have done any benchmarks and observed any significant performance improvements with mps device?
traderpedroso
commented
1 year ago @traderpedroso thanks for creating this issue. The M1 pro already comes with quite a lot of CPU's which distribute training nicely. I was wondering if you have done any benchmarks and observed any significant performance improvements with
mpsdevice?
I must admit that I was profoundly disheartened by the limitations of MPs, particularly due to the lack of support. After conducting numerous tests, I discerned that for reinforcement learning, CPUs have consistently proven to be the optimal choice, or at most, TPUs. However, when it comes to leveraging GPUs from Nvidia, AMD, or MPs, their performance has been largely indistinguishable in my experience. Nevertheless, when combining PyTorch with MPs for NLP and image processing tasks, we are able to witness an exhilarating performance boost, as exemplified below.
import sys
import platform
import torch
import pandas as pd
import sklearn as sk
has_gpu = torch.cuda.is_available()
has_mps = getattr(torch,'has_mps',False)
device = "mps" if getattr(torch,'has_mps',False) \
else "cuda" if torch.cuda.is_available() else "cpu"
print(f"Python Platform: {platform.platform()}")
print(f"PyTorch Version: {torch.__version__}")
print()
print(f"Python {sys.version}")
print(f"Pandas {pd.__version__}")
print(f"Scikit-Learn {sk.__version__}")
print("GPU is", "available" if has_gpu else "NOT AVAILABLE")
print("MPS (Apple Metal) is", "AVAILABLE" if has_mps else "NOT AVAILABLE")
print(f"Target device is {device}")
import torch
from torch import nn, optim
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms
EPOCHS = 5
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4*4*50, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4*4*50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def train(model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
def main():
print("PyTorch version:", torch.__version__)
print("Torchvision version:", torchvision.__version__)
# device = torch.device("mps")
print("Using Device: ", device)
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=64, shuffle=True)
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
for epoch in range(1, EPOCHS + 1):
train(model, device, train_loader, optimizer, epoch)
test(model, device, test_loader)
if __name__ == "__main__":
main()
tty666
commented
1 year ago Same problem here really it's painful to not being able to use the mps correctly today :( At least we should do this modification :
def obs_as_tensor( obs: Union[np.ndarray, Dict[Union[str, int], np.ndarray]], device: th.device ) -> Union[th.Tensor, Dict[Union[str, int], th.Tensor]]: """ Moves the observation to the given device.
:param obs:
:param device: PyTorch device
:return: PyTorch tensor of the observation on a desired device.
"""
if isinstance(obs, np.ndarray):
# Convert the numpy array to float32 before moving it to the device
return th.as_tensor(obs.astype(np.float32), device=device)
elif isinstance(obs, dict):
# Convert the numpy arrays in the dict to float32 before moving them to the device
return {key: th.as_tensor(_obs.astype(np.float32), device=device) for (key, _obs) in obs.items()}
else:
raise Exception(f"Unrecognized type of observation {type(obs)}")To avoid error :
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
but I have no idea if the mps is able to speed up the PPO models ...
gy2256
commented
1 year ago Any updates on this feature request? Is it possible to use MPS with stable baselines3 now?
nize
commented
11 months ago Could the release of MLX play any role for improving performance of stable-baselines3 on Apple silicon? This post discusses how come MLX is not implemented within Pytorch, instead of as an alternative to Pytorch.
Unfortunately (?) it seems like stable-baselines3 would need to support use of MLX in addition to Pytorch to harvest the benefits.
araffin
commented
10 months ago Any updates on this feature request? Is it possible to use MPS with stable baselines3 now?
Please have a look at the PR and the other comments, you can give it a try using device="mps", the main bottleneck was PyTorch so far, please report any issue to the associated PR.
Could the release of MLX play any role for improving performance of stable-baselines3 on Apple silicon?
If you want to have a performance boost (not only on Apple silicon), I would recommend you to have a look at SBX (SB3 + Jax): https://github.com/araffin/sbx
🚀 Feature
PyTorch recently released support for GPU acceleration using the Apple Silicon chips. This should be supported in stable-baselines3 by the
"mps"device (I believe).Minimal Example
The Mac Silicon GPU device is not automatically recognized by stable-baselines at the moment, so it defaults to
"cpu". If you try to force it to use the"mps"device, this stack trace appears.