fix bugs

Signed-off-by: dimapihtar <dpihtar@gmail.com>

nemo/collections/llm/recipes/__init__.py

@@ -5,7 +5,6 @@
     llama3_70b,
     llama3_70b_16k,
     llama3_70b_64k,
-    mistral,
     mistral_7b,
     mistral_nemo_12b,
     mixtral_8x3b,

nemo/collections/llm/recipes/mistral_nemo_12b.py

@@ -27,7 +27,7 @@
 from nemo.collections.llm.gpt.model.mistral import MistralModel, MistralNeMoConfig12B
 from nemo.collections.llm.recipes.log.default import default_log, default_resume, tensorboard_logger
 from nemo.collections.llm.recipes.optim.adam import distributed_fused_adam_with_cosine_annealing
-from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed
+from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed_plugin
 from nemo.utils.exp_manager import TimingCallback
 
 NAME = "mistral_nemo_base_12b"
@@ -117,7 +117,7 @@ def trainer(
         log_every_n_steps=10,
         max_steps=max_steps,
         num_nodes=num_nodes,
-        plugins=bf16_mixed(),
+        plugins=bf16_mixed_plugin(),
         strategy=strategy,
         use_distributed_sampler=False,
         val_check_interval=2000,
@@ -165,7 +165,7 @@ def pretrain_recipe(
             callbacks=[run.Config(TimingCallback)],
         ),
         data=run.Config(MockDataModule, seq_length=4096, global_batch_size=512, micro_batch_size=1),
-        log=default_log(dir=dir, name=name, tensorboard_logger=tensorboard_logger(name=name)),
+        log=default_log(ckpt_dir=dir, name=name, tensorboard_logger=tensorboard_logger(name=name)),
         optim=distributed_fused_adam_with_cosine_annealing(max_lr=3e-4),
         resume=default_resume(),
     )