🐛 Style Transfer: use 1 GPU because NCCL freeze using VLLM inference.

arkhn · Oct 15, 2024 · 1ce052a · 1ce052a
1 parent 7c13bcc
commit 1ce052a
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/lib/style-transfer/configs/rb_gen/dpo/default.yaml b/lib/style-transfer/configs/rb_gen/dpo/default.yaml
@@ -1,14 +1,14 @@
 training_args:
   _target_: trl.DPOConfig
-  per_device_train_batch_size: 4
+  per_device_train_batch_size: 1
   logging_steps: 2
   save_steps: 50
-  gradient_accumulation_steps: 4
+  gradient_accumulation_steps: 16
   gradient_checkpointing: false
   learning_rate: 2e-6
   weight_decay: 1e-7
   eval_strategy: "no"
-  num_train_epochs: 15
+  num_train_epochs: 10
   output_dir: "models/dpo"
   optim: "adafactor"
   save_only_model: true

diff --git a/lib/style-transfer/configs/rb_gen/gen/default.yaml b/lib/style-transfer/configs/rb_gen/gen/default.yaml
@@ -6,9 +6,9 @@ sampling_params:
   top_p: 1.0
 llm:
   _target_: vllm.LLM
-  _partial_: true
   model: "models/merged/"
   speculative_model: "[ngram]"
   num_speculative_tokens: 5
   ngram_prompt_lookup_max: 4
   use_v2_block_manager: true
+  tensor_parallel_size: 1
diff --git a/lib/style-transfer/configs/rb_gen/train.yaml b/lib/style-transfer/configs/rb_gen/train.yaml
@@ -36,4 +36,4 @@ dataset:
   num_generated_samples: 6
 
 seed: 666
-max_steps: 1
+max_steps: 2
diff --git a/lib/style-transfer/style_transfer/rb_gen/steps/generate.py b/lib/style-transfer/style_transfer/rb_gen/steps/generate.py
@@ -58,7 +58,7 @@ def generate(
     gc.collect()
     torch.cuda.empty_cache()
     logging.info("🫧 Building VLLM Pipeline ...")
-    llm = hydra.utils.instantiate(cfg.gen.llm, tensor_parallel_size=torch.cuda.device_count())
+    llm = hydra.utils.instantiate(cfg.gen.llm)
 
     logging.info("🎉 And it's done!")