qwen3.5 lora compatibility (needs SGLang + Megatron-Bridge patches as well)

examples/lora/run-qwen3.5-35B-A3B-megatron-lora.sh

@@ -0,0 +1,172 @@
+#!/bin/bash
+
+# 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)"
+
+source "/root/miles/scripts/models/qwen3.5-35B-A3B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3.5-35B-A3B
+)
+
+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)
+   # BUG: For some reason, gate_up_proj lora adapter still encounters "illegal memory access" errors
+   --target-modules "q_proj,k_proj,v_proj"
+   --megatron-to-hf-mode bridge
+)
+
+
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 1
+   --eval-max-response-len 8000
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 1
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 8
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   # --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+   --no-offload-train
+)
+
+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
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+   # --use-wandb
+   # --wandb-project miles-dev
+   # --wandb-group qwen3.5-4B-test
+   # --wandb-key ${WANDB_KEY}
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.4
+   --sglang-tp-size 1
+   --sglang-ep-size 8
+   --sglang-dtype bfloat16
+
+   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
+   # mtp speculative decoding
+   #--sglang-disable-cuda-graph-padding
+
+
+   --sglang-max-running-requests 512
+   --sglang-moe-runner-backend triton
+   --sglang-lora-backend csgmv
+)
+
+
+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
+   # need to comment this when using model with MLA
+   --attention-backend flash
+   --moe-token-dispatcher-type flex
+   --update-weight-buffer-size 536870912 # 512MB
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --colocate \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${LORA_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]}

miles/backends/megatron_utils/lora_utils.py

@@ -1,5 +1,7 @@
 """LoRA utilities for Megatron backend using Megatron-Bridge PEFT integration."""
 
+from __future__ import annotations
+
 import logging
 import os
 from argparse import Namespace
@@ -231,6 +233,13 @@ def parse_exclude_modules(args: Namespace, lora_type=None) -> list[str]:
         exclude_modules = convert_target_modules_to_megatron(exclude_modules, lora_type=lora_type)
     return exclude_modules
 
+def exclude_mtp_vision_modules(target_modules: list[str]) -> list[str]:
+    """Restrict Qwen3.5-VL LoRA targets to the language model."""
+    return [
+        target if "language_model" in target else f"*language_model.decoder.layers.*.*.{target}"
+        for target in target_modules
+    ]
+
 
 def create_lora_instance(args: Namespace):
     """Create a LoRA or CanonicalLoRA instance based on args.
@@ -249,6 +258,11 @@ def create_lora_instance(args: Namespace):
         lora_cls = LoRA
 
     target_modules = convert_target_modules_to_megatron(args.target_modules, lora_type=lora_cls)
+
+    model_name = args.hf_checkpoint
+    if "Qwen3.5" in model_name:
+        target_modules = exclude_mtp_vision_modules(target_modules)
+
     exclude_modules = parse_exclude_modules(args, lora_type=lora_cls)
 
     lora = lora_cls(
@@ -451,6 +465,7 @@ def load_lora_adapter(
     return False, None
 
 
+
 def _load_training_state(
     adapter_dir: Path,
     optimizer: Any | None,
@@ -503,3 +518,4 @@ def build_lora_sync_config(args: Namespace) -> dict[str, Any]:
         "bias": "none",
         "task_type": "CAUSAL_LM",
     }
+