fix errors with merging GPU specification

Author: angelahzyuan

scripts/combine_generate.py

@@ -17,7 +17,7 @@ def main():
     for j in range(args.pairs):
         results = []
         if args.gpu_ids is not None:
-            gpus = args.gpu_ids.strip("()").split(' ')
+            gpus = args.gpu_ids.strip("()").split(',')
         else:
             gpus = range(args.numgpu)
 

scripts/compute_prob.py

@@ -26,7 +26,7 @@ def from_ranks(args):
 
     scores = [0 for _ in range(len(data))]
     if args.gpu_ids is not None:
-        gpus = args.gpu_ids.strip("()").split(' ')
+        gpus = args.gpu_ids.strip("()").split(',')
     else:
         gpus = range(args.num_gpu)
 

scripts/generate.sh

@@ -62,7 +62,7 @@ all_gen=$!
 
 wait $all_gen
 
-# python3 scripts/combine_generate.py --output_dir "generated/$OUTDIR" --gpu_ids "${AVAILABLE_GPUS[@]}" --pairs $PAIRS
+python3 scripts/combine_generate.py --output_dir "generated/$OUTDIR" --gpu_ids "$(IFS=, ; echo "${AVAILABLE_GPUS[*]}")" --pairs $PAIRS
 
 
 # #####################
@@ -85,4 +85,4 @@ all_rank=$!
 
 wait $all_rank
 
-python3 scripts/compute_prob.py --org $HF_ORG --gpu_ids "${AVAILABLE_GPUS[@]}" --output_dir $OUTDIR --pairs $PAIRS --frac_len $FRAC_LEN --prompts $PROMPTS
+python3 scripts/compute_prob.py --org $HF_ORG --gpu_ids "$(IFS=, ; echo "${AVAILABLE_GPUS[*]}")" --output_dir $OUTDIR --pairs $PAIRS --frac_len $FRAC_LEN --prompts $PROMPTS

setup.py

@@ -42,7 +42,7 @@
 # IMPORTANT: all dependencies should be listed here with their version requirements, if any.
 #   * If a dependency is fast-moving (e.g. transformers), pin to the exact version
 _deps = [
-    "accelerate==0.23.0",
+    "accelerate==0.27.2",
     "bitsandbytes==0.41.2.post2",
     "black==23.1.0",
     "datasets==2.14.6",
@@ -67,7 +67,7 @@
     "tensorboard",
     "torch==2.1.2",
     "transformers==4.42.4",
-    "trl==0.7.10",
+    "trl==0.9.6",
     "jinja2>=3.0.0",
     "tqdm>=4.64.1",
 ]