Fix the serialization of partial functions in nemo 2.0 (#9668)

* fix serialization of partial function * update serialization to handle value.args Signed-off-by: srabhi <[email protected]> * add unit test Signed-off-by: srabhi <[email protected]> * remove redundant code from unit-test Signed-off-by: srabhi <[email protected]> --------- Signed-off-by: srabhi <[email protected]>
NVIDIA · Jul 17, 2024 · f65fea2 · f65fea2
1 parent af4f0ed
commit f65fea2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/nemo/lightning/io/fdl_torch.py b/nemo/lightning/io/fdl_torch.py
@@ -5,7 +5,9 @@
 """
 
 import types
+from functools import partial
 
+import fiddle as fdl
 import libcst as cst
 import torch
 import torch.nn as nn
@@ -110,6 +112,8 @@ def enable():
     def _modified_serialize(self, value, current_path, all_paths=None):
         if isinstance(value, types.BuiltinFunctionType):
             return self._pyref(value, current_path)
+        if isinstance(value, partial):
+            value = fdl.Partial(value.func, *value.args, **value.keywords)
         return self._original_serialize(value, current_path, all_paths)
 
     serialization.Serialization._original_serialize = serialization.Serialization._serialize

diff --git a/tests/lightning/io/test_api.py b/tests/lightning/io/test_api.py
@@ -1,3 +1,6 @@
+from functools import partial
+
+import pytest
 import transformer_engine as te
 from pytorch_lightning.loggers import TensorBoardLogger
 
@@ -7,8 +10,17 @@
 from nemo.lightning import io
 
 
+def dummy_extra(a, b, c=5):
+    return a + b + c
+
+
+@pytest.fixture
+def partial_function_with_pos_and_key_args():
+    return partial(dummy_extra, 10, c=15)
+
+
 class TestLoad:
-    def test_reload_ckpt(self, tmpdir):
+    def test_reload_ckpt(self, tmpdir, partial_function_with_pos_and_key_args):
         trainer = nl.Trainer(
             devices=1,
             accelerator="cpu",
@@ -26,10 +38,13 @@ def test_reload_ckpt(self, tmpdir):
             tokenizer=tokenizer,
         )
 
-        ckpt = io.TrainerContext(model, trainer)
+        ckpt = io.TrainerContext(model, trainer, extra={"dummy": partial_function_with_pos_and_key_args})
         ckpt.io_dump(tmpdir)
         loaded = io.load_context(tmpdir)
 
         assert loaded.model.config.seq_length == ckpt.model.config.seq_length
         assert loaded.model.__io__.tokenizer.vocab_file.startswith(str(tmpdir))
         assert loaded.model.__io__.tokenizer.merges_file.startswith(str(tmpdir))
+
+        loaded_func = loaded.extra["dummy"]
+        assert loaded_func(b=2) == partial_function_with_pos_and_key_args(b=2)