4/5 support kimi 2.5 full + lora: shared-outer grouped-expert LoRA

examples/lora/run-kimi-k25-megatron-lora.sh

@@ -0,0 +1,182 @@
+#!/bin/bash
+
+# Kimi-K2.5 LoRA GRPO — 16 nodes × 8 GPUs (H200), colocated.
+# Inherits the full-param Kimi-K2.5 recipe and only overrides LoRA-specific
+# bits (rank/alpha, target modules, shared-outer adapters, LR, parallelism).
+
+# for rerun the task
+pkill -9 sglang
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+export PYTHONBUFFERED=16
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/kimi-k2-thinking.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint $BASE_DIR/Kimi-K2.5-int4
+   --ref-load $BASE_DIR/Kimi-K2.5-bf16
+   --megatron-to-hf-mode bridge
+   --model-name kimi_k25
+)
+
+LORA_ARGS=(
+   --lora-rank 32                       # LoRA rank (typical values: 8, 16, 32, 64)
+   --lora-alpha 32                      # LoRA alpha (usually equal to rank for RL)
+   --lora-dropout 0.0                   # LoRA dropout (0.0 for RL training)
+   --target-modules "q_a_proj,kv_a_proj_with_mqa,o_proj,gate_proj,up_proj,down_proj"
+   --experts-shared-outer-loras         # shared A on fc1 / shared B on fc2 across experts
+   --no-gradient-accumulation-fusion
+   --sglang-lora-backend triton         # !!! must for moe-lora !!!
+   --sglang-lora-use-virtual-experts    # virtual-experts MoE LoRA path
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --balance-data
+   --rm-type deepscaler
+
+   --num-rollout 20
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 16384
+   --rollout-temperature 1
+
+   --global-batch-size 256
+   --filter-zero-reward-samples
+   --use-dynamic-global-batch-size
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime $BASE_DIR/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 8
+   --sequence-parallel
+   --pipeline-model-parallel-size 2
+   --context-parallel-size 8
+   --expert-model-parallel-size 64
+   --expert-tensor-parallel-size 1
+   --decoder-last-pipeline-num-layers 30
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 4096
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+   # Off-policy IS correction: PPO operates on within-train ratio; TIS clamps
+   # the cross-engine (sglang Marlin int4 vs Megatron fake-QAT bf16) ratio with
+   # a wider bound than PPO's eps_clip, keeping kernel-rounding bias out of
+   # PPO clipping.
+   --use-tis
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-5                            # PEFT tolerates ~10x full-param LR
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+   --use-distributed-optimizer
+)
+
+WANDB_ARGS=(
+   --use-wandb
+   --wandb-project miles-kimi-k25
+   --wandb-group kimi-k25-lora
+   --disable-wandb-random-suffix
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.7
+   --sglang-ep-size 8
+   --sglang-server-concurrency 1024
+   --sglang-cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128
+   --use-rollout-routing-replay
+)
+
+MISC_ARGS=(
+   # default dropout in megatron is 0.1
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   # should be good for model performance
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --no-check-for-nan-in-loss-and-grad
+)
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+    \"NCCL_TIMEOUT\": \"3600\",
+    \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\",
+    \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"32\",
+    \"no_proxy\": \"${no_proxy}\",
+    \"MASTER_ADDR\": \"${MASTER_ADDR}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train.py \
+   --actor-num-nodes 16 \
+   --actor-num-gpus-per-node 8 \
+   --colocate \
+   --use-miles-router \
+   --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024 )) \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${LORA_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]}

miles/backends/megatron_utils/bridge_lora_helpers.py

@@ -93,9 +93,14 @@ def _setup_lora_model_via_bridge(args: Namespace) -> list:
     provider.sequence_parallel = args.sequence_parallel
     provider.virtual_pipeline_model_parallel_size = args.virtual_pipeline_model_parallel_size
     provider.context_parallel_size = args.context_parallel_size
+    provider.gradient_accumulation_fusion = args.gradient_accumulation_fusion
     provider.variable_seq_lengths = True
     provider.moe_token_dispatcher_type = "alltoall"
     provider.moe_router_load_balancing_type = "none"
+    if getattr(args, "decoder_first_pipeline_num_layers", None) is not None:
+        provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers
+    if getattr(args, "decoder_last_pipeline_num_layers", None) is not None:
+        provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers
     provider.finalize()
 
     lora = create_lora_instance(args)

miles/backends/megatron_utils/lora_utils.py

@@ -72,6 +72,19 @@
 
 _HF_MODULE_NAMES = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"}
 
+# 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",
+}
+_MEGATRON_MLA_TO_HF = {v: k for k, v in _MLA_HF_TO_MEGATRON.items()}
+
+# SGLang default get_hidden_dim (lora/utils.py) handles fused_qkv_a_proj_with_mqa via q_a / kv_a mapping,
+# but not separate q_b_proj / kv_b_proj yet — omit from rollout adapter config to avoid init crashes.
+_SGLANG_UNSUPPORTED_HF_TARGETS = frozenset({"q_b_proj", "kv_b_proj"})
+
 
 # ---------------------------------------------------------------------------
 # Core helpers
@@ -181,14 +194,20 @@ def convert_target_modules_to_megatron(
         if hf_modules[0] in ("all", "all-linear", "all_linear"):
             return list(all_modules)
 
-    # Check if already in Megatron format
-    if all(m not in _HF_MODULE_NAMES for m in hf_modules if "*" not in m):
-        return hf_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:
-        megatron_name = hf_to_megatron.get(module, module)
+        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)
 
@@ -204,14 +223,45 @@ def convert_target_modules_to_hf(megatron_modules: list[str]) -> list[str]:
     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:
-        if module in _MEGATRON_TO_HF_MODULES:
-            hf_modules.extend(_MEGATRON_TO_HF_MODULES[module])
+        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:
             hf_modules.append(module)
-    return hf_modules
+    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, with MLA q_b/kv_b dropped (unsupported)."""
+    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
 
 
 # ---------------------------------------------------------------------------
@@ -251,7 +301,7 @@ def create_lora_instance(args: Namespace):
     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 = lora_cls(
+    lora_kwargs = dict(
         target_modules=target_modules,
         exclude_modules=exclude_modules,
         dim=args.lora_rank,
@@ -260,6 +310,12 @@ def create_lora_instance(args: Namespace):
         lora_A_init_method=getattr(args, "lora_A_init_method", "xavier"),
         lora_B_init_method=getattr(args, "lora_B_init_method", "zero"),
     )
+    # Opt-in to SGLang PR #21466's shared-outer grouped-expert LoRA. Only the
+    # standard ``LoRA`` class supports the flag today.
+    if lora_cls is LoRA and getattr(args, "experts_shared_outer_loras", False):
+        lora_kwargs["experts_shared_outer_loras"] = True
+
+    lora = lora_cls(**lora_kwargs)
 
     logger.info(
         f"Created {lora_cls.__name__}: rank={args.lora_rank}, alpha={args.lora_alpha}, "
@@ -490,7 +546,7 @@ def _load_training_state(
 def build_lora_sync_config(args: Namespace) -> dict[str, Any]:
     """Build LoRA config dict for syncing weights to SGLang engines."""
     target_modules_hf = (
-        convert_target_modules_to_hf(list(args.target_modules))
+        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"]
     )