Rlhf hf

benchmarks/rlhf/dev.yaml

@@ -1,5 +1,5 @@
 
-rlhf_:
+_rlhf:
   inherits: _defaults
   definition: .
   install-variant: unpinned
@@ -17,13 +17,12 @@ rlhf_:
 
 
 rlhf-single:
-  inherits: rlhf_
+  inherits: _rlhf
   plan:
     method: per_gpu
 
-
 rlhf-gpus:
-  inherits: rlhf_
+  inherits: _rlhf
   plan:
     method: njobs
     n: 1

benchmarks/rlhf/main.py

@@ -118,11 +118,8 @@ def tokenize(element):
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
     )
+
     trainer.train()
-    trainer.save_model(config.output_dir)
-    if config.push_to_hub:
-        trainer.push_to_hub()
-    trainer.generate_completions()
 
 
 if __name__ == "__main__":

config/base.yaml

@@ -777,7 +777,7 @@ llava-gpus:
     --gradient_accumulation_steps: 1
 
 
-rlhf_:
+_rlhf:
   inherits: _defaults
   definition: ../benchmarks/rlhf
   install-variant: unpinned
@@ -796,12 +796,12 @@ rlhf_:
     --bf16: true
 
 rlhf-single:
-  inherits: rlhf_
+  inherits: _rlhf
   plan:
     method: per_gpu
 
 rlhf-gpus:
-  inherits: rlhf_
+  inherits: _rlhf
   plan:
     method: njobs
     n: 1

scripts/article/run_cuda.sh

@@ -76,9 +76,6 @@ else
     . $MILABENCH_WORDIR/env/bin/activate
 fi
 
-
-export MILABENCH_CONFIG=/home/mila/d/delaunap/milabench/benchmarks/rlhf/dev.yaml
-
 if [ "$MILABENCH_PREPARE" -eq 0 ]; then
     cd $MILABENCH_WORDIR