diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py index 3d6ffd2f56e..e30cb0e1d1a 100644 --- a/megatron/core/model_parallel_config.py +++ b/megatron/core/model_parallel_config.py @@ -246,9 +246,14 @@ class ModelParallelConfig: Defaults to False. """ - cross_entropy_fusion_impl: Literal['native', 'te'] = 'native' - """If 'native', MCore based CE loss fusion is used, if 'te', Parallel CE loss - from Transformer Engine library is used. Defaults to 'native'. + cross_entropy_fusion_impl: Literal['native', 'te', 'linear'] = 'native' + """ + Specifies the implementation of cross-entropy loss fusion. + + Options: + - 'native': Uses MCore-based cross-entropy loss fusion (default). + - 'te': Uses the parallel cross-entropy loss implementation from the Transformer Engine library. + - 'linear': Uses a linear-cross-entropy fusion approach. """ tp_comm_overlap_disable_qkv: bool = False