Skip to content

Commit 8dcf9d2

Browse files
indianspeedsterindianspeedster
committed
test(rocm): register the MI300X test set for stage-c-4-gpu-mi300x
Replaces the earlier blanket registration of every 8-GPU CUDA test with the 8 files selected in review on PR #1606 (9300s total, down from ~10.2h). Five are already 4-GPU on CUDA and take register_rocm_ci() directly: test_qwen3_4B_fsdp_hybrid_shard_r2s2 600s test_qwen3_4B_offload_disk_stream 600s test_glm5_2_744b_a40b_5layer_ci 900s test_inkling_small_4layer_ci 1800s test_deepseek_v4_flash_4layer_ci 1900s Three are 8-GPU on CUDA and cannot run on a 4-GPU runner as written, so each gets a standalone test_amd_<name>.py with its own 4-GPU CaseConfig rather than an IS_HIP branch in the original: test_amd_r3_mtp 1100s cp 2 -> 1 test_amd_mtp1_spec_v2_r3 1600s world 8 -> 4 (dp 2 -> 1) test_amd_deepep_fp8_bridge 800s -> cp2/pp1/tp2/ep4 Each variant keeps the CUDA case's feature flags and metric gates and only narrows parallelism; the per-rank shard is unchanged in every case, so only data parallelism is given up. MI300X carries 192GB per GPU, so 4 of them hold more HBM than the 8x80GB H100 node these were tuned against. The 30B bridge case reuses the 4-GPU shape already proven for that model in test_r3_baseline.py / test_r3_deepep_fp8.py. All 8 carry the `amd` label so --labels amd --match-all-labels selects exactly this set.
1 parent 3bb23e4 commit 8dcf9d2

8 files changed

Lines changed: 192 additions & 5 deletions

File tree

‎tests/e2e/fsdp/test_qwen3_4B_fsdp_hybrid_shard_r2s2.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import os
22

3-
from tests.ci.ci_register import register_cuda_ci
3+
from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
44

55
import miles.utils.external_utils.command_utils as U
66

@@ -9,6 +9,11 @@
99
suite="stage-c-4-gpu-h200",
1010
labels=["fsdp"],
1111
)
12+
register_rocm_ci(
13+
est_time=600,
14+
suite="stage-c-4-gpu-mi300x",
15+
labels=["fsdp", "amd"],
16+
)
1217

1318
NUM_GPUS = 4
1419
DP_REPLICATE_SIZE = 2

‎tests/e2e/megatron/model_scripts/test_deepseek_v4_flash_4layer_ci.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
import os
22

33
from scripts.run_deepseek_v4 import ScriptArgs, _prepare_download, _prepare_single, _prepare_spmd, _train
4-
from tests.ci.ci_register import register_cuda_ci
4+
from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
55
from tests.ci.metric_history import register_ci_gate
66

77
register_cuda_ci(est_time=1900, suite="stage-c-4-gpu-h200", labels=["megatron", "model-scripts"])
8+
register_rocm_ci(
9+
est_time=1900,
10+
suite="stage-c-4-gpu-mi300x",
11+
labels=["megatron", "model-scripts", "amd"],
12+
)
813

914
register_ci_gate(metric_key="train/grad_norm")
1015
register_ci_gate(metric_key="train/ppo_kl")

‎tests/e2e/megatron/model_scripts/test_glm5_2_744b_a40b_5layer_ci.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
_prepare_megatron_ckpt,
99
_validate_glm_checkpoint,
1010
)
11-
from tests.ci.ci_register import register_cuda_ci
11+
from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
1212
from tests.ci.metric_history import register_ci_gate
1313

1414
import miles.utils.external_utils.command_utils as U
@@ -19,6 +19,11 @@
1919

2020

2121
register_cuda_ci(est_time=900, suite="stage-c-4-gpu-h200", labels=["megatron", "model-scripts"])
22+
register_rocm_ci(
23+
est_time=900,
24+
suite="stage-c-4-gpu-mi300x",
25+
labels=["megatron", "model-scripts", "amd"],
26+
)
2227

2328
register_ci_gate(metric_key="train/grad_norm")
2429
register_ci_gate(metric_key="train/ppo_kl")

‎tests/e2e/megatron/model_scripts/test_inkling_small_4layer_ci.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import os
22

33
from scripts.run_inkling import _MODEL_REGISTRY, ScriptArgs, _train
4-
from tests.ci.ci_register import register_cuda_ci
4+
from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
55
from tests.ci.metric_history import register_ci_gate
66

77
import miles.utils.external_utils.command_utils as U
@@ -12,6 +12,11 @@
1212
suite="stage-c-4-gpu-h200",
1313
labels=["megatron", "model-scripts"],
1414
)
15+
register_rocm_ci(
16+
est_time=1800,
17+
suite="stage-c-4-gpu-mi300x",
18+
labels=["megatron", "model-scripts", "amd"],
19+
)
1520

1621
register_ci_gate(metric_key="train/grad_norm")
1722
register_ci_gate(metric_key="train/ppo_kl")

‎tests/e2e/megatron/test_glm47_flash/test_amd_r3_mtp.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
"""AMD 4-GPU variant of test_r3_mtp.py.
2+
3+
Standalone rather than an IS_HIP branch in the original: the MI300X fleet is
4+
split into two 4-GPU runners, so the 8-GPU CUDA case cannot run there as
5+
written, and keeping the variant separate means neither side's parallelism
6+
constrains the other.
7+
8+
Difference from the CUDA case: cp_size 2 -> 1, halving the world size from 8 to
9+
4. TP=2 and PP=2 are unchanged, so each rank holds the same shard as before and
10+
only data parallelism is given up. CP is the safe axis to drop here -- it exists
11+
for long-context throughput, not for the R3 + MTP behaviour under test.
12+
"""
13+
14+
import os
15+
16+
from tests.ci.ci_register import register_rocm_ci
17+
from tests.ci.metric_history import register_ci_gate
18+
from tests.e2e.megatron.test_glm47_flash._common import CaseConfig, execute, prepare
19+
20+
register_rocm_ci(
21+
est_time=1100,
22+
suite="stage-c-4-gpu-mi300x",
23+
labels=["megatron", "amd"],
24+
)
25+
26+
register_ci_gate(metric_key="train/grad_norm")
27+
register_ci_gate(metric_key="train/ppo_kl")
28+
register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
29+
register_ci_gate(metric_key="train/train_rollout_kl")
30+
register_ci_gate(metric_key="rollout/raw_reward")
31+
32+
CASE = CaseConfig(
33+
use_deepep=False,
34+
num_gpus_per_node=4,
35+
cp_size=1,
36+
pp_size=2,
37+
tp_size=2,
38+
ep_size=4,
39+
# GLM-4.7-Flash has 20 attention heads; non-EP SGLang TP must divide it.
40+
rollout_num_gpus_per_engine=4,
41+
)
42+
43+
44+
if __name__ == "__main__":
45+
for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
46+
os.environ.pop(proxy_var, None)
47+
prepare(CASE)
48+
execute(CASE, wandb_file=__file__)

‎tests/e2e/megatron/test_qwen3_30B_A3B/test_amd_deepep_fp8_bridge.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
"""AMD 4-GPU variant of test_deepep_fp8_bridge.py.
2+
3+
Standalone rather than an IS_HIP branch in the original: the MI300X fleet is
4+
split into two 4-GPU runners, so the 8-GPU CUDA case cannot run there as
5+
written, and keeping the variant separate means neither side's parallelism
6+
constrains the other.
7+
8+
Difference from the CUDA case: the parallelism drops to the 4-GPU shape already
9+
proven for this model in test_r3_baseline.py / test_r3_deepep_fp8.py
10+
(cp2/pp1/tp2/ep4), and the rollout engine and SGLang EP follow the world size
11+
down from 8 to 4. The feature flags under test are unchanged: DeepEP on, FP8
12+
rollout on, bridge on, R3 off. max_tokens_per_gpu is left at the CUDA case's
13+
2048 -- that cap was set by host memory, which the split runner does not
14+
change.
15+
"""
16+
17+
import os
18+
19+
from tests.ci.ci_register import register_rocm_ci
20+
from tests.ci.metric_history import register_ci_gate
21+
from tests.e2e.megatron.test_qwen3_30B_A3B._common import CaseConfig, execute, prepare
22+
23+
register_rocm_ci(
24+
est_time=800,
25+
suite="stage-c-4-gpu-mi300x",
26+
labels=["megatron", "amd"],
27+
)
28+
29+
register_ci_gate(metric_key="train/grad_norm")
30+
register_ci_gate(metric_key="train/ppo_kl")
31+
register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
32+
register_ci_gate(metric_key="train/train_rollout_kl")
33+
register_ci_gate(metric_key="rollout/raw_reward")
34+
35+
CASE = CaseConfig(
36+
use_deepep=True,
37+
use_fp8_rollout=True,
38+
use_int4_rollout=False,
39+
use_bridge=True,
40+
use_r3=False,
41+
num_gpus_per_node=4,
42+
cp_size=2,
43+
pp_size=1,
44+
tp_size=2,
45+
ep_size=4,
46+
rollout_num_gpus_per_engine=4,
47+
sglang_ep_size=4,
48+
max_tokens_per_gpu=2048,
49+
)
50+
51+
52+
if __name__ == "__main__":
53+
for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
54+
os.environ.pop(proxy_var, None)
55+
prepare(CASE, need_fp8=CASE.use_fp8_rollout, need_int4=CASE.use_int4_rollout, all_bridge=CASE.use_bridge)
56+
execute(CASE, wandb_file=__file__)

‎tests/e2e/megatron/test_qwen3_4B_offload_disk_stream.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import glob
1313
import os
1414

15-
from tests.ci.ci_register import register_cuda_ci
15+
from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
1616
from tests.ci.metric_history import register_ci_gate
1717

1818
import miles.utils.external_utils.command_utils as U
@@ -27,6 +27,11 @@
2727
suite="stage-c-4-gpu-h200",
2828
labels=["miles-plugin"],
2929
)
30+
register_rocm_ci(
31+
est_time=600,
32+
suite="stage-c-4-gpu-mi300x",
33+
labels=["miles-plugin", "amd"],
34+
)
3035

3136
register_ci_gate(metric_key="train/grad_norm")
3237
register_ci_gate(metric_key="train/ppo_kl")

‎tests/e2e/megatron/test_qwen3_5_35B_A3B_mtp/test_amd_mtp1_spec_v2_r3.py‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
"""AMD 4-GPU variant of test_mtp1_spec_v2_r3.py.
2+
3+
Qwen3.5-35B-A3B: 1 MTP layer + speculative-v2 + R3, on 4 GPUs.
4+
5+
Standalone rather than an IS_HIP branch in the original: the MI300X fleet is
6+
split into two 4-GPU runners, so the 8-GPU CUDA case cannot run there as
7+
written, and keeping the variant separate means neither side's parallelism
8+
constrains the other.
9+
10+
Difference from the CUDA case: num_gpus_per_node 8 -> 4, which drops data
11+
parallelism from 2 to 1. The tp2/pp2/cp1 shape is kept exactly -- TP=4 hits a
12+
Qwen3.5 attention-output-gate sharding bug, CP=1 avoids the memory-heavy
13+
GatedDeltaNet CP backward kernel, and PP=2 halves the resident layers. Because
14+
TP/PP/CP are unchanged, the per-rank shard is identical to the CUDA case; the
15+
original comment notes it fits 8x80GB, and MI300X carries 192GB per GPU, so 4
16+
of them hold more total HBM than the node this was tuned for. The rollout
17+
engine and SGLang EP follow the world size down from 8 to 4.
18+
"""
19+
20+
import os
21+
22+
from tests.ci.ci_register import register_rocm_ci
23+
from tests.ci.metric_history import register_ci_gate
24+
from tests.e2e.megatron.test_qwen3_5_35B_A3B_mtp._common import CaseConfig, execute, prepare
25+
26+
register_rocm_ci(
27+
est_time=1600,
28+
suite="stage-c-4-gpu-mi300x",
29+
labels=["megatron", "qwen35", "amd"],
30+
)
31+
32+
register_ci_gate(metric_key="train/grad_norm")
33+
register_ci_gate(metric_key="train/ppo_kl")
34+
register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff")
35+
register_ci_gate(metric_key="train/train_rollout_kl")
36+
register_ci_gate(metric_key="rollout/raw_reward")
37+
38+
CASE = CaseConfig(
39+
num_gpus_per_node=4,
40+
cp_size=1,
41+
pp_size=2,
42+
tp_size=2,
43+
ep_size=4,
44+
rollout_num_gpus_per_engine=4,
45+
sglang_ep_size=4,
46+
enable_mtp_training=True,
47+
use_r3=True,
48+
# miles has no VLM/vision implementation on the training side, so vision weights are
49+
# never synced; exclude them from the weight-equality check.
50+
check_weight_update_skip_list=("visual",),
51+
)
52+
53+
54+
if __name__ == "__main__":
55+
for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
56+
os.environ.pop(proxy_var, None)
57+
prepare(CASE)
58+
execute(CASE, wandb_file=__file__)

0 commit comments

Comments
 (0)