From 7e0b58241f2b569125daa45432b84a5b4e463b63 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 1 Sep 2026 11:52:52 +0200 Subject: [PATCH 1/3] [PyTorch] Deprecate unused is_cg_capturable in parallel_cross_entropy Since #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 --- transformer_engine/pytorch/cross_entropy.py | 20 ++++++++++--------- .../pytorch/triton/cross_entropy.py | 2 -- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 2e9ed17f0b..7aa4991224 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -31,7 +31,6 @@ def forward( reduce_loss=False, dist_process_group=None, ignore_idx=-100, - is_cg_capturable=False, overwrite_input=False, ): """Compute the loss and save the input and softmax statistics for backward.""" @@ -62,7 +61,6 @@ def forward( ctx.rank = rank ctx.world_size = world_size ctx.ignore_idx = ignore_idx - ctx.is_cg_capturable = is_cg_capturable ctx.overwrite_input = overwrite_input ctx.did_backward = False return loss @@ -90,11 +88,10 @@ def backward(ctx, grad_output): ctx.rank, ctx.world_size, ctx.ignore_idx, - ctx.is_cg_capturable, ) if ctx.overwrite_input: torch.autograd.graph.increment_version(saved_input) - return grad_input, None, None, None, None, None, None, None + return grad_input, None, None, None, None, None, None def _validate_inputs( @@ -136,7 +133,6 @@ def _parallel_cross_entropy_overwrite_input( reduce_loss: bool, dist_process_group: Optional[torch.distributed.ProcessGroup], ignore_idx: int, - is_cg_capturable: bool, ) -> torch.Tensor: """Run destructive cross entropy outside Torch Dynamo's compiled graph.""" @@ -147,7 +143,6 @@ def _parallel_cross_entropy_overwrite_input( reduce_loss, dist_process_group, ignore_idx, - is_cg_capturable, True, ) @@ -191,7 +186,9 @@ def parallel_cross_entropy( ignore_idx : int, default = -100 Target value for ignored rows. is_cg_capturable : bool, default = False - Whether the operation is CUDA graph capturable. + Deprecated and unused. The operation is always CUDA graph capturable + since the backward pass no longer contains a synchronization point. + Kept for backward compatibility; will be removed in a future release. overwrite_input : bool, default = False Allow ``inp`` to be overwritten during backward. The input must be contiguous and cannot be reused afterward. This mode is incompatible with @@ -210,6 +207,13 @@ def parallel_cross_entropy( ) inp = _input + if is_cg_capturable: + warnings.warn( + "The 'is_cg_capturable' parameter is deprecated and has no effect. " + "The operation is always CUDA graph capturable.", + FutureWarning, + ) + _validate_inputs(inp, target, label_smoothing, overwrite_input) if overwrite_input: return _parallel_cross_entropy_overwrite_input( @@ -219,7 +223,6 @@ def parallel_cross_entropy( reduce_loss, dist_process_group, ignore_idx, - is_cg_capturable, ) return CrossEntropyFunction.apply( inp, @@ -228,6 +231,5 @@ def parallel_cross_entropy( reduce_loss, dist_process_group, ignore_idx, - is_cg_capturable, overwrite_input, ) diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index 03c0bc30bb..fda28474c2 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -126,11 +126,9 @@ def cross_entropy_backward( rank: int, world_size: int, ignore_idx: int, - is_cg_capturable: bool = False, ): """Reconstruct the derivative in FP32 and overwrite the saved input buffer.""" - del is_cg_capturable B, SQ, V = saved_input.shape n_rows = B * SQ BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) From 751d0ea4a95456b31453306d6c34cdd314b77ac6 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 1 Sep 2026 11:58:39 +0200 Subject: [PATCH 2/3] Simplify deprecation note in docstring Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/cross_entropy.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 7aa4991224..bd8228362e 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -186,9 +186,8 @@ def parallel_cross_entropy( ignore_idx : int, default = -100 Target value for ignored rows. is_cg_capturable : bool, default = False - Deprecated and unused. The operation is always CUDA graph capturable - since the backward pass no longer contains a synchronization point. - Kept for backward compatibility; will be removed in a future release. + Deprecated and unused. The operation is always CUDA graph capturable. + Will be removed in a future release. overwrite_input : bool, default = False Allow ``inp`` to be overwritten during backward. The input must be contiguous and cannot be reused afterward. This mode is incompatible with From 6e870881698c69c4dc118d93af5a34d3fc7cd751 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 1 Sep 2026 12:00:56 +0200 Subject: [PATCH 3/3] Drop removal promise from deprecation note Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/cross_entropy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index bd8228362e..ebbf30b965 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -187,7 +187,6 @@ def parallel_cross_entropy( Target value for ignored rows. is_cg_capturable : bool, default = False Deprecated and unused. The operation is always CUDA graph capturable. - Will be removed in a future release. overwrite_input : bool, default = False Allow ``inp`` to be overwritten during backward. The input must be contiguous and cannot be reused afterward. This mode is incompatible with