Load master weights from checkpoint (#14072)

* Load master weights from checkpoint

Signed-off-by: kunlunl <[email protected]>

* Change the default behavior to not load the main parameters from the checkpoint.

Signed-off-by: kunlunl <[email protected]>

* Add tests to cover all the added code

Signed-off-by: kunlunl <[email protected]>

* Apply isort and black reformatting

Signed-off-by: ko3n1g <[email protected]>

---------

Signed-off-by: kunlunl <[email protected]>
Signed-off-by: ko3n1g <[email protected]>
Co-authored-by: ko3n1g <[email protected]>

Author: 2 people

nemo/core/optim/mcore_optim.py

@@ -64,11 +64,14 @@ def zero_grad(self, set_to_none: bool = True):
         """
         self.mcore_optimizer.zero_grad(set_to_none)
 
-    def reload_model_params(self):
+    def reload_model_params(self, state_dict=None):
         """
         Reloads model parameters from the optimizer.
         """
-        self.mcore_optimizer.reload_model_params()
+        if state_dict is None:
+            self.mcore_optimizer.reload_model_params()
+        else:
+            self.mcore_optimizer.reload_model_params(state_dict=state_dict)
 
     def state_dict(self):
         """

nemo/lightning/pytorch/strategies/megatron_strategy.py

@@ -170,6 +170,7 @@ class MegatronStrategy(DDPStrategy, io.IOMixin):
         ckpt_type (TrainerCkptProtocol): Checkpoint type. Defaults to TrainerCheckpoint.
         ckpt_load_optimizer (bool): Load optimizer state from trainer.ckpt_path. Defaults to True.
         ckpt_save_optimizer (bool): Save optimizer states in checkpoint. Defaults to True.
+        ckpt_load_main_params (bool): Load main parameters from trainer.ckpt_path. Defaults to False.
         ddp (Union[DDPLiteral, DistributedDataParallelConfig]): DDP configuration. Defaults to "megatron".
         fsdp (Optional[FSDPLiteral]): Option of using torch FSDP2, select from ["megatron", "pytorch"].
             Defaults to None.
@@ -257,6 +258,7 @@ def __init__(
         find_unused_parameters: bool = False,
         ckpt_load_optimizer: bool = True,
         ckpt_save_optimizer: bool = True,
+        ckpt_load_main_params: bool = False,
         ddp: Union[DDPLiteral, DistributedDataParallelConfig] = "megatron",
         fsdp: Optional[FSDPLiteral] = None,
         lazy_init: bool = False,
@@ -319,6 +321,7 @@ def __init__(
         self.lazy_init = lazy_init
         self.ckpt_load_optimizer = ckpt_load_optimizer
         self.ckpt_save_optimizer = ckpt_save_optimizer
+        self.ckpt_load_main_params = ckpt_load_main_params
         self.ckpt_load_strictness = ckpt_load_strictness
         self.use_te_rng_tracker = use_te_rng_tracker
         self.use_sharp = use_sharp
@@ -391,6 +394,9 @@ def __init__(
         else:
             raise ValueError(f"Invalid DDP type: {ddp}")
 
+        if self.ckpt_load_optimizer and self.ckpt_load_main_params:
+            raise ValueError("ckpt_load_optimizer and ckpt_load_main_params cannot be both set to True.")
+
         if isinstance(self.ddp_config, DistributedDataParallelConfig):
             self.ddp_config.num_distributed_optimizer_instances = self.num_distributed_optimizer_instances
 
@@ -1052,7 +1058,13 @@ def load_model_state_dict(self, checkpoint: Mapping[str, Any], strict: bool = Tr
 
         if not 'optimizer' in checkpoint:
             for opt in self.optimizers:
-                opt.reload_model_params()
+                if self.ckpt_load_main_params:
+                    if "state_dict" in checkpoint:
+                        opt.reload_model_params(checkpoint["state_dict"])
+                    else:
+                        opt.reload_model_params(checkpoint)
+                else:
+                    opt.reload_model_params()
 
     @property
     @override

tests/lightning/pytorch/strategies/test_megatron_strategy.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
+
+import pytest
 
 from nemo.lightning.pytorch.strategies import MegatronStrategy
 
@@ -40,3 +42,28 @@ class Dummy: ...
 
         assert first_io != second_io
         assert second_io == strategy2.checkpoint_io
+
+    def test_ckpt_load_main_params_and_ckpt_load_optimizer_both_true(self):
+        # Make sure ckpt_load_optimizer and ckpt_load_main_params cannot be both set to True.
+        with pytest.raises(ValueError):
+            strategy = MegatronStrategy(ckpt_load_optimizer=True, ckpt_load_main_params=True)
+
+    def test_ckpt_load_main_params_with_state_dict(self):
+        # Test ckpt_load_main_params with "state_dict" key.
+        strategy = MegatronStrategy()
+        strategy.ckpt_load_main_params = True
+        strategy.megatron_parallel = MagicMock()
+        strategy.optimizers = [MagicMock()]
+        checkpoint = {"state_dict": {"param": 1}}
+        strategy.load_model_state_dict(checkpoint)
+        strategy.optimizers[0].reload_model_params.assert_called_once_with(checkpoint["state_dict"])
+
+    def test_ckpt_load_main_params_without_state_dict(self):
+        # Test ckpt_load_main_params with "state_dict" key.
+        strategy = MegatronStrategy()
+        strategy.ckpt_load_main_params = True
+        strategy.megatron_parallel = MagicMock()
+        strategy.optimizers = [MagicMock()]
+        checkpoint = {"other": 123}
+        strategy.load_model_state_dict(checkpoint)
+        strategy.optimizers[0].reload_model_params.assert_called_once_with(checkpoint)