Skip to content

Rocm changes to tests/e2e/megatron/test_glm5_744b_a40b_4layer.py #1613

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sreerohi wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Rocm changes to tests/e2e/megatron/test_glm5_744b_a40b_4layer.py #1613

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 27 additions & 15 deletions scripts/run_glm5_744b_a40b.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,13 @@
from pathlib import Path
from typing import Literal

import torch
import typer

import miles.utils.external_utils.command_utils as U

IS_ROCM = hasattr(torch.version, "hip") and torch.version.hip is not None

app = typer.Typer()


Expand Down Expand Up @@ -213,6 +216,10 @@ def _prepare_cp(args: ScriptArgs, skip_existing: bool = False):


def _execute_train(args: ScriptArgs):
if IS_ROCM:
args.use_deepep = False
args.megatron_use_deepep = False

load_save_path = f"{args.output_dir}/{args.run_id}/checkpoints"
hf_name = f"{args.model_name}_fp8" if args.fp8_rollout else args.model_name
ckpt_args = (
Expand Down Expand Up @@ -340,13 +347,13 @@ def _execute_train(args: ScriptArgs):
)
if args.enable_pd:
sglang_args += "--prefill-num-servers 1 "
nsa_backend = "tilelang" if IS_ROCM else "flashmla_sparse"
sglang_args += (
# use flashmla backend for better precision
"--sglang-nsa-decode-backend flashmla_sparse "
"--sglang-nsa-prefill-backend flashmla_sparse "
"--sglang-kv-cache-dtype bf16 "
"--sglang-attention-backend nsa "
"--sglang-page-size 64 "
f"--sglang-nsa-decode-backend {nsa_backend} "
f"--sglang-nsa-prefill-backend {nsa_backend} "
+ ("--sglang-kv-cache-dtype bf16 " if not IS_ROCM else "")
+ "--sglang-attention-backend nsa "
f"--sglang-page-size {1 if IS_ROCM else 64} "
f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} "
# concurrency
f"--sglang-max-running-requests 512 "
Comment on lines 351 to 359

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Mixing implicit string literal concatenation and explicit + operators inside a parenthesized expression is highly unreadable and error-prone. It is much cleaner and more maintainable to split the string construction or use separate additions.

    sglang_args += (
        f"--sglang-nsa-decode-backend {nsa_backend} "
        f"--sglang-nsa-prefill-backend {nsa_backend} "
    )
    if not IS_ROCM:
        sglang_args += "--sglang-kv-cache-dtype bf16 "
    sglang_args += (
        "--sglang-attention-backend nsa "
        f"--sglang-page-size {1 if IS_ROCM else 64} "
        f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} "
        # concurrency
        f"--sglang-max-running-requests 512 "
    )

Expand All @@ -359,10 +366,12 @@ def _execute_train(args: ScriptArgs):
sglang_args += "--sglang-moe-runner-backend flashinfer_trtllm_routed "
else:
sglang_args += "--sglang-moe-runner-backend triton "
sglang_extra_env_vars = {
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": f"{32 if args.enable_pd else 256}",
"SGLANG_NSA_FORCE_MLA": "1",
}

sglang_extra_env_vars = {"SGLANG_NSA_FORCE_MLA": "1"}
if not IS_ROCM:
sglang_extra_env_vars["SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = (
f"{32 if args.enable_pd else 256}"
)

misc_args = (
# default dropout in megatron is 0.1
Expand Down Expand Up @@ -400,16 +409,19 @@ def _execute_train(args: ScriptArgs):
f"{args.extra_args} "
)

extra_env_vars = {
**sglang_extra_env_vars,
"INDEXER_ROPE_NEOX_STYLE": "0",
}
if not IS_ROCM:
extra_env_vars["NVSHMEM_DISABLE_NCCL"] = "1"

U.execute_train(
train_args=train_args,
config=args,
num_gpus_per_node=args.num_gpus_per_node,
megatron_model_type=args.megatron_model_type,
extra_env_vars={
**sglang_extra_env_vars,
"INDEXER_ROPE_NEOX_STYLE": "0",
"NVSHMEM_DISABLE_NCCL": "1",
},
extra_env_vars=extra_env_vars,
megatron_path=args.megatron_path,
)

Expand Down