added yaml (multi)representer for PretrainedConfig object types

Author: speediedan

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [2.5.1] - 2025-03-27
+
+### Added
+
+- Support for Lightning ``2.5.1``
+- added (multi)representer for ``PretrainedConfig`` object types
+
 ## [2.5.0] - 2024-12-20
 
 ### Added

dockers/base-cuda/Dockerfile

@@ -84,7 +84,7 @@ 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/cu124; \
+        pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu124; \
         # ... 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

requirements/base.txt

@@ -1,4 +1,4 @@
-#lightning>=2.5.0,<2.5.1
+lightning>=2.5.0,<2.5.2
 # the below is uncommented when master is targeting a specific pl dev master commit
-git+https://github.com/Lightning-AI/lightning.git@878ecf56b06d5ae3f482e146e78accabc685bfc7#egg=lightning
+# git+https://github.com/Lightning-AI/lightning.git@878ecf56b06d5ae3f482e146e78accabc685bfc7#egg=lightning
 torch>=2.2.0

requirements/standalone_base.txt

@@ -1,4 +1,4 @@
-#pytorch-lightning>=2.5.0,<2.5.1
+pytorch-lightning>=2.5.0,<2.5.2
 # the below is uncommented when master is targeting a specific pl dev master commit
-git+https://github.com/Lightning-AI/pytorch-lightning.git@878ecf56b06d5ae3f482e146e78accabc685bfc7#egg=pytorch-lightning
+# git+https://github.com/Lightning-AI/pytorch-lightning.git@878ecf56b06d5ae3f482e146e78accabc685bfc7#egg=pytorch-lightning
 torch>=2.2.0

setup.py

@@ -131,15 +131,15 @@ def _setup_args(standalone: bool = False) -> Dict[str, Any]:
     )
 
     base_reqs = "standalone_base.txt" if standalone else "base.txt"
-    # install_requires = setup_tools._load_requirements(
-    # _INSTALL_PATHS["require"], file_name=base_reqs, standalone=standalone
-    # )
     install_requires = setup_tools._load_requirements(
-        _INSTALL_PATHS["require"],
-        file_name=base_reqs,
-        standalone=standalone,
-        pl_commit="878ecf56b06d5ae3f482e146e78accabc685bfc7",
+    _INSTALL_PATHS["require"], file_name=base_reqs, standalone=standalone
     )
+    # install_requires = setup_tools._load_requirements(
+    #     _INSTALL_PATHS["require"],
+    #     file_name=base_reqs,
+    #     standalone=standalone,
+    #     pl_commit="878ecf56b06d5ae3f482e146e78accabc685bfc7",
+    # )
     base_setup["install_requires"] = install_requires
     return base_setup
 

src/fts_examples/cfg_utils.py

@@ -3,6 +3,7 @@
 from typing import Dict, List, Any, Union
 from dataclasses import dataclass, field, asdict, fields
 
+from transformers import PretrainedConfig
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 
 
@@ -40,8 +41,19 @@ def optimizer_cfg_mapping_representer(dumper, data):
 def lr_scheduler_cfg_mapping_representer(dumper, data):
     return dumper.represent_mapping('tag:yaml.org,2002:map', asdict(data))
 
-yaml.SafeDumper.add_representer(OptimizerCfg, optimizer_cfg_mapping_representer)
-yaml.SafeDumper.add_representer(LRSchedulerCfg, lr_scheduler_cfg_mapping_representer)
+def pretrained_cfg_mapping_representer(dumper, data):
+    return dumper.represent_mapping('tag:yaml.org,2002:map', data.to_dict())
+
+# Register all custom representers to both base dumper classes
+representers = {
+    PretrainedConfig: pretrained_cfg_mapping_representer,
+    OptimizerCfg: optimizer_cfg_mapping_representer,
+    LRSchedulerCfg: lr_scheduler_cfg_mapping_representer
+}
+
+for dumper_cls in [yaml.Dumper, yaml.SafeDumper]:
+    for cls, representer in representers.items():
+        dumper_cls.add_multi_representer(cls, representer)
 
 def _is_overridden(dataclass_instance) -> bool:
     is_overridden = False