TP=2

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
NVIDIA · Sep 25, 2024 · 465356e · 465356e
1 parent c337c1e
commit 465356e
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
@@ -167,7 +167,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,