[feat] Adding rollout router replay (R3) on tinker path

[kailash109]

Motivation

Stabilizing MoE RL training logprob mismatch according to R3 paper, adding same functionality that exists in SkyRL base training code to Tinker path. Addresses Tinker Engine issues in #1846

Changes

The actual changes for passing router expert indices through to trainer are very small

• skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py: pass routed_experts through in sample() (natively returned by vllm).
• skyrl/tinker/types.py / api.py / engine.py: routing_matrix on GeneratedSequence and LossFnInputs, dtype/shape on TensorData, threading through prepare_model_pass_batch.
• skyrl/backends/skyrl_train_backend.py: add rollout_expert_indices in _to_training_batch and add routing_matrix return to aggregate_sample_results
• skyrl/train/dataset/preprocess.py: extract padding/downcast logic into a shared pad_rollout_expert_indices helper
  extract the padding/downcast logic into a shared pad_rollout_expert_indices helper (classic path behavior unchanged).

Extras:
examples/tinker/moe/: Added Modal launcher (modal_moe_tinker.py) for Tinker server and GSM8K GRPO client (moe_smoke_client.py) with --enable-router-replay that runs Qwen3-30b-a3b. Had to make one change to tinker SDK to reattach routing_matrix in the JSON deserializer (default tinker implementation will just drop unknown response fields)

One design consideration is whether the client should be responsible for receiving the router expert indices and pass it back in the forward_backward call, especially for longer sampling lengths where the indices tensor may grow in size. I decided to go this way for simplicity and because it seems to be what Fireworks does in their tinker service.

Validation

On the GSM8K GRPO run without router replay:

With router replay:

I haven't observed logprob diff collapse in these runs, but the logprob difference for without R3 seems to be > 2x than with it enabled. Will do more runs to see if i can observe a collapse situation

[gemini-code-assist]

Code Review

This pull request introduces support for Mixture of Experts (MoE) router replay, allowing MoE routing matrices to be captured during inference and replayed during training. It adds a Modal example and a smoke client for running LoRA GRPO on GSM8K, alongside updates to the Tinker API, types, and training backend to handle routing matrix data. The feedback highlights several areas for improvement: addressing a potential crash on empty tensors and redundant .max() calls in the padding helper, preventing a TypeError when parsing sequences in the client, avoiding a file descriptor leak by using a context manager for metrics logging, simplifying process termination in the Modal script, and relaxing the dtype validation in TensorData to str | None to prevent future API validation failures.

[gemini-code-assist]

In pad_rollout_expert_indices, padded.max().item() is called twice, which is inefficient. More importantly, if padded is an empty tensor (e.g., if max_total, num_layers, or topk is 0), calling .max() on it will raise a RuntimeError in PyTorch. Checking padded.numel() == 0 first and caching the max value avoids both issues.

    if padded.numel() == 0:
        return padded

    # downcast to uint8 if possible, otherwise int16 to save memory
    padded_max = padded.max().item()
    if padded_max < 2**8:
        return padded.to(torch.uint8)
    if padded_max < 2**15:
        return padded.to(torch.int16)
    return padded

[gemini-code-assist]

In install_routing_matrix_capture, result_dict.get("sequences", []) is used. If the "sequences" key exists in result_dict but its value is None, this will return None, causing zip to raise a TypeError. Using result_dict.get("sequences") or [] is safer. Additionally, using strict=True in zip prevents silent truncation if the lengths of result.sequences and the dictionary sequences list ever mismatch.

Suggested change

	for seq_obj, seq_dict in zip(result.sequences, result_dict.get("sequences", [])):
	for seq_obj, seq_dict in zip(result.sequences, result_dict.get("sequences") or [], strict=True):

[gemini-code-assist]

The metrics_file is opened here and only closed at the very end of the main function (line 326). If an exception or interruption occurs during the training loop, the file handle will be leaked. Consider wrapping the file operations in a with statement or a try...finally block to guarantee proper resource cleanup.

[gemini-code-assist]

In _kill_server, os.killpg(os.getpgid(server.pid), ...) is used. Since the server is started with start_new_session=True (line 199), the process group ID (PGID) is guaranteed to be equal to the process ID (PID) of the server. Therefore, os.killpg(server.pid, ...) can be used directly, which is more robust as it avoids the extra os.getpgid system call that can fail if the process is already dead or in a zombie state.

Suggested change

	os.killpg(os.getpgid(server.pid), signal.SIGTERM)
	except ProcessLookupError:
	return
	try:
	server.wait(timeout=60)
	except subprocess.TimeoutExpired:
	os.killpg(os.getpgid(server.pid), signal.SIGKILL)
	os.killpg(server.pid, signal.SIGTERM)
	except ProcessLookupError:
	return
	try:
	server.wait(timeout=60)
	except subprocess.TimeoutExpired:
	os.killpg(server.pid, signal.SIGKILL)

[gemini-code-assist]

The dtype field in TensorData is restricted to Literal["int64", "float32"] | None. If a client or future SDK version sends other common dtypes (such as "int32", "float16", or "float64"), Pydantic will raise a validation error and block the request. Since dtype is purely metadata and not strictly validated/used for type casting in the backend, using str | None is much more robust and future-proof.

Suggested change

	dtype: Literal["int64", "float32"] | None = None
	dtype: str | None = None

[gemini-code-assist]

The dtype field in TensorData is restricted to Literal["int64", "float32"] | None. To match the robustness improvement in skyrl/tinker/types.py and prevent potential API validation failures, this should be changed to str | None.

Suggested change

	dtype: Literal["int64", "float32"] | None = None
	dtype: str | None = None