update profile.sh

AnswerDotAI · May 23, 2024 · 6e438fd · 6e438fd
1 parent 175cf8f
commit 6e438fd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/profile.sh b/profile.sh
@@ -48,7 +48,6 @@
 #Tested on `python=3.11.9 and torch=2.3.0``
 
 python train.py \
---world_size -1 \
 --model_name "meta-llama/Llama-2-7b-hf" \
 --gradient_accumulation_steps 2 \
 --batch_size 1 \

diff --git a/train.py b/train.py
@@ -1104,7 +1104,8 @@ def fsdp_qlora(
     warmup_steps: int = 0, # Warmup steps when running profiler
     active_steps: int = 5,  # Active steps when running profiler
     repeat: int = 1, #Number of profiler cycles (wait + warmup + active)
-    max_steps: int = -1, # Max number of training steps (in units of batches) per epoch. -1 means no max_steps, otherwise training loop breaks after `max_steps` each epoch.    ):
+    max_steps: int = -1, # Max number of training steps (in units of batches) per epoch. -1 means no max_steps, otherwise training loop breaks after `max_steps` each epoch.
+):
     """
     Train a model with FSDP and QLoRA/QDoRA.