[feat] Ray train integration

README.md

@@ -104,7 +104,8 @@ For more usage see the [NeMo README](./trlx/models)
 #### Use Ray Tune to launch hyperparameter sweep
 
 ```bash
-python -m trlx.sweep --config configs/sweeps/ppo_sweep.yml examples/ppo_sentiments.py
+ray start --head --port=6379
+python -m trlx.sweep --config configs/sweeps/ppo_sweep.yml --accelerate_config configs/accelerate/ddp.yaml --num_gpus 4 examples/ppo_sentiments.py
 ```
 
 #### Benchmark your trlX fork against trlX's `main` branch

configs/accelerate/ddp.yaml

@@ -2,7 +2,7 @@ compute_environment: LOCAL_MACHINE
 deepspeed_config: {}
 distributed_type: MULTI_GPU
 downcast_bf16: no
-dynamo_backend: 'NO'
+dynamo_config: {}
 fsdp_config: {}
 gpu_ids: all
 machine_rank: 0

configs/accelerate/zero2-bf16.yaml

@@ -9,7 +9,7 @@ deepspeed_config:
   zero_stage: 2
 distributed_type: DEEPSPEED
 downcast_bf16: no
-dynamo_backend: 'NO'
+dynamo_config: {}
 fsdp_config: {}
 machine_rank: 0
 main_training_function: main

configs/accelerate/zero2-fp16.yaml

@@ -9,7 +9,7 @@ deepspeed_config:
   zero_stage: 2
 distributed_type: DEEPSPEED
 downcast_bf16: no
-dynamo_backend: 'NO'
+dynamo_config: {}
 fsdp_config: {}
 machine_rank: 0
 main_training_function: main

configs/accelerate/zero3.yaml

@@ -10,7 +10,7 @@ deepspeed_config:
   zero_stage: 3
 distributed_type: DEEPSPEED
 downcast_bf16: no
-dynamo_backend: 'NO'
+dynamo_config: {}
 fsdp_config: {}
 machine_rank: 0
 main_training_function: main

configs/sweeps/ilql_sweep.yml

@@ -3,17 +3,26 @@ tune_config:
   metric: "metrics/sentiments"
   search_alg: "random"
   scheduler: "fifo"
-  num_samples: 32
+  num_samples: 64
 
-lr:
+# https://docs.ray.io/en/latest/tune/api_docs/search_space.html#tune-sample-docs
+optimizer.kwargs.lr:
   strategy: "loguniform"
-  values: [0.00001, 0.01]
-tau:
+  values: [0.000001, 0.001]
+method.tau:
   strategy: "uniform"
   values: [0.6, 0.9]
-steps_for_target_q_sync:
+method.steps_for_target_q_sync:
   strategy: "choice"
   values: [1, 5, 10]
-alpha:
+method.alpha:
   strategy: "loguniform"
   values: [0.001, 1.0]
+
+# disable checkpointing for storage sake
+train.checkpoint_interval:
+  strategy: "choice"
+  values: [10000000]
+train.save_best:
+  strategy: "choice"
+  values: [false]

configs/sweeps/ppo_sweep.yml

@@ -6,12 +6,26 @@ tune_config:
   num_samples: 32
 
 # https://docs.ray.io/en/latest/tune/api_docs/search_space.html#tune-sample-docs
-lr:
+optimizer.kwargs.lr:
   strategy: "loguniform"
-  values: [0.00001, 0.01]
-init_kl_coef:
-  strategy: "uniform"
-  values: [0, 0.2]
-vf_coef:
-  strategy: "uniform"
-  values: [0.5, 2]
+  values: [0.000001, 0.001]
+method.init_kl_coef:
+  strategy: "loguniform"
+  values: [0.0001, 0.2]
+model.num_layers_unfrozen:
+  strategy: "choice"
+  values: [-1, 2, 6]
+method.num_rollouts:
+  strategy: "choice"
+  values: [32, 128, 512]
+method.target:
+  strategy: "choice"
+  values: [null, 1]
+
+# disable checkpointing for storage sake
+train.checkpoint_interval:
+  strategy: "choice"
+  values: [10000000]
+train.save_best:
+  strategy: "choice"
+  values: [false]

examples/ilql_sentiments.py

@@ -37,7 +37,7 @@ def metric_fn(samples: List[str], **kwargs) -> Dict[str, List[float]]:
     trlx.train(
         samples=imdb["text"],
         rewards=imdb["label"],
-        eval_prompts=["I don't know much about Hungarian underground"] * 64,
+        eval_prompts=["I don't know much about Hungarian underground"] * 256,
         metric_fn=metric_fn,
         config=config,
     )

examples/ppo_sentiments.py

@@ -47,7 +47,7 @@ def reward_fn(samples: List[str], **kwargs) -> List[float]:
     trlx.train(
         reward_fn=reward_fn,
         prompts=prompts,
-        eval_prompts=["I don't know much about Hungarian underground"] * 64,
+        eval_prompts=["I don't know much about Hungarian underground"] * 256,
         config=config,
     )
 

examples/summarize_rlhf/configs/default_accelerate_config.yaml

@@ -6,7 +6,7 @@ deepspeed_config:
   zero3_init_flag: false
 distributed_type: DEEPSPEED
 downcast_bf16: 'no'
-dynamo_backend: 'NO'
+dynamo_config: {}
 fsdp_config: {}
 gpu_ids: null
 machine_rank: 0

scripts/sweep-cw.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+#SBATCH --job-name=trlx-sweep
+#SBATCH --account=trlx
+#SBATCH --partition=a100-cu117
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=1
+#SBATCH --mem=0
+#SBATCH --output=%j
+#SBATCH --exclusive
+
+export NCCL_DEBUG=WARN
+export NCCL_PROTO=simple
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+# export CUDA_LAUNCH_BLOCKING=1
+
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+
+cd $TRLX
+source $TRLX/venv-with-pinned-ray/bin/activate
+
+ray start --head --port=6379 &
+
+export HOSTNAMES=($HOSTNAMES)
+for node in ${HOSTNAMES[@]:1}; do
+    echo "Starting ray worker @ $node"
+    srun --nodes=1 --ntasks=1 -w "$node" ray start --address $MASTER_ADDR:6379 --block &
+done
+
+sleep 10
+ray status
+
+NUM_GPUS=16
+python -m trlx.sweep -y --config configs/sweeps/ppo_sweep.yml --accelerate_config configs/accelerate/zero2-bf16.yaml --num_gpus $NUM_GPUS examples/ppo_sentiments.py
+# python -m trlx.sweep -y --config configs/sweeps/ilql_sweep.yml --default_config configs/ilql_config.yml --accelerate_config configs/accelerate/zero2-bf16.yaml --num_gpus $NUM_GPUS examples/ilql_sentiments.py

setup.cfg

@@ -11,22 +11,24 @@ license = MIT
 [options]
 packages = find:
 install_requires =
-    accelerate>=0.16.0
+    accelerate>=0.17.1
     attrs>=22.1.0
     cattrs>=22.2.0
     datasets
-    deepspeed>=0.7.3
+    deepspeed>=0.8.1
     einops>=0.4.1
     numpy>=1.23.2
     torchtyping
     transformers>=4.21.2
     tqdm
     rich
     wandb>=0.13.5
-    ray>=2.0.1
     tabulate>=0.9.0
     networkx
     tritonclient
+    ray@https://ray-ci-artifact-branch-public.s3.amazonaws.com/42bb0357a6fb13e4994789c824f3623f32869ad8/tmp/artifacts/.whl/ray-3.0.0.dev0-cp38-cp38-manylinux2014_x86_64.whl ; python_version=="3.8"
+    ray@https://ray-ci-artifact-branch-public.s3.amazonaws.com/42bb0357a6fb13e4994789c824f3623f32869ad8/tmp/artifacts/.whl/ray-3.0.0.dev0-cp39-cp39-manylinux2014_x86_64.whl ; python_version=="3.9"
+    ray@https://ray-ci-artifact-branch-public.s3.amazonaws.com/42bb0357a6fb13e4994789c824f3623f32869ad8/tmp/artifacts/.whl/ray-3.0.0.dev0-cp310-cp310-manylinux2014_x86_64.whl ; python_version=="3.10"
 
 [options.extras_require]
 bnb = bitsandbytes

trlx/data/configs.py

@@ -295,13 +295,27 @@ def from_dict(cls, config: Dict):
 
     @classmethod
     def update(cls, baseconfig: Dict, config: Dict):
+        update = {}
+        # unflatten a string variable name into a nested dictionary
+        # key1.key2.key3: value -> {key1: {key2: {key3: value}}}
+        for name, value in config.items():
+            if isinstance(value, dict):
+                update[name] = value
+            else:
+                *layers, var = name.split(".")
+                if layers:
+                    d = update.setdefault(layers[0], {})
+                    for layer in layers[1:]:
+                        d = d.setdefault(layer, {})
+                    d[var] = value
+
         if not isinstance(baseconfig, Dict):
             baseconfig = baseconfig.to_dict()
 
         updates = set()
-        merged = merge(baseconfig, config, updates)
+        merged = merge(baseconfig, update, updates)
 
-        for param in config:
+        for param in update:
             if param not in updates:
                 raise ValueError(f"parameter {param} is not present in the config (typo or a wrong config)")
 

trlx/data/default_configs.py

@@ -27,16 +27,16 @@ def default_ppo_config():
         model=ModelConfig(model_path="lvwerra/gpt2-imdb", num_layers_unfrozen=2),
         tokenizer=TokenizerConfig(tokenizer_path="gpt2", truncation_side="right"),
         optimizer=OptimizerConfig(
-            name="adamw", kwargs=dict(lr=1.0e-4, betas=(0.9, 0.95), eps=1.0e-8, weight_decay=1.0e-6)
+            name="adamw", kwargs=dict(lr=3e-5, betas=(0.9, 0.95), eps=1.0e-8, weight_decay=1.0e-6)
         ),
-        scheduler=SchedulerConfig(name="cosine_annealing", kwargs=dict(T_max=10000, eta_min=1.0e-4)),
+        scheduler=SchedulerConfig(name="cosine_annealing", kwargs=dict(T_max=1e12, eta_min=3e-5)),
         method=PPOConfig(
             name="PPOConfig",
             num_rollouts=128,
             chunk_size=128,
             ppo_epochs=4,
-            init_kl_coef=0.05,
-            target=6,
+            init_kl_coef=0.001,
+            target=None,
             horizon=10000,
             gamma=1,
             lam=0.95,
@@ -61,7 +61,7 @@ def default_ilql_config():
     return TRLConfig(
         train=TrainConfig(
             seq_length=64,
-            batch_size=32,
+            batch_size=128,
             epochs=100,
             total_steps=1000,
             checkpoint_interval=1000,
@@ -75,7 +75,7 @@ def default_ilql_config():
             name="adamw", kwargs=dict(lr=5.0e-5, betas=(0.9, 0.95), eps=1.0e-8, weight_decay=1.0e-6)
         ),
         scheduler=SchedulerConfig(
-            name="cosine_annealing", kwargs=dict(T_max=1000, eta_min=5.0e-5)  # train.total_steps
+            name="cosine_annealing", kwargs=dict(T_max=1e12, eta_min=5.0e-5)  # train.total_steps
         ),
         method=ILQLConfig(
             name="ilqlconfig",
@@ -87,7 +87,7 @@ def default_ilql_config():
             beta=0,
             steps_for_target_q_sync=5,
             two_qs=True,
-            gen_kwargs=dict(max_new_tokens=56, top_k=20, beta=4, temperature=1.0),
+            gen_kwargs=dict(max_new_tokens=56, top_k=20, beta=1, temperature=1.0),
         ),
     )
 
@@ -110,7 +110,7 @@ def default_sft_config():
             name="adamw", kwargs=dict(lr=1.0e-4, betas=(0.9, 0.95), eps=1.0e-8, weight_decay=1.0e-6)
         ),
         scheduler=SchedulerConfig(
-            name="cosine_annealing", kwargs=dict(T_max=10000, eta_min=1.0e-4)  # train.total_steps
+            name="cosine_annealing", kwargs=dict(T_max=1e12, eta_min=1.0e-4)  # train.total_steps
         ),
         method=SFTConfig(
             name="sftconfig",