[RLlib] Introduce experimental larger than GPU train batch size featu…

…re for torch (ray-project#34189)

Signed-off-by: Artur Niederfahrenhorst <[email protected]>

rllib/BUILD

@@ -2525,7 +2525,16 @@ py_test(
  name = "tests/test_gpus",
  tags = ["team:rllib", "tests_dir"],
  size = "large",
- srcs = ["tests/test_gpus.py"]
+ srcs = ["tests/test_gpus.py"],
+ args = ["TestGPUs"]
+)
+
+py_test(
+ name = "tests/test_gpus_large_batch",
+ tags = ["team:rllib", "tests_dir"],
+ size = "large",
+ srcs = ["tests/test_gpus.py"],
+ args = ["TestGPUsLargeBatch"]
 )
 
 py_test(

rllib/algorithms/algorithm_config.py

@@ -429,6 +429,7 @@ def __init__(self, algo_class=None):
  self._disable_action_flattening = False
  self._disable_execution_plan_api = True
  self._disable_initialize_loss_from_dummy_batch = False
+ self._load_only_minibatch_onto_device = False
 
  # Has this config object been frozen (cannot alter its attributes anymore).
  self._is_frozen = False
@@ -964,6 +965,14 @@ def validate(self) -> None:
  f"config.framework({self.framework_str})!"
  )
 
+ if (
+ self.simple_optimizer or self.framework_str != "torch"
+ ) and self._load_only_minibatch_onto_device:
+ raise ValueError(
+ "`load_only_minibatch_onto_device` is only supported for "
+ f"config.framework({self.framework_str}) and without simple_optimizer."
+ )
+
  # Detect if specified env is an Atari env.
  if self.is_atari is None:
  self.is_atari = self._detect_atari_env()
@@ -1625,6 +1634,10 @@ def training(
  _enable_learner_api: Whether to enable the LearnerGroup and Learner
  for training. This API uses ray.train to run the training loop which
  allows for a more flexible distributed training.
+ _load_only_minibatch_onto_device: Whether to load only the minibatch onto
+ the given device. This is useful for larger training batches that
+ don't fit on the given device while the mini-batches and their
+ gradients do. This experimental setting is only supported for torch
 
  Returns:
  This updated AlgorithmConfig object.
@@ -2460,6 +2473,7 @@ def experimental(
  _disable_action_flattening: Optional[bool] = NotProvided,
  _disable_execution_plan_api: Optional[bool] = NotProvided,
  _disable_initialize_loss_from_dummy_batch: Optional[bool] = NotProvided,
+ _load_only_minibatch_onto_device: Optional[bool] = NotProvided,
  ) -> "AlgorithmConfig":
  """Sets the config's experimental settings.
 
@@ -2503,6 +2517,8 @@ def experimental(
  self._disable_initialize_loss_from_dummy_batch = (
  _disable_initialize_loss_from_dummy_batch
  )
+ if _load_only_minibatch_onto_device is not NotProvided:
+ self._load_only_minibatch_onto_device = _load_only_minibatch_onto_device
 
  return self
 

rllib/policy/sample_batch.py

@@ -790,15 +790,32 @@ def _zero_pad_in_place(path, value):
  return self
 
  @ExperimentalAPI
- def to_device(self, device, framework="torch"):
+ def to_device(self, device, framework="torch", copy=False):
+ """Moves tensors inside this batch to a given device.
+
+ Depending on the copy flag, this will either return a new batch
+ or modify this batch in-place.
+
+ Args:
+ device: The device to move the tensors to.
+ framework: The framework to use for the device (e.g. "torch").
+ copy: If False, modify batch in place. If True, return a new batch.
+
+ Returns:
+ A batch with all tensors moved to the given device.
+ """
  """TODO: transfer batch to given device as framework tensor."""
  if framework == "torch":
  assert torch is not None
+ if copy:
+ target = SampleBatch()
+ else:
+ target = self
  for k, v in self.items():
- self[k] = convert_to_torch_tensor(v, device)
+ target[k] = convert_to_torch_tensor(v, device)
  else:
  raise NotImplementedError
- return self
+ return target
 
  @PublicAPI
  def size_bytes(self) -> int:

rllib/policy/torch_policy.py

@@ -536,7 +536,14 @@ def load_batch_into_buffer(
  )
 
  # 3) Load splits into the given buffer (consisting of n GPUs).
- slices = [slice.to_device(self.devices[i]) for i, slice in enumerate(slices)]
+ if not self.config.get("_load_only_minibatch_onto_device", False):
+ # We usually want to load the full batch onto the device here, which is
+ # much faster than loading the batch slice-by-slice.
+ # However, if the batch is too large, it may be favorable to load the
+ # batch slice-by-slice.
+ slices = [
+ slice.to_device(self.devices[i]) for i, slice in enumerate(slices)
+ ]
  self._loaded_batches[buffer_index] = slices
 
  # Return loaded samples per-device.
@@ -604,8 +611,20 @@ def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0):
  )
  batch_fetches[f"tower_{i}"] = {"custom_metrics": custom_metrics}
 
+ # If `_load_only_minibatch_onto_device` is True, then the main batch always
+ # remains on the CPU (it's probably too big to be fit on the GPU). Thus, in
+ # this case, for each individual update step, we need to copy the freshly
+ # determined sub-slice to the GPU. These sub-slices need to be small enough
+ # then to fit on the GPU.
+ if self.config.get("_load_only_minibatch_onto_device", False):
+ copy_batch_to_device = True
+ else:
+ copy_batch_to_device = False
+
  # Do the (maybe parallelized) gradient calculation step.
- tower_outputs = self._multi_gpu_parallel_grad_calc(device_batches)
+ tower_outputs = self._multi_gpu_parallel_grad_calc(
+ device_batches, copy_batch_to_device=copy_batch_to_device
+ )
 
  # Mean-reduce gradients over GPU-towers (do this on CPU: self.device).
  all_grads = []
@@ -1061,7 +1080,9 @@ def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch, device=None):
  return postprocessed_batch
 
  def _multi_gpu_parallel_grad_calc(
- self, sample_batches: List[SampleBatch]
+ self,
+ sample_batches: List[SampleBatch],
+ copy_batch_to_device: bool = False,
  ) -> List[Tuple[List[TensorType], GradInfoDict]]:
  """Performs a parallelized loss and gradient calculation over the batch.
 
@@ -1073,6 +1094,10 @@ def _multi_gpu_parallel_grad_calc(
  Args:
  sample_batches: A list of SampleBatch shards to
  calculate loss and gradients for.
+ copy_batch_to_device: Whether to create a copy of the batch that is then
+ moved to GPU. This is useful if we don't want to move the original
+ batch to the device. In case of a large batch, we can thereby only move
+ mini-batches to GPU one by one and free them after each step.
 
  Returns:
  A list (one item per device) of 2-tuples, each with 1) gradient
@@ -1083,6 +1108,11 @@ def _multi_gpu_parallel_grad_calc(
  results = {}
  grad_enabled = torch.is_grad_enabled()
 
+ if copy_batch_to_device:
+ sample_batches = [
+ batch.to_device(i, copy=True) for i, batch in enumerate(sample_batches)
+ ]
+
  def _worker(shard_idx, model, sample_batch, device):
  torch.set_grad_enabled(grad_enabled)
  try:

rllib/policy/torch_policy_v2.py

@@ -734,7 +734,14 @@ def load_batch_into_buffer(
  )
 
  # 3) Load splits into the given buffer (consisting of n GPUs).
- slices = [slice.to_device(self.devices[i]) for i, slice in enumerate(slices)]
+ if not self.config.get("_load_only_minibatch_onto_device", False):
+ # We usually want to load the full batch onto the device here, which is
+ # much faster than loading the batch slice-by-slice.
+ # However, if the batch is too large, it may be favorable to load the
+ # batch slice-by-slice.
+ slices = [
+ slice.to_device(self.devices[i]) for i, slice in enumerate(slices)
+ ]
  self._loaded_batches[buffer_index] = slices
 
  # Return loaded samples per-device.
@@ -749,7 +756,11 @@ def get_num_samples_loaded_into_buffer(self, buffer_index: int = 0) -> int:
 
  @override(Policy)
  @DeveloperAPI
- def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0):
+ def learn_on_loaded_batch(
+ self,
+ offset: int = 0,
+ buffer_index: int = 0,
+ ):
  if not self._loaded_batches[buffer_index]:
  raise ValueError(
  "Must call Policy.load_batch_into_buffer() before "
@@ -803,8 +814,20 @@ def learn_on_loaded_batch(self, offset: int = 0, buffer_index: int = 0):
  )
  batch_fetches[f"tower_{i}"] = {"custom_metrics": custom_metrics}
 
+ # If `_load_only_minibatch_onto_device` is True, then the main batch always
+ # remains on the CPU (it's probably too big to be fit on the GPU). Thus, in
+ # this case, for each individual update step, we need to copy the freshly
+ # determined sub-slice to the GPU. These sub-slices need to be small enough
+ # then to fit on the GPU.
+ if self.config.get("_load_only_minibatch_onto_device", False):
+ copy_batch_to_device = True
+ else:
+ copy_batch_to_device = False
+
  # Do the (maybe parallelized) gradient calculation step.
- tower_outputs = self._multi_gpu_parallel_grad_calc(device_batches)
+ tower_outputs = self._multi_gpu_parallel_grad_calc(
+ device_batches, copy_batch_to_device=copy_batch_to_device
+ )
 
  # Mean-reduce gradients over GPU-towers (do this on CPU: self.device).
  all_grads = []
@@ -1188,7 +1211,9 @@ def _lazy_tensor_dict(self, postprocessed_batch: SampleBatch, device=None):
  return postprocessed_batch
 
  def _multi_gpu_parallel_grad_calc(
- self, sample_batches: List[SampleBatch]
+ self,
+ sample_batches: List[SampleBatch],
+ copy_batch_to_device: bool = False,
  ) -> List[Tuple[List[TensorType], GradInfoDict]]:
  """Performs a parallelized loss and gradient calculation over the batch.
 
@@ -1200,6 +1225,10 @@ def _multi_gpu_parallel_grad_calc(
  Args:
  sample_batches: A list of SampleBatch shards to
  calculate loss and gradients for.
+ copy_batch_to_device: Whether to create a copy of the batch that is then
+ moved to GPU. This is useful if we don't want to move the original
+ batch to the device. In case of a large batch, we can thereby only move
+ mini-batches to GPU one by one and free them after each step.
 
  Returns:
  A list (one item per device) of 2-tuples, each with 1) gradient
@@ -1210,6 +1239,11 @@ def _multi_gpu_parallel_grad_calc(
  results = {}
  grad_enabled = torch.is_grad_enabled()
 
+ if copy_batch_to_device:
+ sample_batches = [
+ batch.to_device(i, copy=True) for i, batch in enumerate(sample_batches)
+ ]
+
  def _worker(shard_idx, model, sample_batch, device):
  torch.set_grad_enabled(grad_enabled)
  try:

rllib/tests/test_gpus.py

@@ -1,13 +1,19 @@
+import copy
 import unittest
 
+import numpy as np
+import torch
+
 import ray
 from ray import air
+from ray import tune
 from ray.rllib.algorithms.a2c.a2c import A2CConfig
-from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.algorithms.qmix import QMixConfig
+from ray.rllib.policy.torch_policy import TorchPolicy
+from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
+from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.test_utils import framework_iterator
-from ray import tune
-
-torch, _ = try_import_torch()
 
 
 class TestGPUs(unittest.TestCase):
@@ -111,8 +117,94 @@ def test_gpus_in_local_mode(self):
  ray.shutdown()
 
 
+class TestGPUsLargeBatch(unittest.TestCase):
+ def test_larger_train_batch_size_multi_gpu_train_one_step(self):
+ # Tests that we can use a `train_batch_size` larger than GPU memory with our
+ # experimental setting `_load_only_minibatch_onto_device` with
+ # multi_gpu_train_one_step.
+
+ # These values make it so that one large minibatch and the optimizer
+ # variables can fit onto the device, but the whole sample_batch is already too
+ # large for the GPU itself.
+ sgd_minibatch_size = int(1e4)
+ train_batch_size = int(sgd_minibatch_size * 1e5)
+
+ # Fake CartPole episode of n time steps.
+ CARTPOLE_FAKE_BATCH = SampleBatch(
+ {
+ SampleBatch.OBS: np.zeros((train_batch_size, 4), dtype=np.float32),
+ SampleBatch.ACTIONS: np.zeros((train_batch_size,), dtype=np.float32),
+ SampleBatch.PREV_ACTIONS: np.zeros(
+ (train_batch_size,), dtype=np.float32
+ ),
+ SampleBatch.REWARDS: np.zeros((train_batch_size,), dtype=np.float32),
+ SampleBatch.PREV_REWARDS: np.zeros(
+ (train_batch_size,), dtype=np.float32
+ ),
+ "value_targets": np.zeros((train_batch_size,), dtype=np.float32),
+ SampleBatch.TERMINATEDS: np.array([False] * train_batch_size),
+ SampleBatch.TRUNCATEDS: np.array([False] * train_batch_size),
+ "advantages": np.zeros((train_batch_size,), dtype=np.float32),
+ SampleBatch.VF_PREDS: np.zeros((train_batch_size,), dtype=np.float32),
+ SampleBatch.ACTION_DIST_INPUTS: np.zeros(
+ (train_batch_size, 2), dtype=np.float32
+ ),
+ SampleBatch.ACTION_LOGP: np.zeros(
+ (train_batch_size,), dtype=np.float32
+ ),
+ SampleBatch.EPS_ID: np.zeros((train_batch_size,), dtype=np.int64),
+ SampleBatch.AGENT_INDEX: np.zeros((train_batch_size,), dtype=np.int64),
+ }
+ )
+
+ # Test if we can even fail this test due too a GPU OOM
+ try:
+ batch_copy = copy.deepcopy(CARTPOLE_FAKE_BATCH)
+ batch_copy.to_device(0)
+ raise ValueError(
+ "We should not be able to move this batch to the device. "
+ "If this error occurs, this means that this test cannot fail "
+ "inside multi_gpu_train_one_step."
+ )
+ except torch.cuda.OutOfMemoryError:
+ pass
+
+ for config_class in (PPOConfig, QMixConfig):
+ config = (
+ config_class()
+ .environment(env="CartPole-v1")
+ .framework("torch")
+ .resources(num_gpus=1)
+ .rollouts(num_rollout_workers=0)
+ .training(
+ train_batch_size=train_batch_size,
+ num_sgd_iter=1,
+ sgd_minibatch_size=self.sgd_minibatch_size,
+ # This setting makes it so that we don't load a batch of
+ # size `train_batch_size` onto the device, but only
+ # minibatches.
+ _load_only_minibatch_onto_device=True,
+ )
+ )
+
+ algorithm = config.build()
+ policy = algorithm.get_policy()
+
+ # Sanity check if we are covering both, TorchPolicy and TorchPolicyV2
+ if config_class is QMixConfig:
+ assert isinstance(policy, TorchPolicy)
+ elif config_class is PPOConfig:
+ assert isinstance(policy, TorchPolicyV2)
+
+ policy.load_batch_into_buffer(CARTPOLE_FAKE_BATCH)
+ policy.learn_on_loaded_batch()
+
+
 if __name__ == "__main__":
  import pytest
  import sys
 
- sys.exit(pytest.main(["-v", __file__]))
+ # One can specify the specific TestCase class to run.
+ # None for all unittest.TestCase classes in this file.
+ class_ = sys.argv[1] if len(sys.argv) > 1 else None
+ sys.exit(pytest.main(["-v", __file__ + ("" if class_ is None else "::" + class_)]))