[megatron] Enable grouped outer expert LoRA targeting for higher accuracy

[casper-hansen]

Add megatron_config.lora_config.experts_shared_outer_loras: fc1 (gate_up) lora_A and fc2 (down) lora_B are shared across all experts while the inner matrices are trained per expert, matching Megatron-Bridge's LoRA(experts_shared_outer_loras=True).

For merge_lora=false sync, vLLM has no shared-expert-LoRA contract, so the adapter export stacks the per-expert emissions of packed-HF models back into 3D and expands the shared side to all experts in
_convert_moe_experts_lora_to_vllm (per-expert-HF models get indexed-key replication instead). Also patches Megatron-Bridge's PackedPerExpertLinear.sharded_state_dict, which omits the required pg_collection kwarg and breaks dist-checkpoint saving.

Note that we bump megatron bridge to fbf2d3c07f8813b86886b3102a1a9ce0532431c4 which is a newer commit that adds this feature.

Held-out GSM8K accuracy (1319 problems, greedy)

I ran a few examples to demonstrate that it works. We observe below that the grouped expert approach improves faster.

	step 0	step 8	step 16	RL gain (0→16)
OFF seed 42	79.61	82.03	83.70	+4.09
OFF seed 7	79.76	82.87	84.08	+4.32
ON seed 42	79.61	82.87	85.06	+5.45
ON seed 7	79.68	83.47	84.69	+5.01

Checkpoint size

One caveat is that this seems to produce much larger checkpoints.

Measured from the real Qwen3.6-35B-A3B artifacts (rank 32, TP=4/EP=8, 40 layers) for ON, with OFF computed by the same per-tensor accounting.

	shared-outer ON (measured)	OFF, default (computed)	ratio
Distinct trained adapter params	521 M (1.04 GB bf16)	70.5 M (0.14 GB bf16)	7.4×
Per-rank adapter_*.pt in checkpoint	142.8 MB	~20.8 MB	6.9×
All 8 rank files	1.14 GB	~0.17 GB	6.9×
Optimizer state (.distcp shards, fp32 + Adam)	6.7 GB measured	~0.9 GB	~7.4×
Total policy checkpoint dir	7.4 GB measured	~1.0–1.1 GB	~7×

[gemini-code-assist]

Code Review

This pull request introduces support for shared-outer grouped-expert LoRA in MoE models on the Megatron backend, including layout conversion to vLLM's flat PEFT format and a patch for sharded state dict saving. The review feedback highlights several improvement opportunities: retrieving the number of MoE experts from the model configuration instead of the provider to ensure reliability, raising explicit errors instead of silently skipping tensor expansion or replication when the expert count is missing, and validating that this Megatron-specific feature is only enabled when using the Megatron training strategy.

[gemini-code-assist]

Using getattr(self.provider, "num_moe_experts", None) to retrieve the number of MoE experts is unreliable, as the provider instance may not directly expose this attribute. Instead, retrieve it from the model's configuration via get_model_config(self.actor_module[0]), which is the canonical source of truth for the instantiated model's architecture.

Suggested change

	adapter_state = _convert_moe_experts_lora_to_vllm(
	adapter_state, num_moe_experts=getattr(self.provider, "num_moe_experts", None)
	)
	model_config = get_model_config(self.actor_module[0])
	adapter_state = _convert_moe_experts_lora_to_vllm(
	adapter_state, num_moe_experts=getattr(model_config, "num_moe_experts", None)
	)

[gemini-code-assist]

If num_moe_experts is None but we encounter a shared-outer expert LoRA tensor (indicated by tensor.shape[0] == 1), the code will silently skip the expansion. This results in a tensor of shape (rank, in) instead of (rank * E, in), leading to a shape mismatch error in vLLM. We should raise a clear ValueError instead of silently continuing with an invalid shape.

Suggested change

	if (is_gate_up or is_down) and tensor.ndim == 3 and not uses_indexed_expert_keys:
	if num_moe_experts is not None and tensor.shape[0] == 1 and num_moe_experts > 1:
	tensor = tensor.expand(num_moe_experts, -1, -1)
	if (is_gate_up or is_down) and tensor.ndim == 3 and not uses_indexed_expert_keys:
	if tensor.shape[0] == 1:
	if num_moe_experts is None:
	raise ValueError(
	f"Found shared-outer expert LoRA tensor '{key}' with shape[0] == 1, "
	"but `num_moe_experts` is None. Cannot expand to all experts."
	)
	if num_moe_experts > 1:
	tensor = tensor.expand(num_moe_experts, -1, -1)

[gemini-code-assist]

If num_moe_experts is None but we encounter a shared-outer expert LoRA tensor for a per-expert-HF model, the replication block is silently skipped. This leaves the shared keys as-is, which vLLM's PEFT loader will ignore, leading to a silent correctness bug where the shared projections run with uninitialized/zero LoRA weights. We should raise a clear ValueError if num_moe_experts is None.

Suggested change

	if shared_match is not None and tensor.ndim == 3 and tensor.shape[0] == 1 and num_moe_experts is not None:
	# Per-expert-HF model: replicate the shared side into the indexed
	# per-expert keys vLLM's PEFT loader parses.
	insert_pos = key.rindex(".mlp.experts.") + len(".mlp.experts.")
	for expert_idx in range(num_moe_experts):
	converted[f"{key[:insert_pos]}{expert_idx}.{key[insert_pos:]}"] = tensor[0].clone()
	continue
	if shared_match is not None and tensor.ndim == 3 and tensor.shape[0] == 1:
	if num_moe_experts is None:
	raise ValueError(
	f"Found shared-outer expert LoRA tensor '{key}' with shape[0] == 1, "
	"but `num_moe_experts` is None. Cannot replicate to indexed expert keys."
	)
	# Per-expert-HF model: replicate the shared side into the indexed
	# per-expert keys vLLM's PEFT loader parses.
	insert_pos = key.rindex(".mlp.experts.") + len(".mlp.experts.")
	for expert_idx in range(num_moe_experts):
	converted[f"{key[:insert_pos]}{expert_idx}.{key[insert_pos:]}"] = tensor[0].clone()
	continue

[gemini-code-assist]

We should validate that experts_shared_outer_loras is only enabled when trainer.strategy is set to "megatron", as this is a Megatron-specific feature and will be silently ignored under FSDP.

        megatron_lora_cfg = cfg.trainer.policy.megatron_config.lora_config
        if megatron_lora_cfg.experts_shared_outer_loras:
            if cfg.trainer.strategy != "megatron":
                raise ValueError(
                    "`experts_shared_outer_loras` is only supported with `trainer.strategy='megatron'`"
                )
            if megatron_lora_cfg.lora_type != "lora":
                raise ValueError(
                    "`megatron_config.lora_config.experts_shared_outer_loras` is only supported with "
                    f'`lora_type="lora"`, got lora_type="{megatron_lora_cfg.lora_type}"'
                )