fix(lora): deduplicate native adapter checkpoints

miles/backends/megatron_utils/lora_checkpoint.py

@@ -0,0 +1,88 @@
+"""Naming and topology helpers shared by the single- and multi-LoRA checkpoint paths."""
+
+from pathlib import Path
+
+import torch.distributed as dist
+
+from miles.backends.training_utils.parallel import get_parallel_state
+from miles.utils.distributed_utils import get_gloo_group
+
+# Cached by adapter_shard_topology(); the topology is fixed for the run.
+_shard_topology: tuple[bool, tuple[tuple[int, int, int, int], ...]] | None = None
+
+
+def raise_if_any_rank_failed(local_error: Exception | None, operation: str) -> None:
+    """Raise on every rank when any rank reports ``local_error``, so failures stay collective."""
+    message = None if local_error is None else f"{type(local_error).__name__}: {local_error}"
+    if dist.is_initialized():
+        group = get_gloo_group()
+        messages: list[str | None] = [None] * dist.get_world_size(group=group)
+        dist.all_gather_object(messages, message, group=group)
+        message = next((item for item in messages if item is not None), None)
+
+    if message is not None:
+        error = RuntimeError(f"{operation} failed on at least one rank: {message}")
+        if local_error is not None:
+            raise error from local_error
+        raise error
+
+
+def megatron_shard_name(
+    tp_rank: int,
+    pp_rank: int,
+    ep_rank: int = 0,
+    etp_rank: int = 0,
+    *,
+    ep_size: int = 1,
+    etp_size: int = 1,
+    tp_size: int = 1,
+) -> str:
+    """Adapter shard name for one ``(tp, pp, ep, etp)`` coordinate; EP ranks hold different
+    local experts and ETP ranks hold different slices of them. The ep suffix is omitted at
+    ``ep_size == 1``, and the etp suffix at ``etp_size <= tp_size`` where the etp rank is a
+    function of the tp rank, so checkpoints written before either existed stay loadable."""
+    name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}"
+    if ep_size > 1:
+        name += f"_ep{ep_rank}"
+    if etp_size > tp_size:
+        name += f"_etp{etp_rank}"
+    return name + ".pt"
+
+
+def local_shard_name() -> str:
+    """Shard name for this rank's model-parallel coordinate."""
+    parallel_state = get_parallel_state()
+    return megatron_shard_name(
+        parallel_state.tp.rank,
+        parallel_state.pp.rank,
+        parallel_state.ep.rank,
+        parallel_state.etp.rank,
+        ep_size=parallel_state.ep.size,
+        etp_size=parallel_state.etp.size,
+        tp_size=parallel_state.tp.size,
+    )
+
+
+def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int, int], ...]]:
+    """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep, etp) coords)`` via one cached gloo all-gather."""
+    global _shard_topology
+    if _shard_topology is not None:
+        return _shard_topology
+
+    parallel_state = get_parallel_state()
+    coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank, parallel_state.etp.rank)
+    if not dist.is_initialized():
+        _shard_topology = (True, (coords,))
+        return _shard_topology
+
+    current_rank = dist.get_rank()
+    group = get_gloo_group()
+    gathered: list[object] = [None] * dist.get_world_size(group=group)
+    dist.all_gather_object(gathered, (coords, current_rank), group=group)
+    is_writer = current_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords)
+    _shard_topology = (is_writer, tuple(sorted({entry_coords for entry_coords, _ in gathered})))
+    return _shard_topology
+
+
+def all_megatron_checkpoints_exist(step_dir: Path, shard_names) -> bool:
+    return all((step_dir / name).exists() for name in shard_names)

miles/backends/megatron_utils/lora_utils.py

@@ -10,6 +10,12 @@
 import torch
 import torch.distributed as dist
 
+from miles.backends.megatron_utils.lora_checkpoint import (
+    adapter_shard_topology,
+    local_shard_name,
+    megatron_shard_name,
+    raise_if_any_rank_failed,
+)
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.utils.lora import is_lora_enabled, lora_rollout_enabled  # noqa: F401  (re-exported)
 
@@ -404,6 +410,30 @@ def create_lora_instance(args: Namespace):
 # ---------------------------------------------------------------------------
 
 
+def _save_native_adapter_checkpoint(model: Sequence[torch.nn.Module], directory: Path) -> Path | None:
+    is_writer, _ = adapter_shard_topology()
+    path = None
+    adapter_state = {}
+    save_error = None
+    if is_writer:
+        try:
+            adapter_state = {
+                name: param.data.cpu()
+                for model_chunk in model
+                for name, param in model_chunk.named_parameters()
+                if _is_adapter_param_name(name)
+            }
+            path = directory / local_shard_name()
+            torch.save(adapter_state, path)
+        except Exception as error:
+            save_error = error
+
+    raise_if_any_rank_failed(save_error, "Native LoRA checkpoint save")
+    if path is not None:
+        logger.info(f"Saved {len(adapter_state)} adapter tensors (native) to {path}")
+    return path
+
+
 def save_lora_checkpoint(
     model: Sequence[torch.nn.Module],
     args: Namespace,
@@ -419,16 +449,15 @@ def save_lora_checkpoint(
     1. **HF PEFT format** (``adapter_model.bin`` + ``adapter_config.json``) for
        external tool compatibility. Uses Megatron-Bridge's ``export_adapter_weights``
        which correctly handles fused QKV / gate-up weight splitting and TP gathering.
-    2. **Megatron-native format** (``adapter_megatron_rank{global_rank}.pt``) for fast
-       checkpoint resume without name/weight conversion. Each TP/PP rank saves its
-       own shard with original parameter names.
+    2. **Megatron-native format** (``adapter_megatron_tp{tp}_pp{pp}[_ep{ep}][_etp{etp}].pt``)
+       for fast checkpoint resume without name/weight conversion.
 
     When ``optimizer`` is provided, training state (optimizer + LR scheduler) is
     also saved per-rank for checkpoint resume. Base model weights are frozen and
     never change, so they are not saved.
 
-    This function is collective: **all ranks must call it** because the bridge
-    export performs TP all-gather internally. Only ``dp_rank == 0`` writes files.
+    This function is collective: all ranks participate in topology discovery and
+    Bridge export, while one writer stores each model-parallel native shard.
     """
     import json
 
@@ -446,16 +475,7 @@ def save_lora_checkpoint(
     if dist.is_initialized():
         dist.barrier()
 
-    adapter_state: dict[str, torch.Tensor] = {}
-    for model_chunk in model:
-        for name, param in model_chunk.named_parameters():
-            if _is_adapter_param_name(name):
-                adapter_state[name] = param.data.cpu()
-
-    global_rank = dist.get_rank() if dist.is_initialized() else 0
-    native_path = save_path / f"adapter_megatron_rank{global_rank}.pt"
-    torch.save(adapter_state, native_path)
-    logger.info(f"Saved {len(adapter_state)} adapter tensors (native) to {native_path}")
+    _save_native_adapter_checkpoint(model, save_path)
 
     # ---- HF PEFT format (uses bridge for correct name/weight conversion) ----
     # Bridge export is collective: all TP ranks participate in the all-gather,
@@ -496,7 +516,7 @@ def save_lora_checkpoint(
             logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors")
     except Exception as hf_export_err:
         logger.warning(
-            f"HF PEFT adapter export skipped ({hf_export_err}); the per-rank native "
+            f"HF PEFT adapter export skipped ({hf_export_err}); the model-parallel native "
             f"shards + training state are sufficient for training resume."
         )
 
@@ -528,8 +548,8 @@ def load_lora_adapter(
 ) -> tuple[bool, int | None]:
     """Load LoRA adapter weights from a saved checkpoint into the model.
 
-    Attempts to load from Megatron-native format first (per-rank ``.pt`` files),
-    which preserves the exact TP/PP sharding and requires no name conversion.
+    Attempts to load from Megatron-native format first (per-model-parallel ``.pt``
+    files), which preserves the exact TP/PP/EP sharding and requires no name conversion.
     Falls back to HF PEFT ``adapter_model.bin`` if native files are not found
     (not yet implemented for HF PEFT format).
 
@@ -552,17 +572,24 @@ def load_lora_adapter(
         logger.warning(f"LoRA adapter path does not exist: {adapter_dir}")
         return False, None
 
-    tp_rank = get_parallel_state().tp.rank
-    pp_rank = get_parallel_state().pp.rank
+    parallel_state = get_parallel_state()
+    tp_rank = parallel_state.tp.rank
+    pp_rank = parallel_state.pp.rank
+    ep_size = parallel_state.ep.size
 
     # ---- Try Megatron-native format first (fast, no conversion needed) ----
-    global_rank = dist.get_rank() if dist.is_initialized() else 0
-    native_path = adapter_dir / f"adapter_megatron_rank{global_rank}.pt"
+    native_path = adapter_dir / local_shard_name()
     if not native_path.exists():
-        legacy = adapter_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt"
-        if legacy.exists():
-            logger.warning(f"Using legacy tp/pp-named adapter shard {legacy}; only valid when EP<=TP")
-            native_path = legacy
+        global_rank = dist.get_rank() if dist.is_initialized() else 0
+        legacy_rank = adapter_dir / f"adapter_megatron_rank{global_rank}.pt"
+        if legacy_rank.exists():
+            logger.warning(f"Using legacy global-rank adapter shard {legacy_rank}")
+            native_path = legacy_rank
+        elif ep_size > 1:
+            legacy_tp_pp = adapter_dir / megatron_shard_name(tp_rank, pp_rank)
+            if legacy_tp_pp.exists():
+                logger.warning(f"Using legacy tp/pp-named adapter shard {legacy_tp_pp}; only valid when EP<=TP")
+                native_path = legacy_tp_pp
     if native_path.exists():
         state_dict = torch.load(native_path, map_location="cpu", weights_only=True)
         loaded = 0

miles/backends/megatron_utils/multi_lora_utils.py

@@ -9,16 +9,19 @@
 import torch
 import torch.distributed as dist
 
+from miles.backends.megatron_utils.lora_checkpoint import (
+    adapter_shard_topology,
+    all_megatron_checkpoints_exist,
+    local_shard_name,
+    megatron_shard_name,
+)
 from miles.backends.training_utils.parallel import get_parallel_state
 from miles.ray.multi_lora.controller import get_multi_lora_controller
 from miles.utils.adapter_config import AdapterRun
 from miles.utils.distributed_utils import get_gloo_group
 
 logger = logging.getLogger(__name__)
 
-# Cached by adapter_shard_topology(); the topology is fixed for the run.
-_shard_topology: tuple[bool, tuple[tuple[int, int, int], ...]] | None = None
-
 
 def create_multi_lora_instance(args: Namespace):
     """Create a MultiLoRA instance from training args."""
@@ -48,53 +51,20 @@ def create_multi_lora_instance(args: Namespace):
     )
 
 
-def megatron_shard_name(tp_rank: int, pp_rank: int, ep_rank: int, ep_size: int) -> str:
-    """Adapter shard name for one (tp, pp, ep) coordinate; EP ranks hold different local
-    experts. The ep suffix is omitted at ep_size == 1 so legacy checkpoints stay loadable."""
-    name = f"adapter_megatron_tp{tp_rank}_pp{pp_rank}"
-    if ep_size > 1:
-        name += f"_ep{ep_rank}"
-    return name + ".pt"
-
-
-def adapter_shard_topology() -> tuple[bool, tuple[tuple[int, int, int], ...]]:
-    """Return ``(this_rank_writes_its_shard, realized (tp, pp, ep) coords)`` via one cached gloo all-gather."""
-    global _shard_topology
-    if _shard_topology is not None:
-        return _shard_topology
-    parallel_state = get_parallel_state()
-    coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
-    if not dist.is_initialized():
-        _shard_topology = (True, (coords,))
-        return _shard_topology
-
-    current_rank = dist.get_rank()
-    group = get_gloo_group()
-    gathered: list[object] = [None] * dist.get_world_size(group=group)
-    dist.all_gather_object(gathered, (coords, current_rank), group=group)
-    is_writer = current_rank == min(rank for entry_coords, rank in gathered if entry_coords == coords)
-    _shard_topology = (is_writer, tuple(sorted({entry_coords for entry_coords, _ in gathered})))
-    return _shard_topology
-
-
-def all_megatron_checkpoints_exist(step_dir: Path, shard_names) -> bool:
-    return all((step_dir / name).exists() for name in shard_names)
-
-
 def find_latest_checkpoint(ckpt_dir: Path) -> tuple[Path | None, int]:
     _, coords = adapter_shard_topology()
     if not ckpt_dir.exists():
         return None, 0
 
     parallel_state = get_parallel_state()
     ep_size = parallel_state.ep.size
-    my_coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank)
+    sizes = dict(ep_size=ep_size, etp_size=parallel_state.etp.size, tp_size=parallel_state.tp.size)
 
-    expected = {megatron_shard_name(*coord, ep_size) for coord in coords}
-    my_shard = megatron_shard_name(*my_coords, ep_size)
+    expected = {megatron_shard_name(*coord, **sizes) for coord in coords}
+    my_shard = local_shard_name()
     # Legacy pre-expert-adapter layout: no ep suffix; safe for all EP ranks to read (EP-replicated).
-    legacy = {megatron_shard_name(tp, pp, 0, 1) for tp, pp, _ in coords}
-    my_legacy = megatron_shard_name(my_coords[0], my_coords[1], 0, 1)
+    legacy = {megatron_shard_name(tp, pp) for tp, pp, _, _ in coords}
+    my_legacy = megatron_shard_name(parallel_state.tp.rank, parallel_state.pp.rank)
 
     def get_step(d):
         return int(d.name.split("_")[1])
@@ -192,7 +162,7 @@ def save_multi_lora_checkpoints(
     Layout (per adapter)::
 
         {adapter.save}/checkpoints/step_{iteration}/
-        ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}].pt  ← per-rank shard, fast resume
+        ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}][_etp{etp}].pt  ← per-rank shard, fast resume
         ├── adapter_model.safetensors           ← gathered HF, inference / external
         └── adapter_config.json                 ← HF PEFT metadata (r, alpha, ...)
     """
@@ -207,10 +177,11 @@ def save_multi_lora_checkpoints(
     tp_rank = parallel_state.tp.rank
     pp_rank = parallel_state.pp.rank
     ep_rank = parallel_state.ep.rank
-    ep_size = parallel_state.ep.size
-    # Exactly one writer per (tp, pp, ep) shard; see adapter_shard_topology.
+    # Exactly one writer per (tp, pp, ep, etp) shard; see adapter_shard_topology.
     is_shard_writer, _ = adapter_shard_topology()
-    is_global_writer = is_shard_writer and tp_rank == 0 and pp_rank == 0 and ep_rank == 0
+    # All four coordinates: at etp_size > tp_size several etp ranks share (tp, pp, ep), and two
+    # global writers would race to promote the same directory.
+    is_global_writer = is_shard_writer and (tp_rank, pp_rank, ep_rank, parallel_state.etp.rank) == (0, 0, 0, 0)
 
     target_modules_hf = (
         convert_target_modules_to_hf(list(args.target_modules))
@@ -245,7 +216,7 @@ def save_multi_lora_checkpoints(
                     for name, param in batch.named_parameters()
                     if ".adapter." in name
                 }
-                native_path = tmp_dir / megatron_shard_name(tp_rank, pp_rank, ep_rank, ep_size)
+                native_path = tmp_dir / local_shard_name()
                 torch.save(shard, native_path)
                 logger.info(f"{log_prefix} saved Megatron shard " f"({len(shard)} tensors) to {native_path}")
 
@@ -365,7 +336,6 @@ def _deregister_adapter(adapter: AdapterRun, args, model, optimizer) -> None:
 def load_adapters(args, model, optimizer, adapters) -> int:
     """Load adapters into Megatron slots; resumes step counts from checkpoints."""
     from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
-    from miles.utils.distributed_utils import get_gloo_group
 
     if dist.is_initialized():
         dist.barrier(group=get_gloo_group())
@@ -391,7 +361,6 @@ def load_adapters(args, model, optimizer, adapters) -> int:
 def cleanup_adapters(args, model, optimizer, adapters) -> int:
     """Save final ckpt + clear Megatron slot, then free_slot on the controller."""
     from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
-    from miles.utils.distributed_utils import get_gloo_group
 
     if dist.is_initialized():
         dist.barrier(group=get_gloo_group())
@@ -450,7 +419,6 @@ def save_due_adapter_checkpoints(args, model) -> bool:
     without a checkpoint on disk. Rank 0 picks and broadcasts, so the
     collective export lines up. Returns False when nothing is due."""
     from miles.backends.megatron_utils.initialize import is_first_replica_megatron_main_rank
-    from miles.utils.distributed_utils import get_gloo_group
 
     due_buffer = [None]
     if is_first_replica_megatron_main_rank() and args.save_interval is not None: