From 7faa228f8537a36cdd2cb8b6059eb4d576831097 Mon Sep 17 00:00:00 2001 From: Artem Yatsenko Date: Mon, 11 May 2026 06:49:24 +0000 Subject: [PATCH] qwen3.5 lora compatibility --- .../lora/run-qwen3.5-35B-A3B-megatron-lora.sh | 172 ++++++++++++++++++ miles/backends/megatron_utils/lora_utils.py | 16 ++ 2 files changed, 188 insertions(+) create mode 100644 examples/lora/run-qwen3.5-35B-A3B-megatron-lora.sh diff --git a/examples/lora/run-qwen3.5-35B-A3B-megatron-lora.sh b/examples/lora/run-qwen3.5-35B-A3B-megatron-lora.sh new file mode 100644 index 00000000000..ca4920a543f --- /dev/null +++ b/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[@]} diff --git a/miles/backends/megatron_utils/lora_utils.py b/miles/backends/megatron_utils/lora_utils.py index 1a78cad63a2..bdbfb6c2db5 100644 --- a/miles/backends/megatron_utils/lora_utils.py +++ b/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", } +