From fc07c3d22fb889548a74c88e86d5d4e22fdd2b94 Mon Sep 17 00:00:00 2001 From: yuxinshan Date: Tue, 25 Aug 2026 19:42:13 +0800 Subject: [PATCH] support kimi-k2-thinking model Signed-off-by: yuxinshan --- docker/npu_patch/megatron-bridge.patch | 12 +- scripts/models/kimi-k2-thinking-2layer.sh | 1 + scripts/models/kimi-k2-thinking.sh | 4 +- scripts/run-kimi-k2-thinking-npu.sh | 222 ++++++++++++++++++ tests/test_kimi_k2_thinking_npu.py | 163 +++++++++++++ .../quantizer_compressed_tensors.py | 38 ++- .../backends/megatron_utils/model_provider.py | 3 +- .../hf_weight_iterator_bridge.py | 35 ++- vime/utils/megatron_bridge_utils.py | 20 ++ 9 files changed, 486 insertions(+), 12 deletions(-) create mode 100644 scripts/models/kimi-k2-thinking-2layer.sh create mode 100644 scripts/run-kimi-k2-thinking-npu.sh create mode 100644 tests/test_kimi_k2_thinking_npu.py diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch index 0e7798937..f7e107d5a 100644 --- a/docker/npu_patch/megatron-bridge.patch +++ b/docker/npu_patch/megatron-bridge.patch @@ -35,15 +35,21 @@ index 63b4cc37..00394727 100644 # Check the actual parameter name to determine the correct parallelism type if self.megatron_param and ( self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias") -@@ -1313,7 +1318,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): +@@ -1313,9 +1318,10 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): return "replicated" # Check parallel_mode for TELinear - if module_type == "TELinear": + if module_type == "TELinear" or module_type == "MindSpeedTELinear": - if module.parallel_mode == "column": ++ parallel_mode = getattr(module, "parallel_mode", None) +- if module.parallel_mode == "column": ++ if parallel_mode == "column": return "column" - elif module.parallel_mode == "row": +- elif module.parallel_mode == "row": ++ elif parallel_mode == "row": + return "row" + else: + return "replicated" diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py index 7360857d..a6c00358 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py diff --git a/scripts/models/kimi-k2-thinking-2layer.sh b/scripts/models/kimi-k2-thinking-2layer.sh new file mode 100644 index 000000000..fdf4e2b38 --- /dev/null +++ b/scripts/models/kimi-k2-thinking-2layer.sh @@ -0,0 +1 @@ +MODEL_ARGS_NUM_LAYERS=2 source "$(dirname -- "${BASH_SOURCE[0]}")/kimi-k2-thinking.sh" diff --git a/scripts/models/kimi-k2-thinking.sh b/scripts/models/kimi-k2-thinking.sh index b7fceda59..ba53efb11 100644 --- a/scripts/models/kimi-k2-thinking.sh +++ b/scripts/models/kimi-k2-thinking.sh @@ -1,4 +1,4 @@ -NLAYERS=61 +NLAYERS="${MODEL_ARGS_NUM_LAYERS:-61}" FIRST_K_DENSE_REPLACE=1 arr=() @@ -15,7 +15,7 @@ printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")" # kimi-k2-thinking MODEL_ARGS=( --disable-bias-linear - --num-layers 61 + --num-layers $NLAYERS --hidden-size 7168 --ffn-hidden-size 18432 --num-attention-heads 64 diff --git a/scripts/run-kimi-k2-thinking-npu.sh b/scripts/run-kimi-k2-thinking-npu.sh new file mode 100644 index 000000000..13295f503 --- /dev/null +++ b/scripts/run-kimi-k2-thinking-npu.sh @@ -0,0 +1,222 @@ +#!/bin/bash +set -ex +ulimit -u 65535 + +# cleanup +pkill -9 -f "vllm serve" 2>/dev/null || true +sleep 2 +npu-smi info 2>/dev/null | grep rayWorker | awk '{print $4}' | xargs -r kill -9 2>/dev/null || true +sleep 3 + +# Ray isolation: independent temp-dir, ports, and cleanup +export RAY_TMPDIR=/tmp/ray_vime_npu_kimi_k2_thinking +export RAY_PORT=6379 +export RAY_DASHBOARD_PORT=8265 +export RAY_AGENT_PORT=52378 +unset RAY_ADDRESS RAY_REDIS_ADDRESS + +ray stop --force 2>/dev/null || true +rm -rf "${RAY_TMPDIR}" +sleep 2 + +project_name="vime" +exp_name="kimi_k2_thinking" +RAY_DATA_HOME=${RAY_DATA_HOME:-"/root/logs"} +start_time=$(date +"%Y%m%d_%H%M%S") +LOG_DIR=${LOG_DIR:-"${RAY_DATA_HOME}/${project_name}/${exp_name}"} +mkdir -p "${LOG_DIR}" +LOG_FILE="${LOG_DIR}/${start_time}.log" + +echo "Experiment Log will be saved to: ${LOG_FILE}" +VIME_DIR="/root/vime" + +# NPU environment +source /usr/local/Ascend/driver/bin/setenv.bash +source /usr/local/Ascend/ascend-toolkit/set_env.sh +source /usr/local/Ascend/nnal/atb/set_env.sh +export PYTHONPATH="${VIME_DIR}:/root/Megatron-LM:/vllm-workspace/vllm:/vllm-workspace/vllm-ascend:/root/Megatron-Bridge/src:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages:${PYTHONPATH}" +export PYTHONUNBUFFERED=1 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:False +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 +export HCCL_CONNECT_TIMEOUT=7200 +export HCCL_DETERMINISTIC=true +export VLLM_ASCEND_ENABLE_NZ=0 +export ASCEND_COREDUMP_SIGNAL=None +export ATB_MATMUL_SHUFFLE_K_ENABLE=0 +export ATB_LLM_LCOC_ENABLE=0 +export TASK_QUEUE_ENABLE=0 +export RAY_DISABLE_SIGINT_OVERRIDE=1 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export ASCEND_CUSTOM_OPP_PATH=/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer +export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:${LD_LIBRARY_PATH} +export VLLM_DISABLE_COMPILE_CACHE=1 +export TRANSFORMERS_VERBOSITY=error +export RUST_LOG=vllm_router_rs=warn + +NUM_NPUS=256 +source "${VIME_DIR}/scripts/models/kimi-k2-thinking.sh" + +CKPT_ARGS=( + --hf-checkpoint /path/to/Kimi-K2-Thinking + --ref-load /path/to/Kimi-K2-Thinking + --load /path/to/Kimi-K2-Thinking_npu/ + --save /path/to/Kimi-K2-Thinking_npu/ + --save-interval 20 + --no-load-optim + --megatron-to-hf-mode bridge +) + +ROLLOUT_ARGS=( + --prompt-data /path/to/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 200 + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-response-len 16384 + --rollout-temperature 1 + --global-batch-size 64 + --balance-data +) + +EVAL_ARGS=( + --eval-interval 50 + --eval-prompt-data aime /path/to/aime-2024/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 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 16384 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --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 +) + +VLLM_ARGS=( + --rollout-backend vllm + --rollout-num-gpus-per-engine 16 + --vllm-gpu-memory-utilization 0.7 + --vllm-data-parallel-size 8 + --vllm-enable-experet-parallel + --vllm-enable-sleep-mode + --vllm-weight-sync-mode native + --vllm-enforce-eager +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-flash-attn +) + +MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +CURRENT_NODE_IP="${CURRENT_NODE_IP:-$(hostname -I | awk '{print $1}')}" +if [[ "${CURRENT_NODE_IP}" != "${MASTER_ADDR}" ]];then + # launch the other nodes of ray in container + unset https_proxy http_proxy proxy + ray start \ + --address="${MASTER_ADDR}:${RAY_PORT}" + --node-ip-address "${CURRENT_NODE_IP}" \ + + --num-gpus 0 \ + --resources "{\"NPU\": $NUM_NPUS}" \ + --disable-usage-stats \ + --block +else + # launch the master node of ray in container + unset https_proxy http_proxy proxy + ray start --head \ + --temp-dir="${RAY_TMPDIR}" \ + --port="${RAY_PORT}" \ + --dashboard-port="${RAY_DASHBOARD_PORT}" \ + --dashboard-agent-listen-port="${RAY_AGENT_PORT}" \ + --node-ip-address "${MASTER_ADDR}" \ + --num-gpus 0 \ + --resources "{\"NPU\": $NUM_NPUS}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 + + # Build the runtime environment JSON with proper variable substitution + RUNTIME_ENV_JSON=$(cat << 'EOF' + { + "env_vars": { + "PYTHONPATH": "${VIME_DIR}:/root/Megatron-LM:/vllm-workspace/vllm:/vllm-workspace/vllm-ascend:/root/Megatron-Bridge/src:/root/mbridge:/root/MegatronAdaptor:/root/TransformerEngineNPU:/usr/local/Ascend/ascend-toolkit/latest/python/site-packages", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050", + "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050", + "HCCL_CONNECT_TIMEOUT": "7200", + "PYTORCH_NPU_ALLOC_CONF": "expandable_segments:False", + "VLLM_DISABLE_COMPILE_CACHE": "1", + "TRANSFORMERS_VERBOSITY": "error", + "RUST_LOG": "vllm_router_rs=warn", + "ASCEND_CUSTOM_OPP_PATH": "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer", + "LD_LIBRARY_PATH": "/usr/local/Ascend/driver/lib64:/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64/common:/usr/local/Ascend/ascend-toolkit/latest/lib64:/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/:/usr/local/Ascend/nnal/atb/latest/atb/cxx_abi_1/lib:/usr/local/Ascend/cann/lib64:/usr/local/Ascend/cann/aarch64-linux/devlib" + } + } + EOF + ) + + ray job submit --address="http://${MASTER_ADDR}:${RAY_DASHBOARD_PORT}" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --working-dir="${VIME_DIR}" \ + -- python3 -u train.py \ + --train-backend megatron \ + --actor-num-nodes 16 \ + --actor-num-gpus-per-node 16 \ + --colocate \ + --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024 )) + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} \ + 2>&1 | tee "${LOG_FILE}" +fi \ No newline at end of file diff --git a/tests/test_kimi_k2_thinking_npu.py b/tests/test_kimi_k2_thinking_npu.py new file mode 100644 index 000000000..fc6219121 --- /dev/null +++ b/tests/test_kimi_k2_thinking_npu.py @@ -0,0 +1,163 @@ +import json +import os +import shlex + +import vime.utils.external_utils.command_utils as U + + +TEST_ROOT = os.environ.get("HF_HOME") or "/root" +MODEL_DIR = f"{TEST_ROOT}/models/Kimi-K2-Thinking" +DATASET_DIR = f"{TEST_ROOT}/datasets/dapo-math-17k" + + +def prepare(): + models_dir = shlex.quote(f"{TEST_ROOT}/models") + datasets_dir = shlex.quote(f"{TEST_ROOT}/datasets") + model_dir = shlex.quote(MODEL_DIR) + dataset_dir = shlex.quote(DATASET_DIR) + + U.exec_command(f"mkdir -p {models_dir} {datasets_dir}") + U.exec_command(f"hf download moonshotai/Kimi-K2-Thinking --local-dir {model_dir}") + U.exec_command("hf download --repo-type dataset zhuzilin/dapo-math-17k " f"--local-dir {dataset_dir}") + + +def generate_dummy_config(model_dir): + model_config = os.path.join(model_dir, "config.json") + with open(model_config, encoding="utf-8") as f: + data = json.load(f) + + data["num_hidden_layers"] = 2 + data["quantization_config"]["ignore"].extend(["model.embed_tokens", "re:.*mlp\\.gate\\.weights$"]) + + with open(model_config, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + +def execute(): + model_dir = shlex.quote(MODEL_DIR) + prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") + + generate_dummy_config(model_dir) + + # NPU skips torch_dist conversion; HF weights load directly via bridge mode. + checkpoint_args = ( + f"--hf-checkpoint {model_dir} " + f"--load {model_dir} " + f"--ref-load {model_dir} " + "--megatron-to-hf-mode bridge " + "--no-load-optim " + ) + + # Smoke-scaled rollout (num-rollout/batch/n-samples trimmed like test_kimi_k2_thinking_npu). + rollout_args = ( + f"--prompt-data {prompt_data} " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 2 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 2048 " + "--rollout-temperature 1 " + "--global-batch-size 16 " + "--balance-data " + ) + + # TP=8/EP=8 mirrors scripts/run-kimi-k2-thinking-npu.sh. + parallel_args = ( + "--tensor-model-parallel-size 8 " + "--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 " + "--max-tokens-per-gpu 2048 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + "--use-tis" + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--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 " + ) + + vllm_args = ( + "--rollout-backend vllm " + "--rollout-num-gpus-per-engine 16 " + "--vllm-gpu-memory-utilization 0.45 " + "--vllm-enable-expert-parallel " + "--vllm-enable-sleep-mode " + "--vllm-weight-sync-mode native " + "--vllm-enforce-eager " + "--vllm-load-format dummy " + ) + + model_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--use-flash-attn " + ) + + runtime_args = ( + "--train-backend megatron " "--actor-num-nodes 1 " "--actor-num-gpus-per-node 8 " "--rollout-num-gpus 8 " + ) + + train_args = ( + checkpoint_args + + rollout_args + + parallel_args + + grpo_args + + optimizer_args + + vllm_args + + model_args + + runtime_args + ) + # Model architecture (--spec, --attention-output-gate, --moe-shared-expert-gate, + # num-experts, moe-* ...) is injected by sourcing scripts/models/kimi-k2-thinking.sh + # via ${MODEL_ARGS[@]}, so only runtime/training args are passed here. + U.execute_train( + train_args=train_args, + num_gpus_per_node=16, + megatron_model_type="kimi-k2-thinking-2layer", + extra_env_vars={ + "DISABLE_L2_CACHE": "1", + "VLLM_USE_AOT_COMPILE": "0", + "ASCEND_CUSTOM_OPP_PATH": "/vllm-workspace/vllm-ascend/vllm_ascend/_cann_ops_custom/vendors/custom_transformer:/usr/local/Ascend/cann-9.0.0/opp/vendors/fla_npu_transformer", + }, + ) + + +def main(): + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() + + +if __name__ == "__main__": + main() diff --git a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py index ca69df8e6..90a6b7d9e 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py +++ b/vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py @@ -230,14 +230,46 @@ def if_quant(name, patterns): return False +def _fake_int4_quant_fallback(weight, group_shape, sym=True): + """Pure PyTorch fallback for fake_int4_quant_cuda on NPU.""" + group_size = group_shape[1] + orig_shape = weight.shape + rows, cols = orig_shape[0], orig_shape[1] + n_groups = cols // group_size + + w = weight.reshape(rows, n_groups, group_size).to(torch.float32) + + if sym: + w_max = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8) + scale = w_max / 7.0 # int4 symmetric range: [-8, 7], use 7 as scale + q = torch.round(w / scale).clamp(-8, 7) + zp = None + else: + w_min = w.amin(dim=-1, keepdim=True) + w_max = w.amax(dim=-1, keepdim=True) + scale = ((w_max - w_min) / 15.0).clamp(min=1e-8) # int4 asymmetric: [0, 15] + zp = torch.round(-w_min / scale).clamp(0, 15) + q = torch.round(w / scale + zp).clamp(0, 15) + + q = q.reshape(orig_shape).to(torch.float32) + scale = scale.reshape(rows, n_groups).contiguous() + if zp is not None: + zp = zp.reshape(rows, n_groups).contiguous() + + return q, scale, zp + + def pack_layer(weight, group_size, sym=True): - w, scale, zp = fake_int4_quant_cuda.fake_int4_quant_cuda(weight, (1, group_size), sym) + if fake_int4_quant_cuda is not None: + w, scale, zp = fake_int4_quant_cuda.fake_int4_quant_cuda(weight, (1, group_size), sym) + else: + w, scale, zp = _fake_int4_quant_fallback(weight, (1, group_size), sym) w = w.view(weight.shape[0], 1, weight.shape[1] // group_size, group_size) scale = scale.view(weight.shape[0], 1, weight.shape[1] // group_size, 1) - zp = zp.view(weight.shape[0], 1, weight.shape[1] // group_size, 1) if sym: w = w * scale else: + zp = zp.view(weight.shape[0], 1, weight.shape[1] // group_size, 1) w = (w - zp) * scale w = w.view(weight.shape) scale = scale.view(weight.shape[0], -1).contiguous() @@ -283,7 +315,7 @@ def quantize_params_compressed_tensors(converted_named_params, quantization_conf qw, s, zp = pack_layer(param, group_size, is_symmetric) qweight_name = name.replace(".weight", ".weight_packed") scale_name = name.replace(".weight", ".weight_scale") - weight_shape = torch.tensor(param.shape, dtype=torch.int32, device="cuda") + weight_shape = torch.tensor(param.shape, dtype=torch.int32, device=param.device) weight_shape_name = name.replace(".weight", ".weight_shape") if zp is not None: zp_name = name.replace(".weight", ".weight_zero_point") diff --git a/vime/backends/megatron_utils/model_provider.py b/vime/backends/megatron_utils/model_provider.py index 8c2b7a33f..afdd29d0f 100644 --- a/vime/backends/megatron_utils/model_provider.py +++ b/vime/backends/megatron_utils/model_provider.py @@ -17,7 +17,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.arguments import core_transformer_config_from_args -from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config +from vime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_auto_bridge_hf_config_for_model from vime.utils.misc import load_function @@ -87,6 +87,7 @@ def wrapped_model_provider( import vime_plugins.megatron_bridge # noqa: F401 # register custom bridges bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) + bridge = patch_auto_bridge_hf_config_for_model(bridge) provider = bridge.to_megatron_provider(load_weights=False) # TODO: we should not manually set this... provider.tensor_model_parallel_size = args.tensor_model_parallel_size diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py index d9aa7338f..2dccd8b82 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py @@ -1,6 +1,6 @@ import dataclasses - +import torch from vime.utils import megatron_bridge_utils from vime.utils.misc import chunk_named_params_by_size @@ -30,7 +30,10 @@ def _patched(self, task, converted_weights_dict): cpu_dict = {k: v.cpu() for k, v in converted_weights_dict.items()} result = _orig(self, task, cpu_dict) # Move merged result back to GPU for CUDA IPC serialization - return {k: v.cuda() for k, v in result.items()} if result else result + if getattr(torch, "npu", None) and torch.npu.is_available(): + return {k: v.npu() for k, v in result.items()} if result else result + else: + return {k: v.cuda() for k, v in result.items()} if result else result GPTOSSBridge.maybe_modify_converted_hf_weight = _patched GPTOSSBridge._cpu_cache_patched = True @@ -49,6 +52,25 @@ def __init__(self, *args, **kwargs): ) _patch_bridge_expert_cache_to_cpu() + # Patch megatron-bridge to handle None parallelism_type + try: + from megatron.bridge.models.conversion.param_mapping import AutoMapping + + _orig_megatron_to_hf = AutoMapping.megatron_to_hf + + def _patched_megatron_to_hf(self, megatron_weight, megatron_module): + try: + return _orig_megatron_to_hf(self, megatron_weight, megatron_module) + except ValueError as e: + if "Unknown parallelism type: None" in str(e): + hf_param = getattr(self, "hf_param", None) + return {hf_param: megatron_weight} + raise + + AutoMapping.megatron_to_hf = _patched_megatron_to_hf + except (ImportError, AttributeError): + pass + def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name) renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()} @@ -60,12 +82,16 @@ def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Upd def _streaming_quantized(): for hf_param_name, weight, megatron_param_name in named_weights: + if weight is None: + continue processed_weight = postprocess_hf_param( args=self.args, megatron_param_name=megatron_param_name, hf_param_name=hf_param_name, param=weight, ) + if processed_weight is None: + continue converted_named_params = [(hf_param_name, processed_weight)] quantized_batch = quantize_params( args=self.args, @@ -93,7 +119,10 @@ def _handle_one(task): ), f"{weight_dict_key=} not in new_weight_dict ({task.vp_stage=}, {task.param_name=}, {list(new_weight_dict)=})" new_param_weight = new_weight_dict[weight_dict_key] - new_param_weight = new_param_weight.cuda() + if getattr(torch, "npu", None) and torch.npu.is_available(): + new_param_weight = new_param_weight.npu() + else: + new_param_weight = new_param_weight.cuda() return dataclasses.replace(task, param_weight=new_param_weight) return _MapWithLen(_handle_one, vanilla_conversion_tasks) diff --git a/vime/utils/megatron_bridge_utils.py b/vime/utils/megatron_bridge_utils.py index c87fb5b7b..ba8610d43 100644 --- a/vime/utils/megatron_bridge_utils.py +++ b/vime/utils/megatron_bridge_utils.py @@ -38,6 +38,26 @@ def patch_auto_bridge_hf_config(bridge): return bridge +def patch_auto_bridge_hf_config_for_model(bridge): + if bridge is None: + return bridge + + hf_pretrained = getattr(bridge, "hf_pretrained", None) + config = hf_pretrained.config if hasattr(hf_pretrained, "config") else hf_pretrained + + # Kimi K2 model + from megatron.bridge.models.kimi.kimi_bridge import KimiK2Bridge + + if ( + getattr(config, "model_type", "") == "kimi_k2" + and "KimiK2ForCausalLM" not in getattr(config, "architectures", []) + and not isinstance(bridge._model_bridge, KimiK2Bridge) + ): + bridge.__dict__["_causal_lm_architecture"] = "KimiK2ForCausalLM" + + return bridge + + @contextmanager def patch_megatron_model(model): unwrapped_model = unwrap_model(model)[0]