From 3ce3601617d5cea10ed65dcb71546d52d8db55e7 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 25 Jul 2026 11:46:44 -0700 Subject: [PATCH 1/2] perf(glm): add GLM-5.2 GB200 proxy recipes Signed-off-by: yaoyu-33 --- .../perf_recipes/glm_moe_dsa/__init__.py | 2 + .../glm_moe_dsa/gb200/__init__.py | 5 + .../perf_recipes/glm_moe_dsa/gb200/glm5.py | 138 +++++++++++++++++- .../recipes/test_glm5_perf_recipes.py | 97 +++++++++++- 4 files changed, 239 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/glm_moe_dsa/__init__.py b/src/megatron/bridge/perf_recipes/glm_moe_dsa/__init__.py index e6cb1bb5e6..bf47794c90 100644 --- a/src/megatron/bridge/perf_recipes/glm_moe_dsa/__init__.py +++ b/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 ( diff --git a/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/__init__.py b/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/__init__.py index 4fc25d0d3c..be0a58c5d0 100644 --- a/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/__init__.py +++ b/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, +) diff --git a/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/glm5.py b/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/glm5.py index 8c1dcf0b10..ddbc80703a 100644 --- a/src/megatron/bridge/perf_recipes/glm_moe_dsa/gb200/glm5.py +++ b/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() diff --git a/tests/unit_tests/recipes/test_glm5_perf_recipes.py b/tests/unit_tests/recipes/test_glm5_perf_recipes.py index 9a4918a409..05cd687716 100644 --- a/tests/unit_tests/recipes/test_glm5_perf_recipes.py +++ b/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 From ee591d4af2c3746ae109597cc973008b2ea5326b Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Sat, 25 Jul 2026 12:07:30 -0700 Subject: [PATCH 2/2] docs(glm): record GLM-5.2 GB200 proxy performance Signed-off-by: yaoyu-33 --- .../model_verification_cards/glm5-2/card.yaml | 52 +++++++++++++++---- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/examples/model_verification_cards/glm5-2/card.yaml b/examples/model_verification_cards/glm5-2/card.yaml index 833c5aee58..e9a4c3f15a 100644 --- a/examples/model_verification_cards/glm5-2/card.yaml +++ b/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.