Code for Huy

NVIDIA · Jan 25, 2025 · e53cda0 · e53cda0
1 parent 79d8bd1
commit e53cda0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/vla/vla_finetune.py b/scripts/vla/vla_finetune.py
@@ -54,7 +54,7 @@
 def main(args) -> None:
 
     mbs = args.mbs
-    gbs = args.mbs *2
+    gbs = args.mbs
     decoder_seq_length = 287
     decoder_seq_length = 287
 
@@ -94,7 +94,7 @@ def main(args) -> None:
     )
 
     strategy = nl.MegatronStrategy(
-        tensor_model_parallel_size=4,
+        tensor_model_parallel_size=2,
         pipeline_model_parallel_size=1,
         pipeline_dtype=torch.bfloat16,
         sequence_parallel=False,