Skip to content

[bugfix][megatron] align GPU/NPU GKD precision - #10102

Draft
xinyu-0877 wants to merge 3 commits into
modelscope:mainfrom
xinyu-0877:gkd-fp32-grad-precision
Draft

[bugfix][megatron] align GPU/NPU GKD precision#10102
xinyu-0877 wants to merge 3 commits into
modelscope:mainfrom
xinyu-0877:gkd-fp32-grad-precision

Conversation

@xinyu-0877

Copy link
Copy Markdown

What does this PR do?

This PR aligns the numerical precision of Megatron GKD training between NVIDIA GPU
and Ascend NPU backends.

The changes include:

  • Run the GKD JSD loss path in FP32 when SWIFT_GKD_JSD_FP32=1.
  • Cast logits before top-k gather and vocabulary alignment to avoid mixed-precision
    differences between GPU and NPU.
  • Correctly forward Swift's accumulate_allreduce_grads_in_fp32 option to the
    corresponding Megatron DDP fields:
    • grad_reduce_in_fp32
    • reduce_scatter_with_fp32_accumulation
  • Log the effective Megatron DDP gradient precision configuration.

The default behavior remains unchanged when the FP32 options are disabled.

Motivation

The original GKD JSD loss was computed using the logits' original low precision
(e.g. BF16) without explicit FP32 conversion. Numerical differences in
low-precision operations, such as softmax, KL divergence, tensor-parallel gather,
and reduction, can accumulate differently across GPU and NPU backends.

In addition, accumulate_allreduce_grads_in_fp32 was not correctly forwarded
from Swift arguments to the Megatron DDP configuration, so gradient accumulation
and reduction were not guaranteed to use FP32 as intended.

This PR explicitly enables FP32 for the GKD JSD loss path and correctly forwards
the FP32 gradient accumulation/reduction configuration, improving numerical
alignment between GPU and NPU Megatron GKD training.

Validation

  • Tested the uploaded branch on the target server.
  • Compared GPU and NPU Megatron GKD training results.
  • Verified that the loss and training behavior are consistent with the reference
    branch.
  • Verified that the configured Megatron DDP FP32 gradient settings are applied.
  • Relevant syntax checks and GKD smoke tests passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant