[automodel] fallback FP8 + LCE -> FP8 + CE (#13349)

* fix

Signed-off-by: Alexandros Koumparoulis <[email protected]>

* make fp8 tests non-optional

Signed-off-by: Alexandros Koumparoulis <[email protected]>

* switch to gemma

Signed-off-by: Alexandros Koumparoulis <[email protected]>

---------

Signed-off-by: Alexandros Koumparoulis <[email protected]>
Co-authored-by: oliver könig <[email protected]>

Author: akoumpa

.github/workflows/cicd-main-automodel.yml

@@ -84,7 +84,6 @@ jobs:
             script: L2_HF_Transformer_PEFT_2gpu_FSDP2_liger
           - runner: azure-gpu-vm-runner1-h100
             script: L2_HF_Transformer_PEFT_2gpu_FSDP2_fp8
-            is_optional: true
           - runner: self-hosted-azure
             script: L2_HF_Transformer_PEFT_2gpu_FSDP2
           - runner: self-hosted-azure
@@ -95,7 +94,6 @@ jobs:
             script: L2_HF_Transformer_SFT_2gpu_FSDP2
           - runner: azure-gpu-vm-runner1-h100
             script: L2_HF_Transformer_SFT_2gpu_FSDP2_fp8
-            is_optional: true
           - runner: self-hosted-azure
             script: L2_HF_Transformer_SFT_2gpu_nemorun
           - runner: self-hosted-azure

nemo/collections/llm/gpt/model/hf_auto_model_for_causal_lm.py

@@ -265,6 +265,14 @@ def configure_model(self):
 
             te_accelerate(self.model, self.model_accelerator.fp8_autocast)
 
+        if self.use_linear_ce_loss:
+            # scan the model for fp8 layers, if found disable lce
+            for module in self.model.modules():
+                if hasattr(module, 'fp8'):
+                    logging.warning("LCE does not support FP8, switching to regular CE.")
+                    self.use_linear_ce_loss = False
+                    break
+
         if self.enable_grad_ckpt:
             if getattr(self.model, 'supports_gradient_checkpointing', False):
                 self.model.gradient_checkpointing_enable()

tests/functional_tests/L2_HF_Transformer_PEFT_2gpu_FSDP2_fp8.sh

@@ -1,7 +1,7 @@
 export TRANSFORMERS_OFFLINE=1
 export HF_HOME=/home/TestData/automodel/hf_home
 coverage run -a --data-file=/workspace/.coverage --source=/workspace/nemo examples/llm/peft/automodel.py \
-    --model /home/TestData/akoumparouli/hf_mixtral_2l/ \
+    --model /home/TestData/akoumparouli/hf_gemma_38m/ \
     --max-steps 3 \
     --devices 2 \
     --strategy fsdp2 --fp8