Repository: ray-project/ray

Issue:
[rllib] num_gpus=0 on a device with GPU uses GPU for Pytorch
### What is the problem?

Ray will find a GPU and place the model (e.g. FCNet) on the GPU even when `num_gpus=0`.

Stack Trace:

```
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/trainer.py:520: in train
    raise e
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/trainer.py:506: in train
    result = Trainable.train(self)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/tune/trainable.py:260: in train
    result = self._train()
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py:139: in _train
    return self._train_exec_impl()
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py:177: in _train_exec_impl
    res = next(self.train_exec_impl)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/util/iter.py:731: in __next__
    return next(self.built_iterator)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/util/iter.py:752: in apply_foreach
    result = fn(item)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/maml/maml.py:143: in __call__
    fetches = self.workers.local_worker().learn_on_batch(samples)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py:742: in learn_on_batch
    .learn_on_batch(samples)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/policy/torch_policy.py:244: in learn_on_batch
    self._loss(self, self.model, self.dist_class, train_batch))
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/agents/maml/maml_torch_policy.py:329: in maml_loss
    logits, state = model.from_batch(train_batch)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/models/modelv2.py:224: in from_batch
    return self.__call__(input_dict, states, train_batch.get("seq_lens"))
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/models/modelv2.py:181: in __call__
    res = self.forward(restored, state or [], seq_lens)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/models/torch/fcnet.py:118: in forward
    self._features = self._hidden_layers(self._last_flat_in)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/module.py:550: in __call__
    result = self.forward(*input, **kwargs)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/container.py:100: in forward
    input = module(input)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/module.py:550: in __call__
    result = self.forward(*input, **kwargs)
../../miniconda3/envs/ray/lib/python3.7/site-packages/ray/rllib/models/torch/misc.py:110: in forward
    return self._model(x)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/module.py:550: in __call__
    result = self.forward(*input, **kwargs)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/container.py:100: in forward
    input = module(input)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/module.py:550: in __call__
    result = self.forward(*input, **kwargs)
../../miniconda3/envs/ray/lib/python3.7/site-packages/torch/nn/modules/linear.py:87: in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mat1' in call to _th_addmm
```
### Reproduction (REQUIRED)
Run this in `ray/rlllib`:
`pytest -v -s agents/maml/tests/test_maml.py` on a machine with a GPU.

- [x] I have verified my script runs in a clean environment and reproduces the issue.
- [x] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).



Discussion:
@michaelzhiluo Looking at this now ...

Resolving patch:
diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py
--- a/rllib/evaluation/rollout_worker.py
+++ b/rllib/evaluation/rollout_worker.py
@@ -9,9 +9,6 @@
     TYPE_CHECKING, TypeVar
 
 import ray
-from ray.util.debug import log_once, disable_log_once_globally, \
-    enable_periodic_logging
-from ray.util.iter import ParallelIteratorWorker
 from ray.rllib.env.atari_wrappers import wrap_deepmind, is_atari
 from ray.rllib.env.base_env import BaseEnv
 from ray.rllib.env.env_context import EnvContext
@@ -21,17 +18,17 @@
 from ray.rllib.env.vector_env import VectorEnv
 from ray.rllib.evaluation.sampler import AsyncSampler, SyncSampler
 from ray.rllib.evaluation.rollout_metrics import RolloutMetrics
-from ray.rllib.policy.sample_batch import MultiAgentBatch, DEFAULT_POLICY_ID
-from ray.rllib.policy.policy import Policy
-from ray.rllib.policy.tf_policy import TFPolicy
-from ray.rllib.policy.torch_policy import TorchPolicy
+from ray.rllib.models import ModelCatalog
+from ray.rllib.models.preprocessors import NoPreprocessor, Preprocessor
 from ray.rllib.offline import NoopOutput, IOContext, OutputWriter, InputReader
 from ray.rllib.offline.off_policy_estimator import OffPolicyEstimator, \
     OffPolicyEstimate
 from ray.rllib.offline.is_estimator import ImportanceSamplingEstimator
 from ray.rllib.offline.wis_estimator import WeightedImportanceSamplingEstimator
-from ray.rllib.models import ModelCatalog
-from ray.rllib.models.preprocessors import NoPreprocessor, Preprocessor
+from ray.rllib.policy.sample_batch import MultiAgentBatch, DEFAULT_POLICY_ID
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.tf_policy import TFPolicy
+from ray.rllib.policy.torch_policy import TorchPolicy
 from ray.rllib.utils import merge_dicts
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.utils.debug import summarize
@@ -42,6 +39,9 @@
 from ray.rllib.utils.types import EnvType, AgentID, PolicyID, EnvConfigDict, \
     ModelConfigDict, TrainerConfigDict, SampleBatchType, ModelWeights, \
     ModelGradients, MultiAgentPolicyConfigDict
+from ray.util.debug import log_once, disable_log_once_globally, \
+    enable_periodic_logging
+from ray.util.iter import ParallelIteratorWorker
 
 if TYPE_CHECKING:
     from ray.rllib.agents.callbacks import DefaultCallbacks
@@ -399,22 +399,30 @@ def make_env(vector_index):
                         tf1.set_random_seed(seed)
                     self.policy_map, self.preprocessors = \
                         self._build_policy_map(policy_dict, policy_config)
-            if (ray.is_initialized()
-                    and ray.worker._mode() != ray.worker.LOCAL_MODE):
-                if not ray.get_gpu_ids():
-                    logger.debug(
-                        "Creating policy evaluation worker {}".format(
-                            worker_index) +
-                        " on CPU (please ignore any CUDA init errors)")
-                elif not tf1.test.is_gpu_available():
-                    raise RuntimeError(
-                        "GPUs were assigned to this worker by Ray, but "
-                        "TensorFlow reports GPU acceleration is disabled. "
-                        "This could be due to a bad CUDA or TF installation.")
         else:
             self.policy_map, self.preprocessors = self._build_policy_map(
                 policy_dict, policy_config)
 
+        if (ray.is_initialized() and
+                ray.worker._mode() != ray.worker.LOCAL_MODE):
+            # Check available number of GPUs
+            if not ray.get_gpu_ids():
+                logger.debug(
+                    "Creating policy evaluation worker {}".format(
+                        worker_index) +
+                    " on CPU (please ignore any CUDA init errors)")
+            elif (policy_config["framework"] in ["tf2", "tf", "tfe"] and
+                  not tf.config.list_physical_devices("GPU")) or \
+                    (policy_config["framework"] == "torch" and
+                     not torch.cuda.is_available()):
+                raise RuntimeError(
+                    "GPUs were assigned to this worker by Ray, but "
+                    "your DL framework ({}) reports GPU acceleration is "
+                    "disabled. This could be due to a bad CUDA- or {} "
+                    "installation.".format(
+                        policy_config["framework"],
+                        policy_config["framework"]))
+
         self.multiagent: bool = set(
             self.policy_map.keys()) != {DEFAULT_POLICY_ID}
         if self.multiagent:
diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py
--- a/rllib/policy/torch_policy.py
+++ b/rllib/policy/torch_policy.py
@@ -4,6 +4,7 @@
 import time
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
+import ray
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
@@ -98,8 +99,10 @@ def __init__(self,
         """
         self.framework = "torch"
         super().__init__(observation_space, action_space, config)
-        self.device = (torch.device("cuda")
-                       if torch.cuda.is_available() else torch.device("cpu"))
+        if torch.cuda.is_available() and ray.get_gpu_ids():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
         self.model = model.to(self.device)
         self.exploration = self._create_exploration()
         self.unwrapped_model = model  # used to support DistributedDataParallel


Test patch: