add llama

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

Author: akoumpa

tests/functional_tests/llm_pretrain_and_kd/L2_TP_Output_Parity_Minified.sh

@@ -23,7 +23,7 @@ KL_THRESHOLD="${KL_THRESHOLD:-1e-6}"
 
 torchrun --nproc_per_node=2 --nnodes=1 \
     tests/functional_tests/llm_pretrain_and_kd/run_tp_output_parity_minified.py \
-    --models qwen3 qwen3_seq_cls ministral3 \
+    --models qwen3 qwen3_seq_cls ministral3 llama \
     --sequence_parallel both \
     --kl_threshold "${KL_THRESHOLD}"
 

tests/functional_tests/llm_pretrain_and_kd/run_tp_output_parity_minified.py

@@ -53,11 +53,13 @@
 
 from nemo_automodel._transformers.utils import apply_cache_compatibility_patches
 from nemo_automodel.components.distributed.parallelizer import _get_parallel_plan
+from nemo_automodel.components.models.llama.model import LlamaForCausalLM as CustomLlamaForCausalLM
 from nemo_automodel.components.models.mistral3.model import Ministral3Config, Ministral3ForCausalLM
+from transformers import LlamaConfig
 from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
 from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM, Qwen3ForSequenceClassification
 
-ModelKind = Literal["qwen3", "qwen3_seq_cls", "ministral3"]
+ModelKind = Literal["qwen3", "qwen3_seq_cls", "ministral3", "llama"]
 SPMode = Literal["true", "false", "both"]
 
 
@@ -205,6 +207,22 @@ def _build_minified_model(kind: ModelKind):
         )
         return cfg, Qwen3ForSequenceClassification(cfg)
 
+    if kind == "llama":
+        cfg = LlamaConfig(
+            vocab_size=128,
+            hidden_size=64,
+            intermediate_size=256,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            max_position_embeddings=128,
+            use_cache=False,
+            tie_word_embeddings=True,
+            attention_bias=False,
+            attn_implementation="eager",
+        )
+        return cfg, CustomLlamaForCausalLM(cfg)
+
     raise ValueError(f"Unknown model kind: {kind}")
 
 
@@ -273,8 +291,8 @@ def main(argv: Sequence[str] | None = None) -> int:
     parser.add_argument(
         "--models",
         nargs="+",
-        default=["qwen3", "qwen3_seq_cls", "ministral3"],
-        choices=["qwen3", "qwen3_seq_cls", "ministral3"],
+        default=["qwen3", "qwen3_seq_cls", "ministral3", "llama"],
+        choices=["qwen3", "qwen3_seq_cls", "ministral3", "llama"],
         help="Which models to test.",
     )
     parser.add_argument(