Skip to content

fix(fsdp2): guard uninitialized accumulated grads - #2744

Merged
akoumpa merged 1 commit into
mainfrom
akoumpa/fix/fsdp2-unsharded-param
Jun 24, 2026
Merged

fix(fsdp2): guard uninitialized accumulated grads#2744
akoumpa merged 1 commit into
mainfrom
akoumpa/fix/fsdp2-unsharded-param

Conversation

@akoumpa

@akoumpa akoumpa commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Adds a common FSDP2 guard so post-backward skips params whose lazy unsharded tensor was never created.

This fixes checkpoint robustness crashes for text-only Ministral3 runs that load through the VLM implementation and leave modality towers unused during backward.

Changelog

  • Guard FSDPParam.to_accumulated_grad_if_needed() when _unsharded_param is absent.
  • Add a focused unit test for the guard behavior and idempotence.

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? N/A

Additional Information

Short stack trace segment from the failing nemo-ci job:

/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py:504: in post_backward
    fsdp_param.to_accumulated_grad_if_needed()
/usr/local/lib/python3.12/dist-packages/torch/distributed/fsdp/_fully_shard/_fsdp_param.py:600: in to_accumulated_grad_if_needed
    or self._unsharded_param.grad is None
E   AttributeError: 'FSDPParam' object has no attribute '_unsharded_param'

Verification:

@akoumpa
akoumpa requested a review from a team as a code owner June 23, 2026 18:41
@copy-pr-bot

copy-pr-bot Bot commented Jun 23, 2026

Copy link
Copy Markdown

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.

@akoumpa
akoumpa force-pushed the akoumpa/fix/fsdp2-unsharded-param branch 2 times, most recently from 1fa2fa2 to b5b5313 Compare June 23, 2026 19:00
FSDPParam.to_accumulated_grad_if_needed = guarded


_patch_fsdp_accumulated_grad_guard()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

this modifies global behavior that downstream libraries may assume, so I'm not 100% sure will move forward with it.

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
@akoumpa
akoumpa force-pushed the akoumpa/fix/fsdp2-unsharded-param branch from b5b5313 to ced1acc Compare June 23, 2026 19:07
@akoumpa

akoumpa commented Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ced1acc

@akoumpa akoumpa added the r0.5.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge. label Jun 23, 2026
@akoumpa
akoumpa merged commit 62a9898 into main Jun 24, 2026
80 checks passed
@akoumpa
akoumpa deleted the akoumpa/fix/fsdp2-unsharded-param branch June 24, 2026 05:56
akoumpa added a commit that referenced this pull request Jun 24, 2026
…r0.5.0` (#2752)

fix(fsdp2): guard uninitialized accumulated grads (#2744)

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

r0.5.0 Auto-cherrypick to release branch. Apply before merge; cherrypick happens after merge.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant