chore: make megatron e2e CaseConfig topology explicit instead of inferred, and tune e2e test time

[guapisolo]

What

Make the Megatron e2e CaseConfig topology explicit instead of inferred, in test_glm47_flash and test_qwen3_30B_A3B.

__post_init__ used to derive tp_size/ep_size (and, for Qwen, rollout_num_gpus / rollout_num_gpus_per_engine / sglang_ep_size), so a test file did not show the parallel sizes it actually trains with — readers had to trace the derivation. Those per-case numbers are now required CaseConfig fields each test_*.py passes explicitly, and __post_init__ is reduced to validation only (divisibility checks; the colocate→rollout_num_gpus resolution is dropped, with the rollout-pool check reading num_gpus_per_node directly for colocated cases).

Behavior-preserving: generated train_args (compared as a sorted arg multiset), the num_gpus_per_node passed to execute_train, and extra_env_vars are unchanged for all 9 existing cases.

Test

• Before/after train_args arg-multiset + execute_train kwargs equality across all 9 cases × MILES_TEST_TIGHT_HOST_MEMORY {0,1} × MILES_TEST_ENABLE_EVAL {0,1} (only the non-deterministic --wandb-group timestamp differs).
• python -m py_compile on both suites.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates the end-to-end Megatron test configurations for GLM-4.7-Flash and Qwen3-30B-A3B to require explicit topology parameters (such as tp_size, ep_size, and rollout_num_gpus_per_engine) instead of inferring them dynamically. All corresponding test cases have been updated to pass these parameters explicitly. The reviewer recommends adding validation logic in the __post_init__ methods of both test suites to ensure that the explicitly provided topology values are mathematically consistent, which helps prevent confusing errors during test execution.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since the topology values (tp_size, ep_size, etc.) are now passed explicitly rather than being inferred, we should add validation in __post_init__ to ensure they are mathematically consistent with num_gpus_per_node, cp_size, and pp_size. This prevents misconfigured test cases from causing confusing errors deep inside Megatron/SGLang.

    def __post_init__(self):
        if self.num_gpus_per_node % self.pp_size != 0:
            raise ValueError(
                f"num_gpus_per_node ({self.num_gpus_per_node}) must be divisible by pp_size ({self.pp_size})"
            )
        gpus_per_stage = self.num_gpus_per_node // self.pp_size
        if self.tp_size * self.cp_size != gpus_per_stage:
            raise ValueError(
                f"tp_size * cp_size ({self.tp_size} * {self.cp_size} = {self.tp_size * self.cp_size}) "
                f"must equal num_gpus_per_node // pp_size ({gpus_per_stage})"
            )
        if self.ep_size != gpus_per_stage:
            raise ValueError(
                f"ep_size ({self.ep_size}) must equal num_gpus_per_node // pp_size ({gpus_per_stage})"
            )
        if self.use_deepep and self.sglang_ep_size is None:
            raise ValueError("sglang_ep_size must be set explicitly when use_deepep is True")

[gemini-code-assist]

Since the topology values (tp_size, ep_size, etc.) are now passed explicitly rather than being inferred, we should add validation in __post_init__ to ensure they are mathematically consistent with num_gpus_per_node, cp_size, and pp_size. This prevents misconfigured test cases from causing confusing errors deep inside Megatron/SGLang.

    def __post_init__(self):
        # Validation only — topology values are passed explicitly, not inferred.
        if self.num_gpus_per_node % (self.cp_size * self.pp_size) != 0:
            raise ValueError(
                "num_gpus_per_node must be divisible by cp_size * pp_size: "
                f"{self.num_gpus_per_node=} {self.cp_size=} {self.pp_size=}"
            )
        gpus_per_stage = self.num_gpus_per_node // self.pp_size
        if self.tp_size * self.cp_size != gpus_per_stage:
            raise ValueError(
                f"tp_size * cp_size ({self.tp_size} * {self.cp_size} = {self.tp_size * self.cp_size}) "
                f"must equal num_gpus_per_node // pp_size ({gpus_per_stage})"
            )
        if self.ep_size != gpus_per_stage:
            raise ValueError(
                f"ep_size ({self.ep_size}) must equal num_gpus_per_node // pp_size ({gpus_per_stage})"
            )
        if self.use_deepep and self.sglang_ep_size is None:
            raise ValueError("sglang_ep_size must be set explicitly when use_deepep is True")
        if not self.colocate and self.rollout_num_gpus is None:
            raise ValueError("rollout_num_gpus must be set when colocate is False")
        rollout_pool = self.num_gpus_per_node if self.colocate else self.rollout_num_gpus
        if rollout_pool % self.rollout_num_gpus_per_engine != 0:
            raise ValueError(
                "rollout pool must be divisible by rollout_num_gpus_per_engine: "
                f"{rollout_pool=} {self.rollout_num_gpus_per_engine=}"
            )
        if self.update_weight_transfer_mode is not None:
            assert self.update_weight_transfer_mode == "broadcast"