feat(deepseek-v4): add native LoRA construction and conversion bridge

miles/backends/megatron_utils/bridge_lora_helpers.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import copy
 import logging
 from argparse import Namespace
 from dataclasses import dataclass
@@ -19,6 +20,21 @@
 
 logger = logging.getLogger(__name__)
 
+_DEEPSEEK_V4_MAIN_ATTENTION_TARGETS = frozenset({"wq_a", "wq_b", "wkv", "wo_a", "wo_b"})
+
+
+def _qualify_deepseek_v4_lora_targets(target_modules):
+    """Disambiguate V4's main attention leaves from nested DSA modules."""
+
+    if not target_modules:
+        return target_modules
+    is_string = isinstance(target_modules, str)
+    modules = [target_modules] if is_string else list(target_modules)
+    qualified = [
+        f"*.self_attention.{module}" if module in _DEEPSEEK_V4_MAIN_ATTENTION_TARGETS else module for module in modules
+    ]
+    return qualified[0] if is_string else qualified
+
 
 @dataclass
 class _BridgeWrapperConfig:
@@ -128,6 +144,11 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
     from megatron.bridge.training.config import DistributedDataParallelConfig
 
     hf_config = load_hf_config(args.hf_checkpoint)
+    from miles_plugins.megatron_bridge.deepseek_v4 import is_deepseek_v4_config
+
+    if is_deepseek_v4_config(hf_config):
+        return _setup_deepseek_v4_lora_model(args)
+
     bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
     provider = bridge.to_megatron_provider(load_weights=False)
 
@@ -207,3 +228,70 @@ def apply_lora_hook(model_chunks):
 
     model = provider.provide_distributed_model(wrap_with_ddp=True, ddp_config=ddp_config)
     return model
+
+
+def _setup_deepseek_v4_lora_model(args: Namespace) -> list:
+    """Build native DeepSeek-V4, applying PEFT before DDP wrapping."""
+
+    if is_multi_lora_enabled(args):
+        raise NotImplementedError("DeepSeek-V4 native construction does not yet support Multi-LoRA")
+
+    from megatron.bridge.models.model_provider import ModelProviderMixin, get_model as get_bridge_model
+    from megatron.bridge.training.config import DistributedDataParallelConfig
+    from megatron.core.enums import ModelType
+    from megatron.core.process_groups_config import ProcessGroupCollection
+
+    from .lora_utils import create_lora_instance
+    from .model_provider import get_model_provider_func
+
+    native_provider_func = get_model_provider_func(
+        args,
+        role="actor",
+        use_bridge_provider=False,
+    )
+
+    class _MilesDeepSeekV4Provider(ModelProviderMixin):
+        virtual_pipeline_model_parallel_size = args.virtual_pipeline_model_parallel_size
+        fp16 = args.fp16
+        bf16 = args.bf16
+
+        def provide(self, pre_process=None, post_process=None, vp_stage=None):
+            return native_provider_func(
+                pre_process=pre_process,
+                post_process=post_process,
+                vp_stage=vp_stage,
+            )
+
+    # The generic MLA aliases map bare ``wq_b`` to the DSA indexer's
+    # ``linear_wq_b``. V4 owns both that nested module and a main-attention
+    # ``wq_b``, so qualify the five main projections before PEFT's suffix
+    # matcher sees them. Explicit indexer targets remain unchanged.
+    lora_args = copy.copy(args)
+    lora_args.target_modules = _qualify_deepseek_v4_lora_targets(args.target_modules)
+    lora_args.exclude_modules = _qualify_deepseek_v4_lora_targets(getattr(args, "exclude_modules", None))
+    lora = create_lora_instance(lora_args)
+
+    def apply_lora_hook(model_chunks):
+        transformed = lora(model_chunks, training=True)
+        lora.set_params_to_save(transformed)
+        return transformed
+
+    use_distributed_optimizer = "muon" not in (args.optimizer or "").lower()
+    ddp_config = DistributedDataParallelConfig(
+        use_distributed_optimizer=use_distributed_optimizer,
+        grad_reduce_in_fp32=args.accumulate_allreduce_grads_in_fp32,
+    )
+    ddp_config.finalize()
+
+    if args.offload_train:
+        patch_param_grad_buffer_for_colocate_mode_lora()
+
+    return get_bridge_model(
+        _MilesDeepSeekV4Provider(),
+        ddp_config=ddp_config,
+        model_type=ModelType.encoder_or_decoder,
+        bf16=args.bf16,
+        fp16=args.fp16,
+        pre_wrap_hook=apply_lora_hook,
+        pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+    )

miles/backends/megatron_utils/hf_export.py

@@ -102,6 +102,15 @@ def _get_hf_bridge(hf_checkpoint: str):
     return AutoBridge.from_hf_pretrained(hf_checkpoint, trust_remote_code=True)
 
 
+def _uses_adapter_only_export(args, model: Sequence[DDP]) -> bool:
+    if not is_lora_model(model):
+        return False
+
+    from miles_plugins.megatron_bridge.deepseek_v4 import is_deepseek_v4_config
+
+    return is_deepseek_v4_config(load_hf_config(args.hf_checkpoint))
+
+
 def save_hf_model(
     args,
     rollout_id: int,
@@ -112,12 +121,17 @@ def save_hf_model(
 ) -> None:
     """Save Megatron model in HuggingFace format.
 
-    For LoRA models this saves both:
+    For most LoRA models this saves both:
     - A **merged** HF model (adapter weights folded into base) at ``{path}/``
       so it can be loaded directly with ``AutoModelForCausalLM.from_pretrained``.
     - An **adapter-only** HF PEFT checkpoint at ``{path}/adapter/``
       so it can be loaded with ``PeftModel.from_pretrained``.
 
+    DeepSeek-V4 uses an adapter-only portable layout. Its model is constructed
+    natively and the Bridge integration is conversion-only, so the complete
+    reload contract is the pinned base checkpoint plus ``{path}/adapter/``
+    rather than a partially merged base.
+
     This function is collective — all ranks must call it. On success, global rank 0
     writes a ``.complete`` marker file.
 
@@ -130,12 +144,22 @@ def save_hf_model(
     """
     should_log = get_parallel_state().effective_dp_cp.rank == 0 and get_parallel_state().tp.rank == 0
     path = Path(path if path is not None else args.save_hf.format(rollout_id=rollout_id))
+    adapter_only = _uses_adapter_only_export(args, model)
 
     try:
         if should_log:
             logger.info(f"Saving model in HuggingFace format to {path}")
 
-        if args.megatron_to_hf_mode == "raw" and not is_lora_model(model):
+        if adapter_only:
+            path.mkdir(parents=True, exist_ok=True)
+            if torch.distributed.get_rank() == 0:
+                (path / HF_EXPORT_COMPLETE_MARKER).unlink(missing_ok=True)
+            if should_log:
+                logger.info(
+                    "DeepSeek-V4 portable export is adapter-only; "
+                    "the frozen base remains the pinned native checkpoint"
+                )
+        elif args.megatron_to_hf_mode == "raw" and not is_lora_model(model):
             # LoRA keeps the bridge (adapter merging).
             hf_config = load_hf_config(args.hf_checkpoint)
             export_hf_model_direct(
@@ -164,7 +188,7 @@ def save_hf_model(
                         f"bridge likely has no mapping for this model architecture."
                     )
 
-        if should_log:
+        if should_log and not adapter_only:
             logger.info(f"Successfully saved merged HuggingFace model to {path}")
     except Exception as e:
         if raise_on_error:
@@ -179,7 +203,12 @@ def save_hf_model(
             adapter_path = path / "adapter"
             if should_log:
                 logger.info(f"Saving LoRA adapter (HF PEFT format) to {adapter_path}")
-            save_lora_checkpoint(model, args, str(adapter_path))
+            save_lora_checkpoint(
+                model,
+                args,
+                str(adapter_path),
+                require_hf_export=adapter_only,
+            )
             if should_log:
                 logger.info(f"Successfully saved LoRA adapter to {adapter_path}")
         except Exception as e:

miles/backends/megatron_utils/lora_utils.py

@@ -414,6 +414,7 @@ def save_lora_checkpoint(
     optimizer: Any | None = None,
     opt_param_scheduler: Any | None = None,
     iteration: int | None = None,
+    require_hf_export: bool = False,
 ) -> str:
     """Save LoRA adapter checkpoint to disk.
 
@@ -474,6 +475,9 @@ def save_lora_checkpoint(
             ):
                 lora_state_dict[hf_name] = weight
 
+        if not lora_state_dict:
+            raise RuntimeError("Megatron-Bridge exported no HF PEFT adapter tensors")
+
         if is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0:
             torch.save(lora_state_dict, save_path / "adapter_model.bin")
 
@@ -497,6 +501,8 @@ def save_lora_checkpoint(
             os.sync()
             logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors")
     except Exception as hf_export_err:
+        if require_hf_export:
+            raise RuntimeError("Required HF PEFT adapter export failed") from hf_export_err
         logger.warning(
             f"HF PEFT adapter export skipped ({hf_export_err}); the per-rank native "
             f"shards + training state are sufficient for training resume."

miles/backends/megatron_utils/model_provider.py

@@ -132,6 +132,8 @@ def forward(
 def get_model_provider_func(
     args: argparse.Namespace,
     role: Literal["actor", "critic"] = "actor",
+    *,
+    use_bridge_provider: bool = True,
 ):
     # Support custom model provider path (similar to --custom-rm-path for reward models)
     if getattr(args, "custom_model_provider_path", None):
@@ -161,7 +163,7 @@ def wrapped_model_provider(
 
         return wrapped_model_provider
 
-    if args.megatron_to_hf_mode == "bridge":
+    if use_bridge_provider and args.megatron_to_hf_mode == "bridge":
         from megatron.bridge import AutoBridge
 
         bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)

miles_plugins/megatron_bridge/__init__.py

@@ -75,3 +75,8 @@ def broadcast_obj_from_pp_rank(self, obj, name=None):
     from . import nemotron_h  # noqa: F401
 except Exception as _e:  # pragma: no cover - defensive
     logger.warning("miles nemotron_h plugin failed to load: %s", _e)
+
+try:
+    from . import deepseek_v4  # noqa: F401
+except Exception as _e:  # pragma: no cover - defensive
+    logger.warning("miles deepseek_v4 plugin failed to load: %s", _e)

miles_plugins/megatron_bridge/deepseek_v4.py

@@ -0,0 +1,112 @@
+"""Conversion-only Megatron-Bridge support for Miles' native DeepSeek-V4."""
+
+from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
+from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
+from megatron.bridge.models.conversion.param_mapping import (
+    AutoMapping,
+    ColumnParallelMapping,
+    ReplicatedMapping,
+)
+from megatron.bridge.models.deepseek.common import get_common_mapping_list
+from megatron.bridge.models.deepseek.deepseek_v3_bridge import DeepSeekV3Bridge
+from megatron.bridge.models.mla_provider import MLAModelProvider
+from megatron.core.models.gpt.gpt_model import GPTModel
+
+
+_DSV4_ATTENTION_MAPPINGS = {
+    "decoder.layers.*.self_attention.wq_a.weight": "model.layers.*.self_attn.wq_a.weight",
+    "decoder.layers.*.self_attention.q_norm.weight": "model.layers.*.self_attn.q_norm.weight",
+    "decoder.layers.*.self_attention.wq_b.weight": "model.layers.*.self_attn.wq_b.weight",
+    "decoder.layers.*.self_attention.wkv.weight": "model.layers.*.self_attn.wkv.weight",
+    "decoder.layers.*.self_attention.kv_norm.weight": "model.layers.*.self_attn.kv_norm.weight",
+    "decoder.layers.*.self_attention.wo_a.weight": "model.layers.*.self_attn.wo_a.weight",
+    "decoder.layers.*.self_attention.wo_b.weight": "model.layers.*.self_attn.wo_b.weight",
+}
+
+_DSV4_REPLICATED_MAPPINGS = {
+    "decoder.layers.*.hc_attn_fn": "model.layers.*.hc_attn_fn",
+    "decoder.layers.*.hc_attn_base": "model.layers.*.hc_attn_base",
+    "decoder.layers.*.hc_attn_scale": "model.layers.*.hc_attn_scale",
+    "decoder.layers.*.hc_ffn_fn": "model.layers.*.hc_ffn_fn",
+    "decoder.layers.*.hc_ffn_base": "model.layers.*.hc_ffn_base",
+    "decoder.layers.*.hc_ffn_scale": "model.layers.*.hc_ffn_scale",
+    "decoder.layers.*.self_attention.compressor.ape": "model.layers.*.self_attn.compressor.ape",
+    "decoder.layers.*.self_attention.compressor.wkv.weight": "model.layers.*.self_attn.compressor.wkv.weight",
+    "decoder.layers.*.self_attention.compressor.wgate.weight": "model.layers.*.self_attn.compressor.wgate.weight",
+    "decoder.layers.*.self_attention.compressor.norm.weight": "model.layers.*.self_attn.compressor.norm.weight",
+    "decoder.layers.*.self_attention.indexer.compressor.ape": "model.layers.*.self_attn.indexer.compressor.ape",
+    "decoder.layers.*.self_attention.indexer.compressor.wkv.weight": (
+        "model.layers.*.self_attn.indexer.compressor.wkv.weight"
+    ),
+    "decoder.layers.*.self_attention.indexer.compressor.wgate.weight": (
+        "model.layers.*.self_attn.indexer.compressor.wgate.weight"
+    ),
+    "decoder.layers.*.self_attention.indexer.compressor.norm.weight": (
+        "model.layers.*.self_attn.indexer.compressor.norm.weight"
+    ),
+    "decoder.layers.*.mlp.router.tid2eid": "model.layers.*.mlp.topk.tid2eid",
+    "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias",
+    "decoder.hc_head_params.hc_head_fn": "model.hc_head_fn",
+    "decoder.hc_head_params.hc_head_base": "model.hc_head_base",
+    "decoder.hc_head_params.hc_head_scale": "model.hc_head_scale",
+}
+
+_DSV4_COLUMN_PARALLEL_MAPPINGS = {
+    "decoder.layers.*.self_attention.attn_sink": "model.layers.*.self_attn.attn_sink",
+}
+
+_DSV4_AUTO_MAPPINGS = {
+    "decoder.layers.*.self_attention.indexer.linear_wq_b.weight": "model.layers.*.self_attn.indexer.wq_b.weight",
+    "decoder.layers.*.self_attention.indexer.linear_weights_proj.weight": (
+        "model.layers.*.self_attn.indexer.weights_proj.weight"
+    ),
+}
+
+
+def is_deepseek_v4_config(hf_config) -> bool:
+    architectures = getattr(hf_config, "architectures", None) or []
+    return bool(architectures and architectures[0] == "DeepseekV4ForCausalLM")
+
+
+def _get_dsv4_explicit_mappings():
+    mappings = [
+        AutoMapping(megatron_param=megatron_param, hf_param=hf_param)
+        for megatron_param, hf_param in _DSV4_ATTENTION_MAPPINGS.items()
+    ]
+    mappings.extend(
+        ReplicatedMapping(megatron_param=megatron_param, hf_param=hf_param)
+        for megatron_param, hf_param in _DSV4_REPLICATED_MAPPINGS.items()
+    )
+    mappings.extend(
+        ColumnParallelMapping(megatron_param=megatron_param, hf_param=hf_param)
+        for megatron_param, hf_param in _DSV4_COLUMN_PARALLEL_MAPPINGS.items()
+    )
+    mappings.extend(
+        AutoMapping(megatron_param=megatron_param, hf_param=hf_param)
+        for megatron_param, hf_param in _DSV4_AUTO_MAPPINGS.items()
+    )
+    return mappings
+
+
+@MegatronModelBridge.register_bridge(
+    source="DeepseekV4ForCausalLM",
+    target=GPTModel,
+    provider=MLAModelProvider,
+    model_type="deepseek_v4",
+)
+class MilesDeepSeekV4Bridge(DeepSeekV3Bridge):
+    """Map the native V4 graph for adapter publication and export only."""
+
+    def provider_bridge(self, hf_pretrained):
+        raise RuntimeError(
+            "DeepSeek-V4 model construction must use Miles' native dsv4 provider; "
+            "this bridge is conversion-only"
+        )
+
+    def mapping_registry(self) -> MegatronMappingRegistry:
+        mappings = get_common_mapping_list(hf_config=self.hf_config)
+        mappings.extend(_get_dsv4_explicit_mappings())
+        return MegatronMappingRegistry(*mappings)
+
+
+__all__ = ["MilesDeepSeekV4Bridge", "is_deepseek_v4_config"]