Set TP=2

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Oct 1, 2024 · 3597e1f · 3597e1f
1 parent 5ff5352
commit 3597e1f
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,