From be53e06b7a14e72f0c4a455678705e27559a8c38 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Fri, 22 May 2026 04:11:07 +0000 Subject: [PATCH] Fix paged stashing test submodules lookup --- tests/unit_tests/transformer/moe/test_paged_stashing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/moe/test_paged_stashing.py b/tests/unit_tests/transformer/moe/test_paged_stashing.py index 04dbbaa5ec6..df78214cb2b 100644 --- a/tests/unit_tests/transformer/moe/test_paged_stashing.py +++ b/tests/unit_tests/transformer/moe/test_paged_stashing.py @@ -15,6 +15,7 @@ paged_stash_init_chunk_handler, paged_stash_reset, ) +from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version from megatron.training.initialize import _set_random_seed @@ -136,7 +137,7 @@ def _create_moe_layer(self, layer_number=0): quantization_context = get_fp8_context(self.config, layer_number, is_init=True) with quantization_context: moe_layer = ( - MoELayer(self.config, transformer_layer_spec.submodules.mlp.submodules) + MoELayer(self.config, get_submodules(transformer_layer_spec.submodules.mlp)) .cuda() .to(dtype=self.test_dtype) )