Merge remote-tracking branch 'origin/feat/to_tensorrt' into feat/to_tensorrt

GdoongMathew · GdoongMathew · commit f0762330a724 · 2025-05-24T01:38:46.000+08:00
diff --git a/.github/ISSUE_TEMPLATE/1_bug_report.yaml b/.github/ISSUE_TEMPLATE/1_bug_report.yaml
@@ -46,12 +46,22 @@ body:
     attributes:
       value: "**Note: The rest of this form is optional, but filling it out may help us to provide better support.**"
 
+  - type: input
+    attributes:
+      label: Reproduced in studio
+      description: >
+        Create a new Lightning Studio with code that reproduces the issue and share the link.
+        Also include all the relevant files and data required to reproduce shared issue.
+        In case the code does not crash, please add assert statements to show what is the real and expected output.
+        A simple guide on how to create such a studio can be found [here](https://www.youtube.com/watch?v=YcW-2Zt_bFg&ab_channel=LightningAI).
+      placeholder: https://lightning.ai/live-session/...
+    validations:
+      required: false
   - type: textarea
     attributes:
       label: How to reproduce the bug
       description: >
-        Provide steps and example code here.
-        You can also paste a link to Google Colab (see our [Colab bug report template](https://colab.research.google.com/github/Lightning-AI/lightning/blob/master/examples/pytorch/bug_report/bug_report_model.ipynb)) or adapt this minimal [snippet](https://github.com/Lightning-AI/lightning/blob/master/examples/pytorch/bug_report/bug_report_model.py).
+        In the special case when the issue can't be reproduced in a studio, provide steps and example code here.
       placeholder: |
         ```python
         # Sample code to reproduce the problem
diff --git a/.github/workflows/_legacy-checkpoints.yml b/.github/workflows/_legacy-checkpoints.yml
@@ -113,7 +113,7 @@ jobs:
 
       - run: pip install -r requirements/ci.txt
       - name: Upload checkpoints to S3
-        if: secrets.AWS_REGION != ''
+        if: ${{ secrets[AWS_REGION] != '' }}
         working-directory: ${{ env.LEGACY_FOLDER }}
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY }}
diff --git a/.github/workflows/ci-checkpoints.yml.disabled b/.github/workflows/ci-checkpoints.yml.disabled
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -94,6 +94,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        # adding dome more images as Thunder mainly using python 3.10,
+        # and we need to support integrations as for example LitGPT
+        python_version: ["3.10"]
+        pytorch_version: ["2.6.0", "2.7.0"]
+        cuda_version: ["12.6.3"]
         include:
           # These are the base images for PL release docker images.
           # Make sure the matrix here matches the one above.
diff --git a/.github/workflows/release-pkg.yml b/.github/workflows/release-pkg.yml
@@ -13,7 +13,7 @@ on:
       - "requirements/ci.txt"
       - ".github/actions/pkg-check/*"
       - ".github/actions/pkg-publish/*"
-      - ".github/workflows/_legacy-checkpoints.yml.yml"
+      - ".github/workflows/_legacy-checkpoints.yml"
       - ".github/workflows/_build-packages.yml"
       - ".github/workflows/release-pypi.yml"
 
@@ -179,12 +179,12 @@ jobs:
         with:
           pkg-folder: dist/${{ steps.folder.outputs.pkg }}
           pypi-token: ${{ secrets[format('PYPI_TOKEN_{0}', matrix.name)] }}
-
-  legacy-checkpoints:
-    needs: [build-packages]
-    uses: ./.github/workflows/_legacy-checkpoints.yml
-    with:
-      push_to_s3: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
-      upload_local: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
-      create_pr: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
-    secrets: inherit
+# FIXME: this is not working suddenly, Unrecognized named-value: 'secrets'
+#  legacy-checkpoints:
+#    needs: [build-packages]
+#    uses: ./.github/workflows/_legacy-checkpoints.yml
+#    with:
+#      push_to_s3: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
+#      upload_local: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
+#      create_pr: ${{ startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release' }}
+#    secrets: inherit
diff --git a/.gitignore b/.gitignore
@@ -198,7 +198,7 @@ node_modules/
 **/events.out.tfevents.*
 examples/**/*.png
 
-# instalation artifacts
+# installation artifacts
 requirements/base.txt
 
 # CI
diff --git a/_notebooks b/_notebooks
@@ -1 +1 @@
-Subproject commit 1c1160e543bb56760886a45dcb7e1e03a22f634c
+Subproject commit fd70f5114b21f7f970bd5587b1d3def689507069
diff --git a/requirements/ci.txt b/requirements/ci.txt
@@ -1,7 +1,7 @@
-setuptools <70.1.1
+setuptools <80.7.2
 wheel <0.46.0
 awscli >=1.30.0, <1.41.0
-twine ==6.0.1
+twine ==6.1.0
 importlib-metadata <9.0.0
 wget
 pkginfo ==1.12.1.2
diff --git a/requirements/docs.txt b/requirements/docs.txt
@@ -1,7 +1,7 @@
 sphinx >5.0, <6.0
-myst-parser >=0.18.1, <3.0.0
-nbsphinx >=0.8.5, <=0.9.2
-nbconvert <7.14  # temporary fix for https://github.com/jupyter/nbconvert/issues/2092
+myst-parser >=0.18.1, <4.0.0
+nbsphinx >=0.8.5, <=0.9.7
+nbconvert >7.14, <7.17
 pandoc >=1.0, <=2.3
 docutils >=0.16, <0.22
 sphinxcontrib-fulltoc >=1.0, <=1.2.0
@@ -12,9 +12,9 @@ sphinx-paramlinks >=0.5.1, <=0.6.0
 sphinx-togglebutton >=0.2, <=0.3.2
 sphinx-copybutton >=0.3, <=0.5.2
 sphinx-multiproject
-sphinx-toolbox ==3.5.0
+sphinx-toolbox ==4.0.0
 sphinx-rtd-dark-mode
-sphinxcontrib-video ==0.2.0
+sphinxcontrib-video ==0.4.1
 jinja2 <3.2.0
 
 lightning-utilities >=0.11.1, <0.15.0
diff --git a/requirements/fabric/test.txt b/requirements/fabric/test.txt
@@ -2,8 +2,8 @@ coverage ==7.8.0
 numpy >=1.17.2, <1.27.0
 pytest ==8.3.5
 pytest-cov ==6.1.1
-pytest-timeout ==2.3.1
-pytest-rerunfailures ==12.0
-pytest-random-order ==1.1.0
-click ==8.1.7
+pytest-timeout ==2.4.0
+pytest-rerunfailures ==15.1
+pytest-random-order ==1.1.1
+click ==8.1.8
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
diff --git a/requirements/pytorch/base.txt b/requirements/pytorch/base.txt
@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 torch >=2.1.0, <2.8.0
-tqdm >=4.57.0, <4.67.0
+tqdm >=4.57.0, <4.68.0
 PyYAML >=5.4, <6.1.0
 fsspec[http] >=2022.5.0, <2025.4.0
 torchmetrics >=0.7.0, <1.8.0
diff --git a/requirements/pytorch/docs.txt b/requirements/pytorch/docs.txt
@@ -2,6 +2,6 @@
 
 nbformat  # used for generate empty notebook
 ipython[notebook] <8.19.0
-setuptools<58.0  # workaround for `error in ipython setup command: use_2to3 is invalid.`
+setuptools<81.0  # workaround for `error in ipython setup command: use_2to3 is invalid.`
 
 #-r ../../_notebooks/.actions/requires.txt
diff --git a/requirements/pytorch/extra.txt b/requirements/pytorch/extra.txt
@@ -2,10 +2,10 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 # extended list of package dependencies to reach full functionality
-matplotlib>3.1, <3.9.0
+matplotlib>3.1, <3.10.0
 omegaconf >=2.2.3, <2.4.0
 hydra-core >=1.2.0, <1.4.0
-jsonargparse[signatures] >=4.27.7, <=4.35.0
-rich >=12.3.0, <13.6.0
+jsonargparse[signatures] >=4.28.0, <=4.40.0
+rich >=12.3.0, <14.1.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.45.2,<0.45.3; platform_system != "Darwin"
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -1,15 +1,15 @@
 coverage ==7.8.0
 pytest ==8.3.5
 pytest-cov ==6.1.1
-pytest-timeout ==2.3.1
-pytest-rerunfailures ==12.0
-pytest-random-order ==1.1.0
+pytest-timeout ==2.4.0
+pytest-rerunfailures ==15.1
+pytest-random-order ==1.1.1
 
 # needed in tests
 cloudpickle >=1.3, <3.2.0
 scikit-learn >0.22.1, <1.7.0
 numpy >=1.17.2, <1.27.0
-onnx >=1.12.0, <1.18.0
+onnx >=1.12.0, <1.19.0
 onnxruntime >=1.12.0, <1.21.0
 psutil <7.0.1 # for `DeviceStatsMonitor`
 pandas >1.0, <2.3.0  # needed in benchmarks
diff --git a/src/lightning/fabric/plugins/precision/transformer_engine.py b/src/lightning/fabric/plugins/precision/transformer_engine.py
@@ -171,7 +171,9 @@ def _convert_layers(module: torch.nn.Module) -> None:
         elif isinstance(child, torch.nn.LayerNorm):
             replacement = te.LayerNorm(child.normalized_shape[0], eps=child.eps)
             replacement.weight.data = child.weight.data.clone()
-            replacement.bias.data = child.bias.data.clone()
+            # Check if bias exists before attempting to clone its data
+            if child.bias is not None and replacement.bias is not None:
+                replacement.bias.data = child.bias.data.clone()
             log.debug(f"Replacing layer {name!r} with Transformer Engine equivalent")
             module.__setattr__(name, replacement)
         else:
diff --git a/src/lightning/fabric/utilities/imports.py b/src/lightning/fabric/utilities/imports.py
@@ -37,5 +37,3 @@
 _TORCH_LESS_EQUAL_2_6 = compare_version("torch", operator.le, "2.6.0")
 
 _PYTHON_GREATER_EQUAL_3_10_0 = (sys.version_info.major, sys.version_info.minor) >= (3, 10)
-
-_UTILITIES_GREATER_EQUAL_0_10 = compare_version("lightning_utilities", operator.ge, "0.10.0")
diff --git a/src/lightning/fabric/utilities/rank_zero.py b/src/lightning/fabric/utilities/rank_zero.py
@@ -15,8 +15,7 @@
 
 import logging
 import os
-from functools import wraps
-from typing import Callable, Optional, TypeVar, overload
+from typing import Optional
 
 import lightning_utilities.core.rank_zero as rank_zero_module
 
@@ -29,9 +28,6 @@
     rank_zero_info,
     rank_zero_warn,
 )
-from typing_extensions import ParamSpec
-
-from lightning.fabric.utilities.imports import _UTILITIES_GREATER_EQUAL_0_10
 
 rank_zero_module.log = logging.getLogger(__name__)
 
@@ -48,33 +44,7 @@ def _get_rank() -> Optional[int]:
     return None
 
 
-if not _UTILITIES_GREATER_EQUAL_0_10:
-    T = TypeVar("T")
-    P = ParamSpec("P")
-
-    @overload
-    def rank_zero_only(fn: Callable[P, T]) -> Callable[P, Optional[T]]:
-        """Rank zero only."""
-
-    @overload
-    def rank_zero_only(fn: Callable[P, T], default: T) -> Callable[P, T]:
-        """Rank zero only."""
-
-    def rank_zero_only(fn: Callable[P, T], default: Optional[T] = None) -> Callable[P, Optional[T]]:
-        @wraps(fn)
-        def wrapped_fn(*args: P.args, **kwargs: P.kwargs) -> Optional[T]:
-            rank = getattr(rank_zero_only, "rank", None)
-            if rank is None:
-                raise RuntimeError("The `rank_zero_only.rank` needs to be set before use")
-            if rank == 0:
-                return fn(*args, **kwargs)
-            return default
-
-        return wrapped_fn
-
-    rank_zero_module.rank_zero_only.rank = getattr(rank_zero_module.rank_zero_only, "rank", _get_rank() or 0)
-else:
-    rank_zero_only = rank_zero_module.rank_zero_only
+rank_zero_only = rank_zero_module.rank_zero_only
 
 # add the attribute to the function but don't overwrite in case Trainer has already set it
 rank_zero_only.rank = getattr(rank_zero_only, "rank", _get_rank() or 0)
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -24,7 +24,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed logger_connector has edge case where step can be a float ([#20692](https://github.com/Lightning-AI/pytorch-lightning/issues/20692))
 
 
 ---
diff --git a/src/lightning/pytorch/cli.py b/src/lightning/pytorch/cli.py
@@ -45,11 +45,18 @@
         Namespace,
         class_from_function,
         register_unresolvable_import_paths,
-        set_config_read_mode,
     )
 
     register_unresolvable_import_paths(torch)  # Required until fix https://github.com/pytorch/pytorch/issues/74483
-    set_config_read_mode(fsspec_enabled=True)
+
+    try:
+        from jsonargparse import set_parsing_settings
+
+        set_parsing_settings(config_read_mode_fsspec_enabled=True)
+    except ImportError:
+        from jsonargparse import set_config_read_mode
+
+        set_config_read_mode(fsspec_enabled=True)
 else:
     locals()["ArgumentParser"] = object
     locals()["Namespace"] = object
diff --git a/src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py b/src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py
@@ -106,12 +106,13 @@ def log_metrics(self, metrics: _OUT_DICT, step: Optional[int] = None) -> None:
         scalar_metrics = convert_tensors_to_scalars(metrics)
 
         if step is None:
-            step = scalar_metrics.pop("step", None)
-
-        if step is None:
-            # added metrics for convenience
-            scalar_metrics.setdefault("epoch", self.trainer.current_epoch)
-            step = self.trainer.fit_loop.epoch_loop._batches_that_stepped
+            step_metric = scalar_metrics.pop("step", None)
+            if step_metric is not None:
+                step = int(step_metric)
+            else:
+                # added metrics for convenience
+                scalar_metrics.setdefault("epoch", self.trainer.current_epoch)
+                step = self.trainer.fit_loop.epoch_loop._batches_that_stepped
 
         # log actual metrics
         for logger in self.trainer.loggers:
diff --git a/tests/tests_fabric/plugins/precision/test_transformer_engine.py b/tests/tests_fabric/plugins/precision/test_transformer_engine.py
@@ -115,3 +115,35 @@ class TELayerNormMock(Mock): ...
     assert isinstance(model.l1, TELinearMock)
     assert isinstance(model.l2, TELayerNormMock)
     assert isinstance(model.l3.l, TELinearMock)
+
+
+def test_convert_module_handles_linear_without_bias(monkeypatch):
+    module = lightning.fabric.plugins.precision.transformer_engine  # Set up mock transformer_engine
+    monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True)
+
+    transformer_engine_mock = Mock()
+    monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock)
+    monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", transformer_engine_mock.pytorch)
+    monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", transformer_engine_mock.recipe)
+
+    class TELinearMock(torch.nn.Linear):  # Mock the Linear replacement class
+        def __init__(self, in_features, out_features, bias=True):
+            super().__init__(in_features, out_features, bias)
+
+    transformer_engine_mock.pytorch.Linear = TELinearMock
+    transformer_engine_mock.pytorch.LayerNorm = torch.nn.LayerNorm
+    transformer_engine_mock.recipe.DelayedScaling.return_value = None
+
+    class BiaslessModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(16, 32, bias=False)  #  This was causing the bug
+
+    model = BiaslessModel()
+    precision = TransformerEnginePrecision(weights_dtype=torch.float16)
+    precision.replace_layers = True
+
+    precision.convert_module(model)  # This should no longer raise AttributeError
+
+    assert isinstance(model.linear, TELinearMock)
+    assert model.linear.bias is None
diff --git a/tests/tests_pytorch/loggers/test_wandb.py b/tests/tests_pytorch/loggers/test_wandb.py
@@ -26,7 +26,6 @@
 from lightning.pytorch.demos.boring_classes import BoringModel
 from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from tests_pytorch.test_cli import _xfail_python_ge_3_11_9
 
 
 def test_wandb_project_name(wandb_mock):
@@ -645,7 +644,6 @@ def test_wandb_logger_download_artifact(wandb_mock, tmp_path):
     wandb_mock.Api().artifact.assert_called_once_with("test_artifact", type="model")
 
 
-@_xfail_python_ge_3_11_9
 @pytest.mark.parametrize(("log_model", "expected"), [("True", True), ("False", False), ("all", "all")])
 def test_wandb_logger_cli_integration(log_model, expected, wandb_mock, monkeypatch, tmp_path):
     """Test that the WandbLogger can be used with the LightningCLI."""
diff --git a/tests/tests_pytorch/test_cli.py b/tests/tests_pytorch/test_cli.py
diff --git a/tests/tests_pytorch/trainer/connectors/test_logger_connector.py b/tests/tests_pytorch/trainer/connectors/test_logger_connector.py
