use module prefix during modelopt restore (#14384)

jenchen13 · web-flow · commit e839eca6ec1c · 2025-08-05T09:46:51.000-07:00
Signed-off-by: Jennifer Chen &lt;jennifchen@nvidia.com&gt;
diff --git a/nemo/collections/llm/modelopt/model_utils.py b/nemo/collections/llm/modelopt/model_utils.py
@@ -234,6 +234,7 @@ def restore_modelopt_state(
         mto.plugins.restore_sharded_modelopt_state(
             [core_model],
             ckpt_to_weights_subdir(path, is_saving=False),
+            prefix="module.",
         )
     if mto.ModeloptStateManager.is_converted(core_model):
         logging.info("Restored Model Optimizer state from checkpoint.")

Original file line number	Diff line number	Diff line change
`@@ -234,6 +234,7 @@ def restore_modelopt_state(`
`234`	`234`	`mto.plugins.restore_sharded_modelopt_state(`
`235`	`235`	`[core_model],`
`236`	`236`	`ckpt_to_weights_subdir(path, is_saving=False),`
	`237`	`+ prefix="module.",`
`237`	`238`	`)`
`238`	`239`	`if mto.ModeloptStateManager.is_converted(core_model):`
`239`	`240`	`logging.info("Restored Model Optimizer state from checkpoint.")`