From eb73662bf5fefc691196906b5e3c8f906134c4b4 Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Mon, 27 Apr 2026 18:31:25 +0200 Subject: [PATCH 1/2] Make regex capture better Signed-off-by: Asha Anoosheh --- modelopt/torch/distill/plugins/megatron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index dbfad6fb6bb..a460204f512 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 From c0312ef26bd8ed91ed27cdd071fc9e3c967950eb Mon Sep 17 00:00:00 2001 From: Asha Anoosheh Date: Mon, 27 Apr 2026 19:24:07 +0200 Subject: [PATCH 2/2] Smarter string replacement Signed-off-by: Asha Anoosheh --- modelopt/torch/distill/plugins/megatron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index a460204f512..9a98eee9c77 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -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}"'