Skip to content

Make gradient clipping reproducible without changing configuration names - #21

Draft
zrr1999 wants to merge 1 commit into
PFCCLab:mainfrom
zrr1999:fix/restore-accuracy-mode-20260916
Draft

zrr1999 wants to merge 1 commit into
PFCCLab:mainfrom
zrr1999:fix/restore-accuracy-mode-20260916

Conversation

@zrr1999

@zrr1999 zrr1999 commented Sep 16, 2026

Copy link
Copy Markdown
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Make positive-threshold gradient clipping reproducible under the existing accuracy-compatible mode. Accumulate device-computed FP32 squares in integer bins across the existing dense/expert ownership groups, then round once before taking the square root. A missing fused tensor-scale kernel uses the existing local scaling implementation.

Keep the existing configuration names and optimizer-selection behavior, including norm_accuracy_compatible, router_accuracy_compatible, dsa_accuracy_compatible, and native_unfused_adamw. Read the existing USE_ACCURACY_COMPATIBLE mode; no switch unification or model-path rewrite is included. Ordinary mode and zero-threshold clipping retain their existing norm path.

Validation: 14 focused tests passed on each of two GPU ranks, covering FP32 rounding, non-finite values, partition/layout invariance, dense/expert ownership groups, actual clipping, disabled clipping, and a native optimizer update. Black/isort and diff checks passed. Pylint returned zero with an existing bad-builtins configuration warning. Full remote precision CI has not passed; the paired local results and installation blocker are detailed below.

The companion Swift change must preserve the requested clipping threshold. This PR does not change CI workflows, dependency sources, model case scripts, or comparison tolerances.

Paired local validation uses Megatron 6b4771e50, Swift d3d0b1a87, and Fleet d889f0a1: the original GLM52 CI profile passed 100 bitwise-identical loss steps and 187 identical canonical checkpoint tensors; the acceptance-profile native entrypoints separately passed 100 steps with strict loss, provenance and checkpoint oracles. Both sides also match their prior local baselines. The original two-rank GLM45 10-step regression with clip_grad=1.0 passed all 40 raw per-token/final-loss hash records and both same-side baseline comparisons.

These runs use local Paddle 3.4.0.post20260808+733f3454aa0 and Torch 2.12.1+cu129; they do not establish remote CI equivalence. The current PR precision job failed before model execution because Fleet develop cannot resolve its pinned Paddle 3.4.0.post20260907+4d8f64b7049. The dependency-index fix is in Fleet NVIDIA#1961. Evidence: experiments/ops/restore-pair-20260916/{ci-terminal,terminal}.json and experiments/ops/restore-pair-regression-20260916/glm45-pair-r0/{protocol,source-binding,comparison}.json.

Companion PR: PFCCLab/ms-swift#14.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

  • New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
  • Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Related: #17, #20 and PaddleFleet #1961.

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

  1. When your PR is ready, click Ready for Review.
  2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
    • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

Signed-off-by: Zhan Rongrui <me@zrr.dev>

@morirun morirun left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

温和核对(Draft):

  1. 把可复现 L2 norm 收成 get_reproducible_grad_norm_bins + ReproducibleL2Norm,并经 use_accuracy_compatible and clip_grad > 0 门控,方向清楚;非 L2 直接 ValueError 也合理。
  2. clip_grad_by_total_norm_fp32 在缺 multi_tensor_scale_tensor_impl 时回退到 clip_coeff.item() 的路径值得留着,避免 accuracy 路径在部分构建上硬断言。
  3. 请确认空 grads_for_normgrad_stats_parallel_group is None 时与旧 get_grad_norm_fp32 行为一致(当前空列表走 device=None 的 zeros,看起来有兜底)。
  4. PFCCLab/ms-swift#14 是配对变更:合并顺序或版本门槛最好在说明里写清,避免下游先合后炸。

当前 alignment_model_accuracy 红,Draft 状态下先把该检查或依赖说明对齐即可。不代推。

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.

2 participants