[BugFix] Fix GRPO tests and runs

ghstack-source-id: a5b7078
Pull-Request: #3213

Author: vmoens

.gitignore

@@ -189,3 +189,11 @@ log
 Roms
 
 scratch/*
+
+# Large directories from git history that should not be committed
+dev/
+main/
+*.html
+
+# Additional cache directories
+.ruff_cache/

sota-implementations/expert-iteration/expert-iteration-async.py

@@ -12,7 +12,7 @@
 
 import hydra
 
-from torchrl import torchrl_logger
+from torchrl import merge_ray_runtime_env, torchrl_logger
 from torchrl.data.llm.history import History
 from torchrl.record.loggers.wandb import WandbLogger
 from torchrl.weight_update.llm import get_model_metadata
@@ -397,19 +397,9 @@ def main(cfg):
             if not k.startswith("_")
         }
 
-        # Add computed GPU configuration
+        # Add computed GPU configuration and merge with default runtime_env
         ray_init_config["num_gpus"] = device_config["ray_num_gpus"]
-        # Ensure runtime_env and env_vars exist
-        if "runtime_env" not in ray_init_config:
-            ray_init_config["runtime_env"] = {}
-        if not isinstance(ray_init_config["runtime_env"], dict):
-            ray_init_config["runtime_env"] = dict(ray_init_config["runtime_env"])
-        if "env_vars" not in ray_init_config["runtime_env"]:
-            ray_init_config["runtime_env"]["env_vars"] = {}
-        if not isinstance(ray_init_config["runtime_env"]["env_vars"], dict):
-            ray_init_config["runtime_env"]["env_vars"] = dict(
-                ray_init_config["runtime_env"]["env_vars"]
-            )
+        ray_init_config = merge_ray_runtime_env(ray_init_config)
         torchrl_logger.info(f"Ray init config: {ray_init_config=}")
         ray.init(**ray_init_config)
 

sota-implementations/expert-iteration/expert-iteration-sync.py

@@ -12,7 +12,7 @@
 
 import hydra
 
-from torchrl import torchrl_logger
+from torchrl import merge_ray_runtime_env, torchrl_logger
 from torchrl.data.llm.history import History
 from torchrl.record.loggers.wandb import WandbLogger
 from torchrl.weight_update.llm import get_model_metadata
@@ -398,19 +398,9 @@ def main(cfg):
             if not k.startswith("_")
         }
 
-        # Add computed GPU configuration
+        # Add computed GPU configuration and merge with default runtime_env
         ray_init_config["num_gpus"] = device_config["ray_num_gpus"]
-        # Ensure runtime_env and env_vars exist
-        if "runtime_env" not in ray_init_config:
-            ray_init_config["runtime_env"] = {}
-        if not isinstance(ray_init_config["runtime_env"], dict):
-            ray_init_config["runtime_env"] = dict(ray_init_config["runtime_env"])
-        if "env_vars" not in ray_init_config["runtime_env"]:
-            ray_init_config["runtime_env"]["env_vars"] = {}
-        if not isinstance(ray_init_config["runtime_env"]["env_vars"], dict):
-            ray_init_config["runtime_env"]["env_vars"] = dict(
-                ray_init_config["runtime_env"]["env_vars"]
-            )
+        ray_init_config = merge_ray_runtime_env(ray_init_config)
         torchrl_logger.info(f"Ray init config: {ray_init_config=}")
         ray.init(**ray_init_config)
 

sota-implementations/grpo/grpo-async.py

@@ -13,7 +13,7 @@
 
 import hydra
 
-from torchrl import torchrl_logger
+from torchrl import merge_ray_runtime_env, torchrl_logger
 from torchrl.data.llm.history import History
 from torchrl.record.loggers.wandb import WandbLogger
 from torchrl.weight_update.llm import get_model_metadata
@@ -319,19 +319,9 @@ def main(cfg):
             if not k.startswith("_")
         }
 
-        # Add computed GPU configuration
+        # Add computed GPU configuration and merge with default runtime_env
         ray_init_config["num_gpus"] = device_config["ray_num_gpus"]
-        # Ensure runtime_env and env_vars exist
-        if "runtime_env" not in ray_init_config:
-            ray_init_config["runtime_env"] = {}
-        if not isinstance(ray_init_config["runtime_env"], dict):
-            ray_init_config["runtime_env"] = dict(ray_init_config["runtime_env"])
-        if "env_vars" not in ray_init_config["runtime_env"]:
-            ray_init_config["runtime_env"]["env_vars"] = {}
-        if not isinstance(ray_init_config["runtime_env"]["env_vars"], dict):
-            ray_init_config["runtime_env"]["env_vars"] = dict(
-                ray_init_config["runtime_env"]["env_vars"]
-            )
+        ray_init_config = merge_ray_runtime_env(ray_init_config)
         torchrl_logger.info(f"Ray init config: {ray_init_config=}")
         ray_managed_externally = os.environ.get("RAY_CLUSTER_MANAGED_EXTERNALLY")
         if ray_managed_externally:

sota-implementations/grpo/grpo-sync.py

@@ -12,7 +12,7 @@
 
 import hydra
 
-from torchrl import torchrl_logger
+from torchrl import merge_ray_runtime_env, torchrl_logger
 from torchrl.data.llm.history import History
 from torchrl.record.loggers.wandb import WandbLogger
 from torchrl.weight_update.llm import get_model_metadata
@@ -319,19 +319,9 @@ def main(cfg):
             if not k.startswith("_")
         }
 
-        # Add computed GPU configuration
+        # Add computed GPU configuration and merge with default runtime_env
         ray_init_config["num_gpus"] = device_config["ray_num_gpus"]
-        # Ensure runtime_env and env_vars exist
-        if "runtime_env" not in ray_init_config:
-            ray_init_config["runtime_env"] = {}
-        if not isinstance(ray_init_config["runtime_env"], dict):
-            ray_init_config["runtime_env"] = dict(ray_init_config["runtime_env"])
-        if "env_vars" not in ray_init_config["runtime_env"]:
-            ray_init_config["runtime_env"]["env_vars"] = {}
-        if not isinstance(ray_init_config["runtime_env"]["env_vars"], dict):
-            ray_init_config["runtime_env"]["env_vars"] = dict(
-                ray_init_config["runtime_env"]["env_vars"]
-            )
+        ray_init_config = merge_ray_runtime_env(ray_init_config)
         torchrl_logger.info(f"Ray init config: {ray_init_config=}")
         ray_managed_externally = os.environ.get("RAY_CLUSTER_MANAGED_EXTERNALLY")
         if ray_managed_externally:

sota-implementations/grpo/requirements_gsm8k.txt

@@ -1,13 +1,11 @@
-torch==2.7.0
-transformers==4.52.4
-peft==0.15.2
-bitsandbytes==0.46.0
-datasets==3.6.0
-wandb==0.19.11
-hydra-core==1.3.2
-ray==2.46.0
-tqdm==4.67.1
-tensordict==0.9.0
-vllm==0.9.0.1
-accelerate==1.7.0
-xformers==0.0.30 
+vllm==0.11.0
+peft
+bitsandbytes
+datasets
+wandb
+hydra-core
+ray
+tqdm
+tensordict
+accelerate
+xformers

sota-implementations/grpo/requirements_ifeval.txt

@@ -1,16 +1,16 @@
-torch==2.7.0
-transformers==4.52.4
-peft==0.15.2
-bitsandbytes==0.46.0
-datasets==3.6.0
-wandb==0.19.11
-hydra-core==1.3.2
-ray==2.46.0
-tqdm==4.67.1
-tensordict==0.9.0
-vllm==0.9.0.1
-accelerate==1.7.0
-xformers==0.0.30 
-nltk==3.9.1
-langdetect==1.0.9
-immutabledict==4.2.1
+vllm==0.11.0
+torch
+transformers
+peft
+bitsandbytes
+datasets
+wandb
+hydra-core
+ray
+tqdm
+tensordict
+accelerate
+xformers
+nltk
+langdetect
+immutabledict

test/llm/test_data.py

@@ -13,7 +13,7 @@
 import torch
 from tensordict import lazy_stack, set_list_to_stack, TensorDict
 
-from torchrl import torchrl_logger
+from torchrl import logger as torchrl_logger
 
 from torchrl.data import (
     History,