[ksivaman]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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:

• Change A
• Change B

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

[ksivaman]

/te-ci

[pggPL]

I've looked on all the changes and they all LGTM, I run subagents throughout TE to clean some mess and they suggested following changes:

With PyTorch 2.8 now required, there are a few remaining compatibility cleanups worth considering:

1. In transformer_engine/pytorch/__init__.py:136–139, replace the try/except AttributeError around error_on_nested_jit_trace with the direct assignment. The fallback is for an API added in 2.2.

2. In transformer_engine/pytorch/distributed.py:24–29, import torch.distributed._symmetric_memory directly and remove HAS_TORCH_SYMMETRIC, its checks at lines 1978 and 2042, and the corresponding skip in tests/pytorch/distributed/test_symmetric_memory.py:69. The module exists in 2.8; retain checks for actual backend capabilities, GPU count, and peer access.

3. In transformer_engine/pytorch/utils.py:34–37, import comptime directly and remove the _comptime is not None condition at line 52. Keep the torch.compiler.is_compiling() distinction between compile-time and eager warnings.

4. In tests/pytorch/distributed/ep_reference.py:44–48,132, replace the dtype-availability lookups with direct references to torch.float8_e4m3fn, torch.float8_e8m0fnu, and torch.float4_e2m1fn_x2. These attributes exist in 2.8. Preserve the format/dtype validation and support for packed uint8 NVFP4 data.

5. For the CUDA/NVIDIA tests in tests/pytorch/attention/test_cp_utils.py:1022,1043, remove the skips for a missing shared_memory_per_block property. It is exposed by PyTorch 2.8 CUDA device properties. This is a CUDA-specific guarantee, not a guarantee for every PyTorch backend.

6. Use the public import paths available in 2.8: torch.distributed.tensor.DTensor in transformer_engine/pytorch/optimizers/{fused_adam,multi_tensor_apply}.py, and torch.distributed.fsdp.fully_shard in the three tests/pytorch/distributed/fsdp2_tests/run_fsdp2_*.py files and examples/pytorch/quantized_model_init/fully_shard.py. Also update the fused_adam.py:557 docstring to reference torch.amp.GradScaler.

7. The graph-safe RNG compatibility branches can also be simplified, but I would prefer a separate follow-up because they span RNG tracking, checkpointing, and CUDA Graph capture. The APIs checked by graph_safe_rng_available() are present in 2.8. Remove that availability check and the unreachable old-version branches in distributed.py:95,114,142,886–901,925–927, graph.py:485,1715–1722,1743–1747, and the assertion in attention/dot_product_attention/dot_product_attention.py:2333. Preserve the explicit graph_safe=False path and RNG state save/restore behavior; API availability does not make those paths obsolete. Validate the cleanup with focused checkpointing and CUDA Graph RNG/replay tests.

8. Optional wording cleanup: replace the obsolete “NVFusion” description in docs/envvars.rst:409; remove the old TorchScript conditioning limitation from the docstrings in pytorch/jit.py:116,144; rename test_sanity_amp_and_nvfuser to describe AMP and internal fusions; and remove the stale bias_gelu_nvfusion reference in the commented-out FSDP code in pytorch/module/layernorm_mlp.py:1375. Keep the existing environment-variable names for compatibility. In pytorch/tensor/float8_blockwise_tensor.py:620, retain the private-API warning but drop the outdated “tested with PyTorch 2.5+” claim; do not replace it with a new testing claim without validation.

I would treat the Dynamo counters fallbacks in tests/pytorch/test_torch_compile.py:20–23 and tests/pytorch/distributed/run_layer_with_overlap.py:22–25 separately. The import exists in 2.8, but the test helper explicitly describes graceful handling of future private-API changes. These are not simply pre-2.8 compatibility branches. I would retain them unless we deliberately choose to fail tests when that API changes; if so, update the associated optional counter checks as well, rather than changing only the imports.

[greptile-apps]

Confidence Score: 5/5

[High impact] The PR appears safe to merge.

Summary

Raises the minimum supported versions to PyTorch 2.8 and JAX 0.5.3.

• Transformer Engine's PyTorch package now requires PyTorch 2.8 or newer.
• PyTorch helpers now use current autocast, compile, and export APIs directly.
• JAX builds now require JAX 0.5.3 or newer.
• NCCL communicator borrowing supports both PyTorch process-group backends.

No actionable issues found. Tests were not run.

_{Reviews (1) · Last reviewed commit: "Merge branch 'main' into update_minimum_..." · Reviewed by Greptile}