Skip to content

[Main] Fix TE version check for retain_pinned_cpu_buffers in cpu offload - #4267

Merged
yaox12 merged 3 commits into
NVIDIA:mainfrom
BestJuly:lit/fix_te_version_for_cpu_offload_main
Apr 15, 2026
Merged

[Main] Fix TE version check for retain_pinned_cpu_buffers in cpu offload#4267
yaox12 merged 3 commits into
NVIDIA:mainfrom
BestJuly:lit/fix_te_version_for_cpu_offload_main

Conversation

@BestJuly

@BestJuly BestJuly commented Apr 13, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Fix incorrect TransformerEngine version gate for retain_pinned_cpu_buffers parameter in get_cpu_offload_context
  • The retain_pinned_cpu_buffers parameter was added in TE v2.10 ([PyTorch] Activation offloading refactor TransformerEngine#1762), but the version check was gating on 2.5.0, causing a TypeError for TE 2.5–2.9 users
  • Add an intermediate branch for TE 2.5–2.9 that passes double_buffering but not retain_pinned_cpu_buffers
  • PR for dev branch.

🤖 Generated with Claude Code

The retain_pinned_cpu_buffers parameter was added to TransformerEngine's
get_cpu_offload_context in TE v2.10 (PR NVIDIA#1762), not v2.5. The previous
version gate of 2.5.0 caused a TypeError for users running TE 2.5-2.9.

Add an intermediate branch for TE 2.5-2.9 that passes double_buffering
but not retain_pinned_cpu_buffers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@BestJuly
BestJuly requested review from a team as code owners April 13, 2026 05:47
@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft April 13, 2026 05:47
@github-actions

Copy link
Copy Markdown
Contributor

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

@copy-pr-bot

copy-pr-bot Bot commented Apr 13, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 13, 2026
@BestJuly BestJuly changed the title Fix TE version check for retain_pinned_cpu_buffers in cpu offload [Main] Fix TE version check for retain_pinned_cpu_buffers in cpu offload Apr 13, 2026
@BestJuly
BestJuly marked this pull request as ready for review April 13, 2026 05:50
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Apr 15, 2026
@yaox12
yaox12 enabled auto-merge April 15, 2026 01:56
@yaox12
yaox12 added this pull request to the merge queue Apr 15, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit c2d1a8f Apr 15, 2026
64 of 65 checks passed
yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026
…oad (NVIDIA#4267)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026
…oad (NVIDIA#4267)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
jon-barker pushed a commit to jon-barker/Megatron-LM that referenced this pull request Jul 10, 2026
…oad (NVIDIA#4267)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
…oad (NVIDIA#4267)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: low

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants