diff --git a/docs/advanced/lora.md b/docs/advanced/lora.md index 38484d31c0b..a7593e73dff 100644 --- a/docs/advanced/lora.md +++ b/docs/advanced/lora.md @@ -102,12 +102,12 @@ in PR #1792 are not released on `main` yet. FSDP does not currently implement LoRA training. -`all-linear` expands to Q/K/V/O and gate/up/down projections, and conditionally -adds MLA Q/KV projections based on the HF config. It does not literally wrap -every linear layer. GDN and other model-specific projections require an explicit -target list. Current GLM recipes validate models that contain DSA while leaving -the DSA indexer unadapted; current hybrid-model recipes also leave MTP blocks and -vision towers unadapted. Use the model launcher as the source of truth. +Omitting `--target-modules` or passing `all-linear` uses the model defaults from +`miles/utils/lora/hf_lora_targets.py`: attention + MLP, with model-specific exclusions +and output-head defaults. Multi-LoRA without explicit targets selects all three +training groups. Use `--target-modules attn,mlp,unembed` to select groups explicitly. +Ordinary LoRA also accepts specific HF targets mixed with group names, such as +`--target-modules attn,lm_head`; Tinker accepts only group names. ### Core arguments @@ -117,8 +117,8 @@ vision towers unadapted. Use the model launcher as the source of truth. | `--lora-alpha` | `16` | Adapter scaling factor. | | `--lora-dropout` | `0.0` | Dropout on the adapter path. | | `--lora-type` | `lora` | `lora` uses fused Megatron projections; `canonical_lora` uses split Q/K/V and gate/up projections. The canonical path is implemented and covered by fast name-mapping tests, but has no maintained recipe or E2E validation. | -| `--target-modules` | none | Required with a positive rank. Accepts `all-linear`, HF leaf names, Megatron names, or model-specific wildcard paths. | -| `--exclude-modules` | none | Comma-separated exact entries removed from the resolved targets. | +| `--target-modules` | none | Uses HF model defaults when omitted. Accepts `all-linear`, `attn/mlp/unembed` groups, HF leaf names or scoped HF paths; Bridge also accepts Megatron selectors. | +| `--exclude-modules` | none | Comma-separated HF leaf names or scoped HF paths removed after selection; Bridge also accepts Megatron selectors. | | `--lora-adapter-path` | none | Warm-start/resume path. Also provide the matching positive rank, alpha, and target modules. Bridge training resume currently requires miles' per-rank adapter shards and the same parallel topology; an HF PEFT-only adapter cannot yet be loaded directly into the Bridge model. Inkling native has its own HF adapter loader. | | `--lora-base-cpu-backup` | off | Colocated mode only: keep a CPU mirror of the frozen SGLang base and avoid re-sending base weights. This trades host RAM for faster and more reliable pause/resume. | | `--lora-train-only` | off | Train the adapter while keeping ordinary rollout engines on the frozen base policy. | @@ -126,10 +126,88 @@ vision towers unadapted. Use the model launcher as the source of truth. | `--check-lora-weight-equal` | off | On the colocated path, verify each synchronized adapter tensor with SHA-256. | | `--update-weights-interval` | `1` | Publish new weights every N rollout/train iterations. This is not LoRA-specific, but it controls when the live adapter is synchronized. | -This argument table describes the general Bridge surface. Current native Inkling -uses a fixed model-specific adapter schema: `--target-modules` does not select -individual training modules, `--exclude-modules` is not applied, and -`canonical_lora` is not implemented. Use the Inkling launcher defaults. +### HF target source of truth + +`miles/utils/lora/hf_lora_targets.py` owns **HF target groups and defaults** for all +backends. It derives attention, MLP, and output-head paths from the model config, +including nested text models, optional MLA projections, expert layouts, and +hybrid attention. Vision towers, routers, norms, and GDN convolutions are excluded. + +| Model layout | Attention | MLP | Ordinary LoRA default | +|---|---|---|---| +| Llama, Qwen2/2.5, Qwen3 | Q/K/V/O | Dense gate/up/down | Attention + MLP | +| Qwen3 MoE | Q/K/V/O | Routed experts; dense layers when configured | Attention + MLP | +| Qwen3-Next | Q/K/V/O + GDN qkvz/ba/out | Routed + shared experts; configured dense layers | Attention + MLP | +| Qwen3.5/3.6, text and multimodal | Q/K/V/O + GDN qkv/z/b/a/out | Dense or packed routed + shared experts | Attention + MLP | +| GPT-OSS | Q/K/V/O | Packed gate_up/down experts | Attention + MLP | +| DeepSeek V2/V3, Kimi K2/K2.5 | MLA | Dense + routed + shared experts as configured | Attention + MLP | +| DeepSeek V3.2, GLM-5/5.1/5.2 | MLA + DSA indexer | Dense + routed + shared experts as configured | Attention + MLP, excluding indexer | +| GLM-4 MoE | Q/K/V/O | Dense + routed + shared experts as configured | Attention + MLP | +| Inkling | Q/K/V/R/O | Dense + packed routed + shared expert projections | Attention + MLP + output head | + +`--target-modules` accepts comma-separated `attn`, `mlp`, and `unembed` groups, +explicit HF targets, or a mix of both. Group names expand through the model's HF +definition, so `mlp` includes dense, shared, and routed expert projections, including +packed `gate_up_proj` weights. Only exact group names expand; full paths remain +literal target patterns. Duplicates are removed, then `--exclude-modules` applies. +Explicit module selections must not overlap exclusions; defaults and group +selections may still be narrowed with exclusions. + +A matching pair of `gate_proj` and `up_proj` selectors also selects +`gate_up_proj` when the HF model has packed projections, with a warning. +Scoped pairs retain their path prefix; selecting only gate or only up does not +expand to the packed projection. Split selectors remain where the model has +split modules and are replaced where it has only packed modules. + +Without explicit targets, ordinary LoRA uses its model defaults and multi-LoRA +selects all three groups. `all-linear` selects the ordinary model defaults and +must be used alone. Explicit lists replace the defaults. + +Packed expert entries identify HF parameters rather than `nn.Linear` modules. +Inkling entries follow the native HF model namespace. Its native Megatron LoRA +implementation requires the complete fixed layout. SGLang discovers its adapter +modules with `all`, and online weight-sync config uses `all-linear`; adapter +checkpoints derive explicit names from exported tensors. Adapter factors, tensor +packing, and checkpoint names are unchanged. The pinned Transformers version does +not yet include native Inkling, so its HF structure is not covered by the native +meta-model tests. A layout entry is not a backend support claim. + +Ordinary LoRA and Tinker share these HF target groups. HF targets retain their +meaning throughout training and serving. `miles/utils/hf_utils/weight_mapping.py` +uses Transformers conversion rules and a meta model's parameter names to relate +checkpoint keys to the current HF model namespace, without loading base weights. +This resolves renamed and packed parameter names. Custom models absent from +native Transformers keep their existing checkpoint namespace. + +Bridge resolves mappings against parameters that actually exist across PP/EP +ranks before checking fused selections. Explicit Megatron selectors retain a +compatibility conversion at startup. Adapter factories receive only Megatron +targets; they do not overwrite the HF selection. Arbitrary layer/expert subsets +within a registry template are not implemented and are rejected. +Standard LoRA requires all projections of a fused weight together; +`canonical_lora` supports individual Q/K/V and dense gate/up selections. +Grouped-expert FC1 and GDN input projections remain fused and require all of +their HF projections even in canonical mode. + +Tinker uses `--target-modules attn,mlp,unembed` by default; for example, +`--target-modules attn,mlp` disables output-head training. Client SDK flags must +match the selected server groups. Tinker rejects explicit module names, +`all-linear`, and `--exclude-modules` because the SDK describes only whole groups. + +For Bridge, SGLang receives the selected HF paths and normalizes them into buffer types +(for example, Q/K/V become `qkv_proj`); it does not own the selection policy. +Adapter checkpoints derive their concrete `target_modules` from exported tensor +keys rather than a separate Megatron-to-HF name table. Online adapter registration +uses the resolved HF selection for Bridge and `all-linear` for native Inkling, +without an extra tensor export. +This requires [SGLang's HF-path normalization support](https://github.com/sgl-project/sglang/pull/40242), +including GDN split names. FSDP LoRA injection remains +unsupported. + +Native Inkling +uses a fixed model-specific adapter schema: defaults select that complete schema, +and partial/custom layouts or `canonical_lora` are rejected before injection. +Use the Inkling launcher defaults. ### Rollout topology @@ -160,7 +238,7 @@ LORA_ARGS=( --lora-rank 32 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "gate_proj,up_proj,down_proj" + --target-modules "gate_up_proj,down_proj" --sglang-lora-backend triton --megatron-to-hf-mode bridge ) diff --git a/docs/examples/multi-lora.md b/docs/examples/multi-lora.md index 44adaa07158..5b6c8017deb 100644 --- a/docs/examples/multi-lora.md +++ b/docs/examples/multi-lora.md @@ -18,11 +18,10 @@ One 8-GPU node, disaggregated (multi-LoRA forbids `--colocate`): (`linear_qkv`, `linear_proj`), the per-expert MoE projections (`linear_fc1`, `linear_fc2`), and the output layer (`output_layer`) so the cookbook's default `train_unembed=True` is servable. -The gateway currently resolves Tinker training groups for `qwen3` and `qwen3_moe`. -`--tinker-train-attn`, `--tinker-train-mlp`, and `--tinker-train-unembed` default to enabled; -use `--no-tinker-train-attn`, `--no-tinker-train-mlp`, or `--no-tinker-train-unembed` to disable a group. -Every client's corresponding SDK flags must match the server layout. Tinker startup rejects -`--target-modules` and `--exclude-modules`; native Miles training still accepts them. +The example enables attention, MLP, and output-head training. Client SDK flags +must match the server's selected groups. See [LoRA target selection](/advanced/lora#hf-target-source-of-truth) +for `--target-modules attn,mlp,unembed`; use `attn,mlp` to disable output-head +training. Tinker accepts only these group names and rejects `--exclude-modules`. ## Run diff --git a/docs/models/kimi/kimi-k3.md b/docs/models/kimi/kimi-k3.md index b42a3b02e49..f3da932175d 100644 --- a/docs/models/kimi/kimi-k3.md +++ b/docs/models/kimi/kimi-k3.md @@ -96,6 +96,8 @@ torchrun --nnodes=8 --nproc-per-node=4 ... \ ## 4. Launch +The launcher uses `all-linear`, which selects Kimi K3's HF defaults: MLA query/KV down projections, attention output projections, dense/shared MLPs, and routed experts. `attn,mlp` selects the same set. Explicit HF targets can omit `language_model.model.layers.*.block_sparse_moe.experts.*.w2` to leave routed-expert down projections frozen; other layouts remain unsupported by the native backend. + `--train-mode lora` (default) or `full`. Validated on **16 nodes × 4 GPUs**: one container per node, a ray cluster across them, `export MILES_SCRIPT_EXTERNAL_RAY=1`, then: diff --git a/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh b/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh index 098a0e76865..402e50f7d64 100644 --- a/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh +++ b/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh @@ -29,7 +29,7 @@ LORA_ARGS=( --lora-rank 32 # LoRA rank (typical values: 8, 16, 32, 64) --lora-alpha 32 # LoRA alpha (usually 2x rank) --lora-dropout 0.0 # LoRA dropout (0.0 for RL training) - --target-modules "gate_proj,up_proj,down_proj" + --target-modules "gate_up_proj,down_proj" --sglang-lora-backend triton # !!! must for moe-lora !!!, else display "Current LoRA backend does not support LoRA on MoE layers; skipping MoE layer" ) diff --git a/examples/multi_lora/README.md b/examples/multi_lora/README.md index 9aa3d2f1e21..8fecd60892e 100644 --- a/examples/multi_lora/README.md +++ b/examples/multi_lora/README.md @@ -15,11 +15,10 @@ One 8-GPU node, disaggregated (multi-LoRA forbids `--colocate`): (`linear_qkv`, `linear_proj`), the per-expert MoE projections (`linear_fc1`, `linear_fc2`), and the output layer (`output_layer`) so the cookbook's default `train_unembed=True` is servable. -The gateway currently resolves Tinker training groups for `qwen3` and `qwen3_moe`. -`--tinker-train-attn`, `--tinker-train-mlp`, and `--tinker-train-unembed` default to enabled; -use `--no-tinker-train-attn`, `--no-tinker-train-mlp`, or `--no-tinker-train-unembed` to disable a group. -Every client's corresponding SDK flags must match the server layout. Tinker startup rejects -`--target-modules` and `--exclude-modules`; native Miles training still accepts them. +The example enables attention, MLP, and output-head training. Client SDK flags +must match the server's selected groups. See [LoRA target selection](../../docs/advanced/lora.md#hf-target-source-of-truth) +for `--target-modules attn,mlp,unembed`; use `attn,mlp` to disable output-head +training. Tinker accepts only these group names and rejects `--exclude-modules`. ## Run diff --git a/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py b/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py index 1b7352dc44a..e1f62a8ba6c 100644 --- a/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py +++ b/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py @@ -29,9 +29,7 @@ class ScriptArgs(U.ExecuteTrainConfig): lora_rank: int = 32 lora_alpha: int = 64 n_adapters: int = 4 - tinker_train_attn: bool = True - tinker_train_mlp: bool = True - tinker_train_unembed: bool = True + target_modules: str = "attn,mlp,unembed" tinker_port: int = 10613 rollout_num_gpus_per_engine: int = 2 @@ -69,14 +67,11 @@ def serve(args: ScriptArgs): f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout 0.0 " "--no-gradient-accumulation-fusion " f"--multi-lora-n-adapters {args.n_adapters} " + f'--target-modules "{args.target_modules}"' ) tinker_args = f"--tinker-server-port {args.tinker_port} " f"--tinker-checkpoint-root {args.save_dir}/{args.run_id}" - for group in ("attn", "mlp", "unembed"): - enabled = getattr(args, f"tinker_train_{group}") - tinker_args += f" --{'' if enabled else 'no-'}tinker-train-{group}" - # initial config only; AdamParams come per optim_step request optimizer_args = "--optimizer adam --lr 1e-4 " diff --git a/examples/swe-agent-harbor-docker/run_glm52_lora_tb2_daytona.py b/examples/swe-agent-harbor-docker/run_glm52_lora_tb2_daytona.py index 470767005d4..9c7ee9c43d3 100644 --- a/examples/swe-agent-harbor-docker/run_glm52_lora_tb2_daytona.py +++ b/examples/swe-agent-harbor-docker/run_glm52_lora_tb2_daytona.py @@ -49,7 +49,7 @@ # tilelang the indexer adapter gets no gradient at all — and EXCLUDING the MLP/MoE # leaves (gate_proj/up_proj/down_proj): this is the MoE-LoRA-off ablation, so all # experts (routed + shared) and the dense MLP stay frozen. -_DEFAULT_TARGET_MODULES = "q_proj,k_proj,v_proj,o_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" +_DEFAULT_TARGET_MODULES = "o_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" @dataclass diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 4ff48a55c7d..3a1e4d7c627 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -13,7 +13,7 @@ from torch_memory_saver import torch_memory_saver from miles.backends.megatron_utils.ft.types import TrainStepOutput -from miles.backends.megatron_utils.lora.utils import build_lora_sync_config, is_lora_enabled, lora_rollout_enabled +from miles.backends.megatron_utils.lora.utils import is_lora_enabled, lora_rollout_enabled from miles.backends.megatron_utils.rematerialize_utils import build_main_cast_context from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator from miles.backends.training_utils.weight_update.updater import WeightUpdater @@ -28,6 +28,7 @@ from miles.utils.distributed_utils import get_gloo_group from miles.utils.ft_utils.indep_dp import IndepDPInfo from miles.utils.hf_utils.config import load_hf_config +from miles.utils.lora.utils import build_lora_config from miles.utils.memory_utils import clear_memory, print_memory from miles.utils.processing_utils import load_tokenizer from miles.utils.ray_utils import Box @@ -287,7 +288,7 @@ def _init_training_state(self) -> None: iterator_factory=get_hf_weight_iterator, parallel_state=get_parallel_state(), is_lora=is_lora, - lora_sync_config=build_lora_sync_config(args) if is_lora else None, + lora_sync_config=build_lora_config(args, target_modules=args.lora_adapter_targets) if is_lora else None, ) def _clear_quantized_weight_workspaces(self) -> None: diff --git a/miles/backends/megatron_utils/lora/actor.py b/miles/backends/megatron_utils/lora/actor.py index 8e16f702d8c..6b84bebe7cd 100644 --- a/miles/backends/megatron_utils/lora/actor.py +++ b/miles/backends/megatron_utils/lora/actor.py @@ -4,12 +4,11 @@ from miles.backends.megatron_utils.lora import checkpoint as lora_checkpoint from miles.backends.megatron_utils.lora import model as lora_model from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer -from miles.backends.megatron_utils.lora.utils import build_lora_sync_config from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator from miles.backends.training_utils.data import get_rollout_data from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement from miles.backends.training_utils.weight_update.snapshot_publisher import WeightPublisher -from miles.utils.lora.utils import AdapterSpec +from miles.utils.lora.utils import AdapterSpec, build_lora_config from miles.utils.ray_utils import Box from miles.utils.tracking_utils.structured_log import with_logs @@ -25,7 +24,9 @@ def _init_training_state(self) -> None: model_name=type(self.hf_config).__name__.lower() if args.model_name is None else args.model_name, quantization_config=getattr(self.hf_config, "quantization_config", None), ) - self.weight_publisher = WeightPublisher(iterator, build_lora_sync_config(args)) + self.weight_publisher = WeightPublisher( + iterator, build_lora_config(args, target_modules=args.lora_adapter_targets) + ) @with_logs def forward_backward(self, batch_id: int, rollout_data_ref: Box) -> dict: diff --git a/miles/backends/megatron_utils/lora/bridge.py b/miles/backends/megatron_utils/lora/bridge.py index f05aafc5ca2..2a4db260289 100644 --- a/miles/backends/megatron_utils/lora/bridge.py +++ b/miles/backends/megatron_utils/lora/bridge.py @@ -10,14 +10,21 @@ from argparse import Namespace from dataclasses import dataclass +import torch.distributed as dist +from megatron.core.tensor_parallel import ColumnParallelLinear from megatron.core.utils import get_attr_wrapped_model +from miles.backends.megatron_utils.lora.slots import create_multi_lora_instance +from miles.backends.megatron_utils.lora.target_modules import resolve_megatron_lora_targets +from miles.backends.megatron_utils.lora.utils import ( + create_lora_instance, + patch_param_grad_buffer_for_colocate_mode_lora, +) from miles.utils.hf_utils.config import load_hf_config +from miles.utils.hf_utils.weight_mapping import HfWeightMapping from miles.utils.lora.utils import is_multi_lora_enabled, targets_expert_leaves from miles.utils.megatron_bridge_utils import apply_dsa_backend_args -from .utils import convert_target_modules_to_hf, patch_param_grad_buffer_for_colocate_mode_lora - logger = logging.getLogger(__name__) @@ -75,7 +82,7 @@ def _validate_multi_lora_moe_support(args: Namespace, provider) -> None: 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): + if not targets_expert_leaves(args.hf_lora_targets): logger.info("[multilora] MoE model with no expert leaves in --target-modules; experts stay frozen") return @@ -94,7 +101,9 @@ def _validate_multi_lora_moe_support(args: Namespace, provider) -> None: "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))) + served = {target.rsplit(".", 1)[-1] for target in args.hf_lora_targets} + if "gate_up_proj" in served: + served.update(("gate_proj", "up_proj")) expert_pair = {"gate_proj", "up_proj", "down_proj"} if served & expert_pair: assert expert_pair <= served, ( @@ -126,6 +135,7 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: List of DDP-wrapped model chunks with LoRA applied. """ from megatron.bridge import AutoBridge + from megatron.bridge.models.conversion.model_bridge import _megatron_local_name_to_global from megatron.bridge.training.config import DistributedDataParallelConfig hf_config = load_hf_config(args.hf_checkpoint) @@ -149,7 +159,7 @@ 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): + if is_multi_lora_enabled(args) and targets_expert_leaves(args.hf_lora_targets): # 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): @@ -168,15 +178,35 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list: if is_multi_lora_enabled(args): _validate_multi_lora_moe_support(args, provider) - from miles.backends.megatron_utils.lora.slots import create_multi_lora_instance - - lora = create_multi_lora_instance(args) - else: - from .utils import create_lora_instance - - lora = create_lora_instance(args) + create_adapter = create_multi_lora_instance if is_multi_lora_enabled(args) else create_lora_instance + assert not ( + is_multi_lora_enabled(args) and args.lora_type == "canonical_lora" + ), "MultiLoRA requires --lora-type lora; it does not implement canonical split adapters" + model_bridge = bridge._model_bridge + model_bridge.hf_pretrained = bridge.hf_pretrained + hf_mapping = HfWeightMapping.from_config(hf_config) def apply_lora_hook(model_chunks): + parameter_names = set(model_bridge._megatron_global_param_names_all_pp_ranks(model_chunks)) + # Tied output layers own no weight but can still carry an independent adapter. + for vp_stage, chunk in enumerate(model_chunks): + for name, module in chunk.named_modules(): + if isinstance(module, ColumnParallelLinear) and module.weight is None: + parameter_names.add( + _megatron_local_name_to_global(model_chunks, chunk.config, f"{name}.weight", vp_stage) + ) + # Include tied layers and experts owned by other PP/EP ranks without changing Bridge's weight inventory. + names_by_rank = [None] * dist.get_world_size() + dist.all_gather_object(names_by_rank, parameter_names) + parameter_names = set().union(*names_by_rank) + adapter_targets = resolve_megatron_lora_targets( + args.hf_lora_targets, + model_bridge.mapping_registry().get_all_mappings(), + parameter_names=parameter_names, + hf_mapping=hf_mapping, + canonical=args.lora_type == "canonical_lora", + ) + lora = create_adapter(args, target_modules=adapter_targets) transformed = lora(model_chunks, training=True) lora.set_params_to_save(transformed) return transformed diff --git a/miles/backends/megatron_utils/lora/slots.py b/miles/backends/megatron_utils/lora/slots.py index 4f51795b4ef..cce9b8af952 100644 --- a/miles/backends/megatron_utils/lora/slots.py +++ b/miles/backends/megatron_utils/lora/slots.py @@ -7,29 +7,14 @@ import torch - logger = logging.getLogger(__name__) -def create_multi_lora_instance(args: Namespace): - """Create a MultiLoRA instance from training args.""" +def create_multi_lora_instance(args: Namespace, *, target_modules): from megatron.bridge.peft.multi_lora import MultiLoRA - from miles.backends.megatron_utils.lora.utils import convert_target_modules_to_megatron - - lora_type_name = getattr(args, "lora_type", "lora").lower() - if lora_type_name == "canonical_lora": - from megatron.bridge.peft.canonical_lora import CanonicalLoRA - - lora_cls = CanonicalLoRA - else: - from megatron.bridge.peft.lora import LoRA - - 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), + target_modules=target_modules, n_adapters=args.multi_lora_n_adapters, dim=args.lora_rank, alpha=args.lora_alpha, diff --git a/miles/backends/megatron_utils/lora/target_modules.py b/miles/backends/megatron_utils/lora/target_modules.py new file mode 100644 index 00000000000..6dfc3451628 --- /dev/null +++ b/miles/backends/megatron_utils/lora/target_modules.py @@ -0,0 +1,149 @@ +import re + +from miles.utils.hf_utils.weight_mapping import HfWeightMapping +from miles.utils.lora.utils import matches_lora_target + +_CANONICAL_PROJECTIONS = { + "q_proj": "linear_q", + "k_proj": "linear_k", + "v_proj": "linear_v", + "gate_proj": "linear_fc1_gate", + "up_proj": "linear_fc1_up", +} + + +def _canonical_adapter_module(module, checkpoint_parameter): + leaf = checkpoint_parameter.removesuffix(".weight").rsplit(".", 1)[-1] + assert leaf in _CANONICAL_PROJECTIONS, f"CanonicalLoRA has no split adapter for {checkpoint_parameter!r}" + return f"{module.rsplit('.', 1)[0]}.{_CANONICAL_PROJECTIONS[leaf]}" + + +def _checkpoint_parameters(mapping): + return {mapping.hf_param} if isinstance(mapping.hf_param, str) else set(mapping.hf_param.values()) + + +def _select_checkpoint_parameters(checkpoint_parameters, targets, *, hf_mapping, megatron_module, exclude_modules): + hf_parameters = {name: hf_mapping.model_parameter(name) for name in checkpoint_parameters} + # Bridge can expose auxiliary layers, such as MTP, absent from the HF model. + selected = { + checkpoint_parameter + for checkpoint_parameter, hf_parameter in hf_parameters.items() + if (not hf_mapping.parameter_names or hf_parameter in hf_mapping.parameter_names) + and any(matches_lora_target(hf_parameter.removesuffix(".weight"), target) for target in targets) + } + return { + checkpoint_parameter + for checkpoint_parameter in selected + if not any( + matches_lora_target(hf_parameters[checkpoint_parameter].removesuffix(".weight"), target) + or matches_lora_target(megatron_module, target) + for target in exclude_modules + ) + } + + +def _resolve_adapter_targets(megatron_module, checkpoint_parameters, selected, *, canonical): + if ( + canonical + and len(checkpoint_parameters) > 1 + and ".experts." not in megatron_module + and megatron_module.rsplit(".", 1)[-1] in ("linear_qkv", "linear_fc1") + ): + return list(dict.fromkeys(_canonical_adapter_module(megatron_module, name) for name in selected)) + assert ( + selected == checkpoint_parameters + ), f"LoRA on fused module {megatron_module!r} requires all HF targets with this adapter" + if canonical and ".experts." in megatron_module and megatron_module.endswith(".linear_fc1"): + # CanonicalLoRA requires split aliases even for fused expert adapters. + return [f"{megatron_module}_gate", f"{megatron_module}_up"] + return [megatron_module] + + +def resolve_megatron_lora_targets(targets, mappings, *, parameter_names, hf_mapping, canonical, exclude_modules=()): + adapter_targets = {} + matched_megatron_parameters = set() + for mapping in mappings: + megatron_module, weight = mapping.megatron_param.rsplit(".", 1) + if weight not in ("weight", "weight*"): + continue + regex = re.compile(re.escape(mapping.megatron_param).replace(r"\*", "(.*)")) + selections = [] + for megatron_parameter in sorted(parameter_names - matched_megatron_parameters): + match = regex.fullmatch(megatron_parameter) + if match is None: + continue + matched_megatron_parameters.add(megatron_parameter) + checkpoint_parameters = _checkpoint_parameters(mapping.resolve(match.groups())) + selected = _select_checkpoint_parameters( + checkpoint_parameters, + targets, + hf_mapping=hf_mapping, + megatron_module=megatron_module, + exclude_modules=exclude_modules, + ) + selections.append((checkpoint_parameters, selected)) + if not any(selected for _, selected in selections): + continue + # Template injection must select the same projections for every layer/expert it matches. + assert all( + selected for _, selected in selections + ), f"LoRA cannot select a subset of parameters in {mapping.megatron_param!r}" + adapter_selections = set() + for checkpoint_parameters, selected in selections: + selected_adapters = _resolve_adapter_targets( + megatron_module, checkpoint_parameters, selected, canonical=canonical + ) + adapter_selections.add(frozenset(selected_adapters)) + adapter_targets.update(dict.fromkeys(selected_adapters)) + assert ( + len(adapter_selections) == 1 + ), f"LoRA cannot select different projections across parameters in {mapping.megatron_param!r}" + assert adapter_targets, "LoRA targets have no Megatron modules" + return list(adapter_targets) + + +def normalize_lora_targets_to_hf(hf_checkpoint, target_modules, *, canonical, exclude_modules, explicit_targets=()): + # Only legacy Megatron selectors need Bridge before trainer creation. + from megatron.bridge import AutoBridge + + bridge = AutoBridge.from_hf_pretrained(hf_checkpoint, trust_remote_code=True) + hf_mapping = HfWeightMapping.from_config(bridge.hf_pretrained.config) + model_bridge = bridge._model_bridge + model_bridge.hf_pretrained = bridge.hf_pretrained + selected_hf_modules, matched_targets = set(), set() + for mapping in model_bridge.mapping_registry().get_all_mappings(): + megatron_module, weight = mapping.megatron_param.rsplit(".", 1) + if weight not in ("weight", "weight*"): + continue + for checkpoint_parameter in _checkpoint_parameters(mapping): + hf_module = hf_mapping.model_parameter(checkpoint_parameter).removesuffix(".weight") + if hf_mapping.parameter_names and not any( + matches_lora_target(name.removesuffix(".weight"), hf_module) for name in hf_mapping.parameter_names + ): + continue + for target in target_modules: + matches = matches_lora_target(hf_module, target) or matches_lora_target(megatron_module, target) + if ( + canonical + and megatron_module.rsplit(".", 1)[-1] in ("linear_qkv", "linear_fc1") + and ".experts." not in megatron_module + ): + matches |= matches_lora_target( + _canonical_adapter_module(megatron_module, checkpoint_parameter), target + ) + if matches: + matched_targets.add(target) + excluded = any( + matches_lora_target(hf_module, exclude) or matches_lora_target(megatron_module, exclude) + for exclude in exclude_modules + ) + assert not ( + excluded and target in explicit_targets + ), f"Explicit LoRA target {target!r} overlaps --exclude-modules at {hf_module!r}" + if not excluded: + selected_hf_modules.add(hf_module) + assert ( + set(target_modules) <= matched_targets + ), f"LoRA targets have no Bridge mapping: {sorted(set(target_modules) - matched_targets)}" + assert selected_hf_modules, "No LoRA targets remain after applying --exclude-modules" + return sorted(selected_hf_modules) diff --git a/miles/backends/megatron_utils/lora/utils.py b/miles/backends/megatron_utils/lora/utils.py index 78ab6cc8409..2700eb3c261 100644 --- a/miles/backends/megatron_utils/lora/utils.py +++ b/miles/backends/megatron_utils/lora/utils.py @@ -12,116 +12,12 @@ from miles.backends.training_utils.parallel import get_parallel_state from miles.utils.distributed_utils import get_gloo_group -from miles.utils.lora.utils import is_lora_enabled, lora_rollout_enabled # noqa: F401 (re-exported) +from miles.utils.lora.utils import is_lora_enabled # noqa: F401 (re-exported) +from miles.utils.lora.utils import lora_rollout_enabled # noqa: F401 (re-exported) +from miles.utils.lora.utils import build_lora_config, get_adapter_target_modules logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# Unified HF <-> Megatron module name mappings -# --------------------------------------------------------------------------- - -# Standard LoRA: merged Q/K/V and merged up/gate -_STANDARD_LORA_HF_TO_MEGATRON = { - "q_proj": "linear_qkv", - "k_proj": "linear_qkv", - "v_proj": "linear_qkv", - "o_proj": "linear_proj", - "gate_proj": "linear_fc1", - "up_proj": "linear_fc1", - "down_proj": "linear_fc2", - "lm_head": "output_layer", - # GDN (Qwen3.5/Qwen3-Next): both slices live in the single fused megatron in_proj - "in_proj_qkvz": "in_proj", - "in_proj_ba": "in_proj", -} - -_STANDARD_LORA_ALL_MODULES = ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"] - -# CanonicalLoRA: Split Q/K/V and up/gate -_CANONICAL_LORA_HF_TO_MEGATRON = { - "q_proj": "linear_q", - "k_proj": "linear_k", - "v_proj": "linear_v", - "o_proj": "linear_proj", - "gate_proj": "linear_fc1_gate", - "up_proj": "linear_fc1_up", - "down_proj": "linear_fc2", - "lm_head": "output_layer", - "in_proj_qkvz": "in_proj", - "in_proj_ba": "in_proj", -} - -_CANONICAL_LORA_ALL_MODULES = [ - "linear_q", - "linear_k", - "linear_v", - "linear_proj", - "linear_fc1_up", - "linear_fc1_gate", - "linear_fc2", -] - -# Megatron -> HF (inverse mapping, one-to-many) -# Covers both standard LoRA (merged) and CanonicalLoRA (split) module names. -_MEGATRON_TO_HF_MODULES = { - # Standard LoRA (merged layers) - "linear_qkv": ["q_proj", "k_proj", "v_proj"], - "linear_proj": ["o_proj"], - "linear_fc1": ["gate_proj", "up_proj"], - "linear_fc2": ["down_proj"], - "output_layer": ["lm_head"], - # CanonicalLoRA (split layers) - "linear_q": ["q_proj"], - "linear_k": ["k_proj"], - "linear_v": ["v_proj"], - "linear_fc1_gate": ["gate_proj"], - "linear_fc1_up": ["up_proj"], - # GDN linear attention: SGLang serves the fused in_proj as two modules - "in_proj": ["in_proj_qkvz", "in_proj_ba"], -} - -_HF_MODULE_NAMES = { - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - "lm_head", - "in_proj_qkvz", - "in_proj_ba", -} - -# DeepSeek / Kimi MLA (HF names on checkpoint; Megatron uses linear_* from Megatron-Bridge mappings). -_MLA_HF_TO_MEGATRON = { - "q_a_proj": "linear_q_down_proj", - "kv_a_proj_with_mqa": "linear_kv_down_proj", - "q_b_proj": "linear_q_up_proj", - "kv_b_proj": "linear_kv_up_proj", - # DSA indexer (GLM-5 / DeepSeek-V3.2): HF/SGLang leaf names vs Megatron-Bridge linear_* names. - "wq_b": "linear_wq_b", - "wk": "linear_wk", - "weights_proj": "linear_weights_proj", -} -_MEGATRON_MLA_TO_HF = {v: k for k, v in _MLA_HF_TO_MEGATRON.items()} - -# Empty: dropping a module here makes sglang silently skip its shipped adapter tensors. -_SGLANG_UNSUPPORTED_HF_TARGETS = frozenset() - - -# --------------------------------------------------------------------------- -# Core helpers -# --------------------------------------------------------------------------- - - -def sglang_lora_target_all_sentinel(args) -> bool: - """Hand SGLang the ``"all"`` shorthand so it auto-detects module names (required for Inkling).""" - from miles.utils.chat_template_utils.inkling import is_inkling_checkpoint - - return is_inkling_checkpoint(getattr(args, "hf_checkpoint", None) or "") - - _marked_lora_grad_params_cache: dict[int, list] = {} @@ -222,144 +118,13 @@ def _patched_init(self, *args, **kwargs): logger.info("Patched _ParamAndGradBuffer.__init__ for LoRA colocate mode (disable cpu backup)") -# --------------------------------------------------------------------------- -# Module name conversion -# --------------------------------------------------------------------------- - - -def _get_lora_class_name(lora_type: type | object | None) -> str: - """Resolve LoRA type to its class name string.""" - if lora_type is None: - return "CanonicalLoRA" - if isinstance(lora_type, type): - return lora_type.__name__ - return type(lora_type).__name__ - - -def convert_target_modules_to_megatron( - hf_modules: str | list[str], - lora_type: type | object | None = None, -) -> list[str]: - """Convert HuggingFace LoRA target module names to Megatron format. - - HF: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj, lm_head - Megatron (LoRA): linear_qkv, linear_proj, linear_fc1, linear_fc2, output_layer - Megatron (CanonicalLoRA): linear_q, linear_k, linear_v, linear_proj, - linear_fc1_up, linear_fc1_gate, linear_fc2 - - Special values: "all", "all-linear", "all_linear" -> all standard linear modules. - If input is already in Megatron format, returns as-is. - """ - class_name = _get_lora_class_name(lora_type) - is_canonical = class_name == "CanonicalLoRA" - - all_modules = _CANONICAL_LORA_ALL_MODULES if is_canonical else _STANDARD_LORA_ALL_MODULES - hf_to_megatron = _CANONICAL_LORA_HF_TO_MEGATRON if is_canonical else _STANDARD_LORA_HF_TO_MEGATRON - - # Handle special "all-linear" variants - if isinstance(hf_modules, str): - if hf_modules in ("all", "all-linear", "all_linear"): - return list(all_modules) - hf_modules = [hf_modules] - elif isinstance(hf_modules, list) and len(hf_modules) == 1: - if hf_modules[0] in ("all", "all-linear", "all_linear"): - return list(all_modules) - - if isinstance(hf_modules, tuple): - hf_modules = list(hf_modules) - - # Check if already in Megatron format (standard / canonical / Kimi MLA linear_*). - if all(m not in _HF_MODULE_NAMES and m not in _MLA_HF_TO_MEGATRON for m in hf_modules if "*" not in m): - return list(hf_modules) - - # Convert HF names to Megatron names (dedup while preserving order) - megatron_modules: list[str] = [] - for module in hf_modules: - if module in _MLA_HF_TO_MEGATRON: - megatron_name = _MLA_HF_TO_MEGATRON[module] - else: - megatron_name = hf_to_megatron.get(module, module) - if megatron_name not in megatron_modules: - megatron_modules.append(megatron_name) - - return megatron_modules - - -def convert_target_modules_to_hf(megatron_modules: list[str]) -> list[str]: - """Convert Megatron LoRA target module names to HuggingFace format. - - Supports both standard LoRA and CanonicalLoRA module names. - - Megatron standard: linear_qkv, linear_proj, linear_fc1, linear_fc2, output_layer - Megatron canonical: linear_q, linear_k, linear_v, linear_proj, - linear_fc1_up, linear_fc1_gate, linear_fc2 - HF: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj - Kimi MLA Megatron: linear_q_down_proj -> q_a_proj, linear_kv_down_proj -> kv_a_proj_with_mqa, ... - - Wildcards (``*.layers.2.mlp.experts.linear_fc1``) get the last dotted - segment mapped to an HF leaf name; SGLang uses the result to choose - adapter-buffer types, not to scope by layer. - """ - if isinstance(megatron_modules, tuple): - megatron_modules = list(megatron_modules) - hf_modules: list[str] = [] - for module in megatron_modules: - lookup_key = module.rsplit(".", 1)[-1] if "." in module else module - if lookup_key in _MEGATRON_MLA_TO_HF: - hf_modules.append(_MEGATRON_MLA_TO_HF[lookup_key]) - elif lookup_key in _MEGATRON_TO_HF_MODULES: - hf_modules.extend(_MEGATRON_TO_HF_MODULES[lookup_key]) - else: - # same-name passthrough; SGLang needs the leaf, not a path or pattern - hf_modules.append(lookup_key) - seen: set[str] = set() - unique: list[str] = [] - for m in hf_modules: - if m not in seen: - seen.add(m) - unique.append(m) - return unique - - -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 [] - hf = convert_target_modules_to_hf(raw) - out = [m for m in hf if m not in _SGLANG_UNSUPPORTED_HF_TARGETS] - dropped = set(hf) - set(out) - if dropped: - logger.warning( - "target_modules_hf_for_sglang_rollout: omitting %s for SGLang (unsupported by default " - "get_hidden_dim); Megatron should not train LoRA on these if rollout sync is required.", - sorted(dropped), - ) - return out - - # --------------------------------------------------------------------------- # Model setup helpers (used by model.py) # --------------------------------------------------------------------------- -def parse_exclude_modules(args: Namespace, lora_type=None) -> list[str]: - """Parse and convert exclude_modules argument.""" - exclude_modules: list[str] = [] - raw = getattr(args, "exclude_modules", None) - if raw: - if isinstance(raw, str): - exclude_modules = [m.strip() for m in raw.split(",")] - else: - exclude_modules = list(raw) - exclude_modules = convert_target_modules_to_megatron(exclude_modules, lora_type=lora_type) - return exclude_modules - - -def create_lora_instance(args: Namespace): - """Create a LoRA or CanonicalLoRA instance based on args. - - Returns: - A LoRA/CanonicalLoRA dataclass instance ready to be applied to a model. - """ +def create_lora_instance(args: Namespace, *, target_modules): + """Create a LoRA adapter with resolved Megatron target modules.""" from megatron.bridge.peft.canonical_lora import CanonicalLoRA from megatron.bridge.peft.lora import LoRA @@ -370,12 +135,8 @@ def create_lora_instance(args: Namespace): else: lora_cls = LoRA - target_modules = convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls) - exclude_modules = parse_exclude_modules(args, lora_type=lora_cls) - lora_kwargs = dict( target_modules=target_modules, - exclude_modules=exclude_modules, dim=args.lora_rank, alpha=args.lora_alpha, dropout=args.lora_dropout, @@ -393,8 +154,7 @@ def create_lora_instance(args: Namespace): logger.info( f"Created {lora_cls.__name__}: rank={args.lora_rank}, alpha={args.lora_alpha}, " - f"dropout={args.lora_dropout}, target_modules={target_modules}, " - f"exclude_modules={exclude_modules}" + f"dropout={args.lora_dropout}, target_modules={target_modules}" ) return lora @@ -461,17 +221,11 @@ def save_lora_checkpoint( if native: if global_rank == 0: - config = { - "peft_type": "LORA", - "r": args.lora_rank, - "lora_alpha": args.lora_alpha, - "target_modules": convert_target_modules_to_hf(list(args.target_modules)), - "lora_dropout": args.lora_dropout, - "bias": "none", - "task_type": "CAUSAL_LM", - "experts_shared_outer_loras": bool(args.experts_shared_outer_loras), - "format": "megatron_rank_sharded", - } + config = build_lora_config(args, target_modules=args.lora_adapter_targets) + config.update( + experts_shared_outer_loras=bool(args.experts_shared_outer_loras), + format="megatron_rank_sharded", + ) with open(save_path / "adapter_config.json", "w") as f: json.dump(config, f, indent=2) os.sync() @@ -486,6 +240,8 @@ def save_lora_checkpoint( from miles.utils import megatron_bridge_utils bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) + # Megatron is optional until the HF export path is entered. + from miles.backends.megatron_utils.update_weight.hf_weight_iterator import _gather_pp_full_adapter lora_state_dict: dict[str, torch.Tensor] = {} with megatron_bridge_utils.patch_megatron_model(model): @@ -496,33 +252,21 @@ def save_lora_checkpoint( ): lora_state_dict[hf_name] = weight - if is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0: - torch.save(lora_state_dict, save_path / "adapter_model.bin") - - target_modules_hf = ( - convert_target_modules_to_hf(list(args.target_modules)) - if args.target_modules - else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - ) - config = { - "peft_type": "LORA", - "r": args.lora_rank, - "lora_alpha": args.lora_alpha, - "target_modules": target_modules_hf, - "lora_dropout": args.lora_dropout, - "bias": "none", - "task_type": "CAUSAL_LM", - } - with open(save_path / "adapter_config.json", "w") as f: - json.dump(config, f, indent=2) - - os.sync() - logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors") + lora_state_dict = dict(_gather_pp_full_adapter(list(lora_state_dict.items()))) except Exception as 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." ) + else: + if is_dp_cp_rank_0 and tp_rank == 0 and pp_rank == 0: + torch.save( + {name: weight.cpu() for name, weight in lora_state_dict.items()}, save_path / "adapter_model.bin" + ) + config = build_lora_config(args, target_modules=get_adapter_target_modules(lora_state_dict)) + (save_path / "adapter_config.json").write_text(json.dumps(config, indent=2)) + os.sync() + logger.info(f"Saved HF PEFT adapter to {save_path} with {len(lora_state_dict)} tensors") # ---- Training state (iteration + scheduler, and the optimizer unless opted out) ---- if optimizer is not None: @@ -666,29 +410,3 @@ def _load_training_state( if iteration is not None: logger.info(f"Resuming LoRA training from iteration {iteration}") return iteration - - -# --------------------------------------------------------------------------- -# LoRA config dict for weight sync to SGLang -# --------------------------------------------------------------------------- - - -def build_lora_sync_config(args: Namespace) -> dict[str, Any]: - """Build LoRA config dict for syncing weights to SGLang engines.""" - if sglang_lora_target_all_sentinel(args): - target_modules_hf: Any = "all-linear" - else: - target_modules_hf = ( - target_modules_hf_for_sglang_rollout(args) - if args.target_modules - else ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - ) - return { - "peft_type": "LORA", - "r": args.lora_rank, - "lora_alpha": args.lora_alpha, - "target_modules": target_modules_hf, - "lora_dropout": args.lora_dropout, - "bias": "none", - "task_type": "CAUSAL_LM", - } diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 9619838696f..36d1622fe14 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -156,7 +156,8 @@ def setup_model_and_optimizer( else: provider_func = get_model_provider_func(args, role) if is_lora_enabled(args) and role == "actor": - if "inkling" in (getattr(args, "custom_model_provider_path", None) or ""): + if "inkling" in (args.custom_model_provider_path or ""): + assert args.lora_type == "lora", "Native Inkling does not implement --lora-type canonical_lora" from miles_plugins.models.inkling.lora import wrap_model_provider_with_inkling_lora provider_func = wrap_model_provider_with_inkling_lora(provider_func, args) @@ -169,6 +170,11 @@ def setup_model_and_optimizer( provider_func = wrap_model_provider_with_kimi_k3_lora(provider_func, args) if args.offload_train: patch_param_grad_buffer_for_colocate_mode_lora() + else: + raise AssertionError( + "Native LoRA injection is only implemented for Inkling and Kimi K3; " + "use --megatron-to-hf-mode bridge" + ) model = get_model(provider_func, ModelType.encoder_or_decoder) if args.debug_disable_optimizer: diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py index 98cb667c1bf..aab504c5696 100644 --- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py +++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator.py @@ -54,10 +54,8 @@ def _iter_hf_adapter_units(self, adapter, *, materialize): if not materialize: return if not named_tensors: - raise RuntimeError( - f"LoRA weight sync failed: the adapter export produced zero tensors{f' for adapter {adapter!r}' if adapter is not None else ''}. This usually means the Megatron-Bridge or SGLang version is incompatible." - ) - if not any(is_lora_weight_name(name) for name, _tensor in named_tensors): + raise RuntimeError("LoRA weight sync failed: the adapter export produced zero tensors") + if not any(is_lora_weight_name(name) for name, _ in named_tensors): raise RuntimeError("LoRA weight sync failed: the adapter export contains no lora_A/lora_B names.") while named_tensors: hf_name, tensor = named_tensors.pop(0) diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 6b93d671a5e..ad7c5c1b4f1 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -6,7 +6,6 @@ from sglang.srt.server_args import ServerArgs -from miles.backends.megatron_utils.lora.utils import convert_target_modules_to_hf, sglang_lora_target_all_sentinel from miles.backends.sglang_utils.server_args_utils import server_args_to_argv from miles.utils.lora.utils import ( LORA_ADAPTER_NAME, @@ -148,15 +147,16 @@ def _compute_server_args( kwargs["enable_lora"] = True kwargs["max_loras_per_batch"] = args.multi_lora_n_adapters kwargs["max_lora_rank"] = max(getattr(args, "lora_rank", 0), 1) - kwargs["lora_target_modules"] = convert_target_modules_to_hf(args.target_modules) + kwargs["lora_target_modules"] = ( + ["all"] if args.lora_adapter_targets == "all-linear" else args.lora_adapter_targets + ) elif lora_rollout_enabled(args): kwargs["enable_lora"] = True kwargs["max_loras_per_batch"] = 1 kwargs["max_lora_rank"] = max(getattr(args, "lora_rank", 0), 1) - if sglang_lora_target_all_sentinel(args): - kwargs["lora_target_modules"] = ["all"] - else: - kwargs["lora_target_modules"] = convert_target_modules_to_hf(args.target_modules) + kwargs["lora_target_modules"] = ( + ["all"] if args.lora_adapter_targets == "all-linear" else args.lora_adapter_targets + ) if engine_loads_adapter_from_disk(args): kwargs["lora_paths"] = [f"{LORA_ADAPTER_NAME}={args.lora_adapter_path}"] diff --git a/miles/backends/training_utils/weight_update/snapshot_publisher.py b/miles/backends/training_utils/weight_update/snapshot_publisher.py index c9e4bc0cdb1..439e48a4980 100644 --- a/miles/backends/training_utils/weight_update/snapshot_publisher.py +++ b/miles/backends/training_utils/weight_update/snapshot_publisher.py @@ -7,7 +7,7 @@ from miles.backends.training_utils.checkpoint_io import write_checkpoint_dir from miles.backends.training_utils.weight_update.hf_weight_iterator import HfWeightIteratorBase -from miles.utils.lora.utils import AdapterSpec +from miles.utils.lora.utils import AdapterSpec, get_adapter_target_modules class WeightPublisher: @@ -28,7 +28,11 @@ def write_shards(tmp_dir: Path): data = safetensors.torch.save(tensors) if is_writer else None if is_writer: - config = self._adapter_config | {"r": adapter.rank, "lora_alpha": adapter.alpha} + config = self._adapter_config | { + "r": adapter.rank, + "lora_alpha": adapter.alpha, + "target_modules": get_adapter_target_modules(tensors), + } (tmp_dir / "adapter_config.json").write_text(json.dumps(config)) (tmp_dir / "adapter_model.safetensors").write_bytes(data) diff --git a/miles/tinker/arguments.py b/miles/tinker/arguments.py index 24a35ca2ae2..b51d8fc479e 100644 --- a/miles/tinker/arguments.py +++ b/miles/tinker/arguments.py @@ -1,6 +1,4 @@ -import argparse - -from miles.utils.hf_utils.config import load_hf_config +from miles.utils.lora.hf_lora_targets import LORA_TARGET_GROUPS, parse_lora_targets def add_tinker_arguments(parser): @@ -19,39 +17,15 @@ def add_argument(name, **kwargs): "checkpoint-root", help="Directory for tinker:// checkpoints (default: /tinker)", ) - add_argument("train-attn", action=argparse.BooleanOptionalAction, default=True) - add_argument("train-mlp", action=argparse.BooleanOptionalAction, default=True) - add_argument("train-unembed", action=argparse.BooleanOptionalAction, default=True) return parser def configure_tinker_args(args): assert args.train_backend == "megatron", "Tinker requires the Megatron backend" - assert ( - args.target_modules is None and args.exclude_modules is None - ), "Tinker uses --tinker-train-attn/mlp/unembed; --target-modules and --exclude-modules are not supported" - modules = _resolve_target_modules( - load_hf_config(args.hf_checkpoint), - train_attn=args.tinker_train_attn, - train_mlp=args.tinker_train_mlp, - train_unembed=args.tinker_train_unembed, - ) - # The common LoRA validator parses and validates this before trainer/engine initialization. - args.target_modules = ",".join(modules) - - -def _resolve_target_modules(hf_config, *, train_attn, train_mlp, train_unembed): - # Other architectures need their own complete attention/MLP mapping. - assert hf_config.model_type in ( - "qwen3", - "qwen3_moe", - ), f"Tinker target layout is not defined for model_type={hf_config.model_type!r}" - modules = [] - if train_attn: - modules.extend(("q_proj", "k_proj", "v_proj", "o_proj")) - if train_mlp: - modules.extend(("gate_proj", "up_proj", "down_proj")) - if train_unembed: - modules.append("lm_head") - assert modules, "Tinker requires at least one trainable LoRA module group" - return modules + assert args.exclude_modules is None, "Tinker selects complete training groups; --exclude-modules is not supported" + groups = parse_lora_targets(args.target_modules) + if groups is None: + groups = list(LORA_TARGET_GROUPS) + assert set(groups) <= set(LORA_TARGET_GROUPS), "Tinker --target-modules accepts only attn,mlp,unembed groups" + args.tinker_lora_groups = groups + args.target_modules = groups diff --git a/miles/tinker/core/input_validation.py b/miles/tinker/core/input_validation.py index 63b4f49fc5c..26da0f1a3c0 100644 --- a/miles/tinker/core/input_validation.py +++ b/miles/tinker/core/input_validation.py @@ -25,7 +25,7 @@ def validate_model_config(lora_config: dict, config: GatewayConfig) -> None: if requested is not None and requested != layout_trains: raise UserInputError( f"lora_config.{field}={requested} conflicts with this gateway's adapter layout " - f"({field}={layout_trains}); the layout is fixed by --tinker-train-attn/mlp/unembed at server start" + f"({field}={layout_trains}); the layout is fixed by the gateway target selection at server start" ) rank = lora_config.get("rank", 32) if type(rank) is not int or not 1 <= rank <= config.max_lora_rank: diff --git a/miles/utils/hf_utils/weight_mapping.py b/miles/utils/hf_utils/weight_mapping.py new file mode 100644 index 00000000000..70f340f8c74 --- /dev/null +++ b/miles/utils/hf_utils/weight_mapping.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass + +import torch +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText +from transformers.conversion_mapping import get_model_conversion_mapping +from transformers.core_model_loading import WeightConverter, WeightRenaming +from transformers.models.auto.auto_factory import _get_model_class + + +@dataclass(frozen=True) +class HfWeightMapping: + parameter_names: frozenset[str] + conversions: tuple = () + + @classmethod + def from_config(cls, config): + # VLMs may also register a CausalLM compatibility class that drops the vision/text namespace. + for auto_model in (AutoModelForImageTextToText, AutoModelForCausalLM): + if type(config) in auto_model._model_mapping: + # HF's lazy mapping also matches remote-code configs by class name. + config_class = _get_model_class(config, auto_model._model_mapping).config_class + if type(config) is not config_class: + config = config_class.from_dict(config.to_dict()) + # Only structure is needed; never allocate or load base weights. + with torch.random.fork_rng(devices=[]), torch.device("meta"): + model = auto_model.from_config(config, attn_implementation="eager") + parameter_names = frozenset( + name for name, param in model.named_parameters(remove_duplicate=False) if param.ndim in (2, 3) + ) + return cls(parameter_names, tuple(get_model_conversion_mapping(model, add_legacy=False))) + # Custom HF implementations without native conversion rules retain their checkpoint namespace. + return cls(frozenset()) + + def model_parameter(self, checkpoint_name): + if checkpoint_name.removesuffix(".weight") in self.parameter_names: + checkpoint_name = checkpoint_name.removesuffix(".weight") + if checkpoint_name in self.parameter_names: + return checkpoint_name + for conversion in self.conversions: + if isinstance(conversion, WeightRenaming): + checkpoint_name, _ = conversion.rename_source_key(checkpoint_name) + for conversion in self.conversions: + if isinstance(conversion, WeightConverter): + target, source_pattern = conversion.rename_source_key(checkpoint_name) + if source_pattern is not None: + assert ( + len(conversion.target_patterns) == 1 + ), f"HF target binding does not support one-to-many conversion of {checkpoint_name!r}" + return target + return checkpoint_name diff --git a/miles/utils/lora/arguments.py b/miles/utils/lora/arguments.py index 7884143e89e..83c549fd181 100644 --- a/miles/utils/lora/arguments.py +++ b/miles/utils/lora/arguments.py @@ -2,6 +2,19 @@ from typing import Any from miles.utils.hf_utils.config import load_hf_config +from miles.utils.hf_utils.weight_mapping import HfWeightMapping +from miles.utils.lora.hf_lora_targets import ( + LORA_TARGET_GROUPS, + exclude_hf_lora_targets, + expand_hf_lora_targets, + expand_packed_hf_lora_targets, + get_hf_lora_targets, + parse_lora_targets, + resolve_hf_lora_targets, +) +from miles.utils.lora.utils import is_lora_enabled, matches_lora_target, targets_expert_leaves +from miles_plugins.models.inkling.lora import resolve_inkling_adapter_targets +from miles_plugins.models.kimi_k3.lora import resolve_kimi_k3_adapter_targets logger = logging.getLogger(__name__) @@ -37,8 +50,9 @@ def add_lora_arguments(parser): "--target-modules", type=str, default=None, - help="Target modules for LoRA. Use 'all-linear' or comma-separated module names " - "(e.g., 'q_proj,k_proj,v_proj,o_proj' for HF naming or 'linear_qkv,linear_proj' for Megatron naming)", + help="LoRA targets: omit or use 'all-linear' for the HF model defaults, or provide " + "comma-separated groups (attn,mlp,unembed), HF targets, or a mix of both. " + "Megatron module names are also accepted by the Bridge backend.", ) parser.add_argument( "--exclude-modules", @@ -65,10 +79,7 @@ def add_lora_arguments(parser): help=( "LoRA + colocate: keep SGLang-side CPU mirror of base weights " "and skip per-step base sync. Trades host RAM for faster " - "onload/offload. Ignored unless --colocate and LoRA are both on. " - "Also needs 'weight' in --offload-rollout-level: SGLang populates " - "the mirror during release_weights_occupation, so with the weights " - "never released the mirror is never built and the flag does nothing." + "onload/offload. Ignored unless --colocate and LoRA are both on." ), ) parser.add_argument( @@ -123,53 +134,86 @@ def add_lora_arguments(parser): def validate_lora_args(args): - # Parse LoRA target modules - if args.lora_rank > 0: - assert args.target_modules is not None, "'--target-modules' is required when LoRA is enabled." + validate_multi_lora_args(args) + if not is_lora_enabled(args): + return + assert args.train_backend == "megatron", "LoRA injection is not implemented for FSDP; use --train-backend megatron" + assert args.lora_rank > 0, "LoRA requires a positive --lora-rank, including when loading an adapter" + hf_config = load_hf_config(args.hf_checkpoint) + args.hf_lora_targets, args.lora_adapter_targets = _resolve_lora_targets(args, hf_config) - if args.target_modules == "all-linear": - # MLA projections are HF-config-gated (SGLang sizes LoRA buffers per module name; - # listing them on a dense model crashes the engine). The DSA indexer stays excluded. - modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - hf_config = load_hf_config(args.hf_checkpoint) - if getattr(hf_config, "kv_lora_rank", None): - modules += ["kv_a_proj_with_mqa", "kv_b_proj"] - if getattr(hf_config, "q_lora_rank", None): - modules += ["q_a_proj", "q_b_proj"] - elif "," in args.target_modules: - modules = [m.strip() for m in args.target_modules.split(",")] - else: - modules = [args.target_modules] + # Training and serving must agree on shared-outer grouped-expert LoRA (expert_dim=1). + if args.experts_shared_outer_loras and hasattr(args, "sglang_experts_shared_outer_loras"): + args.sglang_experts_shared_outer_loras = True + assert args.experts_shared_outer_loras == bool( + getattr(args, "sglang_experts_shared_outer_loras", args.experts_shared_outer_loras) + ), "experts_shared_outer_loras and sglang_experts_shared_outer_loras must agree" + if targets_expert_leaves(args.hf_lora_targets): + logger.warning( + "MoE-expert LoRA layout: %s (--experts-shared-outer-loras).", + "shared-outer" if args.experts_shared_outer_loras else "per-expert", + ) - if args.exclude_modules: - exclude_set = ( - set(m.strip() for m in args.exclude_modules.split(",")) - if "," in args.exclude_modules - else {args.exclude_modules} - ) - modules = [m for m in modules if m not in exclude_set] - args.target_modules = modules +def _resolve_lora_targets(args, hf_config): + """Return final HF and adapter targets without changing the CLI selectors in `args`.""" + hf_mapping = HfWeightMapping.from_config(hf_config) + hf_modules = [name.removesuffix(".weight") for name in hf_mapping.parameter_names] + targets = parse_lora_targets(args.target_modules) + exclusions = parse_lora_targets(args.exclude_modules) or [] + explicit_targets = [] + if targets is not None: + expanded = expand_packed_hf_lora_targets(targets, hf_modules) + explicit_targets = [ + target for target in dict.fromkeys(targets + expanded) if target not in (*LORA_TARGET_GROUPS, "all-linear") + ] + targets = expanded + if exclusions: + conflicts = { + module + for module in hf_modules + explicit_targets + exclusions + if any(matches_lora_target(module, target) for target in explicit_targets) + and any(matches_lora_target(module, exclude) for exclude in exclusions) + } + assert not conflicts, f"Explicit LoRA targets overlap --exclude-modules: {sorted(conflicts)}" + if targets is None and args.multi_lora: + targets = list(LORA_TARGET_GROUPS) + targets = resolve_hf_lora_targets(hf_config.to_dict(), target_modules=targets) + targets = exclude_hf_lora_targets(targets, exclusions) - # Training and serving must agree on shared-outer grouped-expert LoRA - # (expert_dim=1 buffers in SGLang). - if args.experts_shared_outer_loras and hasattr(args, "sglang_experts_shared_outer_loras"): - args.sglang_experts_shared_outer_loras = True - assert args.experts_shared_outer_loras == bool( - getattr(args, "sglang_experts_shared_outer_loras", args.experts_shared_outer_loras) - ), "experts_shared_outer_loras and sglang_experts_shared_outer_loras must agree" + if all(any(matches_lora_target(module, target) for module in hf_modules) for target in targets + exclusions): + hf_targets = targets + if exclusions: + selected = [ + module for module in hf_modules if any(matches_lora_target(module, target) for target in targets) + ] + hf_targets = exclude_hf_lora_targets(selected, exclusions) + elif args.megatron_to_hf_mode == "bridge": + # Only legacy Megatron selectors need Bridge before trainer creation. + from miles.backends.megatron_utils.lora.target_modules import normalize_lora_targets_to_hf - # the two MoE-expert adapter layouts are not checkpoint-compatible; say which one runs - _expert_leaves = ("linear_fc1", "linear_fc2", "gate_proj", "up_proj", "down_proj") - if any(leaf in str(tm) for tm in modules for leaf in _expert_leaves): - logger.warning( - "MoE-expert LoRA layout: %s (--experts-shared-outer-loras).", - "shared-outer" if args.experts_shared_outer_loras else "per-expert", - ) + hf_targets = normalize_lora_targets_to_hf( + args.hf_checkpoint, + targets, + canonical=args.lora_type == "canonical_lora", + exclude_modules=exclusions, + explicit_targets=explicit_targets, + ) + else: + layout = get_hf_lora_targets(hf_config.to_dict()) + hf_targets = exclude_hf_lora_targets(expand_hf_lora_targets(targets, layout), exclusions) - # Sets args.multi_lora, then validates/defaults the multi-LoRA arg surface - # (adapter configs themselves are loaded later by the controller). - validate_multi_lora_args(args) + adapter_targets = list(hf_targets) + if args.megatron_to_hf_mode == "raw": + if hf_config.model_type in ("inkling_model", "inkling_mm_model", "inkling_text"): + adapter_targets = resolve_inkling_adapter_targets(hf_config.to_dict(), hf_targets) + elif hf_config.model_type == "kimi_k3": + adapter_targets = resolve_kimi_k3_adapter_targets( + hf_targets, + canonical=args.lora_type == "canonical_lora", + experts_shared_outer_loras=args.experts_shared_outer_loras, + ) + return hf_targets, adapter_targets def validate_multi_lora_args(args: Any) -> None: @@ -180,7 +224,6 @@ def validate_multi_lora_args(args: Any) -> None: return 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, "context_parallel_size", 1) == 1, ( diff --git a/miles/utils/lora/hf_lora_targets.py b/miles/utils/lora/hf_lora_targets.py new file mode 100644 index 00000000000..52345a1f33c --- /dev/null +++ b/miles/utils/lora/hf_lora_targets.py @@ -0,0 +1,303 @@ +"""LoRA projection groups in the HF model namespace; backend layouts are resolved separately.""" + +import logging +from collections.abc import Callable +from dataclasses import dataclass + +from miles.utils.lora.utils import matches_lora_target + +logger = logging.getLogger(__name__) + +LORA_TARGET_GROUPS = ("attn", "mlp", "unembed") + + +@dataclass(frozen=True) +class HfLoraTargets: + attention: tuple[str, ...] + mlp: tuple[str, ...] + unembed: tuple[str, ...] + default_train_unembed: bool = False + + +@dataclass(frozen=True) +class _HfLoraModelSpec: + build_groups: Callable[[dict], tuple[tuple[str, ...], tuple[str, ...]]] + layer_prefix: str = "model.layers.*" + unembed: str = "lm_head" + unwrap_text_config: bool = False + default_train_unembed: bool = False + + +def _prefix_paths(prefix, *paths): + return tuple(f"{prefix}.{path}" for path in paths) + + +_QKVO_ATTENTION = _prefix_paths("self_attn", "q_proj", "k_proj", "v_proj", "o_proj") +_DENSE_MLP = _prefix_paths("mlp", "gate_proj", "up_proj", "down_proj") +_ROUTED_EXPERTS = _prefix_paths("mlp.experts.*", "gate_proj", "up_proj", "down_proj") +_SHARED_EXPERTS = _prefix_paths("mlp.shared_experts", "gate_proj", "up_proj", "down_proj") +_QWEN_SHARED_EXPERT = _prefix_paths("mlp.shared_expert", "gate_proj", "up_proj", "down_proj") +_PACKED_EXPERTS = _prefix_paths("mlp.experts", "gate_up_proj", "down_proj") +_GDN_NEXT = _prefix_paths("linear_attn", "in_proj_qkvz", "in_proj_ba", "out_proj") +_GDN_35 = _prefix_paths("linear_attn", "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a", "out_proj") + + +def _dense_targets(config): + return _QKVO_ATTENTION, _DENSE_MLP + + +def _gpt_oss_targets(config): + return _QKVO_ATTENTION, _PACKED_EXPERTS + + +def _deepseek_mlp_targets(config): + mlp = [] + if config["first_k_dense_replace"] > 0: + mlp.extend(_DENSE_MLP) + if config["first_k_dense_replace"] < config["num_hidden_layers"]: + mlp.extend(_PACKED_EXPERTS) + if config["n_shared_experts"]: + mlp.extend(_SHARED_EXPERTS) + return tuple(mlp) + + +def _mla_attention_targets(config): + query = ("q_proj",) if config["q_lora_rank"] is None else ("q_a_proj", "q_b_proj") + return _prefix_paths("self_attn", *query, "kv_a_proj_with_mqa", "kv_b_proj", "o_proj") + + +def _deepseek_targets(config): + return _mla_attention_targets(config), _deepseek_mlp_targets(config) + + +def _kimi_targets(config): + # Kimi's custom HF implementation retains per-expert modules and interleaved dense layers. + num_moe_layers = sum( + bool(config["n_routed_experts"]) + and layer_id >= config["first_k_dense_replace"] + and layer_id % config.get("moe_layer_freq", 1) == 0 + for layer_id in range(config["num_hidden_layers"]) + ) + mlp = [] + if num_moe_layers < config["num_hidden_layers"]: + mlp.extend(_DENSE_MLP) + if num_moe_layers: + mlp.extend(_ROUTED_EXPERTS) + if config["n_shared_experts"]: + mlp.extend(_SHARED_EXPERTS) + return _mla_attention_targets(config), tuple(mlp) + + +def _kimi_k3_targets(config): + attention = _prefix_paths("self_attn", "o_proj", "q_a_proj", "kv_a_proj_with_mqa") + num_moe_layers = sum( + config["num_experts"] is not None + and layer >= config["first_k_dense_replace"] + and layer % config["moe_layer_freq"] == 0 + for layer in range(config["num_hidden_layers"]) + ) + mlp = list(_DENSE_MLP) if num_moe_layers < config["num_hidden_layers"] else [] + if num_moe_layers: + mlp.extend(_prefix_paths("block_sparse_moe.experts.*", "w1", "w2", "w3")) + if config["num_shared_experts"] is not None: + mlp.extend(_prefix_paths("block_sparse_moe.shared_experts", "gate_proj", "up_proj", "down_proj")) + return attention, tuple(mlp) + + +def _glm4_moe_targets(config): + return _QKVO_ATTENTION, _deepseek_mlp_targets(config) + + +def _glm_dsa_targets(config): + mlp = [] + if "dense" in config["mlp_layer_types"]: + mlp.extend(_DENSE_MLP) + if "sparse" in config["mlp_layer_types"]: + mlp.extend(_PACKED_EXPERTS) + if config["n_shared_experts"]: + mlp.extend(_SHARED_EXPERTS) + return _mla_attention_targets(config), tuple(mlp) + + +def _qwen_moe_mlp_targets(config, *, shared_expert=False): + # Qwen3MoE serializes its num_experts alias as num_local_experts. + num_experts = config["num_experts"] if "num_experts" in config else config["num_local_experts"] + # These optional HF fields allow dense layers inside an otherwise MoE model. + num_moe_layers = sum( + bool(num_experts) + and layer_id not in config.get("mlp_only_layers", []) + and (layer_id + 1) % config.get("decoder_sparse_step", 1) == 0 + for layer_id in range(config["num_hidden_layers"]) + ) + mlp = [] + if num_moe_layers < config["num_hidden_layers"]: + mlp.extend(_DENSE_MLP) + if num_moe_layers: + mlp.extend(_PACKED_EXPERTS) + if shared_expert and config["shared_expert_intermediate_size"]: + mlp.extend(_QWEN_SHARED_EXPERT) + return tuple(mlp) + + +def _hybrid_attention_targets(config, linear_attention): + attention = [] + if "full_attention" in config["layer_types"]: + attention.extend(_QKVO_ATTENTION) + if "linear_attention" in config["layer_types"]: + attention.extend(linear_attention) + return tuple(attention) + + +def _qwen3_moe_targets(config): + return _QKVO_ATTENTION, _qwen_moe_mlp_targets(config) + + +def _qwen3_next_targets(config): + return _hybrid_attention_targets(config, _GDN_NEXT), _qwen_moe_mlp_targets(config, shared_expert=True) + + +def _qwen3_5_targets(config): + return _hybrid_attention_targets(config, _GDN_35), _DENSE_MLP + + +def _qwen3_5_moe_targets(config): + mlp = _PACKED_EXPERTS + if config["shared_expert_intermediate_size"]: + mlp += _QWEN_SHARED_EXPERT + return _hybrid_attention_targets(config, _GDN_35), mlp + + +def _inkling_targets(config): + attention = _prefix_paths("self_attn", "q_proj", "k_proj", "v_proj", "r_proj", "o_proj") + # Older Inkling configs encode leading dense layers instead of mlp_layer_types. + layer_types = config.get("mlp_layer_types") + if layer_types is None: + layer_types = [ + "dense" if layer < config["dense_mlp_idx"] else "sparse" for layer in range(config["num_hidden_layers"]) + ] + mlp = [] + if "dense" in layer_types: + mlp.extend(_DENSE_MLP) + if "sparse" in layer_types: + mlp.extend(_PACKED_EXPERTS) + if config["n_shared_experts"]: + mlp.extend(_SHARED_EXPERTS) + return attention, tuple(mlp) + + +_HF_LORA_MODELS = { + "llama": _HfLoraModelSpec(_dense_targets), + "qwen2": _HfLoraModelSpec(_dense_targets), + "qwen3": _HfLoraModelSpec(_dense_targets), + "qwen3_moe": _HfLoraModelSpec(_qwen3_moe_targets), + "qwen3_next": _HfLoraModelSpec(_qwen3_next_targets), + "qwen3_5_text": _HfLoraModelSpec(_qwen3_5_targets), + "qwen3_5_moe_text": _HfLoraModelSpec(_qwen3_5_moe_targets), + "qwen3_5": _HfLoraModelSpec( + _qwen3_5_targets, layer_prefix="model.language_model.layers.*", unwrap_text_config=True + ), + "qwen3_5_moe": _HfLoraModelSpec( + _qwen3_5_moe_targets, layer_prefix="model.language_model.layers.*", unwrap_text_config=True + ), + "gpt_oss": _HfLoraModelSpec(_gpt_oss_targets), + "deepseek_v2": _HfLoraModelSpec(_deepseek_targets), + "deepseek_v3": _HfLoraModelSpec(_deepseek_targets), + "deepseek_v32": _HfLoraModelSpec(_deepseek_targets), + "kimi_k2": _HfLoraModelSpec(_kimi_targets), + "kimi_k25": _HfLoraModelSpec( + _kimi_targets, + layer_prefix="language_model.model.layers.*", + unembed="language_model.lm_head", + unwrap_text_config=True, + ), + "kimi_k3": _HfLoraModelSpec( + _kimi_k3_targets, + layer_prefix="language_model.model.layers.*", + unembed="language_model.lm_head", + unwrap_text_config=True, + ), + "glm4_moe": _HfLoraModelSpec(_glm4_moe_targets), + "glm_moe_dsa": _HfLoraModelSpec(_glm_dsa_targets), + "inkling_text": _HfLoraModelSpec(_inkling_targets, default_train_unembed=True), + "inkling_model": _HfLoraModelSpec(_inkling_targets, default_train_unembed=True), + "inkling_mm_model": _HfLoraModelSpec( + _inkling_targets, + layer_prefix="model.language_model.layers.*", + unwrap_text_config=True, + default_train_unembed=True, + ), +} + + +def get_hf_lora_targets(hf_config: dict) -> HfLoraTargets: + model_type = hf_config["model_type"] + assert model_type in _HF_LORA_MODELS, f"HF LoRA target layout is not defined for model_type={model_type!r}" + spec = _HF_LORA_MODELS[model_type] + text_config = hf_config["text_config"] if spec.unwrap_text_config else hf_config + attention, mlp = spec.build_groups(text_config) + return HfLoraTargets( + attention=_prefix_paths(spec.layer_prefix, *attention), + mlp=_prefix_paths(spec.layer_prefix, *mlp), + unembed=(spec.unembed,), + default_train_unembed=spec.default_train_unembed, + ) + + +def resolve_hf_lora_targets(hf_config: dict, *, target_modules: list[str] | None = None) -> list[str]: + if target_modules is None or target_modules == ["all-linear"]: + layout = get_hf_lora_targets(hf_config) + target_modules = ["attn", "mlp", *(["unembed"] if layout.default_train_unembed else [])] + assert target_modules and "all-linear" not in target_modules, "Use all-linear alone or provide explicit targets" + if not any(target in LORA_TARGET_GROUPS for target in target_modules): + return list(target_modules) + + layout = get_hf_lora_targets(hf_config) + groups = dict(zip(LORA_TARGET_GROUPS, (layout.attention, layout.mlp, layout.unembed), strict=True)) + targets = [module for target in target_modules for module in groups.get(target, (target,))] + assert targets, "At least one trainable LoRA module group is required" + return list(dict.fromkeys(targets)) + + +def expand_packed_hf_lora_targets(targets: list[str], modules: list[str]) -> list[str]: + added, replaced = [], set() + for gate in targets: + if gate.rsplit(".", 1)[-1] != "gate_proj": + continue + prefix = gate.removesuffix("gate_proj") + up, packed = prefix + "up_proj", prefix + "gate_up_proj" + if up not in targets or not any(matches_lora_target(module, packed) for module in modules): + continue + if not any(matches_lora_target(packed, target) for target in targets): + added.append(packed) + replaced.update( + target for target in (gate, up) if not any(matches_lora_target(module, target) for module in modules) + ) + if added: + logger.warning("Expanding paired gate_proj/up_proj LoRA targets to packed HF targets: %s", added) + return list(dict.fromkeys([target for target in targets if target not in replaced] + added)) + + +def parse_lora_targets(value: str | list[str] | None) -> list[str] | None: + if value is None: + return None + targets = value.split(",") if isinstance(value, str) else value + targets = [target.strip() for target in targets] + assert targets and all(targets), "LoRA target lists must not contain empty entries" + return list(dict.fromkeys(targets)) + + +def exclude_hf_lora_targets(targets: list[str], exclusions: list[str]) -> list[str]: + selected = [ + target for target in targets if not any(matches_lora_target(target, pattern) for pattern in exclusions) + ] + assert selected, "LoRA target selection is empty after --exclude-modules" + return selected + + +def expand_hf_lora_targets(targets: list[str], layout: HfLoraTargets) -> list[str]: + available = layout.attention + layout.mlp + layout.unembed + for target in targets: + assert any( + matches_lora_target(module, target) for module in available + ), f"LoRA target {target!r} is not an HF target of this model" + return [module for module in available if any(matches_lora_target(module, target) for target in targets)] diff --git a/miles/utils/lora/utils.py b/miles/utils/lora/utils.py index 3492e470f1f..34962c2de19 100644 --- a/miles/utils/lora/utils.py +++ b/miles/utils/lora/utils.py @@ -1,12 +1,17 @@ import json from argparse import Namespace from dataclasses import dataclass +from fnmatch import fnmatchcase from pathlib import Path from typing import Any LORA_ADAPTER_NAME = "miles_lora" +def matches_lora_target(module: str, target: str) -> bool: + return fnmatchcase(module if "." in target else module.rsplit(".", 1)[-1], target) + + def is_lora_weight_name(name: str) -> bool: """Check if an HF weight name corresponds to a LoRA adapter weight.""" return ".lora_A." in name or ".lora_B." in name @@ -46,6 +51,22 @@ def save_adapter_to_disk(out_dir, config: dict, tensors: dict) -> None: safetensors.torch.save_file(tensors, str(out / "adapter_model.safetensors")) +def build_lora_config(args, *, target_modules): + return { + "peft_type": "LORA", + "r": args.lora_rank, + "lora_alpha": args.lora_alpha, + "target_modules": target_modules if isinstance(target_modules, str) else list(target_modules), + "lora_dropout": args.lora_dropout, + "bias": "none", + "task_type": "CAUSAL_LM", + } + + +def get_adapter_target_modules(weight_names): + return sorted({name.removeprefix("base_model.model.").rsplit(".lora_", 1)[0] for name in weight_names}) + + @dataclass(frozen=True) class AdapterSpec: """The slot and scaling needed to export one adapter.""" @@ -61,7 +82,7 @@ def is_multi_lora_enabled(args: Any) -> bool: # 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"}) +_EXPERT_LEAF_NAMES = frozenset({"linear_fc1", "linear_fc2", "gate_proj", "up_proj", "gate_up_proj", "down_proj"}) _ALL_MODULE_ALIASES = frozenset({"all", "all-linear", "all_linear"}) diff --git a/miles_plugins/models/inkling/lora.py b/miles_plugins/models/inkling/lora.py index 9eee20da7f8..9e1fbc54374 100644 --- a/miles_plugins/models/inkling/lora.py +++ b/miles_plugins/models/inkling/lora.py @@ -10,9 +10,19 @@ import torch.nn as nn import torch.nn.functional as F +from miles.utils.lora.hf_lora_targets import resolve_hf_lora_targets + logger = logging.getLogger(__name__) +def resolve_inkling_adapter_targets(hf_config, targets): + expected = resolve_hf_lora_targets(hf_config) + assert set(targets) == set( + expected + ), "Native Inkling LoRA requires its complete adapter layout; omit --target-modules and --exclude-modules" + return "all-linear" + + class InklingLoRAAdapter(nn.Module): """One module's LoRA params keyed for HF export; invisible to Megatron dist-checkpointing.""" diff --git a/miles_plugins/models/kimi_k3/lora.py b/miles_plugins/models/kimi_k3/lora.py index 54af99bb0f9..e4d4dde5f0f 100644 --- a/miles_plugins/models/kimi_k3/lora.py +++ b/miles_plugins/models/kimi_k3/lora.py @@ -14,17 +14,22 @@ # export collation below duplicate miles_plugins/models/inkling/lora.py _SUPPORTED_TARGET_SUFFIXES = { - "self_attention.o_proj", - "self_attention.q_a_proj", - "self_attention.kv_a_proj_with_mqa", - "mlp.linear_fc1", - "mlp.linear_fc2", - "mlp.experts.linear_fc1", - "mlp.experts.linear_fc2", + "self_attn.o_proj", + "self_attn.q_a_proj", + "self_attn.kv_a_proj_with_mqa", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + "block_sparse_moe.shared_experts.gate_proj", + "block_sparse_moe.shared_experts.up_proj", + "block_sparse_moe.shared_experts.down_proj", + "block_sparse_moe.experts.*.w1", + "block_sparse_moe.experts.*.w2", + "block_sparse_moe.experts.*.w3", } # the routed-expert down-proj may be omitted: its EP-shared w2_lora_B dominates adapter growth (#1559) -_OPTIONAL_TARGET_SUFFIXES = {"mlp.experts.linear_fc2"} +_OPTIONAL_TARGET_SUFFIXES = {"block_sparse_moe.experts.*.w2"} class KimiK3LoRAAdapter(nn.Module): @@ -113,9 +118,8 @@ def _grouped_linear(inputs: torch.Tensor, weights: torch.Tensor, tokens_per_expe return torch.cat([F.linear(segment, weights[idx]) for idx, segment in enumerate(segments)], dim=0) -def _validate_targets(args) -> None: - targets = set(args.target_modules) - suffixes = {target.split("decoder.layers.*.", 1)[-1] for target in targets} +def resolve_kimi_k3_adapter_targets(targets, *, canonical, experts_shared_outer_loras): + suffixes = {target.removeprefix("language_model.model.layers.*.") for target in targets} unsupported = suffixes - _SUPPORTED_TARGET_SUFFIXES missing = _SUPPORTED_TARGET_SUFFIXES - suffixes - _OPTIONAL_TARGET_SUFFIXES if unsupported or missing: @@ -123,8 +127,11 @@ def _validate_targets(args) -> None: "Kimi K3 native LoRA currently requires the verified target set; " f"unsupported={sorted(unsupported)}, missing={sorted(missing)}" ) - if not args.experts_shared_outer_loras: + assert not canonical, "Kimi K3 native LoRA does not implement canonical_lora" + if not experts_shared_outer_loras: raise NotImplementedError("Kimi K3 native LoRA currently requires --experts-shared-outer-loras") + # Shared-outer export stores the expert dimension in each tensor, not in its name. + return [target.replace(".experts.*.", ".experts.") for target in targets] def _enable_full_recompute_input_grads(model) -> None: @@ -425,7 +432,6 @@ def apply_kimi_k3_lora(model, args): from .layers import KimiK3Attention - _validate_targets(args) rank = int(args.lora_rank) if rank <= 0: raise ValueError("apply_kimi_k3_lora requires --lora-rank > 0") @@ -451,7 +457,7 @@ def apply_kimi_k3_lora(model, args): layer_idx, scale, dropout, - include_fc2=any(target.endswith("mlp.experts.linear_fc2") for target in args.target_modules), + include_fc2=any(target.endswith("block_sparse_moe.experts.*.w2") for target in args.hf_lora_targets), ) if layer.mlp.shared_experts is None: raise RuntimeError(f"Kimi K3 MoE layer {layer_idx} is missing shared experts") diff --git a/scripts/run_glm5_1_744b_a40b_lora.py b/scripts/run_glm5_1_744b_a40b_lora.py index be5846261a9..6f8979e077b 100644 --- a/scripts/run_glm5_1_744b_a40b_lora.py +++ b/scripts/run_glm5_1_744b_a40b_lora.py @@ -54,11 +54,6 @@ "GLM-5.1-6layer": "glm5.1-744B-A40B_6layer_lora", } -# Standard attn + MLA + MLP/MoE, EXCLUDING the DSA indexer (wq_b/wk/weights_proj). -_DEFAULT_TARGET_MODULES = ( - "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" -) - @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -89,7 +84,7 @@ class ScriptArgs(U.ExecuteTrainConfig): lora_rank: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.0 - target_modules: str = _DEFAULT_TARGET_MODULES + target_modules: str = "all-linear" # required for true on-policy under colocate (OFF -> KL ~1.0 vs ~1e-4); opt out only # when host RAM cannot take the ~372 GB/node mirror on the full model lora_base_cpu_backup: bool = True @@ -179,11 +174,8 @@ def _train(args: ScriptArgs): # the full rollout config applies to the toys too (same glm_moe_dsa serving path) _is_full = True - _tm = args.target_modules # KEEP_MOE_LORA=0 drops the expert projections (attention-only LoRA) _keep_moe_lora = os.environ.get("KEEP_MOE_LORA", "1") != "0" - if _is_full and not _keep_moe_lora: - _tm = ",".join(m for m in _tm.split(",") if m.strip() not in ("gate_proj", "up_proj", "down_proj")) # the MOE_LORA_LAYERS subset feature is disabled; warn so it is not silently ignored _moe_lora_layers = os.environ.get("MOE_LORA_LAYERS", "").strip() if _moe_lora_layers: @@ -191,7 +183,12 @@ def _train(args: ScriptArgs): f"[run_glm5_1_744b_a40b_lora] WARNING: MOE_LORA_LAYERS={_moe_lora_layers} is SET but the subset-rewrite " "feature is DISABLED (commented out for debugging) -> MoE-expert LoRA stays on ALL layers." ) - lora_args = f'--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout {args.lora_dropout} --target-modules "{_tm}" ' + lora_args = ( + f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout {args.lora_dropout} " + f'--target-modules "{args.target_modules}" ' + ) + if not _keep_moe_lora: + lora_args += "--exclude-modules gate_proj,up_proj,gate_up_proj,down_proj " if _keep_moe_lora and args.experts_shared_outer_loras: lora_args += "--experts-shared-outer-loras " if _is_full: diff --git a/scripts/run_glm5_2_744b_a40b_lora.py b/scripts/run_glm5_2_744b_a40b_lora.py index 69a13a8f63b..4df6a87db24 100644 --- a/scripts/run_glm5_2_744b_a40b_lora.py +++ b/scripts/run_glm5_2_744b_a40b_lora.py @@ -63,11 +63,6 @@ "GLM-5.2_5layer": "glm5.2-744B-A40B_5layer_lora", } -# Standard attn + MLA + MLP/MoE, EXCLUDING the DSA indexer (wq_b/wk/weights_proj). -_DEFAULT_TARGET_MODULES = ( - "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" -) - @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -98,7 +93,7 @@ class ScriptArgs(U.ExecuteTrainConfig): lora_rank: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.0 - target_modules: str = _DEFAULT_TARGET_MODULES + target_modules: str = "all-linear" # required for true on-policy under colocate (OFF -> KL ~1.0 vs ~1e-4); opt out only # when host RAM cannot take the ~372 GB/node mirror on the full model lora_base_cpu_backup: bool = True @@ -192,11 +187,8 @@ def _train(args: ScriptArgs): # the full rollout config applies to the toys too (same glm_moe_dsa serving path) _is_full = True - _tm = args.target_modules # KEEP_MOE_LORA=0 drops the expert projections (attention-only LoRA) _keep_moe_lora = os.environ.get("KEEP_MOE_LORA", "1") != "0" - if _is_full and not _keep_moe_lora: - _tm = ",".join(m for m in _tm.split(",") if m.strip() not in ("gate_proj", "up_proj", "down_proj")) # the MOE_LORA_LAYERS subset feature is disabled; warn so it is not silently ignored _moe_lora_layers = os.environ.get("MOE_LORA_LAYERS", "").strip() if _moe_lora_layers: @@ -204,7 +196,12 @@ def _train(args: ScriptArgs): f"[run_glm5_2_744b_a40b_lora] WARNING: MOE_LORA_LAYERS={_moe_lora_layers} is SET but the subset-rewrite " "feature is DISABLED (commented out for debugging) -> MoE-expert LoRA stays on ALL layers." ) - lora_args = f'--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout {args.lora_dropout} --target-modules "{_tm}" ' + lora_args = ( + f"--lora-rank {args.lora_rank} --lora-alpha {args.lora_alpha} --lora-dropout {args.lora_dropout} " + f'--target-modules "{args.target_modules}" ' + ) + if not _keep_moe_lora: + lora_args += "--exclude-modules gate_proj,up_proj,gate_up_proj,down_proj " if _keep_moe_lora and args.experts_shared_outer_loras: lora_args += "--experts-shared-outer-loras " if _is_full: diff --git a/scripts/run_kimi_k3.py b/scripts/run_kimi_k3.py index cb48caa53fd..3f59a1db69e 100644 --- a/scripts/run_kimi_k3.py +++ b/scripts/run_kimi_k3.py @@ -34,19 +34,6 @@ _NUM_ATTENTION_HEADS = 96 _VALIDATED_FULL_MODEL_GPUS = 64 -_LAYERS = "decoder.layers.*" -_DEFAULT_TARGET_MODULES = ",".join( - [ - f"{_LAYERS}.self_attention.o_proj", - f"{_LAYERS}.self_attention.q_a_proj", - f"{_LAYERS}.self_attention.kv_a_proj_with_mqa", - f"{_LAYERS}.mlp.linear_fc1", - f"{_LAYERS}.mlp.linear_fc2", - f"{_LAYERS}.mlp.experts.linear_fc1", - f"{_LAYERS}.mlp.experts.linear_fc2", - ] -) - @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -82,7 +69,7 @@ class ScriptArgs(U.ExecuteTrainConfig): lora_rank: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.0 - target_modules: str = _DEFAULT_TARGET_MODULES + target_modules: str = "all-linear" experts_shared_outer_loras: bool = True reward_model: Literal["deterministic_random", "deepscaler", "math"] | None = None diff --git a/scripts/run_qwen3_5_35b_a3b_lora.py b/scripts/run_qwen3_5_35b_a3b_lora.py index bd34d217beb..5dc062047bc 100644 --- a/scripts/run_qwen3_5_35b_a3b_lora.py +++ b/scripts/run_qwen3_5_35b_a3b_lora.py @@ -47,21 +47,6 @@ "Qwen3.6-35B-A3B": "qwen3.6-35B-A3B_lora", } -# Anchored below decoder.layers: keeps LoRA off the MTP block and the vision tower. -_LAYERS = "language_model.decoder.layers.*" -_DEFAULT_TARGET_MODULES = ",".join( - [ - f"{_LAYERS}.self_attention.linear_qkv", - f"{_LAYERS}.self_attention.linear_proj", - f"{_LAYERS}.mlp.experts.linear_fc1", - f"{_LAYERS}.mlp.experts.linear_fc2", - f"{_LAYERS}.mlp.shared_experts.linear_fc1", - f"{_LAYERS}.mlp.shared_experts.linear_fc2", - f"{_LAYERS}.self_attention.in_proj", - f"{_LAYERS}.self_attention.out_proj", - ] -) - @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -85,7 +70,7 @@ class ScriptArgs(U.ExecuteTrainConfig): lora_rank: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.0 - target_modules: str = _DEFAULT_TARGET_MODULES + target_modules: str = "all-linear" # required for true on-policy under colocate (OFF -> KL ~1.0 vs ~1e-3) lora_base_cpu_backup: bool = True # MoE-expert LoRA layout: shared-outer when True, per-expert when False diff --git a/serve_tinker.py b/serve_tinker.py index f9a0f1e5edc..a2161c38148 100644 --- a/serve_tinker.py +++ b/serve_tinker.py @@ -27,7 +27,7 @@ async def serve(args): assert args.load == args.hf_checkpoint, "Tinker trainers and engines must load the same frozen HF base" checkpoint_root = args.tinker_checkpoint_root or (args.save and f"{args.save}/tinker") assert checkpoint_root, "set --tinker-checkpoint-root (or --save to derive /tinker)" - hf_config = load_hf_config(args.hf_checkpoint) + hf_config = load_hf_config(args.hf_checkpoint).get_text_config() max_tokens_per_datum = hf_config.max_position_embeddings if args.max_tokens_per_gpu is not None: # The trainer pads each packed microbatch to this multiple. @@ -63,9 +63,9 @@ async def serve(args): max_tokens_per_datum=max_tokens_per_datum, lora_alpha=args.lora_alpha, max_lora_rank=args.lora_rank, - trains_attn=args.tinker_train_attn, - trains_mlp=args.tinker_train_mlp, - trains_unembed=args.tinker_train_unembed, + trains_attn="attn" in args.tinker_lora_groups, + trains_mlp="mlp" in args.tinker_lora_groups, + trains_unembed="unembed" in args.tinker_lora_groups, ) router_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" actor_world_size = args.actor_num_nodes * args.actor_num_gpus_per_node diff --git a/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py b/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py index 419952d46ff..a4b4eaea17a 100644 --- a/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py +++ b/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py @@ -43,7 +43,7 @@ def execute(shared_outer: bool, virtual_experts: bool): "--lora-rank 32 " "--lora-alpha 32 " "--lora-dropout 0.0 " - '--target-modules "gate_proj,up_proj,down_proj" ' + '--target-modules "gate_up_proj,down_proj" ' "--sglang-lora-backend triton " f"{'--experts-shared-outer-loras ' if shared_outer else ''}" f"{'' if virtual_experts else '--no-sglang-lora-use-virtual-experts '}" diff --git a/tests/fast/backends/megatron_utils/lora/test_target_modules.py b/tests/fast/backends/megatron_utils/lora/test_target_modules.py new file mode 100644 index 00000000000..aa135185315 --- /dev/null +++ b/tests/fast/backends/megatron_utils/lora/test_target_modules.py @@ -0,0 +1,122 @@ +from types import SimpleNamespace + +import pytest + +from miles.backends.megatron_utils.lora.target_modules import resolve_megatron_lora_targets +from miles.utils.hf_utils.weight_mapping import HfWeightMapping +from miles.utils.lora.hf_lora_targets import resolve_hf_lora_targets + + +class _Mapping(SimpleNamespace): + def resolve(self, captures): + def expand(name): + for value in captures: + name = name.replace("*", value, 1) + return name + + hf = [self.hf_param] if isinstance(self.hf_param, str) else self.hf_param.values() + return _mapping(expand(self.megatron_param), *(expand(name) for name in hf)) + + +def _mapping(megatron, *hf): + return _Mapping(megatron_param=megatron, hf_param=hf[0] if len(hf) == 1 else dict(enumerate(hf))) + + +def _resolve(targets, mappings, parameter_names, *, canonical=False, hf_mapping=None): + return resolve_megatron_lora_targets( + targets, + mappings, + parameter_names=set(parameter_names), + hf_mapping=hf_mapping or HfWeightMapping(frozenset()), + canonical=canonical, + ) + + +_QKV = _mapping( + "decoder.layers.*.self_attention.linear_qkv.weight", + *(f"model.layers.*.self_attn.{p}_proj.weight" for p in ("q", "k", "v")), +) + + +@pytest.mark.parametrize( + "mtp_source", + ["mtp.layers.0.self_attn.o_proj.weight", "model.layers.1.self_attn.o_proj.weight"], + ids=["separate-namespace", "appended-layer"], +) +def test_scoped_attention_excludes_mtp(mtp_source): + targets = resolve_hf_lora_targets({"model_type": "qwen3"}, target_modules=["attn"]) + output = _mapping("decoder.layers.*.self_attention.linear_proj.weight", "model.layers.*.self_attn.o_proj.weight") + mtp = _mapping("mtp.layers.*.self_attention.linear_proj.weight", mtp_source) + megatron_parameters = [ + "decoder.layers.0.self_attention.linear_qkv.weight", + "decoder.layers.0.self_attention.linear_proj.weight", + "mtp.layers.0.self_attention.linear_proj.weight", + ] + hf_mapping = HfWeightMapping(frozenset(f"model.layers.0.self_attn.{p}_proj.weight" for p in ("q", "k", "v", "o"))) + selected = _resolve(targets, [_QKV, output, mtp], megatron_parameters, hf_mapping=hf_mapping) + assert set(selected) == { + "decoder.layers.*.self_attention.linear_qkv", + "decoder.layers.*.self_attention.linear_proj", + } + + +def test_fused_selection_cannot_silently_expand(): + targets = ["model.layers.*.self_attn.q_proj"] + with pytest.raises(AssertionError, match="requires all HF targets"): + _resolve(targets, [_QKV], ["decoder.layers.0.self_attention.linear_qkv.weight"]) + adapter_targets = _resolve(targets, [_QKV], ["decoder.layers.0.self_attention.linear_qkv.weight"], canonical=True) + assert adapter_targets == ["decoder.layers.*.self_attention.linear_q"] + + +@pytest.mark.parametrize("grouped", [True, False], ids=["grouped", "sequential"]) +def test_expert_representations_are_alternatives(grouped): + target = "model.layers.*.mlp.experts.*.down_proj" + mappings = [ + _mapping("decoder.layers.*.mlp.experts.linear_fc2.weight*", target + ".weight"), + _mapping("decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", target + ".weight"), + ] + module = "decoder.layers.0.mlp.experts." + ("linear_fc2" if grouped else "local_experts.0.linear_fc2") + parameter = module + (".weight0" if grouped else ".weight") + selected = _resolve([target], mappings, [parameter]) + assert len(selected) == 1 + assert ("local_experts" in next(iter(selected))) != grouped + with pytest.raises(AssertionError, match="no Megatron modules"): + _resolve([target], mappings, ["decoder.layers.0.mlp.linear_fc2.weight"]) + + +def test_missing_hf_mapping_is_not_a_megatron_passthrough(): + with pytest.raises(AssertionError, match="no Megatron modules"): + _resolve( + ["model.layers.*.self_attn.unknown_proj"], [_QKV], ["decoder.layers.0.self_attention.linear_qkv.weight"] + ) + + +def test_one_to_one_mapping_keeps_bridge_module_name(): + target = "model.layers.*.self_attn.o_proj" + mapping = _mapping("decoder.layers.*.self_attention.output_projection.weight", target + ".weight") + adapter_targets = _resolve( + [target], [mapping], ["decoder.layers.0.self_attention.output_projection.weight"], canonical=True + ) + assert adapter_targets == ["decoder.layers.*.self_attention.output_projection"] + + +def test_absent_fused_alternative_does_not_reject_selection(): + mappings = [ + _mapping( + "decoder.layers.*.self_attention.fused_qkv.weight", + *(f"model.layers.*.self_attn.{p}_proj.weight" for p in ("q", "k", "v")), + ), + _mapping("decoder.layers.*.self_attention.linear_q.weight", "model.layers.*.self_attn.q_proj.weight"), + ] + selected = _resolve(["q_proj"], mappings, ["decoder.layers.0.self_attention.linear_q.weight"]) + assert set(selected) == {"decoder.layers.*.self_attention.linear_q"} + + +def test_canonical_selection_cannot_expand_across_layers(): + with pytest.raises(AssertionError, match="different projections"): + _resolve( + ["model.layers.0.self_attn.q_proj", "model.layers.1.self_attn.k_proj"], + [_QKV], + [f"decoder.layers.{layer}.self_attention.linear_qkv.weight" for layer in range(2)], + canonical=True, + ) diff --git a/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py b/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py index e8359957816..fe05f475168 100644 --- a/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py +++ b/tests/fast/backends/megatron_utils/test_bridge_value_model_hook.py @@ -71,6 +71,7 @@ def stub(name: str, attrs: dict | None = None, is_package: bool = False): try: stub("megatron", is_package=True) stub("megatron.core", is_package=True) + stub("megatron.core.tensor_parallel", {"ColumnParallelLinear": type("ColumnParallelLinear", (), {})}) stub("megatron.core.utils", {"get_attr_wrapped_model": lambda *a, **k: None}) install("megatron.core.parallel_state", parallel_state) diff --git a/tests/fast/backends/megatron_utils/test_lora_model_branches.py b/tests/fast/backends/megatron_utils/test_lora_model_branches.py index f486871a3e0..313256fc533 100644 --- a/tests/fast/backends/megatron_utils/test_lora_model_branches.py +++ b/tests/fast/backends/megatron_utils/test_lora_model_branches.py @@ -7,6 +7,8 @@ from argparse import Namespace from unittest.mock import MagicMock, patch +import pytest + # --------------------------------------------------------------------------- # _ensure_model_list # --------------------------------------------------------------------------- @@ -67,6 +69,7 @@ def _make_args(self, lora_rank=32, role="actor", mode="bridge"): return Namespace( lora_rank=lora_rank, lora_adapter_path=None, + custom_model_provider_path=None, megatron_to_hf_mode=mode, model_name=None, moe_use_upcycling=False, @@ -151,7 +154,7 @@ def test_non_lora_skips_lora_setup(self, mock_lora_setup, mock_provider, mock_ge @patch(f"{_MODEL_MODULE}.get_megatron_optimizer") @patch(f"{_MODEL_MODULE}.get_model") @patch(f"{_MODEL_MODULE}._setup_lora_model_via_bridge") - def test_lora_raw_mode_skips_bridge(self, mock_lora_setup, mock_get_model, mock_opt, mock_sched): + def test_non_inkling_lora_raw_mode_is_rejected(self, mock_lora_setup, mock_get_model, mock_opt, mock_sched): from miles.backends.megatron_utils.model import setup_model_and_optimizer mock_get_model.return_value = [MagicMock()] @@ -159,10 +162,11 @@ def test_lora_raw_mode_skips_bridge(self, mock_lora_setup, mock_get_model, mock_ mock_sched.return_value = MagicMock() args = self._make_args(lora_rank=32, role="actor", mode="raw") - setup_model_and_optimizer(args, role="actor") + with pytest.raises(AssertionError, match="Native LoRA injection is only implemented for Inkling"): + setup_model_and_optimizer(args, role="actor") mock_lora_setup.assert_not_called() - mock_get_model.assert_called_once() + mock_get_model.assert_not_called() # --------------------------------------------------------------------------- diff --git a/tests/fast/backends/megatron_utils/test_lora_utils.py b/tests/fast/backends/megatron_utils/test_lora_utils.py index c7a1fa9975e..3edba786278 100644 --- a/tests/fast/backends/megatron_utils/test_lora_utils.py +++ b/tests/fast/backends/megatron_utils/test_lora_utils.py @@ -1,200 +1,23 @@ -"""Unit tests for miles.backends.megatron_utils.lora.utils. - -Tests cover module name conversion, LoRA detection helpers, parameter identification, -exclude-module parsing, and LoRA sync config building — all without GPU. -""" +"""LoRA detection, adapter parameters, and training checkpoint state.""" import sys import types from argparse import Namespace from types import SimpleNamespace -from unittest.mock import MagicMock, Mock +from unittest.mock import Mock import pytest import torch import miles.backends.megatron_utils.lora.utils as lora_utils from miles.backends.megatron_utils.lora.utils import ( - _get_lora_class_name, _is_adapter_param_name, - build_lora_sync_config, - convert_target_modules_to_hf, - convert_target_modules_to_megatron, is_lora_enabled, load_lora_adapter, - parse_exclude_modules, save_lora_checkpoint, ) from miles.utils.lora.utils import LORA_ADAPTER_NAME, is_lora_weight_name -# --------------------------------------------------------------------------- -# _get_lora_class_name -# --------------------------------------------------------------------------- - - -class TestGetLoraClassName: - def test_none_returns_canonical(self): - assert _get_lora_class_name(None) == "CanonicalLoRA" - - def test_type_returns_class_name(self): - class FakeLoRA: - pass - - assert _get_lora_class_name(FakeLoRA) == "FakeLoRA" - - def test_instance_returns_class_name(self): - class FakeLoRA: - pass - - assert _get_lora_class_name(FakeLoRA()) == "FakeLoRA" - - -# --------------------------------------------------------------------------- -# convert_target_modules_to_megatron -# --------------------------------------------------------------------------- - - -def _make_lora_type(name: str): - """Helper to create a mock lora_type whose class name matches *name*.""" - mock = MagicMock() - type(mock).__name__ = name - return mock - - -class TestConvertTargetModulesToMegatron: - def test_gdn_hf_names_collapse_to_fused_in_proj(self): - lora = _make_lora_type("LoRA") - result = convert_target_modules_to_megatron(["in_proj_qkvz", "in_proj_ba", "out_proj"], lora_type=lora) - assert result == ["in_proj", "out_proj"] - - # --- "all-linear" variants ------------------------------------------------ - - @pytest.mark.parametrize("shorthand", ["all", "all-linear", "all_linear"]) - def test_all_linear_string_canonical(self, shorthand): - result = convert_target_modules_to_megatron(shorthand, lora_type=None) - assert result == [ - "linear_q", - "linear_k", - "linear_v", - "linear_proj", - "linear_fc1_up", - "linear_fc1_gate", - "linear_fc2", - ] - - @pytest.mark.parametrize("shorthand", ["all", "all-linear", "all_linear"]) - def test_all_linear_string_standard(self, shorthand): - lora_type = _make_lora_type("LoRA") - result = convert_target_modules_to_megatron(shorthand, lora_type=lora_type) - assert result == ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"] - - @pytest.mark.parametrize("shorthand", ["all", "all-linear", "all_linear"]) - def test_all_linear_single_element_list(self, shorthand): - result = convert_target_modules_to_megatron([shorthand], lora_type=None) - assert len(result) == 7 # CanonicalLoRA has 7 modules - - # --- HF -> Megatron conversion (standard LoRA) ---------------------------- - - def test_hf_to_megatron_standard_dedup(self): - lora = _make_lora_type("LoRA") - result = convert_target_modules_to_megatron(["q_proj", "k_proj", "v_proj"], lora_type=lora) - assert result == ["linear_qkv"] - - def test_hf_to_megatron_standard_all_modules(self): - lora = _make_lora_type("LoRA") - modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - result = convert_target_modules_to_megatron(modules, lora_type=lora) - assert result == ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"] - - # --- HF -> Megatron conversion (CanonicalLoRA) ---------------------------- - - def test_hf_to_megatron_canonical_split(self): - result = convert_target_modules_to_megatron(["q_proj", "k_proj", "v_proj"], lora_type=None) - assert result == ["linear_q", "linear_k", "linear_v"] - - def test_hf_to_megatron_canonical_gate_up(self): - result = convert_target_modules_to_megatron(["gate_proj", "up_proj"], lora_type=None) - assert result == ["linear_fc1_gate", "linear_fc1_up"] - - # --- Already in Megatron format ------------------------------------------- - - def test_megatron_format_passthrough(self): - modules = ["linear_qkv", "linear_proj"] - result = convert_target_modules_to_megatron(modules, lora_type=None) - assert result == modules - - def test_megatron_format_passthrough_canonical(self): - modules = ["linear_q", "linear_k", "linear_v"] - result = convert_target_modules_to_megatron(modules, lora_type=None) - assert result == modules - - # --- Single string input -------------------------------------------------- - - def test_single_hf_string_input(self): - lora = _make_lora_type("LoRA") - result = convert_target_modules_to_megatron("o_proj", lora_type=lora) - assert result == ["linear_proj"] - - -# --------------------------------------------------------------------------- -# convert_target_modules_to_hf -# --------------------------------------------------------------------------- - - -class TestConvertTargetModulesToHf: - def test_standard_linear_qkv(self): - assert convert_target_modules_to_hf(["linear_qkv"]) == ["q_proj", "k_proj", "v_proj"] - - def test_standard_linear_proj(self): - assert convert_target_modules_to_hf(["linear_proj"]) == ["o_proj"] - - def test_standard_linear_fc1(self): - assert convert_target_modules_to_hf(["linear_fc1"]) == ["gate_proj", "up_proj"] - - def test_standard_linear_fc2(self): - assert convert_target_modules_to_hf(["linear_fc2"]) == ["down_proj"] - - def test_gdn_in_proj_expands_to_sglang_modules(self): - assert convert_target_modules_to_hf(["in_proj"]) == ["in_proj_qkvz", "in_proj_ba"] - - @pytest.mark.parametrize( - "module,expected", - [ - ("out_proj", ["out_proj"]), # same-name passthrough - ("language_model.decoder.layers.*.self_attention.out_proj", ["out_proj"]), # wildcard path - ("language_model.decoder.layers.0.self_attention.out_proj", ["out_proj"]), # dotted path, passthrough - ( - "language_model.decoder.layers.0.self_attention.linear_qkv", # dotted path, table-mapped - ["q_proj", "k_proj", "v_proj"], - ), - ], - ) - def test_paths_reduce_to_leaf_before_mapping(self, module, expected): - assert convert_target_modules_to_hf([module]) == expected - - def test_canonical_split_modules(self): - result = convert_target_modules_to_hf(["linear_q", "linear_k", "linear_v"]) - assert result == ["q_proj", "k_proj", "v_proj"] - - def test_canonical_fc1_gate_up(self): - result = convert_target_modules_to_hf(["linear_fc1_gate", "linear_fc1_up"]) - assert result == ["gate_proj", "up_proj"] - - def test_unknown_module_passthrough(self): - assert convert_target_modules_to_hf(["some_custom_module"]) == ["some_custom_module"] - - def test_roundtrip_canonical_all_linear(self): - megatron = convert_target_modules_to_megatron("all-linear", lora_type=None) - hf = convert_target_modules_to_hf(megatron) - assert set(hf) == {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"} - - def test_roundtrip_standard_all_linear(self): - lora = _make_lora_type("LoRA") - megatron = convert_target_modules_to_megatron("all-linear", lora_type=lora) - hf = convert_target_modules_to_hf(megatron) - assert set(hf) == {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"} - - # --------------------------------------------------------------------------- # is_lora_enabled # --------------------------------------------------------------------------- @@ -277,83 +100,6 @@ def test_negative(self, name): assert _is_adapter_param_name(name) is False -# --------------------------------------------------------------------------- -# parse_exclude_modules -# --------------------------------------------------------------------------- - - -class TestParseExcludeModules: - def test_none(self): - args = Namespace(exclude_modules=None) - assert parse_exclude_modules(args) == [] - - def test_single_module_string(self): - args = Namespace(exclude_modules="o_proj") - result = parse_exclude_modules(args, lora_type=_make_lora_type("LoRA")) - assert result == ["linear_proj"] - - def test_comma_separated(self): - args = Namespace(exclude_modules="o_proj, down_proj") - result = parse_exclude_modules(args, lora_type=_make_lora_type("LoRA")) - assert set(result) == {"linear_proj", "linear_fc2"} - - def test_list_input(self): - args = Namespace(exclude_modules=["o_proj", "down_proj"]) - result = parse_exclude_modules(args, lora_type=_make_lora_type("LoRA")) - assert set(result) == {"linear_proj", "linear_fc2"} - - def test_missing_attr(self): - args = Namespace() - assert parse_exclude_modules(args) == [] - - -# --------------------------------------------------------------------------- -# build_lora_sync_config -# --------------------------------------------------------------------------- - - -class TestBuildLoraSyncConfig: - def test_basic_config(self): - args = Namespace( - lora_rank=32, - lora_alpha=32, - lora_dropout=0.0, - target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"], - ) - config = build_lora_sync_config(args) - assert config["peft_type"] == "LORA" - assert config["r"] == 32 - assert config["lora_alpha"] == 32 - assert config["lora_dropout"] == 0.0 - assert config["bias"] == "none" - assert config["task_type"] == "CAUSAL_LM" - assert set(config["target_modules"]) == { - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - } - - def test_no_target_modules_uses_default(self): - args = Namespace(lora_rank=16, lora_alpha=16, lora_dropout=0.0, target_modules=None) - config = build_lora_sync_config(args) - assert len(config["target_modules"]) == 7 - - def test_canonical_target_modules(self): - args = Namespace( - lora_rank=8, - lora_alpha=8, - lora_dropout=0.1, - target_modules=["linear_q", "linear_k"], - ) - config = build_lora_sync_config(args) - assert config["target_modules"] == ["q_proj", "k_proj"] - assert config["r"] == 8 - - # --------------------------------------------------------------------------- # LORA_ADAPTER_NAME constant # --------------------------------------------------------------------------- @@ -363,29 +109,6 @@ def test_lora_adapter_name_constant(): assert LORA_ADAPTER_NAME == "miles_lora" -class TestBuildLoraSyncConfigUnderMultiLora: - @staticmethod - def _args(**overrides): - return Namespace( - **{ - "lora_rank": 8, - "lora_alpha": 8, - "lora_dropout": 0.0, - "target_modules": ["linear_qkv"], - "multi_lora": False, - **overrides, - } - ) - - def test_a_single_adapter_on_an_inkling_checkpoint_publishes_the_shorthand(self, monkeypatch): - """The engine was launched to auto-detect its targets, so the adapter must say the same.""" - monkeypatch.setattr( - "miles.backends.megatron_utils.lora.utils.sglang_lora_target_all_sentinel", lambda _a: True - ) - - assert build_lora_sync_config(self._args())["target_modules"] == "all-linear" - - class _AdapterModel(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/fast/backends/sglang_utils/conftest.py b/tests/fast/backends/sglang_utils/conftest.py index 4e4ae7c2a3e..6c11102f525 100644 --- a/tests/fast/backends/sglang_utils/conftest.py +++ b/tests/fast/backends/sglang_utils/conftest.py @@ -49,6 +49,7 @@ def make_engine_args(**overrides: Any) -> Namespace: use_rollout_indexer_replay=False, fp16=False, lora_rank=0, + lora_adapter_targets=[f"model.layers.*.self_attn.{projection}_proj" for projection in ("q", "k", "v")], sglang_api_key=None, lora_adapter_path=None, debug_rollout_only=False, diff --git a/tests/fast/backends/sglang_utils/test_compute_server_args.py b/tests/fast/backends/sglang_utils/test_compute_server_args.py index 697d94a5290..b4fb679be3b 100644 --- a/tests/fast/backends/sglang_utils/test_compute_server_args.py +++ b/tests/fast/backends/sglang_utils/test_compute_server_args.py @@ -28,7 +28,7 @@ def make_args(**overrides: object) -> SimpleNamespace: debug_rollout_only=False, debug_skip_weight_update=False, multi_lora_n_adapters=1, - target_modules=["linear_qkv"], + lora_adapter_targets=[f"model.layers.*.self_attn.{projection}_proj" for projection in ("q", "k", "v")], ) defaults.update(overrides) return SimpleNamespace(**defaults) diff --git a/tests/fast/backends/sglang_utils/test_sglang_engine.py b/tests/fast/backends/sglang_utils/test_sglang_engine.py index 88e8ce2f60a..367b2859c02 100644 --- a/tests/fast/backends/sglang_utils/test_sglang_engine.py +++ b/tests/fast/backends/sglang_utils/test_sglang_engine.py @@ -10,6 +10,7 @@ from miles.backends.sglang_utils.server_args_utils import parse_server_args_argv from miles.backends.sglang_utils.sglang_engine import compute_engine_launch_cmd +from miles.utils.lora.utils import build_lora_config def _cmd( @@ -98,55 +99,26 @@ def test_the_command_carries_the_api_key_from_args(self): class TestLoraTargetModules: - @staticmethod - def _parsed_lora_targets(target_modules: list[str]): - args = make_engine_args(lora_rank=16, target_modules=target_modules) - return parse_server_args_argv(shlex.split(_cmd(args=args))[3:]).lora_target_modules - - def test_spellable_targets_are_named_one_by_one(self): - """Naming the exact modules keeps SGLang from allocating adapter buffers for the rest.""" - targets = self._parsed_lora_targets(["layers.*.self_attention.linear_qkv"]) - - assert sorted(targets) == ["k_proj", "q_proj", "v_proj"] - - def test_gdn_attention_targets_are_named_one_by_one(self): - """Qwen3.5 GDN adapters must reach the engine as the exact fused slices, not as the - auto-detecting shorthand that would cover every compatible module instead.""" - targets = self._parsed_lora_targets(["layers.*.self_attention.in_proj"]) - - assert sorted(targets) == ["in_proj_ba", "in_proj_qkvz"] - - def test_an_inkling_checkpoint_asks_sglang_to_discover_the_names(self, monkeypatch: pytest.MonkeyPatch): - """Inkling exposes module names the megatron-to-HF mapping cannot produce, so it is the - one family that hands SGLang the shorthand instead of naming its targets.""" - monkeypatch.setattr( - "miles.backends.sglang_utils.sglang_engine.sglang_lora_target_all_sentinel", lambda _args: True - ) - - targets = self._parsed_lora_targets(["layers.*.self_attention.linear_qkv"]) - - assert set(targets) == {"all"} - - def test_a_multi_lora_inkling_launch_still_names_its_targets(self, monkeypatch: pytest.MonkeyPatch): - """Several adapters share one slot budget here, so discovering every compatible module - sizes that budget off the base model instead of off what the adapters fill.""" - monkeypatch.setattr( - "miles.backends.sglang_utils.sglang_engine.sglang_lora_target_all_sentinel", lambda _args: True - ) + @pytest.mark.parametrize( + "adapter_targets", + [ + [f"model.layers.*.self_attn.{projection}_proj" for projection in ("q", "k", "v")], + [f"model.layers.*.linear_attn.in_proj_{projection}" for projection in ("qkv", "z", "b", "a")], + "all-linear", + ], + ids=["qkv", "gdn", "inkling"], + ) + @pytest.mark.parametrize("multi_lora", [False, True], ids=["single", "multi"]) + def test_adapter_selection_reaches_engine_and_sync_config(self, adapter_targets, multi_lora): args = make_engine_args( lora_rank=16, - target_modules=["layers.*.self_attention.linear_qkv"], - multi_lora=True, + lora_alpha=32, + lora_dropout=0.0, + lora_adapter_targets=adapter_targets, + hf_lora_targets=["model.language_model.layers.*.self_attn.q_proj"], + multi_lora=multi_lora, multi_lora_n_adapters=4, ) - targets = parse_server_args_argv(shlex.split(_cmd(args=args))[3:]).lora_target_modules - - assert sorted(targets) == ["k_proj", "q_proj", "v_proj"] - - def test_asking_for_every_module_is_still_honoured(self): - """SGLang accepts the shorthand as a target name, so a run that spelled it out itself - is not the substitution this refuses.""" - targets = self._parsed_lora_targets(["all"]) - - assert set(targets) == {"all"} + assert set(targets) == ({"all"} if adapter_targets == "all-linear" else set(adapter_targets)) + assert build_lora_config(args, target_modules=adapter_targets)["target_modules"] == adapter_targets diff --git a/tests/fast/miles_plugins/models/inkling/test_inkling_lora_targets.py b/tests/fast/miles_plugins/models/inkling/test_inkling_lora_targets.py new file mode 100644 index 00000000000..b3657e8bcac --- /dev/null +++ b/tests/fast/miles_plugins/models/inkling/test_inkling_lora_targets.py @@ -0,0 +1,52 @@ +from types import SimpleNamespace + +import pytest +import torch + +from miles.utils.lora.hf_lora_targets import resolve_hf_lora_targets +from miles.utils.lora.utils import get_adapter_target_modules +from miles_plugins.models.inkling.lora import _export_dense_mlp, _export_experts, resolve_inkling_adapter_targets + + +class _LocalGather: + def add(self, group, tensor, dim): + return SimpleNamespace(get=lambda: tensor) + + +@pytest.mark.parametrize("multimodal", [False, True], ids=["text", "multimodal"]) +def test_hf_mlp_selection_matches_existing_native_export(multimodal): + config = dict(model_type="inkling_text", mlp_layer_types=["dense", "sparse"], n_shared_experts=0) + if multimodal: + config = dict(model_type="inkling_mm_model", text_config=config) + hf_targets = resolve_hf_lora_targets(config) + assert resolve_inkling_adapter_targets(config, hf_targets) == "all-linear" + prefix = "model.language_model" if multimodal else "model" + assert f"{prefix}.layers.*.mlp.gate_proj" in hf_targets + assert f"{prefix}.layers.*.mlp.experts.gate_up_proj" in hf_targets + tensor = torch.ones(2, 2) + dense = SimpleNamespace( + hf_prefix="language_model.layers.0.mlp.", + load_meta={"i_loc": 1}, + fc1_A=tensor, + fc1_B=tensor, + fc2_A=tensor, + fc2_B=tensor, + ) + experts = SimpleNamespace( + hf_prefix="language_model.layers.1.mlp.experts.", + **{f"w{projection}_{factor}": tensor for projection in (1, 2, 3) for factor in ("A", "B")}, + ) + plan = _export_dense_mlp(dense, _LocalGather()) + _export_experts(experts, _LocalGather()) + weights = {name: value() if callable(value) else value for name, value in plan} + assert "language_model.layers.0.mlp.gate_up_proj" in get_adapter_target_modules(weights) + assert weights["language_model.layers.0.mlp.gate_up_proj.lora_A.weight"] is tensor + assert torch.equal(weights["language_model.layers.0.mlp.gate_up_proj.lora_B.weight"], tensor) + with pytest.raises(AssertionError, match="complete adapter layout"): + resolve_inkling_adapter_targets(config, [target for target in hf_targets if not target.endswith(".up_proj")]) + + +def test_legacy_config_selects_the_same_native_adapters(): + legacy = dict(model_type="inkling_model", dense_mlp_idx=1, num_hidden_layers=2, n_shared_experts=1) + native = dict(model_type="inkling_text", mlp_layer_types=["dense", "sparse"], n_shared_experts=1) + assert resolve_hf_lora_targets(legacy) == resolve_hf_lora_targets(native) + resolve_inkling_adapter_targets(legacy, resolve_hf_lora_targets(legacy)) diff --git a/tests/fast/models/test_kimi_k3_lora.py b/tests/fast/models/test_kimi_k3_lora.py index c19a52279c8..e0648b8eca5 100644 --- a/tests/fast/models/test_kimi_k3_lora.py +++ b/tests/fast/models/test_kimi_k3_lora.py @@ -5,11 +5,14 @@ import torch.nn as nn from torch.utils.checkpoint import checkpoint +from miles.utils.lora.hf_lora_targets import resolve_hf_lora_targets +from miles.utils.lora.utils import get_adapter_target_modules, matches_lora_target from miles_plugins.models.kimi_k3.lora import ( KimiK3LoRAAdapter, _enable_full_recompute_input_grads, _grouped_linear, export_kimi_k3_lora_hf_chunks, + resolve_kimi_k3_adapter_targets, ) @@ -108,7 +111,8 @@ def _model_with_adapters(*, include_shared_experts=True): return model -def test_native_export_is_chunked_by_adapter(monkeypatch): +@pytest.mark.parametrize("include_fc2", [True, False]) +def test_native_export_is_chunked_by_adapter(monkeypatch, include_fc2): """A wrong expert dim is a shape mismatch in SGLang's LoRA pool; a wrong HF name is silently dropped.""" from megatron.core import parallel_state @@ -116,7 +120,17 @@ def test_native_export_is_chunked_by_adapter(monkeypatch): monkeypatch.setattr(parallel_state, "get_expert_model_parallel_world_size", lambda: 1) model = _model_with_adapters() + targets = _default_hf_targets() + model.adapters[3].include_fc2 = include_fc2 + if not include_fc2: + targets = [target for target in targets if not target.endswith(".experts.*.w2")] + del model.adapters[3].w2_lora_A + del model.adapters[3].w2_lora_B + adapter_targets = resolve_kimi_k3_adapter_targets(targets, canonical=False, experts_shared_outer_loras=True) chunks = list(export_kimi_k3_lora_hf_chunks([model])) + exported = get_adapter_target_modules(name for chunk in chunks for name, _ in chunk) + assert all(any(matches_lora_target(name, target) for target in adapter_targets) for name in exported) + assert all(any(matches_lora_target(name, target) for name in exported) for target in adapter_targets) assert len(chunks) == 5 attention = dict(chunks[0]) @@ -131,8 +145,11 @@ def test_native_export_is_chunked_by_adapter(monkeypatch): prefix = "language_model.model.layers.4.block_sparse_moe.experts." assert experts[f"{prefix}w1.lora_A.weight"].shape == (1, 2, 8) assert experts[f"{prefix}w1.lora_B.weight"].shape == (3, 5, 2) - assert experts[f"{prefix}w2.lora_A.weight"].shape == (3, 2, 5) - assert experts[f"{prefix}w2.lora_B.weight"].shape == (1, 8, 2) + if include_fc2: + assert experts[f"{prefix}w2.lora_A.weight"].shape == (3, 2, 5) + assert experts[f"{prefix}w2.lora_B.weight"].shape == (1, 8, 2) + else: + assert not any(".w2." in name for name in experts) shared_experts = dict(chunks[4]) prefix = "language_model.model.layers.4.block_sparse_moe.shared_experts." @@ -144,6 +161,43 @@ def test_native_export_is_chunked_by_adapter(monkeypatch): assert shared_experts[f"{prefix}down_proj.lora_B.weight"].shape == (8, 2) +def _default_hf_targets(): + return resolve_hf_lora_targets( + dict( + model_type="kimi_k3", + text_config=dict( + num_hidden_layers=5, + first_k_dense_replace=4, + moe_layer_freq=1, + num_experts=3, + num_shared_experts=1, + ), + ) + ) + + +@pytest.mark.parametrize("case", ["missing", "extra", "canonical", "per-expert"]) +def test_native_target_resolution_rejects_unsupported_layout(case): + targets = _default_hf_targets() + if case == "missing": + targets = [target for target in targets if not target.endswith(".q_a_proj")] + elif case == "extra": + targets.append("language_model.model.layers.*.self_attn.q_b_proj") + error = AssertionError if case == "canonical" else NotImplementedError + message = { + "missing": "missing=", + "extra": "unsupported=", + "canonical": "canonical_lora", + "per-expert": "shared-outer", + }[case] + with pytest.raises(error, match=message): + resolve_kimi_k3_adapter_targets( + targets, + canonical=case == "canonical", + experts_shared_outer_loras=case != "per-expert", + ) + + def test_native_export_rejects_missing_shared_expert_adapter(monkeypatch): """A layer without its adapter must fail the export, not ship a partial adapter SGLang accepts.""" from megatron.core import parallel_state diff --git a/tests/fast/utils/hf_utils/test_weight_mapping.py b/tests/fast/utils/hf_utils/test_weight_mapping.py new file mode 100644 index 00000000000..8a3485688c5 --- /dev/null +++ b/tests/fast/utils/hf_utils/test_weight_mapping.py @@ -0,0 +1,37 @@ +import pytest +from transformers import Qwen3MoeConfig + +from miles.utils.hf_utils.weight_mapping import HfWeightMapping + + +@pytest.fixture(scope="module") +def hf_mapping(): + config = Qwen3MoeConfig( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=4, + num_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=4, + ) + return HfWeightMapping.from_config(config) + + +_PREFIX = "model.layers.0.mlp.experts" +_TARGETS = [f"{_PREFIX}.gate_up_proj", f"{_PREFIX}.down_proj"] + + +def _unpacked_names(): + return { + f"{_PREFIX}.{expert}.{projection}_proj.weight" for expert in range(2) for projection in ("gate", "up", "down") + } + + +def test_checkpoint_formats_resolve_to_the_same_hf_parameters(hf_mapping): + unpacked = _unpacked_names() + assert {hf_mapping.model_parameter(name) for name in unpacked} == set(_TARGETS) + assert {hf_mapping.model_parameter(name) for name in _TARGETS} == set(_TARGETS) diff --git a/tests/fast/utils/lora/test_hf_lora_model_targets.py b/tests/fast/utils/lora/test_hf_lora_model_targets.py new file mode 100644 index 00000000000..42277327ae8 --- /dev/null +++ b/tests/fast/utils/lora/test_hf_lora_model_targets.py @@ -0,0 +1,212 @@ +import argparse +import importlib +import shlex +from fnmatch import fnmatchcase + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, PretrainedConfig + +from miles.utils.hf_utils.weight_mapping import HfWeightMapping +from miles.utils.lora.hf_lora_targets import ( + _HF_LORA_MODELS, + exclude_hf_lora_targets, + get_hf_lora_targets, + parse_lora_targets, + resolve_hf_lora_targets, +) + +_NATIVE_MODELS = ( + "llama", + "qwen2", + "qwen3", + "qwen3_moe", + "qwen3_next", + "qwen3_5_text", + "qwen3_5_moe_text", + "qwen3_5", + "qwen3_5_moe", + "gpt_oss", + "deepseek_v2", + "deepseek_v3", + "glm4_moe", + "glm_moe_dsa", +) +_MULTIMODAL_MODELS = {"qwen3_5", "qwen3_5_moe"} + + +def _small_config(model_type, overrides): + text = dict( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=64, + ) + if model_type in {"qwen3_moe", "qwen3_next", "qwen3_5_moe_text", "qwen3_5_moe"}: + text.update(num_experts=2, num_experts_per_tok=1, moe_intermediate_size=16) + if model_type in {"qwen3_moe", "qwen3_next"}: + text.update(decoder_sparse_step=2, mlp_only_layers=[]) + if model_type in {"qwen3_next", "qwen3_5_moe_text", "qwen3_5_moe"}: + text.update(shared_expert_intermediate_size=16) + if model_type == "qwen3_next" or model_type.startswith("qwen3_5"): + text.update( + layer_types=["linear_attention", "full_attention"], + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_num_key_heads=2, + linear_num_value_heads=4, + ) + if model_type == "gpt_oss": + text.update(num_local_experts=2, num_experts_per_tok=1) + if model_type in {"deepseek_v2", "deepseek_v3", "glm4_moe", "glm_moe_dsa"}: + text.update( + first_k_dense_replace=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + moe_intermediate_size=16, + n_group=1, + topk_group=1, + ) + if model_type in {"deepseek_v2", "deepseek_v3", "glm_moe_dsa"}: + text.update(q_lora_rank=8, kv_lora_rank=8, qk_nope_head_dim=4, qk_rope_head_dim=4, v_head_dim=8) + if model_type == "glm_moe_dsa": + text.update(mlp_layer_types=["dense", "sparse"], index_n_heads=2, index_head_dim=8, index_topk=4) + text.update(overrides) + if model_type in _MULTIMODAL_MODELS: + return AutoConfig.for_model( + model_type, + text_config=text, + vision_config=dict( + depth=1, + hidden_size=32, + intermediate_size=64, + num_heads=4, + out_hidden_size=32, + num_position_embeddings=16, + patch_size=2, + temporal_patch_size=1, + spatial_merge_size=1, + ), + ) + return AutoConfig.for_model(model_type, **text) + + +def test_native_model_coverage(): + # These entries use custom code, a Miles alias, or a version newer than the pinned Transformers. + non_native = { + "deepseek_v32", + "kimi_k2", + "kimi_k25", + "kimi_k3", + "inkling_model", + "inkling_mm_model", + "inkling_text", + } + assert set(_NATIVE_MODELS) == set(_HF_LORA_MODELS) - non_native + + +@pytest.mark.parametrize( + "model_type,overrides", + [pytest.param(name, {}, id=name) for name in _NATIVE_MODELS] + + [ + pytest.param(name, {"q_lora_rank": None}, id=f"{name}-uncompressed-query") + for name in ("deepseek_v2", "deepseek_v3") + ] + + [ + pytest.param(name, {"first_k_dense_replace": count}, id=f"{name}-{kind}") + for name in ("deepseek_v2", "deepseek_v3", "glm4_moe") + for kind, count in (("dense", 2), ("moe", 0)) + ] + + [ + pytest.param(name, {"decoder_sparse_step": 1, "mlp_only_layers": layers}, id=f"{name}-{kind}") + for name in ("qwen3_moe", "qwen3_next") + for kind, layers in (("dense", [0, 1]), ("moe", [])) + ] + + [ + pytest.param(name, {"layer_types": [kind, kind]}, id=f"{name}-{kind}") + for name in ("qwen3_next", "qwen3_5_text", "qwen3_5_moe_text") + for kind in ("full_attention", "linear_attention") + ] + + [ + pytest.param("glm_moe_dsa", {"mlp_layer_types": [kind, kind]}, id=f"glm-dsa-{kind}") + for kind in ("dense", "sparse") + ], +) +def test_targets_match_native_hf_model(model_type, overrides): + config = _small_config(model_type, overrides) + model_cls = AutoModelForImageTextToText if model_type in _MULTIMODAL_MODELS else AutoModelForCausalLM + with torch.device("meta"): + model = model_cls.from_config(config, attn_implementation="eager") + + # Keep tied output heads; packed expert projections are parameters without a .weight suffix. + projections = { + name.removesuffix(".weight"): param + for name, param in model.named_parameters(remove_duplicate=False) + if param.ndim in (2, 3) + } + assert projections + assert all(param.is_meta for param in model.parameters()) + hf_mapping = HfWeightMapping.from_config(config) + assert hf_mapping.parameter_names == { + name for name, param in model.named_parameters(remove_duplicate=False) if param.ndim in (2, 3) + } + layout = get_hf_lora_targets(config.to_dict()) + for group in (layout.attention, layout.mlp, layout.unembed): + assert group + for target in group: + matches = [name for name in projections if fnmatchcase(name, target)] + assert matches, f"{model_type}: HF model has no projection matching {target!r}" + assert all(projections[name].numel() > 0 for name in matches) + assert all(".indexer." not in name for name in matches) + + defaults = resolve_hf_lora_targets(config.to_dict()) + assert set(defaults) == set(layout.attention + layout.mlp) + all_groups = resolve_hf_lora_targets(config.to_dict(), target_modules=["attn", "mlp", "unembed"]) + assert set(all_groups) == set(layout.attention + layout.mlp + layout.unembed) + + +@pytest.mark.parametrize("launcher", ["run_glm5_1_744b_a40b_lora", "run_glm5_2_744b_a40b_lora"]) +def test_glm_launcher_ablation_selects_only_attention(monkeypatch, launcher): + module = importlib.import_module(f"scripts.{launcher}") + train_commands = [] + monkeypatch.setenv("KEEP_MOE_LORA", "0") + monkeypatch.delenv("MOE_LORA_LAYERS", raising=False) + monkeypatch.setattr(module.U, "execute_train", lambda **kwargs: train_commands.append(kwargs["train_args"])) + module._train(module.ScriptArgs(enable_wandb=False)) + + parser = argparse.ArgumentParser() + parser.add_argument("--target-modules") + parser.add_argument("--exclude-modules") + assert len(train_commands) == 1 + args, _ = parser.parse_known_args(shlex.split(train_commands[0])) + config = _small_config("glm_moe_dsa", {}) + selected = exclude_hf_lora_targets( + resolve_hf_lora_targets(config.to_dict(), target_modules=parse_lora_targets(args.target_modules)), + parse_lora_targets(args.exclude_modules) or [], + ) + assert set(selected) == set(get_hf_lora_targets(config.to_dict()).attention) + selected_parameters = { + name + for name in HfWeightMapping.from_config(config).parameter_names + if any(fnmatchcase(name.removesuffix(".weight"), target) for target in selected) + } + assert selected_parameters + assert all(".self_attn." in name and ".indexer." not in name for name in selected_parameters) + + +def test_remote_config_with_native_class_name(): + config = _small_config("deepseek_v2", {}) + remote_config_class = type("DeepseekV2Config", (PretrainedConfig,), {"model_type": "deepseek_v2"}) + remote_fields = config.to_dict() + remote_fields.pop("head_dim") + remote_config = remote_config_class(**remote_fields) + + assert ( + HfWeightMapping.from_config(remote_config).parameter_names + == HfWeightMapping.from_config(config).parameter_names + ) diff --git a/tests/fast/utils/lora/test_hf_lora_targets.py b/tests/fast/utils/lora/test_hf_lora_targets.py new file mode 100644 index 00000000000..c125edee97c --- /dev/null +++ b/tests/fast/utils/lora/test_hf_lora_targets.py @@ -0,0 +1,80 @@ +"""LoRA selection uses model defaults and preserves explicit user choices.""" + +import pytest + +from miles.utils.lora.hf_lora_targets import ( + exclude_hf_lora_targets, + expand_packed_hf_lora_targets, + parse_lora_targets, + resolve_hf_lora_targets, +) + + +@pytest.mark.parametrize("selection", [None, ["all-linear"], ["attn", "mlp"], ["attn", "mlp", "attn"]]) +def test_missing_targets_and_all_linear_use_model_defaults(selection): + config = {"model_type": "qwen3"} + targets = resolve_hf_lora_targets(config, target_modules=selection) + assert targets == resolve_hf_lora_targets(config, target_modules=["all-linear"]) + assert "model.layers.*.self_attn.q_proj" in targets + assert "model.layers.*.mlp.gate_proj" in targets + assert "lm_head" not in targets + + +@pytest.mark.parametrize( + "config,value,expected", + [ + ({"model_type": "custom"}, "q_proj, k_proj", ["q_proj", "k_proj"]), + ({"model_type": "custom"}, ["q_proj", "k_proj"], ["q_proj", "k_proj"]), + ({"model_type": "qwen3"}, "unembed,lm_head", ["lm_head"]), + ( + {"model_type": "qwen3"}, + "attn, lm_head", + [f"model.layers.*.self_attn.{name}_proj" for name in ("q", "k", "v", "o")] + ["lm_head"], + ), + ( + {"model_type": "gpt_oss"}, + "mlp", + ["model.layers.*.mlp.experts.gate_up_proj", "model.layers.*.mlp.experts.down_proj"], + ), + ({"model_type": "custom"}, "model.mlp", ["model.mlp"]), + ], +) +def test_explicit_targets_preserve_user_choices(config, value, expected): + assert resolve_hf_lora_targets(config, target_modules=parse_lora_targets(value)) == expected + + +def test_exclude_leaf_applies_to_scoped_model_defaults(): + targets = resolve_hf_lora_targets({"model_type": "qwen3"}) + selected = exclude_hf_lora_targets(targets, parse_lora_targets("o_proj, down_proj")) + assert "model.layers.*.self_attn.o_proj" not in selected + assert "model.layers.*.mlp.down_proj" not in selected + assert "model.layers.*.self_attn.q_proj" in selected + + +def test_exclude_all_rejects_empty_selection(): + with pytest.raises(AssertionError, match="empty after"): + exclude_hf_lora_targets(["q_proj", "k_proj"], ["q_proj", "k_proj"]) + + +def test_nonexistent_exclusion_does_not_change_selection(): + assert exclude_hf_lora_targets(["q_proj", "k_proj"], ["nonexistent"]) == ["q_proj", "k_proj"] + + +def test_empty_selector_is_rejected(): + with pytest.raises(AssertionError, match="empty entries"): + parse_lora_targets("q_proj,,k_proj") + + +@pytest.mark.parametrize("dense", [True, False]) +@pytest.mark.parametrize("scoped", [True, False]) +def test_paired_targets_include_packed_experts(dense, scoped, caplog): + modules = ["model.layers.0.mlp.experts.gate_up_proj", "model.layers.0.mlp.experts.down_proj"] + if dense: + modules += [f"model.layers.0.mlp.{name}_proj" for name in ("gate", "up", "down")] + prefix = "model.layers.*.mlp.experts." if scoped else "" + targets = [prefix + name + "_proj" for name in ("gate", "up", "down")] + expected = targets if dense and not scoped else [prefix + "down_proj"] + assert expand_packed_hf_lora_targets(targets, modules) == expected + [prefix + "gate_up_proj"] + assert "Expanding paired gate_proj/up_proj" in caplog.text + assert expand_packed_hf_lora_targets(targets[:1], modules) == targets[:1] + assert expand_packed_hf_lora_targets(targets, ["model.layers.0.mlp.gate_proj"]) == targets diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 1026333c9af..7d446d84ccd 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -2,6 +2,7 @@ import json import logging import sys +from fnmatch import fnmatchcase from pathlib import Path from types import SimpleNamespace from unittest.mock import patch @@ -25,6 +26,7 @@ ) from miles.utils.ft_utils.health_checker import SimpleHealthCheckerConfig from miles.utils.function_registry import function_registry +from miles.utils.hf_utils.weight_mapping import HfWeightMapping from miles.utils.run_uuid import RUN_UUID_LENGTH, validate_run_uuid PATH_ARGS = ["--rollout-function-path", "--custom-generate-function-path"] @@ -1006,6 +1008,21 @@ def test_critic_rejects_reward_level_kl(tmp_path): class TestMultiLoRAValidation: + @pytest.fixture(autouse=True) + def _bridge_selector(self, monkeypatch): + monkeypatch.setattr( + "miles.utils.lora.arguments.HfWeightMapping.from_config", + lambda config: SimpleNamespace(parameter_names=frozenset()), + ) + monkeypatch.setattr( + "miles.utils.lora.arguments.load_hf_config", + lambda path: SimpleNamespace(model_type="qwen3", to_dict=lambda: {"model_type": "qwen3"}), + ) + monkeypatch.setattr( + "miles.backends.megatron_utils.lora.target_modules.normalize_lora_targets_to_hf", + lambda checkpoint, targets, **kwargs: list(targets), + ) + def _parse(self, extra): parser = argparse.ArgumentParser() get_miles_extra_args_provider()(parser) @@ -1024,6 +1041,68 @@ def _parse(self, extra): + REQUIRED_ARGS ) + def test_hf_selection_does_not_resolve_through_bridge(self, monkeypatch): + monkeypatch.setattr( + "miles.utils.lora.arguments.HfWeightMapping.from_config", + lambda config: SimpleNamespace(parameter_names=frozenset({"model.layers.0.self_attn.q_proj.weight"})), + ) + + def unexpected_bridge(*args, **kwargs): + pytest.fail("HF selection must not initialize Bridge") + + monkeypatch.setattr( + "miles.backends.megatron_utils.lora.target_modules.normalize_lora_targets_to_hf", + unexpected_bridge, + ) + args = self._parse(["--target-modules", "q_proj"]) + miles_validate_args(args) + assert args.hf_lora_targets == ["q_proj"] + assert args.target_modules == "q_proj" + + @pytest.mark.parametrize("exclusion", ["model.layers.*.self_attn.o_proj", "model.layers.0.self_attn.o_proj"]) + @pytest.mark.parametrize("targets", ["o_proj,down_proj", "attn,mlp", "all-linear"]) + def test_scoped_exclusion_resolves_without_bridge(self, monkeypatch, exclusion, targets): + hf_mapping = HfWeightMapping( + frozenset( + f"model.layers.{layer}.{module}.weight" + for layer in range(2) + for module in ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ) + ) + ) + monkeypatch.setattr("miles.utils.lora.arguments.HfWeightMapping.from_config", lambda config: hf_mapping) + + def unexpected_bridge(*args, **kwargs): + pytest.fail("HF exclusions must not initialize Bridge") + + monkeypatch.setattr( + "miles.backends.megatron_utils.lora.target_modules.normalize_lora_targets_to_hf", + unexpected_bridge, + ) + args = self._parse(["--target-modules", targets, "--exclude-modules", exclusion]) + if targets == "o_proj,down_proj": + with pytest.raises(AssertionError, match="overlap --exclude-modules"): + miles_validate_args(args) + return + miles_validate_args(args) + + expected = { + name.removesuffix(".weight") + for name in hf_mapping.parameter_names + if not fnmatchcase(name.removesuffix(".weight"), exclusion) + } + assert args.target_modules == targets + assert args.exclude_modules == exclusion + assert set(args.hf_lora_targets) == expected + assert args.lora_adapter_targets == args.hf_lora_targets + def test_rejects_multiple_tokenizer_workers(self): # Each sglang tokenizer worker holds its own LoRA registry, so per-step # upserts fail non-deterministically; fail at launch, not first push. diff --git a/tests/fast/utils/test_lora_arguments.py b/tests/fast/utils/test_lora_arguments.py deleted file mode 100644 index f4e78839e90..00000000000 --- a/tests/fast/utils/test_lora_arguments.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Unit tests for LoRA-related argument parsing in miles.utils.arguments. - -Covers the target-module expansion and exclude-module filtering logic -inside miles_validate_args (lines 1634-1653 of arguments.py). -We isolate the LoRA parsing logic to avoid triggering unrelated validations. -""" - -from argparse import Namespace -from copy import deepcopy - -import pytest - - -def _apply_lora_arg_parsing(args: Namespace) -> Namespace: - """Extract and apply only the LoRA target-module parsing logic from - miles_validate_args, avoiding unrelated assertions.""" - args = deepcopy(args) - if args.lora_rank > 0: - assert args.target_modules is not None, "'--target-modules' is required when LoRA is enabled." - - if args.target_modules == "all-linear": - modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - elif "," in args.target_modules: - modules = [m.strip() for m in args.target_modules.split(",")] - else: - modules = [args.target_modules] - - if args.exclude_modules: - exclude_set = ( - set(m.strip() for m in args.exclude_modules.split(",")) - if "," in args.exclude_modules - else {args.exclude_modules} - ) - modules = [m for m in modules if m not in exclude_set] - - args.target_modules = modules - return args - - -# --------------------------------------------------------------------------- -# Target modules expansion -# --------------------------------------------------------------------------- - - -class TestLoraTargetModuleParsing: - def test_all_linear_expands_to_seven_modules(self): - args = Namespace(lora_rank=32, target_modules="all-linear", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] - - def test_comma_separated_split(self): - args = Namespace(lora_rank=16, target_modules="q_proj, k_proj, v_proj", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj", "v_proj"] - - def test_comma_separated_no_spaces(self): - args = Namespace(lora_rank=16, target_modules="q_proj,k_proj", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj"] - - def test_single_module(self): - args = Namespace(lora_rank=8, target_modules="q_proj", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj"] - - def test_lora_rank_zero_skips_parsing(self): - args = Namespace(lora_rank=0, target_modules="all-linear", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == "all-linear" # unchanged - - def test_missing_target_modules_asserts(self): - args = Namespace(lora_rank=32, target_modules=None, exclude_modules=None) - with pytest.raises(AssertionError, match="--target-modules"): - _apply_lora_arg_parsing(args) - - -# --------------------------------------------------------------------------- -# Exclude modules filtering -# --------------------------------------------------------------------------- - - -class TestLoraExcludeModules: - def test_single_exclude(self): - args = Namespace(lora_rank=32, target_modules="all-linear", exclude_modules="o_proj") - result = _apply_lora_arg_parsing(args) - assert "o_proj" not in result.target_modules - assert len(result.target_modules) == 6 - - def test_multiple_exclude_comma_separated(self): - args = Namespace(lora_rank=32, target_modules="all-linear", exclude_modules="o_proj, down_proj") - result = _apply_lora_arg_parsing(args) - assert "o_proj" not in result.target_modules - assert "down_proj" not in result.target_modules - assert len(result.target_modules) == 5 - - def test_exclude_all_results_in_empty(self): - args = Namespace( - lora_rank=32, - target_modules="q_proj,k_proj", - exclude_modules="q_proj,k_proj", - ) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == [] - - def test_exclude_nonexistent_module_no_effect(self): - args = Namespace(lora_rank=32, target_modules="q_proj,k_proj", exclude_modules="nonexistent") - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj"] - - def test_no_exclude_modules(self): - args = Namespace(lora_rank=32, target_modules="q_proj,k_proj", exclude_modules=None) - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj"] - - def test_empty_string_exclude(self): - """Empty string is truthy; should be treated as a single (non-matching) exclude.""" - args = Namespace(lora_rank=32, target_modules="q_proj,k_proj", exclude_modules="") - result = _apply_lora_arg_parsing(args) - assert result.target_modules == ["q_proj", "k_proj"] diff --git a/tests/manual/models/test_lora_target_resolution.py b/tests/manual/models/test_lora_target_resolution.py new file mode 100644 index 00000000000..c4419f57985 --- /dev/null +++ b/tests/manual/models/test_lora_target_resolution.py @@ -0,0 +1,111 @@ +"""Run in the Miles training image: pytest tests/manual/models/test_lora_target_resolution.py -v.""" + +from types import SimpleNamespace + +import pytest +import torch +from megatron.bridge.models.qwen.qwen3_moe_bridge import Qwen3MoEBridge +from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge +from transformers import AutoConfig + +from miles.backends.megatron_utils.lora.slots import create_multi_lora_instance +from miles.backends.megatron_utils.lora.target_modules import resolve_megatron_lora_targets +from miles.backends.megatron_utils.lora.utils import create_lora_instance +from miles.utils.hf_utils.weight_mapping import HfWeightMapping + + +@pytest.fixture(params=["qkv", "experts", "gdn"]) +def target_case(request): + case = request.param + config = AutoConfig.for_model( + "qwen3_next" if case == "gdn" else "qwen3_moe", + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=16, + shared_expert_intermediate_size=16, + decoder_sparse_step=1, + layer_types=["linear_attention", "full_attention"], + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_num_key_heads=2, + linear_num_value_heads=4, + ) + bridge = Qwen3NextBridge() if case == "gdn" else Qwen3MoEBridge() + targets, module, parameters = { + "qkv": ( + ["model.layers.*.self_attn.q_proj"], + "decoder.layers.0.self_attention.linear_qkv", + ["weight"], + ), + "experts": ( + ["model.layers.*.mlp.experts.gate_up_proj"], + "decoder.layers.0.mlp.experts.linear_fc1", + ["weight0", "weight1"], + ), + "gdn": ( + ["model.layers.*.linear_attn.in_proj_qkvz", "model.layers.*.linear_attn.in_proj_ba"], + "decoder.layers.0.self_attention.in_proj", + ["weight"], + ), + }[case] + return ( + case, + bridge.mapping_registry().get_all_mappings(), + HfWeightMapping.from_config(config), + targets, + module, + parameters, + ) + + +@pytest.mark.parametrize("mode", ["lora", "canonical_lora", "multi_lora"]) +def test_registry_to_adapter_matcher(target_case, mode): + case, mappings, hf_mapping, targets, module, parameters = target_case + canonical = mode == "canonical_lora" + + def resolve(selection): + # Explicit parameter fixtures exercise registry binding without constructing a distributed trainer. + return resolve_megatron_lora_targets( + selection, + mappings, + parameter_names={f"{module}.{parameter}" for parameter in parameters}, + hf_mapping=hf_mapping, + canonical=canonical, + ) + + if case == "qkv" and not canonical: + with pytest.raises(AssertionError, match="requires all HF targets"): + resolve(targets) + targets = [f"model.layers.*.self_attn.{projection}_proj" for projection in ("q", "k", "v")] + if case == "gdn": + with pytest.raises(AssertionError, match="requires all HF targets"): + resolve(targets[:1]) + + selected = resolve(targets) + template = module.replace(".0.", ".*.") + expected = {template} + if canonical and case == "qkv": + expected = {template.removesuffix("linear_qkv") + "linear_q"} + elif canonical and case == "experts": + expected = {template + "_gate", template + "_up"} + assert set(selected) == expected + + args = SimpleNamespace(lora_type=mode, lora_rank=2, lora_alpha=4, lora_dropout=0.0, multi_lora_n_adapters=2) + create_adapter = create_multi_lora_instance if mode == "multi_lora" else create_lora_instance + adapter = create_adapter(args, target_modules=selected) + adapter._init_target_match_state() + linear = torch.nn.Linear(2, 2, device="meta") + matched = set() + for candidate in (module, "decoder.layers.0.mlp.router", "output_layer"): + prefix, name = candidate.rsplit(".", 1) if "." in candidate else ("", candidate) + if adapter.match(linear, name=name, prefix=prefix): + matched.add(candidate) + assert matched == {module} + assert all(adapter._alias_matches[target] == {module} for target in selected) diff --git a/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt index 5922bbcadc5..690f224d183 100644 --- a/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt +++ b/tests/snapshots/launch_scripts/py/examples/multi_lora/serve_qwen3_30b_a3b_tinker.py/serve.txt @@ -51,12 +51,10 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-alpha 64 --lora-dropout 0.0 --no-gradient-accumulation-fusion - --multi-lora-n-adapters 4 + --multi-lora-n-adapters 4 + --target-modules "attn,mlp,unembed" --tinker-server-port 10613 --tinker-checkpoint-root /root/shared_data/checkpoints/260101-000000-000 - --tinker-train-attn - --tinker-train-mlp - --tinker-train-unembed --optimizer adam --lr 1e-4 --tensor-model-parallel-size 2 diff --git a/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/full_train.txt b/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/full_train.txt index b9fc962d3e0..03681a48b73 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/full_train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/full_train.txt @@ -73,7 +73,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/train.txt b/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/train.txt index 28c4eefb87e..50b8f2066c3 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_glm5_1_744b_a40b_lora.py/train.txt @@ -61,7 +61,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/full_train.txt b/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/full_train.txt index aa2cf96c978..ef6fcb7c9dc 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/full_train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/full_train.txt @@ -73,7 +73,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/train.txt b/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/train.txt index ab198a29e1f..38a7e6085ad 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_glm5_2_744b_a40b_lora.py/train.txt @@ -61,7 +61,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/py/scripts/run_kimi_k3.py/train.txt b/tests/snapshots/launch_scripts/py/scripts/run_kimi_k3.py/train.txt index 6fdb50cd1b8..46075a8655f 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_kimi_k3.py/train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_kimi_k3.py/train.txt @@ -68,7 +68,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "decoder.layers.*.self_attention.o_proj,decoder.layers.*.self_attention.q_a_proj,decoder.layers.*.self_attention.kv_a_proj_with_mqa,decoder.layers.*.mlp.linear_fc1,decoder.layers.*.mlp.linear_fc2,decoder.layers.*.mlp.experts.linear_fc1,decoder.layers.*.mlp.experts.linear_fc2" + --target-modules "all-linear" --no-gradient-accumulation-fusion --lora-base-cpu-backup --experts-shared-outer-loras diff --git a/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/full_train.txt b/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/full_train.txt index 2a35f365873..0ea75c6dda5 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/full_train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/full_train.txt @@ -68,7 +68,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "language_model.decoder.layers.*.self_attention.linear_qkv,language_model.decoder.layers.*.self_attention.linear_proj,language_model.decoder.layers.*.mlp.experts.linear_fc1,language_model.decoder.layers.*.mlp.experts.linear_fc2,language_model.decoder.layers.*.mlp.shared_experts.linear_fc1,language_model.decoder.layers.*.mlp.shared_experts.linear_fc2,language_model.decoder.layers.*.self_attention.in_proj,language_model.decoder.layers.*.self_attention.out_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/train.txt b/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/train.txt index 9f268780e0c..0db50e7b769 100644 --- a/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/train.txt +++ b/tests/snapshots/launch_scripts/py/scripts/run_qwen3_5_35b_a3b_lora.py/train.txt @@ -56,7 +56,7 @@ export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit --lora-rank 16 --lora-alpha 32 --lora-dropout 0.0 - --target-modules "language_model.decoder.layers.*.self_attention.linear_qkv,language_model.decoder.layers.*.self_attention.linear_proj,language_model.decoder.layers.*.mlp.experts.linear_fc1,language_model.decoder.layers.*.mlp.experts.linear_fc2,language_model.decoder.layers.*.mlp.shared_experts.linear_fc1,language_model.decoder.layers.*.mlp.shared_experts.linear_fc2,language_model.decoder.layers.*.self_attention.in_proj,language_model.decoder.layers.*.self_attention.out_proj" + --target-modules "all-linear" --experts-shared-outer-loras --no-gradient-accumulation-fusion --lora-base-cpu-backup diff --git a/tests/snapshots/launch_scripts/sh/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh.txt b/tests/snapshots/launch_scripts/sh/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh.txt index cfc63b109a2..b4af5508350 100644 --- a/tests/snapshots/launch_scripts/sh/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh.txt +++ b/tests/snapshots/launch_scripts/sh/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh.txt @@ -123,7 +123,7 @@ "--lora-dropout" "0.0" "--target-modules" -"gate_proj,up_proj,down_proj" +"gate_up_proj,down_proj" "--sglang-lora-backend" "triton" "--hf-checkpoint"