bump lightning dev sha, set pytorch 2.9 as base docker image, pin jsonargparse < 4.42.0 until parsing issue addressed, disable cpu offload config with certain FSDP 2 tests on PT 2.9+ until upstream issues addressed, use new sphinx theme to address mobile dev table of contents issue

speediedan · speediedan · commit 4a5aa6709988 · 2025-10-19T17:46:19.000-07:00
diff --git a/.azure-pipelines/gpu-tests.yml b/.azure-pipelines/gpu-tests.yml
@@ -104,7 +104,7 @@ jobs:
         python -m coverage html
         curl -Os https://uploader.codecov.io/latest/linux/codecov
         chmod +x codecov
-        ./codecov --token= $CODECOV_TOK --commit=$(Build.SourceVersion) --flags=gpu,pytest --name="GPU-coverage" --env=linux,azure
+        ./codecov -t $CODECOV_TOK --commit=$(Build.SourceVersion) --flags=gpu,pytest --name="GPU-coverage" --env=linux,azure
       env:
         CODECOV_TOK: $(CODECOV_TOKEN)  # explicit mapping required for secret azure pipeline variables
       displayName: 'Statistics'
diff --git a/dockers/base-cuda/Dockerfile b/dockers/base-cuda/Dockerfile
@@ -87,11 +87,11 @@ RUN \
     else \
         # or target a specific cuda build, by specifying a particular index url w/...
         # ... default channel
-        #pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128; \
+        pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128; \
         # ... pytorch patch version
         # pip install torch==1.11.1+cu113 torchvision==0.11.3+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html; \
         # ... pytorch nightly dev version
-        pip install --pre torch==2.9.0.dev20250811 --index-url https://download.pytorch.org/whl/nightly/cu128; \
+        # pip install --pre torch==2.9.0.dev20250811 --index-url https://download.pytorch.org/whl/nightly/cu128; \
         # ... test channel
         # pip install --pre torch==2.9.0 --index-url https://download.pytorch.org/whl/test/cu128; \
     fi && \
diff --git a/dockers/docker_images_release.sh b/dockers/docker_images_release.sh
@@ -41,7 +41,7 @@ maybe_build(){
 
 build_eval(){
 	# latest PyTorch image supported by release
-	declare -A iv=(["cuda"]="12.8.1" ["python"]="3.12" ["pytorch"]="2.8.0" ["lightning"]="2.5" ["cust_build"]="0")
+	declare -A iv=(["cuda"]="12.8.1" ["python"]="3.12" ["pytorch"]="2.9.0" ["lightning"]="2.5" ["cust_build"]="0")
 	export latest_pt="base-cu${iv["cuda"]}-py${iv["python"]}-pt${iv["pytorch"]}-pl${iv["lightning"]}"
 	export latest_azpl="py${iv["python"]}-pt${iv["pytorch"]}-pl${iv["lightning"]}-azpl-init"
 	maybe_build iv "${latest_pt}" "${latest_azpl}"
diff --git a/requirements/cli.txt b/requirements/cli.txt
@@ -1,3 +1,3 @@
-jsonargparse[signatures]>=4.27.7
+jsonargparse[signatures]>=4.27.7,<4.42.0
 omegaconf>=2.1.0
 hydra-core>=1.1.0
diff --git a/requirements/docs.txt b/requirements/docs.txt
@@ -11,4 +11,4 @@ sphinx-togglebutton>=0.2
 sphinx-copybutton>=0.3
 typing-extensions  # already in `base.txt` but the docs CI job does not install it
 jinja2>=3.0.0,<3.1.0
-git+https://github.com/speediedan/lightning_sphinx_theme.git@3f124e96e7f035c3391db2a3d601faf11530cd81#egg=pt_lightning_sphinx_theme
+git+https://github.com/speediedan/lightning_sphinx_theme.git@ad28149a5c27ed4bfb3f812760e179202024b2c9#egg=pt_lightning_sphinx_theme
diff --git a/requirements/lightning_pin.txt b/requirements/lightning_pin.txt
@@ -1 +1 @@
-8d1a73475e0f11ed3b3c5907da54fc993717d492
+c943c05c1425fb73cdaef8dee1bca65dabff626b
diff --git a/scripts/build_fts_env.sh b/scripts/build_fts_env.sh
@@ -98,9 +98,9 @@ base_env_build(){
                 if [[ -n ${torchvision_dev_ver} ]]; then
                     torchvision_dev_ver=${torch_dev_ver}
                 fi
-                pip install ${pip_install_flags} --pre torch==2.9.0.${torch_dev_ver} --index-url https://download.pytorch.org/whl/nightly/cu128
+                pip install ${pip_install_flags} --pre torch==2.10.0.${torch_dev_ver} --index-url https://download.pytorch.org/whl/nightly/cu128
             elif [[ $torch_test_channel -eq 1 ]]; then
-                pip install ${pip_install_flags} --pre torch==2.9.0 --index-url https://download.pytorch.org/whl/test/cu128
+                pip install ${pip_install_flags} --pre torch==2.10.0 --index-url https://download.pytorch.org/whl/test/cu128
             else
                 pip install ${pip_install_flags} torch torchvision --index-url https://download.pytorch.org/whl/cu128
             fi
@@ -129,6 +129,10 @@ base_env_build(){
             clear_activate_env python3.12
             pip install ${pip_install_flags} torch==2.8.0 torchvision --index-url https://download.pytorch.org/whl/cu128
             ;;
+        fts_latest_pt2_9_x | fts_release_pt2_9_x)
+            clear_activate_env python3.12
+            pip install ${pip_install_flags} torch==2.9.0 torchvision --index-url https://download.pytorch.org/whl/cu128
+            ;;
         *)
             echo "no matching environment found, exiting..."
             exit 1
diff --git a/scripts/gen_fts_coverage.sh b/scripts/gen_fts_coverage.sh
@@ -75,8 +75,8 @@ tmp_coverage_dir="/tmp"
 coverage_session_log="${tmp_coverage_dir}/gen_fts_coverage_${target_env_name}_${d}.log"
 
 # Define arrays of supported versions
-supported_fts_latest=(fts_latest_pt2_5_x fts_latest_pt2_6_x fts_latest_pt2_7_x fts_latest_pt2_8_x)
-supported_fts_release=(fts_release_pt2_5_x fts_release_pt2_6_x fts_release_pt2_7_x fts_release_pt2_8_x)
+supported_fts_latest=(fts_latest_pt2_5_x fts_latest_pt2_6_x fts_latest_pt2_7_x fts_latest_pt2_8_x fts_latest_pt2_9_x)
+supported_fts_release=(fts_release_pt2_5_x fts_release_pt2_6_x fts_release_pt2_7_x fts_release_pt2_8_x fts_release_pt2_9_x)
 
 # Enable extended globbing for pattern matching
 shopt -s extglob
diff --git a/src/finetuning_scheduler/strategy_adapters/model_parallel.py b/src/finetuning_scheduler/strategy_adapters/model_parallel.py
@@ -265,7 +265,8 @@ def _resolve_cfg_aliases(config_dict: Dict) -> Dict:
         for k, v in list(config_dict.items()):
             # currently adding a `cpu_offload_policy` option alias for convenience
             # open a GitHub issue if you think other poliy alias options would be useful
-            if k == "cpu_offload_policy":
+            # TODO: renable coverage below when upstream cpu offload issue addressed
+            if k == "cpu_offload_policy":  # pragma: no cover
                 config_dict["offload_policy"] = CPUOffloadPolicy(**v)
                 del config_dict[k]
             elif k == "act_ckpt":
@@ -283,7 +284,8 @@ def _validate_fsdp_plan(self) -> None:
         """
         if not self.fsdp_plan:
             return
-        if self.fsdp_default_kwargs:
+        # TODO: renable coverage below when upstream cpu offload issue addressed
+        if self.fsdp_default_kwargs:  # pragma: no cover
             self.fsdp_default_kwargs = ModelParallelStrategyAdapter._resolve_cfg_aliases(self.fsdp_default_kwargs)
         named_modules = dict(self.pl_module.named_modules()).keys()
         resolved_modules: Dict[str, Dict] = {}
diff --git a/src/fts_examples/model_parallel/config/fts_fsdp_auto_plan.yaml b/src/fts_examples/model_parallel/config/fts_fsdp_auto_plan.yaml
@@ -9,7 +9,7 @@ trainer:
         fsdp_default_kwargs:
           reshard_after_forward: True  # default value of a normal ``fully_shard`` kwarg
           act_ckpt: ['composable']  # use composable AC with default kwargs
-          cpu_offload_policy: {}  # apply default cpu offload policy
+          # cpu_offload_policy: {}  # apply default cpu offload policy
         fsdp_plan: {'model.output': {}, 'model.layers.\d*$': {}}
   strategy:
     init_args:
diff --git a/src/fts_examples/model_parallel/config/fts_fsdp_profiling.yaml b/src/fts_examples/model_parallel/config/fts_fsdp_profiling.yaml
@@ -8,7 +8,7 @@ trainer:
       strategy_adapter_cfg:
         fsdp_default_kwargs:
           act_ckpt: ['composable']
-          cpu_offload_policy: {}
+          # cpu_offload_policy: {}
         fsdp_plan: {'model.output': {}, 'model.layers.\d*$': {}, 'model': {}}
   strategy:
     init_args:
diff --git a/src/fts_examples/test_examples.py b/src/fts_examples/test_examples.py
@@ -104,12 +104,12 @@ def test_advanced_examples_fts_superglue(monkeypatch, recwarn, tmpdir, config_fi
 @pytest.mark.parametrize(
     "config_files",
     [
-        ("fts_fsdp_auto_plan.yaml",),
+        # ("fts_fsdp_auto_plan.yaml",), # TODO: debug mixed DTensor/Non-DTensor parameter group issue with PT 2.9+
         ("fts_tp_plan.yaml",),
         ("fts_fsdp_profiling.yaml", "profiling/memprofiler_demo.yaml"),
     ],
     ids=[
-        "fts_fsdp_auto_plan",
+        # "fts_fsdp_auto_plan",
         "fts_tp_plan",
         "fts_fsdp_profiling",
     ],
diff --git a/tests/helpers/expected_warns.py b/tests/helpers/expected_warns.py
@@ -51,6 +51,7 @@ def extend_warns(base_warns, additional_warns):
 ]
 
 distributed_warnings = [
+        "using the device under current context",  # required w/ PT 2.9 stable
         "Using the current device set by the user",  # required starting with PT 2.7
         "of Tensor.pin_memory",
         "Tensor.is_pinned",
@@ -63,7 +64,6 @@ def extend_warns(base_warns, additional_warns):
     "FSDP.state_dict_type",
     "Deallocating Tensor ",
     "`_get_pg_default_device` will be deprecated",
-
 ]
 
 DISTRIBUTED_WARNS = extend_warns(BASE_EXPECTED_WARNS, distributed_warnings)
@@ -78,6 +78,7 @@ def extend_warns(base_warns, additional_warns):
 
 # Example warnings
 EXAMPLE_BASE_WARNS = [
+    "using the device under current context",  # required w/ PT 2.9 stable
     "Using the current device set by the user",  # required starting with PT 2.7
     "sentencepiece tokenizer that you are converting",
     "co_lnotab is deprecated, use co_lines instead.",
diff --git a/tests/test_model_parallel.py b/tests/test_model_parallel.py
@@ -490,10 +490,13 @@ def __post_init__(self):
     ModParallelTestCfg(model_cfg_key="fsdp_cm_fp16", model_cls=cm_mod_parallel, precision_opts=fp16,
                        model_cfg=fsdp_cm_only, runif_alias="alone"),
     # FSDP2 auto plan tests
+    # TODO: debug upstream issues currently preventing use of `fsdp_auto_cpuoffld` cfg with fsdp_auto_* tests
+    # ModParallelTestCfg(model_cfg_key="fsdp_auto", model_cls=fsdp_auto_mod_parallel, model_cfg=fsdp_auto,
+    #                    strategy_adapter_cfg=fsdp_auto_cpuoffld, runif_alias="alone",
+    #                    expected_results=ExpectedResults(expected_state=path_fsdp)),
     ModParallelTestCfg(model_cfg_key="fsdp_auto", model_cls=fsdp_auto_mod_parallel, model_cfg=fsdp_auto,
-                       strategy_adapter_cfg=fsdp_auto_cpuoffld, runif_alias="alone",
+                       strategy_adapter_cfg=fsdp_auto_only, runif_alias="alone",
                        expected_results=ExpectedResults(expected_state=path_fsdp)),
-    # TODO: upstream bug currently preventing use of `fsdp_auto_cpuoffld` here w/ fp16
     ModParallelTestCfg(model_cfg_key="fsdp_auto_fp16", model_cls=fsdp_auto_mod_parallel, model_cfg=fsdp_auto,
                        precision_opts=fp16, strategy_adapter_cfg=fsdp_auto_only, runif_alias="alone",
                        expected_results=ExpectedResults(expected_state=path_fsdp)),

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-8d1a73475e0f11ed3b3c5907da54fc993717d492`
	`1`	`+c943c05c1425fb73cdaef8dee1bca65dabff626b`