Add cp per token loss check (#14282)

* add cp per token loss check

Signed-off-by: Chen Cui <[email protected]>

* Apply isort and black reformatting

Signed-off-by: cuichenx <[email protected]>

---------

Signed-off-by: Chen Cui <[email protected]>
Signed-off-by: cuichenx <[email protected]>
Co-authored-by: cuichenx <[email protected]>
Signed-off-by: Charlie Truong <[email protected]>

Author: cuichenx

nemo/collections/llm/api.py

@@ -35,6 +35,7 @@
     EvaluationTarget,
     MisconfigurationError,
 )
+from nemo.collections.llm.gpt.data.fine_tuning import FineTuningDataModule
 from nemo.collections.llm.modelopt import (
     DistillationGPTModel,
     ExportConfig,
@@ -1359,6 +1360,10 @@ def _validate_config(
                     assert (
                         model.config.seq_length % (trainer.strategy.context_parallel_size * 2) == 0
                     ), 'Sequence length must be divisible by 2 * context parallel size if context parallel is used.'
+                if isinstance(data, FineTuningDataModule):
+                    assert model.config.calculate_per_token_loss, (
+                        "When finetuning with CP>1, " "model.config.calculate_per_token_loss must be True"
+                    )
 
         # EP validation
         if trainer.strategy.expert_model_parallel_size > 1: