[OPD] Stabilize two-node GB200 HybridEP training

[kaixih]

Summary

This is a stacked draft PR on top of #1488. It adds the Miles-side fixes and
recipe support needed to run Qwen3.5-35B-A3B pure OPD on two 4-GPU GB200 nodes
with Megatron HybridEP.

Root cause

Dynamic THD microbatches can contain different numbers of packed token rows on
different expert-parallel ranks. HybridEP metadata collectives require every
rank to contribute the same number of rows. The routing-replay tensor must use
the same alignment as the model input.

Without EP-wide alignment, the run can hit HYBRID-EP ALLGATHER TIMEOUT or
continue with invalid training signals, including order-of-magnitude jumps in
grad norm and importance-sampling metrics.

Changes

• Add optional EP-wide max padding to packed THD model inputs.
• Enable that padding for forward-only teacher scoring and training when using
  flex + hybridep.
• Apply the same EP-wide row alignment to routed-expert replay data, using -1
  for padded replay entries.
• Make the two-node GB200 Phase-2 recipe configurable for HybridEP, token
  budget, eval/save cadence, and CUDA graph batch sizes.
• Compose the torch-memory-saver preload with narrow runtime
  diagnostic/compatibility libraries.
• Document the two-node GB200 workflow and ownership boundaries.

No files under lab/ are included.

Validation

Validated in real two-node GB200 pure-OPD runs with 8 Megatron/SGLang ranks.
With the alignment fixes, HybridEP completed teacher scoring and training
updates without all-gather timeouts, and healthy update metrics stayed near
their expected ranges.

A final two-update run, with the separately tracked TMS/SGLang runtime fixes,
completed with:

Metric	Update 0	Update 1
Mean response length	18,736.00	6,699.26
OPD reverse-KL	0.044982	0.013613
Grad norm	0.313641	0.154027
OIS	1.0000002	1.0000001
ESS	0.9999980	0.9999998
TIS	1.0000207	0.9999925

Local checks:

• python3 -m py_compile on the four changed Python modules
• bash -n on phase2_gb200.sh
• git diff --check

Stack and external fixes

• Base: [OPD] Add Qwen3.5-35B-A3B single-node self-distillation example #1488
• torch-memory-saver PR Fix fsdp prob diff computation when tis is disabled #82 fixes CUDA VMM allocation sizing.
• SGLang #27140 handles post-update CUDA graph recapture.
• SGLang #30895 handles Qwen MoE graph-only padding rows.
• The narrow PyTorch shm_unlink(ENOENT) compatibility shim is still
  container-specific.

Those external fixes are documented here but are not vendored into this Miles
PR.

[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