Skip to content

[PyTorch] Deprecate unused is_cg_capturable in parallel_cross_entropy - #41

Closed
pggPL wants to merge 3 commits into
mainfrom
cross_entropy_deprecate_cg_flag
Closed

[PyTorch] Deprecate unused is_cg_capturable in parallel_cross_entropy#41
pggPL wants to merge 3 commits into
mainfrom
cross_entropy_deprecate_cg_flag

Conversation

@pggPL

@pggPL pggPL commented Sep 1, 2026

Copy link
Copy Markdown
Owner

Description

Since NVIDIA#3273 the fused cross entropy backward reconstructs the derivative from saved statistics and no longer contains the torch.equal-based synchronization point. The is_cg_capturable flag therefore has no effect — the operation is always CUDA graph capturable.

This PR marks the flag as deprecated on the public parallel_cross_entropy API (docstring note + FutureWarning when set) and removes the dead internal plumbing. The parameter is kept in the signature for backward compatibility, since Megatron-LM passes it explicitly.

Should also be mentioned as deprecated in the v2.19 release notes.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Deprecate is_cg_capturable in parallel_cross_entropy: docstring now states it is unused, and passing True emits a FutureWarning.
  • Remove the flag from CrossEntropyFunction, _parallel_cross_entropy_overwrite_input and the Triton cross_entropy_backward wrapper, where it was already a no-op.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

Since NVIDIA#3273 the cross entropy backward no longer contains a
synchronization point, so the operation is always CUDA graph capturable
and the flag has no effect. Keep accepting it on the public API for
backward compatibility (with a FutureWarning), drop the dead internal
plumbing.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
@pggPL pggPL closed this Sep 1, 2026
@pggPL pggPL reopened this Sep 1, 2026
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
@pggPL

pggPL commented Sep 1, 2026

Copy link
Copy Markdown
Owner Author

Superseded by NVIDIA#3455

@pggPL pggPL closed this Sep 1, 2026
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.

1 participant