[Feature][Training] Add Nanbeige recurrent RL training

[0z5a]

Purpose

Add Nanbeige4.2-3B PPO, GRPO, Decoupled PPO and categorical Flow-DPPO through the standard Ray rollout/update contract. Preserve independent loop attention, physical parameter sharing, checkpoint head dimensions and complete policy publication, with separate critic checkpoints.

The current #460/#454 base is merged. Recurrent forward paths preserve draft-feature collection. Native SGD supports Megatron's Torch fallback. TMS imports are confined to offload paths; executing the published offload recipe still requires native TMS.

Validation

Fresh Linux x86_64 qualification also passes all 101 original tests and all five dependency/source/import/test stages on Torch 2.12.1+cu130. All 137 dependency versions are verified. Command-scoped GNU compilers and fresh JIT caches resolve seven missing-libimf.so failures without changing source, test math or pins. Independent review verifies 202 evidence files, raw XML and nine generated ELF libraries. CUDA remains uninitialized; no RL speed, accuracy or reward convergence result is claimed.

Runtime 0ca61635 passes 101 CPU tests with zero failures, errors or skips, including all four tiny three-update/fresh-resume cases. Actor module import succeeds with CUDA uninitialized. Pinned Ruff, Black and isort pass. Unmodified native TMS builds two AArch64 CUDA 13 libraries with CUDA 13.2. Actual recipe argument parsing passes for all four algorithms with FP32 and offload preserved; Ray launch is intercepted.

Algorithm	Tiny CPU training/resume	Recipe argument parsing	Official-weight GPU e2e	Matched speed improvement
PPO	PASS	PASS	Pending	Not measured
GRPO	PASS	PASS	Pending	Not measured
DPPO	PASS	PASS	Pending	Not measured
Flow-DPPO	PASS	PASS	Pending	Not measured

See reproduction recipe and results in the branch. Public validation summaries retain necessary configuration, results and checksums; original logs remain private. Official-weight GPU e2e, CUDA offload/recovery, reward convergence and matched speed measurements remain incomplete. The original two-physical-GPU FP32 campaign and selected-logprob bound of 0.03 are unchanged. Keep this PR Draft.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #35006872 | 📁 Comparing b31c99c against latest (791f101)

  🔍 Preview build  

51 files changed · + 6 added · ± 45 modified

+ Added

• advanced/draft-feature-collection.html
• _examples_synced/looped_grpo/README.html
• _examples_synced/looped_ppo/README.html
• _examples_synced/nanbeige/README.html
• _examples_synced/nanbeige/RESULTS.html
• _examples_synced/ouro/README.html

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...

[gemini-code-assist]

Code Review

This pull request introduces a native vllm-rlt rollout backend as an alternative to the HTTP-based vLLM server, adding support for fixed-depth recurrent policy training (PPO, GRPO, DPPO, and Flow-DPPO) for Ouro and Nanbeige models. It also implements a categorical Flow-DPPO extension with full-vocabulary KL gating and adds corresponding integration and unit tests. The code review identified several critical issues that need to be addressed: a runtime TypeError caused by an invalid enable_gqa argument in F.scaled_dot_product_attention, potential device mismatches in the Flow-DPPO loss, a CUDA-only limitation with fused=True in torch.optim.SGD that will fail on CPU, missing validation for the Hugging Face checkpoint path, a shallow copy issue that could leak sample mutations, and a potential KeyError in the launcher script for unsupported model types.

[gemini-code-assist]

The F.scaled_dot_product_attention function in PyTorch does not accept an enable_gqa argument. Passing this argument will result in a TypeError at runtime. Please remove enable_gqa=True as PyTorch automatically handles Grouped Query Attention (GQA) when the query and key/value tensors have different numbers of heads.

            F.scaled_dot_product_attention(
                q.transpose(0, 1).unsqueeze(0),
                k.transpose(0, 1).unsqueeze(0),
                v.transpose(0, 1).unsqueeze(0),
                is_causal=True,
            )

[gemini-code-assist]

To prevent potential RuntimeError due to device mismatch, explicitly move actions and advantages to log_probs.device. Since old_log_probs is explicitly moved to the GPU device, other input tensors like advantages (which might reside on CPU) should also be moved to the same device before performing element-wise operations.

Suggested change

	log_probs = F.log_softmax(logits.float(), dim=-1)
	old_log_probs = old_log_probs.detach().to(log_probs.device)
	log_probs = F.log_softmax(logits.float(), dim=-1)
	old_log_probs = old_log_probs.detach().to(log_probs.device)
	actions = actions.to(log_probs.device)
	advantages = advantages.to(log_probs.device)

[gemini-code-assist]

Using fused=True for torch.optim.SGD is only supported on CUDA devices. If the optimizer is run on CPU (for example, during CPU-only integration tests or in environments without a GPU), this will raise a RuntimeError. Using fused=torch.cuda.is_available() safely falls back to the non-fused implementation on CPU.

Suggested change

	yield
	megatron_optimizer.SGD = partial(torch.optim.SGD, fused=torch.cuda.is_available())

[gemini-code-assist]

If args.hf_checkpoint is None or the path does not contain config.json, this line will raise a TypeError or FileNotFoundError with an unhelpful stack trace. Adding explicit validation checks provides a much clearer error message during argument validation.

Suggested change

	config = json.loads((Path(args.hf_checkpoint) / "config.json").read_text())
	checkpoint_path = getattr(args, "hf_checkpoint", None)
	if not checkpoint_path:
	raise ValueError("hf_checkpoint is required but not specified or empty")
	config_file = Path(checkpoint_path) / "config.json"
	if not config_file.exists():
	raise ValueError(f"config.json not found in hf_checkpoint path: {checkpoint_path}")
	config = json.loads(config_file.read_text())

[gemini-code-assist]

Using copy.copy(sample) performs a shallow copy. If _prepare_prompt_ids mutates nested structures like multimodal_train_inputs in-place, those changes will leak to the original sample object, defeating the purpose of keeping changes local until validation. Using copy.deepcopy(sample) ensures complete isolation of the original sample.

Suggested change

	prepared_sample = copy.copy(sample)
	prepared_sample = copy.deepcopy(sample)

[gemini-code-assist]

If family is not one of the supported types ("ouro" or "nanbeige"), accessing providers[family] will raise a KeyError. Adding an explicit check with parser.error provides a clean error message to the user.

Suggested change

	providers = {"ouro": "ouro", "nanbeige": "nanbeige"}
	provider = providers[family]
	providers = {"ouro": "ouro", "nanbeige": "nanbeige"}
	if family not in providers:
	parser.error(f"Unsupported model type: {family}. Supported types are: {list(providers.keys())}")
	provider = providers[family]