feat(megatron): add LoRA-only IQuest LoopCoder support

[j316chuck]

• Enables vLLM's native IQuest LoopCoder model for LoRA hot-loading by declaring its packed QKV and MLP adapter mappings.
• Adds a Megatron Bridge model/provider that reuses one physical decoder and its LoRA adapters across both recurrent passes, including loop-1 KV reuse and loop-2 local/global gating.
• Restricts the initial path to standalone LoRA adapters, two loops, pipeline parallel size 1, and MLP-only activation recomputation; the frozen loop gate remains a base-model parameter.

Testing

uv run pytest -q tests/patches/test_iquest_loopcoder_lora.py tests/patches/test_iquest_loopcoder_megatron.py

5 tests passed, covering vLLM LoRA capability detection and recurrent KV mixing/gradients.

uv run pre-commit run --files skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py skyrl/backends/skyrl_train/patches/vllm_iquest_loopcoder_lora.py skyrl/backends/skyrl_train/workers/megatron/iquest_loopcoder.py skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py skyrl/backends/skyrl_train/workers/megatron/model_bridges.py tests/patches/test_iquest_loopcoder_lora.py tests/patches/test_iquest_loopcoder_megatron.py

Ruff, Black, and hardcoded-secret checks passed. AutoBridge dispatch was also verified against IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct; full 40B GPU training/serving was not run locally.

---

Note

Medium Risk
New recurrent attention training path with strict recompute/LoRA config requirements; mistakes could break gradients or inference adapter loading, but scope is limited to LoRA and a single model family.

Overview
Adds LoRA-only support for IQuest LoopCoder across Megatron training and vLLM serving.

On inference, a startup patch marks vLLM’s native IQuestLoopCoderForCausalLM as LoRA-capable (QKV + gate/up packed mappings) and is wired into the vLLM worker bootstrap alongside the existing Kimi patch.

On training, a new Megatron path runs the decoder twice with shared weights: pass 1 does full attention and caches KV; pass 2 blends that global attention with sliding-window local attention via a frozen per-head gate. A Megatron-Bridge registration maps HF LoopCoder weights (including gate projections) into this stack.

Megatron worker changes enforce LoopCoder constraints: LoRA with merge_lora=false, auto-downgrade full activation recompute to MLP-only selective recompute (to preserve loop-1 KV gradients), and provider-side guards on PP=1 and recompute settings.

Tests cover the vLLM LoRA patch and recurrent KV mixing/backprop behavior.

^{Reviewed by Cursor Bugbot for commit 52fd627. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces support for LoRA-only training and serving of IQuest LoopCoder models using Megatron-Core and vLLM. It adds a vLLM patch to enable LoRA capability, implements a custom two-pass Megatron-Core transformer block and attention mechanism to reuse physical layer weights, and integrates these components into the Megatron worker and model bridges. Feedback on the changes highlights a critical shape mismatch in the loop gate projection, a bug in the worker configuration check that would break PPO/RLHF reference workers, and an opportunity to optimize GPU memory usage by clearing captured key-value tensors after the second attention pass.

[gemini-code-assist]

There is a shape mismatch bug here. In Megatron-Core, core_attention modules (such as DotProductAttention) return attention context tensors of shape [..., num_heads, head_dim] (e.g., [seq_len, batch_size, num_heads, head_dim]). However, the current implementation reshapes the gate to [..., num_heads * head_dim], which will cause a RuntimeError due to shape mismatch when performing element-wise multiplication with local_output and global_output in LoopCoderCoreAttention.forward.

Removing the .reshape(*query.shape[:-2], -1) call ensures the gate retains the correct [..., num_heads, head_dim] shape matching the attention outputs.

Suggested change

	return gate.expand(*gate.shape[:-1], self.head_dim).reshape(*query.shape[:-2], -1)
	return gate.expand(*gate.shape[:-1], self.head_dim)

[gemini-code-assist]

This check will completely break PPO/RLHF training for IQuest LoopCoder. During training, a reference worker (MegatronRefWorkerBase) is initialized to compute reference logprobs. The reference worker runs the frozen base model and does not use LoRA adapters, so its lora_config is None. Because this check is enforced globally for any worker initializing an iquestloopcoder model, it will raise a ValueError when initializing the reference worker.

Additionally, using self.cfg.policy.megatron_config directly is incorrect because the reference worker uses self.cfg.ref.megatron_config. We should use the passed megatron_config argument instead.

We should restrict this check to only run for the policy worker (where training actually occurs) by checking isinstance(self, PolicyWorkerBase).

Suggested change

	if getattr(hf_config_original, "model_type", None) == "iquestloopcoder":
	if lora_config is None or lora_config.rank <= 0:
	raise ValueError("IQuest LoopCoder is supported only for LoRA training")
	if self.cfg.policy.megatron_config.lora_config.merge_lora:
	raise ValueError(
	"IQuest LoopCoder requires policy.megatron_config.lora_config.merge_lora=false "
	"so the shared physical-layer adapters are served through vLLM's LoRA path"
	)
	if getattr(hf_config_original, "model_type", None) == "iquestloopcoder" and isinstance(self, PolicyWorkerBase):
	if lora_config is None or lora_config.rank <= 0:
	raise ValueError("IQuest LoopCoder is supported only for LoRA training")
	if megatron_config.lora_config.merge_lora:
	raise ValueError(
	"IQuest LoopCoder requires policy.megatron_config.lora_config.merge_lora=false "
	"so the shared physical-layer adapters are served through vLLM's LoRA path"
	)

[gemini-code-assist]

The attributes self._shared_key and self._shared_value store references to the first-pass key and value tensors. Because these are stored as instance attributes on LoopCoderCoreAttention, they will remain in GPU memory even after the forward and backward passes of the current training step are completed. This can lead to unnecessary memory retention and potential fragmentation.

Since PyTorch's autograd automatically captures and saves any tensors needed for the backward pass during the forward execution of global_attention, we can safely clear these references at the end of the second-pass forward method to free up GPU memory.

Suggested change

	global_output = self.global_attention(
	query,
	self._shared_key,
	self._shared_value,
	attention_mask,
	**kwargs,
	)
	local_output = self.local_attention(query, key, value, attention_mask, **kwargs)
	gate = self.gate_projection(query)
	return local_output * (1.0 - gate) + global_output * gate
	global_output = self.global_attention(
	query,
	self._shared_key,
	self._shared_value,
	attention_mask,
	**kwargs,
	)
	self._shared_key = None
	self._shared_value = None
	local_output = self.local_attention(query, key, value, attention_mask, **kwargs)
	gate = self.gate_projection(query)
	return local_output * (1.0 - gate) + global_output * gate

[greptile-apps]

Confidence Score: 4/5

The implementation has no established behavioral blocker, but it is not ready to merge until the repository-required model comparison test and runnable training example are added.

Findings

1.  Missing model comparison test ▶
2.  Missing model training example ▶

Diagram

sequenceDiagram
    participant HF as HF LoopCoder checkpoint
    participant Bridge as IQuestLoopCoderBridge
    participant MCore as Shared Megatron decoder
    participant LoRA as LoRA adapter
    participant vLLM as vLLM LoopCoder

    HF->>Bridge: Load base configuration and weights
    Bridge->>MCore: Map QKV, MLP, norms, and frozen loop gates
    LoRA->>MCore: Attach adapters to shared physical projections
    MCore->>MCore: Pass 1: global attention and cache KV
    MCore->>MCore: Pass 2: local attention plus cached global attention
    MCore->>MCore: Mix outputs through frozen per-head gate
    MCore->>LoRA: Export adapter-only weights
    LoRA->>vLLM: Hot-load packed QKV and MLP adapters

Loading

_{Reviews (1) · Last reviewed commit: "feat(megatron): add LoRA-only IQuest Loo..."}

[greptile-apps]

Missing model comparison test

This adds new Megatron model support, but the tests only cover LoRA protocol declarations and isolated attention mixing. The repository's new-model checklist requires a generation-and-logprobs test comparing Hugging Face and Megatron behavior. Without that test, checkpoint mapping, recurrent model fidelity, and training-to-inference weight synchronization remain untested. This repository requirement must be satisfied before merging; add an IQuest case to the Megatron model test matrix, using a tiny checkpoint if needed.

Context Used: CLAUDE.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Missing model training example

This registers IQuest LoopCoder as a new Megatron model without the README and training script required by the repository's new-model checklist. This requirement must be satisfied before merging. Add a runnable examples/train/<model>/ example that documents the LoRA-only, unmerged-adapter, two-loop, and pipeline-parallel-size-one configuration.

Context Used: CLAUDE.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 52fd627. Configure here.}

[cursor]

Gate layer mishandles sequence parallel

High Severity

LoopGateProjection builds ColumnParallelLinear from the shared transformer config, so sequence_parallel stays enabled whenever TP is greater than 1. Queries reaching core attention are already full-sequence, and that extra dim-0 all-gather reshapes the gate to tp times the local hidden width. Loop-2 mixing then fails on the intended 40B path; the new tests never construct this layer.

Additional Locations (1)

• skyrl/backends/skyrl_train/workers/megatron/iquest_loopcoder.py#L43-L53

 

^{Reviewed by Cursor Bugbot for commit 52fd627. Configure here.}