diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py
index dbfad6fb6bb..9a98eee9c77 100644
--- a/modelopt/torch/distill/plugins/megatron.py
+++ b/modelopt/torch/distill/plugins/megatron.py
@@ -163,7 +163,7 @@ def setup_distillation_config(
 
 def _adjust_layer_index_for_pp(submodule_name, model_cfg):
     """Adjust any sequence-based layer indices found in a submodule name for Pipeline Parallelism."""
-    match = re.search(r"(?<=\.)\d+(?=\.)", submodule_name)
+    match = re.search(r"(?<=\.)\d+(?=\.|$)", submodule_name)
     if not match:
         return submodule_name
 
@@ -172,7 +172,7 @@ def _adjust_layer_index_for_pp(submodule_name, model_cfg):
     if new_layer_idx < 0:
         raise ValueError(f"Layer {submodule_name} does not fall on final PP rank.")
 
-    new_submodule_name = submodule_name.replace(match.group(0), str(new_layer_idx))
+    new_submodule_name = submodule_name.replace(f".{match.group(0)}", f".{new_layer_idx}")
     if parallel_state.get_tensor_and_context_parallel_rank() == 0:
         logger.info(
             f'Distillation: Renamed layer "{submodule_name}" on final PP rank to "{new_submodule_name}"'