Add calculate_per_token_loss support for Context Parallel by BoxiangW · Pull Request #14065 · NVIDIA-NeMo/NeMo

BoxiangW · 2025-06-30T00:22:00Z

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Collection: [Note which collection this PR will affect]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you add or update any necessary documentation?
Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
- Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

New Feature
Bugfix
Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Related to # (issue)

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com>

akoumpa

LGTM, thanks.

github-actions · 2025-07-16T02:12:06Z

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

github-actions · 2025-07-23T02:12:22Z

This PR was closed because it has been inactive for 7 days since being marked as stale.

…n_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com>

BoxiangW · 2025-07-29T21:34:49Z

/ok to test 843bc58

yspMing · 2025-07-30T04:02:02Z

This works after setting below configuraton for ddp_config and model config for sft training with CP, thanks a lot~

setattr(trainer.strategy.ddp_config, 'average_in_collective', False)
setattr(model.config, 'calculate_per_token_loss', True)

github-actions · 2025-07-30T10:46:11Z

[🤖]: Hi @BoxiangW 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

//cc @chtruong814 @ko3n1g @pablo-garay @thomasdhc

chtruong814 · 2025-07-31T21:36:38Z

Moving forward with merging to ensure this important fix is included.

…o#14065) * Add calculate_per_token_loss for CP Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Add check for average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> * Fix CP issue by adding check on trainer.strategy.ddp_config.average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> --------- Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> Co-authored-by: BoxiangW <BoxiangW@users.noreply.github.com>

* Add calculate_per_token_loss for CP Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Add check for average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> * Fix CP issue by adding check on trainer.strategy.ddp_config.average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> --------- Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> Co-authored-by: BoxiangW <BoxiangW@users.noreply.github.com> Signed-off-by: Charlie Truong <chtruong@nvidia.com>

#14065) (#14282) into `r2.4.0` (#14448) * Add cp per token loss check (#14282) * add cp per token loss check Signed-off-by: Chen Cui <chcui@nvidia.com> * Apply isort and black reformatting Signed-off-by: cuichenx <cuichenx@users.noreply.github.com> --------- Signed-off-by: Chen Cui <chcui@nvidia.com> Signed-off-by: cuichenx <cuichenx@users.noreply.github.com> Co-authored-by: cuichenx <cuichenx@users.noreply.github.com> Signed-off-by: Charlie Truong <chtruong@nvidia.com> * Add calculate_per_token_loss support for Context Parallel (#14065) * Add calculate_per_token_loss for CP Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Add check for average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> * Fix CP issue by adding check on trainer.strategy.ddp_config.average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> --------- Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> Co-authored-by: BoxiangW <BoxiangW@users.noreply.github.com> Signed-off-by: Charlie Truong <chtruong@nvidia.com> --------- Signed-off-by: Chen Cui <chcui@nvidia.com> Signed-off-by: cuichenx <cuichenx@users.noreply.github.com> Signed-off-by: Charlie Truong <chtruong@nvidia.com> Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> Co-authored-by: Chen Cui <chcui@nvidia.com> Co-authored-by: cuichenx <cuichenx@users.noreply.github.com> Co-authored-by: Boxiang Wang <boxiangw@nvidia.com> Co-authored-by: BoxiangW <BoxiangW@users.noreply.github.com>

…o#14065) * Add calculate_per_token_loss for CP Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Add check for average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> * Fix CP issue by adding check on trainer.strategy.ddp_config.average_in_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> * Apply isort and black reformatting Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> --------- Signed-off-by: Boxiang Wang <boxiangw@nvidia.com> Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com> Co-authored-by: BoxiangW <BoxiangW@users.noreply.github.com> Signed-off-by: Guyue Huang <guyueh@nvidia.com>

BoxiangW added 2 commits June 29, 2025 17:18

Add calculate_per_token_loss for CP

f5fc30d

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

Add check for average_in_collective

c70563d

Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW self-assigned this Jun 30, 2025

Apply isort and black reformatting

aa15586

Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com>

BoxiangW enabled auto-merge (squash) June 30, 2025 19:08

BoxiangW added the Run CICD label Jun 30, 2025

BoxiangW temporarily deployed to test June 30, 2025 19:10 — with GitHub Actions Inactive

akoumpa previously approved these changes Jun 30, 2025

View reviewed changes

github-actions bot removed the Run CICD label Jul 1, 2025

github-actions bot added the stale label Jul 16, 2025

github-actions bot closed this Jul 23, 2025

auto-merge was automatically disabled July 23, 2025 02:12
Pull request was closed

BoxiangW reopened this Jul 29, 2025

BoxiangW removed the stale label Jul 29, 2025

BoxiangW added 2 commits July 29, 2025 12:13

Merge branch 'main' into boxiangw/cp-per-token-loss

8b104c6

Fix CP issue by adding check on trainer.strategy.ddp_config.average_i…

5fa9e39

…n_collective Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>

BoxiangW dismissed akoumpa’s stale review via 5fa9e39 July 29, 2025 19:18

Apply isort and black reformatting

843bc58

Signed-off-by: BoxiangW <BoxiangW@users.noreply.github.com>

BoxiangW added the Run CICD label Jul 29, 2025

BoxiangW had a problem deploying to test July 29, 2025 19:21 — with GitHub Actions Error

akoumpa approved these changes Jul 29, 2025

View reviewed changes

akoumpa added Run CICD and removed Run CICD labels Jul 29, 2025

akoumpa temporarily deployed to test July 29, 2025 21:48 — with GitHub Actions Inactive

github-actions bot removed the Run CICD label Jul 30, 2025

chtruong814 approved these changes Jul 31, 2025

View reviewed changes

chtruong814 added the r2.4.0 Pick this label for auto-cherry-picking into r2.4.0 label Jul 31, 2025

chtruong814 merged commit 681601e into main Jul 31, 2025
259 of 266 checks passed

chtruong814 deleted the boxiangw/cp-per-token-loss branch July 31, 2025 21:36

chtruong814 mentioned this pull request Aug 10, 2025

Cherrypick calculate_per_token_loss requirement for context parallel (#14065) (#14282) into r2.4.0 #14448

Merged

8 tasks

ananthsub mentioned this pull request Aug 11, 2025

Add calculate_per_token_loss support for Context Parallel NVIDIA-NeMo/Megatron-Bridge#374

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add calculate_per_token_loss support for Context Parallel#14065

Add calculate_per_token_loss support for Context Parallel#14065
chtruong814 merged 6 commits intomainfrom
boxiangw/cp-per-token-loss

BoxiangW commented Jun 30, 2025

Uh oh!

akoumpa left a comment

Uh oh!

github-actions bot commented Jul 16, 2025

Uh oh!

github-actions bot commented Jul 23, 2025

Uh oh!

BoxiangW commented Jul 29, 2025

Uh oh!

yspMing commented Jul 30, 2025

Uh oh!

github-actions bot commented Jul 30, 2025

Uh oh!

chtruong814 commented Jul 31, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

BoxiangW commented Jun 30, 2025

What does this PR do ?

Changelog

Usage

GitHub Actions CI

Before your PR is "Ready for review"

Who can review?

Additional Information

Uh oh!

akoumpa left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 16, 2025

Uh oh!

github-actions bot commented Jul 23, 2025

Uh oh!

BoxiangW commented Jul 29, 2025

Uh oh!

yspMing commented Jul 30, 2025

Uh oh!

github-actions bot commented Jul 30, 2025

Uh oh!

chtruong814 commented Jul 31, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants