[lora] declarative layouts, support-status registry

miles/backends/megatron_utils/lora_utils.py

@@ -168,7 +168,7 @@ def _adapter_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, *, ep_sharded:
     hold routed-expert adapter state; native state is EP-invariant
     (``native_adapter_shard_name``).
     """
-    from miles_plugins.lora.codec.checkpoint import native_adapter_shard_name
+    from miles_plugins.lora.checkpointing import native_adapter_shard_name
 
     name = native_adapter_shard_name(tp_rank, pp_rank)
     if ep_sharded and ep_rank > 0:
@@ -178,7 +178,7 @@ def _adapter_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, *, ep_sharded:
 
 def _adapter_shards_are_ep_sharded(model: Sequence[torch.nn.Module]) -> bool:
     """Derive the shard policy from the modules that were actually attached."""
-    from miles_plugins.lora.codec.checkpoint import has_native_adapters
+    from miles_plugins.lora.checkpointing import has_native_adapters
 
     return not has_native_adapters(model)
 
@@ -192,7 +192,7 @@ def _non_native_adapter_load_plan(model, state_dict):
     before any copy; the native codec remains the sole owner of strict,
     chunk-qualified checkpoint semantics.
     """
-    from miles_plugins.lora.codec.checkpoint import AdapterLoadPlan
+    from miles_plugins.lora.checkpointing import AdapterLoadPlan
 
     assignments = []
     shape_mismatches = []
@@ -366,8 +366,8 @@ def target_modules_hf_for_sglang_rollout(args: Namespace) -> list[str]:
     """HF target_modules for SGLang LoRA init/sync (minus _SGLANG_UNSUPPORTED_HF_TARGETS, currently empty)."""
     raw = list(args.target_modules) if args.target_modules else []
     if uses_builtin_native_lora_provider(args):
-        from miles_plugins.lora.codec.sglang import expand_sglang_target_modules
         from miles_plugins.lora.config import LoRAConfig
+        from miles_plugins.lora.serving import expand_sglang_target_modules
 
         hf_checkpoint = getattr(args, "hf_checkpoint", None)
         if hf_checkpoint:
@@ -576,7 +576,7 @@ def save_lora_checkpoint(
                     if _is_adapter_param_name(name):
                         adapter_state[name] = parameter.detach().cpu()
         else:
-            from miles_plugins.lora.codec.checkpoint import native_adapter_state_dict
+            from miles_plugins.lora.checkpointing import native_adapter_state_dict
 
             adapter_state = native_adapter_state_dict(model)
 
@@ -615,7 +615,7 @@ def save_lora_checkpoint(
                 else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
             )
         else:
-            from miles_plugins.lora.codec.hf import target_modules_from_hf_names
+            from miles_plugins.lora.hf_adapter import target_modules_from_hf_names
 
             target_modules_hf = target_modules_from_hf_names(lora_state_dict)
         config = {
@@ -694,7 +694,7 @@ def load_lora_adapter(
     ep_rank = get_parallel_state().ep.rank
 
     # ---- Try Megatron-native format first (fast, no conversion needed) ----
-    from miles_plugins.lora.codec.checkpoint import native_adapter_load_plan
+    from miles_plugins.lora.checkpointing import native_adapter_load_plan
 
     # Shard key follows the provider: only bridge/custom adapters are EP-sharded.
     ep_sharded_provider = _adapter_shards_are_ep_sharded(model)

miles/backends/megatron_utils/megatron_to_hf/__init__.py

@@ -58,7 +58,9 @@ def _convert_to_hf_core(args, model_name, name, param):
         converted_named_tensors = convert_mimo_to_hf(args, name, param)
     elif "kimivl" in model_name:
         converted_named_tensors = convert_kimivl_to_hf(args, name, param)
-    elif "kimi_k25" in model_name:
+    elif "kimi_k25" in model_name or "kimik25" in model_name:
+        # KimiK25Config lowercases to "kimik25config"; the underscore spelling
+        # covers callers that pass the HF model_type instead.
         converted_named_tensors = convert_kimi_k25_to_hf(args, name, param)
     else:
         raise ValueError(f"Unsupported model: {model_name}")

miles/backends/megatron_utils/megatron_to_hf/kimi_vl.py

@@ -28,11 +28,23 @@ def convert_kimi_k25_to_hf(args, name, param):
 
 
 def convert_language_model_to_hf(args, name, param):
-    if name == "module.module.language_model.embedding.word_embeddings.weight":
+    # The VL trainer prefixes the text stack with language_model.; the raw
+    # text-only provider (native-LoRA K2.5 runs) does not. HF names always
+    # carry the multimodal shell's language_model. prefix.
+    if name in (
+        "module.module.language_model.embedding.word_embeddings.weight",
+        "module.module.embedding.word_embeddings.weight",
+    ):
         return [("language_model.model.embed_tokens.weight", param)]
-    if name == "module.module.language_model.output_layer.weight":
+    if name in (
+        "module.module.language_model.output_layer.weight",
+        "module.module.output_layer.weight",
+    ):
         return [("language_model.lm_head.weight", param)]
-    if name == "module.module.language_model.decoder.final_layernorm.weight":
+    if name in (
+        "module.module.language_model.decoder.final_layernorm.weight",
+        "module.module.decoder.final_layernorm.weight",
+    ):
         return [("language_model.model.norm.weight", param)]
 
     try:

miles_plugins/lora/__init__.py

@@ -8,18 +8,21 @@
 - Import ``miles``/``megatron`` inside functions, never at module level.
 """
 
-from miles_plugins.lora.codec.sglang import export_lora_sglang_named
 from miles_plugins.lora.lora import (
     apply_native_lora,
     export_lora_hf_named,
     load_lora_adapter_hf,
     wrap_model_provider_with_lora,
 )
+from miles_plugins.lora.registry import default_target_modules, preflight_native_lora
+from miles_plugins.lora.serving import export_lora_sglang_named
 
 __all__ = [
     "apply_native_lora",
+    "default_target_modules",
     "export_lora_hf_named",
     "export_lora_sglang_named",
     "load_lora_adapter_hf",
+    "preflight_native_lora",
     "wrap_model_provider_with_lora",
 ]

miles_plugins/lora/codec/checkpoint.py → miles_plugins/lora/checkpointing.py

@@ -18,7 +18,6 @@
 
 from miles_plugins.lora.modules.linear import iter_adapters
 
-
 _MODEL_CHUNK_PREFIX = "_miles_model_chunks."
 
 
@@ -190,7 +189,7 @@ def has_native_adapters(model_chunks: Sequence[nn.Module]) -> bool:
 
     Concrete ``NativeLoRAAdapter`` classes are deliberately not enough: a
     custom provider may compose one with additional adapter types that the
-    native checkpoint codec cannot serialize. ``apply_native_lora`` marks every
+    native checkpoint helpers cannot serialize. ``apply_native_lora`` marks every
     successful chunk, including legitimate PP/VPP chunks with zero local
     targets, so save and load share one explicit provider contract.
     """

miles_plugins/lora/distributed.py

@@ -124,23 +124,24 @@ def flush(self) -> None:
             offset += size
 
 
-_marked_lora_grad_params_cache: dict[int, list] = {}
-
-
 def reduce_marked_lora_grads(model: Sequence[nn.Module]) -> None:
     """Sum partial gradients for replicated native-LoRA parameters."""
     from megatron.core import parallel_state as ps
 
-    key = id(model[0]) if model else 0
-    marked = _marked_lora_grad_params_cache.get(key)
+    if not model:
+        return
+    # Cache the marked-parameter scan on the leading chunk so the entry's
+    # lifetime is the model's lifetime (an id()-keyed module-global here would
+    # outlive rebuilt models and can collide with recycled ids).
+    marked = getattr(model[0], "_miles_lora_marked_grad_params", None)
     if marked is None:
         marked = []
         for chunk in model:
             for param in chunk.parameters():
                 group_name = getattr(param, "_lora_grad_sum_group", None)
                 if group_name is not None and param.requires_grad:
                     marked.append((param, group_name))
-        _marked_lora_grad_params_cache[key] = marked
+        model[0]._miles_lora_marked_grad_params = marked
     if not marked:
         return