[OPD] Stabilize two-node GB200 HybridEP training

[kaixih]

Summary

This PR is stacked on #1488 and enables its two-node GB200 recipe to use Megatron HybridEP. Because dynamic THD batches can produce different token-row counts across EP ranks, it pads both the model-input and routing-replay buffers to the EP-wide maximum before HybridEP collectives.

Before

flowchart LR
    N0["Node 0 / EP rank<br/>THD input: 8,192 rows<br/>Routing replay: 8,192 rows"]
    N1["Node 1 / EP rank<br/>THD input: 6,144 rows<br/>Routing replay: 6,144 rows"]
    AG["HybridEP metadata<br/>all-gather"]
    H["Different row counts<br/>timeout / hang"]

    N0 --> AG
    N1 --> AG
    AG --> H

Loading

After

flowchart LR
    MAX["EP-wide maximum<br/>8,192 rows"]
    N0["Node 0 / EP rank<br/>8,192 rows<br/>no padding"]
    N1["Node 1 / EP rank<br/>6,144 real + 2,048 padded rows<br/>replay padding uses -1"]
    AG["HybridEP metadata<br/>all-gather"]
    OK["Equal row counts<br/>collective completes"]

    MAX --> N0
    MAX --> N1
    N0 --> AG
    N1 --> AG
    AG --> OK

Loading

[gemini-code-assist]

Code Review

This pull request adds support for the HybridEP backend in the Qwen3.5-35B-A3B self-distillation pipeline on GB200. It implements padding for model inputs and routing replay data to align with the EP-wide maximum token-row count, preventing timeouts and mismatches. Additionally, it updates the Ray actor group to preload compatibility shims and adds comprehensive workflow documentation. The review comments provide valuable suggestions to prevent device mismatch errors by using the tensor's device instead of the current device context, and to avoid accidental tracing by only propagating the CUMEM_TRACE_ALL environment variable when it is non-empty.

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]

Using torch.cuda.current_device() directly can lead to device mismatch errors if replay_data is already on a specific CUDA device that differs from the active current device context. To make this more robust, we should use replay_data.device if it is a CUDA tensor, and fall back to torch.cuda.current_device() otherwise.

Suggested change

	target = torch.tensor(
	[local_target],
	dtype=torch.int64,
	device=torch.cuda.current_device(),
	)
	target = torch.tensor(
	[local_target],
	dtype=torch.int64,
	device=replay_data.device if replay_data.is_cuda else torch.cuda.current_device(),
	)

[gemini-code-assist]

Setting CUMEM_TRACE_ALL to an empty string ("") when it is not set or empty in os.environ can still trigger tracing in C shims that check for the presence of the environment variable using getenv("CUMEM_TRACE_ALL") (since getenv returns a non-NULL pointer to an empty string). To prevent accidental tracing, we should only propagate CUMEM_TRACE_ALL to env_vars if it is set and non-empty in os.environ.

Suggested change

	trace_lib = os.environ.get("TMS_CUMEM_TRACE_LIB")
	if trace_lib:
	assert os.path.exists(trace_lib), f"cuMemCreate trace library {trace_lib} does not exist."
	preload_libs.append(trace_lib)
	env_vars["CUMEM_TRACE_ALL"] = os.environ.get("CUMEM_TRACE_ALL", "")
	trace_lib = os.environ.get("TMS_CUMEM_TRACE_LIB")
	if trace_lib:
	assert os.path.exists(trace_lib), f"cuMemCreate trace library {trace_lib} does not exist."
	preload_libs.append(trace_lib)
	cumem_trace_all = os.environ.get("CUMEM_TRACE_ALL")
	if cumem_trace_all:
	env_vars["CUMEM_TRACE_ALL"] = cumem_trace_all