Fix quantize.py script and support packed sequences in pretrain_gpt.py by AAnoosheh · Pull Request #3564 · NVIDIA/Megatron-LM

AAnoosheh · 2026-02-24T17:20:25Z

What does this PR do ?

Fix quantize.py example script and bug with eval logging
Allow pretrain_gpt.py to support packed sequences, bringing it up to par with pretrain_mamba
Fix bug with CP in both training eval logging and Packed-Sequence data initialization

⚠️ 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 the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

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

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2026-02-24T17:20:30Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

AAnoosheh · 2026-03-07T00:03:39Z

/ok to test f43f37c

- Add is_packed_sequence flag (True when args.sft is set) to get_batch and is_dataset_built_on_rank, mirroring pretrain_mamba.py behavior - Middle PP stages with packed sequences return (None×5, PackedSeqParams) carrying cu_seqlens/max_seqlen for attention masking, instead of early- returning None×6 - Import PackedSeqParams from megatron.core.packed_seq_params - is_dataset_built_on_rank: packed sequences build dataset on all TP rank-0 stages (not just first/last PP), consistent with THD format requirements - Add detailed docstring explaining packed sequence flow and differences from pretrain_mamba.py (return format, middle-stage handling, CP, MTP) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

AAnoosheh · 2026-03-09T15:46:05Z

/ok to test 0e375ba

AAnoosheh · 2026-03-11T15:46:43Z

/ok to test 2acb2a8

AAnoosheh · 2026-03-11T16:11:51Z

/ok to test 326dac9

cu_seqlens_q_padded may not be strictly monotonic when context parallelism slices sequences across ranks, or when padded cumulative lengths exceed total_tokens. The diff can go negative, causing torch.repeat_interleave to fail with "repeats can not be negative". Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

AAnoosheh · 2026-03-12T10:59:54Z

/ok to test 1e064f3

AAnoosheh · 2026-03-12T17:03:12Z

/ok to test 1e064f3

chtruong814 · 2026-03-13T03:06:19Z

/ok to test ae7cd6b

svcnvidia-nemo-ci · 2026-03-16T21:40:22Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23167215405

svcnvidia-nemo-ci · 2026-03-16T23:01:40Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23169985301

svcnvidia-nemo-ci · 2026-03-16T23:06:18Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23170131822

NVIDIA#3564) Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com> Co-authored-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com>

NVIDIA#3564) Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com> Co-authored-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

NVIDIA#3564) Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com> Co-authored-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

NVIDIA#3564) Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com> Co-authored-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Charlie Truong <chtruong@nvidia.com>

AAnoosheh self-assigned this Feb 24, 2026

AAnoosheh requested a review from ChenhanYu February 24, 2026 17:20

ChenhanYu approved these changes Feb 24, 2026

View reviewed changes

AAnoosheh force-pushed the aanoosheh/modelopt-example-import-order branch from 27a6b72 to bae176b Compare February 25, 2026 12:56

AAnoosheh marked this pull request as ready for review March 6, 2026 23:55

svcnvidia-nemo-ci requested a review from a team March 6, 2026 23:55

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Mar 6, 2026

AAnoosheh added 2 commits March 6, 2026 16:02

Fix import order

17a8206

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

Prevent NaNs in CP when rank has all zero mask

f43f37c

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

AAnoosheh force-pushed the aanoosheh/modelopt-example-import-order branch from bae176b to f43f37c Compare March 7, 2026 00:02

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 7, 2026

copy-pr-bot Bot temporarily deployed to test March 7, 2026 00:05 Inactive

copy-pr-bot Bot temporarily deployed to test March 9, 2026 15:47 Inactive

AAnoosheh added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Mar 9, 2026

AAnoosheh requested a review from jaredcasper March 9, 2026 18:13

Phlip79 removed the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Mar 10, 2026

chg: update

2acb2a8

AAnoosheh changed the title ~~Fix quantize.py example script and bug with eval logging~~ Fix quantize.py script and support packed sequences in pretrain_gpt.py Mar 10, 2026

jenchen13 reviewed Mar 10, 2026

View reviewed changes

Comment thread examples/post_training/modelopt/convert_model.py Outdated

copy-pr-bot Bot temporarily deployed to test March 11, 2026 15:47 Inactive

Merge branch 'main' into aanoosheh/modelopt-example-import-order

326dac9

copy-pr-bot Bot temporarily deployed to test March 11, 2026 16:13 Inactive

ChenhanYu requested a review from a team as a code owner March 12, 2026 00:04

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage and removed Approved All necessary approvals have been made labels Mar 12, 2026

chg: update restore changes to finetune.py

1e064f3

copy-pr-bot Bot temporarily deployed to test March 12, 2026 11:00 Inactive

yaoyu-33 approved these changes Mar 12, 2026

View reviewed changes

jaredcasper approved these changes Mar 12, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Mar 12, 2026

Merge branch 'main' into aanoosheh/modelopt-example-import-order

ae7cd6b

copy-pr-bot Bot temporarily deployed to test March 13, 2026 03:07 Inactive

chtruong814 added this pull request to the merge queue Mar 16, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Mar 16, 2026

chtruong814 added this pull request to the merge queue Mar 16, 2026

Merged via the queue into NVIDIA:main with commit f89744b Mar 17, 2026
77 of 81 checks passed

asolergi-nv mentioned this pull request Mar 24, 2026

megatron.core.datasets SFTDataset and get_batch* functions refactor #3542

Closed

8 tasks

AAnoosheh deleted the aanoosheh/modelopt-example-import-order branch March 25, 2026 23:45

xiaoyao0115 mentioned this pull request Apr 9, 2026

Minor improvements for Dynamic-cp #4226

Merged

5 tasks

Navigation Menu

Uh oh!

Fix quantize.py script and support packed sequences in pretrain_gpt.py - #3564

Fix quantize.py script and support packed sequences in pretrain_gpt.py#3564
chtruong814 merged 8 commits into
NVIDIA:mainfrom
AAnoosheh:aanoosheh/modelopt-example-import-order

AAnoosheh commented Feb 24, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Feb 24, 2026

Uh oh!

AAnoosheh commented Mar 7, 2026

Uh oh!

AAnoosheh commented Mar 9, 2026

Uh oh!

Uh oh!

AAnoosheh commented Mar 11, 2026

Uh oh!

AAnoosheh commented Mar 11, 2026

Uh oh!

AAnoosheh commented Mar 12, 2026

Uh oh!

AAnoosheh commented Mar 12, 2026

Uh oh!

chtruong814 commented Mar 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

AAnoosheh commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Feb 24, 2026

Uh oh!

AAnoosheh commented Mar 7, 2026

Uh oh!

AAnoosheh commented Mar 9, 2026

Uh oh!

Uh oh!

AAnoosheh commented Mar 11, 2026

Uh oh!

AAnoosheh commented Mar 11, 2026

Uh oh!

AAnoosheh commented Mar 12, 2026

Uh oh!

AAnoosheh commented Mar 12, 2026

Uh oh!

chtruong814 commented Mar 13, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

AAnoosheh commented Feb 24, 2026 •

edited

Loading

(Step 1): Add PR label `Expert Review`