diff --git a/miles/backends/megatron_utils/bridge_lora_helpers.py b/miles/backends/megatron_utils/bridge_lora_helpers.py index 9114a14c5b..ae2e6a880e 100644 --- a/miles/backends/megatron_utils/bridge_lora_helpers.py +++ b/miles/backends/megatron_utils/bridge_lora_helpers.py @@ -6,15 +6,18 @@ from __future__ import annotations +import logging from argparse import Namespace from dataclasses import dataclass from megatron.core.utils import get_attr_wrapped_model from miles.utils.hf_config import load_hf_config -from miles.utils.multi_lora import is_multi_lora_enabled +from miles.utils.multi_lora import is_multi_lora_enabled, targets_expert_leaves -from .lora_utils import patch_param_grad_buffer_for_colocate_mode_lora +from .lora_utils import convert_target_modules_to_hf, patch_param_grad_buffer_for_colocate_mode_lora + +logger = logging.getLogger(__name__) @dataclass @@ -66,6 +69,46 @@ def _get_model_config_from_wrapped(model): return get_attr_wrapped_model(model, "config", allow_none=False) +def _validate_multi_lora_moe_support(args: Namespace, provider) -> None: + """Reject MoE configs the multi-slot grouped-expert adapter cannot serve (checked + post-finalize because they depend on the resolved provider, not the CLI).""" + if not getattr(provider, "num_moe_experts", None): + return + if not targets_expert_leaves(args.target_modules): + logger.info("[multilora] MoE model with no expert leaves in --target-modules; experts stay frozen") + return + + # Checked on the provider: --expert-tensor-parallel-size stays None until Megatron resolves it. + expert_tp = getattr(provider, "expert_tensor_parallel_size", 1) or 1 + assert expert_tp == 1, ( + f"Multi-LoRA on MoE experts requires expert_tensor_parallel_size=1 (resolved to " + f"{expert_tp}); set --expert-tensor-parallel-size 1." + ) + assert getattr(provider, "moe_grouped_gemm", False), ( + "Multi-LoRA on MoE experts requires moe_grouped_gemm=True (SequentialMLP expert " + "linears are skipped, so the experts would train no adapter)." + ) + assert not getattr(provider, "fp8", None) and not getattr(provider, "fp4", None), ( + "Multi-LoRA on MoE experts does not support fp8/fp4 experts (quantization padding " + "desynchronizes the dispatched token order)." + ) + # sglang only wraps a fused MoE layer when both expert projections are targeted. + served = set(convert_target_modules_to_hf(list(args.target_modules))) + expert_pair = {"gate_proj", "up_proj", "down_proj"} + if served & expert_pair: + assert expert_pair <= served, ( + f"Multi-LoRA on MoE experts requires all of {sorted(expert_pair)} in " + f"--target-modules (got {sorted(served & expert_pair)}); a one-sided expert " + f"target is dropped at rollout time." + ) + assert not getattr( + provider, "moe_pad_expert_input_to_capacity", False + ), "Multi-LoRA on MoE experts does not support --moe-pad-expert-input-to-capacity." + assert not getattr( + provider, "moe_permute_fusion", False + ), "Multi-LoRA on MoE experts requires moe_permute_fusion=False." + + def _setup_lora_model_via_bridge(args: Namespace) -> list: """Build Megatron model with LoRA using Megatron-Bridge. @@ -104,6 +147,15 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: provider.variable_seq_lengths = True provider.moe_token_dispatcher_type = "alltoall" provider.moe_router_load_balancing_type = "none" + if is_multi_lora_enabled(args) and targets_expert_leaves(args.target_modules): + # Expert adapters cannot replay the fused permute's row_id_map, and most bridge + # MoE providers default the fusion on — so turn it off rather than refuse to build. + if getattr(provider, "moe_permute_fusion", False): + logger.info( + "[multilora] disabling moe_permute_fusion: expert adapters replay the " + "dispatcher's permutation, which the fused kernel does not expose" + ) + provider.moe_permute_fusion = False if getattr(args, "decoder_first_pipeline_num_layers", None) is not None: provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers if getattr(args, "decoder_last_pipeline_num_layers", None) is not None: @@ -113,6 +165,8 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: provider.finalize() if is_multi_lora_enabled(args): + _validate_multi_lora_moe_support(args, provider) + from miles.backends.megatron_utils.multi_lora_utils import create_multi_lora_instance lora = create_multi_lora_instance(args) diff --git a/miles/backends/megatron_utils/multi_lora_utils.py b/miles/backends/megatron_utils/multi_lora_utils.py index 707541fa4e..4d6d768909 100644 --- a/miles/backends/megatron_utils/multi_lora_utils.py +++ b/miles/backends/megatron_utils/multi_lora_utils.py @@ -12,9 +12,13 @@ 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.""" @@ -32,6 +36,7 @@ def create_multi_lora_instance(args: Namespace): lora_cls = LoRA + # exclude_modules was already folded into target_modules during arg validation. return MultiLoRA( target_modules=convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls), n_adapters=args.multi_lora_n_adapters, @@ -43,21 +48,53 @@ def create_multi_lora_instance(args: Namespace): ) -def all_megatron_checkpoints_exist(step_dir: Path, tp_size, pp_size) -> bool: - return all( - (step_dir / f"adapter_megatron_tp{tp}_pp{pp}.pt").exists() for tp in range(tp_size) for pp in range(pp_size) - ) +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() - tp_size = parallel_state.tp.size - pp_size = parallel_state.pp.size - tp_rank = parallel_state.tp.rank - pp_rank = parallel_state.pp.rank + ep_size = parallel_state.ep.size + my_coords = (parallel_state.tp.rank, parallel_state.pp.rank, parallel_state.ep.rank) + + expected = {megatron_shard_name(*coord, ep_size) for coord in coords} + my_shard = megatron_shard_name(*my_coords, ep_size) + # 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) def get_step(d): return int(d.name.split("_")[1]) @@ -69,8 +106,11 @@ def get_step(d): ) for step_dir in step_dirs: step = get_step(step_dir) - if all_megatron_checkpoints_exist(step_dir, tp_size, pp_size): - return step_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt", step + if all_megatron_checkpoints_exist(step_dir, expected): + return step_dir / my_shard, step + if ep_size > 1 and all_megatron_checkpoints_exist(step_dir, legacy): + logger.info(f"[multilora] resuming from pre-expert-parallel shard layout in {step_dir}") + return step_dir / my_legacy, step return None, 0 @@ -116,20 +156,28 @@ def zero_optimizer_state_for_adapter(optimizer, model, idx: int) -> None: def slice_lora_to_rank(hf_name: str, tensor: torch.Tensor, adapter_rank: int) -> torch.Tensor: - if "lora_A" in hf_name and adapter_rank < tensor.shape[0]: - remainder = tensor[adapter_rank:] - assert remainder.abs().max() == 0, ( - f"lora_A padded dims are non-zero: {hf_name}, " - f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" - ) - return tensor[:adapter_rank] - if "lora_B" in hf_name and adapter_rank < tensor.shape[1]: - remainder = tensor[:, adapter_rank:] - assert remainder.abs().max() == 0, ( - f"lora_B padded dims are non-zero: {hf_name}, " - f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" - ) - return tensor[:, :adapter_rank] + """Trim a max-rank-padded LoRA tensor to ``adapter_rank`` on the rank axis, addressed + from the end so packed grouped-expert exports are not sliced on the expert axis.""" + if "lora_A" in hf_name: + rank_dim = tensor.ndim - 2 + if adapter_rank < tensor.shape[rank_dim]: + remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) + assert remainder.abs().max() == 0, ( + f"lora_A padded dims are non-zero: {hf_name}, " + f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" + ) + return tensor.narrow(rank_dim, 0, adapter_rank) + return tensor + if "lora_B" in hf_name: + rank_dim = tensor.ndim - 1 + if adapter_rank < tensor.shape[rank_dim]: + remainder = tensor.narrow(rank_dim, adapter_rank, tensor.shape[rank_dim] - adapter_rank) + assert remainder.abs().max() == 0, ( + f"lora_B padded dims are non-zero: {hf_name}, " + f"max={remainder.abs().max().item():.6e}, shape={tensor.shape}, rank={adapter_rank}" + ) + return tensor.narrow(rank_dim, 0, adapter_rank) + return tensor return tensor @@ -144,7 +192,7 @@ def save_multi_lora_checkpoints( Layout (per adapter):: {adapter.save}/checkpoints/step_{iteration}/ - ├── adapter_megatron_tp{tp}_pp{pp}.pt ← per-rank shard, fast resume + ├── adapter_megatron_tp{tp}_pp{pp}[_ep{ep}].pt ← per-rank shard, fast resume ├── adapter_model.safetensors ← gathered HF, inference / external └── adapter_config.json ← HF PEFT metadata (r, alpha, ...) """ @@ -158,11 +206,11 @@ def save_multi_lora_checkpoints( parallel_state = get_parallel_state() tp_rank = parallel_state.tp.rank pp_rank = parallel_state.pp.rank - # One writer per (tp, pp) shard: LoRA params are replicated across DP AND - # CP, so gate on the combined dp×cp group. Gating on intra_dp alone left - # every CP rank writing the same shard file and racing the os.replace. - is_dp_cp_rank_0 = parallel_state.intra_dp_cp.rank == 0 - is_global_writer = is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0 + ep_rank = parallel_state.ep.rank + ep_size = parallel_state.ep.size + # Exactly one writer per (tp, pp, ep) 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 target_modules_hf = ( convert_target_modules_to_hf(list(args.target_modules)) @@ -183,21 +231,21 @@ def save_multi_lora_checkpoints( final_dir = config.save / "checkpoints" / f"step_{iteration}" tmp_dir = config.save / "checkpoints" / f"_tmp_step_{iteration}" - if is_dp_cp_rank_0: + if is_shard_writer: tmp_dir.mkdir(parents=True, exist_ok=True) if dist.is_initialized(): dist.barrier() with expose_adapter_slot(model, adapter.slot): # Megatron checkpoints - if is_dp_cp_rank_0: + if is_shard_writer: shard: dict[str, torch.Tensor] = { name: param.data.cpu() for batch in model for name, param in batch.named_parameters() if ".adapter." in name } - native_path = tmp_dir / f"adapter_megatron_tp{tp_rank}_pp{pp_rank}.pt" + native_path = tmp_dir / megatron_shard_name(tp_rank, pp_rank, ep_rank, ep_size) torch.save(shard, native_path) logger.info(f"{log_prefix} saved Megatron shard " f"({len(shard)} tensors) to {native_path}") diff --git a/miles/utils/multi_lora.py b/miles/utils/multi_lora.py index a195f73169..67ff62d15f 100644 --- a/miles/utils/multi_lora.py +++ b/miles/utils/multi_lora.py @@ -46,6 +46,23 @@ def define_new_adapter_metrics(snapshot: dict) -> None: define_step_key_metric_group(prefix=f"{name}/perf", step_key="rollout/step") +# Leaf module names that can live inside MoE experts (they also name the dense MLP +# projections); the bulk aliases expand to them during target-module resolution. +_EXPERT_LEAF_NAMES = frozenset({"linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj"}) +_ALL_MODULE_ALIASES = frozenset({"all", "all-linear", "all_linear"}) + + +def targets_expert_leaves(target_modules: Any) -> bool: + """Whether ``target_modules`` can put adapters on MoE expert linears.""" + if isinstance(target_modules, str): + target_modules = [target_modules] + entries = [str(tm).strip().lower() for tm in (target_modules or [])] + if any(entry in _ALL_MODULE_ALIASES for entry in entries): + return True + # Map each entry (possibly a dotted or wildcard path) to its leaf module name. + return any(entry.split(".")[-1] in _EXPERT_LEAF_NAMES for entry in entries) + + def validate_multi_lora_args(args: Any) -> None: """Set ``args.multi_lora``, then validate and default the multi-LoRA arg surface. Called from ``miles_validate_args``; a no-op for normal runs.""" @@ -68,6 +85,23 @@ def validate_multi_lora_args(args: Any) -> None: assert args.lora_rank > 0, "--lora-rank must be set when --multi-lora-n-adapters > 0" assert args.target_modules is not None, "--target-modules must be set when --multi-lora-n-adapters > 0" assert args.train_backend == "megatron", "Multi-LoRA currently requires --train-backend megatron" + # Adapter routing is only recompute-safe without pipelining; enforce at launch. + assert getattr(args, "pipeline_model_parallel_size", 1) == 1, ( + "Multi-LoRA requires --pipeline-model-parallel-size 1: no single rank holds a " + "complete adapter to push to the rollout engines, and a pipelined schedule would " + "recompute activations against a later micro-batch's adapter routing." + ) + # Per-slot token spans assume sequence-major contiguous sample packing, which only 'thd' provides. + assert getattr(args, "qkv_format", "thd") == "thd", ( + "Multi-LoRA requires --qkv-format thd: per-adapter token spans assume the " + f"micro-batch packs samples contiguously, which bshd does not (got {args.qkv_format!r})." + ) + assert not getattr(args, "experts_shared_outer_loras", False), ( + "Multi-LoRA does not support --experts-shared-outer-loras; MoE expert adapters " + "use the per-expert layout. Drop the flag (and --sglang-experts-shared-outer-loras)." + ) + # Expert-parallel sizes are checked post-finalize in _validate_multi_lora_moe_support: + # --expert-tensor-parallel-size stays None until Megatron's own validate_args resolves it. assert "muon" not in str(getattr(args, "optimizer", "")).lower(), ( "Multi-LoRA does not support Muon: per-adapter decoupled stepping is only " "implemented for Adam-family per-slot optimizers" diff --git a/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py new file mode 100644 index 0000000000..97ab9f3c5e --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_multi_lora_checkpoint_naming.py @@ -0,0 +1,46 @@ +"""Adapter shards are keyed by (tp, pp, ep): EP ranks hold different local experts, and +the realized coordinates are not the tp x pp x ep cross product when ETP < TP.""" + +from miles.backends.megatron_utils.multi_lora_utils import all_megatron_checkpoints_exist, megatron_shard_name + + +def _names(coords, ep_size): + return {megatron_shard_name(*coord, ep_size) for coord in coords} + + +def test_shard_name_omits_ep_suffix_without_expert_parallelism(): + # Checkpoints written before expert adapters existed must stay loadable. + assert megatron_shard_name(0, 0, 0, ep_size=1) == "adapter_megatron_tp0_pp0.pt" + assert megatron_shard_name(1, 2, 0, ep_size=1) == "adapter_megatron_tp1_pp2.pt" + + +def test_shard_name_is_unique_per_expert_parallel_rank(): + names = {megatron_shard_name(0, 0, ep, ep_size=4) for ep in range(4)} + assert len(names) == 4 + assert megatron_shard_name(0, 0, 2, ep_size=4) == "adapter_megatron_tp0_pp0_ep2.pt" + + +def test_completeness_check_requires_every_realized_shard(tmp_path): + coords = [(0, 0, 0), (0, 0, 1), (0, 0, 2)] + for coord in coords[:2]: + (tmp_path / megatron_shard_name(*coord, 3)).touch() + + assert not all_megatron_checkpoints_exist(tmp_path, _names(coords, 3)) + + (tmp_path / megatron_shard_name(*coords[2], 3)).touch() + assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 3)) + + +def test_completeness_ignores_unrealized_coordinates(tmp_path): + # TP=2, EP=2, ETP=1: only (0,0,0) and (1,0,1) exist; a cross-product check + # would demand four shards and never resume. + coords = [(0, 0, 0), (1, 0, 1)] + for coord in coords: + (tmp_path / megatron_shard_name(*coord, 2)).touch() + + assert all_megatron_checkpoints_exist(tmp_path, _names(coords, 2)) + + +def test_completeness_check_with_single_shard(tmp_path): + (tmp_path / "adapter_megatron_tp0_pp0.pt").touch() + assert all_megatron_checkpoints_exist(tmp_path, _names([(0, 0, 0)], 1)) diff --git a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py index 6b22665766..2df14f54b1 100644 --- a/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py +++ b/tests/fast/backends/megatron_utils/test_slice_lora_to_rank.py @@ -46,3 +46,38 @@ def test_full_rank_tensor_passes_through(): def test_non_lora_names_pass_through(): tensor = torch.ones(32, 8) assert slice_lora_to_rank("x.some_other.weight", tensor, 16) is tensor + + +# Packed grouped-expert exports ([E, rank, in] / [E, out, rank]) must be sliced on the +# rank axis — dim 0/1 slicing would drop experts while reporting the requested shape. + + +def test_packed_expert_lora_a_is_sliced_on_the_rank_dim(): + # 4 experts, max rank 32, in_features 8; only the first 16 ranks are live. + tensor = torch.zeros(4, 32, 8) + tensor[:, :16] = 1.0 + out = slice_lora_to_rank("base_model.experts.gate_proj.lora_A.weight", tensor, 16) + assert out.shape == (4, 16, 8) + assert torch.equal(out, tensor[:, :16]) + + +def test_packed_expert_lora_b_is_sliced_on_the_rank_dim(): + tensor = torch.zeros(4, 8, 32) + tensor[..., :16] = 1.0 + out = slice_lora_to_rank("base_model.experts.gate_proj.lora_B.weight", tensor, 16) + assert out.shape == (4, 8, 16) + assert torch.equal(out, tensor[..., :16]) + + +def test_packed_expert_nonzero_padding_is_rejected(): + tensor = torch.ones(4, 32, 8) + with pytest.raises(AssertionError, match="padded dims are non-zero"): + slice_lora_to_rank("x.experts.gate_proj.lora_A.weight", tensor, 16) + + +def test_packed_expert_fewer_experts_than_rank_is_not_confused(): + # 2 experts, max rank 32, sliced to 4: dim-0 slicing would silently pass the shape check. + tensor = torch.zeros(2, 32, 8) + tensor[:, :4] = 1.0 + out = slice_lora_to_rank("x.experts.gate_proj.lora_A.weight", tensor, 4) + assert out.shape == (2, 4, 8) diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 4f84cb5dd4..bf6e710076 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -253,6 +253,37 @@ def test_rejects_experimental_ft_trainer(self, monkeypatch): with pytest.raises(AssertionError, match="MILES_EXPERIMENTAL_FT_TRAINER"): miles_validate_args(args) + def test_rejects_pipeline_parallelism(self): + # Adapter routing is not recompute-safe under a pipelined schedule. + args = self._parse([]) + args.pipeline_model_parallel_size = 2 + with pytest.raises(AssertionError, match="pipeline-model-parallel-size 1"): + miles_validate_args(args) + + def test_rejects_bshd_qkv_format(self): + # bshd interleaves samples in the sequence-major flattening the spans assume. + args = self._parse([]) + args.qkv_format = "bshd" + with pytest.raises(AssertionError, match="qkv-format thd"): + miles_validate_args(args) + + def test_rejects_shared_outer_expert_loras(self): + # Per-expert layout only; the flag would switch sglang to a layout training never produces. + args = self._parse([]) + args.experts_shared_outer_loras = True + with pytest.raises(AssertionError, match="experts-shared-outer-loras"): + miles_validate_args(args) + + def test_accepts_expert_leaf_targets_without_expert_tp_flag(self): + # --expert-tensor-parallel-size stays None until Megatron's own validate_args; + # comparing the raw value here rejected every run that omitted the flag. + args = self._parse(["--target-modules", "gate_proj,up_proj,down_proj"]) + args.expert_tensor_parallel_size = None + + miles_validate_args(args) + + assert args.multi_lora is True + class TestResolveFtComponents: def test_disabled_with_no_components_returns_empty_without_warning(self, caplog) -> None: diff --git a/tests/fast/utils/test_targets_expert_leaves.py b/tests/fast/utils/test_targets_expert_leaves.py new file mode 100644 index 0000000000..9b2844a0e7 --- /dev/null +++ b/tests/fast/utils/test_targets_expert_leaves.py @@ -0,0 +1,36 @@ +"""targets_expert_leaves gates the MoE-specific multi-LoRA handling (permute-fusion +off, expert validations); a false negative silently skips those.""" + +from miles.utils.multi_lora import targets_expert_leaves + + +def test_mlp_leaf_names_target_experts(): + # These names match the dense MLP and the routed experts alike. + assert targets_expert_leaves(["gate_proj", "up_proj", "down_proj"]) + assert targets_expert_leaves(["linear_fc1"]) + assert targets_expert_leaves(["linear_fc2"]) + + +def test_expert_scoped_wildcards_target_experts(): + assert targets_expert_leaves(["*.layers.*.mlp.experts.linear_fc1"]) + + +def test_attention_only_targets_do_not(): + assert not targets_expert_leaves(["linear_qkv", "linear_proj"]) + assert not targets_expert_leaves(["q_proj", "k_proj", "v_proj", "o_proj"]) + + +def test_bulk_aliases_target_experts(): + # "all" is only resolved by the later target-module conversion, so the alias itself counts. + for alias in ("all", "all-linear", "all_linear", "ALL"): + assert targets_expert_leaves([alias]), alias + + +def test_bare_string_is_accepted(): + assert targets_expert_leaves("gate_proj") + assert not targets_expert_leaves("linear_qkv") + + +def test_empty_targets_do_not(): + assert not targets_expert_leaves(None) + assert not targets_expert_leaves([])