Skip to content

fix(recipes): remove deprecated fully_shard() kwargs#1585

Open
svc-bionemo wants to merge 2 commits into
NVIDIA-BioNeMo:mainfrom
svc-bionemo:svc-bionemo/fix-nightly-20260601-fully-shard-api
Open

fix(recipes): remove deprecated fully_shard() kwargs#1585
svc-bionemo wants to merge 2 commits into
NVIDIA-BioNeMo:mainfrom
svc-bionemo:svc-bionemo/fix-nightly-20260601-fully-shard-api

Conversation

@svc-bionemo
Copy link
Copy Markdown
Collaborator

Summary

Remove deprecated fully_shard() keyword arguments that were removed in a recent megatron-core update.

Changes

  • vit/train.py: Remove grad_reduce_in_fp32 and preserve_fp32_weights from fully_shard() call
  • vit/config/defaults.yaml + vit_base_patch16_224.yaml: Remove orphaned config keys
  • esm2_native_te/hydra_config/defaults.yaml: Remove check_for_nan_in_grad, grad_reduce_in_fp32, preserve_fp32_weights
  • geneformer_native_te_mfsdp_fp8/hydra_config/defaults.yaml: Same

Root Cause

The installed megatron-core in CI has updated fully_shard() API, removing these parameters:

  • grad_reduce_in_fp32
  • preserve_fp32_weights
  • check_for_nan_in_grad

These were simple removals — no replacement mechanism needed for the affected recipes.

Fixes #1584 (partially — the nvidia-resiliency-ext version mismatch is a container image issue requiring an image rebuild)

Testing

  • Pre-commit passes on all changed files
  • Import verification confirms fully_shard loads without the removed kwargs

Remove grad_reduce_in_fp32, preserve_fp32_weights, and check_for_nan_in_grad
kwargs from fully_shard() calls in vit, esm2_native_te, and
geneformer_native_te_mfsdp_fp8 recipes. These parameters were removed in
a recent megatron-core update.

Fixes: NVIDIA-BioNeMo#1584
Signed-off-by: svc-bionemo <267129667+svc-bionemo@users.noreply.github.com>
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Jun 1, 2026

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.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented Jun 1, 2026

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 6a2f2b08-a518-438f-843c-f9f5043ef198

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@pstjohn
Copy link
Copy Markdown
Collaborator

pstjohn commented Jun 2, 2026

/ok to test 7fe5c2d

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.

Nightly CI: fully_shard() API breakage + nvidia-resiliency-ext version mismatch

2 participants