fix(recipe): partition GLM-4.5 pipeline layers (#5551)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/recipes/glm/h100/glm45.py

@@ -48,6 +48,8 @@ def glm45_355b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.tensor_model_parallel_size = 2
     cfg.model.pipeline_model_parallel_size = 8
     cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.num_layers_in_first_pipeline_stage = 10
+    cfg.model.num_layers_in_last_pipeline_stage = 10
     cfg.model.pipeline_dtype = None
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
@@ -179,6 +181,8 @@ def glm45_air_106b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.tensor_model_parallel_size = 1
     cfg.model.pipeline_model_parallel_size = 4
     cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.num_layers_in_first_pipeline_stage = 11
+    cfg.model.num_layers_in_last_pipeline_stage = 11
     cfg.model.pipeline_dtype = None
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
@@ -314,6 +318,8 @@ def glm45_355b_sft_128gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.pipeline_model_parallel_layout = None
     cfg.model.tensor_model_parallel_size = 2
     cfg.model.pipeline_model_parallel_size = 8
+    cfg.model.num_layers_in_first_pipeline_stage = 10
+    cfg.model.num_layers_in_last_pipeline_stage = 10
     cfg.model.pipeline_dtype = None
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
@@ -449,6 +455,8 @@ def glm45_air_106b_sft_32gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.pipeline_model_parallel_layout = None
     cfg.model.tensor_model_parallel_size = 1
     cfg.model.pipeline_model_parallel_size = 4
+    cfg.model.num_layers_in_first_pipeline_stage = 11
+    cfg.model.num_layers_in_last_pipeline_stage = 11
     cfg.model.pipeline_dtype = None
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1

tests/unit_tests/recipes/test_glm45_recipes.py

@@ -54,8 +54,11 @@ def __init__(self):
         # Set default attributes that recipes might set
         self.tensor_model_parallel_size = 1
         self.pipeline_model_parallel_size = 1
+        self.pipeline_model_parallel_layout = None
         self.pipeline_dtype = None
         self.virtual_pipeline_model_parallel_size = None
+        self.num_layers_in_first_pipeline_stage = None
+        self.num_layers_in_last_pipeline_stage = None
         self.context_parallel_size = 1
         self.expert_model_parallel_size = 1
         self.expert_tensor_parallel_size = None
@@ -103,6 +106,23 @@ def from_hf_pretrained(hf_path: str, **kwargs):
         return _FakeBridge()
 
 
+class _RealDepthFakeBridge(_FakeBridge):
+    """Fake AutoBridge that preserves the selected GLM-4.5 model depth."""
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+
+    def to_megatron_provider(self, load_weights: bool = False):
+        model = _FakeModelCfg()
+        model.num_layers = self.num_layers
+        return model
+
+    @staticmethod
+    def from_hf_pretrained(hf_path: str, **kwargs):
+        num_layers = 46 if hf_path.endswith("Air") else 92
+        return _RealDepthFakeBridge(num_layers)
+
+
 class _FakeTokenizer:
     """Fake HuggingFace tokenizer for testing."""
 
@@ -439,6 +459,35 @@ def test_glm45_air_106b_pretrain_defaults(monkeypatch: pytest.MonkeyPatch):
     assert hasattr(cfg.model, "mtp_loss_scaling_factor")
 
 
+@pytest.mark.parametrize(
+    "recipe_name,num_layers",
+    [
+        ("glm45_355b_pretrain_config", 92),
+        ("glm45_355b_sft_config", 92),
+        ("glm45_air_106b_pretrain_config", 46),
+        ("glm45_air_106b_sft_config", 46),
+    ],
+)
+def test_glm45_pipeline_partition_covers_real_model_depth(
+    recipe_name: str,
+    num_layers: int,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """GLM-4.5 recipe defaults should partition every real decoder layer."""
+    from megatron.core.transformer.transformer_block import get_num_layers_to_build
+
+    recipe_func = getattr(_glm_module, recipe_name)
+    patch_recipe_module_global(monkeypatch, recipe_func, "AutoBridge", _RealDepthFakeBridge)
+
+    cfg = recipe_func()
+    layers_per_stage = [
+        get_num_layers_to_build(cfg.model, pp_rank=pp_rank)
+        for pp_rank in range(cfg.model.pipeline_model_parallel_size)
+    ]
+
+    assert sum(layers_per_stage) == num_layers
+
+
 def test_glm45_sft_offline_packing_is_disabled(monkeypatch: pytest.MonkeyPatch):
     """Test that unsupported offline packing remains disabled for GLM-4.5."""
     from megatron.bridge.recipes.glm import glm45_355b_sft_config