Work around grouped-expert LoRA export and checkpoint load in pinned Miles/Bridge

[kevintli]

Summary

While trying to get gpt-oss-20b working on Spindle, I noticed a few bugs in Miles and Megatron-Bridge which caused our sampler policy to diverge from the trainer policy over the course of training. This showed up as unusually large KL and worse reward curves compared to the original blog post results. This PR implements a workaround on the Spindle side before we upstream the actual bugfixes to those repos.

Observed bugs

1. Export (Megatron-Bridge gpt-oss integration): Due to a bug with expert parallelism shape handling, Megatron-Bridge was exporting bad expert LoRA weights and SGLang would throw them away silently at load time. This means rollout workers were essentially running with untrained expert weights.
2. Load (Miles lora/checkpoint.load_slot): Miles unintentionally throws away the result of dist_checkpointing.load, which means the first training step uses uninitialized weights (until optim_step copies over the correct ones from the optimizer state), causing a temporary spike in KL / instability.

Workaround

This PR implements miles_runtime/expert_lora_compat.py, which runs only when EXPERT_LORA_COMPAT=1. That compatibility wrapper does the following:

1. At checkpoint write time: all-gather the expert LoRA tensors over EP groups and rewrite the checkpoint in the correct format
2. At load time: Rebuild the state dict and correctly copy the params into the live weights

Safety checks:

• The rewrite requires the published tensor to be bit-identical to a prefix of the gathered experts, with an optional leading singleton axis. It raises on any other layout.
• It leaves tensors alone that are already (E, ...), and logs publish_noop and restore_noop when upstream is already correct.
• It asserts PP=1 and expert-TP=1.

Removal: once upstream fixes both bugs, delete the module and the two enabled() branches.

Validation

Setup: gpt-oss-20b, TP=8 EP=8, sec-search-rl, full 24-update runs

The plots below show kl_sample_train_v2 before and after this bug fix. Notice that without the fix:

• KL grows steadily over time, since the trainer diverges from the untrained-expert-LoRA rollout weights
• When we restart the run from a checkpoint, KL temporarily drops before climbing steadily -- this is because with bug 2 above, the trainer temporarily has untrained LoRAs (matching the bug on the LoRA side) before it gets fixed on the first optimizer step

Tests:

• tests/backends/test_expert_lora_compat.py covers the 4→32 rewrite, the no-op on complete exports, rejection of mismatched or missing tensors, the param copy, and the master refresh
• Full pytest (with torch installed) gives 714 passed, 1 skipped
• Ruff check and format are clean on the new files

Link to Devin session: https://modal.devinenterprise.com/sessions/f53cfabb210146de8f0338fe388d7973
Open in Devin Desktop: https://modal.devinenterprise.com/desktop/session/f53cfabb210146de8f0338fe388d7973?variant=devin
Requested by: @kevintli

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access that start with 'Devin'.

• Disable automatic comment, CI, and merge conflict monitoring