Set TP=2

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Sep 30, 2024 · 533ce3d · 533ce3d
1 parent fd82d76
commit 533ce3d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/llm/recipes/mistral_nemo_12b.py b/nemo/collections/llm/recipes/mistral_nemo_12b.py
@@ -169,7 +169,7 @@ def pretrain_recipe(
         fn,
         model=model(),
         trainer=trainer(
-            tensor_parallelism=1,
+            tensor_parallelism=2,
             pipeline_parallelism=1,
             pipeline_parallelism_type=None,
             virtual_pipeline_parallelism=None,