[CI] linux://doc/source/ray-air/examples:gptj_serving is failing/flaky on master.

[justinvyu]

• 077addb FAILED Buildkite 🚋 ml: train gpu tests
• 40073ee FAILED Buildkite 🚋 ml: train gpu tests
• ab1532c FAILED Buildkite 🚋 ml: train gpu tests
• 9afdf6d FAILED Buildkite 🚋 ml: train gpu tests
• fa4963b FAILED Buildkite 🚋 ml: train gpu tests
• 55ae63b FAILED Buildkite 🚋 ml: train gpu tests
• d432bd6 FAILED Buildkite 🚋 ml: train gpu tests
• 7b33cd8 FAILED Buildkite 🚋 ml: train gpu tests
• d3ac31e FAILED Buildkite 🚋 ml: train gpu tests
• 28fdcb6 FAILED Buildkite 🚋 ml: train gpu tests
• 3eeb3de FAILED Buildkite 🚋 ml: train gpu tests

....
Generated from flaky test tracker. Please do not edit the signature in this section.
DataCaseName-linux://doc/source/ray-air/examples:gptj_serving-END
....

[anyscalesam]

@justinvyu can we upgrade this to p0 given it's a linux release test and something when it gets to ray210 release we would block correct?

[justinvyu]

@anyscalesam Yes let's fix it. I think it's caused by some external change (GPU out of memory error without any code change on our side).

(ServeReplica:default:PredictDeployment pid=12306)     attn_outputs = self.attn(
--
  | (ServeReplica:default:PredictDeployment pid=12306)   File "/opt/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  | (ServeReplica:default:PredictDeployment pid=12306)     return forward_call(*args, **kwargs)
  | (ServeReplica:default:PredictDeployment pid=12306)   File "/opt/miniconda/lib/python3.8/site-packages/accelerate/hooks.py", line 165, in new_forward
  | (ServeReplica:default:PredictDeployment pid=12306)     output = old_forward(*args, **kwargs)
  | (ServeReplica:default:PredictDeployment pid=12306)   File "/tmp/ray/session_2023-11-29_11-26-47_341554_8580/runtime_resources/pip/9fc91d65ff21920e7899b7316ff121f0058d7e36/virtualenv/lib/python3.8/site-packages/transformers/models/gptj/modeling_gptj.py", line 254, in forward
  | (ServeReplica:default:PredictDeployment pid=12306)     present = (key.to(hidden_states.dtype), value)
  | (ServeReplica:default:PredictDeployment pid=12306) torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 14.56 GiB total capacity; 11.48 GiB already allocated; 1.50 MiB free; 11.54 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF