[lora] refactor - native LoRA

[yushengsu-thu]

Refactor PR into #1792's branch (zhichen/lora-native), so the native-LoRA extraction ships as a plugin from the start. Adapter math is unchanged — a side-by-side harness against the base monolith shows bit-identical adapter parameters, forward outputs, backward gradients, and HF exports for full-target GQA (TP=1, TP=2 ± SP) and MLA configs. Bridge-mode LoRA is untouched.

Includes #2072: a function-level redesign on top of the extraction, aimed at one-declaration support for new models, plus the fixes found while validating four models end-to-end.

Layout

miles_plugins/lora/
    __init__.py
    lora.py              # Single public entry point and lifecycle orchestration
    config.py            # LoRAConfig: rank, alpha, dropout, and target modules
    registry.py          # model_type -> ModelEntry(spec, status); launcher helpers
    hf_adapter.py        # HF/PEFT adapter naming, import, and export
    serving.py           # SGLang serving export (fused-buffer zero-fill, target expansion)
    checkpointing.py     # Adapter checkpoint state dicts and per-rank shard naming
    distributed.py       # TP/EP/SP communication and gradient handling

    spec/
      base.py            # Contracts: ProjectionSpec, AttachContext, protocols, enums
      layout.py          # Declarative layouts: dim resolvers, bindings, attach walk, spec bases
      attention.py       # GQA / MLA / GDN / hybrid attention spec classes
      mlp.py             # Fused gate/up + down spec class
      moe.py             # Routed-expert per-layer policy (shared-outer aware)

    modules/
      linear.py          # LoRALinear + LoRASplitAdapter (SplitQKV / SplitFC1)
      moe.py             # Future routed/grouped-expert modules

    kernel/              # Future fused and grouped-expert kernels

Design (from #2072)

Classes carry the taxonomy; every fact is declared once and derived everywhere else.

• Declarative layouts (spec/layout.py): an architecture spec is a class holding a ModuleLayout table — which projections exist, on which physical linear, with which shard geometry. One shared attach_layout() implements the existence-check → target-filter → guard → build → hook walk; the dim-resolver section is the only spec-side reader of MCore attribute names. A new layout is ~15 lines of facts.
• Spec hierarchy (LayoutSpec → AttentionSpecBase → GQA/MLA/GDN, HybridGQAGDN(GQAAttentionSpec), SharedOuterExpertMoESpec(GeneralExpertMoESpec)): supported_targets, the canonical --target-modules order, and SGLang's fused-family expansion all derive from the layout declaration; a new architecture's fused groups reach serving-side target expansion automatically. ShardLayout/AttentionFamily str-enums replace bare strings; no parallel name tables in the exporters.
• ModelEntry(spec, status, reason) with VALIDATED / STRUCTURAL / UNSTABLE: support status is one first-class field; UNSTABLE entries carry their reason and warn at attach.
• Public export descriptors: adapters expose exports() -> ProjectionExport (with SGLangFusedGroup metadata); serving.py consumes only that API.
• Launcher helpers: default_target_modules(hf_checkpoint) and preflight_native_lora(...) (registry ∧ mbridge bridge ∧ model-args audit, no GPU) — launch scripts stop re-declaring per-family target strings and fail in seconds when a conversion bridge is missing.

Fixes found by running four models natively (from #2072)

• miles_plugins/mbridge/kimi_k25.py (new): raw-mode K2.5 conversion needs an mbridge bridge (bridge mode uses megatron.bridge's own KimiK25VL bridge and never converts). Thin DeepseekV3Bridge subclass: text_config descent + language_model. prefix remap.
• megatron_to_hf kimi dispatch: _convert_to_hf_core matched only kimi_k25, but the direct weight iterator passes kimik25config (dead branch); the converter also lacked the raw text-only provider's unprefixed top-level weight names.
• convert_kimi_int4_to_bf16.py strips quantization_config: a surviving one sends SGLang through its CompressedTensors path, which serves the BF16 checkpoint with a degenerate context-free forward → garbage rollouts, train/rollout logprob_abs_diff ~2.2. Verified by a megatron-vs-HF-vs-SGLang three-way logprob comparison.
• torch_dsa_topk clamps k to the key length: short sequences (< index_topk keys) crashed torch.topk in raw-mode log-prob compute.
• run_glm5_2_744b_a40b_lora_native.py (new): raw-mode GLM-5.2 launcher; its toy conversion must be single-rank (any PP split of the 5-layer toy starts a stage on a DSA skip layer).

Validation

20-rollout native GRPO runs (wandb ch271828n-team/miles_lora_native), acceptance ≈ 0.01 train/rollout logprob_abs_diff:

model	logprob_abs_diff
Qwen3-8B	0.0099
Qwen3.5-35B-A3B	0.0105 (grad_norm ~1e-2 throughout; the recorded GDN raw-backward divergence did not reproduce → qwen3_5_moe VALIDATED)
GLM-5.2 5-layer	0.0096
Kimi-K2.5 2-layer	0.0124 (after the quantization_config fix)

tests/fast/miles_plugins/lora: 47 passed. Behavior-preservation smoke: a 2-rollout Qwen3-8B run reproduces the baseline's first-rollout logprob_abs_diff (0.00967 vs 0.00978). Adapter attribute names are pinned, so checkpoint keys are unchanged.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[yushengsu-thu]

fix them

[yushengsu-thu]

To-do

• dsk-v4 support: because we have not support dsk-v4 lora on sglang side yet. Thus, I reomve dsk-v4 part and we can add it after sglang side support.

• Refer to megatron-bridge to compelete the features at below layouts:

• • miles_plugins/lora/spec/moe.py
• • miles_plugins/lora/distributed.py

• [I added] moe lora: in native branch, it only supports shared moe lora (bth, that's tml special design moe lora). Another general moe lora is non-shared and this is implemented in sglang and megatron-bridge already as well. We can refer to that part when we refactor megatron-bridge lora in miles. (Refer to different type of lora in megatron-bridge).