From 43efe7599bdd5a732f70e8c3f80d99e7512ab33b Mon Sep 17 00:00:00 2001 From: albertimff Date: Fri, 31 Jul 2026 17:19:26 +0800 Subject: [PATCH 1/2] Train all Qwen3.5 assistant turns --- tests/chat_template/test_chat_template.py | 16 ++- tests/datasets/test_qwen35_vl_tokenize_fn.py | 38 ++++--- tests/test_qwen35_chat.py | 102 +++++++++++++++++++ xtuner/v1/data_proto/messages/qwen35_chat.py | 12 +-- 4 files changed, 143 insertions(+), 25 deletions(-) create mode 100644 tests/test_qwen35_chat.py diff --git a/tests/chat_template/test_chat_template.py b/tests/chat_template/test_chat_template.py index 6ed85fc014..cb66bbb2ab 100644 --- a/tests/chat_template/test_chat_template.py +++ b/tests/chat_template/test_chat_template.py @@ -248,12 +248,22 @@ def test_qwen35vl_template(self): for j, data in enumerate(all_data): if j in [13,14]: # video 肯定和 hf 对不上 continue - gt_token_ids, gt_labels = qwen35_tokenize_fn_slowspeed(tokenizer, data['messages'], tools=data.get('tools'), add_vision_id=True) _messages = Qwen35ChatMessages(messages=data["messages"], tools=data.get("tools")) tokenized = _messages.tokenize(tokenizer, chat_template, add_vision_id=True) decode_str = tokenizer.decode(tokenized['input_ids'], skip_special_tokens=False) - - if j!=15 and j!=16: + has_multiple_user_queries = sum(msg["role"] == "user" for msg in data["messages"]) > 1 + + if has_multiple_user_queries: + # XTuner 内部渲染会保留所有 assistant turn，不再与会丢弃历史 thinking 的 HF 模板对齐。 + for message in data["messages"]: + if message["role"] == "assistant": + self.assertIn(message["content"].strip(), decode_str) + if message.get("reasoning_content"): + self.assertIn(message["reasoning_content"].strip(), decode_str) + elif j!=15 and j!=16: + gt_token_ids, gt_labels = qwen35_tokenize_fn_slowspeed( + tokenizer, data["messages"], tools=data.get("tools"), add_vision_id=True + ) self.assertEqual(tokenized['input_ids'], gt_token_ids) self.assertEqual(tokenized['labels'], gt_labels) diff --git a/tests/datasets/test_qwen35_vl_tokenize_fn.py b/tests/datasets/test_qwen35_vl_tokenize_fn.py index 3f8cbf5295..52ae1a896b 100644 --- a/tests/datasets/test_qwen35_vl_tokenize_fn.py +++ b/tests/datasets/test_qwen35_vl_tokenize_fn.py @@ -42,22 +42,36 @@ def test_qwen35vl_text(self): for j, data in enumerate(all_data): if j>=12: break - gt_token_ids, gt_labels = qwen35_tokenize_fn_slowspeed(self.tokenizer, data['messages'], tools=data.get('tools'), add_vision_id=True) ret = tokenize_fn(data) input_ids_xtuner = ret['input_ids'] labels_xtuner = ret['labels'] - self.assertEqual(input_ids_xtuner, gt_token_ids) - self.assertEqual(labels_xtuner, gt_labels) - - enable_thinking = any("reasoning_content" in msg for msg in data['messages']) decode_str = self.tokenizer.decode(input_ids_xtuner, skip_special_tokens=False) - hf_text = self.tokenizer.apply_chat_template(data['messages'], - tools=data.get('tools'), - add_vision_id=True, - tokenize=False, - enable_thinking=enable_thinking, - add_generation_prompt=False) - self.assertEqual(decode_str, hf_text) + has_multiple_user_queries = sum(msg["role"] == "user" for msg in data["messages"]) > 1 + + if has_multiple_user_queries: + # XTuner 内部渲染会保留所有 assistant turn，不再与会丢弃历史 thinking 的 HF 模板对齐。 + for message in data["messages"]: + if message["role"] == "assistant": + self.assertIn(message["content"].strip(), decode_str) + if message.get("reasoning_content"): + self.assertIn(message["reasoning_content"].strip(), decode_str) + else: + gt_token_ids, gt_labels = qwen35_tokenize_fn_slowspeed( + self.tokenizer, data["messages"], tools=data.get("tools"), add_vision_id=True + ) + self.assertEqual(input_ids_xtuner, gt_token_ids) + self.assertEqual(labels_xtuner, gt_labels) + + enable_thinking = any("reasoning_content" in msg for msg in data["messages"]) + hf_text = self.tokenizer.apply_chat_template( + data["messages"], + tools=data.get("tools"), + add_vision_id=True, + tokenize=False, + enable_thinking=enable_thinking, + add_generation_prompt=False, + ) + self.assertEqual(decode_str, hf_text) @parametrize.parametrize("add_vision_id", [(True,), (False,)]) def test_qwen35_vl_sft_single_image(self, add_vision_id): diff --git a/tests/test_qwen35_chat.py b/tests/test_qwen35_chat.py new file mode 100644 index 0000000000..0254b77c12 --- /dev/null +++ b/tests/test_qwen35_chat.py @@ -0,0 +1,102 @@ +from xtuner.v1.data_proto.messages.qwen35_chat import qwen35_tokenize_fn_fastspeed + + +class _CharacterTokenizer: + def __init__(self): + self.text = "" + + def __call__(self, text, return_offsets_mapping=False, add_special_tokens=False): + assert not add_special_tokens + self.text = text + output = {"input_ids": [ord(char) for char in text]} + if return_offsets_mapping: + output["offset_mapping"] = [(index, index + 1) for index in range(len(text))] + return output + + +def _labels_for(text, labels, value): + start = text.index(value) + return labels[start : start + len(value)] + + +def _is_trained(text, input_ids, labels, value): + start = text.index(value) + end = start + len(value) + return labels[start:end] == input_ids[start:end] + + +def _is_masked(text, labels, value): + return all(label == -100 for label in _labels_for(text, labels, value)) + + +def test_qwen35_trains_every_assistant_turn(): + messages = [ + {"role": "system", "content": "SYSTEM_TEXT"}, + {"role": "user", "content": "FIRST_USER_TEXT"}, + { + "role": "assistant", + "reasoning_content": "FIRST_THINK_TEXT", + "content": "FIRST_ANSWER_TEXT", + }, + {"role": "user", "content": "SECOND_USER_TEXT"}, + { + "role": "assistant", + "reasoning_content": "SECOND_THINK_TEXT", + "content": "SECOND_ANSWER_TEXT", + }, + { + "role": "assistant", + "reasoning_content": "DISABLED_THINK_TEXT", + "content": "DISABLED_ANSWER_TEXT", + "loss": False, + }, + ] + + tokenizer = _CharacterTokenizer() + input_ids, labels = qwen35_tokenize_fn_fastspeed(messages, tokenizer=tokenizer) + text = tokenizer.text + + assert _is_trained(text, input_ids, labels, "FIRST_THINK_TEXT") + assert _is_trained(text, input_ids, labels, "FIRST_ANSWER_TEXT") + assert _is_trained(text, input_ids, labels, "SECOND_THINK_TEXT") + assert _is_trained(text, input_ids, labels, "SECOND_ANSWER_TEXT") + + assert _is_masked(text, labels, "SYSTEM_TEXT") + assert _is_masked(text, labels, "FIRST_USER_TEXT") + assert _is_masked(text, labels, "SECOND_USER_TEXT") + assert _is_masked(text, labels, "DISABLED_THINK_TEXT") + assert _is_masked(text, labels, "DISABLED_ANSWER_TEXT") + + +def test_qwen35_trains_tool_calls_but_not_tool_responses(): + messages = [ + {"role": "system", "content": "SYSTEM_TEXT"}, + {"role": "user", "content": "USER_TEXT"}, + { + "role": "assistant", + "reasoning_content": "TOOL_THINK_TEXT", + "content": "TOOL_ANSWER_TEXT", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "UNIQUE_TOOL_NAME", + "arguments": {"argument": "UNIQUE_ARGUMENT_VALUE"}, + }, + } + ], + }, + {"role": "tool", "content": "UNIQUE_TOOL_RESPONSE"}, + {"role": "assistant", "content": "FINAL_ANSWER_TEXT"}, + ] + + tokenizer = _CharacterTokenizer() + input_ids, labels = qwen35_tokenize_fn_fastspeed(messages, tokenizer=tokenizer) + text = tokenizer.text + + assert _is_trained(text, input_ids, labels, "TOOL_THINK_TEXT") + assert _is_trained(text, input_ids, labels, "TOOL_ANSWER_TEXT") + assert _is_trained(text, input_ids, labels, "UNIQUE_TOOL_NAME") + assert _is_trained(text, input_ids, labels, "UNIQUE_ARGUMENT_VALUE") + assert _is_trained(text, input_ids, labels, "FINAL_ANSWER_TEXT") + assert _is_masked(text, labels, "UNIQUE_TOOL_RESPONSE") diff --git a/xtuner/v1/data_proto/messages/qwen35_chat.py b/xtuner/v1/data_proto/messages/qwen35_chat.py index dcb09d5d51..fccb2f95c2 100644 --- a/xtuner/v1/data_proto/messages/qwen35_chat.py +++ b/xtuner/v1/data_proto/messages/qwen35_chat.py @@ -175,16 +175,8 @@ def _append(text: str, is_loss: bool) -> None: sys_content = _render(messages[0]["content"], False).strip() _append(f"<|im_start|>system\n{sys_content}<|im_end|>\n", False) - # ── 计算 last_query_index ───────────────────────────────────────────── - multi_step_tool = True - last_query_index = len(messages) - 1 - for i in range(len(messages) - 1, -1, -1): - msg = messages[i] - if multi_step_tool and msg["role"] == "user": - content_str = _render(msg["content"], False).strip() - if not (content_str.startswith("") and content_str.endswith("")): - multi_step_tool = False - last_query_index = i + # 强制所有 assistant turn 都走训练分支，保留并训练历史 thinking/content。 + last_query_index = -1 # ── 主循环 ──────────────────────────────────────────────────────────── for idx, message in enumerate(messages): From bdaf8cc68326bebc7133539fc97698734730c001 Mon Sep 17 00:00:00 2001 From: albertimff Date: Mon, 3 Aug 2026 16:01:56 +0800 Subject: [PATCH 2/2] Add Qwen3.6 template and tokenize debug examples --- example_run/sft_qwen35_35b_256k_32gpu.py | 198 +++++ example_run/submit_qwen35_sft_16gpu_sp8.sh | 153 ++++ example_run/submit_qwen35_sft_32gpu.sh | 141 ++++ example_run/submit_qwen36_sft_16gpu_sp8.sh | 160 ++++ example_run/tokenize_sample_viewer.py | 797 ++++++++++++++++++ tests/datasets/test_tokenize_debug.py | 171 ++++ tests/test_qwen36_chat.py | 54 ++ xtuner/v1/data_proto/messages/__init__.py | 3 +- xtuner/v1/data_proto/messages/qwen35_chat.py | 12 +- xtuner/v1/data_proto/messages/qwen36_chat.py | 73 ++ xtuner/v1/data_proto/templates/__init__.py | 6 + xtuner/v1/datasets/config.py | 6 + .../mllm_tokenize_fn/qwen3_vl_tokenize_fn.py | 44 +- xtuner/v1/datasets/sft_tokenize_fn/openai.py | 4 +- xtuner/v1/datasets/tokenize_debug.py | 282 +++++++ 15 files changed, 2084 insertions(+), 20 deletions(-) create mode 100644 example_run/sft_qwen35_35b_256k_32gpu.py create mode 100755 example_run/submit_qwen35_sft_16gpu_sp8.sh create mode 100755 example_run/submit_qwen35_sft_32gpu.sh create mode 100755 example_run/submit_qwen36_sft_16gpu_sp8.sh create mode 100644 example_run/tokenize_sample_viewer.py create mode 100644 tests/datasets/test_tokenize_debug.py create mode 100644 tests/test_qwen36_chat.py create mode 100644 xtuner/v1/data_proto/messages/qwen36_chat.py create mode 100644 xtuner/v1/datasets/tokenize_debug.py diff --git a/example_run/sft_qwen35_35b_256k_32gpu.py b/example_run/sft_qwen35_35b_256k_32gpu.py new file mode 100644 index 0000000000..8c2cad0fa1 --- /dev/null +++ b/example_run/sft_qwen35_35b_256k_32gpu.py @@ -0,0 +1,198 @@ +import json +import os +import shutil +from pathlib import Path +from typing import Any + +from xtuner.v1.config import FSDPConfig, LRConfig, MuonConfig +from xtuner.v1.datasets import ( + PretrainTokenizeFunctionConfig, + Qwen3VLTokenizeFnConfig, +) +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.datasets.mllm_tokenize_fn import OSSLoaderConfig +from xtuner.v1.loss import CELossConfig +from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config +from xtuner.v1.model.compose.qwen3_vl.modeling_qwen3_vl import ( + QWEN3VL_COMPILE_CFG, +) +# from xtuner.v1.model.moe.moe import MTPConfig +from xtuner.v1.train import ResumeConfig, TrainerConfig + + +# This vision-layer compile rule is incompatible with Qwen3.5-35B-A3B. +QWEN3VL_COMPILE_CFG.pop( + "xtuner.v1.model.compose.qwen3_vl.modeling_vision." + "Qwen3VLVisionLayer.forward", + None, +) + + +def _get_int_env(name: str, default: int) -> int: + return int(os.getenv(name, str(default))) + + +def _get_float_env(name: str, default: float) -> float: + return float(os.getenv(name, str(default))) + + +def _get_bool_env(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.lower() in {"1", "true", "yes", "y", "on"} + + +# Paths are provided by one of the submit scripts in this directory. +ceph_config = os.getenv("CEPH_CONFIG", "") +meta_data_path = Path(os.environ["META_DATA_PATH"]) +model_path = Path(os.environ["MODEL_PATH"]) +work_dir = Path(os.environ["WORK_DIR"]) +tokenizer_cache_dir = os.environ["TOKENIZER_CACHE_DIR"] +chat_template_name = os.getenv("CHAT_TEMPLATE_NAME", "qwen3.5-vl") + +work_dir.mkdir(parents=True, exist_ok=True) +shutil.copy2(__file__, work_dir) + +# 256k-context SFT defaults for 32 GPUs (4 nodes x 8 GPUs). +sample_max_length = _get_int_env("SAMPLE_MAX_LENGTH", 256 * 1024) +pack_max_length = _get_int_env("PACK_MAX_LENGTH", 256 * 1024) +rand_video_max_frames = _get_int_env("RAND_VIDEO_MAX_FRAMES", 24) +num_workers = _get_int_env("NUM_WORKERS", 4) +global_batch_size = _get_int_env("GLOBAL_BATCH_SIZE", 8) +total_epoch = _get_int_env("TOTAL_EPOCH", 1) +hf_interval = _get_int_env("HF_INTERVAL", 500) +hf_max_keep = _get_int_env("HF_MAX_KEEP", 2) +checkpoint_interval = _get_int_env("CHECKPOINT_INTERVAL", 500) +checkpoint_maxkeep = _get_int_env("CHECKPOINT_MAXKEEP", 2) + +lr = _get_float_env("LR", 2e-5) +lr_min = _get_float_env("LR_MIN", 1e-6) +weight_decay = _get_float_env("WEIGHT_DECAY", 0.05) +warmup_ratio = _get_float_env("WARMUP_RATIO", 0.1) +recompute_ratio = _get_float_env("RECOMPUTE_RATIO", 1.0) +loss_reduction = os.getenv("LOSS_REDUCTION", "square") +max_pixels = _get_int_env("MAX_PIXELS", 16_777_216) + +sp_size = _get_int_env("SP_SIZE", 4) +ep_size = _get_int_env("EP_SIZE", 1) +tp_size = _get_int_env("TP_SIZE", 1) +torch_compile = _get_bool_env("TORCH_COMPILE", True) + +# Qwen3.5-35B-A3B model settings. +model_cfg = Qwen3_5_VLMoE35BA3Config() + +with (model_path / "config.json").open("r", encoding="utf-8") as file: + model_hf_config: dict[str, Any] = json.load(file) + +model_cfg.text_config.vocab_size = model_hf_config["text_config"]["vocab_size"] +# model_cfg.text_config.mtp_config = [ +# MTPConfig( +# name="normal", +# mask_type=None, +# num_layers=4, +# share_weights=True, +# loss_scaling_factor=1.0, +# ), +# ] + +if ep_size > 1: + model_cfg.text_config.ep_size = ep_size + model_cfg.text_config.dispatcher = "deepep" + +# Dataset recipe: META_DATA_PATH points to a metadata JSON file. +oss_loader_cfg = ( + OSSLoaderConfig(backend_kwargs={"conf_path": ceph_config}) + if ceph_config + else None +) + +ds_collections: dict[str, Any] = json.loads( + meta_data_path.read_text(encoding="utf-8") +) +has_pretrain = any( + data.get("text_pretrain", False) for data in ds_collections.values() +) +dataset_config: list[dict[str, Any]] = [] + +for name, data in ds_collections.items(): + is_pretrain = data.get("text_pretrain", False) + if is_pretrain: + tokenize_fn = PretrainTokenizeFunctionConfig(hash=data.get("hash")) + else: + tokenize_fn = Qwen3VLTokenizeFnConfig( + chat_template=chat_template_name, + llm_pack_weight=-3.2, + visual_pack_weight=5.0, + max_length=sample_max_length, + processor_path=str(model_path), + rand_video_max_frames=rand_video_max_frames, + oss_loader_cfg=oss_loader_cfg, + max_pixels=max_pixels, + debug=True, + ) + + dataset_config.append( + { + "dataset": DatasetConfig( + name=name, + anno_path=data["annotation"], + media_root=data.get("media_root") or "", + sample_ratio=data.get("sample_ratio", 1.0), + class_name="JsonlDataset" if is_pretrain else "VLMJsonlDataset", + enable_sequential_sampler=True, + cache_tag="xtuner_train_v2", + cache_dir=tokenizer_cache_dir, + ), + "tokenize_fn": tokenize_fn, + } + ) + +dataloader_config = DataloaderConfig( + dataset_config_list=dataset_config, + pack_max_length=pack_max_length, + pack_level="mllm_hybrid" if has_pretrain else "soft", + pack_to_max_length=True, + collator="qwen3_vl_sft_collator", + num_workers=num_workers, + pack_extra_buffer_size=_get_int_env("PACK_EXTRA_BUFFER_SIZE", 20), +) + +optim_cfg = MuonConfig(lr=lr, weight_decay=weight_decay) +lr_cfg = LRConfig( + lr_type="cosine", + warmup_ratio=warmup_ratio, + lr_min=lr_min, +) +fsdp_cfg = FSDPConfig( + tp_size=tp_size, + ep_size=ep_size, + recompute_ratio=recompute_ratio, + torch_compile=torch_compile, + checkpoint_preserve_rng_state=False, +) + +trainer = TrainerConfig( + sp_size=sp_size, + load_from=str(model_path), + resume_cfg=ResumeConfig(auto_resume=True), + tokenizer_path=str(model_path), + fsdp_cfg=fsdp_cfg, + exp_tracker="tensorboard", + model_cfg=model_cfg, + optim_cfg=optim_cfg, + dataloader_cfg=dataloader_config, + lr_cfg=lr_cfg, + loss_cfg=CELossConfig( + mode="chunk", + chunk_size=1024, + loss_reduction=loss_reduction, + ), + global_batch_size=global_batch_size, + total_epoch=total_epoch, + hf_interval=hf_interval, + checkpoint_interval=checkpoint_interval, + checkpoint_maxkeep=checkpoint_maxkeep, + hf_max_keep=hf_max_keep, + work_dir=work_dir, +) diff --git a/example_run/submit_qwen35_sft_16gpu_sp8.sh b/example_run/submit_qwen35_sft_16gpu_sp8.sh new file mode 100755 index 0000000000..e2ff377d07 --- /dev/null +++ b/example_run/submit_qwen35_sft_16gpu_sp8.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +# 16 GPUs = 2 nodes x 8 GPUs per node. +# +# Muon currently does not support EP > 1 in this XTuner version, so use +# EP=1 and shard the language model over all 16 ranks with FSDP. SP=8 splits +# every packed 256K sequence over 8 ranks (32K tokens per rank). +gpu_group="${GPU_GROUP:?Set GPU_GROUP to the rjob charged group}" +namespace="${NAMESPACE:?Set NAMESPACE to the rjob namespace}" +gpus_per_node=8 +num_nodes=2 + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +xtuner_path="$(cd -- "${script_dir}/.." && pwd)" +# The Python config reads all world-size-dependent values from the environment, +# so the same config can be used for both 16-GPU and 32-GPU submissions. +config_file="${script_dir}/sft_qwen35_35b_256k_32gpu.py" +meta_data_path="${META_DATA_PATH:-${script_dir}/meta.json}" + +model_path="${MODEL_PATH:?Set MODEL_PATH to the Qwen3.5 model snapshot directory}" +output_root="${OUTPUT_ROOT:?Set OUTPUT_ROOT to the training output directory}" +tokenizer_cache_dir="${TOKENIZER_CACHE_DIR:?Set TOKENIZER_CACHE_DIR to the tokenizer cache directory}" +log_dir="${LOG_DIR:-${output_root}/logs}" +ceph_config="${CEPH_CONFIG:-}" + +image="${IMAGE:?Set IMAGE to the training image}" +meta_name=$(basename "${meta_data_path}") +meta_name=${meta_name%.*} +run_tag="${meta_name}-16gpu-ep1-sp8-muon-$(date +%Y%m%d-%H%M%S)" +job_name="qwen35-sft-${run_tag}" +work_dir="${output_root}/${run_tag}" + +required_files=( + "${config_file}" + "${meta_data_path}" + "${model_path}/config.json" + "${model_path}/tokenizer_config.json" + "${model_path}/model.safetensors.index.json" +) +for required_file in "${required_files[@]}"; do + if [[ ! -f "${required_file}" ]]; then + echo "ERROR: required file does not exist: ${required_file}" >&2 + exit 2 + fi +done + +submit_mode_args=() +if [[ "${PREDICT_ONLY:-false}" == "true" ]]; then + submit_mode_args+=(--predict-only=true) +fi + +mount_args=() +if [[ -n "${RJOB_MOUNTS:-}" ]]; then + IFS=',' read -r -a mount_specs <<< "${RJOB_MOUNTS}" + for mount_spec in "${mount_specs[@]}"; do + [[ -n "${mount_spec}" ]] && mount_args+=(--mount="${mount_spec}") + done +fi + +rjob submit \ + "${submit_mode_args[@]}" \ + --name="${job_name}" \ + --task_name t0 \ + --gpu="${gpus_per_node}" \ + --memory=1500000 \ + --cpu=50 \ + --charged-group="${gpu_group}" \ + --namespace="${namespace}" \ + --private-machine=group \ + -P "${num_nodes}" \ + --image="${image}" \ + "${mount_args[@]}" \ + --host-network=true \ + --gang-start=true \ + --custom-resources=rdma/mlnx_shared=8 \ + --custom-resources=mellanox.com/mlnx_rdma=1 \ + -e DISTRIBUTED_JOB=true \ + -e XTUNER_PATH="${xtuner_path}" \ + -e CONFIG_FILE="${config_file}" \ + -e MODEL_PATH="${model_path}" \ + -e META_DATA_PATH="${meta_data_path}" \ + -e CEPH_CONFIG="${ceph_config}" \ + -e WORK_DIR="${work_dir}" \ + -e TOKENIZER_CACHE_DIR="${tokenizer_cache_dir}" \ + -e XTUNER_TOKENIZE_DEBUG_SAMPLES="${XTUNER_TOKENIZE_DEBUG_SAMPLES:-0}" \ + -e LOG_DIR="${log_dir}" \ + -e GPUS_PER_NODE="${gpus_per_node}" \ + -e TORCHRUN_NNODES="${num_nodes}" \ + -e SAMPLE_MAX_LENGTH=262144 \ + -e PACK_MAX_LENGTH=262144 \ + -e GLOBAL_BATCH_SIZE=8 \ + -e SP_SIZE=8 \ + -e TP_SIZE=1 \ + -e EP_SIZE=1 \ + -e NUM_WORKERS=4 \ + -e PACK_EXTRA_BUFFER_SIZE=20 \ + -e RAND_VIDEO_MAX_FRAMES=24 \ + -e MAX_PIXELS=16777216 \ + -e LR=2e-5 \ + -e LR_MIN=1e-6 \ + -e WEIGHT_DECAY=0.05 \ + -e WARMUP_RATIO=0.1 \ + -e RECOMPUTE_RATIO=1.0 \ + -e LOSS_REDUCTION=square \ + -e TORCH_COMPILE=true \ + -e TOTAL_EPOCH=1 \ + -e HF_INTERVAL=500 \ + -e HF_MAX_KEEP=2 \ + -e CHECKPOINT_INTERVAL=500 \ + -e CHECKPOINT_MAXKEEP=2 \ + -- bash -lc ' + set -euo pipefail + set -x + + export PYTHONPATH="${XTUNER_PATH}:${PYTHONPATH:-}" + export TORCHRUN_NODE_RANK="${NODE_RANK:-${RANK:-}}" + export MASTER_PORT="${MASTER_PORT:-29500}" + export LOG_FILE="${LOG_DIR}/qwen35-sft-16gpu-ep1-sp8-muon-node${TORCHRUN_NODE_RANK:-unknown}.log" + + if [ -z "${TORCHRUN_NODE_RANK}" ] || [ -z "${MASTER_ADDR:-}" ]; then + echo "ERROR: NODE_RANK/RANK and MASTER_ADDR are required." + env | sort | grep -E \ + "^(NODE_RANK|RANK|NODE_COUNT|WORLD_SIZE|MASTER_ADDR|MASTER_PORT|HOSTNAME)=" \ + || true + exit 2 + fi + + if [ "${TORCHRUN_NNODES}" != "1" ] && { + [ "${MASTER_ADDR}" = "127.0.0.1" ] || + [ "${MASTER_ADDR}" = "localhost" ] + }; then + echo "ERROR: MASTER_ADDR=${MASTER_ADDR} is invalid for a multi-node job." + exit 2 + fi + + mkdir -p "${LOG_DIR}" + exec > >(tee "${LOG_FILE}") 2>&1 + + cd "${XTUNER_PATH}" + ls -l "${CONFIG_FILE}" "${META_DATA_PATH}" + python -c "import xtuner; print(\"xtuner import:\", xtuner.__file__)" + + torchrun \ + --nproc-per-node="${GPUS_PER_NODE}" \ + --nnodes="${TORCHRUN_NNODES}" \ + --node_rank="${TORCHRUN_NODE_RANK}" \ + --master_addr="${MASTER_ADDR}" \ + --master_port="${MASTER_PORT}" \ + xtuner/v1/train/cli/sft.py \ + --config="${CONFIG_FILE}" + ' diff --git a/example_run/submit_qwen35_sft_32gpu.sh b/example_run/submit_qwen35_sft_32gpu.sh new file mode 100755 index 0000000000..59b0c6e653 --- /dev/null +++ b/example_run/submit_qwen35_sft_32gpu.sh @@ -0,0 +1,141 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +# 32 GPUs = 4 nodes x 8 GPUs per node. +gpu_group="${GPU_GROUP:?Set GPU_GROUP to the rjob charged group}" +namespace="${NAMESPACE:?Set NAMESPACE to the rjob namespace}" +gpus_per_node=8 +num_nodes=4 + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +xtuner_path="$(cd -- "${script_dir}/.." && pwd)" +config_file="${script_dir}/sft_qwen35_35b_256k_32gpu.py" +meta_data_path="${META_DATA_PATH:-${script_dir}/meta.json}" + +model_path="${MODEL_PATH:?Set MODEL_PATH to the model snapshot directory}" +output_root="${OUTPUT_ROOT:?Set OUTPUT_ROOT to the training output directory}" +tokenizer_cache_dir="${TOKENIZER_CACHE_DIR:?Set TOKENIZER_CACHE_DIR to the tokenizer cache directory}" +log_dir="${LOG_DIR:-${output_root}/logs}" +ceph_config="${CEPH_CONFIG:-}" + +image="${IMAGE:?Set IMAGE to the training image}" +meta_name=$(basename "${meta_data_path}") +meta_name=${meta_name%.*} +run_tag="${meta_name}-$(date +%Y%m%d-%H%M)" +job_name="qwen35-sft-${run_tag}" +work_dir="${output_root}/${run_tag}" + +required_files=( + "${config_file}" + "${meta_data_path}" + "${model_path}/config.json" + "${model_path}/tokenizer_config.json" + "${model_path}/model.safetensors.index.json" +) +for required_file in "${required_files[@]}"; do + if [[ ! -f "${required_file}" ]]; then + echo "ERROR: required file does not exist: ${required_file}" >&2 + exit 2 + fi +done + +mount_args=() +if [[ -n "${RJOB_MOUNTS:-}" ]]; then + IFS=',' read -r -a mount_specs <<< "${RJOB_MOUNTS}" + for mount_spec in "${mount_specs[@]}"; do + [[ -n "${mount_spec}" ]] && mount_args+=(--mount="${mount_spec}") + done +fi + +rjob submit \ + --name="${job_name}" \ + --task_name t0 \ + --gpu="${gpus_per_node}" \ + --memory=1500000 \ + --cpu=50 \ + --charged-group="${gpu_group}" \ + --namespace="${namespace}" \ + --private-machine=group \ + -P "${num_nodes}" \ + --image="${image}" \ + "${mount_args[@]}" \ + --host-network=true \ + --gang-start=true \ + --custom-resources=rdma/mlnx_shared=8 \ + --custom-resources=mellanox.com/mlnx_rdma=1 \ + -e DISTRIBUTED_JOB=true \ + -e XTUNER_PATH="${xtuner_path}" \ + -e CONFIG_FILE="${config_file}" \ + -e MODEL_PATH="${model_path}" \ + -e META_DATA_PATH="${meta_data_path}" \ + -e CEPH_CONFIG="${ceph_config}" \ + -e WORK_DIR="${work_dir}" \ + -e TOKENIZER_CACHE_DIR="${tokenizer_cache_dir}" \ + -e XTUNER_TOKENIZE_DEBUG_SAMPLES="${XTUNER_TOKENIZE_DEBUG_SAMPLES:-0}" \ + -e LOG_DIR="${log_dir}" \ + -e GPUS_PER_NODE="${gpus_per_node}" \ + -e TORCHRUN_NNODES="${num_nodes}" \ + -e SAMPLE_MAX_LENGTH=262144 \ + -e PACK_MAX_LENGTH=262144 \ + -e GLOBAL_BATCH_SIZE=8 \ + -e SP_SIZE=8 \ + -e TP_SIZE=1 \ + -e EP_SIZE=1 \ + -e NUM_WORKERS=4 \ + -e PACK_EXTRA_BUFFER_SIZE=20 \ + -e RAND_VIDEO_MAX_FRAMES=24 \ + -e MAX_PIXELS=16777216 \ + -e LR=2e-5 \ + -e LR_MIN=1e-6 \ + -e WEIGHT_DECAY=0.05 \ + -e WARMUP_RATIO=0.1 \ + -e RECOMPUTE_RATIO=1.0 \ + -e LOSS_REDUCTION=square \ + -e TORCH_COMPILE=true \ + -e TOTAL_EPOCH=1 \ + -e HF_INTERVAL=500 \ + -e HF_MAX_KEEP=2 \ + -e CHECKPOINT_INTERVAL=500 \ + -e CHECKPOINT_MAXKEEP=2 \ + -- bash -lc ' + set -euo pipefail + set -x + + export PYTHONPATH="${XTUNER_PATH}:${PYTHONPATH:-}" + export TORCHRUN_NODE_RANK="${NODE_RANK:-${RANK:-}}" + export MASTER_PORT="${MASTER_PORT:-29500}" + export LOG_FILE="${LOG_DIR}/qwen35-sft-node${TORCHRUN_NODE_RANK:-unknown}.log" + + if [ -z "${TORCHRUN_NODE_RANK}" ] || [ -z "${MASTER_ADDR:-}" ]; then + echo "ERROR: NODE_RANK/RANK and MASTER_ADDR are required." + env | sort | grep -E \ + "^(NODE_RANK|RANK|NODE_COUNT|WORLD_SIZE|MASTER_ADDR|MASTER_PORT|HOSTNAME)=" \ + || true + exit 2 + fi + + if [ "${TORCHRUN_NNODES}" != "1" ] && { + [ "${MASTER_ADDR}" = "127.0.0.1" ] || + [ "${MASTER_ADDR}" = "localhost" ] + }; then + echo "ERROR: MASTER_ADDR=${MASTER_ADDR} is invalid for a multi-node job." + exit 2 + fi + + mkdir -p "${LOG_DIR}" + exec > >(tee "${LOG_FILE}") 2>&1 + + cd "${XTUNER_PATH}" + ls -l "${CONFIG_FILE}" "${META_DATA_PATH}" + python -c "import xtuner; print(\"xtuner import:\", xtuner.__file__)" + + torchrun \ + --nproc-per-node="${GPUS_PER_NODE}" \ + --nnodes="${TORCHRUN_NNODES}" \ + --node_rank="${TORCHRUN_NODE_RANK}" \ + --master_addr="${MASTER_ADDR}" \ + --master_port="${MASTER_PORT}" \ + xtuner/v1/train/cli/sft.py \ + --config="${CONFIG_FILE}" + ' diff --git a/example_run/submit_qwen36_sft_16gpu_sp8.sh b/example_run/submit_qwen36_sft_16gpu_sp8.sh new file mode 100755 index 0000000000..9ec9011324 --- /dev/null +++ b/example_run/submit_qwen36_sft_16gpu_sp8.sh @@ -0,0 +1,160 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +# 16 GPUs = 2 nodes x 8 GPUs per node. +# +# Qwen3.6 reuses the existing Qwen3.5 XTuner model/tokenize implementation. +# MTP stays disabled in sft_qwen35_35b_256k_32gpu.py, so only the compatible +# Qwen3.6 main model, vision tower, and projector weights are trained. +# +# Muon currently does not support EP > 1 in this XTuner version, so use +# EP=1 and shard the language model over all 16 ranks with FSDP. SP=8 splits +# every packed 256K sequence over 8 ranks (32K tokens per rank). +gpu_group="${GPU_GROUP:?Set GPU_GROUP to the rjob charged group}" +namespace="${NAMESPACE:?Set NAMESPACE to the rjob namespace}" +gpus_per_node=8 +num_nodes=2 + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +xtuner_path="$(cd -- "${script_dir}/.." && pwd)" +# The Python config reads all world-size-dependent values from the environment. +# Its Qwen3.5 name is retained intentionally: Qwen3.6 has the same non-MTP +# architecture; CHAT_TEMPLATE_NAME selects the dedicated internal Qwen3.6 +# renderer without changing the remaining model or training path. +config_file="${script_dir}/sft_qwen35_35b_256k_32gpu.py" +meta_data_path="${META_DATA_PATH:-${script_dir}/meta.json}" + +model_path="${MODEL_PATH:?Set MODEL_PATH to the Qwen3.6 model snapshot directory}" +output_root="${OUTPUT_ROOT:?Set OUTPUT_ROOT to the training output directory}" +tokenizer_cache_dir="${TOKENIZER_CACHE_DIR:?Set TOKENIZER_CACHE_DIR to the tokenizer cache directory}" +log_dir="${LOG_DIR:-${output_root}/logs}" +ceph_config="${CEPH_CONFIG:-}" + +image="${IMAGE:?Set IMAGE to the training image}" +meta_name=$(basename "${meta_data_path}") +meta_name=${meta_name%.*} +run_tag="${meta_name}-16gpu-ep1-sp8-muon-$(date +%Y%m%d-%H%M%S)" +job_name="qwen36-sft-no-mtp-${run_tag}" +work_dir="${output_root}/${run_tag}" + +required_files=( + "${config_file}" + "${meta_data_path}" + "${model_path}/config.json" + "${model_path}/tokenizer_config.json" + "${model_path}/model.safetensors.index.json" +) +for required_file in "${required_files[@]}"; do + if [[ ! -f "${required_file}" ]]; then + echo "ERROR: required file does not exist: ${required_file}" >&2 + exit 2 + fi +done + +submit_mode_args=() +if [[ "${PREDICT_ONLY:-false}" == "true" ]]; then + submit_mode_args+=(--predict-only=true) +fi + +mount_args=() +if [[ -n "${RJOB_MOUNTS:-}" ]]; then + IFS=',' read -r -a mount_specs <<< "${RJOB_MOUNTS}" + for mount_spec in "${mount_specs[@]}"; do + [[ -n "${mount_spec}" ]] && mount_args+=(--mount="${mount_spec}") + done +fi + +rjob submit \ + "${submit_mode_args[@]}" \ + --name="${job_name}" \ + --task_name t0 \ + --gpu="${gpus_per_node}" \ + --memory=1500000 \ + --cpu=50 \ + --charged-group="${gpu_group}" \ + --namespace="${namespace}" \ + --private-machine=group \ + -P "${num_nodes}" \ + --image="${image}" \ + "${mount_args[@]}" \ + --host-network=true \ + --gang-start=true \ + --custom-resources=rdma/mlnx_shared=8 \ + --custom-resources=mellanox.com/mlnx_rdma=1 \ + -e DISTRIBUTED_JOB=true \ + -e XTUNER_PATH="${xtuner_path}" \ + -e CONFIG_FILE="${config_file}" \ + -e MODEL_PATH="${model_path}" \ + -e META_DATA_PATH="${meta_data_path}" \ + -e CEPH_CONFIG="${ceph_config}" \ + -e WORK_DIR="${work_dir}" \ + -e TOKENIZER_CACHE_DIR="${tokenizer_cache_dir}" \ + -e CHAT_TEMPLATE_NAME=qwen3.6-vl \ + -e XTUNER_TOKENIZE_DEBUG_SAMPLES="${XTUNER_TOKENIZE_DEBUG_SAMPLES:-0}" \ + -e LOG_DIR="${log_dir}" \ + -e GPUS_PER_NODE="${gpus_per_node}" \ + -e TORCHRUN_NNODES="${num_nodes}" \ + -e SAMPLE_MAX_LENGTH=262144 \ + -e PACK_MAX_LENGTH=262144 \ + -e GLOBAL_BATCH_SIZE=8 \ + -e SP_SIZE=8 \ + -e TP_SIZE=1 \ + -e EP_SIZE=1 \ + -e NUM_WORKERS=4 \ + -e PACK_EXTRA_BUFFER_SIZE=20 \ + -e RAND_VIDEO_MAX_FRAMES=24 \ + -e MAX_PIXELS=16777216 \ + -e LR=2e-5 \ + -e LR_MIN=1e-6 \ + -e WEIGHT_DECAY=0.05 \ + -e WARMUP_RATIO=0.1 \ + -e RECOMPUTE_RATIO=1.0 \ + -e LOSS_REDUCTION=square \ + -e TORCH_COMPILE=true \ + -e TOTAL_EPOCH=1 \ + -e HF_INTERVAL=500 \ + -e HF_MAX_KEEP=2 \ + -e CHECKPOINT_INTERVAL=500 \ + -e CHECKPOINT_MAXKEEP=2 \ + -- bash -lc ' + set -euo pipefail + set -x + + export PYTHONPATH="${XTUNER_PATH}:${PYTHONPATH:-}" + export TORCHRUN_NODE_RANK="${NODE_RANK:-${RANK:-}}" + export MASTER_PORT="${MASTER_PORT:-29500}" + export LOG_FILE="${LOG_DIR}/qwen36-sft-no-mtp-16gpu-ep1-sp8-muon-node${TORCHRUN_NODE_RANK:-unknown}.log" + + if [ -z "${TORCHRUN_NODE_RANK}" ] || [ -z "${MASTER_ADDR:-}" ]; then + echo "ERROR: NODE_RANK/RANK and MASTER_ADDR are required." + env | sort | grep -E \ + "^(NODE_RANK|RANK|NODE_COUNT|WORLD_SIZE|MASTER_ADDR|MASTER_PORT|HOSTNAME)=" \ + || true + exit 2 + fi + + if [ "${TORCHRUN_NNODES}" != "1" ] && { + [ "${MASTER_ADDR}" = "127.0.0.1" ] || + [ "${MASTER_ADDR}" = "localhost" ] + }; then + echo "ERROR: MASTER_ADDR=${MASTER_ADDR} is invalid for a multi-node job." + exit 2 + fi + + mkdir -p "${LOG_DIR}" + exec > >(tee "${LOG_FILE}") 2>&1 + + cd "${XTUNER_PATH}" + ls -l "${CONFIG_FILE}" "${META_DATA_PATH}" + python -c "import xtuner; print(\"xtuner import:\", xtuner.__file__)" + + torchrun \ + --nproc-per-node="${GPUS_PER_NODE}" \ + --nnodes="${TORCHRUN_NNODES}" \ + --node_rank="${TORCHRUN_NODE_RANK}" \ + --master_addr="${MASTER_ADDR}" \ + --master_port="${MASTER_PORT}" \ + xtuner/v1/train/cli/sft.py \ + --config="${CONFIG_FILE}" + ' diff --git a/example_run/tokenize_sample_viewer.py b/example_run/tokenize_sample_viewer.py new file mode 100644 index 0000000000..72c7afc696 --- /dev/null +++ b/example_run/tokenize_sample_viewer.py @@ -0,0 +1,797 @@ +#!/usr/bin/env python3 +"""Read-only browser for XTuner tokenize debug samples. + +The page only receives natural-language fields. Token IDs, integer labels, +token pieces, and other tensor payloads are intentionally removed by the +server before a sample is returned to the browser. +""" + +from __future__ import annotations + +import argparse +import json +import mimetypes +import re +import sys +from dataclasses import dataclass +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + + +DEFAULT_INPUT = Path("tokenize_debug_samples") + + +@dataclass(frozen=True) +class SampleRef: + key: str + label: str + path: Path | None = None + offset: int | None = None + length: int | None = None + + +class SampleStore: + """Index a debug-sample directory, one JSON file, or a JSONL file.""" + + def __init__(self, input_path: Path) -> None: + self.input_path = input_path.expanduser().resolve() + if not self.input_path.exists(): + raise FileNotFoundError(f"input does not exist: {self.input_path}") + + self._refs: list[SampleRef] = [] + self._by_key: dict[str, SampleRef] = {} + if self.input_path.is_dir(): + self._index_directory() + self.kind = "directory" + elif self.input_path.suffix.lower() == ".jsonl": + self._index_jsonl() + self.kind = "jsonl" + elif self.input_path.suffix.lower() == ".json": + self._index_json() + self.kind = "json" + else: + raise ValueError("input must be a sample directory, .json, or .jsonl file") + + if not self._refs: + raise ValueError(f"no samples found in {self.input_path}") + + @staticmethod + def _key_from_filename(path: Path) -> str: + match = re.fullmatch(r"sample_(\d+)", path.stem) + return str(int(match.group(1))) if match else path.stem + + def _add_ref(self, ref: SampleRef) -> None: + key = ref.key + if key in self._by_key: + suffix = 2 + while f"{key}-{suffix}" in self._by_key: + suffix += 1 + ref = SampleRef( + key=f"{key}-{suffix}", + label=f"{ref.label} ({suffix})", + path=ref.path, + offset=ref.offset, + length=ref.length, + ) + self._refs.append(ref) + self._by_key[ref.key] = ref + + def _index_directory(self) -> None: + files: list[Path] = [] + manifest_path = self.input_path / "manifest.json" + if manifest_path.is_file(): + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + for filename in manifest.get("files", []): + if not isinstance(filename, str) or Path(filename).name != filename: + continue + path = (self.input_path / filename).resolve() + if path.parent == self.input_path and path.is_file(): + files.append(path) + if not files: + files = sorted(self.input_path.glob("sample_*.json")) + + for path in files: + key = self._key_from_filename(path) + self._add_ref(SampleRef(key=key, label=f"sample {key}", path=path)) + + def _index_json(self) -> None: + record = json.loads(self.input_path.read_text(encoding="utf-8")) + index = record.get("source", {}).get("sample_index", 0) + key = str(index) + self._add_ref(SampleRef(key=key, label=f"sample {key}", path=self.input_path)) + + def _index_jsonl(self) -> None: + with self.input_path.open("rb") as file: + line_number = 0 + while True: + offset = file.tell() + line = file.readline() + if not line: + break + if not line.strip(): + continue + self._add_ref( + SampleRef( + key=str(line_number), + label=f"line {line_number}", + offset=offset, + length=len(line), + ) + ) + line_number += 1 + + def list_samples(self) -> list[dict[str, str]]: + return [{"key": ref.key, "label": ref.label} for ref in self._refs] + + def read(self, key: str) -> dict[str, Any]: + try: + ref = self._by_key[key] + except KeyError as error: + raise KeyError(f"unknown sample key: {key}") from error + + if ref.path is not None: + return json.loads(ref.path.read_text(encoding="utf-8")) + + assert ref.offset is not None and ref.length is not None + with self.input_path.open("rb") as file: + file.seek(ref.offset) + payload = file.read(ref.length) + return json.loads(payload) + + def public_info(self) -> dict[str, Any]: + return { + "input": str(self.input_path), + "kind": self.kind, + "sample_count": len(self._refs), + } + + +def _natural_record(record: dict[str, Any]) -> dict[str, Any]: + """Remove all token-ID-shaped data before sending a record to the UI.""" + + rendered = record.get("rendered") or {} + rendered_spans = [] + for span in rendered.get("loss_character_spans") or []: + rendered_spans.append( + { + "start": span.get("start"), + "end": span.get("end"), + } + ) + + loss = record.get("loss") or {} + loss_spans = [] + for position, span in enumerate(loss.get("token_spans") or [], start=1): + loss_spans.append( + { + "number": position, + "start": span.get("start"), + "end": span.get("end"), + "decoded_text": span.get("decoded_text", ""), + } + ) + + tokenized = record.get("tokenized") or {} + raw_data = record.get("raw_data") or {} + return { + "source": record.get("source") or {}, + "raw_data": { + "messages": raw_data.get("messages") or [], + "tools": raw_data.get("tools") or [], + }, + "rendered": { + "text": rendered.get("text", ""), + "character_count": rendered.get("character_count"), + "loss_character_count": rendered.get("loss_character_count"), + "loss_character_spans": rendered_spans, + }, + "tokenized": { + "token_count": tokenized.get("token_count"), + "decoded_text": tokenized.get("decoded_text", ""), + }, + "loss": { + "token_count": loss.get("token_count"), + "masked_token_count": loss.get("masked_token_count"), + "token_spans": loss_spans, + }, + "error": record.get("error"), + } + + +HTML = r""" + + + + + XTuner Sample Viewer + + + +
+ XTuner Sample Viewer + + + + +
+ +
+
正在读取样本列表…
+ + + +
+ + + + +""" + + +class ViewerHandler(BaseHTTPRequestHandler): + store: SampleStore + + def log_message(self, format_string: str, *args: Any) -> None: + sys.stderr.write(f"[viewer] {self.address_string()} {format_string % args}\n") + + def _send_bytes( + self, + payload: bytes, + *, + content_type: str, + status: HTTPStatus = HTTPStatus.OK, + ) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(payload))) + self.send_header("Cache-Control", "no-store") + self.send_header("X-Content-Type-Options", "nosniff") + self.end_headers() + self.wfile.write(payload) + + def _send_json(self, data: Any, status: HTTPStatus = HTTPStatus.OK) -> None: + payload = json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + self._send_bytes(payload, content_type="application/json; charset=utf-8", status=status) + + def _send_text(self, text: str, status: HTTPStatus) -> None: + self._send_bytes( + text.encode("utf-8"), + content_type="text/plain; charset=utf-8", + status=status, + ) + + def do_GET(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API + parsed = urlparse(self.path) + if parsed.path == "/": + self._send_bytes(HTML.encode("utf-8"), content_type="text/html; charset=utf-8") + return + + if parsed.path == "/api/samples": + self._send_json( + { + "store": self.store.public_info(), + "samples": self.store.list_samples(), + } + ) + return + + prefix = "/api/sample/" + if parsed.path.startswith(prefix): + key = unquote(parsed.path[len(prefix) :]) + try: + record = self.store.read(key) + except KeyError as error: + self._send_text(str(error), HTTPStatus.NOT_FOUND) + return + except (OSError, ValueError, json.JSONDecodeError) as error: + self._send_text(str(error), HTTPStatus.INTERNAL_SERVER_ERROR) + return + self._send_json(_natural_record(record)) + return + + if parsed.path == "/favicon.ico": + self._send_bytes(b"", content_type=mimetypes.types_map.get(".ico", "image/x-icon")) + return + self._send_text("not found", HTTPStatus.NOT_FOUND) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Serve a natural-text-only viewer for XTuner tokenize debug samples." + ) + parser.add_argument( + "--input", + type=Path, + default=DEFAULT_INPUT, + help="sample_*.json directory, one sample JSON, or a JSONL file", + ) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8765) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + store = SampleStore(args.input) + except (OSError, ValueError, json.JSONDecodeError) as error: + print(f"ERROR: {error}", file=sys.stderr) + return 2 + + handler = type("ConfiguredViewerHandler", (ViewerHandler,), {"store": store}) + server = ThreadingHTTPServer((args.host, args.port), handler) + print(f"Reading: {store.input_path}") + print(f"Samples: {len(store.list_samples())}") + print(f"Open: http://{args.host}:{args.port}") + try: + server.serve_forever() + except KeyboardInterrupt: + print("\nStopped.") + finally: + server.server_close() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/datasets/test_tokenize_debug.py b/tests/datasets/test_tokenize_debug.py new file mode 100644 index 0000000000..c852bdf785 --- /dev/null +++ b/tests/datasets/test_tokenize_debug.py @@ -0,0 +1,171 @@ +import json +from types import SimpleNamespace + +import torch + +from xtuner.v1.datasets.tokenize_debug import ( + _build_record, + _sample_indices, + maybe_dump_tokenize_debug_samples, +) + + +class _FakeTokenizer: + def convert_ids_to_tokens(self, token_ids): + return [f"token_{token_id}" for token_id in token_ids] + + def decode(self, token_ids, **kwargs): + return "|".join(str(token_id) for token_id in token_ids) + + +class _FakeDataset: + media_root = "" + + def __init__(self, path): + self.path = str(path) + self.offsets = [0] + + def __len__(self): + return 1 + + +class _FakeTokenizeFunction: + chat_template_name = "qwen3.5-vl" + chat_template = SimpleNamespace(default_system=None) + add_vision_id = False + tokenizer = _FakeTokenizer() + state = "runtime" + + def set_state(self, state): + self.state = state + + def __call__(self, raw_data, **kwargs): + return { + "input_ids": [10, 11, 12], + "labels": [-100, 11, 12], + "num_tokens": 3, + } + + +class _FakeQwen36TokenizeFunction(_FakeTokenizeFunction): + chat_template_name = "qwen3.6-vl" + + +def test_build_record_reverse_decodes_contiguous_loss_spans(): + record = _build_record( + tokenizer=_FakeTokenizer(), + raw_data={"messages": [{"role": "assistant", "content": "AB"}]}, + rendered_text="promptAB", + character_loss_mask=[False] * 6 + [True, True], + tokenized_data={ + "input_ids": [10, 11, 12, 13, 14], + "labels": [-100, 11, 12, -100, 14], + "position_ids": torch.zeros(3, 5), + "num_tokens": 5, + }, + dataset_name="demo", + dataset_path="/tmp/demo.jsonl", + sample_index=0, + ) + + assert record["rendered"]["loss_character_spans"] == [ + {"start": 6, "end": 8, "text": "AB"} + ] + assert record["loss"]["token_count"] == 3 + assert record["loss"]["token_spans"] == [ + { + "start": 1, + "end": 3, + "token_ids": [11, 12], + "token_pieces": ["token_11", "token_12"], + "decoded_text": "11|12", + }, + { + "start": 4, + "end": 5, + "token_ids": [14], + "token_pieces": ["token_14"], + "decoded_text": "14", + }, + ] + assert record["tokenized"]["other_fields"]["position_ids"] == { + "kind": "torch.Tensor", + "shape": [3, 5], + "dtype": "torch.float32", + } + + +def test_sample_indices_are_seeded_unique_and_reproducible(): + first = _sample_indices(dataset_size=100, sample_count=10) + second = _sample_indices(dataset_size=100, sample_count=10) + + assert first == second + assert first != list(range(10)) + assert len(first) == len(set(first)) == 10 + assert first == sorted(first) + + +def test_single_environment_variable_dumps_one_sample(monkeypatch, tmp_path): + data_path = tmp_path / "demo.jsonl" + data_path.write_text( + json.dumps( + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + } + ) + + "\n" + ) + monkeypatch.setenv("XTUNER_TOKENIZE_DEBUG_SAMPLES", "1") + monkeypatch.setenv("WORK_DIR", str(tmp_path / "work_dir")) + + maybe_dump_tokenize_debug_samples( + dataset=_FakeDataset(data_path), + tokenize_fn=_FakeTokenizeFunction(), + dataset_name="demo", + ) + + sample_files = list((tmp_path / "work_dir" / "tokenize_debug_samples").rglob("sample_000000.json")) + assert len(sample_files) == 1 + record = json.loads(sample_files[0].read_text()) + assert record["raw_data"]["messages"][0]["content"] == "question" + assert record["rendered"]["text"].startswith("<|im_start|>user\nquestion") + assert record["loss"]["token_spans"][0]["decoded_text"] == "11|12" + + manifest_files = list((tmp_path / "work_dir" / "tokenize_debug_samples").rglob("manifest.json")) + assert len(manifest_files) == 1 + manifest = json.loads(manifest_files[0].read_text()) + assert manifest["sampling"] == "random_without_replacement" + assert manifest["random_seed"] == 42 + assert manifest["sample_indices"] == [0] + + +def test_qwen36_debug_uses_template_specific_output_directory(monkeypatch, tmp_path): + data_path = tmp_path / "demo.jsonl" + data_path.write_text( + json.dumps( + { + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "answer"}, + ] + } + ) + + "\n" + ) + monkeypatch.setenv("XTUNER_TOKENIZE_DEBUG_SAMPLES", "1") + monkeypatch.setenv("WORK_DIR", str(tmp_path / "work_dir")) + + maybe_dump_tokenize_debug_samples( + dataset=_FakeDataset(data_path), + tokenize_fn=_FakeQwen36TokenizeFunction(), + dataset_name="demo", + ) + + manifest_path = next( + (tmp_path / "work_dir" / "tokenize_debug_samples" / "qwen3.6-vl").rglob("manifest.json") + ) + manifest = json.loads(manifest_path.read_text()) + assert manifest["chat_template_name"] == "qwen3.6-vl" diff --git a/tests/test_qwen36_chat.py b/tests/test_qwen36_chat.py new file mode 100644 index 0000000000..40cf390d51 --- /dev/null +++ b/tests/test_qwen36_chat.py @@ -0,0 +1,54 @@ +from xtuner.v1.data_proto.messages.qwen35_chat import qwen35_tokenize_fn_fastspeed +from xtuner.v1.data_proto.messages.qwen36_chat import qwen36_tokenize_fn_fastspeed + + +def _render(render_fn, arguments): + messages = [ + {"role": "user", "content": "call the tool"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "demo", "arguments": arguments}}], + }, + ] + text, _ = render_fn(messages, return_labels=False) + return text + + +def test_qwen36_only_changes_tool_argument_serialization(): + arguments = { + "string": "raw text", + "truth": True, + "falsehood": False, + "nothing": None, + "number": 3, + "mapping": {"enabled": True}, + "sequence": [False, None], + } + + qwen35_text = _render(qwen35_tokenize_fn_fastspeed, arguments) + qwen36_text = _render(qwen36_tokenize_fn_fastspeed, arguments) + + assert qwen35_text.replace("True", "true").replace("False", "false").replace("None", "null") == qwen36_text + assert "\nraw text\n" in qwen36_text + assert "\ntrue\n" in qwen36_text + assert "\nfalse\n" in qwen36_text + assert "\nnull\n" in qwen36_text + assert '\n{"enabled": true}\n' in qwen36_text + assert "\n[false, null]\n" in qwen36_text + + +def test_qwen36_preserves_qwen35_non_tool_rendering(): + messages = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "first question"}, + {"role": "assistant", "reasoning_content": "first thought", "content": "first answer"}, + {"role": "user", "content": "second question"}, + {"role": "assistant", "reasoning_content": "second thought", "content": "second answer"}, + ] + + qwen35_text, qwen35_loss_mask = qwen35_tokenize_fn_fastspeed(messages, return_labels=False) + qwen36_text, qwen36_loss_mask = qwen36_tokenize_fn_fastspeed(messages, return_labels=False) + + assert qwen36_text == qwen35_text + assert qwen36_loss_mask == qwen35_loss_mask diff --git a/xtuner/v1/data_proto/messages/__init__.py b/xtuner/v1/data_proto/messages/__init__.py index 9b95ad8296..ffa5f0dae1 100644 --- a/xtuner/v1/data_proto/messages/__init__.py +++ b/xtuner/v1/data_proto/messages/__init__.py @@ -2,6 +2,7 @@ from .base import BaseMessages from .chat import ChatMessages from .qwen35_chat import Qwen35ChatMessages +from .qwen36_chat import Qwen36ChatMessages -__all__ = ["BaseMessages", "ChatMessages", "Qwen35ChatMessages"] +__all__ = ["BaseMessages", "ChatMessages", "Qwen35ChatMessages", "Qwen36ChatMessages"] diff --git a/xtuner/v1/data_proto/messages/qwen35_chat.py b/xtuner/v1/data_proto/messages/qwen35_chat.py index fccb2f95c2..7167d8ac3f 100644 --- a/xtuner/v1/data_proto/messages/qwen35_chat.py +++ b/xtuner/v1/data_proto/messages/qwen35_chat.py @@ -138,6 +138,7 @@ def qwen35_tokenize_fn_fastspeed( add_generation_prompt=False, add_vision_id=False, return_labels=True, + _tool_call_args_renderer=_render_tool_call_args, ): enable_thinking = any("reasoning_content" in msg for msg in messages) @@ -248,7 +249,7 @@ def _append(text: str, is_loss: bool) -> None: _append(f"\n\n\n", body_is_loss) if isinstance(tc_args, dict): - _append(_render_tool_call_args(tc_args), body_is_loss) + _append(_tool_call_args_renderer(tc_args), body_is_loss) _append("\n", body_is_loss) _append("<|im_end|>\n", body_is_loss) @@ -382,6 +383,11 @@ class Qwen35ChatMessages(BaseModel): messages: List[dict] # 暂时不做校验 tools: Optional[List[Dict]] = None + def _tokenize_chat(self, tokenizer, add_vision_id): + return qwen35_tokenize_fn_fastspeed( + self.messages, tokenizer, self.tools, add_vision_id=add_vision_id, return_labels=True + ) + def tokenize( self, tokenizer: PreTrainedTokenizer, chat_template: HybridChatTemplate, add_vision_id=False, **kwargs ) -> Dict: @@ -407,7 +413,5 @@ def tokenize( else: self.messages.insert(0, {"role": "system", "content": chat_template.default_system}) - token_ids, label_ids = qwen35_tokenize_fn_fastspeed( - self.messages, tokenizer, self.tools, add_vision_id=add_vision_id, return_labels=True - ) + token_ids, label_ids = self._tokenize_chat(tokenizer, add_vision_id) return {"input_ids": token_ids, "labels": label_ids} diff --git a/xtuner/v1/data_proto/messages/qwen36_chat.py b/xtuner/v1/data_proto/messages/qwen36_chat.py new file mode 100644 index 0000000000..54fdf774f3 --- /dev/null +++ b/xtuner/v1/data_proto/messages/qwen36_chat.py @@ -0,0 +1,73 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import json +from typing import Dict, List + +from xtuner.v1.data_proto.messages.qwen35_chat import ( + Qwen35ChatMessages, + qwen35_tokenize_fn_fastspeed, + qwen35_tokenize_fn_slowspeed, +) + + +def _render_tool_call_args(arguments: dict) -> str: + """Render tool arguments with the Qwen3.6 JSON scalar semantics.""" + parts = "" + for key, value in arguments.items(): + parts += f"\n" + if isinstance(value, str): + parts += value + else: + parts += json.dumps(value, ensure_ascii=False) + parts += "\n\n" + return parts + + +def qwen36_tokenize_fn_fastspeed( + messages, + tokenizer=None, + tools=None, + add_generation_prompt=False, + add_vision_id=False, + return_labels=True, +): + """Use the Qwen3.5 training renderer with Qwen3.6 tool serialization.""" + return qwen35_tokenize_fn_fastspeed( + messages, + tokenizer=tokenizer, + tools=tools, + add_generation_prompt=add_generation_prompt, + add_vision_id=add_vision_id, + return_labels=return_labels, + _tool_call_args_renderer=_render_tool_call_args, + ) + + +def qwen36_tokenize_fn_slowspeed( + tokenizer, + messages: List[Dict[str, str]], + tools=None, + add_vision_id=False, + **kwargs, +): + """Reference renderer for comparison with the modified HF template.""" + kwargs["preserve_thinking"] = True + return qwen35_tokenize_fn_slowspeed( + tokenizer, + messages, + tools=tools, + add_vision_id=add_vision_id, + **kwargs, + ) + + +class Qwen36ChatMessages(Qwen35ChatMessages): + """Qwen3.6 messages; all behavior except tool serialization is Qwen3.5.""" + + def _tokenize_chat(self, tokenizer, add_vision_id): + return qwen36_tokenize_fn_fastspeed( + self.messages, + tokenizer, + self.tools, + add_vision_id=add_vision_id, + return_labels=True, + ) diff --git a/xtuner/v1/data_proto/templates/__init__.py b/xtuner/v1/data_proto/templates/__init__.py index 7d7a7a78df..ccf7958ed1 100644 --- a/xtuner/v1/data_proto/templates/__init__.py +++ b/xtuner/v1/data_proto/templates/__init__.py @@ -14,6 +14,12 @@ image_context_token="<|image_pad|>", video_context_token="<|video_pad|>", ), + "qwen3.6-vl": HybridChatTemplate( + image_start_token="<|vision_start|>", + image_end_token="<|vision_end|>", + image_context_token="<|image_pad|>", + video_context_token="<|video_pad|>", + ), "intern-s1": HybridChatTemplate( system="<|im_start|>system\n{system}<|im_end|>\n", user="<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n", diff --git a/xtuner/v1/datasets/config.py b/xtuner/v1/datasets/config.py index f9783369f4..b49e1e6e50 100644 --- a/xtuner/v1/datasets/config.py +++ b/xtuner/v1/datasets/config.py @@ -31,6 +31,7 @@ from .preset_pack import PresetPackDataset from .preset_sampler import PresetSampler from .sampler import LengthGroupedSampler, ParallelSampler +from .tokenize_debug import maybe_dump_tokenize_debug_samples from .utils import CachableTokenizeFunction, tokenizer_xxhash from .vlm_jsonl import VLMJsonlDataset @@ -132,6 +133,11 @@ def build_datasets( _tokenize_fn = _tokenize_fn_name.build(tokenizer, tokenizer_hash=tokenizer_hash, anno_name=anno_name) _dataset = _dataset_config.build(_tokenize_fn) if get_rank() == 0: + maybe_dump_tokenize_debug_samples( + dataset=_dataset, + tokenize_fn=_tokenize_fn, + dataset_name=_dataset_config.name, + ) logger.info( f"[Dataset] (Original) {_dataset_config.name}/{os.path.basename(anno_path)}: {len(_dataset)} samples." ) diff --git a/xtuner/v1/datasets/mllm_tokenize_fn/qwen3_vl_tokenize_fn.py b/xtuner/v1/datasets/mllm_tokenize_fn/qwen3_vl_tokenize_fn.py index d9571b4a1d..9a3ae89137 100644 --- a/xtuner/v1/datasets/mllm_tokenize_fn/qwen3_vl_tokenize_fn.py +++ b/xtuner/v1/datasets/mllm_tokenize_fn/qwen3_vl_tokenize_fn.py @@ -17,7 +17,7 @@ import transformers from transformers import AutoProcessor, PreTrainedTokenizer from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize -from xtuner.v1.data_proto.messages import ChatMessages, Qwen35ChatMessages +from xtuner.v1.data_proto.messages import ChatMessages, Qwen35ChatMessages, Qwen36ChatMessages from xtuner.v1.data_proto.templates import CHAT_TEMPLATE_MAP, HybridChatTemplate from xtuner.v1.utils import get_logger @@ -40,6 +40,10 @@ # when multiple datasets share the same processor. # Keyed by processor_path to support different model paths. _PROCESSOR_CACHE: dict[str, tuple] = {} +_QWEN_INTERNAL_MESSAGE_CLASSES = { + "qwen3.5-vl": Qwen35ChatMessages, + "qwen3.6-vl": Qwen36ChatMessages, +} logger = get_logger() @@ -409,8 +413,10 @@ def __init__( ) def calc_num_tokens_pure_text_get_item(self, data_item) -> CacheItem: - if self.chat_template_name == "qwen3.5-vl": - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) else: messages = ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) tokenized = messages.tokenize(self.tokenizer, self.chat_template, add_vision_id=self.add_vision_id) @@ -440,8 +446,10 @@ def _truncated_data_item( def pure_text_get_item(self, data_item: dict) -> QwenVL3DataItem: is_pretrain = False - if self.chat_template_name == "qwen3.5-vl": - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) if len(data_item["messages"]) == 1 and data_item["messages"][0]["role"] == "pretrain": is_pretrain = True else: @@ -497,8 +505,10 @@ def calc_num_tokens_multi_modal_get_item(self, data_item: dict) -> CacheItem: return {"num_tokens": 0, "num_img_tokens": [0]} # type: ignore is_pretrain = False - if self.chat_template_name == "qwen3.5-vl": - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) if len(data_item["messages"]) == 1 and data_item["messages"][0]["role"] == "pretrain": is_pretrain = True tokenized = messages.tokenize(self.tokenizer, self.chat_template, add_vision_id=self.add_vision_id) @@ -554,8 +564,10 @@ def multi_modal_get_item(self, data_item: dict, media_root: str = "") -> QwenVL3 grid_thw_merged = [merged_thw.prod().item() // self.merge_length for merged_thw in grid_thw] # type: ignore is_pretrain = False - if self.chat_template_name == "qwen3.5-vl": - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) if len(data_item["messages"]) == 1 and data_item["messages"][0]["role"] == "pretrain": is_pretrain = True tokenized = messages.tokenize(self.tokenizer, self.chat_template, add_vision_id=self.add_vision_id) @@ -793,9 +805,11 @@ def calc_num_tokens_video_get_item(self, data_item: dict) -> CacheItem: total_sum_media_grid_thw += sum_media_grid_thw is_pretrain = False - if self.chat_template_name == "qwen3.5-vl": + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: replace_video_timestamps_and_num_frame(data_item["messages"], num_image_token_list, timestamps_list) - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) if len(data_item["messages"]) == 1 and data_item["messages"][0]["role"] == "pretrain": is_pretrain = True tokenized = messages.tokenize(self.tokenizer, self.chat_template, add_vision_id=self.add_vision_id) @@ -930,9 +944,11 @@ def video_get_item(self, data_item: dict, media_root: str = "") -> QwenVL3DataIt total_sum_media_grid_thw += sum_media_grid_thw is_pretrain = False - if self.chat_template_name == "qwen3.5-vl": + if self.chat_template_name in _QWEN_INTERNAL_MESSAGE_CLASSES: replace_video_timestamps_and_num_frame(data_item["messages"], num_image_tokens_list, timestamps_list) - messages = Qwen35ChatMessages(messages=data_item["messages"], tools=data_item.get("tools")) + messages = _QWEN_INTERNAL_MESSAGE_CLASSES[self.chat_template_name]( + messages=data_item["messages"], tools=data_item.get("tools") + ) if len(data_item["messages"]) == 1 and data_item["messages"][0]["role"] == "pretrain": is_pretrain = True tokenized = messages.tokenize(self.tokenizer, self.chat_template, add_vision_id=self.add_vision_id) @@ -1020,7 +1036,7 @@ class Qwen3VLTokenizeFnConfig(BaseMLLMTokenizeFnConfig): add_vision_id: bool = True trim_memory_interval: int = 1 - chat_template: Literal["qwen3-vl", "qwen3-vl-rl", "qwen3.5-vl"] = "qwen3-vl" + chat_template: Literal["qwen3-vl", "qwen3-vl-rl", "qwen3.5-vl", "qwen3.6-vl"] = "qwen3-vl" def build( self, tokenizer, tokenizer_hash: str | None = None, anno_name: str = "", **kwargs diff --git a/xtuner/v1/datasets/sft_tokenize_fn/openai.py b/xtuner/v1/datasets/sft_tokenize_fn/openai.py index 5ce9c8cfdf..bff11a049c 100644 --- a/xtuner/v1/datasets/sft_tokenize_fn/openai.py +++ b/xtuner/v1/datasets/sft_tokenize_fn/openai.py @@ -7,7 +7,7 @@ from pydantic import BaseModel, ConfigDict from transformers import PreTrainedTokenizer -from xtuner.v1.data_proto.messages import ChatMessages, Qwen35ChatMessages +from xtuner.v1.data_proto.messages import ChatMessages, Qwen35ChatMessages, Qwen36ChatMessages from xtuner.v1.data_proto.templates import CHAT_TEMPLATE_MAP from xtuner.v1.datasets.data_item import CacheItem, DataItem from xtuner.v1.utils import get_logger @@ -47,6 +47,8 @@ def __call__(self, item: dict | list, **kwargs) -> DataItem | CacheItem: if self.chat_template_name == "qwen3.5-vl": messages = Qwen35ChatMessages(messages=item, tools=tools) + elif self.chat_template_name == "qwen3.6-vl": + messages = Qwen36ChatMessages(messages=item, tools=tools) else: messages = ChatMessages(messages=item, tools=tools) tokenized = messages.tokenize(self.tokenizer, self.chat_template) diff --git a/xtuner/v1/datasets/tokenize_debug.py b/xtuner/v1/datasets/tokenize_debug.py new file mode 100644 index 0000000000..87abb1d8c3 --- /dev/null +++ b/xtuner/v1/datasets/tokenize_debug.py @@ -0,0 +1,282 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Environment-controlled tokenizer snapshots for data-pipeline debugging.""" + +import copy +import hashlib +import json +import os +import random +import re +from pathlib import Path +from typing import Any + +from xtuner.v1.utils import get_logger + + +logger = get_logger() +_DEBUG_SAMPLE_ENV = "XTUNER_TOKENIZE_DEBUG_SAMPLES" +_DEBUG_SAMPLE_SEED = 42 + + +def _true_spans(mask: list[bool]) -> list[tuple[int, int]]: + spans = [] + start = None + for index, enabled in enumerate(mask): + if enabled and start is None: + start = index + elif not enabled and start is not None: + spans.append((start, index)) + start = None + if start is not None: + spans.append((start, len(mask))) + return spans + + +def _decode_ids(tokenizer, token_ids: list[int]) -> str: + try: + return tokenizer.decode( + token_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ) + except TypeError: + return tokenizer.decode(token_ids, skip_special_tokens=False) + + +def _token_pieces(tokenizer, token_ids: list[int]) -> list[str]: + pieces = tokenizer.convert_ids_to_tokens(token_ids) + if isinstance(pieces, str): + return [pieces] + return [str(piece) for piece in pieces] + + +def _summarize_field(value: Any) -> Any: + value_module = type(value).__module__ + if value_module.startswith("torch") and hasattr(value, "shape"): + return { + "kind": "torch.Tensor", + "shape": list(value.shape), + "dtype": str(value.dtype), + } + if value_module.startswith("numpy") and hasattr(value, "shape"): + if len(value.shape) == 0 and hasattr(value, "item"): + return value.item() + return { + "kind": "numpy.ndarray", + "shape": list(value.shape), + "dtype": str(value.dtype), + } + return value + + +def _render_qwen_internal(tokenize_fn, raw_data: dict) -> tuple[str, list[bool]]: + from xtuner.v1.data_proto.messages.qwen35_chat import qwen35_tokenize_fn_fastspeed + from xtuner.v1.data_proto.messages.qwen36_chat import qwen36_tokenize_fn_fastspeed + + renderers = { + "qwen3.5-vl": qwen35_tokenize_fn_fastspeed, + "qwen3.6-vl": qwen36_tokenize_fn_fastspeed, + } + + messages = copy.deepcopy(raw_data["messages"]) + chat_template = tokenize_fn.chat_template + if chat_template.default_system is not None: + if messages[0]["role"] == "system": + messages[0]["content"] = chat_template.default_system + else: + messages.insert(0, {"role": "system", "content": chat_template.default_system}) + + return renderers[tokenize_fn.chat_template_name]( + messages, + tools=copy.deepcopy(raw_data.get("tools")), + add_vision_id=tokenize_fn.add_vision_id, + return_labels=False, + ) + + +def _build_record( + *, + tokenizer, + raw_data: dict, + rendered_text: str, + character_loss_mask: list[bool], + tokenized_data: dict, + dataset_name: str, + dataset_path: str, + sample_index: int, +) -> dict: + input_ids = [int(token_id) for token_id in tokenized_data["input_ids"]] + labels = [int(label) for label in tokenized_data["labels"]] + if len(input_ids) != len(labels): + raise ValueError(f"input_ids length {len(input_ids)} != labels length {len(labels)}") + + loss_mask = [label != -100 for label in labels] + loss_spans = [] + for start, end in _true_spans(loss_mask): + token_ids = labels[start:end] + loss_spans.append( + { + "start": start, + "end": end, + "token_ids": token_ids, + "token_pieces": _token_pieces(tokenizer, token_ids), + "decoded_text": _decode_ids(tokenizer, token_ids), + } + ) + + return { + "source": { + "dataset_name": dataset_name, + "dataset_path": dataset_path, + "sample_index": sample_index, + }, + "raw_data": raw_data, + "rendered": { + "text": rendered_text, + "character_count": len(rendered_text), + "loss_character_count": sum(character_loss_mask), + "loss_character_spans": [ + { + "start": start, + "end": end, + "text": rendered_text[start:end], + } + for start, end in _true_spans(character_loss_mask) + ], + }, + "tokenized": { + "input_ids": input_ids, + "labels": labels, + "decoded_text": _decode_ids(tokenizer, input_ids), + "token_count": len(input_ids), + "other_fields": { + key: _summarize_field(value) + for key, value in tokenized_data.items() + if key not in {"input_ids", "labels"} + }, + }, + "loss": { + "token_count": sum(loss_mask), + "masked_token_count": len(labels) - sum(loss_mask), + "token_spans": loss_spans, + }, + } + + +def _write_json(path: Path, data: dict) -> None: + tmp_path = path.with_suffix(path.suffix + ".tmp") + with tmp_path.open("w", encoding="utf-8") as file: + json.dump(data, file, ensure_ascii=False, indent=2) + file.write("\n") + os.replace(tmp_path, path) + + +def _safe_path_component(value: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._") or "dataset" + + +def _read_raw_data(dataset, sample_index: int) -> dict: + with open(dataset.path) as file: + file.seek(dataset.offsets[sample_index]) + return json.loads(file.readline()) + + +def _sample_indices(dataset_size: int, sample_count: int) -> list[int]: + count = min(sample_count, dataset_size) + return sorted(random.Random(_DEBUG_SAMPLE_SEED).sample(range(dataset_size), count)) + + +def maybe_dump_tokenize_debug_samples(*, dataset, tokenize_fn, dataset_name: str) -> None: + """Dump the first N records when XTUNER_TOKENIZE_DEBUG_SAMPLES is set.""" + raw_sample_count = os.environ.get(_DEBUG_SAMPLE_ENV, "0") + try: + sample_count = int(raw_sample_count) + except ValueError as error: + raise ValueError(f"{_DEBUG_SAMPLE_ENV} must be an integer, got {raw_sample_count!r}") from error + if sample_count <= 0: + return + + chat_template_name = getattr(tokenize_fn, "chat_template_name", None) + if chat_template_name not in {"qwen3.5-vl", "qwen3.6-vl"}: + logger.warning( + f"[Dataset] Skip tokenizer debug snapshots for [{dataset_name}]: " + "only chat_template='qwen3.5-vl' or 'qwen3.6-vl' is supported." + ) + return + + dataset_path = str(dataset.path) + path_hash = hashlib.sha256(dataset_path.encode()).hexdigest()[:8] + output_root = Path(os.environ.get("WORK_DIR", os.getcwd())) / "tokenize_debug_samples" + output_dir = ( + output_root + / _safe_path_component(chat_template_name) + / f"{_safe_path_component(dataset_name)}__{_safe_path_component(Path(dataset_path).name)}__{path_hash}" + ) + output_dir.mkdir(parents=True, exist_ok=True) + + previous_state = tokenize_fn.state + tokenize_fn.set_state("runtime") + generated_files = [] + sample_indices = _sample_indices(len(dataset), sample_count) + try: + for sample_index in sample_indices: + output_path = output_dir / f"sample_{sample_index:06d}.json" + raw_data = _read_raw_data(dataset, sample_index) + try: + rendered_text, character_loss_mask = _render_qwen_internal(tokenize_fn, raw_data) + tokenized_data = tokenize_fn( + copy.deepcopy(raw_data), + media_root=getattr(dataset, "media_root", ""), + ) + record = _build_record( + tokenizer=tokenize_fn.tokenizer, + raw_data=raw_data, + rendered_text=rendered_text, + character_loss_mask=character_loss_mask, + tokenized_data=tokenized_data, + dataset_name=dataset_name, + dataset_path=dataset_path, + sample_index=sample_index, + ) + except Exception as error: + record = { + "source": { + "dataset_name": dataset_name, + "dataset_path": dataset_path, + "sample_index": sample_index, + }, + "raw_data": raw_data, + "error": { + "type": type(error).__name__, + "message": str(error), + }, + } + logger.exception( + f"Failed to dump tokenizer debug sample {sample_index} " + f"from [{dataset_name}]{dataset_path}" + ) + + _write_json(output_path, record) + generated_files.append(output_path.name) + finally: + tokenize_fn.set_state(previous_state) + + _write_json( + output_dir / "manifest.json", + { + "environment_variable": _DEBUG_SAMPLE_ENV, + "sampling": "random_without_replacement", + "random_seed": _DEBUG_SAMPLE_SEED, + "dataset_name": dataset_name, + "dataset_path": dataset_path, + "chat_template_name": chat_template_name, + "requested_sample_count": sample_count, + "generated_sample_count": len(generated_files), + "sample_indices": sample_indices, + "files": generated_files, + }, + ) + logger.info( + f"[Dataset] Dumped {len(generated_files)} tokenizer debug samples " + f"from [{dataset_name}]{dataset_path} to {output_dir}." + )