perf(glm): add GLM-5.2 GB200 debugging proxy recipes

examples/model_verification_cards/glm5-2/card.yaml

@@ -3,16 +3,21 @@
 
 title: glm5_2
 summary: >
-  Performance disclaimer: this model has not been performance-tuned; reported
-  timing and throughput metrics are sanity checks, not optimized performance
-  results. GLM-5.2 H100 support verification currently covers exact GPU
-  checkpoint round-trip on CPU and GPU, deterministic Megatron inference,
-  HF-to-Megatron forward equivalence, bounded full SFT with export/reload
-  inference, 200K-context CP32 SFT, bounded PEFT, and bounded pretraining with
-  a reloadable middle checkpoint and verified continuation. GB200 support
-  verification currently covers bounded pretraining, bounded full SFT, 128K
-  packed long-context SFT, bounded PEFT, and bounded checkpoint resume;
-  post-SFT export/reload inference remains pending.
+  Performance scope: only pretrain_performance.GB200 uses the tuned canonical
+  dense-equivalent performance proxy; its reduced depth and expert count, mock
+  data, forced router balancing, reduced DSA indexer top-k, and disabled
+  indexer loss make it an upper-throughput benchmark rather than
+  training-equivalent GLM-5.2 evidence. Timing and throughput metrics from
+  functional training items remain sanity checks rather than optimized
+  performance results. GLM-5.2 H100 support verification currently covers
+  exact GPU checkpoint round-trip on CPU and GPU, deterministic Megatron
+  inference, HF-to-Megatron forward equivalence, bounded full SFT with
+  export/reload inference, 200K-context CP32 SFT, bounded PEFT, and bounded
+  pretraining with a reloadable middle checkpoint and verified continuation.
+  GB200 support verification currently covers bounded pretraining, bounded
+  full SFT, 128K packed long-context SFT, bounded PEFT, bounded checkpoint
+  resume, and the dense-equivalent MXFP8 performance proxy; post-SFT
+  export/reload inference remains pending.
 verification_index:
   model_level:
     verified:
@@ -28,6 +33,8 @@ verification_index:
       unverified: [sft_export_inference]
     H100:
       verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume]
+  performance:
+    GB200: verified
 model:
   hf_id: zai-org/GLM-5.2
   hf_revision: 4d67f66cc64d3219133b767c253b2ad1425c6c88  # pragma: allowlist secret
@@ -581,3 +588,28 @@ items:
         absolute difference of 0.021295 within the allowed 0.081712. Final grad
         norm is 1.840; the last 10 steps average 67,700.250 ms and 90.710
         TFLOPS/GPU.
+
+  pretrain_performance:
+    GB200:
+      status: verified
+      precision: fp8_mx
+      bridge_commit: 3ce3601617d5cea10ed65dcb71546d52d8db55e7  # pragma: allowlist secret
+      command: >
+        ./scripts/training/train.sh --nodes 2 --gpus-per-node 4
+        --recipe glm52_50b_pretrain_8gpu_gb200_fp8mx_config
+        --mode pretrain --max_steps 20
+      last_verified: 2026-07-25
+      metrics:
+        initial_loss: 13.187000
+        final_loss: 0.02891495
+        last_10_steps_step_time_ms_avg: 5540.930
+        last_10_steps_model_tflops_per_gpu_avg: 977.880
+      expected_result: >
+        On two nodes with 8x GB200, the exact 50B dense-equivalent MXFP8 proxy
+        completes exactly 20 steps with finite LM and MTP losses and no skipped
+        or NaN iterations. The benchmark uses 18 layers (3 dense and 15 MoE),
+        64 experts, mock data, forced router balancing, HybridEP, DSA indexer
+        top-k 512, and zero indexer-loss coefficient, so its throughput is not
+        train-equivalent GLM-5.2 evidence. Over steps 11-20, the run averages
+        5,540.930 ms and 977.880 model TFLOP/s/GPU; an exact-recipe repeat
+        averages 5,562.800 ms and 974.030 model TFLOP/s/GPU.

src/megatron/bridge/perf_recipes/glm_moe_dsa/__init__.py

@@ -16,6 +16,8 @@
 
 from megatron.bridge.perf_recipes.glm_moe_dsa.gb200.glm5 import (
     glm51_sft_192gpu_gb200_bf16_config,
+    glm52_50b_pretrain_8gpu_gb200_bf16_config,
+    glm52_50b_pretrain_8gpu_gb200_fp8mx_config,
     glm52_sft_192gpu_gb200_bf16_config,
 )
 from megatron.bridge.perf_recipes.glm_moe_dsa.h100.glm5 import (

src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/__init__.py

@@ -11,3 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from megatron.bridge.perf_recipes.glm_moe_dsa.gb200.glm5 import (
+    glm52_50b_pretrain_8gpu_gb200_bf16_config,
+    glm52_50b_pretrain_8gpu_gb200_fp8mx_config,
+)

src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/glm5.py

@@ -11,16 +11,150 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""GB200 performance recipes for GLM-5.1 and GLM-5.2 SFT."""
+"""GB200 performance recipes for GLM-5.1 and GLM-5.2 pretraining and SFT."""
+
+import torch
 
 from megatron.bridge import AutoBridge
 from megatron.bridge.perf_recipes._common import _benchmark_common, _perf_precision
 from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS
-from megatron.bridge.recipes.common import _sft_common
+from megatron.bridge.recipes.common import _pretrain_common, _sft_common
 from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+from megatron.bridge.training.comm_overlap import CommOverlapConfig
 from megatron.bridge.training.config import ConfigContainer
 
 
+_GLM52_MODEL_ID = "zai-org/GLM-5.2"
+_GLM52_MODEL_REVISION = "4d67f66cc64d3219133b767c253b2ad1425c6c88"  # pragma: allowlist secret
+
+
+def _glm52_50b_proxy_pretrain_config(compute_dtype: str) -> ConfigContainer:
+    """Build the shared 8-GPU GLM-5.2 50B dense-equivalent performance proxy."""
+    cfg = _pretrain_common()
+
+    cfg.model = AutoBridge.from_hf_pretrained(
+        _GLM52_MODEL_ID,
+        revision=_GLM52_MODEL_REVISION,
+    ).to_megatron_provider(load_weights=False)
+    cfg.mixed_precision = _perf_precision(compute_dtype)
+
+    cfg.tokenizer.tokenizer_type = "NullTokenizer"
+    cfg.tokenizer.tokenizer_model = None
+    cfg.tokenizer.vocab_size = DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
+
+    cfg.dataset.seq_length = 8192
+    cfg.dataset.num_dataset_builder_threads = 1
+
+    cfg.model.seq_length = 8192
+    cfg.model.num_layers = 18
+    cfg.model.moe_layer_freq = [0, 0, 0] + [1] * 15
+    cfg.model.num_moe_experts = 64
+    cfg.model.qk_pos_emb_head_dim = 64
+    # This benchmark is a hardware-performance proxy rather than a train-equivalent
+    # GLM-5.2 recipe. Reducing sparse-attention top-k and disabling its auxiliary
+    # indexer loss exposes the throughput of the remaining training stack.
+    cfg.model.dsa_indexer_topk = 512
+    cfg.model.dsa_indexer_loss_coeff = 0.0
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 1
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.expert_model_parallel_size = 8
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.sequence_parallel = False
+
+    cfg.train.global_batch_size = 64
+    cfg.train.micro_batch_size = 1
+
+    cfg.model.transformer_impl = "transformer_engine"
+    cfg.model.attention_backend = "auto"
+    cfg.model.gradient_accumulation_fusion = True
+    cfg.model.moe_permute_fusion = True
+    cfg.model.moe_grouped_gemm = True
+    cfg.model.moe_router_fusion = True
+    cfg.model.moe_router_force_load_balancing = True
+    cfg.model.moe_token_dispatcher_type = "flex"
+    cfg.model.moe_flex_dispatcher_backend = "hybridep"
+    cfg.model.moe_router_dtype = "fp32"
+    cfg.model.moe_shared_expert_overlap = False
+    cfg.model.persist_layer_norm = True
+    cfg.model.bias_dropout_fusion = True
+    cfg.model.bias_activation_fusion = True
+    cfg.model.calculate_per_token_loss = True
+    cfg.model.dsa_kernel_backend = "cudnn"
+    cfg.model.mtp_num_layers = 1
+
+    cfg.model.recompute_granularity = None
+    cfg.model.recompute_method = None
+    cfg.model.recompute_num_layers = None
+    cfg.model.recompute_modules = None
+
+    cfg.model.cuda_graph_impl = "none"
+    cfg.model.cuda_graph_scope = []
+    cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = False
+
+    cfg.ddp.use_distributed_optimizer = True
+    # The indexer parameters intentionally receive no gradient in this proxy.
+    # Non-overlapped reduction supports those unused parameters.
+    cfg.ddp.overlap_grad_reduce = False
+    cfg.ddp.average_in_collective = False
+    cfg.ddp.grad_reduce_in_fp32 = False
+    cfg.optimizer.use_distributed_optimizer = True
+    cfg.optimizer.use_precision_aware_optimizer = True
+    cfg.optimizer.exp_avg_dtype = torch.bfloat16
+    cfg.optimizer.exp_avg_sq_dtype = torch.bfloat16
+    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False, overlap_grad_reduce=False)
+    cfg.dist.distributed_timeout_minutes = 30
+
+    _benchmark_common(cfg, cross_entropy_impl="native")
+    # A fixed low learning rate keeps longer mock-data benchmark windows finite
+    # without changing the measured kernel mix.
+    cfg.optimizer.lr = 3e-5
+    cfg.optimizer.min_lr = 3e-5
+    cfg.scheduler.lr_warmup_iters = 0
+    # The shared benchmark helper defaults HybridEP to 32 SMs; 16 is faster for this 8-GPU proxy.
+    cfg.model.moe_hybridep_num_sms = 16
+    cfg.checkpoint.load = None
+    cfg.model.apply_rope_fusion = False
+    cfg.env_vars = {
+        **COMMON_PERF_ENV_VARS,
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        "NCCL_GRAPH_REGISTER": 0,
+        "NCCL_NVLS_ENABLE": 0,
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
+        "USE_MNNVL": 1,
+    }
+    return cfg
+
+
+def glm52_50b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
+    """GLM-5.2 50B dense-equivalent proxy: 8× GB200, BF16, 8K mock data, EP=8."""
+    return _glm52_50b_proxy_pretrain_config("bf16")
+
+
+def glm52_50b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
+    """GLM-5.2 50B dense-equivalent proxy: 8× GB200, MXFP8, 8K mock data, EP=8."""
+    cfg = _glm52_50b_proxy_pretrain_config("fp8_mx")
+    cfg.model.use_transformer_engine_op_fuser = True
+    cfg.model.moe_mlp_glu_interleave_size = 32
+    cfg.model.fp8_output_proj = True
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = True
+    cfg.env_vars.update(
+        {
+            "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
+            "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+            "NVTE_NORM_BWD_USE_CUDNN": 1,
+            "NVTE_NORM_FWD_USE_CUDNN": 1,
+        }
+    )
+    return cfg
+
+
 def glm51_sft_192gpu_gb200_bf16_config() -> ConfigContainer:
     """GLM-5.1 SFT: 192× GB200, BF16, 128K packed THD, CP=32, cuDNN DSA."""
     cfg = _sft_common()

tests/unit_tests/recipes/test_glm5_perf_recipes.py

@@ -20,16 +20,21 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
+import torch.nn.functional as F
 from megatron.core.transformer.enums import LayerType
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 
 from megatron.bridge.perf_recipes.glm_moe_dsa import (
     glm51_sft_192gpu_gb200_bf16_config,
     glm51_sft_416gpu_h100_bf16_config,
+    glm52_50b_pretrain_8gpu_gb200_bf16_config,
+    glm52_50b_pretrain_8gpu_gb200_fp8mx_config,
     glm52_sft_192gpu_gb200_bf16_config,
     glm52_sft_416gpu_h100_bf16_config,
 )
 from megatron.bridge.training.config import ConfigContainer
+from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory
 
 
 pytestmark = pytest.mark.unit
@@ -63,7 +68,9 @@ def __init__(self, model_id: str) -> None:
         self.model_id = model_id
 
     @classmethod
-    def from_hf_pretrained(cls, model_id: str) -> "_FakeAutoBridge":
+    def from_hf_pretrained(cls, model_id: str, revision: str | None = None) -> "_FakeAutoBridge":
+        if model_id.endswith("GLM-5.2") and revision is not None:
+            assert len(revision) == 40
         return cls(model_id)
 
     def to_megatron_provider(self, load_weights: bool = False) -> SimpleNamespace:
@@ -72,6 +79,23 @@ def to_megatron_provider(self, load_weights: bool = False) -> SimpleNamespace:
         return SimpleNamespace(
             model_id=self.model_id,
             num_layers=78,
+            hidden_size=6144,
+            ffn_hidden_size=12288,
+            num_attention_heads=64,
+            num_query_groups=64,
+            kv_channels=192,
+            qk_pos_emb_head_dim=64,
+            num_moe_experts=256,
+            moe_ffn_hidden_size=2048,
+            moe_shared_expert_intermediate_size=2048,
+            moe_layer_freq=[0, 0, 0] + [1] * 75,
+            moe_router_topk=8,
+            gated_linear_unit=True,
+            activation_func=F.silu,
+            share_embeddings_and_output_weights=False,
+            vocab_size=154880,
+            make_vocab_size_divisible_by=1280,
+            should_pad_vocab=True,
             experimental_attention_variant="dsa",
             dsa_indexer_topk_freq=4 if is_glm52 else 1,
             dsa_indexer_skip_topk_offset=3 if is_glm52 else 0,
@@ -101,6 +125,77 @@ def _dsa_source_layer_id(layer_id: int, *, skip_topk_offset: int, topk_freq: int
     return layer_number - ((layer_number - sharing_offset) % topk_freq) - 1
 
 
+@pytest.mark.parametrize(
+    ("recipe_func", "fp8_recipe"),
+    [
+        (glm52_50b_pretrain_8gpu_gb200_bf16_config, None),
+        (glm52_50b_pretrain_8gpu_gb200_fp8mx_config, "mxfp8"),
+    ],
+    ids=["bf16", "fp8mx"],
+)
+def test_glm52_50b_proxy_pretrain_recipes(
+    recipe_func: Callable[[], ConfigContainer],
+    fp8_recipe: str | None,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The GB200 proxy preserves GLM-5.2 dimensions while scaling depth and expert count."""
+    cfg = _build_recipe(recipe_func, monkeypatch)
+
+    assert cfg.dataset.blend is None
+    assert cfg.dataset.seq_length == 8192
+    assert cfg.model.seq_length == 8192
+    assert cfg.model.num_layers == 18
+    assert cfg.model.moe_layer_freq == [0, 0, 0] + [1] * 15
+    assert cfg.model.num_moe_experts == 64
+    assert cfg.model.moe_router_topk == 8
+    assert cfg.model.dsa_indexer_topk == 512
+    assert cfg.model.dsa_indexer_loss_coeff == 0.0
+    assert cfg.model.hidden_size == 6144
+    assert cfg.model.qk_pos_emb_head_dim == 64
+    assert cfg.model.mtp_num_layers == 1
+    assert cfg.model.tensor_model_parallel_size == 1
+    assert cfg.model.pipeline_model_parallel_size == 1
+    assert cfg.model.context_parallel_size == 1
+    assert cfg.model.expert_model_parallel_size == 8
+    assert cfg.model.expert_tensor_parallel_size == 1
+    assert cfg.train.global_batch_size == 64
+    assert cfg.train.micro_batch_size == 1
+    assert cfg.model.moe_router_force_load_balancing is True
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "hybridep"
+    assert cfg.model.moe_hybridep_num_sms == 16
+    assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 8
+    assert cfg.env_vars["NUM_OF_TOKENS_PER_CHUNK_COMBINE_API"] == 128
+    assert cfg.dist.distributed_timeout_minutes == 30
+    assert cfg.ddp.overlap_grad_reduce is False
+    assert cfg.ddp.average_in_collective is False
+    assert cfg.optimizer.use_precision_aware_optimizer is True
+    assert cfg.optimizer.exp_avg_dtype == torch.bfloat16
+    assert cfg.optimizer.exp_avg_sq_dtype == torch.bfloat16
+    assert cfg.optimizer.lr == 3e-5
+    assert cfg.optimizer.min_lr == 3e-5
+    assert cfg.scheduler.lr_warmup_iters == 0
+    assert cfg.comm_overlap.tp_comm_overlap is False
+    assert cfg.comm_overlap.overlap_grad_reduce is False
+    assert cfg.mixed_precision.fp8_recipe == (fp8_recipe or "tensorwise")
+    if fp8_recipe is None:
+        assert cfg.mixed_precision.fp8 is None
+        assert cfg.model.use_transformer_engine_op_fuser is False
+        assert cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag is False
+    else:
+        assert cfg.mixed_precision.fp8 == "e4m3"
+        assert cfg.model.use_transformer_engine_op_fuser is True
+        assert cfg.model.moe_mlp_glu_interleave_size == 32
+        assert cfg.model.fp8_output_proj is True
+        assert cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag is True
+        assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
+        assert cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] == 1
+        assert cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] == 1
+
+    parameter_count = estimate_training_memory(cfg, include_activation=False).total_parameters
+    assert 45e9 < parameter_count < 55e9
+
+
 @pytest.mark.parametrize("recipe_func", _RECIPES, ids=lambda recipe: recipe.__name__)
 def test_glm5_perf_recipes_are_flat_and_preserve_bridge_dsa_fields(
     recipe_func: Callable[[], ConfigContainer], monkeypatch: pytest.MonkeyPatch