Remove unused VLLM import from train_rl & rename install_tunix_vllm_requirement.sh

Author: SurbhiJainUSC

docs/tutorials/posttraining/rl.md

@@ -54,7 +54,7 @@ Next, run the following bash script to get all the necessary installations insid
 This will take few minutes. Follow along the installation logs and look out for any issues!
 
 ```
-bash ~/maxtext/src/MaxText/examples/install_tunix_vllm_requirement.sh
+bash tools/setup/setup_post_training_requirements.sh
 ```
 
 Primarily, it installs `vllm-tpu` which is [vllm](https://github.com/vllm-project/vllm) and [tpu-inference](https://github.com/vllm-project/tpu-inference) and thereby providing TPU inference for vLLM, with unified JAX and PyTorch support.

src/MaxText/examples/sft_train_and_evaluate.py

@@ -21,8 +21,11 @@
 
 ## Example command to run on single-host TPU:
 ```
-# Install dependencies in virtual environment
-bash setup.sh && bash src/MaxText/examples/install_tunix_vllm_requirement.sh
+# Install dependencies in virtual environment:
+# https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-pypi-recommended
+
+# Install post-training dependencies in virtual environment:
+bash tools/setup/setup_post_training_requirements.sh
 
 # Environment configurations
 export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)

src/MaxText/rl/train_rl.py

@@ -64,7 +64,6 @@
 from tunix.rl.rollout import base_rollout
 from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner
 from tunix.sft import metrics_logger, profiler
-from vllm.outputs import PoolingRequestOutput  # pylint: disable=unused-import
 
 # for vLLM we can skip JAX precompilation with this flag, it makes startup faster
 os.environ["SKIP_JAX_PRECOMPILE"] = "1"

…amples/install_tunix_vllm_requirement.sh …etup/setup_post_training_requirements.shsrc/MaxText/examples/install_tunix_vllm_requirement.sh renamed to tools/setup/setup_post_training_requirements.sh src/MaxText/examples/install_tunix_vllm_requirement.sh renamed to tools/setup/setup_post_training_requirements.sh

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# This script installs the dependencies for running GRPO with MaxText+Tunix+vLLM on TPUs
+# This script installs the dependencies for running post-training with MaxText+Tunix+vLLM on TPUs
 
 set -e
 set -x