From f6ac12d5bdfd0928d855e3303b5651d4fe588b36 Mon Sep 17 00:00:00 2001 From: sreerohi Date: Thu, 28 May 2026 14:51:07 -0500 Subject: [PATCH] megatron: support GLM-4.7-Flash R3+MTP test on ROCm Mirrors the substitutions from PR #1126 (And related to #1153): - Set SGLANG_USE_AITER=0 (bypass aiter MoE silu_and_mul crash on gfx950) - Add --use-miles-router (sglang router strips return_routed_experts; see sgl-project/sglang#20696) - Add --ci-disable-logprobs-checker (ROCm numerical drift) - Set MILES_TEST_R3_THRESHOLD=1.0 (tolerate MoE topk routing drift) --- tests/e2e/megatron/test_glm47_flash/_common.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/e2e/megatron/test_glm47_flash/_common.py b/tests/e2e/megatron/test_glm47_flash/_common.py index e31844d1eb..a881a4d885 100644 --- a/tests/e2e/megatron/test_glm47_flash/_common.py +++ b/tests/e2e/megatron/test_glm47_flash/_common.py @@ -1,12 +1,15 @@ import os from dataclasses import dataclass +import torch + import miles.utils.external_utils.command_utils as U MODEL_NAME = "GLM-4.7-Flash" MODEL_TYPE = "glm4.7-flash" TIGHT_HOST_MEMORY = bool(int(os.environ.get("MILES_TEST_TIGHT_HOST_MEMORY", "1"))) +IS_ROCM = torch.version.hip is not None @dataclass @@ -125,6 +128,8 @@ def build_train_args(case: CaseConfig, *, wandb_file: str) -> str: "--sglang-speculative-eagle-topk 1 " "--sglang-speculative-num-draft-tokens 3 " ) + if IS_ROCM: + sglang_args += "--use-miles-router " if case.use_deepep: sglang_args += "--sglang-moe-a2a-backend deepep --sglang-deepep-mode auto " @@ -134,6 +139,8 @@ def build_train_args(case: CaseConfig, *, wandb_file: str) -> str: mtp_args = "--enable-mtp-training " "--mtp-loss-scaling-factor 0.2 " ci_args = "--ci-test " + if IS_ROCM: + ci_args += "--ci-disable-logprobs-checker " misc_args = ( "--attention-dropout 0.0 " @@ -173,8 +180,16 @@ def execute(case: CaseConfig, *, wandb_file: str) -> None: train_args = build_train_args(case, wandb_file=wandb_file) + extra_env_vars = None + if IS_ROCM: + extra_env_vars = { + "SGLANG_USE_AITER": "0", + "MILES_TEST_R3_THRESHOLD": "1.0", + } + U.execute_train( train_args=train_args, num_gpus_per_node=case.num_gpus_per_node, megatron_model_type=MODEL_TYPE, + extra_env_vars=extra_env_vars, )