Skip to content

[megatron] update cross_entropy_fusion_impl default value 'te' -> 'native' (#9578) #9578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jintao-Huang merged 2 commits into from
Jun 17, 2026
Merged

[megatron] update cross_entropy_fusion_impl default value 'te' -> 'native' (#9578) #9578

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7da7f9a
fix cross_entropy_fusion_impl
Jintao-Huang Jun 16, 2026
f5d54a0
fix
Jintao-Huang Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion docs/source/Megatron-SWIFT/Command-line-parameters.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@
- apply_rope_fusion: 默认为False。用于开启rope融合。该参数为megatron-core参数透传。注意：并不是所有情况都支持rope融合，例如：MLA、mrope等不支持。
- gradient_accumulation_fusion: 默认为True。用于开启梯度累加融合。
- 🔥cross_entropy_loss_fusion: 启动交叉熵损失计算融合。默认为True。
- cross_entropy_fusion_impl: 交叉熵损失融合的实现。可选为'native'和'te'。默认为None，如果是cuda设置为'te'，npu设置为'native'。
- cross_entropy_fusion_impl: 交叉熵损失融合的实现。可选为'native'和'te'。默认为'native'。
- **"ms-swift>=4.3.1"默认值从"te"修改为"native"**，原因查看[这个PR](https://github.com/NVIDIA/Megatron-LM/pull/5115)，这可能会导致更多的显存占用。
- calculate_per_token_loss: 根据全局批次中的非填充token数量来对交叉熵损失进行缩放。默认为None，`task_type`为'causal_lm'且为预训练/微调时，默认为True，否则默认为False。
- 🔥attention_backend: 使用的注意力后端 (flash、fused、unfused、local、auto、flash_2、flash_3、flash_4)。默认为 flash。
- 如果安装'flash_attention_4/3'，`--attention_backend flash`则优先使用fa4/fa3。如果要显式设置，你可以设置`--attention_backend flash_2/3/4`。训练脚本参考[这里](https://github.com/modelscope/ms-swift/tree/main/examples/train/flash_attention_3)。多模态模型的vit部分要使用flash_attention_4/3，请设置`--attn_impl flash_attention_4/3`。
Expand Down
3 changes: 2 additions & 1 deletion docs/source_en/Megatron-SWIFT/Command-line-parameters.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@
- apply_rope_fusion: Defaults to False. Used to enable RoPE fusion. This parameter is passed through from megatron-core. Note: RoPE fusion is not supported in all cases, for example: MLA, mrope, etc. are not supported.
- gradient_accumulation_fusion: Defaults to True. Used to enable gradient accumulation fusion.
- 🔥cross_entropy_loss_fusion: Enables cross-entropy loss computation fusion. Defaults to True.
- cross_entropy_fusion_impl: Implementation for cross-entropy loss fusion. Options are 'native' and 'te'. Defaults to None. Automatically set to 'te' for CUDA and 'native' for NPU.
- cross_entropy_fusion_impl: The implementation of cross-entropy loss fusion. Options are `'native'` and `'te'`. Defaults to `'native'`.
- **The default value in `ms-swift>=4.3.1` has been changed from `'te'` to `'native'`**, for the reason see [this PR](https://github.com/NVIDIA/Megatron-LM/pull/5115). This may result in higher GPU memory usage.
- calculate_per_token_loss: Scales the cross-entropy loss according to the number of non-padding tokens in the global batch. Defaults to None. When `task_type` is 'causal_lm' and during pretraining/fine-tuning, it defaults to True; otherwise, it defaults to False.
- 🔥attention_backend: The attention backend to use (flash, fused, unfused, local, auto, flash_2, flash_3, flash_4). Defaults to `flash`.
- If `flash_attention_4/3` is installed, `--attention_backend flash` will prioritize fa4/fa3. To explicitly specify a version, you can set `--attention_backend flash_2/3/4`. For training script examples, refer to [here](https://github.com/modelscope/ms-swift/tree/main/examples/train/flash_attention_3). To use flash_attention_4/3 for the ViT part of multimodal models, please set `--attn_impl flash_attention_4/3`.
Expand Down
2 changes: 1 addition & 1 deletion examples/megatron/fp8/benchmark.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ megatron sft \
--num_train_epochs 1 \
--finetune true \
--cross_entropy_loss_fusion true \
--cross_entropy_fusion_impl te \
--cross_entropy_fusion_impl native \
--lr 1e-5 \
--lr_warmup_fraction 0.05 \
--min_lr 1e-6 \
Expand Down
11 changes: 1 addition & 10 deletions swift/megatron/arguments/megatron_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,7 +440,7 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
apply_rope_fusion: bool = False
gradient_accumulation_fusion: bool = True
cross_entropy_loss_fusion: bool = True
cross_entropy_fusion_impl: Optional[Literal['native', 'te']] = None
cross_entropy_fusion_impl: Literal['native', 'te'] = 'native'

@gemini-code-assist gemini-code-assist Bot Jun 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Changing the type of cross_entropy_fusion_impl to Literal['native', 'te'] (non-optional) and removing the None check in __post_init__ can lead to runtime errors if the argument is explicitly passed as None or deserialized as None from older configurations or external parsers. Keeping it as Optional and adding a fallback to 'native' in __post_init__ is safer and more robust.

Suggested change
cross_entropy_fusion_impl: Literal['native', 'te'] = 'native'
cross_entropy_fusion_impl: Optional[Literal['native', 'te']] = 'native'

calculate_per_token_loss: Optional[bool] = None
attention_backend: str = 'flash' # flash, fused, unfused, local, auto
optimizer: Literal['adam', 'sgd', 'muon', 'dist_muon'] = 'adam'
Expand Down Expand Up @@ -720,15 +720,6 @@ def __post_init__(self):
self._init_vpp_size()
if self.vit_gradient_checkpointing is None:
self.vit_gradient_checkpointing = not self.freeze_vit
if self.cross_entropy_fusion_impl is None:
if is_torch_npu_available():
self.cross_entropy_fusion_impl = 'native'
else:
import transformer_engine
if version.parse(transformer_engine.__version__) >= version.parse('2.8.0'):
self.cross_entropy_fusion_impl = 'te'
else:
self.cross_entropy_fusion_impl = 'native'
if isinstance(self.report_to, str):
self.report_to = [self.report_to]
Comment on lines 723 to 724

@gemini-code-assist gemini-code-assist Bot Jun 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent potential runtime issues when cross_entropy_fusion_impl is explicitly set or parsed as None, we should add a defensive check in __post_init__ to default it to 'native' if it is None.

Suggested change
if isinstance(self.report_to, str):
self.report_to = [self.report_to]
if self.cross_entropy_fusion_impl is None:
self.cross_entropy_fusion_impl = 'native'
if isinstance(self.report_to, str):
self.report_to = [self.report_to]

self.model_info, self.model_meta = get_model_info_meta(
Expand Down
Loading